Free Energy Mixer: Value-Aware Attention Readout for Efficient Sequence Models
Jiecheng Lu ⋅ Shihao Yang
Abstract
Standard attention stores keys/values losslessly but reads them via a per-head convex average, blocking channel-wise selection. We propose the Free Energy Mixer (FEM): a free-energy (log-sum-exp) read that applies a value-driven, per-channel log-linear tilt to a fast prior (e.g., from queries/keys in standard attention) over indices. Unlike methods that attempt to improve and enrich the $(q,k)$ scoring distribution, FEM treats it as a prior and yields a value-aware posterior read at unchanged complexity, smoothly moving from averaging to per-channel selection as the learnable inverse temperature increases, while still preserving parallelism and the original asymptotic complexity ($O(T^2)$ for softmax; $O(T)$ for linearizable variants). We instantiate a two-level gated FEM that is plug-and-play with standard and linear attention, linear RNNs and SSMs. FEM is not an efficiency method: it does not change the underlying mixer's complexity. Instead, for a small readout-side cost added on top of a linear-time prior (GLA, Mamba, linear RNN/SSM), it narrows the quality gap to quadratic softmax attention while keeping the prior's complexity class, parallel training, and streaming inference. It consistently improves strong baselines on NLP, vision, and time-series at matched parameter budgets, and we report FEM-GLA/FEM-Mamba results alongside L40S latency and throughput to show the added readout cost is modest.
Chat is not available.
Successful Page Load