Question about RMSNorm scale handling in weight-merged LaRoSA

Hello LaRoSA authors,

First of all, thank you for sharing this very interesting paper and the public implementation. I found the idea of using layer-wise rotations for activation sparsity very insightful, and I have been reading both the paper and the EfficientAI/LaRoSA code with great interest.

I had a question about how the RMSNorm scale parameter is handled when the rotation matrix is merged into the weights. As I understand it, the paper explains that the orthogonal rotation matrix $Q$ can be merged into the weight matrix to avoid additional runtime rotation cost. However, from the public code, it seems that the evaluation path applies the rotation to the activation after RMSNorm, rather than using a fully weight-merged inference path. I may be misunderstanding the code, but it looks closer to an accuracy/benchmarking implementation than a fully merged runtime implementation.

My question is about the learned scale parameter in LLaMA-style RMSNorm. An orthogonal rotation preserves the RMS value, so the normalization part itself seems fine. However, RMSNorm also applies a learned per-channel scale. Since this scale is channel-wise, it generally does not commute with the rotation.

In other words, if $\Gamma = \mathrm{diag}(\gamma)$ is the RMSNorm scale, then generally:

$$
Q\Gamma \neq \Gamma Q
$$

So I am not sure whether the following equivalence still holds exactly when the learned RMSNorm scale is included:

$$
\mathrm{RMSNorm}(xQ)Q^T \overset{?}{=} \mathrm{RMSNorm}(x)
$$

If the implementation explicitly computes the rotated activation after RMSNorm, i.e., something like:

$$
\mathrm{RMSNorm}(x)Q
$$

then this issue does not arise, because the RMSNorm scale has already been applied before the rotation.

However, if the residual stream is kept in the rotated basis and the rotation is merged into the weights as described in the paper, I am not sure how the learned RMSNorm scale is handled.

Could you clarify whether the fully weight-merged version assumes that the RMSNorm scale is folded into adjacent weights, ignored in the commutation argument, or handled in some other way?

Thank you again for the great work and for making the code available. I would really appreciate any clarification on this point.
Please also excuse any awkward wording, as English is not my first language and I used AI assistance to translate my question.

Start by reading the public implementation's evaluation path, especially where RMSNorm output is rotated, and compare it with the paper's fully weight-merged inference description. Trace how the learned RMSNorm scale is treated in the rotated-basis case. Done means documenting whether the equivalence is exact and explaining how the scale is handled; no file or test is named in the issue.

Written by the indexing model from the issue text.