CompleteP Transformer Parameterization

Introduction: The Problem of HP Transfer Across Model Sizes

A fundamental challenge in training Large Language Models (LLMs) is determining how to scale hyperparameters (HPs) when changing model architecture. Modern deep learning’s core paradigm—that larger models yield better performance¹—creates enormous pressure to scale model width and depth efficiently.

The Scaling Challenge: When practitioners scale models from small experiments to production deployments, they face two critical problems:

1. HP Retuning Overhead: Traditional parameterizations require re-tuning learning rates and other HPs for each model size, which is computationally expensive
2. Sub-optimal Training: When re-tuning is prohibitive, models are trained with sub-optimal HPs, wasting valuable compute resources

The maximal update parameterization (µP)² was developed to address HP transfer when scaling width. However, extending this to simultaneous width and depth scaling remained an open problem.

This article introduces CompleteP³, the parameterization that achieves both:

• Depth-wise HP transfer: Optimal base HPs remain approximately constant when scaling model depth
• Non-lazy learning in all layers: Every layer learns genuinely non-linear features, not just linearizations

CompleteP enables 12-34% compute efficiency improvements over prior state-of-the-art parameterizations.

---

The Lazy Learning Regime in Deep Networks

What is Lazy Learning?

In deep networks, lazy learning refers to a regime where layers learn features that remain very close to their linearization around initialization⁴. This phenomenon occurs when the network’s non-linear components contribute minimally to the learned representations.

Mathematical Definition: A layer is in the lazy learning regime if:

$$\Vert f_{\theta }(x)-f_{{\theta }_0}(x)-{\nabla }_{\theta }{f}_{{\theta }_0}(x)\cdot (\theta -{\theta }_0)\Vert \ll \Vert \theta -{\theta }_0\Vert$$

where $f_{\theta }$ is the network output with parameters $\theta$, and ${\theta }_0$ is the initialization.

Why Lazy Learning is Problematic

When networks operate in the lazy regime:

Aspect	Impact
Feature Diversity	Layers learn similar (linear) features, wasting depth
Non-linearity Unused	Activation functions contribute little to learning
Depth Inefficiency	Adding more layers provides diminishing returns
Compute Waste	Training compute is spent without proportional benefit

Previous work⁵ argued that HP transfer was not truly possible at any depth scaling factor, and that $\alpha =0.5$ was optimal in practice. This suggested that lazy learning was an unavoidable trade-off.

---

CompleteP Parameterization Design

Core Design Principle

CompleteP sets the depth scaling exponent $\alpha =1$, which gives the parameterization its name. This is the unique value of $\alpha$ that ensures complete feature learning (non-lazy learning in all layers) while maintaining HP transfer.

The Residual Connection Formula

The key difference between parameterizations lies in how the transformer’s residual block outputs are scaled:

$${\mathbf{h}}^{\ell +1}={\mathbf{h}}^{\ell }+L^{-\alpha }\cdot {\mathcal{F}}_{\ell }({\mathbf{h}}^{\ell }),\ell \in \{1,\dots ,L\}$$

where:

• ${\mathbf{h}}^{\ell }$ is the residual stream at layer $\ell$
• ${\mathcal{F}}_{\ell }$ is the $\ell$-th residual block (MLP or attention)
• $L$ is the total number of layers
• $\alpha \in [0.5,1]$ controls depth-dependent re-scaling

Parameterization Comparison

Parameterization	$\alpha$ Value	Key Properties
SP (Standard)	N/A	No depth scaling; unstable at depth
µP	N/A	Width scaling only; $\alpha$ not applicable
Mean-Field $\alpha =0.5$	$0.5$	Partial depth scaling; still lazy
CompleteP	$1.0$	Full depth scaling; non-lazy learning

CompleteP Scaling Rules

CompleteP extends µP with principled depth scaling. The complete scaling rules for a Pre-LN transformer³:

Width Multipliers

Component	SP	µP	CompleteP ($\alpha =1$)
Embedding Init Var	${\sigma }_{base}^2$	${\sigma }_{base}^2$	${\sigma }_{base}^2$
Embedding LR	${\eta }_{base}$	${\eta }_{base}$	${\eta }_{base}$
Hidden Init Var	${\sigma }_{base}^2$	${\sigma }_{base}^2/m_N$	${\sigma }_{base}^2/m_N$
Hidden LR (Adam)	${\eta }_{base}$	${\eta }_{base}/m_N$	${\eta }_{base}\cdot m_L/m_N$

where:

• $m_N=N/N_{base}$ (width multiplier)
• $m_L=L/L_{base}$ (depth multiplier)

