LARS: Layer-wise Adaptive Rate Scaling

• LARS is a layer-wise adaptive rate scaling algorithm that adjusts learning rates based on individual layer weight and gradient norms.
• It enhances training stability and generalization in large-batch deep learning by ensuring updates are proportional to parameter magnitudes.
• Empirical results show LARS significantly improves performance in distributed training and convolutional architectures on datasets like ImageNet and MNIST.

Layer-wise Adaptive Rate Scaling (LARS) is an optimization algorithm introduced to address the generalization and stability challenges of large-batch training in deep neural networks. By adaptively scaling the step size for each layer based on the relative magnitudes of the weights and their gradients, LARS enables the use of batch sizes and learning rates that would otherwise induce divergence or significant drops in test accuracy when using conventional Stochastic Gradient Descent (SGD). The algorithm's utility is demonstrated across convolutional architectures, with notable advances in scaling up training bandwidth on distributed systems and mitigating optimization barriers at very large batch scales.

1. Motivation and Background

Classic SGD relies on a scalar global learning rate, which poses difficulties in deep networks where the magnitudes of weights and gradients can differ by orders of magnitude between layers. Empirically, as the mini-batch size $B$ increases, linear scaling of the learning rate $\gamma$ (i.e., $\gamma \sim B$) initially preserves or sometimes accelerates convergence, but above a threshold batch size (commonly $B>8\,\mathrm{K}$ for vision models), it results in sharp drops in accuracy or even divergence. This failure is partially attributed to the mismatch in per-layer update-to-parameter ratios: layers with small weights or gradients are disproportionately affected, inhibiting balanced optimization progress (You et al., 2017).

LARS was introduced to maintain a consistent ratio between the norm of a layer's update and the norm of its parameters, ensuring each layer receives an appropriately scaled step. This preserves stability and generalization at unprecedented batch sizes by controlling both over- and under-shooting across the parameter manifold (You et al., 2017, Chowdhury et al., 2021, You et al., 2019).

2. Core Algorithmic Formulation

At the core of LARS is a per-layer "trust ratio" that modulates the base learning rate. For parameters $w_\ell$ of layer $\ell$ with gradient $\nabla L(w_\ell)$, the LARS update is given by

$$w^{t+1}_\ell = w^t_\ell - \gamma\,\eta\,\frac{\|w^t_\ell\|_2}{\|\nabla L(w^t_\ell)\|_2 + \beta\|w^t_\ell\|_2 + \epsilon}\;\nabla L(w^t_\ell)$$

where $\gamma$ is the global learning rate, $\eta$ is the trust coefficient, $\beta$ is weight decay, and $\epsilon$ is a small constant for numerical stability (Chowdhury et al., 2021, You et al., 2017). The layer-wise factor

$$\lambda_\ell = \eta\, \frac{\|w_\ell\|_2}{\|\nabla L(w_\ell)\|_2 + \beta\|w_\ell\|_2 + \epsilon}$$

ensures the magnitude of each update $\|\Delta w_\ell\|_2$ is proportional to $\|w_\ell\|_2$.

When $\beta=0,\, \epsilon \to 0^+$, this becomes $$\lambda_\ell = \eta \|w_\ell\|_2/\|\nabla L(w_\ell)\|_2$$, closely matching the original LARS formulation (You et al., 2017). In practice, a momentum buffer and explicit weight decay are included as in standard SGD, with the LARS scaling applied to the (regularized) gradient.

3. Hyperparameterization and Scheduling

LARS introduces minimal hyperparameters beyond those in SGD:

• Base learning rate ($\gamma$): Tuned according to batch size scaling rules. For example, $\gamma = 0.01$ for MNIST (Chowdhury et al., 2021), $0.1$ for ResNet-50/ImageNet with $B=8\,\mathrm{K}$ (You et al., 2017).
• Trust coefficient ($\eta$): Governs the aggressiveness per layer (typical $\eta=0.001$ to $0.01$). Smaller values yield conservative updates; grid search in $[10^{-4},10^{-2}]$ is effective (Chowdhury et al., 2021, You et al., 2017).
• Weight decay ($\beta$): Default $10^{-4}$ or $5 \times 10^{-4}$ as L2 regularization.
• $\epsilon$: Small value, e.g., $10^{-6}$ or $10^{-8}$, for stability (Chowdhury et al., 2021, You et al., 2017).
• Momentum: Commonly set to $0.9$.

To avoid initial divergence with large batch sizes, a warm-up phase is often introduced, linearly increasing $\gamma$ from zero to its target over 5–10 epochs (You et al., 2017, Chowdhury et al., 2021). The most recent developments critique linear warm-up for causing stagnation in sharp minima and propose time-varying scaling schedules (e.g., inverse-sigmoid via TVLARS) for improved early exploration and generalization (Do et al., 2023).

4. Empirical Performance and Scaling Results

LARS has been empirically validated in both single-node and distributed settings across a variety of deep architectures. Notable findings include:

