Neural Dimensional Analysis: Foundations of Dimensionless and Scale-Consistent Computation

Summary (Overview)

• Unified framework: Introduces Neural Dimensional Analysis (NDA), a framework based on dimensional consistency that unifies scale-control mechanisms in Transformers (normalization, residual scaling, attention temperatures, loss weighting, and adaptive updates) under one set of local typing rules.
• Key theoretical contribution: Defines neural units as transformation laws under a rescaling group, proving that forward unit declarations fix gradient transformation laws (Reverse-mode dual typing, Prop. 1), enabling unit-consistent diagnostics and network design.
• Novel diagnostics and architecture: Introduces dimensionless Unit Violation Indices (UVI) for diagnosing scale mismatches and the Gauge-Canonical Transformer (GCT), which computes every block on dimensionless coordinates with orthogonalized residual additions.
• Strong empirical results: GCT remains stable across a 16× learning-rate range (51M–1B parameters), achieves superior depth transfer (3.903 vs 4.243 nats for nGPT at L=48), and shows exceptional quantization robustness (only +0.009 nats loss increase under per-tensor int8 at 1B, vs up to +5.51 for baselines).
• Broad applicability: Validates NDA principles beyond language modeling via vision benchmarks and physics-informed neural networks (PINNs), where unit-aware loss aggregation preserves accuracy across a 10⁴ unit sweep.

Introduction and Theoretical Foundation

Background and Motivation

Modern Transformers rely on a broad family of scale-control mechanisms:

• Normalization: BatchNorm, LayerNorm, RMSNorm normalize activations
• Residual scaling: DeepNet, ReZero, LayerScale, Fixup control residual addition scales
• Attention regulation: Attention scaling and QK-Norm regulate softmax inputs
• Loss balancing: Loss-balancing methods and adaptive optimizers manage scale imbalances

The paper observes that all these mechanisms manage how quantities of different scales interact within a neural network, and interprets them through the lens of dimensional consistency—a fundamental principle from physics.

Theoretical Foundation: Neural Units

The core insight is to define neural units as transformation laws governing how internal tensors rescale when the computation graph is rescaled. Unlike physical units, neural units describe internal scaling behavior even when data have no physical units.

Rescaling group and representations:

Let $\mathcal{G} := (\mathbb{R}_+)^K$ be a rescaling group with one factor per independent internal scale degree of freedom. A neural unit of node $v$ is a representation:

$$\rho_v : \mathcal{G} \to \text{GL}(V_v), \quad x_v \mapsto \rho_v(g) x_v$$

A tensor is dimensionless if and only if $\rho_v \equiv I$. For scalar units, we use continuous characters:

$$\chi_u(g) := \prod_{k=1}^{K} g_k^{u_k}, \quad u \in \mathbb{R}^K$$

Unit consistency for an operation $y := f(x_1, \ldots, x_n)$ requires the equivariance identity:

$$f(\rho_1(g) x_1, \ldots, \rho_n(g) x_n) = \rho_y(g) f(x_1, \ldots, x_n) \quad \forall g \in \mathcal{G}$$

Key Theoretical Results

Proposition 1 (Reverse-mode dual typing): If the loss node is typed dimensionless and every primitive operation is unit-consistent, then every reverse-mode gradient node transforms by the dual law $\rho_v^{-\top}$.

Under scalar rescaling $x \mapsto cx$, we get $\nabla_x L \mapsto c^{-1}\nabla_x L$, meaning the primal–dual RMS product $\Pi_x := \text{RMS}(x) \cdot \text{RMS}(\nabla_x L)$ is comparable across rescalings.

Local Typing Rules

Calibrated attention temperature:

$$\tau_{\ell,h} := \sqrt{d_h} \, s_{Q,\ell,h} \, s_{K,\ell,h}$$

where $s_Q$ and $s_K$ are per-head coordinate scales, and the standard $\sqrt{d_h}$ is the $s_Q = s_K = 1$ special case.

Loss weight factorization:

$$\lambda_i = \alpha_i / S_i, \quad u(S_i) = u(L_i), \quad u(\alpha_i) = 1$$

where $1/S_i$ converts each loss to dimensionless value and $\alpha_i$ sets task preference.

Residual Scale Dynamics

For residual addition $h_{\ell+1} = h_\ell + r_\ell$ with step size $\alpha_\ell := \text{RMS}(r_\ell)/\text{RMS}(h_\ell)$:

$$s_{\ell+1}^2 = s_\ell^2 (1 + \alpha_\ell^2 + 2\alpha_\ell c_\ell), \quad c_\ell := \frac{\langle h_\ell, r_\ell \rangle}{\|h_\ell\| \|r_\ell\|}$$

For persistent positive alignment ($c_\ell \geq c_0 > 0$), bounded scale requires $\sum_\ell \alpha_\ell = O(1)$, giving $O(L^{-1})$ per layer. Orthogonalization removes the linear term, permitting $\sum_\ell \alpha_\ell^2 = O(1)$ with uniform $\Theta(L^{-1/2})$ steps.

Methodology

Unit Violation Indices (UVI)

UVI measures dimensionless log-deviations from target values or log-dispersion across depth:

$$\text{UVI}_{\text{res}}(\ell) := \log \frac{\alpha_\ell + \epsilon}{\alpha^\star + \epsilon}, \quad \alpha^\star := \frac{1}{\sqrt{2L}}$$ $$\text{UVI}_{\text{attn}}(\ell, h) := \log \frac{\bar{v}_{\ell,h} + \epsilon}{v^\star + \epsilon}, \quad v^\star := 1$$

