Distilling to Hybrid Attention Models via KL-Guided Layer Selection
Authors: Yanhong Li
Songlin Yang
Shawn Tan
Mayank Mishra
Rameswar Panda
Jiawei Zhou
Yoon Kim
Venue: ICLR 2026 Poster (2026)
Keywords: Linear attention
Hybrid architectures
Distillation
Layer selection
Inference efficiency
Primary area: foundation or frontier models, including LLMs

Abstract: Distilling pretrained softmax attention Transformers into more efficient hybrid architectures that interleave softmax and linear attention layers is a promising approach for improving the inference efficiency of LLMs  without requiring expensive pretraining from scratch. A critical factor in the conversion process is  layer selection, i.e., deciding on which layers to convert to linear attention variants. This paper describes a simple and efficient recipe for layer selection that uses layer importance scores derived from a small amount of training on generic text data.  Once the layers have been selected we use a  recent pipeline for the distillation process itself \citep[RADLADS;][]{goldstein2025radlads}, which consists of attention weight transfer, hidden state alignment, KL-based distribution matching, followed by a small amount of finetuning. We find that this approach is more effective than existing  approaches for layer selection, including heuristics that uniformly interleave linear attentions based on a fixed ratio, as well as more involved approaches that rely on specialized diagnostic datasets.

