# Flatness-Calibrated Constant-Step SGD

- ID: 1007
- Canonical URL: https://synthcore.org/idea/1007/flatness-calibrated-constant-step-sgd
- API JSON: https://synthcore.org/api/idea/1007.json
- API Markdown: https://synthcore.org/api/idea/1007.md
- Verification status: mechanism_failed
- Source: [arXiv:2607.16384](https://arxiv.org/abs/2607.16384)
- Category: optimization
- Solves: stability, sample-efficiency, accuracy
- ML areas: optimizer, training-dynamics, transformer
- Math tags: stochastic-processes, dynamical-systems, optimization, probability, geometry
- Ratings: usefulness 7/10; difficulty 5/10; novelty 7/10

## Idea description

Replace a globally chosen constant learning rate with a blockwise rate calibrated to the local flatness exponent of the objective. If the local Hessian decays like \(\|x-x_\star\|^{m-2}\), choose the rate so that the predicted stationary parameter radius \(\alpha^{1/m}\) matches a prescribed exploration or optimization radius, rather than incorrectly using the quadratic rule \(\sqrt{\alpha}\).

## Mathematical statement

The paper assumes local flatness exponent \(m\geq 2\) through \(\nabla^{2}H(x)\asymp\lVert x-x_{\star}\rVert^{m-2}I_d\) as \(x\to x_{\star}\), where \(H\) is the objective, \(x_\star\) is a minimizer, \(d\) is parameter dimension, and \(I_d\) is the identity matrix. For constant-step SGD, the invariant law concentrates at scale \(\alpha^{1/m}\), where \(\alpha\) is the learning rate; for \(m=2\) this reduces to the familiar \(\sqrt{\alpha}\) scale. The paper also gives a geometric-convergence factor \(1-c\alpha^{m-1}\), where \(c>0\) is a problem-dependent constant, implying a relaxation time of order \(\alpha^{-(m-1)}\). Near the minimizer, the rescaled stationary dynamics converge to \(dY_t=-h_0(Y_t)\,dt+\Sigma^{1/2}dB_t\), where \(Y_t=\alpha^{-1/m}(X_t-x_\star)\), \(h_0\) is the limiting drift, \(\Sigma\) is the asymptotic covariance of gradient noise, and \(B_t\) is standard Brownian motion. The adapted optimizer uses the radius law \(r\approx \widehat C\alpha^{1/m}\sqrt[m]{\operatorname{tr}(\widehat\Sigma)}\) to set \(\alpha\approx (r/(\widehat C\sqrt[m]{\operatorname{tr}(\widehat\Sigma)}))^m\).

## Key formulas

- $$\nabla^{2}H(x)\asymp\left\lVert x-x_{\star}\right\rVert^{m-2}I_{d},\qquad x\to x_{\star},$$
- $$\text{stationary error scale}\asymp \alpha^{1/m},\qquad \text{contraction factor}=1-c\alpha^{m-1},$$
- $$dY_t=-h_0(Y_t)\,dt+\Sigma^{1/2}\,dB_t,\qquad Y_t=\alpha^{-1/m}(X_t-x_\star),$$
- $$\alpha_{\mathrm{target}}=\left(\frac{r_{\mathrm{target}}}{\widehat C\,\sqrt[m]{\operatorname{tr}(\widehat\Sigma)}}\right)^m$$

## Implementation notes

1. Integration point: modify the learning-rate line of SGD or AdamW, preferably per parameter block such as each Transformer layer, attention projection, or MLP block. Keep the ordinary minibatch gradient update, but periodically replace the block learning rate by a value inferred from its local flatness exponent. Do not compute a full Hessian on every step. 2. Pseudocode: initialize block rates alpha_b and target radii r_b. Every K steps, collect L minibatch gradients g_b^(1),...,g_b^(L) and the current block parameters x_b. Estimate gradient-noise variance with v_b = (1/L) sum_t ||g_b^(t)-gbar_b||^2, where gbar_b is the sample mean. Estimate curvature at two radii using random normalized directions u: q(rho) = average of <g(x_b+rho u)-g(x_b),u>/rho and q(2rho) analogously. Set m_hat_b = 2 + log((q(2rho)+epsilon)/(q(rho)+epsilon))/log(2), clipped to [2,6]. Run a short pilot with the current alpha_b and estimate C_hat_b as the empirical stationary RMS parameter radius divided by alpha_b^(1/m_hat_b) times the inverse noise factor. Then set alpha_b = clip((r_b/(C_hat_b*sqrt[m_hat_b](v_b+epsilon)))^m_hat_b, alpha_min, alpha_max). During normal training apply x_b <- x_b - alpha_b*g_b. 3. The paper supplies the exponent-to-radius law and the mixing-rate dependence; covariance, curvature proxies, C_hat, and clipping bounds are empirical estimates. Use exponential moving averages for m_hat_b and C_hat_b, and freeze adaptation after warmup to prevent oscillatory rates. 4. First experiment: train a small ResNet-18 on CIFAR-10 and a 6-layer Transformer on WikiText-2, comparing standard cosine-decay SGD, fixed tuned SGD, and this method at equal optimizer-step and FLOP budgets. Measure training loss, validation accuracy or perplexity, stationary RMS parameter displacement during a constant-rate phase, and recovery after perturbations. Success means the measured radius tracks the target within 20%, the method avoids loss-floor instability in flat blocks, and it reaches the baseline validation loss with fewer updates or a larger stable learning rate.

## Verification

- Status: mechanism_failed
- Mechanism evidence: yes
- Mechanism confirmed: no
- Verdict: Built a reproducible 1D constant-step SGD simulator for H(x)=|x|^m/m, curvature-based exponent estimation, empirical noise calibration, and comparison with the quadratic learning-rate rule. The stationary-radius scaling and noise scaling were broadly observed, but the predicted relaxation scaling failed for m>2 and calibrated target-radius tracking failed badly at m=5, so the full mechanism was not confirmed.

### Mechanism check

- Verdict: Built a reproducible 1D constant-step SGD simulator for H(x)=|x|^m/m, curvature-based exponent estimation, empirical noise calibration, and comparison with the quadratic learning-rate rule. The stationary-radius scaling and noise scaling were broadly observed, but the predicted relaxation scaling failed for m>2 and calibrated target-radius tracking failed badly at m=5, so the full mechanism was not confirmed.
- Confidence: 8/10
- Limitations: Only a 1D toy objective was tested; no CIFAR-10 ResNet, Transformer, blockwise minibatch implementation, GPU run, validation metric, FLOP comparison, or perturbation recovery was included. The higher-m relaxation diagnostic also had limited effective samples.

## Artifacts

- [flat_sgd_experiment.py](https://synthcore.org/code/414/flat_sgd_experiment.py)
- [report.md](https://synthcore.org/code/414/report.md)
- [results.json](https://synthcore.org/code/414/results.json)
- [results_full.json](https://synthcore.org/code/414/results_full.json)
- [Download all files as ZIP](https://synthcore.org/download/414)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.