Transformer updates separate direction and scale for robust learning
Disentangling Representation Evolution in Transformers through Directional Decomposition
Computation and LanguageMachine Learning
Summary
Transformers learn by changing their internal data representations step by step. The authors show that these changes can be split into two parts: one that keeps the same direction and one that changes it. They find that the part keeping the direction is surprisingly important and behaves differently in different parts of the model. By understanding this, they can make the model’s training more stable and fix errors better when compressing the model. This insight helps improve how transformers learn and are edited.
What this means in practice
- •For machine learning engineers: Improve transformer training by suppressing certain update components to achieve better validation loss and downstream performance.
- •For model compression teams: Diagnose and reduce update errors during compression by analyzing component-wise update differences to maintain model accuracy.
Authors
Shwai He, Haichao Zhang, Shen Yan
Abstract
Transformer representations evolve through learned additive transformations that either preserve their current direction or redirect it. We study this evolution as a functional geometry, decomposing learned updates into parallel and perpendicular components. Across pretrained models, we find substantial parallel components beyond the residual identity path. We then apply the decomposition in two spaces: to attention and MLP updates relative to the hidden state, and to attention value aggregation relative to the current token's value. Targeted edits reveal a strongly space-dependent asymmetry: exclude-self value-space parallel manipulation is markedly more robust than residual-space and perpendicular counterparts, preserving the direct self message while scaling only the non-self aggregate. The same decomposition gives a component-resolved description of compression-induced update error: perpendicular error separates compression methods more clearly than parallel error. Extensive experiments further demonstrate that full-aggregate parallel suppression during from-scratch pretraining lowers validation-loss trajectories and improves downstream averages, with the value-space variant strongest. Together, these results connect representation geometry to editing robustness, compression diagnosis, and training-time intervention. Code is available in the \href{https://github.com/Shwai-He/Transformer-Geometry}{project repository}.