---
title: Split Weight Decay in Deep Learning
url: https://www.emergentmind.com/topics/split-weight-decay
type: topic
---

# Split Weight Decay in Deep Learning

Split weight decay encompasses a family of techniques for regularization in deep learning where the traditional isotropic (uniform) decay penalty is decomposed, modulated, or selectively applied based on parameter geometry, module structure, or step-by-step dynamics. Unlike standard approaches that penalize all weights equally in $\ell_2$ norm, split weight decay exploits architectural symmetries and empirical evidence that meaningful regularization should often occur along specific subspaces (e.g., radial direction of weights), or be tailored module-wise, or switched adaptively. This paradigm underlies advances such as orthogonally decoupled optimizers, layerwise spectral-adaptive schemes, and geometry-aware decay for fine-tuning, leading to demonstrably improved generalization, hyperparameter robustness, and sometimes sparser or more robust solutions in both vision and language models.

## 1. Geometric Motivation: Decomposition of Weight Dynamics

Recent analysis shows that standard decoupled weight decay (such as AdamW) induces a "Radial Tug-of-War" in adaptive optimization. For parameter vector $w \in \mathbb{R}^d$ and gradient $g = \nabla L(w)$, gradient steps often increase the norm $\|w\|_2$ (expanding model capacity), while decay attempts to shrink $\|w\|_2$. This interaction injects high-variance radial oscillations into adaptive optimizer statistics, particularly corrupting estimates of second moments and thus impairing feature learning in tangential (orthogonal) directions [2602.05136].

To resolve this, the weight-update vector is decomposed into radial (parallel to $w$) and tangential (orthogonal) components:
\[
u := \frac{w}{\|w\|_2},\quad
g_r := \langle g, u \rangle u,\quad
g_t := g - g_r
\]
The update can be further cast as separate projections:
\[
\phi_r^{(w)}(z) := \frac{ \langle z, w \rangle }{ \langle w, w \rangle } w,\quad
\phi_{\phi}^{(w)}(z) := z - \phi_r^{(w)}(z)
\]
where $g_r = \phi_r^{(w)}(g)$, $g_t = \phi_{\phi}^{(w)}(g)$.

## 2. Algorithmic Instantiations: Radial-only and Module-wise Decay

**Radial-only decay (AdamO):**  
The full optimizer maintains distinct moment buffers for the decomposed subspaces and applies decay exclusively along the radial. The AdamO algorithm, for each step $t$, updates as follows [2602.05136]:

- **Radial direction:** Simple SGD update with learning rate $\eta_{\rho,t}$ (adaptive to local curvature), no adaptive moment.
- **Tangential direction:** Adam-style preconditioning (first and second moments) with fixed learning rate $\eta_\theta$.
- **Decay:** Applied as $w \leftarrow (1 - \eta_{\rho,t}\lambda) w$, strictly in the radial component. Tangential direction is untouched by decay.

For biases, LayerNorm, or small dimensionalities, a default Adam-style update is used. For scale-invariant layers (e.g., BatchNorm, LayerNorm), the radial step and decay are omitted.

**Module-wise decay (AlphaDecay):**  
Instead of uniform decay, AlphaDecay estimates the spectral "heavy-tailedness" of each module's weight correlation matrix via the empirical spectral density and Hill estimator of the power-law tail index $\alpha$ [2506.14562]. Modules with smaller $\alpha$ (heavier tails, more principal directions) receive weaker decay, while lighter-tailed modules get stronger decay:
\[
\lambda^m_t = \eta \left[ \frac{(\alpha^m_t - \alpha_{\min})}{(\alpha_{\max} - \alpha_{\min})} (s_2 - s_1) + s_1 \right]
\]
This schedule is updated every $T_{\mathrm{dec}}$ steps and implemented via grouped parameter groups in the optimizer.