Published as a conference paper at ICLR 2026
DISTILLING TO HYBRID ATTENTION MODELS VIA KL-GUIDED LAYER SELECTION
Yanhong Li Songlin Yang" Shawn Tan
Mayank Mishra
Rameswar Panda Jiawei Zhou
Yoon Kim?
Allenn Institute for MMT MIT-IBM Watson Al Lab
*Stony Brook University
yangs166lmit..ed
ABSTRACT
Distilling pretrained softmax attention Transformers into more efficient hybrid architectures that interleave softmax and linear attention layers is promising approach for improving the inference efficiency of LLMs without requiring ex- pensive pretraining from scratch. A critical factor in the conversion process is layer selection, i.c., deciding on which layers to convert to linear attention vari- ants. This paper describes simple and efficient recipe for layer selection that uses layer importance scores derived from a small amount of training on generic text data. Once the layers have been selected tion process itself (RADLADS; Goldstein et al. 2025). )), which consists of attention weight transfer, hidden state alignment, Kl.-based distribution matching, followed by small amount of finetuning. We find that this approach is more effective than existing approaches for layer selection, including heuristics that uniformly inter- leave linear attentions based on a fixed ratio, as well as more involved approaches that rely on specialized diagnostic datasets
a recent pipeline for the distilla-
INTRODUCTION
Linear attention (Katharopoulos et al. 2020 Peng et al., 2021 Yang et al. 2023
i.a.) and state-
space models (Gu et [002 Gu & Dao 2024; Dao & Gu 2024 i.a.) have gained significant trac- tion recently due to their high inference speed and competitive performance. However, most existing pretrained models are still purely based on softmax attention, and pretraining such linear attention models from scratch is resource-intensiie This has motivated the approaches for crsssarcchtecctre distillation, a process that converts pretrained Transformer checkpoints into more efficient linear attention counterparts (Kasai al. 2021 Wang et al. 2024
Bick et al. 2025 i.a).
This distillation process involves two key deci- sions: (1) the student architecture, and (2) the optimal distillation recipe once the architecture has been selected. For the second question, recent work has shown the effectiveness of multi-stage pipeline over pure continued fine- tuning approaches (Bick et al. 22025] Goldstein jet al., (022). This pipeline involves an initial stage of per-layer output alignment with an L2 loss, followed by second stage of end-to-end knowledge distillation. What student architec- ture to distill to, however, remains open. Prior efforts to distill Transformers into purely sub- quadratic models have often resulted in perfor- mance degradation (Zhang al., 2024a b. Mer- cat et al. 2024). More recently, models in- corporating a sliding window attention (SWA)
Performance vs. Sliding Window Size
0.8
Performanc
0.6
0.4
0.2
26
32
64
120
256
512
1024
2048
Sliding Window Size
ATTN)
Common Sense
Common Sense (Full ATTN]
Figure 1: Performance of bidding-windww atten- tion model (distilled from Dwen2.553BB-Instmee across different window sizes on RULER and commonsense reasoning tasks.
mechanism have shown surprisingly strong results across various benchmarks (Lan et al. 2025
"Equal contribution. Work conducted while YL. was a visiting student at MIT.
Code available at atthttss:/// http::////tttttootootttttttttttt http://e///gghiaa..cccc//ooaaa//
on

---

Published as a conference paper at ICLR 2026
Zhang et al. 2023). However, these evaluations have primarily focused on knowledge-ing inttengg-iintensiee commonsense reasoning tasks, where in-context recall plays a lesser role. Indeed, Figure shows that even small sliding window of size 16 is sufficient for distilled SWA model to recover strong performance on such tasks. In contrast, performance on in-context recall benchmarks like RULER (Hsieh et al.) 2024) is highly dependent the sliding window size (Figure []. This is perhaps un- surprising, as it reflects the well-documented limitations of fixed-state models in in-context recall (Wen al. 2025 Arora al. 2024a b).
A simple yet effective solution is to incorporate a few global (softmax) attention layers, resulting in hybrid architecture. This approach has been successfully adopted in recent models pretrained from scratch, such as Jamba (Lenz et al., (225)), MiniMax-01 (MiniMax et al.) 2025). Falcon-HI (Zuo et al., 2025), and Qwen3-Next. These models typically interleave global and linear attention layers at fixed ratio (e.g., one global layer for every three or seven linear layers) (Wang et al.) 2025ap. Following this trend, some distillation works have also adopted fixed interleaving strategy (Wang et al. 2024). However, our preliminary experiments show this uniform approach remains suboptimal for in-context recall, presumably due to the fundamental difference between pretraining and distillation. This observation has been recognized in recent work (Gu al. 2025 Yang al., 2025 Hoshino et al. 2025), which also explore various criteria for selecting global attention layers.
In this work, we adopt simple global attention selection criterion based on the distillation KL diver- gence loss: intuitively, the more critical a global attention layer is, the more it reduces the resulting distillation KI. loss. Our experiments demonstrate the effectiveness of our selective hybrid distilla- tion, which achieves strong in-context retrieval performance while maintaining efficiency. Our work paves the path for future work lest-time compute scaling for distilled hybrid models (Paliotta 2025 Wang et al. 2025b), where in-context retrieval remains key bottleneck (Chaudhry 2025).
et al.
al.
PRELIMINARIES
Notation. Let i : xT] E Txx be sequence of token embeddings with model width d. We use L pre-norm Transformer blocks indexed by l E {l,..., L}. and h attention heads with per-head width d, so d dh. A Transformer block then given by
U ( Mix (IN(((())
( U FFN(*) (IN(U(9))
where M-x is sequence mixing operation (i.e., softmax linear attention) for layer f. When not essential, we omit LN and residuals for readability. We write M for the (additive) attention mask, which encodes causality and any positional encoding (e.g.. RoPE/Alibi) standard.
Softmax attention. For single head (we suppress head indices) softmax attention proceeds by computing the query, key and value matrices
Q = XW K= XW V XWv,
where Wo. WK, Wv Rd are learnable parameters. The output is given by (with mask M)
- Softmax
Vdn OK +M)
and multi-head concatenates per-head outputs which is transformed by a linear layer Wo E R (hds)xxd During autoregressive inference, the same operation admits recurrent view:
Qgii exp
A Val q. k:).
Σ iS!
Ci 1.
iSt
The memory cost of softmax attention grows linearly with respect to sequence length due to the KV cache, which can result in substantial slowdowns as generation length grows due to increasing data movement across the memory hierarchy.
Linear attention. Linear attention layers have been proposed to address the above inefficiencies of softmax attention during decoding. While many variants exist, they generally adopt the following recurrent form:
o, S.S., S. = =M,S k.v,,
(3)
2

---

Published as a conference paper at ICLR 2026
where M, is data-dependent and time-varying transition matrix that is function of x,- Setting M, = diag (+)) where a E R is a function of x, recovers recent gated linear attention (GLA)
variants Yang al. 2023 Katsch 2023 Qin al. 2024 Peng et al. 2024). Alternatively, using k,k, recovers the (gated) DeltaNet family of models (Schlag et al. 2021 Yang
M,
02 The structure of M, enables efficient parallel training via chunking mechanism.
al.
2024b
Linear attention compresses the entire history into the hidden state matrix S, and thus the memory cost is constant with respect to generation length, leading to much more efficient decoding compared to softmax attention. However, this hidden state bottleneck is a fundamental limitation when comes to crucial capabilities such as performing associative recall over a given context.
Hybrid attention. A common strategy for maintaining the capabilities of softmax attention while realizing some of the efficiency benefits of linear attention is to use a hybrid model. This approach partitions the of layer indices into Ssotmax and Stincar such that Ssoftmax U Sjincar (I....,4]) L. Then the mequence-mmxing layer is given by
SoftmaxAttn!", LinearAttn())
E Ssoftmax:
Mix 1
8 Stinear-
Recent works have shown that architectures that use a fixed ratio of linear to softmax attention layers performs well when pretrained from scratch (Lenz et al.) 2025 MiniMax et al. (025). However, such a uniform strategy may be suboptimal for distilling hybrid attention models from pretrained softmax attention models, motivating the present work on layer selection for distillation.
3 LAYER SELECTION FOR DISTILLING HYBRID ATTENTION
For distilling a pretrained softmax attention LLM into hybrid attention model, we seek to find set Con for a given budget [Conn] K such that converting all the other layers into linear attention has minimal performance degradation. Solving this exactly would require combinatorial search over all possible K -siized subsets of [L]. which would be intractable. Our key idea is to measure a layer's marginal utility by restoring exactly that layer (and only that layer) to softmax in otherwise all-linear student, then distilling briefly and scoring how much the teacher-student KI. improves.
3.1 INITIAL DISTILLATION TO AN ALL-LINEAR STUDENT
We first distill to an all-linear student model, adopting the first two stages of the distillation pipeline from RADLADS (Goldstein al. 2025). Let Mieacher be the original teacher model and Mall-linar be an all-linear student model, where the linear attention parameters are initialized from the teacher's parameters, i.c., (Wq. Q: K. Wv. Wo). The other parameters of the linear attention layer (in particular the parameters of linear layer for the data-dependent gating term a.) are initialized randomly. Then distillation proceeds as follows:
Stage 1: Hidden-state alignment. For given token sequence x = IT, the attention hidden states from the all-linear student model {U(() all-linear are trained match the teacher's hidden states
(U teacher [E([)]
Lrididen (Mall linear, ar) Σ EE(L]
( T teacher
1
U(O)
aall-tineur/
Here, RADLADS only trains the parameters of the student's linear attention layer while freezing FFN's parameters. The targets are produced by the teacher model and remain fixed.
Distribution matching. In stage 2, RADLADS minimizes a temperature-scaled KL be- tween teacher logits Lsencher,t E V and student logits Lall-linarrt IR with respect to all student parameters (i.e., including the student's FFN layers)
Stage 2:
LKL. CRI(M (Mall
(Mall-inner linear; x))
?
T
KT. KT. (Softmax
(aauthrr
Softmax
(aasincess )),
(5)
DeltaNet also multiplies the additive term k.v. with Be, which we omit for simplicity.
3

---

Published as a conference paper at ICLR 2026
where smoothing term that provides stronger gradient signal non-argmax tokens. (The func- tions Liideen and LKL are obviously functions of Mwacher but we omit them for readability.)
Stage uses 100M tokens while stage 2 uses 600M tokens. All subsequent applications of the stagewise pipeline (i.c., in $3. and $3.3 use the same number tokens
3.2 DERIVING LAYERWISE IMPORTANCE SCORES
With the all-linear model Mallainear derived from the above process, we now describe our layer selection strategy. Let M all-incar be model derived from Mall-linea where the f-th block has been restored back into the E-th layer of Mwacher- We run stage I and stage of the above process again
(-t)
to finetune the student M M which now has one softmax attention layer. We define I(E), the layer importance for layer E, as the KL divergence between and the teacher model, i.c.,
I(e) E (M
(-t) (M all-linear x>)
(6)
Higher Z(6) means larger KL reduction (i.e., greater marginal utility under our objective). Because the baseline student and neighbors fixed, I(E) is hybrid-aware and variant-aware.
3.3 LAYER SELECTION AND FINAL DISTILLATION
Algorithm KL-guided Layer Selection for Hybrid Attention Distillation
Require: Teacher Mteacheri dataset D (DCLM); temperature т; target budget K
1: Distill into pure linear attention model Mall-linea ($3.1)
2: for l = 1 to L in parallel do ($3.2)
3: Obtain M (-f) by changing č-th layer of Mall-liner to E-th layer of Mreacher
4: Stage 1: align all linear blocks by Lhia on D.
5: Stage 2: distill by LKL on D.
6: Compute I(t) [[cu
held-out slice of D.
7: end for
8: Select: Sanftmax +- top-K layers by I(() ($33)
9: Final hybrid: instantiate hybrid based on Ssofimux and linear on layers [Z]] \ Swotmaxi train with
the two-stage distillation pipeline.
Given budget of K softmax attention layers that we can keep, we now take the top- K most impor- tant layers and convert the result into linear attention i.e.,
Syatmax top-K(I(()) Sincar {l..... L} Ssotmmax
Denoting the above hybrid model with K softmax attention layers Maybria-K we run final distillation pipeline by rerunning stages and 2 with this hybrid model. Our full algorithm is given in Algorithm
4 EXPERIMENTS
Having introduced our method, we now present a series of experiments designed to build a com- prehensive case effectiveness. We begin by establishing why hybrid models are essential for maintaining long-context capabilities ($4.) 44.1 We then demonstrate that our KL-guided approach
outperforms a wide range baselines
4.3
4.1
THE CASE FOR HYBRID MODELS
There has been flurry of recent work on distilling to pure linear attention models (Chen et al.,
2024
Mercat al.) 2024 Zhang et al.) 2025 Goldstein et al. 2025 Wang et al.] 2024 Yueyu
ct al.
2025
Lan et al. 2025 Bick et al. 2025). These works generally report that pure linear
For our main GA-S2 selector, the final hyhrid model reuses the Stage 1-aligned linear attention layers from Mall linea and therefore only Stage in the last distillation step.
4

---

Published as a conference paper at ICLR 2026
Qwen22553B-Instrucc Parformance
Accuracy
0.8
RULER (Student)
RULER (ibacher)
(Student)
Sense (Teacher)
0.40
15
35
Number Full Altention Layers
Llama-3.2 2-38--nstruct Performance
0.8
Accuracy
RULER (Student)
RULER (leacher)
(Student)
(Teacher)
25
Number Full Altentinn Layers
Figure 2: Performance on recall-intensive vs. commonsense tasks as the number of full-attention layers is varied for Qwen2.5-3B 1amm2..-3B-Instruue Instructtt (top) and 1llmm--.22-3B 3.2-3B-Instruct (bottom). Recall ability is highly sensitive to the softmax budget, while commonsense reasoning is not.
attention can maintain the performance of pretrained softmax attention baselines with the right dis- tillation process. However, this conclusion often based on comparing performance on tasks such as MMLU and Commonsense Reasoning, whose context lengths are short; it is unclear the extent to which such pure linear attention models can maintain performance on benchmarks which require understanding and performing recall over longer contexts. To analyze this, we construct a series of hybrid models based on our approach where the number softmax layers ranges from to L 1. We then evaluate these models RULER (Hsich al., (002)), diagnostic benchmark designed to probe the long-context capabilities of LLMs. We also evaluate these models on short-cone chott-context commonsense reasoning benchmarks evaluated by previous methods, including PIQA, ARC-Easy, ARC-Challenge, HellaSwag and WinoGrande (we report the average).
The results in Figure 2 reveal stark dichotomy. Performance the long-context RULER bench- mark is highly sensitive to the number of softmax layers (K), growing monotonically and confirming that global context aggregation is critical for in-context retrieval. In contrast, commonsense reason- ing performance is almost entirely insensitive to K; models with even single softmax layer achieve near-teacher--evel performance, suggesting these local tasks are well-handled by linear attention. Ironically, the efficiency benefits of linear attention are minimal on precisely these short-context tasks. This dichotomy motivates our work: the central challenge in distilling hybrid models to preserve cong-context recall. This requires method that can judiciously allocate a limited budget of expensive softmax layers to the positions where they are most impactful.
4.2 EXPERIMENTAL SETUP
Having established the importance of selection, we now evaluate our KL-guided method against the a suite of baselines.
Model and data. We evaluate two 3B-class decoder-only teachers: Qwen2.5-3BBInntruue and Alama-3.2-3B-Instruee
For each architecture we take the checkpoint's native depth L and re- port to match target softmax:
ratio. We target four ratios 1:8, 1:3, 1:2, (thus K K (4,9, 12, 18) when L-366: if L differs, we use the nearest integer K). All selection and distil- lation runs use the DCLM (Li et al. 2025) generic-text mixture. As noted 3.1 each instance of stage uses 100M tokens while stage uses 600M tokens.
Baselines. We compare one-swap selector to the baselines below. Each returns set of K soft- max layers and is trained with the same two-stage distillation and token budget as ours ($3.I): (1) Uniform interleave (UNIFORM). Pick K layers by evenly spacing them across depth (one roughly
5

---

Published as a conference paper at ICLR 2026
1aan2.5-3B-Instree
1a2ma-3.2-3B-Immse
Porformance (RU1, 1.ERR
Porformance (RULER)
0.8
0.0
0.6
0.6
0.4
0.4
Linemer-Linear Ratio (%)
Softmax:Line Linear Ratio (%)
1.0
1.0
Performance (SWDE)
Purformance (SWDE)
0.4
0.4
0.2
0.2
0.0
50
0.0
10
20
30
40
50
Softmax:Lineat (%)
Softmax: Lanear Ratio (%)
---
Uniform
AR-MH
LM-PPL
---
Ours
Teacher
AR
AcI-MSE
SMART
X
PostNAS
All-Linear
Figure 3: Performance comparison of various layer selection methods on RULER (top) and SWDE (left) and Llama-3.2-3B-Instee (right) into hybrid
(bottom) for distilling Qwen2.5-33-Insttne GDN-based models. Performance is plotted against the percentage of softmax layers retained. The dashed line indicates the performance of the all-softmax teacher model.
every LL/K] blocks), as adopted by Wang et al. (0224). (2) Task-guided selectors. AR (Associa- tive Recall): bypass each layer and measure the drop a synthetic key-value recall task and then rank layer importance by drop in performance (Chaudhry et al., (025). AR-MH (Associative Recall Multihop); same AR but with multi-hop alias chains, which makes the task more diffi- cult. (3) Model-signal selectors. ACT-MSE: layer importance is derived from zero-ing out a layer and measuring increase in activation MSE vs. the baseline. LM-PPL: same Act-MSE, but derived from measuring an increase in LM perplexity on held-out data. (4) SMART (Yang et al.) 2023). A sensitivity aware strategy: (i) score each layer by the reduction in seacher-student KL. when swapping an global layer into an otherwise linear baseline; (ii) preserve high-score layers near input/output (so-called "terminal preservation"), (iii) choose the rest from near-uniform candidates to maximize total sensitivity. We also compare against PostNAS (Gu et al.) 2025). contemporane- ous work that uses a more complex search procedure. Their method involves training a once-for-all SuperNet and then using beam search to find the optimal K softmax layers for specific down- stream task. This process computationally intensive, requiring 50B training tokens, whereas our selection pipeline uses only 5-6B tokens. Fortunately, PostNAS released their selected layers for the Qwen2.5 model. To ensure a fair comparison, we take their publicly released layer set and distill using our own pipeline and token budget. More baselines descriptions are included in Table5in the Appendix A
4.3
MAIN RESULTS
We use gated DeltaNet (GDN) for our linear attention layer and evaluate our proposed layer lection method against the baselines for Qwen2.5-3B-Instre and Llama-3.2-3BB- -nnssruutttt teachers. The results two long-context, recall-intensive benchmarks, RULER and SWDE, are presented in Figure| Our central finding is that our selection method consistently and substantially out- performs all other baselines across both models and tasks. This demonstrates the effectiveness of using a brief, KL--ivergencc--uuide distillation to derive model-intrings layer importance scores for creating hybrid architectures.

---

Published as a conference paper at ICLR 2026
25% softmax
33% softmax
Teacher
Teacher Performance
SMART
GA-S2 (Ours)
SMART
GA-S2 (Ours)
Qwen2.5-1.5B Qwen2.5-7B
0.8742 0.9445
0.5098 0,8158
0.5408 0.8584
0.6479 0.8949
0.6953
0.9110
Table I: Experiments on different model sizes on RULER at fixed hybrid budgets.
Stage (MSE-based)
Stage (KL-baaed)
Model
GR-SI
GA-SI
AVG-S1
GR-S2
GA-S2 (Ours)
AVG-S2
t-attttrettt 0,4508
0,4193
0.4233
0,4950
0.7539
0,5580
Qwen2.5-3B-Instee
0,4827
0.5408
0.4933
0.8259
0.8631
0.8205
Table 2: Ablation on layer selection strategies for fixed 25% softmax ratio. We compare Greedy Addition (GA), Greedy Removal (GR), and Averaged (AVG) search using either Stage-1 (MSE) or Stage-2 (KL) importance metric.
A key advantage of our approach is particularly evident in the low-budget regime, where only small fraction of layers are kept as full softmax attention. For instance, on RULER with Qwen2.5 12.5% softmax budget (5 softmax layers), GA-S2 reaches 0.662, outperforming the strongest baseline (AR at 0.542) by +0.12 and the standard UNIFORM interleaving strategy (0.441) by +0.22. This pronounced gap at low softmax ratios highlights our method's efficiency in identifying the most critical layers for preserving long-context recall, enabling significant performance gains with minimal computational overhead from expensive attention layers.
As the budget for softmax layers increases, our method continues to maintain a performance advan- tage, approaching the teacher model's performance more rapidly than competing approaches. For both models, hybrid with 50% of its layers selected by our method recovers a vast majority of the teacher's performance on these challenging recall tasks. Similar performance trends were observed on other benchmarks, including FDA and SQuADv2; these results are detailed in the Appendix A
To test whether the gains of KL-guided selection persist beyond the 3B setting, we distill hy- brid students from 1wvn2..511.5B. Instruct and Qwen2.5-7B-Innstrue
25% and 33% softmax budgets (same distillation recipe and data). Table [[] shows GA-S2 consistently outperforms the strongest baseline (SMART) across both scales, with improvements of +0.031/+0.047 (1.5B) and +0.043/+0.016 (7B) at 25% Full results (including all baselines) are showed in Section
Inn.-ruuttt
5 ANALYSIS
In this section, we conduct series of ablation studies to deconstruct our method ($5.I) understand its architectural sensitivities (5.2), and validate its practical efficiency ($ 5533
5.1 THE IMPORTANCE OF KL AND GREEDY ADDITION STRATEGY
Our proposed layer selection method involves two key design choices: (1) we use the stage-2 (S2) knowledge distillation (KL-based) loss as the importance metric for each layer in the one-swap setting of $3.2 and (2) given these layerwise scores, we select the top- softmax layers in a greedy addition fashion (GA), i.e., we keep the K layers that yield the largest marginal KL reduction relative to the all- linear baseline. There are natural alternatives: we could use the stage-1 (SI) hidden-state alignment (MSE-based) metric as our layer importance; we could also use a greedy removal (GR) search strategy, which starts from an all-softmax model and greedily converts the least important layer to a linear attention layer. both GA and GR (AVG). Note that our main proposed method corresponds to GA-S2.
also possible to average the layer importance rankings from
The ablation results, presented Table 2 show that the Stage-2 (KL-based) methods consistently and dramatically outperform their Stage-T (MSE-based) counterparts, and our greedy addition strat- egy (GA-S2) is more effective than greedy removal (GR-S2). This suggests that identifying the single most impactful layer to add from an all-linear base is more robust signal than identifying the least harmful layer to remove. Full layer-wise importance rankings for all selectors are provided in Appendix C
7

---

Published as a conference paper at ICLR 2026
Llama-3.2-3B
Qwen2.5-3B
Ratio
GDN
GLA
GDN
GLA
12.5%
0.5731
0.5166
0.6617
0.6074
25%
0.7539
0.6498
0.8631
0.6921
33%
0.8118
0.7967
0.8733
0.878
50%
0.8619
0.8487
0.9061
0,9069
Table 3: Final RULER performance using architecture-speciie selections
Agrecment of Top Layer Selections Over Time 3.2-33-3.2-3B)
Agreement of Top Layer Selections Over Timee (Qwen2.5-38)
1.0
Jacued Stailatte
0.4
lscred
0.2
64
3200
2033
4302
3300
2003
20000
6303
Solooton St suning Stop
Selectona Sranning Stop
(a) Llama-3.2-3B (Mean Similarity: 0.65)
(b) Qwen2.5-33 (Mean Similarity: 0.54)
Figure 4: Jaccard similarity of top-K layer selections between GDN and GLA variants the selection pass. Llama shows higher agreement, suggesting its layer importance is less student- dependent.
5.2 THE IMPORTANCE ARCHITECTURE CONSISTENCY
Our layer selection approach is sensitive to the type of linear attention layer employed. To what extent is this selection approach wwchitecture-gnnosiic- i.e., is our method simply finding a fixed set of "important layers" in the teacher, is it adapting its selection to the specific architecture of the student's s linear layers? To test this, we run the selection process independently for both GDN and GLA students and analyze the results.
The results in Figure 4 and Table reveal an interesting architectural dependence. At fixed 25% softmax budget (K -9,, the GDN and GLA selection trajectories exhibit only moderate overlap: the mean Jaccard similarity is 0.65 for Llama-3.233--Instrue and 0.54 for Qwen2.5-3B-Insttreee which corresponds to roughly ~7/9 and ~6/9 layers overlapping on average (Figure 4 Thus, the two student variants typically disagree only 1-3 layers, yet these small differences can have an outsized impact on long-context recall. Concretely, in this low-budget regime, the architecture- specific GDN-GDN models substantially outperform the architecture-specine GI GLA-GI models on RULER: 0.7539 vs. 0.6498 for Illama and 0.8631 vs. 0.6921 for Qwen (Table
Most surprisingly, when we test transferability by using the GDN-selected layers to distill a GLA student, we achieve strong RULER performance (0.6927 for Llama and 0.8407 for Qwen; Table J This result is not only far better than all baselines, but is also significantly better than the score from the specialized GLA-GLA process (0.6498 for Llama and 0.6921 for Qwen). This reveals a key finding: the choice of linear attention variant used during the selection pass acts as a "probe", and some probes are better than others at identifying a robust set of important layers for a given teacher architecture. In particular, using GDN the probe yields layer set that transfers well to both GDN and GL. students in the low-budget regime. This demonstrates that our method's strength is not just in specialization, but its ability leverage different student architectures identify the most impactful softmax layers for preserving long-context recall.
5.3 How MANY TOKENS ARE REALLY NECESSARY FOR LAYER SELECTION?
We used 100M tokens for stage I and 600M tokens for stage 2 following the recipe recommended in Goldstein et al. (0025). However, it is possible that the layer selection process could even more token-efficient To investigate this, tracked the top- K layer set chosen by our selector throughout the Stage-2 training process (at a 1:3 softmax ratio for both models). We measured stability over time using bolling-windore Jaccard similarity and the size of the intersection between consecutive
8

---

Published as a conference paper at ICLR 2026
Model
Student
UNIFORM
AR
AR-MH
MSE
PPL
SMART
Ours
GDN
0.461
0.59
0.5512
0.4899
0.5011
0.6274
0.7539
Llama
GLA
0.4342
0.5404
0.4838
0.4436
0.4266
0.5851
0.6927
GDN
0.6904
0.6385
0.716
0,4421
0.51
0.6401
0.8631
Qwen
GLA
0.633
0.574
0.6628
0.3961
0.4225
0.6007
0.8407
Table 4: Performance RULER for GDN- and GLA-based hybrid students at fixed 25% softmax ratio. For both student variants, the layer set for our method (Ours) was selected using GDN-based
process to test for transferability. Note that Llama refers to Llama-3.2 I-3Butnsttuutt and Qwen refers to Brren..5-3B-Instruee
Pafforennee
0.66
M
0.84
0.82
0.80
Step
Figure 5: The evolution of RULER performance during the Stage-2 selection process for Qwen2.5- 3B-Instruct.
sets (the "backbone"). For both teacher models, we find that the set of selected layers stabilizes long before the full training budget is consumed. A nearly complete "backbone" of K - 1 layers is typically identified within the first 25-40% of training. Continuing training beyond this point only refines the choice for the final
two slots, with negligible impact the final model's RULER performance (a difference of less than 0.01 absolute points). This observation suggests that simple stability-based rule can dramatically improve efficiency. For instance, conservative early stopping point for our runs would have reduced the token budget for the selection pass by 58-74%. The effectiveness of this early stopping rule is backed by our empirical observation: for Qwen, the RULER performance during Stage-2 stabilizes around step 1500, as shown in Figure 5 For more details, please refer Appendix B
6 RELATED WORK
In-context recall presents a significant challenge for subquadratic models, a difficulty often attributed to the perplexity gap between them and standard transformers (Arora al.) [2024a). One promising approach to address this is the development of linear attention variants with superior recall capa- bilities. The seminal work on DeltaNet (Schlag et al.) 202T Yang et al. 2024b) and its successors Yang et al.) 2024a Siems et al., 2025 Grazzi et al. 2025) has demonstrated great success in this area. Nevertheless, these recurrent approaches are fundamentally limited associative recall by their fixed-size state (Wen et al. 2025 Arora et al. 2024a). Highlighting the importance of this problem, recent work reveals a connection between in-context recall and test-time scaling perfor- mance, arguably making one of the most critical research directions in efficient sequence model design (Chaudhry et al., 2025). Other notable efforts to improve recall include reading inputs twice (Arora et al.) 2024c), dynamic state allocation (Ben-Kish al. 2025), and dynamic caching for hard-to-meman meoomizee items (Nguyen et al.) 2025).
Hybrid attention architectures, which combine the complementary strengths of global attention (for accurate retrieval) and linear attention (for fast local processing), can theoretically overcome these state-size limitations (Wen et al. 2025 Arora al.] 2024b). While most hybrid models adopt an inter-layer strategy, interleaving global and linear attention layers (Ren et al., 2025) MiniMax al.) 2025 Lenz et al. (025). we also note the potential intra-layer hybridization schemes for efficient time mixing (Irie et al. 2025 Dong et al. 2024 Zuo et al.) 2025 Zancato al. 2024). However, pretraining these linear and hybrid models from scratch is computationally expensive. An effective alternative is to distill pretrained softmax attention model into linear attention- sreettio--bsedd one. This concept was first proposed by Kasai et al. (202T). Subsequent work has emphasized preserving or mimicking the softmax operator during distillation maintain performance while achieving linear

---

Published as a conference paper at ICLR 2026
complexity Peng al. (2022); Zhang et al. (2024b a). Research work shows that sliding window attention with window size 64 works well in many benchmarks Lan et al. (2025): Zhang et al. (2025), though we show in this work that such strategies still perform poorly in-context recall.
In the context of distilling into a hybrid of global and linear attention, key question has emerged: how to select which global attention patterns preserve. Some methods rely on downstream bench- mark performance to determine importance Gu et al. (2023), while others use speculative decoding diagnostic tool to identify redundant attention layers Hoshino et al. 2025). In contrast, our work focuses on a simple strategy using an unsupervised learning loss and provides extensive analysis that goes beyond prior research Yangg al. 2025).
7 CONCLUSION
In this work, we introduced a simple and effective method for selecting which softmax attention lay- ers to retain when distilling pretrained Transformer into more efficient hybrid architecture. While our selection process is more efficient than complex bearch-based alternatives, future work could explore even cheaper proxies for layer importance, potentially derived directly from the teacher model's s activations gradients. Other promising directions include extending this selection frame- work from the layer level to a more fine-grained, head-level hybridization.
ACKNOWLEDGMENTS
This work was supported by National Science Foundation under CAREER Award No. 2441872, MIT-IBM Watson Al Lab, and gift from Jane Street.
STATEMENT ON LLM USAGE
We acknowledge the use of Large Language Models (LLMs) to assist in the preparation of this manuscript. Specifically, LL.Ms were utilized to improve grammar and clarity, aid in literature discovery, and generate boilerplate code snippets for our experiments and testing scripts. The authors have carefully reviewed and edited all LLM-generated outputs and take full responsibility for the final content and scientific integrity of this work.
REFERENCES
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations, 2024a. URL https: /ttpss///eeeeeeeeeeeeeeee net Fottt//orrrmimd
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall-throughput tradeoff. In Forty-first International Conference on Machine Learning, 2024b. URL https: //openreview. wwt.//fooommeee. hett/forrmmm...ee33fffff.e
Simran Arora, Aman Timalsina, Aaryan Singhal, Benjamin Spector, Sabri Eyuboglu, Xinyi Zhao, Ashish Rao, Atri Rudra, and Christopher Ré. Just read twice: closing the recall gap for recurrent language models, 2024c.
Assaf Kish,, Itamar Zimerman, Muhammad Jehanzeb Mirza, Lior Wolf, James R. Glass, Leonid Karlinsky, and Raja Giryes. Overflow prevention enhances long-context recurrent LLMs. In Sec- ond Conference on Language Modeling, 2025. URL.https://ppeeeevvies wwwp://wwwieeeeeeeeeeeeeemm
id=h99hJ10997
Aviv Bick, Tobias Katsch, Nimit Sohoni, Arjun Desai, and Albert Gu. Llamba: Scaling distilled recurrent models for efficient language processing, 2025. URL httpt: //tpii//avivv.i/.a org. abs/
2502. 14458
10

---

Published as a conference paper at ICLR 2026
Hamza Tahir Chaudhry, Mohit Kulkarni, and Cengiz Pehlevan. Test-time scaling meets associative memory: Challenges in subquadratic models. In New Frontiers in Associative Memories, 2025.
URL httpss:///penreeeeeett wwt.eeewwwwwwweeeeeeeeeeeeeeeeee http://www..eeeeeemeeeeeeeeeeeee D0tt//ooooomdddd02222bo0000
Hanting Chen, Zhicheng Liu, Xutao Wang, Yuchuan Tian, and Yunhe Wang. Dijiang: Efficient large language models through compact kernelization. arXiv preprint arXiv:2403. 19928, 2024.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint arXiv:2405.21066 2024.
Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, and Pavlo Molchanov. Hymba: A hybrid-head architecture for small language models, 2024. URL https::/aaxiv.oooo abs/ 2411.44111 13676
Daniel Goldstein, Eric Alcaide, Janna Lu, and Fugene Cheah. Radlads: Rapid attention distillation to linear attention decoders at scale. arXiv preprint arXiv:2500.030005 2025.
Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, and Massimiliano Pon- til. Unlocking state-tracking in linear RNNs through negative eigenvalues. The Thirteenth International Conference on Learning Representations, 2025.
Albert Gu and Tri Dao. Mamba: Linear-tima sequence modeling with selective state spaces. In Proceedings of CoLM, 2024.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In Proceedings of ICLR, 2022.
Yuxian Gu, Qinghao Hu, Shang Yang, Haocheng Xi, Junyu Chen, Song Han, and Han Cai. Jet- nemotron: Efficient language model with post neural architecture search, 2025. URL https:
//arxiv.orq/abtt wtwpr///abs//5a/ 15884
Yuichiro Hoshino, Hideyuki Tachibana, Muneyoshi Inahara, and Hiroto Takegawa. Rad: Redundancy-aware distillation for hybrid models via self-speculative decoding. 2025. URL
https::/arxiz: org/ /25. 22135
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654 2024.
Kazuki Irie, Morris Yau, and Samuel J. Gershman. Blending complegeettmeeeeettrnenyoeo memory systems in hy- brid quadratic-linear transformers, 2025, UR/ptppss//ers....orr/aaaa/ 0074
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into RNNs. In Marie- Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pp. 10630-10643, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 00.18653/v1/022 wwt.222e..mmnlpmmmmm wwwnlp-miinnn mapp..mmiin.3000 URL attps://aclanthol// org/ 2021 www.eennnmmal...m min... 8307
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In Proceedings of ICML, 2020.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01927, 2023.
Disen Lan, Weigao Sun, Jiaxi Hu, Jusen Du, and Yu Cheng. Liger: Linearizing large language models to gated recurrent structures. arXiv preprint arXiv:2503.01496, 2025.
11

---

Published as a conference paper at ICLR 2026
Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman. Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhl- gay, Dor Zimberg, Edden M. Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedi- gos, Jhonathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zus- man, Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, Oded Fried, Ohad Leshno, Omer Antverg, Omri Abend, Or Dagan, Orit Cohavi, Raz Alon, Ro'i Belson, Roi Cohen, Rom Gilad, Roman Glozman, Shahar Lev, Shai Shalev-Shward Shaked Haim Meirom, Tal Delbari, Tal Ness, Tomer Asida, Tom Ben Gal, Tom Braude, Uriya Pumerantz, Josh Cohen, Yonatan Belinkov, Yuval Globerson, Yuval Peleg Levy, and Yoav Shoham. Jamba: Hybrid transformer-mambe language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https: http:://ppeeeeeeeeeeemmom net peppaarrerreee.net//ooooroooo-pp /ooommmdddd-FPPa PBD
Jeffrey Li, Alex Fang, Georgios Smymis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Rein- hard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Al- balak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng-Yu Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang. Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel II- harco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic. Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Se- woong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kol- lar. Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: In scarch of the next generation of training sets for language models, 2025. URL https: www.iv.vvvvv..///vv//aaa/a.v.vvv o
halgggaabss. abs 240. 11794
Jean Mercat, Igor Vasiljevic, Sedrick Scott Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. In First Conference on Language Modeling, net/ http//orummmmmm wtt.foorrmmm..socccccooooo.
2024. URL /ttppttts////pneeettt wwtp:/wwwwwww.eeeeeeeeeeeeeeeeee
MiniMax, Aonian Li, Bangwei Gong, Bo Yang. Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, Enwei Jiao, Gengxin Li, Guojun Zhang, Haohai Sun, Houze Dong, Jiadai Zhu, Jiaqi Zhuang, Jiayuan Song, Jin Zhu, Jingtao Han, Jingyang Li, Junbin Xie, Junhao Xu, Junjie Yan, Kaishun Zhang, Kecheng Xiao, Kexi Kang, Le Han, Leyang Wang. Lian- fei Yu, Liheng Feng, Lin Zheng, Linbo Chai, Long Xing, Meizhi Ju, Mingyuan Chi, Mozhi Zhang, Peikai Huang, Pengcheng Niu, Pengfei Li, Pengyu Zhao, Qi Yang, Qidi Xu, Qiexiang Wang, Qin Wang, Qiuhui Li, Ruitao Leng, Shengmin Shi, Shuqi Yu. Sichen Li, Songquan Zhu, Tao Huang, Tianrun Liang. Weigao Sun, Weixuan Sun, Weiyu Cheng, Wenkai Li, Xiangjun Song, Xiao Su, Xiaodong Han, Xinjie Zhang, Xinzhu Hou, Xu Min, Xun Zou, Xuyang Shen, Yan Gong, Yingjie Zhu, Yipeng Zhou, Yiran Zhong, Yongyi Hu, Yuanxiang Fan, Yue Yu, Yufeng Yang, Yuhao Li, Yunan Huang, Yunji Li, Yunpeng Huang, Yunzhi Xu, Yuxin Mao, Zehan Li, Zekang Li, Zewei Tao. Zewen Ying. Zhaoyang Cong, Zhen Qin, Zhenhua Fan, Zhihang Yu, Zhuo Jiang, and Zijia Wu. Minimax-0l. Scaling foundation models with lightning attention, 2025. URL https: //arxiv.v.vv////////a/eaa/aawav org/
abs abws//51s11. 2t///551..00 08313
Chien Van Nguyen, Ruiyi Zhang, Hanich Deilamsalchy, Puneet Mathur, Viet Dac Lai, Haoliang Wang, Jayakumar Subramanian, Ryan A. Rossi, Trung Bui, Nikos Vlassis, Franck Dernoncourt, and Thien Huu Nguyen. Lizard: An efficient linearization framework for large language models, 2025. URL httttt:tt:////o //ariiv./axi....///a www.vvoorg//baae abzg/abs/220 2250 09025
Daniele Paliotta, Junxiong Wang. Matteo Pagliardini, Kevin Y. Li, Aviv Bick, J. Zico Kolter, Albert Gu, François Fleuret, and Tri Dao. Thinking slow, fast: Scaling inference compute with distilled reasoners, 2025. URL httttttt////xiiiii
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint
3, 2024.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. In Proceedings of ICLR, 2021.
12

---

Published as a conference paper at ICLR 2026
Hao Peng. Jungo Kasai, Nikolaos Pappas, Dani Yogatama, Zhaofeng Wu, Lingpeng Kong, Roy Schwartz, and Noah A. Smith. ABC: Attention with bounded memory control. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Anmual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7469-7483, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 10.18653/v1/2022 acl-long.515. URL htttps::/aclaattollggt.to
org/2022 www.::22...c1.oo0o a.c. acn-long..............oooo 515/
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong- HGRN2: Gated Linear RNNs with State Expansion. In Proceedings of CoLM, 2024.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling, 2025. URL 07522
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers. In Proceedings of ICML, 2021.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state tracking in linear mns via householder products, 2025.
URL httttss://aaxii.rrrrrrrtrs attrrg/abs/502..... 10297
Dustin Wang, Rui-Jic Zhu, Steven Abreu, Yong Shan, Taylor Kergan, Yuqi Pan, Yuhong Chou, Zheng Li, Ge Zhang, Wenhao Huang, and Jason Eshraghian. A systematic analysis of hybrid org/ abs / 2507 06457
linear attention, 2025a. URL httphhttts://atee
Junxiong Wang. Daniele Paliotta, Avner May. Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. In The ehirth- Annual Conference on Neural Information Processing Systems, 2024. Id=uAzhOD JALU
www.:ew..eeeeeeeeeeeeeeeeeeeeeee forum?
Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M. Rush, and Tri Dao. M1: Towards scalable time compute with mamba reasoning models, 2025b. URL https:
//axsiv.orc//as//250444 10449
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. RNNs are not transformers (yet): The key bottleneck on in-context retrieval. In The Thirteenth International Conference on Learning Representations, 2025. URL htt.:htttpss/omttte htwp://wwwwww..aeeeeee.eeeeeeeee wtt.:/frrmmmiico dww Buck
Mingyu Yang, Mehdi Rezagholizadeh, Guihong Li, Vikram Appia, and Emad Barsoum. Zebra- Ilam:: Towards extremely efficient hybrid models, 2025. URL http::: /ttp:////axvvv.... www. abs/ 17272
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with harfware-efficient training. arXiv preprint arXiv:2312.06635, 2023.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arX1v:2412.06464, 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. In Proceedings of NeurIPS, 2024b.
Lin Yueyu, Li Zhiyuan, Peter Yue, and Liu Xiao. Arwkv: Pretrain is not what we need, attention-based language model born from transformer. arXiv preprint arXiv:2501. 2025.
Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, and Stefano Soatto. B mojo: Hybrid state space realizations of foundation models with eidetic and fading memory, 2024. URL
abs/ / 2407 06324
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The Hedgehog & the Por- cupine: Expressive Linear Attentions with Softmax Mimiery, 2024a. eprint: 0rX4v.2402.04347.
13

---

Published as a conference paper at ICLR 2026
Michael Zhang, Simran Arora, Rahul Chalamala, Benjamin Frederick Spector, Alan Wu, Krithik Ramesh, Aaryan Singhal, and Christopher Re. LoLCATs: On low-rank linearizing of large lan- guage models. In The Thirteenth International Conference Learning Representations, 2025.
htt.httts///pppneeeee
net /f forumoormmm/// Bit
Yu Zhang, Songlin Yang. Rui-Jie Zhu, Yue Zhang. Leyang Cui, Yiqiao Wang, Bolun Wang. Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, and Guohong Fu.
Gated slot attention for efficient linear-time sequence modeling. In NeurIPS,
/ 2024b. URL htw.aa/ch.........SSSS2oo---o-n. html]
G3139e51F5f634FF
Jingwei Zuo, Maksim Velikanov, Ilyas Chahed, Younes Belkada, Dhia Eddine Rhayem, Guillaume Kunsch, Hakim Hacid, Hamza Yous, Brahim Farhat, Ibrahim Khadraoui, Mugariya Farooq, Giu- lia Campesan, Ruxandra Cojocaru, Yasser Djilali, Shi Hu, Iheb Chaabane, Puneesh Khanna, Mohamed EI Amine Seddik, Ngoc Dung Huynh, Phuc Le Khac, Leen AlQadi, Billel Moked- dem, Mohamed Chami, Abdalgader Abubaker, Mikhail Lubinets, Kacper Piskorski, and Slim Frikha. Falcon-hl: A family of hybrid-head language models redefining efficiency and perfor- mance, 2025. URL https::/araae
org/ abs / 2200 22448
14

---

Published as a conference paper at ICLR 2026
A
COMPLETE RESULTS ON RECALL-INTENSIVE BENCHMARKS
Tag
Selector
Signal One-Line Procedure
UNIFORM
Uniform Interleave
Selects layers by evenly interleaving softmax layers at the target ratio.
Task-Guided Search (Heuristic-Rasee
KV
KV Retrieval
Importance from performance drop on a synthetic key-value dic- tionary lookup task when layer bypassed.
AR
Associative Recall
Importance from performance drop on task to sum the values of prompted keys when layer is bypassed.
AR-MH
Assoc. Recall- Mullil--ulti-hoo
As above, but with alias chains requiring multi-hop reasoning; performance drop defines importance.
VT
Variable Tracking
Importance from exact-set accuracy drop on a pointer-chasing task over shuffled assignments.
CWE
Common Words Extraction
Importance from set-match accuracy drop task identify the most frequent words long text.
ACT-MSE
Activation MSE
Mean-squared error on generic between final hidden states of a haseline vs. layer-bypassed model.
I.-PPI.
LM Perplexity
Measures the increase in perplexity held-out corpus when a layer bypassed.
Greedy Structural Search (Learning-Basee
GR-S1
Greedy Removal (S1)
Starts with all softmax; greedily converts the layer to lincar that hurts performance least after brief Stage-I I adaptation.
GR-S2
Greedy Removal (S2)
above, but using brief Stage-2 knowledge distillation for adaptation at each step.
GA-S1
Greedy Addition (S1)
Starts with all linear; greedily converts the layer to softmax that helps performance most after brief Stage-1 adaptation. As above, but using brief Stage-2 knowledge distillation for adaptation each step.
GA-S2
Greedy Addition (S2)
AVG-S1
Rank-Avg Greedy (S1)
Averages the layer importance rankings from GR S1 and GA S1 before selecting the top-K layers. Averages the layer importance rankings from GR-S2 and GA-S2 before selecting the top-K layers.
AvG-S2
Rank-Avg Greedy (S2)
Table 5: Layer-selection baselines and the tags used in figures. Layer bypass means applying an identity residual connection across the block's mixing sublayer.
15

---

Published as a conference paper at ICLR 2026
Llama-332-3B-Instee
Owee22.5-B-Insttreee
Selector
12.5%
25%
33%
50%
12.5%
25%
33%
50%
Heuristic-Rased
UNIFORM
0,4586
0.4610
0.5717
0.7262
0.4412
0.6904
0.6464
0.8031
KV
0.2029
0.6051
0.6626
0.7538
0.2543
0.7539
0.7552
0,8257
AR
0,3684
0.5900
0.6707
0.8347
0.5417
0.6385
0.6749
0.8717
VT
0.1839
0.2012
0.4334
0.7538
0.2922
0.4780
0.5359
0.7409
CWE
0.3129
0.3579
0.6752
0.8394
0.2900
0.4907
0.7065
0.8444
ACT-MSE
0,3154
0.4899
0,5557
0.6023
0.3827
0.4421
0,5965
0,7175
LM-PPL
0.3767
0,5011
0.5138
0.7152
0,4129
0.5100
0.6907
0.6997
SMART
0,4476
0.6274
0.5595
0.7728
0.4089
0.6401
0.8126
0.8190
AR-MH
0.4622
0.5512
0.6329
0.8231
0.4371
0.7160
0.7318
0.7983
Learning- Boaning-Bassed (SI MSE)
GR-S1
0.2903
0.4508
0.5214
0.6435
0.3563
0.4827
0.6743
0.8209
GA-S1
0.3092
0.4193
0.4892
0.6569
0.3843
0.5408
0.6657
0.7873
AVG-S1
0,3108
0,4233
0.5355
0.6390
0.3960
0.4933
0.6441
0.8226
Learning-Basee Based (S2 KI.J
GR-S2
0.3084
0.4950
0.6991
0.7662
0,5804
0.8259
0.8541
0.8869
GA-S2
0,5731
0.7539
0.8118
0,8619
0.6617
0.8631
0.8733
0.9061
AvG-S2
0.4764
0.5580
0.6786
0.8111
0.7075
0.8205
0.8704
0.9051
Table 6: RULER performance for various layer selection strategies across different softmax ratios, for GDN- bosed hybrid students. The all-linear (0%) baselines are 0.0427 for Llama-3.2 and 0.1236 for Qwen2.5. The all-softmax teacher scores are 0.8934 and 0.9174, respectively.
Llama-3.2-3B-Instru
Qwen2.5-3B-Instrucc
Selector
12.5%
25%
33%
50%
12.5 %
25%
33%
50%
Heuristic-Based
UNIFORM
0.3648
0.4011
0.4728
0.6642
0.2523
0.6334
0.3684
0.7178
KV
0.2069
0.6461
0.6760
0.6942
0.1370
0.6788
0.6261
0.7096
AR
0,5000
0.7042
0.7505
0.7051
0,4601
0.6016
0.6379
0,7432
VT
0.1978
0.3385
0.3648
0.6960
0.1588
0.4183
0.4809
0.6279
CWE
0.3149
0.3258
0.6207
0.6779
0.0789
0.2087
0.5345
0.6842
ACT-MSE
0.3058
0.5045
0.5726
0.5672
0.1579
0.2405
0.3385
0.5771
LM-PPI.
0.4201
0.5000
0.5299
0.7613
0.2523
0.2541
0.4809
0.5508
SMART
0.3276
0,7033
0.4374
0.7169
0.1515
0.3448
0.7359
0.6679
AR-MH
0.4773
0.5572
0.6833
0.7505
0.1588
0.6416
0.6679
0.7069
Learning- Baaeing-Baseddieee (SI
MSE)
GR-S1
0.2015
0.3548
0.5644
0.6007
0.2677
0.4465
0.5100
0.6697
GA-SI
0.2105
0.4365
0.4746
0.5563
0.2532
0.4247
0.5163
0.6234
AVG-S1
0.1951
0.4074
0.4628
0.6443
0.2414
0.4165
0.5227
0.6751
Learning- -Based (S2 KLJ
GR-S2
0.3303
0.5054
0.6933
0.6633
0.3612
0.6860
0.7459
0.7468
GA-S2
0.7314
0.7514
0.7514
0.7595
0.6025
0.7559
0.7822
0.8004
AvG-S2
0.6588
0.6806
0.7241
0.7060
0.5880
0.7532
0.7196
0.7641
Table FDA performance for various layer selection strategies across different softmax ratios, for GDN-based based hybrid students.
16

---

Published as a conference paper at ICLR 2026
Llama-3.2-3B-Immtt
Owen2.5533BIIsstrue
Selector
12.5%
25%
33%
50%
12.5%
25%
33%
50%
Heuristie-Bar Based
UNIFORM
0.7615
0.7678
0.8236
0.8479
0.6760
0.8515
0.8380
0.8740
KV
0.4671
0.7894
0.8110
0.8515
0.5311
0.8074
0.8101
0.8272
AR
0.6850
0.8245
0.8326
0.8263
0.7498
0.7777
0.7984
0.8839
VT
0,4761
0.6688
0.6895
0.8569
0.5572
0.7255
0.7507
0.8587
CWE
0.5878
0.6598
0.8290
0.8569
0.5302
0.7192
0.8020
0.8956
ACT-MSE
0.6121
0.7687
0.8254
0.8380
0.5725
0.6742
0.7939
0.8398
LM-PPL
0.6985
0.7714
0.7795
0.8443
0.6535
0.7381
0.8155
0,8353
SMART
0.7669
0.8155
0.8074
0,8659
0.6733
0.8335
0.8767
0.8902
AR-MH
0.7678
0.8254
0.8335
0.8551
0.6445
0.7921
0.7930
0.8299
Learning Based (S1
MSE)
GR-S1
0.5779
0.6958
0.7480
0.8254
0.6688
0.7831
0,8326
0.8821
GA-S1
0.5707
0.7282
0.8146
0.8344
0.6553
0.8047
0,8569
0.8668
AVG-S1
0,5671
0,7192
0.7957
0.8254
0.6670
0.7975
0.8506
0.8866
Learning Based (S2
KLJ
GR-S2
0.6301
0.8110
0.8245
0.8425
0.8299
0.8875
0.8749
0.8929
GA-S2
0.7885
0,8506
0.8614
0.8587
0.8398
0.8839
0.8929
0.8992
AvG-S2
0.7885
0.8137
0,8565
0,8704
0.8128
0.8848
0.9001
0.9109
Table SWDE performance for various layer selection strategies across different softmax ratios, for GDN-based hybrid students.
Llama-32-3BB-Intruuu
Qwen2.5-3B-Instre
Selector
12.5%
25%
33%
50%
12.5%
25%
33%
50%
Heusistic-Based
UNIFORM
18.1553
21.5138
22.1316
24.1098
10.3260
13,2595
11.3725
16,9006
KV
17.4030
25.5568
26.3946
29.5483
6.6478
10.8318
16.4796
15.2550
AR
19.5522
25.7810
28.0807
30,8078
10.1869
12,3800
10.1677
9.3073
VT
19.0819
24.3118
23,9263
29.9387
7.1499
8.7797
14.3150
18.8876
CWE
23.7679
23.2527
28.0014
30.3961
6.7367
12.9678
9,9249
7.2352
ACT-MSE
17.2452
22.1599
24.1714
24.8185
9,9016
8,5959
8.8570
13.7590
LM-PPL
20,3075
22.2478
22.4090
27.4774
11.1046
10.4351
9.2945
8.1057
SMART
18.9845
28.1036
22.8017
31.0010
14.1967
12,4760
12,8118
16,8912
AR-MH
24.0994
27.0128
28.0295
29.2224
13.7095
11.9901
11.2729
15.8943
Learning- BoenningBBasedan (SI MSE)
GR-S1
13.3918
20.7552
23.2197
27.3407
7.8245
7,0497
9,4220
8.6667
GA-S1
13.6481
17.8867
22.6633
29.2390
8.9412
9.0555
11.1751
9.1234
AvG-S1
15.0889
18.4342
24.3658
28.2178
7.6409
10.6217
10.1589
10.3181
Learning-Bast Based (S2 KL)
GR-S2
18.0648
25.7848
30.4299
30.5907
12.1582
6,4855
7.8482
6.9539
GA-S2
25.1144
30.9408
31 31.4913
33.3520
9.7717
13.4097
13,9486
13.9875
AvG-S2
23,5556
29.2189
31.1063
32.1499
10.6181
6.4121
6.5623
11.3837
Table SQuADv2 (F1) performance for various layer selection strategies across different softmax ratios, for GDN-based hybrid students.
17

---

Published as a conference paper at ICLR 2026
B ELABORATION ON EARLY STOPPING FOR EFFICIENT SELECTION
Protocol. We study the sample efficiency of our one-swap selector ($3.2) at a fixed hybrid ratio of 1:3 (K=9 for Qwen2.5-3B -3ntstructt,,, = for Llama-3.2-3B 1aaaa-3.-3BInssttucctt During Stage-2 we train for 4,550 steps and, every 50 steps, compute the current top-K set of layers (from the one-swap importance scores). This yields 91 snapshot sets per model. To quantify stability we analyze each rolling window of the last R=10 snapshots using complementary views:
Rolling pairwise similarity: the mean pairwise Jaccard over the R sets.
Rolling backbone size: the size of the intersection across the R sets (how many positions are "locked in").
We also relate snapshots to the final selection by reporting the fraction that are within one swap of the final consensus (Jaccard 2 K-1: K+1 i... 0.80 for K= and 0.75 for K=7)
Reliable selections emerge well before 4550 steps. ers:
Two patterns consistent across both teach-
Qwen2.5-3B-Instrue (K=9). The run-best first appears by step 850. From step 1500 onward, 95% of snapshot sets are within one swap of the final consensus; the 10-snapshot rolling Jaccard is high on average 0s95,,, and rises to 0.99 beyond step 2350.
By step 1900, the last snapshots share an 8/9 backbone with at most two candidates for the remaining slot; any one-swap variant at this point attains RULER within 0.007-0.009 absolute points of the run-best (0.8662 vs. 25559200.558/0..55
(K=7). A 6/7 backbone appears by step 750 (mean window Jac- card x 0.91). The near-optimal set that differs by single layer first appears at step 1200; from step 1200 onward, 100% of snapshots are within swap of the final consensus. Stopping here gives RULER 0.6971, within 0.004 absolute of the run-best 0.7011 and comparable to the best late-appearing set.
-lamm--.22-B---me
These observations (i) The selector's rankings stabilize far earlier than the full 4500-step budget; (ii) the windowed sets agree K-1 layers, the remaining degree of freedom is small and can be resolved cheaply: (iii) one-swap neighbors of the eventual best set typically match downstream RULER within 0.1-1.0 absolute points, so stopping once the K- backbone is stable is a sound efficiency-qualle trade-off.
conservative choice (see rule below) would have stopped at 1900 steps for Qwen and 1200 steps for Llama- Llama-consuminn 42% and 27% of the 4550-step budget, respectively (i.e., 58-74% fewer tokens for the selection pass).
Practical recipe (rolling-Jaccard early stop). Let S₁ be the top-K set at step and W. = (S.- S.). Define
2 R(R - 1) K<j
Backbone;
n S, JaccardMean, SEW,
Σ
Jac(S (.,
Stop the first step satisfying:
1. JaccardMean, 2 0.90,
2. Backbonez 2 K 1, and
3. Usew, S) < K + 1 (at most two options for the remaining slot).
(Optional) Stop when (3) first becomes true and S, + S.- to pick the newer of the two candidates.
Foor fixed set size K, replacing one layer yields intersection K-1 and union K+1, hence Jaccard (K 1)/(K + 1).
18

---

Published as a conference paper at ICLR 2026
C COMPLETE LAYER IMPORTANCE RANKINGS
For all methods that produce scalar importance score per layer, we obtain hybrid architectures at target softmax ratios (12.5%, 25%, 33%, 50%) by taking the top-K most important layers according to that ranking (with K determined by the ratio and total depth L). In this section we report the full importance ranking for each such method. Layer indices are sero-based. Methods such as POSTNAS and SMART do not provide layerwise importance scores, so they are omitted here.
C.I QWEN2.5-3B-INSTREN
Selector
Layer indices (most
least important)
KV
[1. 0, 26, 19. 18, 20, 5. 17, 27, 6. 15, 22. 24, 16, 3. 11, 23, 21. 28, 8, 14, 25, 32, 12,
13,9,
10, 31, 34, 35, 30, 33, 71
AR
1. 18, 20, 25, 24, 26,
12, 19, 7,35, 17, 33, 22, 28, 16, 32, 34, 29,
6,
14, 15, 11]
VT
19, 26, 28, 25, 10, 15, 17,3,7,27, 16, 10, 34, 32, 31, 23, 33, 9, 13, 18, 8,
2,, 11, 11,
22,
24, 5,4,
CWE
1,
24, 16. 13. 2, 29, 11, 6, 31, 28, 29, 34, 7,30, 32, 9, 25,
2,,1,
18,
14, 12]
ACT-MSE
[0, 1.35, 34,
28, 10, 29, 11, 7, 14. 13, 26, 25, 16. 32, 12, 2...2..,33.22.00,,,,22233222
2, 4,7, 10, 11, 25, 28, 16, 14, 13, 26, 24, 3,, 34, 31, 33, 30, 27, 12, 20,2222221
20,
LM-PPI.
20, 23, 15, 18, 29,
AR-MH
1,,7,
34, 15, 26, 16, 20, 28, 11, 2,, 18, 6, 13,3,9 9, 22, 8, 17, 35, 19, 4, 25, 12, 30, 7, 29,
GA-SI
32.
34. 31.
19, 30, 24. 16, 23, 26. 12, 17, 18. 20, 14, 25, 10, 3, 29,
61, 13.7,9 15.
GR-S1
[33, 32, 34. 35,
26, 25, 16, 24, 23, 17, 18, 15, 12, 13, 30, 21,
11, 1.
AVG-S1
[33, 32, 31, 28, 29,
20, 30. 19, 16, 24, 26, 23, 17, 25, 18, 12, 14, 20, 10, 11, 13, 8.9.
GA-S2 (OuRS)
[20, 33, 25, 17, 31,,4,, 10, 26, 29, 0, 15, 23, 6, 12, 24,
18, 9,
11, 16, 2,
GR-S2
0, 17, 22, 20, 25, 25, 18, 24. 15, 29, 12, 26, 31, 16, 3, 10, 13, 14, 28, 30, 5. 11, 4, 9,2, 34,
[21. 19, 2,
AVG-S2
[21. 32. 20. 19, 22, 17, 25, 27, 31, 29, 23, 10, 15, 24, 18, 30, 12, 13, 4. 28.
16, 7, 14,
11, 34, 35,
Table 10: Complete layer-importance rankings for Qwen2. 5 3-3B- IBs-ntttuct. Each row lists all 36 layers from most least important.
C.2 CAA-MA--.2-B--INTTTEE
19

---

Published as a conference paper at ICLR 2026
Selector
Layer indices (most + least important)
KV
[0. 7,5, 11, 14, 3, 6, 23, 10, 20, 26, 17, 22, 9, 24, 21, 25, 18, 16, 19, 12, 13, 27,
15]
AR
[0, 16, 11. 14. 14, 7.5,9, 13, 2. 12, 1, 8, 27, 26, 10, 6. 24, 15, 3, 20. 18, 19, 17, 21, 25, 4. 23,
22]
VT
15, 16, 18, 23, 14, 21, 24, 20, 26, 22, 6, 19, 10, 12, 10, 1, 17,9,
71
CWE
[0, 5, 12, 8, 9, 1, 13, 14, 10, 21, 24, 16, 22, 15, 25, 20, 6, 2. 26, 23, 18, 3, 11, 17, 19,
7]
ACT-MSE
2,, 25, 26, 4, 15, 19, 21, 18, 20, 14, 5, 22, 17, 13, 12, 11, 10,
9]
LM-PPL
[0, 2, 3, 26, 25,4, 14, 15, 19, 16, 5, 20, 23, 12, 17, 10, 21, 13, 18, 6, 22, 9, 11,
8]
AR-MH
[0, 13, 12, 16, 11, 7. 14, 1., 25, 9. 19, 17. 6, 4. 1. 18, 24, 15, 22. 20]
GA-S1
24, 13, 20, 23,7, 10, 22, 12, 19,8, 14, 15, 11, 16, 17, 18, 2, 6, 1,
31
GR-SI
226. 24, 25, 23, 12, 22. 13, 14. 10. 19, 11, 15,9, 16. 18, 17, 6, 5, 5, 3,2, 0]
AVG-S1
[26, 27. 24, 25, 23, 22, 12, 20, 14. 19,7.9.21, 15,8, 11, 16. 17, 5, 6, 2, 4, 1.3,
0]
GA-S2 (OuRS)
[14, 8, 5, 12, 15, 13, 26, 24, 16, 17, 18, 21, 10, 25, 19, 23, 22, 9,, 7, 6, 1, 11,
3]
GR-S2
1, 12, 13, S. 10, 14, 7,9, 6, 11, 26, 1S, 16, 4,, 24, 27, 25, 19, 17, 18, 23, 21, 20]
AVG-S2
[12, 5, 14,2, 8, 13, 10, 15, 26, 0, 1, 16, 24, 7, 6, 17, 18, 25, 22, 19, 21, 3, 11, 20, 23]
Table 11: Complete inyrrimmpotancc rankings for L1 aallamm333. -..2......... 2-3B-Instruct. Each row lists
all L = 28 layers from most to least important.
20

---

Published as a conference paper at ICLR 2026
D LAYER-SELENE -EEECSEEN PATTERNS AND SPATIAL ORGANIZATION
We now examine where in depth the selected softmax layers tend to lie, and whether our selector prefers isolated layers or groups of consecutive layers.
Setup. For each teacher we take the GA-S2 ranking R R= (......, lt) from Appendix C ordered from most least important. For a softmax budget K we define the selected set SK (f1,....63 Lx To quantify how much the selected layers cluster in depth, we use the adjacency index
Ax = i E Sx : + Sx }/-
ic., the number pairs of consecutive layers that are both selected. For a uniformly random K- subset of L 1}, the expected value is E[Ax] 2 K(K - 1)/L, values substantially above this baseline indicate more clustering than would be obtained by chance. F shows the selected indices across budgets, and Figure 7 compares observed and expected adjacency counts.
Results and discussion. For Bree2..5-BBInsttruee (L=36), GA-S2 produces selected sets that visibly concentrated in few depth ranges. At 25% budget (K=9), we obtain Ax 4.0 versus random baseline of 2.0; at 33% 12), Ax 7.0 versus 3.68; and at 50% (K (1=18), Ax 11.0 versus 8.49. The plot in Figure show that several of these adjacent pairs occur repeatedly around layers roughly 3-5, 19-22, and 31-33, while the remaining layers are used more sparsely. Thus, the selector does not simply spread the softmax layers uniformly but repeatedly reuses a small number of depth regions the budget increases.
(L=28), the effect weaker but still present. At 25% (K=7), Ax 3.0 versus baseline of 1.50; at 33% (K==)), Ax = 3.0 versus 2.58; and at 50% (K (K=14), AK = 6.0 versus 6.50. The selected layers tend to form one main group in the middle of the network (around layers 12-18), with smaller number of layers near the input and output.
For 3.ama-3.2-3B-Imm.
Overall, both models show some degree of clustering beyond what would be expected from ran- dom K-subset, but the pattern (multiple groups versus single main group) depends the teacher architecture.
21

---

Published as a conference paper at ICLR 2026
Iwen2..-3B-Instruct: Selected Layers Across Budgets (GA-S2)
(50.0%)
Budget Configuration
K=12 (33.3%)
Km9 (25.0%)
K-4 (11.1%)
25
Layer Index
Llama-3.2-28 2.-ma--.--3Btnsstruuct Selected Layers Across Budgets (GA-S2)
K-14 (50.0%)
Budget Configuration
K=9 (32.
(25.0%)
K=4 (14.3%)
10
20
Layer Index
Figure 6: Visualization of selected layers for Seen2.5-3B-Instruue (top) and Llama-3.2 (bottom) across budgets (12.5%, 25%, 33%, 50%). Each vertical tick marks selected layer index.
Llama-3.2-3B-Instruue
Observed
Observed
Expected (Random)
Expected (Random)
8.5
Adjacency Count
Adjacency Count
3
3.0
3.0
2.6
4
4.0
3.7
2
0.3
0.0
K=4
K=9
K=12
K=18 (50.0%)
K=4
K=7
K=9 (32.1%)
K=14 (50.0%)
(11.1%)
(25.0%)
(33.3%)
(14.3%)
(25.0 (25.0%)
Budget (Softmax Layers)
Budget (Softmax Layers)
Figure 7: Observed (solid) vs. random- random-baseline expected (dashed) adjacency counts Ak for Qwen2.5-3B (left) and 1.ama-3.2-3B (right).
22

---

Published as a conference paper at ICLR 2026
E
SELECTION (DIVERSIFICATION ABLATION)
To probe whether clustering redundant, we evaluate a re-weighted greedy rule for selecting K layers:
le 1
I(e I(()
Σ JES
exp
with > 0 Here S is the set of softmax layers selected so far and I(e) is the original softmax layer too close (in
GA-S2 importance score. The exponential term penalizes placing depth) to previously selected ones, nudging the selector toward more spatially diverse configurations without discarding the model-intrinsic KL signal.
We instantiate this diversification for Qwen2.5-3B. 1aee2.5533-Instreee -B-Innstrct with a GDN student at a fixed 25% soft- ratio 99) and sweep À {0.025, 0.05) and [1,2). All other training and evaluation settings are kept identical to the main GA-S2 runs.
RULER (4096)
Selected layers
(GA S2)
0,8713
[20, 32, 33, 21, 22,
25, 17, 19.
0.025
1
0,8509
[20, 32, 25,
17,
22,
5, 33, 10,
0.025
2
0.8244
[20, 32, 25. 17. 10, 33,
22]
0.050
0.8334
[20, 32, 25, 17,
10,
22,
29]
0.050
0,8303
[20, 32, 5, 25, 10, 17, O, 33,
13]
Table 12:
Distance-regulariied GA-S2 selection on Qwen2.5-3B-Inss Instructt with a GDN student at 25% softmax ratio. The λ=0 row corresponds to our default GA-S2 selector without regularization; the last column lists the resulting softmax layer indices.
As shown in Table 12 none of the distance-regularizee variants outperform the unregularized GA- S2 selector. A mild penalty (A=0.025, o-1) yields small degradation (0.8509 vs. 0.8713 RULER). while stronger or more broadly supported penalties lead to larger drops. This suggests that the clustering observed in our selections not merely redundant: forcing softmax layers to spread out in depth tends to remove genuinely useful local groupings. At the same time, the X=0.025, a=1 configuration may be acceptable when a slightly more uniform spatial allocation is desired and modest recall loss (about two points on RULER) is tolerable.
23

---

Published as a conference paper at ICLR 2026
F EXTENDED LONG-CONTEXT EVALUATION VIA NEBDLE-IN-A--AYSTAAC
In the main text, long-context behavior is evaluated primarily through RULER and SWDE ($4 $4.1), whose contexts are below 10k tokens, and our distillation pipeline ($3.1) trained on generic text with comparatively shorter sequence lengths. This leaves open whether the distilled hybrid model recovers teacher-like retrieval ability at substantially longer sequences than those used during distillation and benchmark evaluation. To probe this, we perform an additional needle sondle-in-ahhaysee a--a--htttakk (NiHA) experiment.
We consider the Qwen2.5-3B-Img Innstruuttt teacher and its corresponding hybrid student with 25% softmax 75% GDN configuration selected by our method. For each context length, we construct inputs by embedding single target "needle" span into long filler context and measure retrieval accuracy, defined as the fraction of cases where the model correctly identifies the target span. We evaluate across exponentially increasing context window sizes from 8k to 128k tokens. Results are reported in Table I
Context length (tokens)
Teacher
Hybrid student
8,192
1.000
1.000
16,384
1.000
0.998
32,768
1.000
0.998
65,536
1.000
0.994
131,072
0.954
0.684
Table 13: Needle-in-a-hhysstcck retrieval accuracy as a function of context length for Qwen2.5-3B- Instruct (teacher) and the corresponding hybrid student (25% softmax, 75% GDN layers).
The hybrid model maintains near-perfect retrieval accuracy up to 65,536 tokens, closely tracking the teacher with only minor degradation. At 131,072 tokens both models begin to degrade, with larger drop for the hybrid student. These results indicate that the proposed layer selection and distillation procedure successfully preserves long-conted retrieval well beyond the context lengths used during distillation and primary benchmark evaluations, while leaving further improvements extreme lengths as an interesting direction for future work.
24

---

Published as a conference paper at ICLR 2026
G ADDITIONAL SCALING RESULTS FOR QWEN2.5 TEACHERS
To verify that our KL-guided layer selection method scales across model sizes within a family, we also distill GDN-based hybrid students from additional Qwen2.5 teachers:
Qwen2.5-1.5B-Instrrue with RULER score 0.8742.
Qwen2.5-7B-Instrucc, with RULER score 0.9445.
We use the same DCLM mixture and distillation pipeline as in the main Qwen2.5-3B experiments, and evaluate 25% and 33% softmax ratios. As in the main text, compare against UNI- FORM, AR, AR-MH, ACT-MSE: LM-PPL, and SMART. Our selector GA-S2 remains consis- tently stronger than all baselines, particularly in the low-budget regime.
Model Ratio
UNIFORM
AR
AR-MH
ACT-MSE
LM-PPL
SMART
GA-S2
Qwen2s5-1.B-IIstruutt
(teacher RULER: 0.8742)
25% 33%
0,4778
0,5096
0,4243
0.3807 0.4374
0,4271 0.5056
0.5098 0.6479
0,5408
0.5651
0.5552
0.5229
0.6953
Qwen2.5-77-Instruct
(teacher RULER: 0.9445)
25%
0.7357
0.7453
0.7322
0.6469
0.6544 0.6590
0.8158
0.8584
33%
0.7516
0.8423
0.8533
0.7227
0.8949
0.9110
Table 14: RULER performance of GDN-based hybrid students distilled from smaller (1.5B) and larger (7B) Qwen2.5 teachers at 25% and 33% softmax ratios. Our GA-S2 selector consistently outperforms all baselines across scales.
25


Reviews:
[{"review_id": "1Rvdjrhcjo", "reviewer": "Reviewer_X6VJ", "summary": "This paper proposes a method for distilling pretrained softmax attention Transformers into more efficient hybrid architectures that interleave softmax and linear attention layers. The key aspect of this process is layer selection, i.e., deciding which layers to convert to linear attention variants. \n\nThe authors present a simple and efficient approach for layer selection based on layer importance scores derived from a small amount of training on generic text data. Once the layers have been selected, the authors use a recent distillation pipeline (RADLADS) consisting of attention weight transfer, hidden state alignment, KL-based distribution matching, and a small amount of finetuning. The authors find this approach to be more effective than existing methods for layer selection.", "questions": "* Is there any clear pattern or intuition of which layers get selected for linearization ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Addresses a Practical**. Tackles the challenge of converting pretrained softmax attention Transformers into more efficient hybrid architectures without expensive pretraining from scratch. Focuses on improving inference efficiency of LLMs, which is a critical concern in practical deployments\n* **Intuitive Layer Selection Approach**. Proposes a KL-guided layer selection criterion that is both simple and theoretically motivated. The intuition is clear: layers that are more critical for maintaining performance will show larger KL divergence reduction when kept as global attention layers. This represents an improvement over fixed interleaving strategies used in previous work.\n* **Comprehensive Empirical Analysis and Ablations**. Provides detailed analysis showing the trade-offs between different attention mechanisms on various task types. Demonstrates that sliding window attention works well for common-sense reasoning but struggles with in-context recall tasks (Figure 1). Shows empirical evidence that their approach outperforms existing layer selection methods\n* **Clear Problem Formulation**.\nClearly identifies the two key decisions in distillation: student architecture selection and distillation recipe optimization\nFocuses on the less-explored but equally important architecture selection problem", "weaknesses": "* How does this method compare to a baseline that just randomly selects the layers to replace with linear ones i.e instead of Uniform or any of the fancy methods of selecting the layers to linearize if we just randomly chose K, layers and linearized them, how would that affect the performance. \n* The number of layers to be linearized K seems to be a heuristic or dataset dependent ? It makes the method feel somewhat brittle. Would this automatically transfer to some other task or would you need to do a linear transfer of layers again to get it working?\n* The goal of linearization etc., is to speed up the performance. How does the memory, latency/throughput required to match or be close to the quality of full attention, compare to that of using full 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": "0lAEJtkuLG", "reviewer": "Reviewer_GNir", "summary": "This paper proposes a learned method to choose which layers should be softmax attention layers and which should be linear attention layers in a student model to distil from a teacher model with only softmax attention layers. To achieve this, the authors use distillation and KL-loss to guide which layers to select. They start by distilling a fully linear attention student model. Then they used that distilled model to calculate the importance of each layer being a softmax layer, by distilling the previous model, with its k-th layer replaced by a softmax layer. They then use the KL-loss to the teacher model to calculate the K most important layers, which they switch to softmax layers before doing a final distillation. Their method requires few training tokens (600M unique, 20B total) and outperforms other techiniques.", "questions": "- For Figure 2, it would be nice to have the performance when all the layers are linear.\n- Did you consider the existence of clusters of important layers, which could indicate that only one of those layers is important, rather than all of them being important?\n- Did you try having a penalty for important layers that are close to each other (to avoid selecting them, in case their importance was due to the same factors)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well structured and written, the sections flow naturally from one to another, and the motivation is very clear.\n- The experiments are thorough, and the authors compare their method to various baselines and show that their method outperforms the others in settings where the vast majority (75% or more) of the layers are linear layers rather than softmax layers.\n- The method is simple and effective, and does not require much computing (compared to the total pre-training cost) to create much more efficient models with little to no performance cost. This is especially important because of the reduction in inference cost for such models.\n- The authors do multiple ablations to show that the choice of stage at which to calculate importance and to do addition of softmax layers rather than removal of softmax layers is justified. In addition, they show that their method works for different linear attentions and that calculating importance using one type of linear layer can be translated to another type of linear attention", "weaknesses": "- On line 339, the authors state that their method is iterative; however, this is not totally correct, their method calculates the importance of each softmax layer independently, and then selects the top-K best. This ignores how adding one softmax layer can affect the importance of other softmax layers, which an iterative process (add one, recalculate importance, add another, etc.) would take into account.\n- It would be interesting and helpful for the paper to know which layers were chosen. Whether the layers with the highest importance were near each other or if they were spread out through the model (in other words, checking whether there exist clusters of important layers).", "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": "ZEUOj2b18T", "reviewer": "Reviewer_CzRy", "summary": "This paper tackles the problem of converting pretrained softmax attention Transformers into more efficient hybrid architectures that combine softmax and linear attention layers. While prior works have largely used fixed interleaving ratios (e.g., 1 softmax per 3 linear layers), this paper argues that such heuristics are suboptimal for distillation settings.\n\nThe authors propose a simple yet effective KL-guided layer selection strategy: Each layer’s importance is determined by the reduction in teacher–student KL divergence when that layer alone is restored to softmax attention in an otherwise all-linear student. Layers with the highest KL importance scores are retained as global attention layers in the final hybrid model.\n\nExtensive experiments with Qwen2.5-3B-Instruct and Llama-3.2-3B-Instruct show that this method consistently outperforms heuristic or NAS-based baselines (e.g., UNIFORM, SMART, PostNAS), particularly on long-context recall benchmarks such as RULER and SWDE.\nThe approach requires only ~25–30B tokens, far fewer than PostNAS (400B).", "questions": "Validate the layer-selection behavior across scales (e.g., 1B, 7B, or 14B) to demonstrate scalability and consistency.\n\nProvide intuition or gradient-based analysis showing why KL reduction correlates with information retention or attention selectivity.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Despite its simplicity, the method achieves large gains—especially in the low-softmax regime (e.g., 12.5% ratio).\nThe KL-based layer importance metric is conceptually elegant and empirically grounded.", "weaknesses": "All experiments focus on 3B-class decoder-only models; no results for encoder–decoder or smaller-scale models.\n\nWhile the paper emphasizes efficiency, explicit measurements of inference latency or memory savings are missing.\n\nThe accuracy of layer importance relies heavily on the stability of Stage-2 KL distillation. If the teacher–student mismatch is large, the ranking may become unreliable.", "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": "i3nFsAx9r2", "reviewer": "Reviewer_7f2E", "summary": "This paper proposes a distillation and layer selection method that allows to derive hybrid attention LLMs from classical Transformers-based ones. Their method allows them to turn computationally intensive models into more efficient alternatives without training from scratch, while retaining strong performance on in-context retrieval tasks, where previous methods were less performant.\n\nThe proposed method is based on a 2-step distillation process (first at layer-wise attention representation level, and then at logit level), and on a greedy layer switching exploration of possible hybrid configurations. The best identified configuration is then distilled again.\n\nThe authors conduct extensive evaluations on the RULER and SWDE benchmarks with 3B Llama and Qwen models, showing the efficacy of their approach.", "questions": "- Why is a stage-1 MSE step necessary in the final distillation stage? Would re-using the distilled layers from the previous step make sense?\n- How did the distilled models behave in terms of validation perplexity on the DCLM benchmark? Is the distillation training process stable?\n- Although the Jaccard similarity seems to provide valuable insights on the transferability and similarity of selected layers, I am wondering if a correlation/distillation metric would be better-suited for this analysis. The Jaccard metric fails to differentiate the ranking of the layers by the $\\mathcal{I}$ score, which would be captured by e.g. a Spearman correlation, and seems like a meaningful information to compare two selections. Did you run such analysis?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Overall, this paper clearly states its problem and proposes a simple but efficient method as a solution.\n- **The method is rather straightforward and yields benefits**: the proposed greedy approach is easy to grasp and seems to provide good results on the two benchmarks used, especially for Qwen.\n- **Several variants and combinations are tested**: I appreciate that the authors have tried sensible variants of the Greedy Addition method in Table 1. This provides valuable insights and helps underline the reason for the success of their method. It is also nice to see the interplay between different linear attention variants in section 5.2, and the discussion about early stopping in section 5.3.", "weaknesses": "- **Limited scaling potential**: The method presented in this paper has compute requirements that scale linearly in the number of layers. As model grow larger, they will tend to have more layers and will thus be doubly more expensive (both from the parameter and data viewpoints) to distil into hybrid variants. Given that the authors do not provide a data point at a different size (either smaller or larger than 3B), it is difficult to assess the relevance of this method at a larger scale where it will be more costly to run.\n- **Lack of discussion about selected layers**: There is no discussion or analysis about potential patterns (or the absence thereof) for the selected layers, and the behavior of $\\mathcal{I}(l)$. This would be a valuable insight for subsequent work.\n- **Limited evaluation scenarios**: Although the authors report successful results on RULER and SWDE, it remains unclear whether their approach recovers the same long-context abilities than softmax-based models for every sequence lengths, or whether the hybrid distilled models solely match the performance of the teacher models on the context lengths used in these benchmarks (<10k tokens). A perplexity-based analysis or a needle-in-a-haystack could help provide insights on that question.", "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)