Free Energy Mixer
Authors: Jiecheng Lu
Shihao Yang
Venue: ICLR 2026 Poster (2026)
Keywords: Sequence Modeling
Attention
Transformer
Primary area: foundation or frontier models, including LLMs

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. It consistently outperforms strong baselines on NLP, vision, and time-series at matched parameter budgets.

Published as a conference paper at ICLR 2026
FREE ENERGY MIXER
Jiecheng Lu, Shihao Yang Georgia Institute of Technology 3lu4 1u4414@ga 4egatech.edu. shihao baaoo..angllln wwt.ee..atcc......
ABSTRACT
Standard attention stores keys/values losslessly but reads them via a per-head con- vex average, blocking channel-wis 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 a fast prior (e.g., from queries/keys in standard attention) over indices. Unlike methods that attempt to improve and enrich the (g, k) scoring distribution, FEM treats it as a prior and yields value-aware posterior read 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²) for softmax; O(T) for lineariz- able variants). We instantiate a two-level gated FEM that is plug-and auugaand with standard and linear attention, linear RNNs and SSMs. It consistently outperforms strong basclines on NLP, vision, and time-series at matched parameter budgets. The code implementation is available at this link
INTRODUCTION
Transformers, powered by attention mechanisms, have become the default backbone for sequence
modeling across language, vision, speech, and decision making (Vaswani 2017
Devlin et al. 2019
Radford 2018 Brown et al. 2020 Dosovitskiy et al. 2020 Dong et al. 2018 Chen et al. 2021
Touvron et al. 2023). Their success often linked to selective access to ever- wver-growing key-val value
cache while retaining parallel training and inference. In large language models, this selective abil- ity, composed across multiple attention layers and residual pathways, supports long-range memory retrieval and the algorithmic behaviors associated with in-context learning (for example induction heads and pattern completion), as shown by recent empirical and mechanistic studies (Min et al. 2022
Wei et al. 2023 Xie et al.) 2022 Zhang et al. 2023 Garg et al. 2022; Akyürek et al. 2023
Li et
2023 Dai et al. 2023 Bai et al. 2023 Olsson et al. 2022] Elhage al. (021).
Causal softmax attention combines strong selectivity with parallel efficiency: at each step it forms distribution over past indices and mixes their values, while all steps can computed in parallel. Given (Q, K, V) E Rxxd with rows qt. ki, V;; define masked scores 8t,i q: k; /d for < t and -oo otherwise, and set at, = softmax(s,,) E A where 1 denotes the probability simplex over {l....,,5 f} The step-t read is o, = Eist dt, U; O, E ( J. and stacking all t yields 0 = AV with Ar,; at,i, so single matrix multiply produces all outputs.
The convex-mixture view explains efficiency: outputs are probability-weigheedd averages of the shared value bank, computed in one matrix multiply. Yet this also reveals a lossless-storage versus lossy-processing dilemma (Fig. [p). The KV-cache stores full context, but the read is lossy: each head applies the same weights to all coordinates of v; so O, Eist (, lies in conv (ow,(,,....... ).) and all channels are synchronized. As result, even simple per-channel indexing, such as :; = (1) U...D) (e.g., coordinate-win wise argmax), cannot represented unless all chosen indices coincide. Adding more heads only creates a few synchronized groups, and deeper stacks cannot recover per-channel index identity once the first convex mixing has occurred. This limitation hinders Transformers in long-range modeling with non-sequential or irregular timestep indexing, and in tasks where channel- -wie structure critical, such as multivariate time series mod-
eling TTy al. 2020 Zeng et al. 2023 Nie et al. 2022 Liu et al. 2024 Lu et al. 2025).
Most recent advances in attention aim to improve expressivity and efficiency, typically by designing richer selective distributions but still reading values through token-separable linear combination.

---

Published as a conference paper at ICLR 2026
Query
(1)
Query
Kays
Hightt tarpesider powartor Spacchannt
Free Energy Mixer (FEM)
log E p:60 explst
beaudae Gonamiis
Σ, P1000
Valuee
Kays
Exisdin Fixes Lossy Dilmmn??
Hoden Ctrte
Selection Diambbution
Corme
Combiratio
Selection Distribution
Values
Capurdyy Coriggs
Lonakess
Farory Strege
Crassevvise
Morn Complesiyy?
Vartetoonn Postwier Salection
Values
Mersry Storage
Other Multi dini--iinnniooaal
Special Case: FEM can solve charnel via value-inite inttcid optimal titing of the prior
Lasry Mamory Oiamrmas approdmasing the Simplest Indasing on storagg
Ban Londona Me Jurige
Liwaar 1aN
(a) Classic attention stores past values losslessly but reads them as a single convex com- bination, so channel-wise indexing (e.g.. per-channel argmax) is not representable. (b) Free Energy Mixer (FEM) treats selection as a DV free-energy problem: values tilt the prior to a value-aware posterior with learnable per-channel temperature, enabling low-entropy (point-like) posteriors and channel-wise selection while preserving the prior's time complexity. (c) Common fixes (more heads, deeper stacks, separable mixers, and per-channel scoring) either keep channels synchronized or raise cost / rely on fixed-state storage; none close the lossy-meoory gap that FEM addresses.
Figure 1:
These methods include sparsity (Beltagy et al. 2020 Child et al. 2019 Zaheer et al., 2020 low-
al. 2021 and kernelizable variants with normalization or gating (Katharopoulos et al. 2020 Choromansko et al. [021 Hua et al. 2022 Yang et al. 2024bb work accelerates via factorized implementations (Dao al., 2022 Dao 2023) or replaces the cache with streaming SSM and RNN models of fixed size (Gu & Dao 2023 Sun et al.) (003).. Across these lines, computation faster or the distribution richer, but the read remains a linear mix, so channels share weights, and even simple per-channel indexing (e.g.. argmax) cannot be realized one step. Some recent works explore more complex combinations (e.g.. nonlinear mixing such as log-sum-exp in LASER attention, or hard/top-k selection (Guptal ctt al. Duvvuri & Dhillon 2025) Hashemi et al. 2025). SJ.Ip, D. yet these mainly target training stability or accuracy in specific cases and do not address the lossy processing limitation.
rank projections (Wang et al. |020
Qin et a 2022a b). Efficiency-
Motivated by this gap, we propose the Free Energy Mixer (FEM), which regards lossless processing as the optimal interaction between selection distribution and stored values: for each channel, choose an index distribution that maximizes utility under an information budget. FEM removes the linear read bottleneck and enables per-channel, context-dependent selection from the memory, while keeping causal parallelism and the time complexity of the underlying mechanism. When channel selection is not needed, FEM reduces to the standard expectation; when is, different channels can focus on different past indices in the same step. FEM consists of four components: temperature gating (T), LSE mixing (L.), outer gating (G), and low-rank convolution (C), described in 22.3
Contributions. (1) We identify lossless- memory processing gap in attention: per-head convex mixing cannot realize channel-wise selection from the lossless KV-cache. (2) We propose FEM, which closes this gap by casting the read as a variational free-energy optimization that, per chan- nel, selects an index distribution under an information budget, enabling value-aware channel-wise selection. (3) FEM is agnostic to how the selection distribution is formed (softmax, kernel/lowrrak attention, linear RNNs, SSMs) and preserves the corresponding time complexity. (4) On NLP, vi- sion, and time-series tasks, FEM consistently improves strong baselines at matched parameter sizes.
2 METHODOL.OGY
2.1 PRELIMINARIES: SELECTION DISTRIBUTIONS
To analyze the storage- proccssiig gap. we introduce the following notion of a selection distribu- tion. At step I, we formalize memory selection over past indices I2 (1....,)) by probability
2

---

Published as a conference paper at ICLR 2026
vector pt E A with support M, {i E Ie pe(i) > 0},
pr (i) 2 E' Pc()) =1
(1)
Hard masks such as local window can be encoded by restricting Mr Given values v, E R., the per-step readout is the expectation
0f P:(1) = Ei-p, [e:] E conv (01,..., 0
(2)
Causal softmax self- attention is the case where Pt is a masked row-softmax over logits q, k:/Vd; linear attention arises when pt is a normalized nonnegative kernel, as detailed in $B.1
Lossless storage vs. lossy processing. Unlike RNNs, which compress history into a fixed-size state, softmax attention stores the full VV-cache {(k; v:)) without compression (lossless storage), but the read equation equation@aapplies one weight vector per head to all coordinates, outputs lie in per-head convex hull. This is potentially lossy when different channels should retrieve different indices in the same step. To state the target capability we define the finest-granularity retrieval:
Definition 2.1 (Channel-wise selector). A channel-wise selector at time t is any vector s (% (1......... ,Vb,.D with i, I allowed to differ across j [D).
Lemma 2.2. Let m, (maxist V/,1:.... max,<t V.,D). If m, E conv(u)..... then a single index simultaneously atlains all coordinate maxima. Hence if the arg-max indices differ across coordinates, m. & conv (oi........;
Corollary 2.3. A per-head convex read E, Pe());; cannot realize a generic channel- selector with at least two coordinates selecting different indices.
This geometric limitation above motivates our method. We can see that a single head applies one se- lection distribution all channels step t. synchronizing channel-wise index choices; with H heads the number of realizable head-level arg max patterns is at most t" far below the t patterns needed for lossless per-channel selection when H < D. This gap motivates replacing the expectation read equation 2 with the free-energy read in Section 2.3 Proofs of Lemma 2.2 and Corollary 2.3 the !!
capacity counting are deferred Appendix
2.2 WHY STANDARD REMEDIES FAIL: TOWARD A FAITHFUL, LOSSLESS READ
We revisit common extensions around attention and explain why they do not close the channel-wise lossless-selection gap, as shown in Fig. D Full details and proofs are in Appendix C
(1) More heads. Heads provide H selection distributions per layer but synchronize channels within each head. Hence the step-t head-level argmax capacity is at most ! far below to when H < D.
(h) E A be the distribution of head h E [H]. Across contexts, realized Lemma 2.4. Let oi,-
head-level argmax assignments are at most t and all coordinates controlled by head share at,
Increasing reduces the per-head width d, = D/H, tightening the low-rank bottlenek on the value path; as H approaches D, the cache become well-appproimmated by finite-state linearization, effectively breaking the lossless memooy advantage. See Appendix C.ramddxxC..ceee. for details and analysis.
(2) More depth. After first per-head convex mixing acts at step t, per-channel index identities are no longer available unless a fresh, independent selection acts before that first mixing.
Lemma 2.5. The map (vi)ist + E: (r, is row-stochastic with image in conv (vi...., v.). Any channel-wise selector outside this hull cannot be realized at step by composing coondinate-wise maps and later attentions that only access already mixed tokens.
Proposition 2.6 (Selection budget). With attention-MLL MLP blocks and heads per block, at most H disjoint channel groups receive independent firsi-mixing distributions by step L. A necessary condition for independent per-channel selections at step t is HL 2 D (which is not practical).
(3) Per-dimension queries/keys. Giving each coordinate its own scoring subspace raises capac- ity toward p but raises score parameters and compute from (dd)) to O(Dd) per layer, typically harming value bandwidth or MLP width under fixed budgets.
(4) Richer in-head mixers. The progressive family below still keeps mixing token-separable at, ① vz)
0 = Σ, U;
Ee an (B. vz) Σ. San,, vi);
3

---

Published as a conference paper at ICLR 2026
(A)
Linear Loverank Cosv
(d)
Two-Level Gated Free Energy Mixer
Energy Destrpposiion and Linearized Temperature
Uighiiight
(Final Architecture)
Livter
Cone
Rp Li Pyl0te ;- (1 P1
KLij (Pael) Plok
me (1 Ac/ap A Apar
nomallrr
normilaa
Input (1x))
Prosx
[mx)
Sowinng
Beotng
Linear & Low-rank Local Conv
(
LSE implementaaion of Posteriar Selection
Pv( (/)
A:91
6 Paux
fe
FEM Transformer Block
Prior Distribution Parameterization
A
91
Pell from: eiuu way)
Bexm
Mixinn rvll
Linear Asmn | [Lineor RHK ISM
Variational Posterior
Norm
FEM
pa(f)
Nx
Exyle
Norm
MLP
Aais
Linear
octput
70
Figure 2: Overview of the Two-Level Gated Free Energy Mixer. (a) Lightweight linear & low-rank local convolution for local conditioning. normalizer, while linear RNN/SSM use an operator-induced normalizer. (c) FEM integrated into a Pre-Norm Transformer block. (d) Final architecture: compute mean μ4 and max-temperature branch Frusx, with inner gate X₂ interpolating and outer gate 9t scaling. (e) Free-energy curve: improvement (I) Efficient implementation: one mixing with pt yiclds both Ep. [u] and 8 max log Ep, [e [mmxx)), then gating produces or.
(b) Prior selection: softmax attention uses probability
over br equals (L [[q(0) )BB.
Proposition 2.7 separable mixers are convexly constrained). Linear and coordinate -ise gated variants lie in a convex hull of transformed values; even with pointwise nonlinearity inside the sum, channel-wise selection of the original coordinates is not realizable in general. For general token-separable f. per-channel argmax is impossible in general. Additionally, adding per-channel cross-token competition in ff may break O(TYO(T²) parallelism. Details in Appendix C.4
They offer rich dimension interactions but store history in a fixed-size state, cannot support arbitrary index retrieval at large horizons without lossless storage; see $
(5) Linear RNNs/SSMs.
Takeaway and connections. Prior remedies fall into three buckets: (a) increasing assignment capacity substantial cost (e.g., per-feature score-space inflation to obtain ar,i,). (b) keeping a token-separable convex read (e.g.. in-head pointwise gates), or (c) relying on fixed-state storage (e.g.. linear RNNs/SSMs). None provides per-channel, value-aware cross-token competition before the first mixing while preserving the time complexity. In particular, pushing capacity from tl to- ward iD via per-feature inflation leaves the read token-separable, so the same-step lossless-selection gap persists (Lemma 2.5 Proposition 2.7p; likewise, simply scaling heads/depth or adding in-head gates cannot recover channel-wise index identity once the first convex mix has acted. These gaps motivate single, stronger mixer that performs value-aware competition without changing asymp- totic cost: our FEM via a variational free-energy read. See Appendix C.6 for a mapping of existing designs and Appendix C.7 for more discussion.
2.3 FREE ENERGY MIXER: VALUE-AWARE POSTERIOR SELECTION
Motivation and objective. Classic attention performs a per-head convex read and cannot real- ize same-step channel-wise selectors in general (cf. Lemma 2.2 Corollary 2.3). We therefore cast channel-wise retrieval as an informatiio-cconstrainedd selection problem: at step f, a fast, information- sparse prior Pt (from queries/keys or an operator-induced normalizer) proposes indices on the masked support M,, while values {v,} supply evidence For each channel j we choose q E A(M)
'Somewhat counterintuitively, we treat selection prior and values as evidence because evidence requires log-exp processing while the prior does not; this preserves the time complexity of the selection mechanism.
4

---

Published as a conference paper at ICLR 2026
to maximize expected utility under KL budget relative Pt,
maX B s.t. (T(glpp)) S BL..
(3)
Free Energy Mixer formulation. Introducing Lagrange multiplier Be, > 0 yields the per-channel free energy output
FBr()))) Be. log ΣieM, M,, Pr(i) exp(3e, (3e. vas),
(4)
and the corresponding posterior selection distribution
Pe(i) (xp()))))))) (1) Σrem, Pc(r) Pr(r) exp (B
i e M,.
(5)
Theorem 2.8 (Free-energy selection and budget duality (with 3 as inverse temperature)). The con- strained problem equation 3 has unique solution q" There exists a unique Pt.3 2 0 such that 9¹ = - and Egg [viz] = F8t Equivalently (DV form), for any 8 > 0 the maximizer of Σ. (0)) 1 (Llll) is qr,3* Moreover, B + FL)(3) is continuous and strictly increasing E.2 and Proposition E.3
unless V.j. is P.S.. constant. See Appendix Lemmas E.
Consequences (summary). (i) Reverse-KL improvement over the mean: (B) = Ep [w,g]] +
(T. ( (Proposition E.3 (ii) Value-aware competition: the gradient equals the posterior and
the Hessian is Fisher covariance scaled by 3; thus Fr.j is convex and / 8/2-smooth in v. (Proposi-
tion E.4b. (iii) Channel-wise selection on the prior support: with margin A.j > 0, (.50) concentrates the argmax with exponentially small error in B; (B) )) max; Vs, (Proposition E.5). (iv) Ca- pacity and complexity: across channels, FEM attains the assignment upper bound |.M, whereas H heads offer at most |M,] patterns; computing equation 4 with a fixed temperature adds one masked log- reg-sum-exp per channel and preserves the prior's asymptotic complexity (Theorem E.7 Proposition E.8). (v) Masks and invariances: masking is preserved; shift/scale laws and sensitivity to prior probabilities/logis follow from log-sum-exp structure (Proposition E.
Outputs. FEM exposes two per-channel readouts sharing the same posterior 4g,3-
()) the free energy Vi,g Ei Vi.j. Under B-concentration they coincide at the
(9) and the posterior expectation E, selected value-letting the model smoothly move from averaging to hard indexing without changing the architecture. In $ 2.3.2 we add a lightweight two-level gating and linearized temperature learning that learn dynamic temperature without changing the prior's asymptotic complexity.
Rethinking the Design of Attention. Attention can viewed as simplified subclass of a more general and computationally expensive map-reduce structure o = Σise g(x;, T1:1). The simplifica- tion occurs in the map stage: instead of computing full channel-wise function g(x,, Z14), attention retains the full memory state Fiat but replaces with a channcl-synchronize. form Eise it This avoids materializing large channel-wise weight tensor and reduces the read to a single scalar weight per position, enabling highly efficient parallelization. FEM restores the missing channel interaction not by increasing the internal complexity the map function g(-) (which would raise time complexity), but by enriching the reduce stage Et By replacing the linear read with free-energy read, FEM recovers channel-wise selection ability while preserving the computational efficiency of the original attention. this way, FEM closes the expressiveness gap introduced by the map-stage simplification of classical attention.
2.3.1 EFFICIENT COMPUTATION OF FEM AND LINEARIZED TEMPERATURE LEARNING
Fixed temperature. For a fixed inverse temperature 8 > 0 and channel j. FEM reads
(B) Log Po(i)
(())) e a Ei-p, [v
KL(p, ll (.))
(6)
with posterior selector 9:.6 (t) (t) u on the same support Mr the prior. Evaluating equa- tion requiree a single masked oug-sum-exp (LSE) per channel, so the asymptotic time complexity is identical the prior (e.g., O) for softmax, O(T) for kernel/SSM priors). See Appendix F.1
5

---

Published as a conference paper at ICLR 2026
Why 8 should be dynamic. The decomposition in equation o reveals energy meegy-entrooy trade-off: governs the improvement over the expectation baseline through [p Il Tasks typically need different entropy levels across steps and channels, but directly recomputing equation each learned Bt, would break single-past efficiency.
Linearized Temperature Learning (LTL). (Figure 2D Now we fix a per-channel maximum Smax 0 and define the expectation baseline Mt. Linp, [vs,]] and the high-temperature branch Fu Amax log EiEM, (i(i)emmaeut) A learned gate Ar,j [0, interpolates
()an) (1 Ae.) Pr. + A Fiyxx
requiring only the baseline expectation and single LSE at Smax per step. hence preserving the prior's asymptotic complexity.
Hidden temperature and equivalent reparameterization in LTL. (Figure 2 Let (t. = (0). Then (B) 3- (T Ui
(B) and (B) (3) continuous and strictly increasing on [0, Smax] unless v. is pr-a.s. constant. By the intermediate
lppe) 2 0, Fr, is
value theorem, for each At,5 [0, 1] there exists unique
(A.,) = AT,(Ar. (Smaax) e [0, Smas] such that
Fi, (s,t)) (8)
Therefore, optimizing A is a strictly monotone reparameterization of optimizing hidden temper- ature Pidd for equation 4 see Proposition F2
Final form of FEM and complexity. Collecting terms gives the phannhannel read
FA,( (1- EieM, Pr(i) (r()) Un.j +
At, log Eie M, Pr(i) 2.
equal F., the unique hidden temperature Ii 1,5 (Ar.)). (Ar.)) Both terms can be obtained in one pass by mixing (/.,f em with the same (t)... Hence LTL achieves dynamic temperature control without changing the prior's asymptotic complexity. A KL interpretation appear in
2.3.2 TWO-LEVEL GATED FEM: VALUE-AWARE INNER GATING AND OUTER MODULATION
We present the two-level gated FEM that turns prior selection distribution p: E - into per- channel, value-aware read while preserving the prior's time complexity. All operations below act element-wise over channels j E [D]; @ and Ø denote Hadamard product and division. Let B Bmax Rh be a learnable global maximum inverse temperature, and let X E [o, 1] and 9: E RSo be per-channel gates step parameterized from the current token features. We apply sigmoid and softplus activations, and normalize 9: with RMSNorm so that its modulation does not overly distort the norm of o₁- In what follows, whenever refer to FEM, we default to this two-level gated version. Proofs and details of this section appear in Appendix G
Inner (temperature) gate via one-pass linearized temperature learning. Define the expectation baseline and a single high-temperature branch
1 Σ, (1) vs E RD,
F.nax 9m1 GlogE, exp) (xmmx On.) RRD,
which can be obtained in one pass by mixing [v,, exp(3mmxx Rvi)) with (i). The inner gate as hidden temperature interpolates
FA(A.) (1 x=) + A Funax
(10)
Outer gate and final read. The outer gate modulates the inner read:
O, O. = 2 (1 11-x))pp + 6
(11)
Note that the outer gating can be regarded as applying an scaling after the token mixing in free energy with hidden temperature, i.e., Shale log ΣiEM, (1) exp(Shid,es me.s) 9 For smoother opti- mization, therefore parameterize 9: strictly positive by default. Computing equation 10 I1
matches the asymptotic time complexity of the prior pt (c.g., O(T2) for softmax; ((T) for ker- M priors) as shown in the section above.

---

Published as a conference paper at ICLR 2026
Containment of common mixer families. The two-level gate subsumes several widely used mixers:
(i) setting A = 0 yields per-channel linear reweighting o, = Σ, (i( (g. vv;;; (ii) 0 < X < 1 gives a monotone, convex mean mean->real--e -roaaas interpolation per channel, enabling value-aware thresholding; (iii) letting Xx, 9t depend on (etx, Pt. Mt, Fm ) realizes a broad token-separable class E; f(a,i, 02) while introducing cross-token competition through the log-sum-exp branch.
2.3.3 FEM AND SELECTION DISTRIBUTIONS: PRIOR-AGNOSTIC INTERFACE
FEM only requires a nounegative, normalized selection prior E A over indices Ie = (I. with masked support M, {i S t (c( 0}, and the variational read always en- forces qi < Pt. Any streaming or parallel mechanism that produces nonnegative scores sf (i) 2 0
*, (6) induces a valid prior via the normalization p² Erg. + (r)
(i (1..).
Proposition 2.9 (Complexityprresen) preservvng normalization). If s7 (t) is produced by associative operator (e.g., kernelized/tinear attention, linear RNN, or SSM) that admits an O(1) streaming update per step, then the denominator is obtained by applying the same operator to all-ones stream, so forming Pt. preserves the asymptotic complexity of the underlying mechanism. Under FEM with fixed L.TL-controlled lemperature, the read adds one masked sum-exp per channel on the prior support and thus preserves ( (softmax) or O(T) (linear/SSM) cost (See Fig. 2
Parameter budgeting. Let the input/value width be D and let FEM use working width on the value path. We allocate ratio r > 0 of parameters the prior (e.g., Q, K and, where applicable, a decay gate). Ignoring biases/norms, the per-head linear parameters decompose as 4Dd + Ddr, covering value, output, temperature and outer gates, and the prior block of size D x (rd). To match the classic 4 4D budget in standard attention: (i) d 2
(keeps Q, K width D): (ii) d 2D 3. r = 2 (balanced split). See Appendix H.. for the split and costs. In our experiments we default to (i) since it uses forward pass with exactly the same shape as standard attention. Notably, under (i) the value part of the KV-cache can in principle have half the dimensionality. Subsequent experimental results show that FEM's fine-gr graigna processing allows to achieve superior performance over priors while using an even smaller memory state cache.
Instantiations of 8t (and pt) We adopt the following FEM selection priors as examples. (i) Softmax attention recovers the standard masked row-softmax prior. (ii) Gated linear attention (Yang et al.) 2024b) keeps an associative O(T) form by combining a feature kernel with an input-conditioned decay. (iii) Linear RNNs admit nonnegative bilinear scores with normalization from the same currence. (iv) SSM/Mamba-styLe priors use nonnegative impulse responses; a variant lifts the index set to pairs (i, k) and normalizes per output channel, enabling cross-channel competition. All formulas, streaming recurrences, and complexity details appear in Appendix H
Low-rank convolution. Recent sequence models such as Mamba and DeltaNet (Gu & Dao 2023
Yang et al. 2024c a) variants commonly enhance feature parameterization with local convolutions. We adopt this idea FEM by inserting lightweight adaptive low-rank convolution module that produces local, position-sensitive features. Concretely, it forms a simple time-decay kernel with O(1) streaming updates, so the overall cost only H,) with the low-rank dimension H, D (Hl = d/16 16 by default). The resulting features modulate both the selection prior and the FEM gates, providing local adaptivity. See $ and
K for more details.
FEM as a universal fastht programmer. FEM provides a unified mechanism that upgrades expectation- based reads into value- aware, per-channel posterior selection while preserving the com- plexity. It combines temporal mixing, entropy control, local conditioning, and dual gating, thereby serving as an effective fast-weight programmer Schmidhuber (1992) detailed M
3 EMPIRICAL EVALUATION
We evaluate the two-level gated Free Energy Mixer (FEM) with different selection priors across synthetic, NLP, CV, and time-series tasks. Specifically, we test FEM with softmax attention (FEM- SM), gated linear attention (FEM-GLA), and on selected tasks also with Mamba (FEM-Mamba) and linear RNNs using AFT (Zhai et al. 202T) (FEM-AFT) (see $ H Unless otherwise noted, we use parameter budgeting strategy (i) from
2.3.3 which matches the parameter size of standard

---

Published as a conference paper at ICLR 2026
attention. Under this setting, FEM reuses existing efficient implementations (e.g., FlashAttention, TlashLineurAtteniie for the core prior mixer (see Fig. 21:f) with only minor value- value-path path overhead. Our main focus is algorithmic: exploring improved mathematical structures (see $ $. Due to lim- ited compute and lack of fused CUDA kernels, we scale models modestly but provide fine-grained metrics and extensive ablations to highlight FEM's advantages. For ablation, we denote FEM mod- ules as (C: low-rank convolution, L: LSE mixing. T: linearized temperature learning, G: outer gate). For example, FEM-SM (-G,T) removes outer gating and temperature learning, equivalent to SMAttn (+C,L). Unless specified, default FEM variants include all modules (C,L.T,G). We make sure that every variants have same parameter sizes with the parameter budgeting. Aside from causal autore- gressive FEM shown above, encoder-only use simply removes masking. In all experiments, FEM directly replaces the attention in a Transformer block (Fig. 2 without altering other components (MLPs, embeddings, hyperparameters) More implementation details appear in K datasets in $
MAD. We first evaluate FEM on the synthetic MAD benchmark (Poli ct al. 2024). which probes sequence models in-context tasks. As
Table I: MAD benchmark evaluation results across com-
pression, fuzzy/in-conter recall, memorization, robustness, and selective copying. Bold marks column best.
shown in Table [] FEM-SM out- performs all other baselines (Hyena, DeltaNet, Linear Attention, Mamba2, Gated DeltaNet, Differential Trans- former, Poli et 2023 Yang et al.
Com- Fuzzy In-Ctx Memorize Noisy Selective Copy
Model
press Recall Recall TrainSet Recall
Avg
Hyena
44.8
14.4
99.0
89.4
98.6
93.0
73.2
DeltaNet
42.2
35.7
99.9
52.8
99.9
99.9
71.7
LinAttn
33.1
8.2
91.0
74.9
75.6
93.
62.6
Mamba2
43.6
21.1
96.4
86.9
96.7
95.3
73.0
2024b C Dao & Gu 2024; Yang 2024a; Ye et al. 2025 by
GatedDeltaNet
45.0
29.8
99.9
80,2
99.9
94.3
74,9
et al.
DiffTrans
42.9
39.0
99.9
83.7
97.1
95.8
76.4
clear margin. In particular, different
FEM-SM ME--SO--TLO
44.3
24.5
99.9
85.7
98.5
95.1
74.7
FEM variants show strong gains on the Compress Recall tasks, which heavily rely algorithmic handling of dynamic context and channel inter- actions. On the Compress task, FEM models achieve significant improve- ments over existing methods thanks to their finer-grained processing of context storage. The ablation study further reveals that the two major per- formance jumps over prior baselines occur after introducing +1. (LSF)) and (temperature), corroborating our earlier discussion of FEM's enhanced memory storage processing. More- over, the ablations demonstrate that FEM can elevate linear-time methods such GLA and Mamba (with nor- malized pt to a level comparable with the latest attention-based vari-
FEM SM--aTL
45.0
31.4
99.9
85,5
99,9
96.3
76.3
FEM SMGGT)
(SMAm++C.)
50.3
39.0
99.9
85.4
99.9
98.0
78.8
PEM-SMIG
(SMA
52.3
85.8
79.4
FEM-SM
(SMAB+CI. CT.))
53.1
43.1
99.9
85.9
99.9
99.3
80.2
CG.
(SMABu4))
49.5
26.3
99.9
85.7
97.5
97.5
76.1
FEM-SMc-CO
(SMAMarL.T)
50.7
32.8
99.9
85.7
98.0
97.6
77.5
FEM FEM-SM:O
L.G)
51.2
35.4
99.9
85.9
98.5
99.0
78.3
FEM-SM
(SMA
53.1
99.9
85.9
99.3
80.2
GLAAGLLC
40.2
8.5
91.3
81.3
86.8
76.8
64.2
GLACGTL))
(GLA+C)
47.1
9.4
91.7
83.4
92.5
88.5
68.8
PEM-GLA(GT
(GLAICL, s, (0)
51.2
12.4
92.2
85.1
92.4
89.2
70.4
FEM GLA64)
(GLAICL.T (()
51.9
13:2
97.1
86.1
93.5
91.4
72.2
FEM-GLA
(Gl (GL.CCL...// (P)
53.0
19.1
99.9
86.3
99.9
99.0
74.9
TFM-MAMBALGLTU
(Mamba)
52.7
90.4
89.5
90.1
86.3
69.3
FEM- MAMBAG- nom)
ants.
MMmbaCLT (G.,7(0) (7) FEM-MAMIE
50.5
12.8
93.4
88,9
86.3
92.2
70.7
Language Modeling. We follow the experimental setup of (Yang et al.] 2)244]c). Under the same train- ing environment, we train autoregres- sive language models with 1.3B and 340M parameters on the FineWeb- Edu dataset (Penedo et al.) 2024), using 100B and 15B sampled tokens, respectively. The models are optimized with Adam W (learning rate 4 x 10-4, cosine annealing, 1B-token warmup), weight decay 0. gradient clipping of 1.0, and batch size of 0.5M tokens. We use the LLaMA-2 tokenizer with a 32K vocabulary, and set the training context length to 4096. We adopt the Open LLM Leaderboard protocol and a suite of general-ability tasks, as shown in Tab. 2 See $ for more evaluation details.
(MusbaCLL.G (T)
5L.1
16.8
90.7
89.7
73.0
FEM-AN AFT--GTLO)
50.5
9.15
63
31.1
69.2
90.1
52.2
FEM-AFT (APHELL..
55.5
9.78
90.3
80.1
90.2
93.4
69.9
8

---

Published as a conference paper at ICLR 2026
Table 2: Unified language modeling evaluation results across model families and scales. Abbr: Acc accuracy; EM=exact match; IFE-I/P = IFEval (Inst/Prompt, strict only). Shots: MMLU-P-5, GPQA=0, BBH-3, MATH-4, MuSR=0; others 0-shot. Bold and underline indicate group- wise best and second-bes, results, respectively.
Model suriact
Open LLM Leaderbourd
General
Ranking
MMIIIP
RRH MATH MuSR
IPE-T
IFF.P ARC-C ARC-E
HS
PIQA
RualQ Winoii COPA ORQA
SeiQ
#Topl
(Aee1)
(ARc,1)(AcC, ") (EMM)) (7 (sztict)
(Acce" (Accf) (Acc (hec. ") (Avo.. Rank]
1AB Params- 1000
DeliNet
0.308
0.011
0.417
0.288
0.266
0.502
0.704
0.61
0.541
0.240
0.318
0.761
GSA
0.110
0.270
0.294
0.013
0.438
0.300
0.179
0.287
0.529
0.510
0.712
0.541
0.535
0.760
0.330
0.773
3.38
RatNet
0.110
0.252
0.295
0.001
0.334
0.006
0.024
0.271
0.439
0.480
0.301
0.58
0.533
0.710
0.324
0.756
7,63
HGRN
0.114
0.260
0.207
0.008
0.400
0.253
0.122
0.221
0.518
0.481
0.207
0.584
0.515
0.700
0.326
0.695
5.75
HGRN2
0.115
0.254
0.205
0.002
0.350
0.223
0.129
0.282
0.504
0.317
0.671
0.416
0.522
0.770
0.328
0.378
6.63
Transforeer
(SMA
0.114
0.259
0.296
0.011
0.365
0.270
0.141
0.280
0.492
0.492
0.705
0.621
0.582
0.760
0.318
0.769
4,56
FENESM
0.113
0.263
0.326
0.364
0.519
0.713
0.624
0.534
0.240
0.807
2,06
GLA
0.255
0.272
0.157
0.277
0.485
0.702
0.576
0.721
3.63
(GI.A-CLTT))
0.112
0.207
0.277
0.310
0.482
0.708
0.602
0.35g
0.782
3.85
34M Params-
15B
DiffTrans
0.299
0.266
0.133
0.289
0.408
0.668
0.603
0.734
4.38
GosdDeehueet
0.113
0.200
0.2%
0.019
0.25%
01135
0.226
0.527
0.396
0.662
0.58%
0.521
0,710
0.398
0:735
4.25
DeliaNae
0.112
0.260
0.300
0.009
0.492
0.277
0.150
0.289
0.5 0.502
0.405
0.653
0.519
0.504
0.690
0.316
0.717
5.44
HEM SM--..L.
(SMABO)
0.106
0.367
0.292
0.010
0.586
0.260
0.126
0.273
0.506
0.396
0,650
0.569
0.499
0.720
0.334
0.727
6.50
(SMATIn-c)
0.113
0.254
0.2%6
0.009
0.553
0.246
0.122
0.777
0.507
0.403
0.664
0.583
0.515
0.670
0.330
0.728
6.63
HMM-SM
(SMABInCE)
0.112
0.258
0.208
0
000
0.401
0:254
0.120
0.700
0.518
0.407
0.657
0.595
0.511
0.690
0.342
0.731
4.81
FEM SM-o)
(SMATInCLL1)
0.112
0.297
0.010
0.421
0.266
0.203
0.412
01669
0.593
0,519
0,710
0.338
0.716
3.31
FEMESM
(SM-AII-CLAT))
0.203
0.284
0.409
0.609
0.735
GLA
0.289
0:228
0.247
0.366
0.637
0.547
0.649
9.38
FRM-GLA
(GLA-CLTC)
0,115
0,258
0.297
0.009
0,473
0.241
0.123
0.271
0.493
0.397
0.644
0.592
0,510
0.680
0.331
0.683
6.56
Compared with models of the same scale, using FEM im- proves the overall performance of prior methods such as soft- max and gated linear attention. These gains are most evi- dent in handling longer contextual instructions, tackling more complex reasoning tasks (c.g.. IFEval and ARC), and boost- ing accuracy across multiple QA benchmarks. This reflects FEM's ability to enhance general retrieval and context pro- cessing by extending the originally synchronized head-level prior distribution into richer channel-wise and token-wise interactions. The ablation results further confirm that introducing components like +L and +T leads to substantial performance improvements.
Table 3: Comparative analysis of image classification on ImageNet.
DefT-Tiny
DelT-Small
Model
Top-1 Acc Params
Top-I Acc Params
DeiT TNN
72.20
57M
79.90
0M
23
74.40
6.1M
23.7M
HGRN2
75.39
6.1M
23 8M
FEM-SM
76.70
80.45
22 3M
FEM-GLA
75,80
80.20
22.3M
Image Modeling. We evaluate FEM on the ImageNet- 1K image classification task, following Qin eet al. (2024) by replacing the DeiT architecture's softmax attention with our encoder-only FEM implementation. As presented in Table 3 both FEM-SM and FEM-GLA surpass previous methods QQin et al. 2223a b] 2024) while maintaining parameter budgets.
Time Series Forecasting (TSF). Following Lu & YYang (2025). we evaluate FEM variants on TSF. as shown in Table 4 Across datasets, FEM surpasses both its priors and domain- specific baselines such as iTransformer (Liu et al. 2024) and PatchTST (Ne et al. 2022).
Table Benchmark evaluation of TSF tasks.
Datasct
FEM
SM
GLA
FEM
Mamba
FEM
FEM
AFT
GLA
AFT
iTrans-
former TST
Patch-
DMLinear
Weather
0.223
0.223
0.221
0.233
0.188
0.201
0.198
0.202
0.266
ETThi
0.488
0.421
04188
0.454
0.413
0.422
ETTR2
0.344
0.342
0.342
0.330
0.426
Computational Cost. We evaluate the training and inference speed of FEM on a Nvidia L40S
RTTm2
ETTml
0.343 0.247
0.345
0.346
0.246
0.341
0.357
0.250
0.245
0.373
0.365
0.346
0.247
0.252
0.347
GPU. To avoid confounding factors, we use an 8-layer model with 4 heads and a hidden dimension of 512, tested on randomly generated data with context length of 2K and a batch size of 4. As shown in Table the full FEM-SM achieves comparable efficiency to recent model structures, even without additional engineering designs. See additional efficiency analysis in and Table 7
4 CONCLUSION AND LIMITATION
We proposed the Free Energy Mixer (FEM), which reframes sequence modeling as a context- interactive selection problem overcome the "lossless storage but lossy readout" limitation of clas-
9

---

Published as a conference paper at ICLR 2026
sic attention. FEM enables value-aware, per-channel posterior selection top of any prior (soft- max/linear attention, RNNs, SSMs) and, with log-sum-exp, linearized temperature learning, and two-level gating, interpolates smoothly from averaging to near hard indexing without extra com- plexity. It enhances contextual fast-weight programming in theory and achieves consistent gains across NLP, vision, and time-series tasks equal parameter budgets, with ablations highlighting LSE and temperature control as key. Overall, FEM is a plug-and-play mechanism for fine-grained context processing.
Limitation. Our work focuses on advancing the algorithmic ex- pressivity ($7.7) rather than pursuing engineering optimizations such as custom GPU kernels or acceleration strategies. Due to limited computational resources, we were unable to scale FEM to very large models conduct very long-context evaluations. This constrained but focused scope allowed us highlight FEM's al- gorithmic contributions without heavy reliance on engineering or large-scale compute.
Table 5: Latency & through- put comparison (TPS in K to- kens/s). Lower better for la- tency; higher is better for TPS.
Fed
Train
Fwd
Train
Model
Lat.
(K)
IPS
GatedDelnaNe
0.016
0.042
250.4
97.8
DeliaNet
0.014
0.036
292.5
113.9
HGRNZ
0.009
0.024
420.0
170.7
RWKV6
0.014
0,037
293.9
109.4
RWKVT
0.017
0.000
245.1
82.2
DITTaam
0.018
0.041
233.3
100.6
ACKNOWLEDGMENTS
(G.TLL))
FEM-SM
0,012
0.027
333.1
153.7
This research was supported in part through research cyberinfras- tructure resources and services provided by the Partnership for an Advanced Computing Environment (PACE 2017) at the Georgia Institute of Technology, Atlanta, Georgia, USA.
(4G.LL
0.015
0.033
291.5
124.6
FEM-SM
0,016
0.035
283.7
121.2
(-G)
0.017
0.040
249.7
114.6
FEM
0.017
0.041
246.0
104.1
This work used DeltaAl the National Center for Supercomput- ing Applications through allocation MTH250051 from the Ad- vanced Cyberinfrrattructure Coordination Ecosystem: Services & Support (ACCESS) program [Boerner al. 2023). which is supported by National Science Foundation grants #2138259,
#2138286, #2138307, #2137603, and #2138296.
ETHICS STATEMENT
We evaluate FEM only on publicly available benchmarks under their licenses, without collecting personal or sensitive data. FEM's enhanced retrieval ability could be misused (c.g., surveillance or deceptive content), so responsible deployment requires privacy safeguards, bias checks, and legal compliance. We also report model sizes and training tokens, and encourage energy- experi- mentation.
REPRODUCIBILITY STATEMENT
All experiments were under consistent setup, with FEM modules directly replacing standard attention while keeping other components unchanged. Code, configurations, and instructions are provided in the linked repository enable replication of results. See the code base and $ K $D for more details.
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. TRLlhttps://openrevve net forum?
1dd=0g0X4HeyN4I
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL httptppppss//peeeeeeeeeeeee forrm?
1d=1iMSqUuVq9
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long dong-document transformer. arXiv preprint 1500 2020.
10

---

Published as a conference paper at ICLR 2026
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piga: Reasoning about physical commonsense in natural language. arXiv preprins arXiv:1911.1144 2019. doi: AAA12022
10.48550/nXivv.11111114 URLhttps::/ http:/aaivvv...e..v.eeeeee.ee.ee org/ abs 1911 11641
Timothy J. Boerner, Stephen Deems, Thomas R. Furlani, Shelley I.. Knuth, and John Towns. Ac- cess: Advancing innovation: Nsf's advanced cyberinfrastructure coordination ecosystem: Ser- vices & support. In Practice and Experience in Advanced Research Computing 2023: Com- puting for the Common Good, PEARC '23, pp. 173-176, New York, NY, USA, 2023. Associa- tion for Computing Machinery. ISBN 9781450399852. doi: 10.1145/3569951 3597559.. URL https:://do..rgg// http://.....1111511..133333
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretehen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., (oo__iiii files/paper/2020.
2020. URL http:://pproceeccs....eeeiiiaeia Ftte://57ccorrrr0011105522222222
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In International Conference Learning Representattins, 2021. URL https: httpi/ewwwe...orrrieeeeeee /ttp://ommmmdd-aa2uuu0000
7 //openreview..o net
Christopher Clark, Kenton Lee, Ming-Wor Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ:
Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 2924-2936, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/ v1/N19- URL http://cccaattccoo.....oooo org 100000//
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1883.:1803.0545. 2018. htthhttp::/aat
org abs / 1803. 05457
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models implicitly perform gradient descent as meta-optimizers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023. httttttt:///:eeeeee
ftttt/foummiii fzbHR jAdgU
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXivv2307.08666 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arX1v22405.21060, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35: 35:16344-16359 2022.
11

---

Published as a conference paper at ICLR 2026
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Mingyu Ding, Bin Xiao, Noel Codella, Ping Luo, Jingdong Wang, and Lu Yuan. Davit: Dual attention vision transformers. In European conference on computer vision, pp. 74-92. Springer, 2022.
Linhao Dong, Shuang Xu, and Bo Xu. Specch-transfommer: a no-recurrence sequence-to-sequenec model for speech recognition. 2018 IEEE international conference on acoustics, speech and IEEE, 2018.
signal processing (ICASSP), pp.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:201O.I 19999, 2020.
Sai Surya Duvvuri and Inderjit S Dhillon. Laser: Attention with exponential transformation. In Forty-second International Conference Machine Learning. 2025.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021.
Guoxin Feng. Element-wise attention is all you need. arXiv preprint arXiv:2501.057300 2025.
Leo Gao, Jonathan Tow, Stella Biderman, Shawn Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jasmine Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Version vO. 0.1. Sept, 10:8-9, 2021.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 55:30583-30595 2022.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Shiwei Guo, Ziang Chen, Yupeng Ma. Yunfei Han, and Yi Wang. Scformer: Structured channel- wise transformer with cumulative historical state for multivariate time series forecasting. arXiv preprint arXiv:2505.0263: 2025.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory--fficient transformers via top-k attention. arXiv preprint arXiv:2106. 106.09999, 2021. doi: 10.48550/arXiv. 21 2106.06899. htt.
abs /1... 06899
Baran Hashemi, Kurt Pasque, Chris Teska, and Ruriko Yoshida. Tropical attention: Neural algo- rithmic reasoning for combinatorial algorithms. arXiv preprint arXiv:2505.17199, 2025. URL
https::/arxiis org/ abg/abs///aaooo. 2505. 17190
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint XXXv:21033.08844 2021. htthhtts::/ //tp:vv.aavvvvv.. wtt.oorg/aabbaooooe 03874
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Inter- national conference machine learning, pp. 9099-9117. PMLR, 2022.
Hugging Face. Open Ilm leaderboard: Scores normalization. httpps:/huuging co/
ddocs/ httpe//eeedbbbrrrseeeeeee wwwddd//en//peeee 11m leaderbbaarc//oonallaaaaaaaa 2025. Doc. accessed 2025-09-25.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2026
Guokun Lai, Wei-Cheng Chang, Yiming Yang, and Hanxiao Liu. Modeling long- and short- term temporal patterns with deep neural networks. In Kevyn Collins- Thompson, Qiaozhu Mei, Brian D. Davison, Yiqun Liu, and Emine Yilmaz (eds.). The 41st International ACM SIGIR Conference on Research & Development
Information Retrieval, SIGIR 2018, Ann Arbor, MI, USA, July 08-12, 2018, pp. 90-104. ACM, 2018. doi: N0.1145/2209978.220000 URL 020144///22978... 3210006
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565-19594. PMLR, 2023.
Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. itransformer: Inverted transformers are effective for time series forecasting. In The Twelfth In- ternational Conference Learning Representations, 2024. URL https: cetttp::/openreviiie
Jiecheng Lu and Shihao Yang. Linear transformers as var models: Aligning autoregressive atten- tion mechanisms with autoregressive forecasting, 2025. URL flttps://arkor org/ abs/
2502. 2502.0722 07244]
Jiecheng Lu, Yan Sun, and Shihao Yang. In-context time series predictor. In The Thirteenth Inter- national Conference on Learning Representations, 2025.
Ang Lv, Ruobing Xie, Shuaipeng Li, Jiayi Liao, Xingwu Sun, Zhanhui Kang, Di Wang, and Rui Yan. More expressive attention with negative weights, 2025. URL
3bs/2411.07176
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. 2022.
Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. time series is worth 64 words: Long-term forecasting with transformers. In The Eleventh International Conference on Learning Representations, 2022.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.1 11895, 2022.
PACE. Partnership for Advanced Computing Environment (PACE), 2017. 2017. URL http:///www
pace www.ga.aahhcc........ edu
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances Neural Information Processing Systems, 37:308 11-3081,, 2024.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023. URL BtRLhttps:// pettts://aaxiv /atp/vaavvi//.i////i//eee/aa.ava www//bbg.......ooooo... 2302 10866
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kris- tian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, Ce Zhang, and Ste- fano Massaroli. Mechanistic design and scaling of hybrid architectures, 2024. URL https:
//axxxvv.rrg/ab//2400..3
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong. Nick Barnes, and Yiran Zhong. The devil in linear transformer. arXiv preprint arXiv:2210. 10340, 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Ling- peng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint 08877 2022b.
13

---

Published as a conference paper at ICLR 2026
Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh In- ternational Conference Learning Representations, 2023a. httts::///ppneeeeee
id= w
Zhen Qin, Songlin Yang. and Yiran Zhong. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Process- ing Systems, volume 36, pp- 33202-33221. Curran Associates, Inc., 2023b. URL https: https: / pttpps:/pproceeeinss.eeeee neurips cctes.ce//e httpppppper_..lesspppp///222... 69 3
Hierarchically gated recurrent neu-
ral network for sequence modeling.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li. Weigao Sun, and Yiran Zhong. Hgm2: Gated linear mns with state expansion, 2024. URL https::/arxiv. /bs/
7404.07908
Alec Radford. Improving language understanding by generative pre-training.
Ope-
nAI technical report, 2018. URL http:::://cc...peeeee htt.://////eeeeeeeeeeeeeeee. covers/
language Language-unsuppe www.peeeeiieeeeeaaaaaaaaaaaeeaeu www.:uuttddddnnnnnnnnnm.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Di- rani, Julian Michael, and Samuel Bowman. Gpqa: graduate-level google q&a bench- mark. arXiv preprint arXiv22311.12022 2023. URL https::/arxiv.oom http:///.rrg/////////////
12022
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S. Gordon. Choice of plausible alterna- tives: An evaluation of commonsense causal reasoning. In AAAI Spring Symposium on Logical Formalizations of Commonsense Reasoning, Stanford, CA, 2011.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An ad- versarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Ar- tificial Intelligence, volume 34, pp. 8732-8740, 2020. doi: 10.1669/aaii.v3/el 66999, URL
https:://jss.o. aaai. org/ index php http://///aaaatteteeeeeeeeeeeeee
/ttpa/aaaaaaaaaaaaaaaaaaaaaaaaa/ 6399
Jürgen Schmidhuber. Learning to control west- memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
Tao Shen, Tianyi Zhou, Guodong Long, Jing Jiang, and Chengqi Zhang, Tensorized self- attention: Efficiently modeling pairwise and global dependencies together. arXiv preprint NoXYv:1805000111 2018.
Zayne Sprague, Xi Ye, Kaj Bostrom, Swarat Chaudhuri, and Greg Durrett. Musr: Testing the limits of chain-of-thought with multistep soft reasoning, arXiv preprint arXiv:2310.16049, 2023. URL
https:://rxxisi wtt. abs 2310.16049
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le, Ed H. Chi, Denny Zhou, and Jason Wei. Challenging big- bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261 2022. doi: 104485500axxi.221000266.. URL httts:: wttpx//wo///b///22..0000
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:201 10220001.40066 2020.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL petttps://ariiv
httppsss//arrvv..rrraa////////// wttpv/ooggg/b//////e 13971
14

---

Published as a conference paper at ICLR 2026
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
Yubo Wang, Xueguang Ma, Ge Zhang. Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang. Tianle Li, Max Ku, Kai Wang. Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. arXiv preprint arXiv:2406.01577 2024. URL org/aba 2406. 01574
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303. S0t.::3333003866,, 2023.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsoureing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy
Text, pp. 94-106, Copenhagen, Den- mark, September 2017. Association for Computational Linguisties. doi: 10. 10.11865//11/11744433
URL httph:tts///// www.elaannhhloogo org/ org/w17 4413/
Haixu Wu, Jichui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition transformers with auto-correlation for long-term series forecasting. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 22419-22430, 2021. URL attps: pttppss neurips. http::/iii....o.roooooo..o/ppp/a
cc/paper/2021/
hash/bcc0dd 0028 hrc.a//8eeeeecccceaaaaaac-o-----
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference Learning Representations, 2022.
JUMI
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströnformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 14138-14148, 2021.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.0m 2.6464, 2024a.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-effficien training. In Forty-first International Conference on Machine Learning, 2024b. URL httptttpp://peeeee.eeeeee
net/ for um? 1=1.:i1xx1ax1u1r.r
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. arXiv preprint arXiv:2406. 06484, 2024c.
Tianzhu Ye. Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. In The Thirleenth International Conference Learning Representations, 2025.
URL petttps://oppereeee
wwtn/eeooruiiit. forum? wttpmmmmm=--e d=OooCo1
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 12833733-7 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can ma- chine really finish your sentence? In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Com- putational Linguistics. doi: 10. 10.11853//1/P19 1472. URL.
oorg
P19-144 14727
Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? In Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 11121-11128, 2023.
15

---

Published as a conference paper at ICLR 2026
Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103 2021.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett, Trained transformers learn linear models in-context. arXiv preprini arXiv:2306.0992 2023.
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on
Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium Edu- cational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pp.
11106-11115 AAAI Press, 2021. URL https: //oss///jjss.o aaai. org/ www./nndxxx./ php/ AAΑI/
axticle/view/17325
Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instruction-following evaluation for large language models. arXiv preprint arXiv:2311.07911, 2023. URL
www. awt.aa//......./..
16

---

Published as a conference paper at ICLR 2026
APPENDIX CONTENTS
Statement of LLM Usage
17
Details and proofs for Section 2.1
17
Details and proofs for Section 2.2
18
D
Additional discussion: per-channel score distributions vs. loken-separable mixers
22
Theoretical properties of FEM
23
F
Details for Linearized Temperature Learning
25
Details for Two-level Gated FEM
26
H
Parameterizations of the Prior Selection in FEM
28
I
Low-rank Convolution: Time-Decay Conditioner (TDC)
J
FEM as Universal Fast-Weight Programmer
32
Additional Implementation Details
L
Additional Dataset Description
33
M
FEM-MLP Interaction: An NTK and Fast Weight Perspective
N
Additional Discussion: Significance and Motivation
36
Additional Runtime Efficiency Analysis
P
Toy problem: single-layer channel-wise argmax
38
A
STATEMENT OF LLM USAGE
In this paper, LLMs were mainly used to assist with writing-related tasks, including grammar check- ing, wording adjustments, length reduction, layout reorganization, text formatting, formula format- ting, theoretical derivation formatting, and table template generation.
We also used L.Mss search for existing methods and references in order avoid duplicating and over-claiming However, we did not use LLMs to conduct literature reviews, nor did LLMs replace the authors in studying the cited works. We confirm that all cited literature was read by the authors, not solely by LLMs.
During experiments, LLMs were used assist with generating or refining experimental code and scripts, especially for bug fixing and efficiency optimization.
LLMs were not used for defining research problems, proposing ideas, designing methodologies, providing theoretical insights, or creating algorithms and model architectures.
B DETAILS AND PROOFS FOR SECTION 2.1
B. SELECTION DISTRIBUTIONS, SUPPORT, AND NORMALIZATION
We encode causality by restricting the feasible support to M2 (1. softmax attention,
(i( = explg. k. Vd) 1 S
Ejste exp(y. q: k//vd)
In linear attention nonnegative feature map Φ: _d R and set
(i) (d(g)), (k;)) 1ii <1}
Ejsr (
Nonnegativity guarantees Pt !--1 Row-masking is absorbed by Mc.
17

---

Published as a conference paper at ICLR 2026
B.2 PROOF OF LEMMA 2.2
Let m. = E, any with A 2 0 and Σ A = 1. For any coordinate j. 04, S (mz)y implies
(mi)) Σ Auayy s
Em((m)) (me)).
Thus equality holds termwise: for all with X; 0, vi = (mi)j. Hence every such simultane- ously attains all coordinate maxima, proving the claim.
B.3 PROOF OF COROLLARY
Let ii 1.......... :.....-UUppll..... Vip. with least distinct indices among {iz)- Unless the cho-
sen Vi, coincide on all selected coordinates (a measure-zero degeneracy), Lemma 2.2 implies
of conv(vi..... UrJ. so no pt satisfies E, P; s7.
B.4 HEAD-SYNCHRONOUS ASSIGNMENT CAPACITY
Consider H heads at step &. Let oi, A-11 be head h's selection distribution and 4 arg max;5: (h) Channels routed through head h share the same a their first mixing, so the pattern is determined by (4.....,(H) and fixed partition of channels into heads. The number of realizable patterns is at most (H, versus 12 for fully independent per-channel selection.
B.5 REMARKS STORAGE VERSUS PROCESSING
Softmax attention stores the entire set {(k, bi)lise without compression, but the per-head read equation enforces one weight vector across all coordinates, which is the bottleneck for tasks re- quiring different indices per channel. Pointwise nonlinearities additional depth cannot recover per-channel index identity at the same step unless a new, independent selection distribution acts before the first mixing on those channels.
C DETAILS AND PROOFS FOR SECTION 2.2
C.I MORE HEADS: CAPACITY, LOW-RANK EFFECTS, AND FNITR-FEATURE LINEARIZATION
Bilinear form and rank.
With H heads and d,, D/H,
(1) oi w (b)
(w())) Σ Mi(1) 1Et
rank Lank(w>! (w (w(h)) d,. (12)
y
h=1
Proof of Lemma 2.4 At step f, head h selects arg max; o, the Cartesian product over H heads
(h) has size at most H. Inside head, all output coordinates are linear images of the same o
Finite-feature approximation (value-path erosion). Assuming clipped logits 1q A| < R, a sin- gle softmax head of width d,, admits an E-accurate finite monomial feature approximation with
N + dh
M
N R log(1/s)),
dp
so its read uniformly approximable by a linear streaming state of size M x d,., The full result is below.
Proposition C.1 DDimension-dependee linearization and memory collapse for softmax head). Consider one softmax attention head with query/key width d,, and value width d,, Assume bounded scores and values:
1/b ks < R (i ().
||ai|12 S V.
18

---

Published as a conference paper at ICLR 2026
Fix E (0, 1) and choose N such that
R
S
n!
n=
Define the feature map that collects all monomials up to total degree N,
&N.dd (x) -33a)
2
E ERA,
M
Σ n
n2 dh dh -1) =
(N +h)
JolSN
dp -1
dp
Then, uniformly {1g kl kl S R).
c k dN.dn (1) TONT,,d (k)) S E.
(13)
Define the streaming sufficient statistics
S. Σ &N (ki) 2s Bx.,
Ze =
DoN (ki) E ERM,
i5t
and the linearized readout
6z = ΦN.d ( St
E Rdd
NN (4t)) Z5
If. in addition, e S then the exact softmax output O2 Eige ae with dns x el hs satisfies the uniform (in t) error bound
sup
|| o. ō. 2 S 4V &
(14)
Consequently, a single softmax head is O(e)aapproximable by a linear, streaming state of size M x d, plus one -vvector, where
M
N + +do (N:
0
No d.
N 6 R log log !
di
In particular, when dp we have M = N + dimensional kernel-RNN-like compressed memory with arbitrarily small uniform error as N + cc.
O + log )) the head collapses to a one-
Proof. Multivariate Taylor expansion of e of @N,d, ΦN.d, TN.,,, d, (ls) Σ Eja
gives el D Ejaion
By construction
R" of so the truncation error is the scalar exponen-
tial tail evaluated at 14 k| < R. yielding equation
13 by the choice of N.
Let K2(i) = e k. Ri(i) <(q:) ((k.). Write N, Σ, K,(i)v,, D. E, Ke(i) and N,
>, (1)v,, D. = E; (1). From equation and ||u/|l2 V
IIN, - Nella 5 Σ ||oil12 S EVt,
|D, - D.l S st.
ist
Since |g| < R, we have te R < Dr S R If < 2, then D: - DD: Dill 2 te R, Usinng the standard ratio perturbation bound,
N₂ De
N
| Nellz
+ |Neel22 |D, D,
<
D.
D. Dr
D.1
De Dr ID= Del
Because |N S VD, the RHS at most
+ Vitt RVe which proves equation 14
= - The stated complexity follows from M ( (d) and Stirling's approximation; for da = 1, M
N ++ 1.
L
Any common scaling (c.g.. 1/vd,) in dot-product attention can be absorbed into R. Position biases likewise be included provided the total score remains bounded by R.
Remark.
19

---

Published as a conference paper at ICLR 2026
Numerical illustration (state size under bounded scores). We instantiate Proposition C.1 with two practically relevant score radii: a high quantile R = and an extreme upper bound = 10. For target uniform kernel error E, choose the smallest degree N with En R" /n! S s. The resulting
hidden-state size per head (in the d, d,=1 collapse) is (N++) de ( d across all heads with total value width D Hd it ( D).
Minimal degrees N (exact tail test).
8 10-4
8= 10 6
8 = 10
R = 5
N 19
N 22
N 25
== 10
= 33
= 36
N = =40
These values satisfy the safety condition of Proposition
C.1
(EeR S 111; e.g., ce20 2.2x10-2
8 = 10
Concrete state sizes (per head, d, 11). For ε 10-6,
R 5= = N+1 23 > state (N+1) d, 23 and O(N D) ((2 overall,
R = 10: M=N+1 = + state = +1)+1)ee d, = 37 d. and O(N D) = (36 D) overall.
Thus, under realistic bounded scores, a single softmax head with d, is equivalent (up to uniform error E) to linear streaming memory whose per-head size grows essentially linearly with R and only mildly with E. For d, > 1, the finite-feature dimension becomes M (N (N++di) O(N4, /d,,), explaining the strong dependence on per-head width.
C.2 DEPTH: NO SAME-STEP UNMIXING AND SELECTION BUDGET
E, or, linear, nonnegative, and weight-summing to 1, hence images lie in conv(e;; Composing coordinate-win wise maps keeps outputs convex hull of transformed points and does not reveal per-channel indices used before mixing. Later attentions at step operate finite set of already mixed tokens; selector outside conv(u; is unreachable without a fresh independent selection before the first mixing touching those coordinates.
Proof Lemma 2.5
Ta (w,) + E,
Define a channel group as coordinates whose first attention mix- ing shares the same head at some layer. Across L layers there are most HL groups. Each group gets one independent selection distribution for its first mixing, hence at most H L indepen- dent per-channel selections by step Necessity of HL > D follows; achieving the bound requires avoiding re-synchronization before first attention.
Proof of Proposition 2.6
Accumulation. Layer l writes ( R D to KV. Stored channels scale LD, independently selectable groups as LH; the fraction non-indeppendeetly-seeeettt channels does not vanish unless H scales with D.
C.3 DIMEDIMENSION QUERIES/KEYS: CAPACITY- TRADEOFF
Giving each coordinate j its own scoring subspace increases assignment capacity toward tP, but increases parameters and compute from (d?2) to ((D)) per layer. Under a fixed budget this forces shrinking D (hurting value bandwidth) or the MLP width (hurting global capacity), both detrimental in long-conter regimes.
C.4 TOKEN-SEPARABLE MIXERS REMAIN CONVEXLY CONSTRAINED
We analyze
Σ 0r +
(ar + Ea,,
ai aa Cv) + Σ ((a,,))))
with coordinate-wise
20

---

Published as a conference paper at ICLR 2026
Proposition C.2 (Full statement of Proposition 2.7 (i) The first
linear; images lie in conv(u, and conv (3, O ) up to coordinate-wise scaling. (ü) For Σ, 01,00 a v,), outputs lie in conv (onv(a(B, @ v:);: recovering channel-wise selector the original coordinates is impossible in general unless special degeneracies (e.g. identical selected coordinates across candidates) hold. (iii) For general token-separable J, per-channel argmax over original coordinates is impossible in general.
Proof sketch. (i) Direct. (ii) m (max, U(,,,...) is outside conv (Lemma 2.2 any
convex combination of transformed values cannot map back to m unless σ is globally invertible and aligned simultaneously across all candidates, which fails generically. (ii) Duplication argument in D I: take two identical tokens u at indices * j but target max to prefer one index; any token-sepane -eparabbe E; f (a1,5 U4) invariant under swapping the two, contradicting index-sensitive selection.
Complexity remark. Per-channel cross-token operations (e.g., top-k, per-channel log-sum-exp) introduce non-separable normalizations over and typically break fused ((T)O((T) implementa- tions.
C.5 LINEAR RNNS AND SSMS LACK LOSSLESS STORAGE
Let h, SR.S be fixed-size state updated by (possibly input-dependent) contractive linear op- erator. Classical lower bounds for linear time-invariant systems imply existence of sequences and horizons / where single-token recovery error from h2 is bounded away from 0 for any fixed S. Hence fixed-state models cannot provide lossless storage of {w; Jise for arbitrary index retrieval at step t, in contrast to a KV cache, and thus cannot realize channel-wise selection over all past values.
C.6 CONNECTIONS TO RECENT PER- CHANNEL VARIANTS
The families in Section 2.2 subsume many contemporary designs:
(i) Score-space inflation per feature. Tensorized/mulli iinzorizedmmuttt-ddmenninnal attention and element- wise attention allocate scoring subspace per coordinate to produce Or,i. (Shen et al. 2018 Feng 2025) This moves assignment capacity from ! toward ID,. but the read stays token-separable, hence sub- ject to the convex-hull constraint (Proposition C.2p Moreover, the per-feature distributions are typically prior-only (value-agnostic) at the same step, value-aware cross-token competition is introduced before first mixing (cf. Lemma 2.5) The parameter/compute cost also seales from (cf))
to O Dd) per layer; see Appendix C.3
(ii) More heads/depth. Increasing H adds only H independent selection groups, bounding head- level assignments by 1.H (Lemma 2.4 depth increases storage but not the number of independent first-mixing distributions per step beyond H (Proposition 2.6). Hence the channel-wise lossless- selection gap remains unless Hl scales with D.
(iii) In-head pointwise gating. Adding coordinate-wise gates inside the per-head mixer keeps token separability (the form Σ; fa so outputs remain convex hull of transformed val- ues and cannot realize per-channel argmax of the original coordinates in general (Proposition C.2) g (TT)O((T)) implementations; see Appendix C.4
Making the gates sedex-sensitive requires cross-token competition per channel, which naively breaks
Summary. Across (i)-(iii, either capacity increases at significant parameter/compute cost while the read remains token-separable, or the same convex bottleneck persists, or fixed-state storage limits retrieval. None provides per-channel, value-aware cross-token competition before the first mixing under the prior's asymptotic complexity.
C.7 WHY A STRONGER ALGORITTMIC MIXING STRUCTURE MATTERS
value-aware, A mixer that natively performs value- aware, per-channel cross-token competition at the first mixing step has two practical advantages under fixed budgets:
21

---

Published as a conference paper at ICLR 2026
Separation of roles. The mixer shoulders dynamic fast- fast-weight programming (context-dependent routing/sclection), while MLPs focus on feature synthesis and knowledge consolidation. ker- nel/NTK view, this conessonds to adapting the effective kernel online at the mixing site, reducing the burden on downstream static nonlinearities. More discussion can be found in $M
Parallelism and efficiency. such competition is realized without changing the asymptotic com- plexity of the selection prior (e.g.. by computing a per-channel log-partition over the same masked support), we preserve the O(T)) softmax or O(T) streaming behavior and fused-kernel practical- ity. This is the design objective satisfied by FEM in the next subsection: it introduces awlue- per-channel posterior selection via a variational free-energy read while retaining the prior's time complexity.
D ADDITIONAL DISCUSSION; PER-CHANNEL SCORE DISTRIBUTIONS VS. TOKEN-SEPARATE MIXERS
Many recent variants extend single per-head distribution = at,- to per-channel distributions Qe(c, at,,, E A' yielding
Ot,c Eige Otiie Us,er 0 Eist Diag oi.i,,.... OLDDD) (vv:) = Eigr Wi, ① ((v)).
D:D
(15)
where acts coordinate-wise and Wi,i (o (t,.. Expression equation 15 is token-
separable: the outer sum is over tokens and introduces no cross-token interaction inside the mixer. Consequently, for each channel c, ot. E conv (el.e..... vt,c). and exact coordinate- wise selection at the same step is unattainable unless ar,,,e degenerates point mass (cf. Lemma 2.2 Lemma 2.5 Proposition
Assignment capacity vs. convexity. Per-channel scoring lifts -ynchroncss capacity from ! to the natural upper bound 1D: across contexts, independent argmax patterns (E))))(D) can in principle be realized by {ar,,e} (Shen et al. 2018 Feng 2025). However, the mixer remains convex expectation per channel; without value-aware cross-token competition, the distributions need not concentrate on the value argmax, and the dossless-selection gap remains.
Mapping of representative designs.
Per-dimension score inflation. Tensorizac//malti-diinnnnss and element-wise attentions in- stantiate at,i by combining shared token-to-token term with per-channel terms or by per- channel distances (Shen et al.) [0188 Feng (225).. These methods increase assignment capacity (toward {D) but keep the sekeasseprrbble convex read in equation 15 and are typically prior-only (depending on (q, k) but not r:).
In-head mixer enrichments. Per-channel rescaling, pointwise nonlincarities, or FiLM-style gates fit E, (B, O vz) and remain within Proposition C.2 the image convex hull of transformed values, and no same-step unmixing arises without an additional independent selec- tion before first mixing (cf. Lemma 2.5).
Axis/channel attention and structural re-partition. Methods that attend over channels (or axes) rather than over past indices change the domain of selection but do not produce per-channel distributions across time; thus they do not affect channel-wist wise index capacity over L; see, e.g., channel-token attention in vision and time-channel layouts in forecasting (Ding et al., 2022 Liu
et 2024 Guo al. 2025).
Linear RNNs/SSMs and kernel priors. Streaming fast-weight priors with fixed-size state offer cross-dimension couplings yet lack lossless storage over all past indices; kernelized/linearized priors preserve streaming complexity but still yield expectation reads (Katharopoulos al. 2020)
Chorooanski et al.) 2021 Gu & Dao 2023).
Where FEM differs. FEM preserves the chosen prior Pt (softmax, kernel, RNN/SSM) but up-
grades the read from expectation to the free energy 3- log E, (e) exp(Se,,)) yielding per-
22

---

Published as a conference paper at ICLR 2026
channel, value-aware posteriors Ht.8 (i) e This introduces cross-token competition
-(0)
per channel before first mixing, achieves the M, assignment capacity and admits exponential posterior concentration while retaining the prior's asymptotic time complexity (see $2.3.1
E THEORETICAL PROPERTIES OF FEM
We fix timestep f, a channel j E [D], the prior selection distribution pr with support M, := {i- Pe(t) > 0), and the values {ur,
Notation.
For 3 > 0, define the per-channel free energy and posterior selector
Pe(i) e Erems (r( eBors i E M2, (16)
(i) = log
Σ
(r()) e
(9) 9r,5 (i)
AE M.
and let v. E R Mel collect feijJee..
Standing assumptions. All statements are over the support Mr and assume pi(i) > 0 for i Mr. For 8 < oc, the posterior (9) is unique; in the limit B + co, ties may persist if margins vanish, which does not affect finite-ß claims.
4t
Lemma E.1 (Equivalence of budgeted and penalized forms). Fix t.j and budget B 2 0. The constrained problem equation 3 Ghas a unique maximizer q" E (MM). There exists a unique 3 20 (Tll))p))) conversely, for every B 2 0, the maximizer of the penalized objective solves equation for the budget B = (Iq")))))) The map B + (*)
such that 9" arg max, g (E; (i)mz
continuous and strictly increasing whenever v. is not pr-a.s. constant.
Lemma E.2 (Donsker-Varadhan variational principle and mirror ascent). For every 3 > 0,
Σ q(i) Va KL(llp.)) }
(1)
max (EA(M.)
(17)
with the unique maximizer
qt in equation 16 Equivalently,
(9)
arg min
B (L((l)))) (a. v.g).
(18)
(MM(
i.e. an exponentiated-grann gradient (mirror ascent) step from pt with step 3 along v.g.
Proof. Standard DV identity: log P,e = max, v) - KL(glp)) Divide by B and apply KKT: uniqueness holds on (MM) since the objective is strictly concave in g.
Proposition E.3 (Expectation baseline and monotonicity). Let Pt.j = Ep. [vi,l. Then
Fi) Pt.j + KLpp. || di.3) Z Ms.j
(19)
Moreover, B (B) is continuous and strictly increasing unless v.. is pr-ass constant, with
d
1 ku p+) 2 0, F2)( jt.j + 2
B Varp.(ui.)) 0(33) (B 0). (20)
(i) d3
B2
Proof. equation 19 follows by direct algebra using (9) Be Differentiate B Boglog E, Pve to obtain equation 20 The small-3 expansion is the second cumulant of Ut, under Pt.
Proposition E.4 (Local geometry: gradient, curvature, smoothness). Fe,( convex and t in U.g. with
Vo, e F() = 9ig
(8)
V2 F(,(33 = =3( =D(iag((e)))
2t.5 Aiut
(3) (8)
(000 2 0. (21)
Hence
|V 11441) 112 = and F = : 1|2 I Moreover, Fi, is 8/2- 3/2-shooth in
lzi
( 3 Amaz( - T) 3//2,
(22)
and the bound tight when q is supported coordinates equally, e.g. 9 (1/2,1/2 0, ....,,))
23

---

Published as a conference paper at ICLR 2026
with respect v.j to obtain VFF.((3) (0)
Proof. For equation 21 differentiate equation 16
12 and
(B) 8 (Diag(g) - qq gq where Q = Convexity and smoothness (indeed ( follow from the log-sum-exp structure. The 1 and -orm statements follow since g is probability vector: ||/lla = and ||al13 = E; g. - E; =
For equation 22 write J(q) := Diag(g) 9q This is the covariance matrix of a one-hot random
vector with class prosabbibbaiiistees q. hence J(q) = 0. To bound its spectral norm, apply the Gershgorin disc theorem. Row i has diagonal entry gi q3) and the sum of absolute values of the off-diagonal entries is Σjại 4.4 (1 ge), every eigenvalue lies in
Ula(1 1 q:) - q 4),, q. q:) + q:11 - -=)=-Un 2q.(1 -.)
Therefore {max(((4) < max; 2q, 9:) S 1/2, with equality attained when is supported on
two coordinates equally, e.g. q = 1/.. ..)....).......... (then J(q) has eigenvalues {0, 1/2, o.... ..)....)))).)))))
0}),
Multiplying by 3 gives [[ (3)llop = =8 = S B/2, and the bound tight in the stated
case.
Proposition E.5 (Range, concentration, and finite-B guarantees), Let i arg maxieM, Uü,j and Arj Ue maxiyli Ui 2 0. Then for all
0,
Maj S (B) S vj 1 1- (8) (i*) 1 pe(i") Pe(i*)
S e e- Des, ue = S U.
(23)
In particular, if D.. > 0 then Dest = 6x and (B) t Us exponentially as - xc.
Proof. Upper bound: log E, Bie < B max; U, Lower bounds: F= JKLpplg(0))) 2 μ and Eigir p00 I (1 pr) p=))=(0''=0 give equation 23
Proposition E.6 (Mask preservation; shift/scale; prior sensitivity). (i) (Masking) If (. 0 then
(1) (1) O; restricting Mi can only decrease equation 17
(ūi (Shift/scale) For any e. E R and o > 0,
Fis + c) c (B; (iii) (Prior sensitivity: probabilities) Viewing Fi, as a function of A(M,),
(B; FF.((a
1 8v..
3 0, 1 @n @Bu.ITT
VPF
ViFu = 3 Σ E, Epeeen))))
3 E, Pre
so Ft.j is concave in on the simplex.
(iv) (Prior sensitivity: logits)
For unnormalized weights 8; with wi = log s; and F(3 w) = 1 log E, ewitsmss
1
V
q.
V (Dag(()) - -qq t o.
3
hence f is convex in
For normalized logits b with p softmax(b), writing (r) = Diag(r) - Tr
ViJis } (i. p),
(9)
ViFi ( (
which is in general indefinite; thus Fe, is
of
Proof. (i) is immediate from equation 16 For (ii), add inside the exponent or reparameterize + a5 to obtain the stated identities. For (iii). F(p) = 3 log ea"h) is a log of an affine function in p, hence concave; the displayed derivatives follow by direct differentiation. For (iv), both statements follow from standard properties of log-sum-exp: the unnormalized case is convex in we composing with softmax yields DC form with the given gradient and Hessian.
n
24

---

Published as a conference paper at ICLR 2026
Theorem E.7 (Channel-wiss) assignment capacity over the prior support). Let D be the number of channels and a = (d.....,aD) E M.. If each channel has a positive margin Du.j == VayJ
ma May3 Vi.j > O, then there exist finite temperatres {Bri} such that arg max; Qt.j (91.5) (i)
aj for all j. Hence the set of achievable chanmel-index argmax patterns has cardinality [M, (the natural upper bound). single attention head, in contrast, yields at most [M patterns (all channels
synchronized on one distribution).
E.5 per channel and choose St,j to concentrate posterior mass on Proof sketch. Apply Proposition
aj with any desired margin; counting patterns gives |AMp]
Proposition E.8 (Complexity preservation and stable backpropagation). For fixed B, computing (B) requires one masked log-sum-exp over M, and produces q:3 as the gradient equation 21 Therefore FEM preserves the asymptotic time complexity of the underlying prior (e.g., O or O(T)) while enabling numerically stable forward/backward passes using standard LSE/softmax primitives.
Proof sketch. Convexity in q and Slater's 's condition yield strong duality for equation 3 KKT gives the log-linear form q* Pt @ev with multiplier 3*. 'he monotonicity follows from the derivative d F(B) 3 (KL((()) (1p)..
The
Remark. Lemmas E.8 establish that FEM variationally optimal (DV), value-aware with explicit local geometry, monotone in temperature with variance-comtrooled small-B behavior, mask- preserving, concave in the prior on the simplex, convex in unnormalized log-weights, and DC in normalized logits, capacity-optimal for channel-wise assignment the prior support, and complexity-preservingg with stable gradients.
F DETAILS FOR LINEARIZED TEMPERATURE LEARNING
F.1 FIXED TEMPERATURE: DECOMPOSITION AND COST
Lemma F.1. For fixed B > O, the FEM read satisfies FL)(B) = Pt. (L(ella where (8)
Mtj = Ep. (erg) and Z. (i) x Pe(i)e on M.. Evaluating Fi,(B) adds one masked LSE per channel and preserves the prior's asymptotic complexity.
Proof. Algebra from (()) x p @ yields the identity; cost follows since the support is Mc.
F.2 MONOTONICITY AND HIDDEN TEMPERATURE
Proposition F22 Let (B) 3o1 EiEM, Pe and = (t) Then (8) (lpx)
0, with strict positivity unless v.j is pr-a.s. constant. For any X [0, 1], there exists a unique (A) [0, 3max] such that (1- +M (i. (A)).
(3) --22
Moreover À + Pi,5 is continuous and strictly increasing.
Proof. Differentiate F obtain F(( ( Continuity and strict monotonicity on [0, 3max] imply the claim by the intermediate value theorem; strict increase follows from strict positivity of F" in the nondegenerate case.
Corollary F3 (Reparameterization equivalence). For any smooth loss L. optimizing Àt in C(F,, (Ar,)) is a strictly monotone reparameterization of optimizing 3 in L(Fr)(3)): ac/ax (OL/OF) F² (3*) (08* (ax) with F²(3*) > 0.
25

---

Published as a conference paper at ICLR 2026
F.3 KL-CONTROLLED INTERPRETATION OF THE GATE
From Fir(3) - prj = B -L (o0) and equation
1 kL(p 19:. (A))
(L(pellag.
(A)
3max
so À specifies the fraction of the achievable KL improvement realized at step L.
F.4 ONE-PASS IMPLEMENTATION
Form the augmented value stream Vi.j [usj e and compute EiEM, Pe(O)oes = in one pass. Then Fimex Ama log (i) and equa- tion equation follow.
[Pt.]:
F.5 GROMETRY, STABILITY, AND ADDITIONAL. PROPERTIES
For completeness collect properties proved in Appendix E (i) DV variational form F(B) max, [15[0] - 8- (L(((l))) with maximizer g(8) pe (ii gradient V.F(B) g(
Hessian FB3) 8 ((iiag(g()) - g(0) and 3/2- //2-smoothness: (iii) small-3 expansion = _ (vr,(u) 0(B)); (iv) mask preservation; shift/scale laws; concavity in p; convexity in unnormalized logits; ciference-of-convee
normalized logits; (v) complexity preservation and capacity consequences when 33 is large.
F.6 COMPLEXITY SUMMARY
requires one expectation and one masked LSE at 3max per channel, both over M,. thus matching the prior's asymptotic complexity (O(T²) for softmax; O(T) for kernel/SSM priors) while enabling dynamic temperature control in single pass.
G DETAILS FOR TWO-LEVEL GATED FEM
G.1 INNER GATE AS HIDDEN TEMPERATURE; PROPERTIES AND PROOF
Lemma G.1 (Monotonicity and smoothness). Pt.j is continuous on [0, oo), differentiable (0, oo), and
d
3-2 (9) [| 20. 0,
(24)
d3
with equality iff v.g is Pr-as. conslant. Moreover Feg is corvex and 3 3/2-smooth in v.g-
Proof. Standard Donsker-Varadhan calculus yields FL,(3) = max (E, [ugl]
(1/ lp)). Envelope differentiation gives equation 24 Convexitylsmoothnes follow from (.5)
the Fisher covariance of 9:
Proposition G.2 (Inner gate hidden-temperatuue free energy). For each channel and any Xe.j [o, there exists unique Brid.t.j E [0,, j) such that
Pey(Des) = - log Σ peli) exp (apiiii Va.s)-
Moreover, dej + Buid,ey strictly increasing unless v. a.s. constant.
Proof. Let Dr, Fr., (B) (0). By Lemma G.1 A is continuous, strictly increasing on [0, Bmax,j] unless v. is constant. For any X E {00, 1]. define
Phid,ej AT (r. (
[0, Bmux,j];
which is unique by strict monotonicity. Substituting yields Fr.g (Nr.)) = Fi5 (en(Snid,es) as claimed.
26

---

Published as a conference paper at ICLR 2026
Reverse-KL improvement over the mean. For any B > 0,
1 Fi)(3) Pas + kl(pilla.))
(25)
so (A) improves over Mr.j by controlled reverse- term at the hidden temperature. This explains the mean->soft-max interpolation effect of the inner gate.
G.2 COMPLEXITY, SINGLE-PASS COMPUTATION, AND STREAMING
Proposition G.3 (Complexity preservation). Computing equation TO equation 77 requires one ex- pectation and one masked log-sum-exp at Bmax per channel, hence matches the asymptotic time complexity of the prior Pt (e.g., O(T²) for softmax; O(T) for kernel/SSM priors).
Proof. Compute (t) 0e exp(Smax 2 vz)] once, then split obtain μrr and Fiaax, This quires one expectation and one masked log-sum-exp per channel on the prior support M, matching the prior's asymptotic time (softmax O kernel/SSM O(T)). The outer gate 9t pointwise modulation.
Streaming compatibility. For associative priors (kernel/SSM), the normalized read is computed by the same scan used for pa; concatenating a constant "1" channel yiclds the normalizer and nu- merator in one pass. The LSE branch uses the same support M, and thus preserves streaming.
Numerical stability. We use standard LSE stabilization per channel: subtract max; (8 (8 Usgg) inside the exponential and add it back after the logarithm. Gradient clipping for Bmax prevents overflow when tasks push toward hard selection.
G.3 CONTAINMENT OF MIXER FAMILIES
Proposition G.4 (Formal containment). (i) A; = gives O = Ei (ge Ov)) matching per- channel linear reweighting. (ii) < X <1 yields a monotone, convex aggregator in each channel that interpolates between Paj and max, Vú Arj increases. (iil) Allowing Xt, 9: depend (ctx, Pt, Pt, Fquax) realizes token-separable couplings of the form E; fars, v.) and adds cross- token competition through the log- sum-exp term.
Proof. Direct substitution of the choices for A and identification of limits 3 0 and B oo per channel.
G.4 CAPACITY AND HARD-SELEETTI LIMITS
Proposition G.5 (Capacity and limits on the prior support). With X R and sufficiently large Bmax. the per-channel posterior concentrates its own arg-max over the prior support M, {i Pi()) > 0}, so the achievable channel-inde -indxiidxx assignment capacity attains [Ml). In the limit X = 0, FEM reduces the expectation baseline (the original read of the selection prior).
Proof. Fix channel j and let A.,; = mine (u-. > 0 be the margin at the arg-m max index i*.
For any B 2 3o( (),,), the posterior qt. places least cxp( (xxp(-AA,,) -3A,,) mass i*, and Fr, Ve Across channels, with X 1 and sufficiently large Buax. the joint posterior concentrates independently per channel over Mr. achieving [M,' D distinct index assignments. Setting X recovers pt:
G.5 GRADIENTS AND CURVATURE
For channel j.
(B), Oui.j
82 ())
(8) (i. (1),
= 3 (11 (8) (1)((ir) -4
27

---

Published as a conference paper at ICLR 2026
Thus gradients the posterior weights and the Hessian is Fisher covariance scaled by 8, giving stable, value-aware competition. Backprop through equation 10 is a convex combination of the mean and LSE branches with coefficients 1 X and X
G.6 INVARIANCES AND SENSITIVITY TO PRIORS
For any constants 4j > 0 and bj. Pe,((;aau3, + bb) - aj Pa by. Multiplying prior probabilities by positive scalar and renormalizing leaves Fl.j unchanged; reweighting Pt within Mt shifts the posterior via It,5 Pt exp(Bu.g), which is exploited by the outer and inner gates.
(8)
H PARAMETERIZZATIONS OF THE PRIOR SELECTION IN FEM
Unified interface. At step I, let the accessible index set be It = {I...., O and let a nonnegative score St : Ie - Rz0 define the prior selection by
(1(1)
(i(
i ist,
(26)
with st (i) for i (causal mask). FEM then optimizes, per channel j. the DV free energy (9) B log
( 151
=()) x pc(1) B0ss.
Below we specify (hence p.) for each prior family, along with the streaming recurrences and time complexity. Throughout, M, {i 1: st(1) > 0} is the support carried into FEM (we enforce
< pt).
H.1
SOFTMAX-ATTENTIO PRIOR
Scores and normalization. Given masked scores (i) (g), kz)) b with 4 00 for t,
(1))
(i) cx)()
pc()) Erse cxp()((r))
(27)
This is the standard row-softmax causal scores.
Complexity. Matrix form softmax,ow (QKT + + MA) yields O(T²) time and O
memory (or O time, O(T) KV-cache autoregressive setting).
H.2 GATED LINEAR ATTENTION (GLA) PRIOR
Positional encoding and positivity. We inject relative position with RoPE, then map queries/keys to the nonnegative orthant:
qe = ReLU(RoPE((:)) + R%
k = ReLU(RoPE(k:)
+ Rgo
where > 0 is a small constant for numerical stability.
Decay gating. Let gr 50 be learned (scalar per-head perd//annerr))) gate and define the cumulative envelope
D; = exp Σ
(clipped in practice).
1=1
The causal time-decay factor between index and step / is K2, D,D,1 =exp== 9r)
(0, 1].
Scores and normalization. The nonnegative score and prior are
(i)
se(1) Kes (g,, k;) (i
pr(i)
Z4 Σ st(r).
(28)
Z₁
28

---

Published as a conference paper at ICLR 2026
Equivalently, with an associative scan form,
Z. (Di qe:
D
Σ s((i) =DD Dk (29) ist
rst
Hence the baseline normalized read
(Diqu, At)
(Diqr. Bi)
Streaming recurrences, Both states update in O(1) per step:
B₁ Bi-1 D. ki A = A--1 + Dr (kg 2 vz).
one-pass implementation appends a constant channel to values: T = [e 1]. A, = Ers Dj (k, & Ty)); then [uum, denj = (D.q., At) and μr
num/den. munnumdden.
Complexity. GLA preserves the O streaming complexity (per head), with the same associative-seae scan cost standard linear attention. FEM operates over the same support M, {i pi(1)= > 0} and adds one masked log-sum-exp per channel (at fixed or LTLL-cnnrolledd temperature).
H.3 LINEAR RNN-STYLE PRIORS
(L.RNN-softmax) AFT-style normalized exponential weights. and define
Let k2 E IR be per-step logits
(i) (xp(k: 1tist,
Ze = Σ exp(kr)
(30)
rst
Streaming recurrence:
S = + u Z: Z. = -i-1+
Ep.lus] = Se/Ze.
(We stabilize with k; max,se k, practice.)
(LRNN-decay) Input-conditioned exponential decay. Let gr Tegn be learned generator and define
si( =expl Σ 9- 1(i
(31)
r=i+1
With I' = exp( (xp(Ers(9)) we have se(i) = 1'1 and the streaming form
C -- ur
Σ s()v = = T,
z. ETT.
Thus numerator and denominator share the envelope Tt, preserving O(T) cost. (Conceptually, LRNN-decay recovers the decay portion of GLA without the dot-product features.)
Complexity. Both LRNN-softmax and LRNN-decay are O(T)) with O(1) updates; FEM adds one masked log-sum-exp per channel.
H.4 SSM / MAMBA-STEE -TYLE PRIORS
Positive impulse-respons SSM. impulse Ho(T) 2 0:
Consider causal linear state-space operator with nonnegative
(Sau); = Σ Holt - u) un 251
He(T) = C.AA B BT1 1 D1 = =0}
where (Aa, Ca,Ca D) stable, nonnegative discretizations.
29

---

Published as a conference paper at ICLR 2026
Scores and normalization. Set
(i)
Z. CHolt
Holt -r) (S.1).-
(32)
Both numerator and denominator come from the same scan (once with u, = Us once with u, = 1), the ((T) streaming complexity is preserved. In practice parameterize to ensure Ho(r) 2
(e.g., softplus for (A, B,C,C D) and negative-softplus for the diagonal generator).
H.5 LOCAL CONDITIONING OF THE PRIOR
Let Ct E tee be the output of a learnable, O(T) time-decay conditioner (low-rank causal convolu- tion). We modulate the prior parameters and value-pate gates by
0, G (ct), pr()) pa(-; 0t), 8; (1++))) ")), A x (M) 9. 01++m 9: 01++1+))
where (ov) and r are small MLPs. This preserves the streaming/parallel cost of the chosen prior.
H.6 SUPPORT, MASKING, AND COMPLEXITY SUMMARY
We always enforce (i) for > t and for
indices, hence M, = fi S se(1)
0}. FEM's per-channel variational step operates M, and adds exactly one masked log-sum-exp per channel (at fixed or LTL-controlled temperature), so the asymptotic time complexity matches that of the prior: softmax O GLA/LRNN/SSM O(T).
Prior family
Scores s1(1)
Complexity (per head)
Softmax attention
exp (ge, ki) + b,,} e² -012 5r (āc. Ki)
O(T2d)
Gated linear attention (GLA) L.RNN-softmxx (AFT)
O
exp(kis)
O
LRNN-decay
exp Et - Hp i). (-
S 0
O(Td)
SSM/Mamba
O
Remark (RoPE & positivity mapping). Any invertible positional transform (q, + (Tq, Tk) can precede score evaluation. In our GLA prior we use RoPE followed by a ReLUU +e mapping on both queries and keys to guarantee nonnegative feature vectors (g,, k;)) E Rm before decay gating and normalization.
H.7 WIDTH AND PARAMETER BUDGETING FOR THE PRIOR
Let the input/value width be and let FEM use a working width d the value path. We allocate a parameter ratio > 0 for the prior parameterization (queries/keys and decay gate in GLA), scaled with d. Ignoring biases and norms, the per-head linear parameters decompose into five projections:
D + Dxd + Dxd D +
D (rd) = D Ddr.
value
outer gate 9
temperature
output
prior (QIK decay)
The prior block x (rd) is split among Te, Tis projections and the decay gate. To keep the total parameter count equal to classic attention (4D*), two convenient choices are
(i) d -= or (ii) d 2D 2,
since 4Dd + Ddr 2 in both cases. In (i), the prior (Q/K) runs D-dim width-identical to standard coantionnwhile the value path uses d = 2 In (ii). the value width increases to d = 20 with balanced prior split (e.g., dim(Q) dim(K) K) 1 and the remaining budget supports the decay gate. Both settings preserve the asymptotic time complexity of the chosen prior (softmax OTT2): GLA/ O(T)).
P.
3
30

---

Published as a conference paper at ICLR 2026
1 LOW-RANK CONVOLUTION: TIME-DECAY CONDITIONER (TDC)
I.1 DEFINITION AND STREAMING FORM
Given token features 21.7 E RT let E; (æ+) and choose hidden width He D. Define
8, softplus( ( E RR", u, = 3. W. E Ik"
at = softplus T3 W.) R
with y. Wz, W, E DxH. The positive envelope is
ft = exp
87 E
(element-wise).
T=1
The TDC output the causal, input-conditionee separable convolution
h ft Σ u. 2 fi) =
Σ
exp
ST
oue E ERR.
(33)
K
A calibrated shortcut and projection produce the conditioning features: he (or.Un)))((a)) LN(h:),
C: h.W. E ER?,
where W. RR and norm(-) rescales unit & norm.
Proposition I.1 (Rank-1 1onnk----in-iine kernel and ((1) updates). The kernel equation 33 factors as K = ft @ (ft)- i.e. rank-1 in time for each channel. Hence the convolution admits O(1)
streaming updates:
C== CC-1+ + ue Ø ft,
ft Cr.
The per-sequence cost is Ho) and the per-step memory is O(Ho).
Proof. By definition, Kt,i = exp( Σ Ei=41 8r) = exp( - Ers( $7) C exp( Er
fe O (fi)-1. Substituting into equation 33 yields the stated streaming form.
L
Stability. The softplus parameterization ensures s, 2 0, hence fc E (0, element- wise; this prevents exploding envelopes and ensures well-conditioned division in u./ ff with standard stabi- lization.
1.2 COUPLING TDC TO FEM
We use disjoint slices of e to modulate (i) the parameterization of the prior selection pr(:: and (ii) FEM's value-path gates:
Prior modulation:
0e 0, + A9r
AQ,
(G))
pi(1)
(34a)
Value gate:
6;
④ (1 m
n (e) E (o C Ct:
(34b)
Outer gate: g.
gi
(1 + m
(g)
n (s)
E [n) c Ct
(34c)
Temperature gate:
A = X @ + m
(1)
m. E (n) c cc.
(34d)
FEM then applies equation 11 with (pr, Us. g:. Ax) replaced by (pe, ú, ÿ., 2.), yielding -altionn-wwar, locally conditioned selection without changing the prior's asymptotic complexity.
1.3 COMPLEXITY AND COMPATIBILITY
Proposition I.2 (Complexity preservation). TDC adds O(THe) time and O(H.) memory per layer
and does not alter the asymptotic complexity of FEM's read, which remains ( for softmax priors and (T)) for kernel/SSM priors. The per-step coupling in equation 34 is pointwise t and thus creaming-compatitne
31

---

Published as a conference paper at ICLR 2026
Relation to recent convolutional/SSM designs. TDC follows the spirit of low-rank, input-conditonne time-decay filters used in StM/DeltaNeetstylee models and the local convolutional augmentations commonly paired with Mamba-like architectures. Our use is FiLM-like: TDC lears a compact context e that modulates both the selection prior and FEM gates, providing local adap- tivity while preserving streaming costs.
Implementation notes. We apply standard LSE stabilization per channel in FEM's log-sum-exp branch, and clamp the envelope by computing ft exp(-cunsun(s:)) in log-space with an floor. The projections G(2), [ny[0), [,(*), [n]a] are small MLPs with per-channel outputs; their widths are tuned so that He D.
J FEM AS A UNIVERSAL FAST- WEIGHT PROGRAMMER
Putting the pieces together, the final Free Energy Mixer realizes unified, parallel fast-weight pro- gram:
o. - -r (1-001111)))
xx)))=
Einn. [
+ R
B- max
log iS!
(i(() (ir(t) exp(Smax 3 0b)) (35)
(bigh-entropy)
max free energy (law-entropy)
where the prior Pt and the value-path gates (0; 9t, Ae) are locally conditioned by TDC as in equa- tion 34 Equation equation shows that the mixer is simultaneously:
temporal mixer (log-sum-ens exm across indices, with causal masking and per-channel competi- tion);
entropy mixer (inner temperature via Re: meant max interpolation);
local-feature mixer (position-award modulation injected by TDC);
dual-gated mixer (inner temperature gate over indices i; outer amplitude gate timesteps t).
Crucially, the assignment capacity over the prior support attains the upper bound [Ma] P (per-channel posterior selection), the variational objective solved exactly (DV optimality), and the overall time complexity matches that of the chosen prior (softmax O(T²), kernel/SSM O(T)), up to the O convolution overhead. FEM thus serves as broadly applicable, universal fast-weight programmer that upgrades expectation-based reads to value-aware, memory processing without sacrificing paral- lel efficiency.
J.1 RELATION TO PRIOR POOLING AND SELECTION METHODS
Our Free Energy Mixer (FEM) is related to but distinct from several existing approaches:
Log-Sum-Exp (LSE) pooling. FEM is not simply generalized mean that interpolates between average and max pooling. Instead, from a Donsker- Varadhan variational view, uses values to tilt arbitrary prior distribution Pt- This yields per-channel, value-aware posteriors rather than only adjusting the softness of pooling.
Monsker-Varadhan
Entmax Sparsemax. These operate directly on the scoring distribution over (g, k), changing how probability mass is allocated. FEM instead treats this distribution as prior and introduces cross-token competition through the values. The two directions are comple- mentary and could be combined.
Gumbel-Softmax Top-k. Such methods emphasize hard selection, sampling, or ranking, often requiring non-pannen con-paalllll sampling or offline sorting. In contrast, FEM remains fully dif- ferentiable, parallel in one pass, and preserves the asymptotic complexity of the underlying prior.
32

---

Published as a conference paper at ICLR 2026
K ADDITIONAL IMPLEMENTATION DETAILS
Our detailed experimental setup is available in the linked code repository. All language model- ing experiments, including both training and inference, were conducted on 8x Nvidia H100 GPUs, while all non-language modeling tasks were trained on 8x Nvidia L40S GPUs. We use 42 as the random seed. The training and inference precision is b6 For each task, we replaced the standard Transformer block with FEM Transformer block, substituting the attention layer with FEM-(SM, GLA, Mamba, AFT), while keeping all other settings unchanged to ensure a fully con- sistent experimental environment. Parameter budgeting was carefully applied to keep overall model size and architecture comparable to the baselines. The additional low-rank convolution used in our parameterization introduces less than 1% extra parameters (with H. = d/16)
Special configurations required by the experimental setup when specified. Otherwise, all linear projections are randomly initialized from a centered normal distribution with a standard deviation of 0.02. All biases and embeddings are initialized to zero. For the maximum inverse temperature, we initialize to zero and then apply the parameterization softplus(x+1. 3) to ensure that its initial value is around and remains strictly positive throughout training.
L ADDITIONAL DATASET DESCRIPTION
Language Model Evaluation Setup. We adopt the Open L.IM Leaderboard (OLL) protocol and complementary suite of general-ability tasks. The Open LLM Leaderboard core covers MMLJU Pro (5-shot, accuracy), GPQA (O-shot, normalized accuracy), BBH (3-shot, normalized accuracy), MATH (4-shot, exact match), and MuSR (O-shot, normalized accuracy), plus IFEval for instruction following, where we report striet pass rates for instruction- and prompt-level constraints (Wang
et al. 2024: Rein et al., 2023 Suzgun al. 2022, Hendrycks et al., 2021 SSprague et al.) 2023 Zhou et al. 2023). Following OLL, we use the normalized-accuranl
metric acc, for multiple- choice tasks, which subtracts the random-quess baseline and rescales scores a common range for fair cross-task comparison (Hugging Face 2025). To broaden coverage, we also evaluate on widely used general-ability benchmarks: ARC (Challenge/Easy)), HellaSwag, PIQA, BoolQ. WinoGrande, COPA. OpenBookQA, and SciQ. reporting accuracy or acC, as standard; unless noted, these are evaluated in 0-shot (Clark et al., 2018 Zellers et al. 2019 Bisk af., 2019 Clark al., 2019] Sakaguchi et al. 2020; Roemmele et al.
2011 Welbl et al. 2017). We perform the evaluations with
im-evaluation-harr (Gao et al. 2021
We assess our architecture using the Mechanistic Architecture Design (MAD) framework, recently introduced methodology for cost-efficient evaluation of deep learning models Poli et al. 2024). MAD provides set of capability-focure focussed benchmarks benchmarks-innluddnn in-context recall, fuzzy recall, selective copying, and compression that probe core sequence modeling abilities. It has been validated across more than 500 language models ranging from 70M to 7B parameters, showing a strong correlation between performance on these synthetic tasks and compute-optimal perplexity at scale. By leveraging MAD as a reliable predictor large-scale behavior, we can identify architce- tural advantages without relying the prohibitive compute costs of full-scale training.
MAD
Time Series Forecasting We evaluate our module on several standard time series forecasting benchmarks, following the setup of Lu & Yang (2025). (1) Weather (Wu al. 2021 P 21 me- teorological variables (c.g.. temperature. humidity) collected every 10 minutes in 2020 from a Ger- man weather station. (2) Solar (Lai et al. 2018 Solar power output recorded every 10 minutes
in 2006 from 137 U.S. photovoltaic plants (3) ETT (Zhou et al. 2021 f Transformer load and temperature data from July 2016 to July 2018, sampled at 15-minute (ETTm1/E I/ETT/2)T22) and hourly (ETThI/ETTh2) intervals, covering 7 key operational features.
http: http::///www...e.oooo//////////0
33

---

Published as a conference paper at ICLR 2026
M FEM-MLP INTERACTION: AN NTK AND FAST- WEIGHT PERSPECTIVE
This section provides a mathematical description of how the Free Energy Mixer interacts with the MLP within Transformer block. FEM performs per-channel, value-aware fast-weight update that adapts the effective kernel directly at the mixing site, while the MLP focuses on feature synthesis and processing. This division of responsibilities improves data efficiency and preserves both the parallelism and the asymptotic time complexity.
M.1 NOTATION AND PLACEMENT WITHIN A BLOCK
adopt the column-vector convention. Let h2 IR the hidden state at index S t, and let the value path working width d:
1/2 W T Me, =1 xx,e) Me, + de, Funax
O2 It C (Wou))
where e [d] indexes the value channels. Here
Pr.c (i(()) Vi,0s
f,e
log
Σ (i( Amax
Brnx
ist
ist
and = (u..1.... " E Rd. The matrices Wy E Bxd and Wo E R Dxd are the value and output projections, and we define
by: T:=Wooe Rd
as the j-th output direction. The outer gate 9t E [R..0 is applied elementwise. The selection prior Pt A([1,... +{}) is nonnegative and masked, supplied by mechanisms such as softmax attention, kernelizable/lincar attention, linear RNNs, or SSMs. FEM replaces the per-head convex read with per-channel free-energy read computed the same masked support.
FEM weights and LTL. For channel e E [d]. the inner (temperature) gate Xe, E [0, 1] defines the effective per-channel weights
(i( exp(B vi
(1,c(1) (1 )e,e) p((() + de,c Pi.m (t),
2r,0
(c)
PostP exp vr
where 4ú,8 is the value-aware posterior over the masked support. Linearized temperature learning (LTL) shows that
(1- pe.e dt, Fre Fimx Fr.((Sha,t.e))
for a unique hidden temperature Bii....c E Bmax] that varies strictly monotonically with A Hence optimizing dt, is equivalent to optimizing B while requiring only expectation and single masked log-sum-exp per channel.
M.2 LOCAL FREE-ENERGY GEOMETRY: FAST WEIGHTS AT THE MIXER
Fix step and channel c. The per-channel free energy and its posterior are
pr(1) Buri
(B,) = log iSt
Σ B
Erse Pe(r) eBu.."
The gradient equals the posterior and the Hessian is Fisher covariance scaled by B:
Vo Pe,)(B) - 22 PB. = (Dag(()) 44" 1 x 0,
[r))ll S 8/22
Moreover,
(r) = Ep, [v.cc (T.
(3) 3 20
with exponential posterior concentration under value margins. These properties identify FEM as an exponentiated-graee -radieet (mirror-ascent) ascent) fast-weight update applied to the prior pt along the value direction v.c-
34

---

Published as a conference paper at ICLR 2026
M.3 VALUE-PATH JACOBIAN AND NTK CONTRIBUTION INSIDE TEM-MLL BLOCK
Treat Or and ) as constants with respect to Wv when forming the NTK for the value pa- rameter group. Writing by = Woe ej Rd for the j-th output direction and wx(i) (we,1 (i),.... ...,, (r,a(i)) Rd we have
Don. OWv
Ir.
Σ hi DDag(( (1)) 6,]" E Dxd ist
(f) = =11 (1- pe(d) + de, Ri (t).
Hence the value-path NTK contribution between outputs (t.j) and (8,j") is
K
(V) (1.))1a5))
/ous Sw
do,.y' dwv /F
ΣΣ
(he hT) (w) (1) T (
ist
[(Wo)ss
ΣΣ (r,(()) (r) (he, ist rsa
= g
This form makes explicit both the role the output projection (via by) and the per-channel, value- aware weights (i).
Channel-token rank. Let Welassic W.FEM R collect token weights per channel the value path:
(Welassie) ),,t = (i)),, (WFEM) )c, = wi,((i).
Then rank WWelasssic = 1 (all rows identical), whereas rank(W,FEM) can reach min d, t} because channels select different indices in the same step. This rank gap explains why single vex read cannot realize generic channel-wise selectors, while FEM can approximate them at finite temperature.
M.4 GATE-INDUCED CROSS-TERMS AND PRIOR SELF-CORRECTION
Beyond the direct value-path term above, additional chain-rule terms arise because or, depends on the inner gate At,5 and the outer gate It.j- both of which are parameterized from value-derived statistics:
Dor.j dWv
.or. LoWY
+ Dor.
Ohg Hp,. F!nax)
do,, dg.. a(-)
+
direct
Ahes Oiu, Fmnax
dw
0g..j 0(-) OW
The first term yields K(V); the second and third induce additional NTK components K(-VV) and (())) that propagate value information through the gate parametrizations. These contributions should be accounted for separately in parameter- group NTK decomposition.
FEM also contributes a prior-pate NTK term through the sensitivity of the free energy to the prior logits. Pt moftmax(b)) then
V (L.(3) = (1
9 ( -p)..
Gradients with respect to the query and key parameters therefore move the prior pt toward the value- aware posterior q in single pass, implementing online kernel adaptation the score level.
M.5 COMPOSITION WITH THE MLP: ROLE SEPARATION IN THE NTK
Let denote the MLP the same block, and Jo(o:) its Jacobian. The block-level NTK decomposes as
Krinek (1,,), (8, (or) (L,M( (s.1') Jo(o.)T + (s.
with
KFEM K(O) + R(-eV) R(o) K(QK) + cross-terms.
35

---

Published as a conference paper at ICLR 2026
Since KVV) already encodes value-aware, per-channel cross-token competition, the MLP receives inputs whose coordinates have undergone a selective, information bounmation-budeeted fast-weight update. The MLP therefore concentrate feature synthesis and consolidation, rather than attempting to reconstruct channel-wise index identities that are provably lost after a first convex mixing without FEM.
M.6 CONSEQUENCES FOR THE MIP'S WORKLOAD AND EMPIRICAL SIGNATURES
This analysis suggests two empirical signatures, both observed in our ablations. First, adding the LSE branch (+L) and temperature control (+T) yields large gains, they enable value-aware com- petition at the mixing site and reduce the load on downstream MI.ss Second, when the mixer handles festhweight programming and the MI.P focuses on feature synthesis and knowledge consol- idation, FEM improves data efficiency on retrieval-heavy and algorithmic tasks without increasing parameter budgets.
Summary. FEM upgrades the readout from head-synchronous convex average to a per-channel, value-aware variational update with stable geometry and single-pass temperature control. In an NTK view, this an online kernel adaptation step the mixer, after which the MLP performs feature synthesis on selectively retrieved representation. This yields clean separation of roles that preserves parallelism and asymptotic complexity while expanding the class of functions that can be realized in single block.
N ADDITIONAL DISCUSSION: SIGNIFICANCE AND MOTIVATION
N.1 HIGH-LEVEL PERSPECTIVE FROM LLM SCALING LAWS
Let
f size ability
denote the empirical scaling law that maps model size (under comparable data and optimization) to emergent capability. Viewed through this perspective, research on attention mechanisms naturally splits into two complementary directions.
Efficiency-orientn work. This line of work keeps the semantics of attention essentially un- changed, but makes each point the curve / cheaper realize. Typical approaches include kemel- izable or streaming variants that replace the growing K KV-cache with fixed-width sufficient statistics, enabling efficient scans while trying to preserve the original attention behavior. Such improvements extend the practical regime of the existing scaling law without modifying the underlying function f.
Expressivity-oriented work. This line aims to increase capability at a fixed parameter budget by enlarging the architecture's representable function class. In attention, central bottleneck is the read: standard attention stores values losslessly but mixes them via a token-separable convex aver- age, which synchronizes channels and blocks even simple per-channel indexing. Recent works (like Differential Transformer) has started to relax this constraint by extending convex mixing to affine mixing with signed (including negative) weights, allowing more expressive selection behavior from the same value bank (Ye et al. 2025 Lv et al. 2025). In the same spirit, replacing expectation-style reads with richer, value-aware per-channe mechanisms seeks to upgrade the algorithmic expressive- ness of the architecture without changing its asymptotic complexity. In the context of scaling laws, such improvements aim not merely to move more efficiently along the curve f size - ability, but to fundamentally reshape the function f itself by increasing ability fixed model size.
36

---

Published as a conference paper at ICLR 2026
N.2 A
OFF PATH IN MODEL STRUCTURE
The design of attention mechanisms, from the perspective of the developmental progression of both effective and efficient sequence modeling, can be regarded as follows:
g( x1a) of a at, = softmax(q) k.) + O: = ((g)) TSe
f(11a) (0)
((a))
(1)
(3)
where, at the functional level. both f and g can be realized (up to arbitrary precision on com- pact domains) by flattening their arguments and applying sufficiently wide MLP. and where S. = Eise (kk)) and 2 = Eist ((k;) are width sufficient statistics in linear/kernelized attention. Each rightward step lowers asymptotic cost and, crucially for causal training, strengthens causal training parallelism: the ability to compute all T token-level losses in one parallel for- ward. Nodes (0)-(1) lack efficient suiput-paralle, sharing across timesteps (each target requires its own materialized context or independent map), whereas Nodes (3)-3) obtain output parallelism by sharing global operators across all steps: a masked score matrix applied to a shared value bank for (2), or associative scans over fixed-width states for (3).
Intermediate designs between the nodes. Along the compute-expressiiig frontier based on soft- max attention, there are two movement directions. A leftward move (between Nodes (1) and (2)) accepts modest increase in read-side cost to gain expressivity at fixed parameter budgets; examples include separabee nonlinear mixers such as Σ, Or,, (3, @ v;) and FEM. A rightward move (between Nodes (2) and (3)) accepts some loss in representable functions and memory storage to obtain larger efficiency and stronger causal training parallelism, like gated linear attention (GLA) and grouped query attention (GQA). These designs therefore occupy different, reasonable operating points on the frontier, reflecting distinct trade-offs between computation and expressiveness.
A consolidated comparison. Table summarrissssss compute cost, causal training parallelism, mem- ory states, representable classes, and minimal non represeentaal.... The overall pattern is clear: mov- ing toward more efficient computations (matrix-paralle evaluation or fixed-state streaming/scan) systematically reduces what the model can represent. Conversely, recovering same-step channel- wise selection requires paying for additional capacity, either through richer read-side interactions (such as per-channel, cross-token competition before the first mixing) or through increased architec- tural resources such as larger states or more heads.
N.3 TWo DIRECTIONS FOR ATTENTION: EFFICIENCY AND EXPRESSIVITY
Efficiency direction. This line preserves the expectation semantics of the read and reduces cost via factorization, sparsity/low rank, fused kernels, state-collapsed streaming mechanisms. ex- tends the practicable reach of the existing scaling law by making the same (size) cheaper to realize longer contexts or larger batch sizes. However, inherits the limitations of sekarabee ex- pectation reads and, for fixed-state designs, the consequences of state collapse.
Expressivity direction. This line retains lossless storage and upgrades the read so that the KV cache functions as a dynamic selection database capable of per-channel indexing. From a classic data-structure viewpoint, a per-head convex average cannot implement even basic two-dimensional array indexing in one step; FEM remedies this gap by converting the score-induced prior into value-aware posterior that can concentrate independently per channel, all at the prior's asymptotic cost. In scaling-law terms, the goal is to raise ability at fixed size, effectively lifting or reshaping the function f. rather than only moving more economically along it.
Directional summary. Starting from Node (2), efficiency work tends to move further rightward, pushing streaming and fixed-state implementations while accepting additional expressivity losses. Our study moves leftward from Node (2) toward Node (1): we keep the prior's time complexity and masking semantics but restore per-channel selection at read time, thereby turning the lossless mem- ory storage into enhanced algorithmic expressiveness and general ability under the same parameter budget. This complementary direction targets the mechanism of the scaling function f itself, by increasing ability at fixed size rather than merely reducing the cost of reaching the same ability.
37

---

Published as a conference paper at ICLR 2026
O ADDITIONAL RUNTIME EFFICIENCY ANALYSIS
We replace the attention layer in a GPT-2 style Transformer (hidden size 768, 12 heads, 12 lay- ers) with various alternatives and evaluate all models at sequence length 1024. We experiment with three prior types within our Free Energy Mixer (FEM); softmax attention (SM), gated linear atten- tion (GLA), and Mamba. To isolate the overhead of the free-energy branch, report results for the branch alone (denoted +L,T,G as well as for the full model including the convolutional path (+C,L,T.G), highlighting the impact of these components on computational efficiency.
In FEM, the implementations of softmax attention, gated linear attention, and Mamba are taken from the F1 illashAttentic (flash_attm)) Flash Linear Attent on (fla), and Mamba (mamba.ssm) libraries, respectively (see our code repository for details). In addition, we com- pare against recent basclines: RWKV7, linear attention (LinAttn), HGRN2, PaTH Attention, and Gated Slot Attention, whose implementations are all taken from F Flash Linear Attention (fla). All runs use single NVIDIA L40S GPU, batch size and FP32 precision. We report the mean latency over 500 steps following 50 warm-up steps. Fwd denotes full-sequence inference without KV or hidden-state caches. FLOPs are estimated using the PyTorch profiler on the same configuration: FwdFLOPs measures the total number of floating-poin operations in a single for- ward pass, TrainFLOPs in a single training step (forward + backward), and the per-token metrics are obtained by dividing by the number of tokens in the batch.
As shown in the table below, even without optimally fusing the operations in the free-energy branch with the mixing backend kernels, the additional cost introduced by the +L.T,G branch remains rea- sonable. Moreover, even with the full structure (+C,L,T,G), FEM with softmax attention and Mamba backends still achieves upper-mid performance among recent baselines. Thus, even if we treat fur- ther kernel- and system-level optimizations as future work, the current FEM-SM implementation already offers computational advantage over PaTH Attention, recent compusation-engineered based architecture.
P TOY PROBLEM: SINGLE CHANNEL-WISH ARGMAX
We construct a minimal synthetic task illustrate the inability of single convex read to perform channel-wise selection and to contrast it with FEM.
Data. Each sample contains value matrix V RRx', For every channel j. a winner index aj is drawn uniformly from {1...., T} and we set
Vay. A Ej.
Vigays ~ (O,,
with margin A and std level 0-0.55. The target output is the channel-wis. maximum
0ブ 15iST
max Vig.
Since different channels typically select different winners, y² almost always lies outside the convex hull of [V...} and cannot be produced by a single convex mixture. We use T2- D=512, number of heads H = 4, with 200,000 training and 2,000 validation examples.
Models and training. We compare two single-layer value mixers that map input V E RTxD to an output in R2 taken from the last position directly after the value mixing. 1) Softmax attention causal self-attention selection producing multi-head convex combination of the value vectors above. 2) FEM: single FEM layer with the same softmax prior and free-energy read of the value vectors. We disabled the outer gating and low-rank convolution modules of FEM for more fair comparison. Both models are trained with mean-squared error, Adam W, batch size 64, learning rate 10-2, for 2,000 steps.
Metrics. We report MSE and a per-channel index accuracy measuring whether the correct winner index recovered. Given prediction y E TR. the predicted winner for channel j is
0 arg min (Vi, 3²)² 1SST
Index accuracy is the fraction of channels with à; = aj. averaged over the validation set. Random guessing yields accuracy 1/T (about 0.8% for T = 128).
38

---

Published as a conference paper at ICLR 2026
Validation MSE vs Steps
Per-Channel Index Accuracy
1.0
FEM-IL (val)
(val)
0.6
300
200
0.2
100
FEM-11
SoftmaxAltr-1L
250
500
750
1000
1250
1500 1750 2000
250
500
750 1000
1250
1500
1750 2000
steps
staps
Figure 3: Single-layer toy per-channel argmax task. Left: validation MSE over training steps. Right: per-channel index accuracy. FEM rapidly fits the channel-wise argmax, while a softmax attention layer stays near chance level, reflecting the limitation of convex mixing.
Results. FEM quickly learns the per-channel argmax mapping, while single softmax attention selection does not. FEM's validation MSE decreases steadily and its index accuracy rises from chance level essentially 100% for 128, D=512, recovering the correct winner in every chan- nel. The softmax attention baseline shows almost no improvement: its MSE remains near the initial value and its index accuracy stays around - 2%, close the random baseline 1/T.
39

---

Published as a conference paper at ICLR 2026
Node
Form
Time complexity Memory state and causal training (read-time) parallelism
Representable class vs. strict non-representable
(0) General map
f(11a)
Naively O(T²); no Full Fis (lossless
All measurable causal maps, no exclusion
efficient causal storage and training
readout)
parallelism: each target loss requires materializing distinct, large flattenedppadded
context (no shared value bank across outputs)
(1) Additive map-reduce
01
Naively O with map-reduce; token-map parallel but output-wise reduces are independent: maps cannot be shared across different outputs, yielding weak causal training parallelism
Full 3iæ (lossless storage and readout); Each g can read the full prefix
Additive over tokens at the output: loses the one-hot, flattened positional addressability
Σ, Eigt ( Z14)
since all token contributions additively superposed in the same space, so positional
matching must be recovered through learned positional mechanisms with appropriate
inductive bias;
excludes
non-additive
global functionals. Image per head lies in conv (u;) with channel- synchronized
(2) Value-path linearization
0==
Matrix-parallel O(T²): all T outputs in masked QKT pass and single AV: full causal training parallelism via shared value bank and linear mixing
Requires full lossless Tidt storage; achieves parallel efficiency by replacing per-channel
Eiso 21,5% a
softmax(q k)
weights; excludes same-step channel-wise
dynamic readout with per-head
indexing (c.g.. coordinate wise argmax) unless all selected indices coincide; heads at most t patterns;
convex mixing
(3) Fixed-state kernelization
O: p(g)) St
O(T) with associative scans; full per-pass causal training parallelism via shared fixed-widie width lossless. prefix statistics
Memory storage collapses to fixed-wiie width sufficient statistics (Sc, 2t). Not
Depends only on sufficient statisties; identity-level retrieval
olae) "'""
S. o (k:))o
2 2
impossible at long horizons; Gating may enrich the prior yet remains state-collapsed.
Table 6: Compute-eppresiviity trade-offs across four nodes. Rightward moves reduce cost and increase causal training parallelism, but each step provably removes function families available to carlier nodes.
40

---

Published as a conference paper at ICLR 2026
Table 7: Latency, throughput, and peak GPU memory on GPT-2 (768/12/12) setting, sequence length
1024, batch size 8, FP32 single L40S GPU. Fwd full-sequence inference.
Model
Latency Latency Throughput Throughpur MemPeak MemPeak
Train
Fwd
Train
Fwd (s) Train (s) Fwd (lok/s) Train (lok/s) Fwd (GB) Train
FLOPs
FLOPstok
FLOPs/ok
RWKV7
0.0832 0.2367
98417.56
34610.31
11.08
11.48
1.47T
4.41T
179.35M
338.09M
LinAttn
0.0652 0.1368
125587.32
59881.77
6.01
6.16
1.39T
4.18T
169.94M
509.75M
HGRN2
0.0680 0.1481
120405.72
55326.95
6.66
6.94
1.39T
4.18T
169.89M
509.65M
PaTHAttention
0.0693
0.2098
118161.31
39051.39
5.62
5.82
1.40T
4.21T
171.29M
513.84M
GatedSlotAttention
0.0765
0.1901
107069. 10
43082.50
7.72
8.08
1.51T
4.52T
184.04M
552.11M
SM (Naive)
0.1306 0.3334
62740.89
24574.18
9.61
10.34
1.70T
5.11T
207,84M
623.41M
SM (Flash)
0.0649 0.1471
126304.94
55697.16
4.89
4.99
1.39T
4.18T
169.89M
509.65M
SM (-L,T.G)
0.0693 0,1615
118154.14
50733.90
5.82
5.92
1.39T
4.18T
169.92M
509.75M
SM (+C.L.T.G)
0.0732 0.1853
111955.25
44217.70
7.67
7.77
1.40T
4.21T
171,43M
514.19M
GLA
0.0824 0.2213
99389.45
37016.71
9.95
10.17
1.70T
5.11T
208.03M
623.90M
GLA (+L.T.)
0.0875 0.2344
93674.03
34946.88
10.89
11.08
1.70T
5.11T
208.07M
624.02M
GLA (+C.L.T.G)
0.0934 0.2566
87675.64
31924.23
12.74
5.15T
209.57M
628.46M
Mamba
0,0606
0.1249
135144.49
65603.32
4.60
1.28T
3,84T
156.33M
468.97M
Mamba (+L.T,G)
0.0586
0,1438
139778.42
56978.94
6.77
6.87
1.22T
3.67T
149.29M
447.83M
Mamba (+C.I.T.G)
0.0647
0.1831
126581.53
44743.67
7.97
8.07
1.23T
3.70T
150.44M
451.25M
41


Reviews:
[{"review_id": "NakhoqSmCB", "reviewer": "Reviewer_YvfL", "summary": "The paper introduces the \"Free Energy Mixer\", a novel way of mixing values within the attention block. In the standard attention operation, the $Q$-$K$ dot-product determines the probability distribution over the sequence length, and the output is a weighted-average of the $V$ vectors accordingly. This operation mixes all channels in the same way however, so it cannot do an operation such as per-channel maximum. By posing the problem as a DV free-energy problem, the authors pose the read-out as a mix between the \"low temperature\" extreme (choose the maximum value per channel), and \"high temperature\" extreme (weighted average), controlled by a learnable parameter, thus enabling smooth control between the two regimes. Similarly, the technique can be applied to recurrent architectures such as RNNs and SSMs, by re-interpreting the prior distribution. The authors showcase that, by using this architecture as a drop-in replacement for attention/recurrence, it can improve the performance of the models.", "questions": "- Could the authors clarify the exact FLOP overhead relative to standard attention? While they state that the asymptotic complexity is preserved, a quantitative estimate of constant-factor costs would be informative.\n- Could the authors clarify if the method could be fully fused with the hardware-efficient optimisations, such as Flash Attention? Are there any difficulties that might be an issue?\n- Similarly, a report on the total wall-clock time differences between the methods would be helpful.\n- Minor: I was not familiar with the $p \\in \\Delta$ / $p \\in \\Delta(M)$ notation, the authors could potentially define it in the text.\n- Minor: Bold/underline notation could be clarified in the tables.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- To the best of my knowledge, the authors introduce a novel technique for mixing the value vectors in the attention block, enriching the expressivity of the operation. This approach can be applied to both attention-based models, as well as recurrent variations.\n- The authors do extensive evaluation on reasonable model/training budgets (1.3B parameters/100B tokens, 340M parameters/15B tokens), covering language modelling, image modelling, and time series forecasting.\n- The authors account for the architectural differences by parameter-matching during their evaluations.\n- Figures 1 and 2 are clear and instructive.", "weaknesses": "- I am not entirely convinced that the \"lossy mixing\" of the classical attention is obviously a significant limitation. While the theoretical justification seems compelling, it would be interesting to understand how different the proposed mixing ends up being from the standard attention after training. Nonetheless, this does not detract from the paper’s contribution, as the empirical results indicate consistent benchmark improvements.\n- The proposed method appears to add some constant computational overhead (due to the per-channel log-sum-exp operation), even if the asymptotic complexity remains unchanged. This means that, although the methods might be parameter-matched, they are not necessarily FLOP-matched.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HMKx5rYj3Y", "reviewer": "Reviewer_hSiq", "summary": "This paper identifies a \"lossless-storage versus lossy-processing dilemma\" (Section 1) in standard attention mechanisms. The authors argue that the attention readout, a \"per-head convex average\" , prevents channel-wise selection from the key-value cache, meaning it \"cannot realize a generic channel-wise selector\" (Section 2.1, Corollary 2.3). To address this, the paper proposes the Free Energy Mixer (FEM), a new layer that replaces the standard attention readout. FEM uses a \"free-energy (log-sum-exp) read\". This method treats the standard (q, k) attention distribution as a \"prior\" and applies a \"value-driven, per-channel log-linear tilt\" to generate a \"value-aware posterior read\". This formulation allows the model to \"smoothly moving from averaging to per-channel selection\". The proposed FEM is \"plug-and-play\" with various architectures (e.g., standard attention, linear attention, RNNs, SSMs) and \"preserves the corresponding time complexity\" (Section 2, Contribution 3). The final model is a \"Two-Level Gated FEM\" that uses \"Linearized Temperature Learning (LTL)\" (Section 2.3.1) for efficient dynamic temperature control.", "questions": "1. Regarding Weakness 1, could the authors provide a more fine-grained latency breakdown? The \"LSE Implementation\" (Fig 2f)  computes both the mean and the LSE branch. What is the measured latency and throughput overhead of only the FEM components (L, T, G) when applied to an already-optimized baseline (e.g., FlashAttention), rather than the unoptimized baseline in Table 5?\n\n2. Regarding Weakness 2, please elaborate on the distinction with LASER attention. The claim that it \"do[es] not address the lossy processing limitation\" (Section 1)  is critical. How does FEM's formulation of LSE as a \"value-aware posterior\" derived from a \"prior\" fundamentally differ from LASER's \"Attention with exponential transformation\"  in addressing the channel-wise selection problem?\n\n3. Regarding Weakness 3, how sensitive is the model to the $\\beta_{max}$ hyperparameter? The paper describes it as a \"learnable global maximum inverse temperature\" (Section 2.3.2) and initializes it near 1.0 (Section K). What range of values does $\\beta_{max}$ converge to during training? How does performance vary if $\\beta_{max}$ is fixed at different values (e.g., 1, 5, 10)? (By the way, a more theoretical explaination will be preferred than empirical demonstrations for this question)\n\n4. **KV-Cache Reduction:** The paper states that \"budgeting strategy (i)\" \"reduces the dimension of the value part needed to be stored in the KV-cache by half\" (Section 2.3.3). This seems to be a significant practical benefit (halving the KV-cache memory), but it is not highlighted in the introduction or abstract. Can the authors confirm this interpretation and, if correct, was this memory reduction reflected in the experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper makes a solid contribution to the field.\n\n**Originality**: The primary contribution is the principled reframing of the attention readout as a \"variational free-energy optimization\" (Section 2, Contribution 2). This provides a novel, value-aware alternative to the standard convex-combination readout. The \"Linearized Temperature Learning\" (Section 2.3.1)  is also a clever technique for maintaining efficiency. The originality is slightly tempered by prior work on LSE-based mechanisms, such as \"LASER attention\" (Section 1), which the paper acknowledges but does not sufficiently distinguish itself from.\n\n\n**Significance**: The contribution is significant. The \"plug-and-play\" (Abstract) and \"agnostic\" (Section 2, Contribution 3)  nature of FEM means it could potentially serve as a general-purpose upgrade for many sequence models, including SSMs and linear RNNs. The empirical results, showing consistent gains across \"NLP, vision, and time-series\" (Abstract), validate its value.", "weaknesses": "1. **Practical Efficiency and Latency**: The paper's primary weakness is the gap between theoretical and practical efficiency. While FEM preserves \"asymptotic complexity\" (Abstract) , Table 5 shows a significant practical latency cost: \"FEM-SM\" (0.017s) is substantially slower than the baseline Transformer (\"FEM-SM (-G,T,L,C)\", 0.012s)  in the forward pass. The authors' \"Limitation\" (Section 4) section admits this is due to a \"lack of fused CUDA kernels\" (Section 4). This is a critical barrier for adoption, as the community heavily relies on optimized kernels (e.g., FlashAttention).\n\n2. **Novelty in Relation to Prior Work**: The paper needs to more clearly differentiate itself from prior work using LSE. The paper mentions \"nonlinear mixing such as log-sum-exp in LASER attention\" (Section 1) and \"Tropical attention\". The paper claims these \"do not address the lossy processing limitation\" (Section 1), but this distinction is not elaborated upon. A more detailed comparison is needed to establish why FEM's specific free-energy formulation  succeeds where these other LSE-based methods allegedly fail.\n\n3. **Architectural Complexity**: The final \"Two-Level Gated Free Energy Mixer\" (Section 2.3.2) is a complex module. It combines the core LSE mechanism with a \"low-rank convolution\" (Section 2.3.3) , an \"inner (temperature) gate $\\lambda_{t}$\", and an \"outer gate $g_{t}$\". The ablation study in Table 1 (e.g., \"FEM-SM\" Avg 80.2 vs. \"FEM-SM(-G)\" Avg 79.4) suggests all these components are necessary for optimal performance. This adds many new hyperparameters (e.g., $\\beta_{max}$, $H_{c}$) and increases the difficulty of implementation and tuning compared to standard attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yes1R18dOt", "reviewer": "Reviewer_bNwC", "summary": "This paper proposes the Free Energy Mixer (FEM), a novel inference-time mechanism that addresses a perceived limitation in standard attention: the inability to perform per-channel selection from the key-value cache. The authors argue that the standard convex combination of values within each attention head synchronizes channels, preventing them from independently selecting different past indices. FEM tackles this by reframing the read operation as a variational free-energy optimization. It uses a value-driven, per-channel \"tilt\" on a prior distribution (e.g., from Q/K scores) to produce a posterior, enabling a smooth interpolation from averaging to near-hard selection. The paper presents a two-level gated instantiation of FEM that is plug-and-play with various backbones (softmax/linear attention, RNNs, SSMs) and demonstrates consistent performance gains on NLP, vision, and time-series tasks under matched parameter budgets.", "questions": "For details, please see weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Motivation and Problem Identification: The paper compellingly argues that the standard per-head convex combination of values creates a \"lossy read\" bottleneck, preventing channel-wise selection. This is a well-articulated and non-trivial insight into a potential limitation of standard attention mechanisms.\n\nTheoretical and Empirical Rigor: The proposed FEM method is grounded in a solid theoretical framework derived from the Donsker-Varadhan variational principle. The paper is supported by thorough theoretical analysis and extensive experiments across diverse domains (synthetic tasks via MAD, language modeling, image classification, time-series forecasting) and model architectures (FEM-SM, FEM-GLA, etc.). \n\nPracticality: A key strength is the method's plug-and-play nature and its ability to preserve the asymptotic time complexity of the underlying prior mechanism (e.g., O(T²) for softmax, O(T) for linear attention), as demonstrated by the latency/throughput analysis for FEM-SM.", "weaknesses": "1.  **Baseline Comparisons:** While the empirical results are comprehensive, the baselines could be strengthened by including more recent architectures that also introduce channel-wise inductive biases. A comparison against methods like Mamba (with its data-dependent SSM) or Multi-head Latent Attention (MLA), which inherently manipulate channel interactions, would provide a more rigorous assessment of FEM's unique contribution beyond simply adding channel-specific gating or modulation.\n2.  **Core Motivation and Justification:** The central premise—that enabling full per-channel selection is a critical goal—may be somewhat overstated.\n    *   It potentially creates a false dichotomy between the roles of the attention mechanism (for token mixing) and the subsequent Feed-Forward Network (FFN, for channel mixing). The universal approximation theorem suggests that a standard Transformer block, with MHA and FFN, is already capable of learning complex functions, and the necessity of extreme channel-wise selection within the attention block itself is not fully justified.\n    *   Many studies indicate that attention heads exhibit significant sparsity and redundancy. Therefore, the performance gains observed with FEM might not stem from achieving perfect channel-wise selection, but rather from introducing a more effective and general form of channel-wise *modulation* or *gating*. Reframing the contribution around enhancing channel *utilization efficiency* rather than enabling a theoretically maximal selection capacity might be a more compelling and accurate narrative.\n3.  **Efficiency Analysis:** The computational cost analysis, while provided for FEM-SM, is incomplete. The overhead of FEM when applied to its linear-time variants (e.g., FEM-GLA, FEM-Mamba) is not thoroughly quantified. In large-scale training and inference, even a modest constant-factor overhead for linear-time models can become a significant practical bottleneck. The performance benefits of FEM might diminish when scaled to very large models and datasets, while the computational cost remains.\n\n**Questions and Suggestions for Authors**\n\n*   Could you include comparisons with other modern sequence models that incorporate strong channel-wise interactions (e.g., Mamba, MLA) to better isolate the benefit of FEM's specific variational formulation?\n*   The paper would be strengthened by a more nuanced discussion on the interplay between attention and the FFN. Could the observed gains be interpreted as FEM providing a more powerful form of channel-wise conditioning *before* the FFN, rather than strictly enabling a previously impossible \"channel-wise selector\"?\n*   Please provide a detailed latency/throughput analysis for the linear-time FEM variants (FEM-GLA, FEM-Mamba) to give a complete picture of the method's computational trade-offs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L8ofscflkX", "reviewer": "Reviewer_ajEo", "summary": "This work observes that regular attention, together with most variants, is restricted to outputting a value mixture (before output projection) that lies within the convex hull of values. Certain operations, for example channel-wise argmax, are not fully representable (because the only selection signal comes from the q k product, which is broadcast over the value dimension before a linear operation). The authors propose the free energy mixture (FEM), which is a channel-gated interpolation between a regular weighted mixture of values, and a log-weighted-exp mixture of values with temperature. This can be augmented by a lightweight low-rank convolution at the inputs. Experiments on artificial tasks, language modelling, vision and time series data show the efficacy of this approach, which can be applied on top of a \"prior\" of softmax attention, gated linear attention or Mamba, outperforming regular softmax attention and other variants on a parameter-matched basis.", "questions": "My main questions are in \"specific concerns\" above.\n\n**Clarifying question** I understand the preferred FEM layer (with softmax prior, linearised temperature, conv & gating) to be something like the following pseudocode. I would appreciate any corrections.\n\n```\n# given x\n# given p = softmax((q @ k.T) / sqrt(d))\n# given parameters {W_{c,v,g,lam,o}, w_beta}\n\nc_v, c_g, c_lam = low_rank_conv(x, W_c)\nv = (W_v @ x) * (1 + c_v)\ng = rms_norm(softplus((W_g @ x) * (1 + c_g)))\nlam = sigmoid((W_lam @ x) * (1 + c_lam))\nmu = p @ v\nbeta = softplus(w_beta + 1.8)\nF = log(p @ exp(beta * v)) / beta\no = g * ((1 - lam) * mu + lam * F)\nreturn W_o @ o\n```", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The key idea seems sound, the empirical investigation shows promising results, and the paper contains many theoretical justifications for the proposed design. While many works focus on making attention more efficient, this work has an equally relevant focus of making it more expressive. I particularly appreciated:\n\n - Comprehensive ablations of the four major architectural components in Table 1 and Table 2.\n - The intuition and results concerning when attention mixing mechanisms are constrained to the convex hull of values (2.2-(4) and C.4).\n - Empirical results from multiple domains and with multiple choices of prior.\n - Parameter-matching for a fair comparison across techniques.\n - Ability to reuse existing attention implementations by stacking/unstacking values.", "weaknesses": "While I am generally positive about this paper for the reasons described above, I find it hard to follow, and so may not be as certain as I would like to be. My specific concerns are as follows:\n\n1. Clarity is poor. Details about the FEM design are spread out and interleaved with justifications and theorems. I would find it much easier to follow if the \"final\" (C+L+T+G) function were introduced first & with enough details to implement it, followed by justification & the general framework into which it fits, in a separate section.\n   - For example, Equations 4, 6, 9 and 10 overlap considerably, but with subtle differences.\n   - It wasn't clear to me how $\\beta_{\\text{max}}$ is chosen/parameterised (this detail is in Appendix K, L1681).\n   - The link between $\\beta$, $\\beta_{\\text{hid}}$, $\\beta_{\\text{max}}$ and $\\lambda$ takes some time to understand in the presentation of Section 2.3.\n1. The proposed method is has a privileged basis in the value space. That is, regular attention mixing would be equivariant to rotations of $v$ (i.e. $\\text{mix}(p, v) \\equiv R^{-1} \\text{mix}(p, R v)$ for any rotation $R$), while the proposed method is not equivariant. This is not a problem, but highlights that channel-wise arguments such as \"expressing channel-argmax\" promote methods with this basis-dependence, something that is not obviously necessary, from my perspective.\n1. I am unconvinced by the phrase \"lossy readout\". The readout of a value matrix to a \"mixed\" pre-output vector is inherently compressive. FEM may make the selection function class richer in some ways, but since it shrinks the value dimension when parameter-matched, it loses expressivity elsewhere. I think the paper can be read as \"there are mixing functions that we think are interesting/useful, but cannot be represented with usual convex mixing, while they can with FEM\". So I find it hard to agree that FEM makes the mixing \"less lossy\" in any real sense.\n1. Baseline results in Table 2 don't match Yang et al. (2024), Table 3. E.g. DeltaNet and HGRN2 with 1.3B parameters / 100B tokens, on HellaSwag, ARC-{C,E}, BoolQ. Both use lm-evaluation-harness. Is there a known reason?\n\nMinor comments:\n\n- No discussion about whether/how this is compatible with grouped query attention (Ainslie et al., 2023). I might be concerned that query heads that share a value may become too similar, if $e^{v}$ is sparser/lower-temperature than $e^{q k^T / \\sqrt{D}}$.\n- Since much is made of the channel-argmax limitation, I expected a \"toy problem\" demonstrating the FEM's ability to solve this in a single layer, contrasted with convex mixing.\n- L111 says \"causality ... can be encoded\", but causality is already encoded in the equations above which run \"over past indices\".\n- L156 says \"$H L \\geq D$ (which is not practical)\", but I think $H L$ is commonly $> D$, e.g. Llama 3 8B has 32 blocks and 32 query heads per layer, with head dimension 128.\n- Typo L1635 \"Equation equation 35\"\n- Appendix J.1 would benefit from references, even if these are referenced elsewhere.\n\n---\n\n_Yang, S., Kautz, J. and Hatamizadeh, A., 2024. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464._\n\n_Ainslie, J., Lee-Thorp, J., De Jong, M., Zemlyanskiy, Y., Lebrón, F. and Sanghai, S., 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245._", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)