**Selective per-layer decay (SPD):**  
During fine-tuning of foundation models, split decay may be applied only to layers whose update directions are inconsistent with previous progress [2411.01713]. For layer $l$:
- Compute $c^l_t = -\langle g^l_t, \theta^l_{t-1} - \theta^l_0 \rangle$.
- If $c^l_t < 0$ ("over-exploring"), decay is applied proportional to the fractional overshoot radius; otherwise, no shrinkage.

## 3. Theoretical and Practical Implications

**Separation of directions:**  
AdamO and related splits address the inherent conflict between magnitude and directional regularization. By confining decay to the radial subspace and reserving tangential dynamics for feature learning (via moment-preconditioned adjustment), split decay stabilizes norm oscillations, leading to smoother loss and improved effective capacity utilization [2602.05136].

**Spectral adaptation:**  
Module-wise split decay leverages heavy-tailed self-regularization theory: networks empirically display highly variable principal spectra across different attention and MLP blocks. Uniform decay inadvertently suppresses useful high-variance modes, while AlphaDecay spatially regularizes according to empirically measured spectral exponents [2506.14562].

**Fine-tuning robustness:**  
Selective decay (SPD) restricts shrinkage to layers that depart from their pre-trained trajectory, enhancing retention of in-domain and out-of-distribution model behavior during fine-tuning [2411.01713]. This splits regularization power across parameter groups automatically, balancing flexibility with strong constraint where needed.

**Algorithmic realization:**  
Split decay methods can be implemented using projection operators, grouped parameter schedules, or selective masking within a standard optimizer framework compatible with PyTorch or TensorFlow [2602.05136, 2506.14562, 2411.01713].

## 4. Empirical Results and Benchmark Performance

**AdamO (radial-only):**  
On CIFAR-100 (ResNet-18, BatchNorm, 300 epochs), AdamO achieves 79.74 ± 0.09% accuracy, outperforming AdamW (74.75 ± 0.15%) by ≈5 points. Removing any split component (projection, dimension-aware rule, or curvature adaptation) significantly degrades performance, collapsing results to the AdamW regime if no split decay is used. AdamO exhibits smoother optimization trajectories and greater hyperparameter robustness [2602.05136].

**AlphaDecay (module-wise):**  
For LLaMa-family LLMs (60M–1B parameters), perplexity improvements over uniform decay are observed: e.g., 3.0% lower on LLaMa-60M, and consistent but attenuated benefits at larger scale (0.8% for 1B) [2506.14562].

**SPD (selective, fine-tuning):**  
On DomainNet and ImageNet variants, SPD consistently reduces parameter deviation $\|\theta_t - \theta_0\|$ by 3–5 × vs. AdamW and boosts out-of-distribution (OOD) accuracy by 5–10 points. In PEFT settings, SPD improves commonsense QA scores for LLaMA-7B/13B by 1–2 points [2411.01713].

| Optimizer/Method          | CIFAR-100 Acc (%) | OOD Gain (DomainNet) | Perplexity (LLaMa-60M) |
|--------------------------|-------------------|----------------------|-----------------------|
| Adam                     | 74.48 ± 0.12      | –                    | 32.56                 |
| AdamW                    | 74.75 ± 0.15      | 39.3                 | 32.56                 |
| AdamO (full)             | 79.74 ± 0.09      | –                    | –                     |
| SPD (DomainNet)          | –                 | 45.9                 | –                     |
| AlphaDecay (LLMs)        | –                 | –                    | 31.58                 |

## 5. Algorithmic Details and Pseudocode

