Question about computational invariance in LayerNorm

[Wangguanru1]

Suppose the normalization layer is not RMSNorm, but LayerNorm with the mean term μ(x) subtracted during calculation. Additionally, there is a matrix addition operation after the scaling operation. does this computational equivalence still hold?

[Wangguanru1]

If LayerNorm is replaced by AdaNorm in the diffusion model, and the original fixed learnable scaling and translation coefficients become dynamic with the conditional input, then is it impossible to ensure that there is no additional computational cost while maintaining computational equivalence?

[sashkboos]

Thanks @Wangguanru1

Let me explain this for LayerNorm (LN): In LN, we have the following sequence of operations: mean-subtraction -> normalization -> scaling-operation (see Figure 2 in [1]). The above can be considered as linear ops -> normalization -> linear ops. We can then fuse linear ops into the adjacent layers (see Figure 3 in [1]).

In general, as long as you have linear operations (like the above example), you can fuse them into the adjacent layers and CE is held.

Please let me know if you have any other question

[1] https://arxiv.org/pdf/2401.15024

[Wangguanru1]

Thanks very much for your explanation @sashkboos ! I still have two questions.
Q1. For LayerNorm, I understand that the operation of multiplying X by Q and subtracting the mean can be incorporated into the weight matrix of the previous layer. And these operations do not incur additional computational overhead. However, during normalization, std is calculated using XQ, but due to previous computational fusion operations, the value of XQ is no longer available. How is normalization implemented?
Q2. In the second linear ops, there is usually a bias term. Given that the LayerNorm's scale term and $Q^{-1}$ are pre-incorporated into the subsequent weights, how should this bias term be handled?

[sashkboos]

@Wangguanru1

Please check this function for the detailed implementation:

https://github.com/microsoft/TransformerCompression/blob/6b12cdee6ad51791d7c776b3a046bc408b9e77e9/src/slicegpt/layernorm_fusion.py#L137

[Wangguanru1]

Thank you very much! I understood how to process Layernorm from SliceGPT.
Finally, I'd like to know your opinion on this issue.

In diffusion model, adaLN_modulation is added after LayerNorm , then conditional injection modifies activations between layers, is it impossible to ensure that there is no additional computational cost while maintaining computational equivalence?