PowLU: Taming the Quadratic Beast for Stable LLM Training

TL;DR

Researchers from the Ling Team at Ant Group have proposed PowLU (Power Linear Unit), a drop-in replacement for the popular SwiGLU activation function. While SwiGLU’s approximate quadratic growth ($x^2$) boosts expressivity, it often triggers catastrophic "loss spikes" in large-scale and low-precision (FP8) training. PowLU uses a rational power function to maintain non-linearity while enforcing a more stable, near-linear growth for large inputs, ensuring smooth pre-training of models up to 124B parameters.

The Problem: The Hidden Cost of SwiGLU's Expressivity

In the current LLM landscape, SwiGLU is the de facto standard. Its success stems from its ability to approximate a quadratic function for large positive inputs, providing the "non-linear muscle" needed for complex reasoning.

However, this $x^2$ behavior is a double-edged sword. As models grow deeper and training moves toward lower precision (FP8/FP4), this quadratic amplification:

1. Exacerbates Outliers: Small anomalies in earlier layers are squared and magnified.
2. Triggers Numerical Overflow: Values quickly exceed the narrow dynamic range of low-precision formats.
3. Causes Training Collapse: Uncontrolled gradient and activation spikes manifest as the dreaded "loss spikes," forcing researchers to restart training runs.

Figure 1: SwiGLU (left) shows a wide band of high-magnitude outliers compared to the constrained distribution of PowLU (right).

Methodology: The Logic of Bounded Power

The core innovation of PowLU lies in its mathematical formulation for $x>0$:

$extPowLU(x)=x\cdot x^{\frac{m}{\sqrt{x}+1}}\cdot extsigmoid(x)$

Why this specific form?

• Adaptive Growth: For small $x$, the exponent $1+m/(\sqrt{x}+1)$ remains high, providing strong non-linearity. As $xo\infty$, the exponent approaches 1, meaning the function transitions gracefully to linear growth.
• Differentiability: The $\sqrt{x}+1$ term is a clever inclusion. Without the $+1$, the derivative would approach infinity as $xo0$, leading to new instabilities.
• Smooth Suppression: Unlike SwiGLU-Clip, which simply "chops off" large values (hard truncation), PowLU smoothly suppresses growth, preserving more information in the latent space.

Figure 2: First-order derivatives highlighting the smooth, controlled gradients of PowLU compared to the explosive growth of SwiGLU.

Experiments: Performance Without the Spikes

The authors tested PowLU across several scales, including a heavy-weight 124B MoE model.

1. Stability at Scale

In FP8 training, standard SwiGLU encountered significant stability issues. While SwiGLU-Clip delayed spikes, PowLU maintained a perfectly smooth loss trajectory. Figure 3: Training loss curves. Notice the red line (PowLU) stays flat and stable while others spike.

2. Benchmark Excellence

Stability didn't come at the cost of "intelligence." At 124B parameters:

• ARC-Challenge: SwiGLU (77.29) vs PowLU (83.05).
• MATH: SwiGLU (42.22) vs PowLU (44.98).
• HumanEval (Coding): SwiGLU (54.27) vs PowLU (55.49).

PowLU consistently matched or outperformed SwiGLU, proving that "near-linear" growth for large activations is sufficient for high-level reasoning.

Critical Insight: The Outlier Channel Solution

One of the most profound visualizations in the paper involves Outlier Channels. LLMs often suffer from a few "screaming" channels that carry disproportionately large magnitudes. PowLU significantly homogenizes these channel magnitudes, which is likely the primary reason it works so well for low-precision quantization.

Conclusion

PowLU is a mathematically grounded and empirically robust evolution of the activation function. It addresses the fundamental tension between expressive capacity and numerical stability. For any team currently pre-training LLMs in FP8 or seeking to scale beyond the 100B barrier without constant supervision of loss spikes, PowLU represents a significant step forward.

Takeaway: The era of "bigger is better" is evolving into "more stable is better." PowLU proves that by slightly curbing the wilder tendencies of SwiGLU, we can build models that are both smarter and more reliable.