**AdamO Split Weight Decay Core Loop:**  
```python
# Pseudocode for AdamO core (radial-only split)
g = grad(L, w)
# Decompose gradient
u = w / norm(w)
g_r = dot(g, u) * u
g_t = g - g_r

# Update moments
m_r = beta1_r * proj_r(m_r) + (1-beta1_r) * g_r
m_t = beta1_t * proj_t(m_t) + (1-beta1_t) * g_t
v_t = beta2_t * v_t + (1-beta2_t) * (g_t * g_t)
# Bias correction
mhat_r = m_r / (1 - beta1_r**t)
mhat_t = m_t / (1 - beta1_t**t)
vhat_t = v_t / (1 - beta2_t**t)

# Compute step sizes
eta_r_t = eta_r / (sqrt(curvature/tau_target) + epsilon)

# Updates
Delta_r = eta_r_t * proj_r(mhat_r)      # SGD-style radial
Delta_t = eta_theta * proj_t(mhat_t / (sqrt(vhat_t) + epsilon))  # Adam-style tangential

if is_scale_invariant(w):
    Delta = Delta_t
else:
    Delta = Delta_r + Delta_t

# Radial-only decay
w = (1 - eta_r_t * lambda) * w - Delta
```
**AlphaDecay Module-wise Update (high-level):**
Every $T_{\mathrm{dec}}$ steps, for each module:
- Compute eigenvalues of $C_m = W_m W_m^\top/\operatorname{Tr}(W_m W_m^\top)$.
- Estimate power-law exponent $\alpha^m_t$ with Hill estimator.
- Interpolate decay rate and assign to parameter group in optimizer.

**SPD Update (per layer):**
```python
for each layer l:
    c_t^l = - dot(g_t^l, theta_{t-1}^l - theta_0^l)
    gamma_t^l = norm(theta_tilde^l - theta_0^l)
    gamma_prev^l = norm(theta_{t-1}^l - theta_0^l)
    r_t^l = max(0, gamma_t^l - gamma_prev^l) / (gamma_t^l + epsilon)
    if c_t^l < 0:
        theta^l = theta_tilde^l - lambda * r_t^l * (theta_tilde^l - theta_0^l)
    else:
        theta^l = theta_tilde^l
```
## 6. Comparison to Standard and Decoupled Weight Decay

Traditional $\ell_2$-regularization, linked to "weight decay," applies a penalty $L_{\mathrm{reg}}(\theta) = L(\theta) + \frac{\lambda}{2}\|\theta\|^2$ either as a gradient term or as a multiplicative shrinkage post-update. Loshchilov and Hutter [1711.05101] proved that for adaptive optimizers (e.g., Adam), $\ell_2$ penalty and explicit decay are inequivalent: only fully decoupled decay (AdamW) ensures correct separation of normalization and regularization.

Split weight decay generalizes this separation by further restricting the directions or groups to which decay is applied, providing more nuanced control over optimization geometry. Empirically, isotropic (non-split) decay can suppress beneficial capacity expansion or over-regularize anisotropic modules, outcomes precluded by geometric or selective split decay (AdamO, AlphaDecay, SPD) [2602.05136, 2506.14562, 2411.01713].

## 7. Limitations, Practical Guidelines, and Outlook

Split weight decay requires geometric projections, module-wise eigenanalysis, or parameter grouping, introducing minimal but nonzero computational overhead. For module-wise approaches, measuring spectral densities periodically can increase cost, but update intervals (e.g., every 500 steps) make this manageable in practice [2506.14562]. For finer-grained architectures or decentralization across large-scale pretraining, further research may optimize trade-offs between regularization selectivity and compute efficiency.

Hyperparameter guidelines:
- For AdamO, set $\eta_\rho$ and $\eta_\theta$ by grid search; $\lambda \sim 2 \cdot 10^{-4}$ is common [2602.05136].
- For AlphaDecay, $(s_1, s_2)$ scaling factors between $(0.67, 3)$ and $(0.67, 5)$ are effective [2506.14562].
- For SPD, projection strength $\lambda \in [0.5, 2.5]$ is typical, with $\lambda = 1$ as a robust default [2411.01713].

Split weight decay continues to evolve as neural architectures grow, drawing on geometric, spectral, and empirical insights. It provides a principled approach to directional, modular, and adaptive regularization, with quantifiable improvements in generalization and robustness over isotropic schemes.

Source: https://www.emergentmind.com/topics/split-weight-decay