Block-Based Learning Rate Optimization

• Parameter block-based learning rates are an adaptive optimization strategy that assigns distinct rates to predefined parameter groups based on intra-block curvature.
• They utilize methods such as block-diagonal matrix adaptation, spectrum clipping, and sharpness-aware scheduling to achieve improved training convergence and nearly 2× speedup in large model pre-training.
• These techniques offer practical benefits including computational efficiency, ease of integration with optimizers like AdamW and SGD, and potential for automatic tuning through bandit and learning-rate-free approaches.

Parameter block-based learning rates are adaptive strategies in stochastic optimization that assign distinct learning rate values to predefined groups ("blocks") of model parameters, such as layers or modules, rather than scaling globally or on a per-parameter basis. Blockwise regimes are designed to exploit structural characteristics of modern neural architectures, utilize intra-block curvature information, and enhance both convergence and generalization performance. Technical advancements in this area encompass block-diagonal matrix adaptation, sharpness-guided blockwise scheduling, learning-rate-free blockwise optimization, and dynamic scheduling rooted in lightweight curvature or gradient statistics.

1. Structural Motivation and Block Partitioning

The impetus for parameter block-based learning rates arises from the limitations of scalar or diagonal (per-parameter) learning rate adaptation in high-dimensional models. Classical diagonal adaptation (e.g., Adagrad, RMSprop, Adam) processes each parameter independently, ignoring key cross-coordinate correlations. Conversely, full-matrix adaptation fully models these correlations but is computationally prohibitive in deep networks.

Block-diagonal adaptation partitions the parameters into $r$ groups, typically by structural elements (layers, neurons, attention heads), exploiting architectural modularity. Each block is assigned a separate learning rate, which can be estimated dynamically. This scheme captures inter-parameter curvature within blocks and approximates the loss surface more faithfully than diagonal methods, with superior cost-efficiency compared to full-matrix schemes (Yun et al., 2019).

Method	Granularity	Curvature Modeling
Full-matrix	All parameters	Maximal
Diagonal	Per-parameter	Minimal
Block-diagonal	Groups/blocks	Intermediate

2. Blockwise Matrix Adaptation and Spectrum Clipping

A rigorous framework for blockwise learning rates is block-diagonal matrix adaptation (Yun et al., 2019). Here, a blockwise second-moment matrix $V_t$ is constructed, where per-block gradients $g_t^{(j)}$ (for block $j$) yield blockwise matrices via $g_t^{(j)} [g_t^{(j)}]^T$. The update rule is:

$$x_{t+1} = x_t - \alpha_t \left(V_t^{1/2} + \delta I\right)^{-1} m_t$$

where $m_t$ is a momentum term and $\delta$ is for numerical stability. Each block's inverse square-root is tractable. Exact diagonal adaptation is recovered by setting each group to a singleton.

To mitigate overfitting from aggressive initial learning rates, a spectrum-clipping scheme is proposed: Singular value decomposition within blocks, followed by clipping eigenvalues, interpolates the preconditioner from adaptive matrix towards a scaled identity. This emulates vanilla SGD's generalization and smooths convergence, particularly in later training stages.

3. Sharpness-Disparity-Driven Blockwise Scheduling

Recent investigations (Wang et al., 26 Feb 2025) reveal that blocks within transformer architectures manifest persistent "sharpness disparities." Sharpness, quantified as average diagonal Hessian or squared gradient norm per block parameter,

