# Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

arXiv: https://arxiv.org/abs/2605.07588 · HTML: https://mlsanity.com/paper/2605.07588/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 44 of its 54 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 30 of 44)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [s1: Simple test-time scaling](https://mlsanity.com/paper/2501.19393/) (2025)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Training Large Language Models to Reason in a Continuous Latent Space](https://mlsanity.com/paper/2412.06769/) (2024)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Convolutional Sequence to Sequence Learning](https://mlsanity.com/paper/1705.03122/) (2017)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [Pay Attention to MLPs](https://mlsanity.com/paper/2105.08050/) (2021)
- [Dense Associative Memory for Pattern Recognition](https://mlsanity.com/paper/1606.01164/) (2016)
- [Kimi Linear: An Expressive, Efficient Attention Architecture](https://mlsanity.com/paper/2510.26692/) (2025)
- [Primer: Searching for Efficient Transformers for Language Modeling](https://mlsanity.com/paper/2109.08668/) (2021)
- [COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics](https://mlsanity.com/paper/2202.11705/) (2022)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)

## Abstract (arXiv)

Transformer blocks typically combine multi-head attention (MHA) for token mixing with gated MLPs for token-wise feature transformation, yet many choices in their parameterization remain largely empirical. We introduce Causal Energy Minimization (CEM), a framework that recasts Transformer layers as optimization steps on conditional energy functions while explicitly accounting for layer parameterization. Extending prior energy-based interpretations of attention, CEM shows that weight-tied MHA can be derived as a gradient update on an interaction energy, and that a gated MLP with shared up/down projections can be viewed through an element-wise energy. This perspective identifies a design space for Transformer layers that includes within-layer weight sharing, diagonal-plus-low-rank interactions, lightweight preconditioners, and recursive updates. We evaluate CEM-derived layers in language-modeling experiments at the moderate hundred-million-parameter scale. Despite their constrained parameterizations, these layers train stably and can match corresponding Transformer baselines. Overall, our results suggest that CEM provides a useful lens for understanding Transformer layer parameterization, connecting Transformer architectures to energy-based models and motivating further exploration of energy-guided layer designs.