where $\alpha_\ell$ is the measured residual step size and $\bar{v}_{\ell,h}$ is the batch-averaged row-centered logit variance.

Gauge-Canonical Transformer (GCT)

GCT computes every block on normalized, dimensionless coordinates via the factorization:

$$\text{unitful} = \text{unit scale} \times \text{canonical coordinate}$$

Canonicalizer (with numerical stability floor):

$$C(x) := \frac{x}{\sqrt{\text{RMS}(x)^2 + \epsilon_x}}$$

Orthogonalized residual update:

$$\bar{a}_\ell := \Phi_\ell(\bar{h}_\ell), \quad \hat{a}_\ell := C(P_{\bar{h}_\ell}^{\perp} \bar{a}_\ell)$$ $$\bar{h}_{\ell+} := C(\bar{h}_\ell + \alpha_\ell \hat{a}_\ell)$$

where the orthogonal projector is:

$$P_{\bar{h}}^\perp a := a - \frac{\langle a, \bar{h} \rangle}{\|\bar{h}\|^2} \bar{h}, \quad \bar{h} \neq 0$$

and learnable dimensionless scalars $\alpha_\ell$ are initialized at $1/\sqrt{2L}$.

Key components:

• Queries/keys canonicalized per head: $\bar{q} := C(q)$, $\bar{k} := C(k)$
• Per-head dimensionless gain $g_h$ initialized at 1
• Normalized gate: $\tilde{g} := C(a)\sigma(C(a))$ for gated MLPs
• Dimensionless logits from $C(\bar{h}_L)$ with learnable logit gain

Experimental Setup

• Models: Decoder-only Transformers with 51M, 124M, 350M, and 1B parameters
• Baselines: Pre-LN, LLaMA-style, QK-Norm, nGPT
• Data: FineWeb-Edu with identical data order, token budget, and optimization recipe
• Tests: Learning-rate robustness, depth transfer, internal-rescaling invariance, PTQ robustness, PINN unit invariance

Empirical Validation / Results

Learning-Rate Robustness

GCT and nGPT remain stable across the full 16× learning-rate range and retain low validation loss at the highest rates. At rate $9.6 \times 10^{-3}$:

• GCT and nGPT continue improving
• QK-Norm's loss is 0.197 nats above its best
• Pre-LN and LLaMA-style stopped early (~8% of budget) due to UVI threshold

Depth Transfer (Fixed Budget, 400M Tokens)

Adding the GCT projection to nGPT reduces loss by 0.275 nats at L=48, demonstrating the NDA-derived component's broad applicability.

Post-Training Quantization (PTQ) at 1B

Key finding: NDA predicts that cross-token scale variation consumes $\log_2 \kappa_{\text{scale}}$ bits of effective resolution, where $\kappa_{\text{scale}}$ is the ratio of largest to smallest token RMS. The zero-floor GCT stream fixes $\kappa_{\text{scale}} = 1$.

Weight-Space Fingerprint

Prop. 6 predicts zero radial sensitivity $\langle W, \nabla_W L \rangle_F$ when downstream computation cancels rescaling. All 54 of 54 architecture-by-role predictions match across eight trained 124M models; 34 of 34 cells match at 1B.

PINN Results

Across a $10^4$ unit sweep, NDA's scaled output head and loss weights keep worst relative $L_2$ error at 0.015, whereas equal or adaptive weights degrade for at least one rescaling.

Theoretical and Practical Implications

Theoretical Contributions

1. Unified unit calculus: Provides a graph-level unit calculus covering activations, objectives, gradients, and updates, linking forward computation and backpropagation through a single set of unit rules (Prop. 1).

2. Quantization theory: Establishes that worst-case downstream output error grows as $\Theta(\kappa_{\text{scale}}/q)$ where $q = 2^{b-1} - 1$, formalizing how cross-token scale variation consumes bits of effective resolution.

3. Residual stability analysis: Proves that orthogonalized residual additions permit $\Theta(L^{-1/2})$ per-layer steps while maintaining bounded cumulative squared step size, as in depth-wise $\mu$P.

Practical Implications

1. Diagnostic tool: UVI provides inexpensive, actionable diagnostics on small batches during training, locating scale mismatches at specific interfaces.

2. Architecture design: GCT demonstrates that explicit scale management grounded in dimensional analysis yields:

   • More reliable learning-rate transfer across model sizes
   • Better depth scaling without hyperparameter retuning
   • Superior quantization robustness for deployment
   • Unit-invariant loss aggregation for multi-task and physics-informed learning

3. General principle: The success of NDA across language modeling, vision, and PINNs suggests dimensional consistency as a fundamental neural network design principle, not merely a practical trick.

Conclusion

Neural Dimensional Analysis establishes dimensional consistency as a typing principle for scale management across the entire computation graph. The framework:

1. Unifies existing scale-control mechanisms (normalization, residual scaling, attention temperature, loss weighting, optimizer updates) under one set of local typing rules
2. Provides testable diagnostics (UVI) and a weight-space audit that locate scale mismatches
3. Yields a practical architecture (GCT) that achieves competitive clean performance while providing superior learning-rate robustness, depth transfer, and quantization resilience
4. Extends beyond language modeling to vision and physics-informed neural networks

Future directions suggested by this work include:

• Applying NDA to other architectures (CNNs, graph neural networks)
• Exploring unit-consistent objectives for multi-modal learning
• Extending the framework to structured state spaces and other sequence models
• Investigating connections with neural ODEs and other physics-grounded architectures

The results collectively establish dimensional consistency as a neural network design principle, offering a principled alternative to the ad-hoc scale-control mechanisms currently used in deep learning.