Technical Summary: Learning in Curved Weight Space (SymExpLin)

1. Problem Statement

Standard neural network training relies on additive updates in a linear parameter space. However, many internal network operations are inherently multiplicative (e.g., scaling feature amplitudes, normalization gains, or singular values). In a linear space, relative changes of equal magnitude (e.g., doubling vs. halving a weight) require unequal optimization distances. For instance, moving a weight from 1.0 to 2.0 requires a step of size 1.0, while moving from 1.0 to 0.5 requires a step of 0.5.

This mismatch creates optimization challenges:

• Heavy-Tailed Distributions: Pretrained models often exhibit heavy-tailed weight distributions where a small fraction of parameters grow significantly larger than their initialization scale (e.g., 30× larger), while most remain near initialization.
• Optimizer Invariance: Adaptive optimizers like Adam normalize updates per coordinate, making step sizes invariant to parameter magnitude. Consequently, a small parameter (e.g., 0.004) and a large parameter (e.g., 0.4) receive similar absolute updates. This results in a 100× difference in relative perturbation, potentially causing small parameters to be over-perturbed while large, critical parameters are under-perturbed.
• Static vs. Dynamic Scaling: Existing solutions like Maximal Update Parameterization (µP) address scale mismatches statically based on tensor width at initialization but do not adapt to the dynamic growth of individual parameters during training.

2. Methodology: SymExpLin (SEL)

The authors propose SymExpLin (SEL), a nonlinear weight reparameterization that decomposes effective weights into a symmetric-exponential pathway and a linear pathway. This creates a curved parameter-space geometry where additive updates in the raw space map to magnitude-proportional changes in the effective weight space.

2.1 The Transform

The effective weight $w_{eff}$ is computed from a raw weight $w$ using the following components:

1. Symmetric-Exponential Pathway ($E_\theta$): A learnable, sign-preserving function based on the symmetric exponential (symexp). For small $|w|$, it behaves linearly; for large $|w|$, it grows exponentially. This allows large raw steps to produce larger effective displacements.
   • $E_\theta(|w|) = e_w (\exp(\beta |w| \cdot m - n) - \exp(-n))$
   • Parameters: $\beta$ (curvature), $m$ (inner multiplier), $n$ (offset for slope at origin), $e_w$ (output scale).
2. Linear Pathway ($L_\theta$): An identity-like pathway $L_\theta(w) = w \cdot l_w$ that provides a direct gradient route, hypothesized to stabilize optimization.
3. Sign Wrapping: The sign of the raw weight is applied to the sum of the pathways: $w_{eff} = \text{sign}(w) \cdot (E_\theta(|w|) + L_\theta(w)) + \Delta_\phi$.
4. Residual ($\Delta_\phi$): An optional low-rank (LoRA-style) additive correction.

2.2 Key Design Choices

• Mismatched Initialization: The authors identify a beneficial asymmetry where raw weights are initialized by inverting a "congruent" combination rule (where linear and exponential parts are added outside the sign wrap), but the forward pass uses a "mismatch" rule (where both are wrapped inside the sign). This results in negative weights being initialized at a smaller magnitude than positive weights, acting as an early symmetry-breaking mechanism.
• Structured Scale Parameters: Learnable scales ($e_w, l_w, m, n$) are parameterized using row-column multiplicative patterns (rank-1 vectors) to balance expressivity and computational cost.
• Learning Rate Annealing: A targeted schedule is used for scale parameters. The outer scale $e_w$ starts with a high learning rate to quickly commit to a sign, while the inner multiplier $m$ starts low and anneals upward to refine curvature once the scale stabilizes.
• Curvature Calibration: The curvature parameter $\beta$ is set heuristically based on model width to ensure the exponential regime activates near typical initialization magnitudes (which shrink as width increases).

3. Key Contributions

• SymExpLin (SEL): A weight reparameterization applicable to any learnable matrix or bias, decomposing weights into symmetric-exponential and linear components with independently learnable structured scales.
• Mismatched Initialization: The discovery that using different combination rules for initialization inversion versus the forward pass creates a beneficial asymmetric bias that improves early optimization.
• Structured Scale & Annealing: Introduction of row-column multiplicative scale parameters and a specific learning rate annealing schedule to manage redundant exponential controls.
• Practicality: The method is implementation-friendly. The training transform is element-wise, and post-training, effective weights are precomputed and folded into standard linear layers, resulting in zero inference cost.
• Complementarity: SEL operates on a dynamic, per-parameter axis, complementing static, per-tensor methods like µP.

4. Experimental Results

The method was evaluated on autoregressive language modeling using OpenWebText across nine configurations of width $\{1024, 2048, 3072\}$ and depth $\{12, 24, 36\}$.

• Loss Descent: SEL consistently reached lower validation loss than the baseline (Adam with standard parameterization) across all configurations.
• Step Efficiency: SEL achieved matched validation loss in 1.32× to 1.49× fewer training steps. The gains were most pronounced in larger width models (e.g., 1.49× speedup at 3072×24).
• Wall-Clock Speedup: Despite a per-step overhead of 1.7% (small model) to 5.5% (large model) due to the transform, the reduction in total steps resulted in an estimated 1.29× to 1.37× wall-clock speedup.
• Ablations:
  • Linear Pathway: Removing the linear pathway resulted in a "failed-to-learn" case, confirming its necessity for stability.
  • Mismatch vs. Congruent: Mismatched initialization consistently outperformed congruent initialization, particularly in smaller models.
  • Annealing: The annealed learning rate schedule for scale parameters provided consistent, albeit small, improvements over fixed multipliers.

5. Significance and Claims

The paper claims that SEL addresses the fundamental mismatch between the multiplicative nature of many neural computations and the additive nature of standard optimization. By curving the parameter space, SEL allows the optimizer to take steps that are naturally proportional to the magnitude of the weights, effectively providing relative scaling without requiring per-parameter learning rate tuning.

The authors emphasize that SEL is a practical, drop-in replacement for standard linear layers that:

1. Accelerates convergence, particularly for models with heavy-tailed weight distributions.
2. Requires no architectural changes during inference.
3. Can be inverted, allowing standard pretrained weights to be converted back into SEL space for continued training or fine-tuning.

The work suggests that the "soft log-space" parameterization helps weights reach the large-magnitude targets required by the network's computation, which are often disadvantaged by linear additive updates.