# Critical-Batch Momentum Scaling

- ID: 3080
- Canonical URL: https://synthcore.org/idea/3080/critical-batch-momentum-scaling
- API JSON: https://synthcore.org/api/idea/3080.json
- API Markdown: https://synthcore.org/api/idea/3080.md
- Verification status: unverified
- Source: [arXiv:2609.02728](https://arxiv.org/abs/2609.02728)
- Category: optimization
- Solves: speedup, stability, sample-efficiency
- ML areas: optimizer, scheduler, training-dynamics, transformer
- Math tags: dynamical-systems, optimization, stochastic-processes
- Ratings: usefulness 7/10; difficulty 5/10; novelty 6/10

## Idea description

Replace a fixed momentum and learning-rate schedule with a batch-aware stability controller derived from the paper's critical-learning-rate scalings. Polyak learning rates should scale approximately with B(1-rho), whereas Nesterov learning rates can scale as B^beta(1-rho) until reaching the base stability ceiling; this may allow larger batches without crossing the instability boundary.

## Mathematical statement

The paper defines the one-pass update horizon as T=D/B, where D is the total data budget and B is the batch size. It characterizes the critical learning rate, the largest learning rate preserving stable risk dynamics, as eta_sgd^crit approximately proportional to 1, eta_polyak^crit approximately proportional to min{1,B(1-rho)}, and eta_nesterov^crit approximately proportional to min{1,B^beta(1-rho)}. Here rho is the momentum coefficient in [0,1), beta>1 is the capacity exponent of the power-law kernel model, and the notation approximately proportional means equality up to problem-dependent multiplicative constants. Transfer these relations as relative scaling rules around a calibrated reference batch B0: eta_P(B,rho)=eta0 min{1,cP(B/B0)((1-rho)/(1-rho0))} and eta_N(B,rho)=eta0 min{1,cN(B/B0)^beta_hat((1-rho)/(1-rho0))}. Here eta0 and rho0 are the reference learning rate and momentum, cP and cN are safety factors, and beta_hat is a tuned or empirically estimated exponent.

## Key formulas

- $$T=\frac{D}{B}$$
- $$\eta_{\mathrm{sgd}}^{\mathrm{crit}}\asymp 1,\qquad \eta_{\mathrm{polyak}}^{\mathrm{crit}}\asymp\min\{1,B(1-\rho)\},\qquad \eta_{\mathrm{nesterov}}^{\mathrm{crit}}\asymp\min\{1,B^{\beta}(1-\rho)\}$$
- $$\eta_{\mathrm{P}}(B,\rho)=\eta_0\min\left\{1,c_{\mathrm{P}}\frac{B}{B_0}\frac{1-\rho}{1-\rho_0}\right\},\qquad \eta_{\mathrm{N}}(B,\rho)=\eta_0\min\left\{1,c_{\mathrm{N}}\left(\frac{B}{B_0}\right)^{\hat\beta}\frac{1-\rho}{1-\rho_0}\right\}$$

## Implementation notes

Integrate the method into the optimizer learning-rate scheduler immediately before each parameter update. Inputs are the current batch size B, reference batch B0, base learning rate eta0, reference momentum rho0, current momentum rho, optimizer type, and safety factors cP and cN. For Polyak momentum, maintain a velocity v and use v_t=rho v_(t-1)+g_t followed by theta_(t+1)=theta_t-eta_P v_t. For Nesterov, evaluate the gradient at a look-ahead point theta_t-eta_N rho v_(t-1), then use v_t=rho v_(t-1)+g_t and theta_(t+1)=theta_t-eta_N v_t. Compute the scheduled rate with scale=min(1,cP*(B/B0)*((1-rho)/(1-rho0))) for Polyak, or scale=min(1,cN*(B/B0)^beta_hat*((1-rho)/(1-rho0))) for Nesterov, and set eta=eta0*scale. Start with cP=cN=0.5. Monitor loss, gradient norm, update norm, and non-finite values over windows of 50 to 200 steps. If loss or gradient norms grow persistently, halve the relevant safety factor; if the run remains stable for several windows, increase it by 10 percent. Estimate beta_hat from short Nesterov sweeps at batch ratios q in {2,4,8}: find the largest non-divergent eta at each ratio, then fit the slope of log(eta_crit) against log(q), clipped to [1,2]. The first experiment should use a small Transformer on WikiText-103 or C4 with a fixed token budget, comparing AdamW, fixed-rate SGD momentum, fixed-rate Nesterov, and this controller at B/B0 in {1,2,4,8,16}. Report maximum stable learning rate, validation loss at equal token counts, final one-pass loss, and wall-clock throughput. Success means stable operation at a larger batch with equal or lower iso-token perplexity and faster loss reduction per wall-clock hour.

## Disclaimer

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