• ResNet-50 on ImageNet: Maintains $\sim$0.7% drop (73.0% $\to$ 72.3% top-1) when moving from $B=256$ to $B=32\,\mathrm{K}$, an order of magnitude upscaling without architecture- or batchnorm-specific tuning (You et al., 2017). Comparable results are seen with similar recipes in distributed systems (Chowdhury et al., 2021).
• AlexNet(-BN): Substantial test accuracy improvements with LARS at $B=8\,\mathrm{K}$ (44.8% with linear LR scaling versus 58.8% baseline with SGD; LARS recovers near-baseline accuracy) (You et al., 2017).
• SystemML Case Study: On MNIST, LARS shows marked resilience to increasing batch size: SGD and LARS both reach $\sim99\%$ at $B=1024$, but for $B \geq 16000$, LARS maintains $80{-}90\%$ while SGD drops to $60{-}70\%$; at $B = 32\,\mathrm{K}$, LARS outperforms SGD by 10–15 percentage points in test accuracy (Chowdhury et al., 2021).

Convergence in large-batch regimes remains consistent and wall-clock improvements are substantial, with multi-GPU and parameter-server systems benefiting from fewer required synchronizations (You et al., 2017).

5. Relationship to and Extensions Beyond SGD

LARS generalizes SGD by introducing layerwise normalization, effectively addressing the pathological scaling of updates encountered during large-batch optimization. Standard SGD assigns a global step size, while LARS adaptively matches update magnitudes to parameter norms per layer. The principle contrasts with algorithms such as back-matching propagation-derived layerwise adaptive rates, which employ least-squares backward matching and include normalization based on higher-order statistics and sharing factors, although both techniques deliver layer-specific adaptation (Zhang et al., 2018).

The introduction of the "trust ratio" concept in LARS directly inspired subsequent optimizers such as LAMB and its clipped variant LAMBC, which further stabilize training under extremely large batch or transformer regimes. LAMB, for instance, matches the per-layer update norm to the parameter norm but applies Adam-style moment normalization, targeting architectures like BERT where LARS may be unstable (You et al., 2019, Fong et al., 2020).

6. Limitations, Instability, and Mitigation

Despite its efficacy, LARS can yield unstable or extreme trust ratios, especially in late-stage training or when $\|\nabla L(w_\ell)\|_2$ becomes very small. This results in spurious learning rate amplification and potential divergence (Fong et al., 2020). Empirical studies observe catastrophic drops in performance or loss oscillations for batch sizes exceeding 2\,000 without suppressed trust ratios.

Addressing this, trust ratio clipping mechanisms set lower ($\tau$) and upper ($\mu$) bounds on the ratio $$\rho_{\ell,t} = \|w_t^{(\ell)}\|_2 / \|\nabla L(w_t^{(\ell)})\|_2$$, with recommended $\mu\sim1$ for maximal accuracy and stability (Fong et al., 2020). TVLARS replaces linear warm-up with a configurable sigmoid-based schedule to enable better escape from sharp minima early in training, reporting up to 2–3 percentage points improvement in classification and up to 10 percentage points in self-supervised regimes (Do et al., 2023).

On transformer models and tasks with nonuniform parameter landscapes, LARS often fails to preserve convergence (e.g., BERT pretraining for batches $>16\,\mathrm{K}$), motivating the shift toward LAMB/LAMBC optimizers (You et al., 2019, Fong et al., 2020).

7. Implementation and Practical Considerations

LARS is computationally lightweight: it requires only per-layer $L_2$ norm computations (readily available as framework primitives) and simple algebraic scaling of gradients. In distributed or model-parallel infrastructures (e.g., SystemML on Spark), lambda computation can be performed post-aggregation and scale factors applied in-place to avoid memory duplication (Chowdhury et al., 2021). Special care is needed for parameter groups such as batch-norm statistics or very small biases, which can be exempted from LARS scaling to prevent noisy or degenerate lambda values (You et al., 2017).

Practical steps:

1. Compute $\|w_\ell\|$ and $\|\nabla L(w_\ell)\|$ per layer.
2. Calculate $\lambda_\ell$ (or apply clipping as needed).
3. Scale and update weights.
4. Apply momentum and weight decay as appropriate.

Initialization strategies, learning-rate scheduling, synchronized batch normalization, and trust ratio clamping are recommended for robust behavior with very large batch sizes and distributed training configurations (You et al., 2017, Do et al., 2023).

---

References:

• "Large Batch Training of Convolutional Networks" (You et al., 2017)
• "Evaluating Deep Learning in SystemML using Layer-wise Adaptive Rate Scaling(LARS) Optimizer" (Chowdhury et al., 2021)
• "Large Batch Optimization for Deep Learning: Training BERT in 76 minutes" (You et al., 2019)
• "Improving Layer-wise Adaptive Rate Methods using Trust Ratio Clipping" (Fong et al., 2020)
• "Revisiting LARS for Large Batch Training Generalization of Neural Networks" (Do et al., 2023)
• "Train Feedfoward Neural Network with Layer-wise Adaptive Rate via Approximating Back-matching Propagation" (Zhang et al., 2018)