$$S([b]) = \frac{B \cdot \left\| \nabla_{[b]} L \right\|^2}{\#(b)}$$

can vary substantially between embedding, feedforward, attention, and normalization layers. High-sharpness blocks are governed by stability constraints; low-sharpness blocks often exhibit under-training unless their learning rates are amplified.

Blockwise learning rate scheduling sets the optimizer's base learning rate for the highest-sharpness block, while boosting others proportionally:

$$\eta_b = r(b) \cdot \eta_{\text{base}}, \quad r(b) \propto \frac{S(\text{ref})}{S(b)}$$

Empirically, Blockwise LR delivers <i>lower terminal losses and nearly $2\times$ training speedup</i> in LLM pre-training with negligible overhead. Integration with AdamW and memory-efficient Adam-mini demonstrates joint speed and memory savings (Wang et al., 26 Feb 2025).

4. Blockwise Learning-Rate-Free and Bandit-Based Methods

Learning-rate-free optimization extends naturally to parameter block scenarios (Suh et al., 6 Jan 2024). By interpreting adaptive updates as steepest descent in a parameter-scaled network, block-specific scaling factors effectively modulate learning rates, obviating manual tuning. Algorithms such as Parameter-Scaled Stochastic Polyak Step-size (PS-SPS) and PS-D-Adapt SGD apply Polyak-style or D-Adaptation step-size rules within each block, maintaining competitive performance with tuned baselines.

Separately, Lipschitz bandit algorithms (Priyanka et al., 15 Sep 2024) can optimize per-block learning rates by formulating each block's learning rate search as a continuous-arm bandit problem. Adaptive discretization (e.g., via the Zooming algorithm) efficiently finds optimal learning rates for different blocks with minimal evaluations, especially effective in models where hyperparameter tuning is a bottleneck.

Approach	Block-wise Adaptivity	Tuning Complexity
PS-SPS/PS-DA-SGD	Yes	Minimal
Zooming Bandit	Yes	Automatic

5. Dynamic and Heuristic Blockwise Mechanisms

Binary Forward Exploration (BFE) and its adaptive extension AdaBFE (Cao, 2022) illustrate forward-looking search strategies for per-parameter or per-block learning rate selection. At each step, a binary search over the loss surface detects whether current rates are too aggressive or conservative, adjusting by halving or doubling until a locally "safe" rate is found. AdaBFE uses per-parameter angular deviations of gradients as a blockwise safety criterion. These mechanisms often converge faster than classic SGD or momentum in early training phases.

Other heuristic adjustments, such as learning rate perturbation (Liu et al., 2022), inject Gaussian noise at the per-parameter or block level, promoting convergence toward flatter minima. This stochastic plugin can be annexed to any learning rate schedule, introducing block-wise effective diversity that empirically enhances generalization.

6. Comparison with Traditional and Universal Regimes

Traditional block-based learning rate schemes require manual assignment of rates to parameter groups, sometimes guided by architectural intuition (layer-wise scaling) or empirical trial-and-error. Explainable learning rate regimes based on stochastic second-order cues (Yang, 19 Aug 2025) dispense with grouping altogether: they compute a principled global rate, often similar in effect to blockwise schemes if the loss surface's curvature is dominated by particular blocks. The adjustment is automatic:

$$\alpha_t = \frac{1}{\sqrt{|S_H|} \frac{\|\hat{s}_t\|^2}{\langle \hat{y}_t, \hat{s}_t \rangle + \|\hat{s}_t\|^2}}$$

where $(\hat{s}_t, \hat{y}_t)$ implicitly capture curvature from aggregate gradient changes.

While dedicated blockwise approaches exploit architecture-specific information, universal regimes offer self-tuning learning rates based on stochastic gradient dynamics, often achieving robustness and scalability without grouping.

7. Practical Considerations and Impact

Blockwise learning rate adaptation provides several practical advantages:

• Computational efficiency: In block-diagonal matrix adaptation, the cost scales as $\mathcal{O}(rB^2)$ per iteration (with $B$ block size), significantly less than for full-matrix methods.
• Improved convergence: Empirical studies demonstrate reduced oscillations, more stable loss curves, and better out-of-sample test error compared to diagonal-only or global scalar approaches.
• Generalization: Spectrum clipping and sharpness-aware amplification mitigate overfitting induced by aggressive initial learning rates while accelerating slow blocks.
• Transferability: Dynamic RL-based scheduling and autonomous controllers can extend to blockwise schemes, though careful preprocessing and decision granularity must address subtle differences in loss scales and dynamics across blocks.
• Implementation: Blockwise regimes are compatible with common optimizers (SGD, Adam, AdamW) and training frameworks, often requiring only modest additional memory and code complexity.

The ongoing exploration of sharpness disparity, bandit-driven tuning, and learning-rate-free scaling continues to drive innovation in parameter block-based learning rate methods, particularly for LLMs and architectures with pronounced modularity.

Future directions include further automating block definition, integrating blockwise regularization, and developing universal, explainable learning rate controllers that adapt in real-time to evolving architectural and optimization dynamics.