Depth-Specific Corrections (CompleteP)

CompleteP requires additional corrections for stable training:

1. Residual Scaling Factor:

$$\text{residual\_scaling}=\frac{1}{m_L^{\alpha }}=\frac{1}{m_L}$$

2. LayerNorm Learning Rate: ${\eta }_{LN}={\eta }_{base}\cdot m_L$

3. Weight Decay: ${\lambda }_{hidden}={\lambda }_{base}\cdot m_N/m_L$

4. AdamW $ϵ$:

   • Embedding layers: ${ϵ}_{emb}={ϵ}_{base}/m_N$
   • Hidden layers: ${ϵ}_{hidden}={ϵ}_{base}/(m_N\cdot m_L)$

---

Depth-wise HP Transfer Guarantees

Empirical Evidence

The paper demonstrates that only CompleteP ($\alpha =1$) achieves reliable depth-wise HP transfer³:

Finding 1: With SP, µP, and $\alpha =0.5$, optimal learning rate $\eta$ and initialization standard deviation $\sigma$ do not remain stable as depth $L$ varies.

Finding 2: Only CompleteP ($\alpha =1$) achieves reliable depth-wise HP transfer.

Coordinate Check Results

The paper validates HP transfer using coordinate checks—training with different depth multipliers and verifying that:

• Loss curves align when using transferred HPs
• Optimal HPs remain constant across depths

CompleteP (α=1): Concentric curves with stable minima
α=0.5: Diverging optimal HPs with depth
µP/SP: Unstable training at depth

Compute-Optimal Setting

Even under the compute-optimal setup (20 tokens per parameter), CompleteP shows superior HP transfer:

• Reduced sensitivity to learning rate choice
• Consistent loss improvements across depths
• No additional HP tuning required when scaling depth

---

Non-Lazy Learning in All Layers: Proof Sketch

The Lazy vs Non-Lazy Spectrum

Yang et al. (2023)⁵ argued that $\alpha =0.5$ was optimal based on feature diversity considerations. CompleteP’s key insight is that this argument only holds in the lazy learning regime.

Complete Feature Learning Requirement

CompleteP introduces a refined desideratum for HP transfer that includes complete feature learning³:

Complete Feature Learning: The learned representation in every layer of the model must remain non-lazy (i.e., non-linear) with respect to the parameters in both that layer and all earlier layers.

Mathematical Framework

Consider the network’s representation at layer $\ell$:

$${\mathbf{h}}^{\ell }={\mathbf{h}}^0+\sum _{i=1}^{\ell }{L}^{-\alpha }{\mathcal{F}}_i({\mathbf{h}}^{i-1})$$

Critical Insight: For $\alpha =1$, the contribution from each residual block is scaled by $L^{-1}$, ensuring:

1. Order 1 terms (linear) contribute $\mathcal{O}(L\cdot L^{-1})=\mathcal{O}(1)$
2. Order 2+ terms (non-linear) remain significant at $\mathcal{O}(1)$ scale
3. Lazy regime avoided: Non-linear features cannot be absorbed into linear approximations

Theorem Sketch (Informal)

Theorem: Under CompleteP parameterization ($\alpha =1$), as $N,L\to \infty$ with finite ratio $N/L$:

• Every layer’s learned features are non-lazy with respect to all earlier layers
• The training dynamics remain in the feature learning regime rather than the NTK/lazy regime
• HP transfer holds: optimal base HPs are constant across model sizes

This theorem distinguishes CompleteP from all other $\alpha \in [0.5,1)$ parameterizations.

---

Hardware-Aware Shape Optimization

The N:L Ratio Question

Kaplan et al. (2020)⁶ established that $N:L\approx 100$ is compute-optimal using SP. However, SP doesn’t fairly admit stable width and depth scaling, undermining this conclusion.

CompleteP Enables Wider Design Space

CompleteP’s stable depth scaling enables revisiting the compute-optimal width-to-depth ratio:

N:L Ratio	SP/µP Viability	CompleteP Viability
100:1	Stable	Stable
50:1	Moderate	Stable
20:1	Challenging	Stable
10:1	Poor	Stable
1:1	Unstable	Viable

Hardware Considerations

Different hardware configurations favor different N:L ratios:

Hardware Type	Preferred N:L	Rationale
GPU (memory-bound)	Higher N:L	Better memory utilization
Wafer-scale (Cerebras)	Variable	Massive memory bandwidth enables deeper models
Edge devices	Lower N:L	Memory constraints favor depth efficiency

