vietanh.dev

Squared ReLU: The Tiny Activation Change Used by Primer

· 3 min read · 405 words

Authors

Squared ReLU is exactly what its name suggests:

f(x)=ReLU⁡(x)2=max⁡(0,x)2.f(x) = \operatorname{ReLU}(x)^2 = \max(0, x)^2.

Its derivative is

f′(x)={2x,x>0,0,x≤0.f'(x) = \begin{cases} 2x, & x > 0, \\ 0, & x \le 0. \end{cases}

Compared with ordinary ReLU, small positive values are suppressed while large positive values are amplified. The operation adds no learned parameters and only one elementwise multiplication.

Loading diagram…

Why it became interesting

The Primer architecture search identified squared ReLU and a depthwise convolution after each attention projection as the two changes responsible for most of its language-modeling gains. In the paper's experiments, those changes reduced the compute needed to reach a target quality across several model scales.

That result is evidence for a particular Transformer training setup—not a universal guarantee that squaring ReLU improves every network.

If this replaces the activation inside a feed-forward network,

FFN⁡(x)=W2 f(W1x+b1)+b2,\operatorname{FFN}(x) = W_2\,f(W_1x+b_1)+b_2,

the tensor shapes and parameter count remain unchanged. It can therefore be compared directly with ReLU or GELU without changing model width.

What changes in practice

  • Positive activations above one grow faster than with ReLU.
  • Positive activations between zero and one shrink.
  • The positive-side gradient is no longer constant; it grows with the activation.
  • The function remains zero for negative inputs, so it retains ReLU's hard gate.

The unbounded quadratic branch means activation scale matters. Monitor activation statistics, loss spikes, and mixed-precision overflow when transferring the idea to a new architecture. Normalization and optimizer settings that were stable with GELU are not automatically stable with squared ReLU.

Recommendation

Treat squared ReLU as a cheap experiment, not a default. Compare it under the same parameter budget, training tokens, optimizer, and evaluation protocol. If it wins, verify the gain across seeds and check numerical behavior at the precision used in production.

Primary source