Reproduce “Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Kirato Yoshihara
Research from The University of Osaka revealed that Transformer attention and MLP layers benefit from distinct weight-space geometries during optimization, with Stiefel constraints enhancing stability for attention and DGram constraints being compatible with MLP layers, leading to the lowest stable validation loss of 3.3544 for a GPT-2 model.
Key insights
- Transformer optimization benefits from a heterogeneous application of manifold constraints, with attention layers preferring strong spectral control (Stiefel) and MLP layers tolerating more scale freedom (DGram).
- Unchecked singular value growth in DGram-constrained attention layers amplifies attention logits, leading to softmax saturation and gradient degeneracy, which causes training instability.
- MLP layers tolerate DGram constraints due to their element-wise non-linearities (e.g., GELU) and architectural elements like Layer Normalization, which can mitigate activation scale changes.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.