Finding Alternatives to Normalization via Gradient Approximation
I wonder if you've noticed the recent paper Transformers without Normalization? This paper attempts to replace the Normalization layer in Transformer models with an element-wise operation called DyT, in hopes of improving speed while maintaining performance. This kind of foundational architecture topic is inherently attention-grabbing, and combined with the fact that Kaiming He and Yann LeCun are both listed as authors, the paper drew quite a bit of attention when it was released, with mixed reviews.
Coincidentally, a new paper from last week, The Mathematical Relationship Between Layer Normalization and Dynamic Activation Functions, interprets DyT from the perspective of gradient analysis and differential equations, and proposes a new alternative. Personally I feel that this angle of understanding is quite essential, so I went ahead and studied it and want to share it here.
Preface
DyT stands for Dynamic Tanh, and it replaces the Normalization layer with the following operation:
\begin{equation}\mathop{\text{DyT}}(\boldsymbol{x}) = \boldsymbol{\gamma} \odot \tanh(\alpha \boldsymbol{x}) + \boldsymbol{\beta}\end{equation}more
where $\alpha,\boldsymbol{\beta},\boldsymbol{\gamma}$ are all learnable parameters, and $\boldsymbol{\beta},\boldsymbol{\gamma}$ is already present in the original Normalization layer, so the key here is that $\tanh(\alpha \boldsymbol{x})$ replaces the Normalize operation. $\tanh$ is an element-wise operation, which eliminates the need to compute the two statistics of mean and variance.
Regarding DyT, I once shared some thoughts on Zhihu in How would you evaluate Meta's new paper Transformers without Normalization? — in short, I wasn't very optimistic about it. The reason is that Normalization mindlessly stabilizes the forward propagation of the model, thereby leaving more degrees of freedom and possibilities for other aspects of the model (such as performance). So I don't believe that a general operation simpler than having Normalization can achieve better results (No Free Lunch).
In fact, back in 2021, in A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers we already discussed the topic of removing Normalization, with related work including SkipInit, ReZero, and Fixup. At the time I tried a few schemes and found that even though they could match Normalization in some respects, they still fell short in other respects — for example, pretraining performance might be acceptable, but fine-tuning performance would be poorer. So I didn't dig any deeper into it at the time.
Because of this, I now tend to appreciate work of this kind purely as an exploration of the limits along the "simplification" dimension, much like how nGPT: Normalized Transformer with Representation Learning on the Hypersphere adds Normalize almost everywhere it possibly can — both belong to extreme explorations along a particular direction.
Gradient Computation
Of course, not being optimistic doesn't stop us from studying and analyzing it. If we want to find a replacement or approximation for Normalization, the most direct approach is to start from the gradient, because deep learning ultimately comes down to forward propagation and backpropagation, and backpropagation is just computing gradients — which often play a fairly essential role.
Next we'll only consider RMS Norm. Its key operation is
\begin{equation}\boldsymbol{y} = \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert_{RMS}} = \sqrt{d}\times \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert}\label{eq:rms-norm}\end{equation}
where $\boldsymbol{x}\in\mathbb{R}^d$, and
\begin{equation}\Vert\boldsymbol{x}\Vert_{RMS} = \frac{\Vert\boldsymbol{x}\Vert}{\sqrt{d}},\qquad \Vert\boldsymbol{x}\Vert = \sqrt{\boldsymbol{x}^2} = \sqrt{\sum_{i=1}^d x_i^2}\end{equation}
So finding the gradient of $\boldsymbol{x} / \Vert\boldsymbol{x}\Vert_{RMS}$ is equivalent to finding the gradient of $\boldsymbol{x} / \Vert\boldsymbol{x}\Vert$, which we can compute as follows:
\begin{equation}\frac{\boldsymbol{x}+\Delta\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} = \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} + \frac{\Delta\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} \approx \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} + \frac{\Delta\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert}\label{eq:exp-1}\end{equation}
The more complicated part is expanding $\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert = \sqrt{(\boldsymbol{x}+\Delta\boldsymbol{x})^2}$:
\begin{equation}\begin{aligned} &\,\sqrt{(\boldsymbol{x}+\Delta\boldsymbol{x})^2} \\ \approx&\, \sqrt{\Vert\boldsymbol{x}\Vert^2+2\boldsymbol{x}\cdot\Delta\boldsymbol{x}} \\ =&\, \Vert\boldsymbol{x}\Vert\sqrt{1+2\boldsymbol{x}\cdot\Delta\boldsymbol{x}/\Vert\boldsymbol{x}\Vert^2} \\ =&\, \Vert\boldsymbol{x}\Vert (1+\boldsymbol{x}\cdot\Delta\boldsymbol{x}/\Vert\boldsymbol{x}\Vert^2) \end{aligned} \quad \Rightarrow \quad \begin{aligned} \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} \approx&\, \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert}(1-\boldsymbol{x}\cdot\Delta\boldsymbol{x}/\Vert\boldsymbol{x}\Vert^2) \end{aligned}\end{equation}
Substituting into equation $\eqref{eq:exp-1}$ gives:
\begin{equation}\frac{\boldsymbol{x}+\Delta\boldsymbol{x}}{\Vert\boldsymbol{x}+\Delta\boldsymbol{x}\Vert} - \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert} \approx \frac{\Delta\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert} - \frac{(\boldsymbol{x}\cdot\Delta\boldsymbol{x})\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert^3}\quad\Rightarrow\quad\nabla_{\boldsymbol{x}} \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert} = \frac{\boldsymbol{I}}{\Vert\boldsymbol{x}\Vert} - \frac{\boldsymbol{x}\boldsymbol{x}^{\top}}{\Vert\boldsymbol{x}\Vert^3}\end{equation}
Finally, substituting back into equation $\eqref{eq:rms-norm}$ gives
\begin{equation}\nabla_{\boldsymbol{x}} \boldsymbol{y} = \sqrt{d}\left(\frac{\boldsymbol{I}}{\Vert\boldsymbol{x}\Vert} - \frac{\boldsymbol{x}\boldsymbol{x}^{\top}}{\Vert\boldsymbol{x}\Vert^3}\right) = \frac{1}{\Vert\boldsymbol{x}\Vert_{RMS}}\left(\boldsymbol{I} - \frac{\boldsymbol{y}\boldsymbol{y}^{\top}}{d}\right)\label{eq:rms-norm-grad}\end{equation}
Enter DyT!
Note that $\boldsymbol{x},\boldsymbol{y}$ are both vectors, so $\nabla_{\boldsymbol{x}} \boldsymbol{y}$ is a matrix (the Jacobian matrix). Now let's consider finding an element-wise approximation for RMS Norm, i.e., one whose components are each computed independently:
\begin{equation}f(\boldsymbol{x}) = [f(x_1),f(x_2),\cdots,f(x_d)]\end{equation}
This independence means its Jacobian matrix must be diagonal! We want this approximation to preserve as much of RMS Norm's gradient as possible, so we consider keeping the diagonal part of equation $\eqref{eq:rms-norm-grad}$:
\begin{equation}\frac{dy_i}{dx_i} = \frac{1}{\Vert\boldsymbol{x}\Vert_{RMS}}\left(1 - \frac{y_i^2}{d}\right)\label{eq:ode-1}\end{equation}
If we further assume that $\rho = \Vert\boldsymbol{x}\Vert_{RMS}$ is a constant, then we can directly solve the above differential equation to obtain
\begin{equation}y_i = \sqrt{d}\tanh\left(\frac{x_i}{\rho\sqrt{d}}\right)\end{equation}
This gives us the T in DyT ($\tanh$), where the initial condition chosen for solving the equation is $y_i(0)=0$.
DyT effectively absorbs the preceding $\sqrt{d}$ into the parameter $\boldsymbol{\gamma}$, and then treats the $\frac{1}{\rho\sqrt{d}}$ inside the parentheses as a trainable parameter $\alpha$, which alleviates the restriction imposed by the assumption that "$\rho = \Vert\boldsymbol{x}\Vert_{RMS}$ is a constant." However, in my view, it might be more valuable to explicitly retain $\sqrt{d}$, and simply treat the $\frac{1}{\rho}$ part as a trainable parameter.
DyISRU
I wonder if you've noticed that for RMS Norm we always have $y_i = x_i / \Vert\boldsymbol{x}\Vert_{RMS}$, so in equation $\eqref{eq:ode-1}$, $\Vert\boldsymbol{x}\Vert_{RMS}$ can be replaced with $x_i/y_i$, giving us
\begin{equation}\frac{dy_i}{dx_i} = \frac{y_i}{x_i}\left(1 - \frac{y_i^2}{d}\right)\label{eq:ode-2}\end{equation}
This is an equation involving only $x_i,y_i$, which eliminates the need to approximate $\Vert\boldsymbol{x}\Vert_{RMS}$. Solving this equation gives
\begin{equation}y_i = \frac{\sqrt{d}x_i}{\sqrt{x_i^2 + C}}\end{equation}
where $C$ is an arbitrary constant. This form has a name — ISRU (Inverse Square Root Unit, which we've also previously called SoftSign) — coming from the paper Improving Deep Learning by Inverse Square Root Linear Units (ISRLUs). If we treat $C$ as a trainable parameter, then by analogy with DyT, we can call this DyISRU (Dynamic ISRU).
Going from the gradient $\eqref{eq:rms-norm-grad}$ to the equation $\eqref{eq:ode-1}$ and then to $\eqref{eq:ode-2}$, DyISRU is the best result achievable using an element-wise function, since no additional approximation beyond the diagonal assumption has been introduced. In terms of form, DyISRU is actually also more intuitive than DyT, because $\Vert\boldsymbol{x}\Vert_{RMS}^2$ is $\mathbb{E}[x_i^2]$ — since we're seeking an element-wise operation, we simply need to replace $\mathbb{E}[x_i^2]$ with $x_i^2$, and finally add a multiplication by $C$ and $\sqrt{d}$ as a smoothing operation:
\begin{equation}\frac{x_i}{\sqrt{\color{red}{\frac{1}{d}\sum\limits_{i=1}^d x_i^2}}}\quad\to\quad \frac{x_i}{\sqrt{\color{green}{x_i^2}}}\quad\to\quad \frac{\color{orange}{\sqrt{d}} x_i}{\sqrt{\color{green}{x_i^2} + \color{orange}{C}}}\end{equation}
Related Work
$\tanh$ and ISRU can both be viewed as smooth approximations of the sign function, and based on them, we can construct smooth approximations of the $\mathop{\text{clip}}$ operation, for example
\begin{equation}\mathop{\text{clip}}(x, -t, t) = \left\{ \begin{aligned}t,&\,\,\, x > t \\ x,&\,\,\, x\in[-t,t] \\ -t,&\,\,\, x < -t\end{aligned} \right.\quad\approx\quad t\tanh\left(\frac{x}{t}\right)\triangleq \mathop{\text{softcap}}(x, t)\end{equation}
From this, we can also understand DyT as introducing a (smooth) $\mathop{\text{clip}}$ operation to prevent the forward pass from exploding, thereby stabilizing the model.
$\mathop{\text{softcap}}$ was proposed by Google's Gemma2, where it was originally used on the Attention Logits matrix before the Softmax, to prevent excessively large logit values. However, in our own experiments we found that although the logits after $\mathop{\text{softcap}}$ won't explode, the logits before $\mathop{\text{softcap}}$ still risk exploding — so using $\mathop{\text{softcap}}$ to prevent logit explosion is really just relocating the problem elsewhere, treating the symptom rather than the cause.
I'm not sure whether Google later also became aware of this issue, but in the latest Gemma3, they chose to remove $\mathop{\text{softcap}}$ and switch to QK-norm instead. Our own experiments also show that QK-norm can better suppress the growth of Attention Logits. This change and conclusion actually once again indirectly convey a pessimistic signal: operations like DyT and other $\mathop{\text{softcap}}$-type approaches are unlikely to fully replace Normalization in practice.
Summary
This post analyzed, from the perspective of gradient approximation, what kind of element-wise activation function could (to some extent) replace the Normalization layer, and from this we were able to derive both DyT and a new result.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.