CompleteP’s key advantage: Practitioners can now choose shapes based on hardware characteristics rather than training stability constraints.

---

Compute Efficiency Improvements

FLOP Savings

The paper reports substantial compute efficiency improvements with CompleteP³:

Model Configuration	FLOP Savings vs µP
Optimally-shaped 1.9B model	11.8%
179-layer deep model	34.4%

Training Curves

Models trained with CompleteP achieve:

• Faster convergence (lower loss at same token count)
• Better final performance (lower asymptotic loss)
• Improved scaling with depth

Efficiency Breakdown

The 12-34% improvement comes from three sources:

1. HP Transfer Efficiency: Eliminating re-tuning costs
2. Training Efficiency: Non-lazy learning extracts more value from each layer
3. Shape Efficiency: Deeper models become viable, enabling hardware-optimized shapes

---

Practical Implementation Guidelines

Minimal Implementation

The following shows the key modifications needed to implement CompleteP³:

# In the Block forward pass:
def forward(self, x):
    residual_scaling = 1 / (self.depth_multiplier ** self.depth_alpha_exp)
    x = x + residual_scaling * self.attn(self.ln_1(x))
    x = x + residual_scaling * self.mlp(self.ln_2(x))
    return x

Optimizer Configuration

CompleteP requires separated optimizer groups with depth-dependent learning rate scaling:

# Depth learning rate scaling
depth_lr_scaling = depth_multiplier ** (depth_alpha_exp - 1)  # = 1/depth_multiplier for α=1
width_lr_scaling = 1 / width_multiplier
 
# AdamW epsilon scaling
emb_unemb_adam_eps = adam_eps / width_multiplier
hidden_adam_eps = adam_eps / (width_multiplier * depth_multiplier)
 
optim_groups = [
    {'params': emb_params, 'weight_decay': weight_decay, 'lr_scale': 1.0, 'eps': emb_unemb_adam_eps},
    {'params': hidden_ln_params, 'weight_decay': 0.0, 'lr_scale': depth_lr_scaling, 'eps': hidden_adam_eps},
    {'params': hidden_weight_params, 'weight_decay': weight_decay / width_lr_scaling, 
     'lr_scale': width_lr_scaling * depth_lr_scaling, 'eps': hidden_adam_eps},
    {'params': hidden_bias_params, 'weight_decay': 0.0, 'lr_scale': depth_lr_scaling, 'eps': hidden_adam_eps},
]

Configuration Checklist

When implementing CompleteP:

• Set depth_alpha_exp = 1.0 (the key CompleteP parameter)
• Implement residual scaling: residual_scaling = 1/depth_multiplier
• Separate optimizer groups for LN, hidden weights, biases
• Scale AdamW $ϵ$ by depth and width multipliers
• Scale weight decay inversely with depth

Reference Implementation

A minimal implementation is available at:

github.com/EleutherAI/nanoGPT-mup/tree/completep

---

Summary and Key Takeaways

What CompleteP Achieves

1. Depth-wise HP Transfer: Optimal learning rates and initialization scales transfer across model depths without re-tuning

2. Non-Lazy Learning: Every layer learns genuinely non-linear features, maximizing the value of model depth

3. Hardware Flexibility: Enables a wider range of width-to-depth ratios to be compute-efficient

4. Empirical Gains: 12-34% compute efficiency improvements over prior state-of-the-art

When to Use CompleteP

Scenario	Recommendation
Scaling depth significantly	Essential
Multi-model family training	Highly recommended
Hardware-specific shape optimization	Recommended
Single small model only	µP sufficient

Key Parameters

Parameter	Value for CompleteP	Purpose
$\alpha$ (depth_alpha_exp)	1.0	Core CompleteP parameter
Residual scaling	1/m_L	Depth-dependent residual scaling
Hidden LR scaling	$m_L/m_N$	Learning rate adjustment
AdamW $ϵ$ scaling	$1/(m_N\cdot m_L)$	Numerical stability

---

References

Related Papers

• Depthwise hyperparameter transfer in residual networks — Dalmazi et al., 2023
• Infinite limits of multi-head transformer dynamics — Bordelon et al., 2024
• The practitioner’s guide to the maximal update parameterization — Dey et al., 2024

Footnotes

1. Deep learning scaling is predictable, empirically ↩

2. Tuning large neural networks via zero-shot hyperparameter transfer ↩

3. Don’t be lazy: CompleteP enables compute-efficient deep transformers ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶

4. On lazy training in differentiable programming ↩

5. Feature learning in infinite-depth neural networks ↩ ↩²

6. Scaling Laws for Neural Language Models ↩