PaperlayerView exact revision on arXiv

Source-linked AI summary

Hybrid Linear Attention Done Right: Efficient Distillation and Effective Architectures for Extremely Long Contexts

Yingfa Chen, Zhen Leng Thai, Zihan Zhou, Zhu Zhang, Xingyu Shen, Shuo Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu

arXiv:2601.22156v1cs.CLcs.AIcs.LG

TL;DR

Existing Transformer-to-hybrid distillation methods demand prohibitive data, while pure RNNs underperform attention on some tasks. This paper introduces HALO and HypeNet, whose HyPE encoding enables efficient conversion with superior long-context tradeoffs.

  • Problem

    Existing Transformer-to-hybrid distillation methods require tens to hundreds of billions of tokens, while pure RNNs generally underperform softmax attention on recall-intensive tasks.

  • Method

    The paper proposes HALO for efficient Transformer-to-hybrid distillation and HypeNet, which combines HyPE with architectural improvements for length generalization.

  • Results

    Applying HALO to Qwen3 produces hybrid models with much better performance-throughput tradeoffs and memory efficiency in long-context scenarios.

  • Takeaways & Limitations

    The approach makes empirically validating hybrid architectures at scale cheaper and supports cost-efficient long-context language-model research.

  • Takeaways & Limitations

    Conversion on pre-training-style data may diminish instruction-following and alignment behaviors, and recovering these capabilities remains open.

Abstract

from arXiv · show

Hybrid Transformer architectures, which combine softmax attention blocks and recurrent neural networks (RNNs), have shown a desirable performance-throughput tradeoff for long-context modeling, but their adoption and studies are hindered by the prohibitive cost of large-scale pre-training from scratch. Some recent studies have shown that pre-trained softmax attention blocks can be converted into RNN blocks through parameter transfer and knowledge distillation. However, these transfer methods require substantial amounts of training data (more than 10B tokens), and the resulting hybrid models also exhibit poor long-context performance, which is the scenario where hybrid models enjoy significant inference speedups over Transformer-based models. In this paper, we present HALO (Hybrid Attention via Layer Optimization), a pipeline for distilling Transformer models into RNN-attention hybrid models. We then present HypeNet, a hybrid architecture with superior length generalization enabled by a novel position encoding scheme (named HyPE) and various architectural modifications. We convert the Qwen3 series into HypeNet using HALO, achieving performance comparable to the original Transformer models while enjoying superior long-context performance and efficiency. The conversion requires just 2.3B tokens, less than 0.01% of their pre-training data

1. Introduction

Transformers are prohibitively expensive for long contexts, while hybrid attention–RNN architectures offer better throughput but remain difficult to pre-train and distill effectively. This paper introduces HALO, HyPE, and HypeNet to enable efficient Transformer conversion with stronger long-context generalization.

  • Motivation: Transformers have quadratic context-length complexity, whereas RNNs such as linear attention and state space models have linear complexity and are faster for long-context modeling.Hybrid architectures interleave softmax attention and RNN layers to trade off performance and inference throughput.
  • Limitations: Existing distillation methods require tens to hundreds of billions of training tokens and suffer severe long-context performance degradation despite comparable short-context performance.These limitations constrain academic research and undermine hybrid models in the scenarios where they provide the greatest advantage.
  • Method: HALO converts pre-trained Transformer models into attention–RNN hybrids using fewer than 3B tokens.Its efficient attention-layer selection method determines which attention layers remain unconverted to preserve long-context performance.
  • Method: HyPE combines RoPE and NoPE with attention scaling to provide superior length generalization for hybrid models.The position-encoding scheme is designed specifically for hybrid architectures.
  • Architecture: HypeNet is a hybrid architecture based on HyPE that incorporates multiple architectural improvements when converting a pre-trained Transformer model.The paper evaluates HypeNet as the target architecture of its Transformer-to-hybrid conversion approach.

2. Related Works

Prior hybrid models match Transformers on several tasks while processing long contexts more efficiently, but hybrid distillation remains underexplored and existing methods require substantial data with weak long-context recall. This work uses a novel position encoding scheme and a 2.3B-token distillation procedure to improve long-context performance through stronger length generalization.

  • RNN-Attention Hybrid Models: State-of-the-art hybrid models with up to hundreds of billions of parameters achieve performance comparable to Transformers on commonsense reasoning and recall-intensive tasks while processing long contexts more efficiently.The cited tasks include needle-in-a-haystack recall evaluation.
  • Position Encoding in Hybrid Models: RoPE is standard in Transformers, whereas RNNs typically encode positions through decay or transition matrices and do not use RoPE; hybrid models generally give attention layers RoPE.Kimi-Linear concurrently adopted NoPE in both attention and RNN layers.
  • Position Encoding in Hybrid Models: The proposed model uses a novel position encoding scheme and achieves better long-context performance than typical position encoding methods in existing hybrid models.This contrasts with Kimi-Linear’s use of NoPE in both attention and RNN layers.
  • Distilling Transformers into Hybrid Models: Converting Transformers into hybrid models remains underexplored, and selecting which attention layers to convert to RNNs is critical for preserving performance on tasks difficult for RNNs.Prior work has more commonly converted Transformers into pure RNN models through distillation.
  • Distilling Transformers into Hybrid Models: 2.3B training tokens are sufficient for the proposed distillation procedure, versus more than 10B tokens typically used by prior works, while the resulting architecture provides much stronger long-context performance.Prior works also show poor recall performance relative to Transformers, especially on long contexts.

3. Preliminaries

This section defines the layered Transformer and hybrid-model notation, then characterizes softmax-attention and modern RNN mixers. It also motivates attention-layer selection and positional-encoding choices for performance, efficiency, and length generalization.

  • Notation: Both Transformer and hybrid models are represented as stacks of L layers, with hybrid models specifying an attention-layer index set I_attn.I_attn contains the indices of the L_attn layers retained as attention layers.
  • Softmax Attention Layers: Transformer mixer layers use softmax attention with learnable query, key, value, and output projections together with an attention mask.The formulation uses row-vector representations and treats x^⊤x as an outer product.
  • Modern RNN Layers: Modern RNN layers are described through a transition matrix F_t, recurrent state S_t, update rule, and query rule.The transition matrix is a function of x_t and is typically diagonal or rank-1 for fast parallelization; the formulation includes Mamba2 and Gated DeltaNet.
  • Attention-Layer Selection: Selecting I_attn seeks maximum performance without increasing |I_attn|, because efficiency decreases as the number of retained attention layers grows.The motivation reflects prior findings that RNNs underperform attention models on recall-intensive tasks.
  • Position Encoding and Length Generalization: RoPE commonly improves Transformer language modeling, whereas attention without RoPE, or NoPE, provides superior training-free length generalization.RNNs are inherently position-aware through F_t, but their language-modeling performance and length generalization depend sensitively on update-rule structure and parameterization.

4. HALO: An Efficient Pipeline to Distill Transformers into Hybrids

HALO converts Transformer attention layers into RNN layers through weight transfer, staged distillation, and attention-layer selection. The resulting hybrid model is further trained with end-to-end distillation and long-context finetuning.

  • Pipeline overview: HALO combines attention-weight transfer, three training stages, and attention-layer selection to convert Transformers into hybrid models.The procedure adopts and improves RADLADS, which converts Transformers into pure RNN models.
  • Weight transfer: Pre-trained Wq, Wk, Wv, and Wo projection weights instantiate each RNN layer, while uncovered RNN modules use empirical initialization.The RNN layer inherits the attention layer’s configuration and available projection weights.
  • Stage 1 alignment: Each instantiated RNN layer is independently trained with MSE against its corresponding attention layer’s hidden states, while other weights remain frozen.After stage 1, each attention layer has a student RNN layer that can potentially replace it.
  • Attention layer selection: Attention layers are selected when replacing them with RNNs causes a large recall drop but a small CSR drop, using the resulting importance scores.The method evaluates models with individual attention layers replaced by their stage-1 RNN counterparts.
  • Attention layer selection: 25% of the final model’s layers are attention layers because HALO sets k = ⌊L/4⌋.The selected attention-layer indices are reported in Appendix C.
  • Stages 2–3: 1B training data is used for end-to-end knowledge distillation, followed by 1B training data for long-context finetuning at greater context length and lower learning rate.Stage 2 freezes the Transformer teacher and uses a cosine learning-rate schedule; the final finetuning stage optimizes the hybrid model’s capabilities.

5. HypeNet: An Effective Attention-RNN Hybrid Architecture

HypeNet combines HyPE with architectural modifications that improve performance and length generalization while remaining compatible with modern RNN mixers. Its design is more efficient for long contexts because it uses a much smaller KV cache despite being roughly 10% larger than the distilled model.

  • RNN Mixer Compatibility: HypeNet is compatible with modern RNN mixers that accept QKV inputs, including Lightning attention, Mamba2, GLA, GDN, and RWKV-7.Among mixers tested for Qwen3-1.7B, Lightning Attention provided the best balance between CSR and length generalization.
  • HyPE: HyPE applies RoPE in RNN layers and NoPE in attention layers, combining NoPE’s length generalization with RoPE’s positional information.The scheme is motivated by RNNs’ limited receptive field: RNN layers model short-distance dependencies while attention layers model long-distance dependencies.
  • Architectural Modifications: HypeNet adds QK-normalization to RNN layers, decouples GQA KV heads by cloning projection weights, and adds output gates to RNN and attention layers.These changes are intended to improve hybrid-model performance, preserve RNN expressivity, and provide consistent gains with little inference-cost increase.
  • Efficiency: Roughly 10% larger, HypeNet is much more efficient than its base model because its KV cache is much smaller.The increased model size results from converting GQA to MHA and introducing output gates; reduced KV size is described as more cost-effective for long-context scenarios.

6. Experiments

Experiments evaluate HypeNet + HALO against Qwen3 and prior Transformer-to-hybrid methods on CSR, long-context recall, architectural choices, layer selection, and inference efficiency. HypeNet + HALO delivers superior long-context performance and substantial efficiency gains, while HyPE, Lightning Attention, HALO modifications, and the proposed layer selection method are validated as effective components.

  • 6.1 Experimental Setup: Experiments evaluate Qwen3-derived 1.7B, 4B, and 8B models mainly on zero-shot CSR tasks and NIAH5 long-context recall.HALO training uses randomly sampled data from the 10B-token FineWeb-edu subset.
  • 6.2 Main Results: HypeNet + HALO achieves superior long-context performance versus Transformer-to-hybrid methods despite fewer training tokens, open-source data, and a smaller model than KL-LS (GDN).The comparison includes Qwen3 and state-of-the-art hybrids converted from pre-trained Transformer models.
  • 6.3 HypeNet Design Choices: HyPE substantially outperforms existing position encodings in length generalization, and Lightning Attention outperforms more recent RNN variants despite its simpler update rule.The reported findings identify HyPE and Lightning Attention as essential for HypeNet’s exceptional length generalization; data-independent forget gates are offered as one possible explanation for Lightning Attention’s advantage.
  • 6.4 HALO Architectural Ablations: HALO’s architectural modifications provide effective CSR and NIAH gains, considerably outperforming common approaches for training hybrid architectures.The ablation converts Qwen3-1.7B and evaluates modifications marked ➊, ➋, ➌, and ➍ in Figure 3.
  • 6.5 Attention Layer Selection: The proposed attention layer selection method achieves better overall CSR and recall than Jet-Nemotron, KL-LS, and naive evenly distributed baselines.The comparison replaces HALO’s selection method without running the substantially larger distillation procedures of Jet-Nemotron or KL-LS.
  • 6.6 Inference Efficiency: 3.0× decoding speedup and 3.4× prefilling speedup are achieved on 512K context length before Qwen3-1.7B runs out of GPU memory at 1M.Under 128K context length, HypeNet is more efficient than Qwen3 in memory and throughput, with the advantage increasing as context length grows.

7. Conclusion

The paper introduces HALO for converting pre-trained Transformers into RNN-attention hybrids with less than 3B tokens and HypeNet, whose HyPE encoding enables superior length generalization. Applied to Qwen3, these methods yield hybrid models with improved performance-throughput tradeoffs and memory efficiency for long contexts.

  • Conclusion: HALO converts pre-trained Transformer models into RNN-attention hybrid architectures using less than 3B tokens.It is presented as a novel distillation procedure.
  • Conclusion: HypeNet uses the novel HyPE position-encoding scheme and achieves superior length generalization.HypeNet is introduced as a hybrid architecture based on HyPE.
  • Conclusion: Qwen3 conversion produces hybrid models with much better performance-throughput tradeoffs and memory efficiency for long-context scenarios.The reported benefits concern long-context use cases.

Limitations · A. Complete Formulation of HypeNet

The paper identifies limitations in HALO’s preservation of post-training capabilities and its architecture-specific applicability, then gives HypeNet’s complete hybrid-layer formulation, including attention, RNN, multi-head, and forget-gate components.

  • Limitations: Conversion on FineWeb-Edu may diminish instruction-following and alignment behaviors introduced during post-training.The authors identify efficient recovery of the base model’s capabilities as an open question shared by existing hybrid-conversion distillation methods.
  • Limitations: The conversion protocol is designed for Transformer architectures, so applicability to other model architectures requires further investigation.The authors note that most publicly available LLMs are Transformer-based.
  • A. Complete Formulation of HypeNet: HypeNet consists of L stacked layers, each containing a token mixer and an MLP.The complete formulation is presented for clarity.
  • A. Complete Formulation of HypeNet: Each mixer is either an attention layer or an RNN layer, selected through an attention index set I_attn.Layer index l ranges from 1 to L, and Norm denotes RMSNorm.
  • A. Complete Formulation of HypeNet: The MLP matches the base model, while each mixer uses n_h identical heads except for GQA’s KV-sharing mechanism, with outputs summed across heads.The formulation omits the head index and specifies one head for simplicity.
  • A. Complete Formulation of HypeNet: Attention layers use learnable W_q, W_k, W_v, W_o, and W_z parameters with position-dependent scaling s_t and hyperparameter a.Depending on the base model, attention layers may also include QK-norm.
  • A. Complete Formulation of HypeNet: RNN layers use the same learnable projection parameters, with γ as Lightning Attention’s head-specific, data-independent forget gate and RoPE_t providing position rotation.The formulation is given per head, with the head index omitted for simplicity.
  • A. Complete Formulation of HypeNet: For H = 32 heads, γ_h is head-indexed, and HypeNet omits the original implementation’s layer-specific rescaling because preliminary hybrid-model results showed no performance gain.The forget gate is defined using head index h from 1 through H.

B. HALO Training Configurations … E.1. RNN Mixer Efficiency Measurement

The paper specifies HALO’s training and model configurations, evaluates architectural additions and layer selection, quantifies conversion-stage costs, and compares RNN-mixer efficiency across context lengths. The reported setup uses standardized optimization and precision choices, while Lightning Attention with data-independent forget gates is identified as especially fast.

  • B. HALO Training Configurations: HALO uses AdamW with β=(0.9, 0.95), no weight decay, 50-step linear learning-rate warmup, and BFloat16 precision.These defaults apply across the conversion stages.
  • B. HALO Training Configurations: Table 5 reports hyperparameters for each HALO training stage, including the model-dependent ηstage2 parameter.The model-dependent value is reported in Table 6.
  • C. HypeNet Model Configurations: Table 6 reports HypeNet model hyperparameters, while Table 7 gives the actual attention-layer indices for each HypeNet model.The configurations cover the models studied in the paper.
  • D. Addition Notes on the Model Architecture: Short convolution is a per-channel 1D convolution with a typically 2–4 kernel, but post-training addition gave no 8B performance gain and failed to converge for 1.7B.The paper notes that most Transformer models lack this component.
  • D. Addition Notes on the Model Architecture: Table 7 compares attention-layer selections by importance scores from different methods and identifies the top-k layers retained in the final model.Red boxed indices indicate layers not selected by the authors’ approach.
  • E. Computational Cost of Each Stage: 8.36M tokens are used by HALO’s layer-selection evaluation tasks, whose inference FLOPs per token are notably lower than training FLOPs per token.Stage 3 has higher FLOPs per token because it uses a longer context, while layer selection requires inference without backward passes.
  • E. Computational Cost of Each Stage: 700M × L + 600M tokens are required by KL-guided layer selection, where L is the number of layers in the base model.The paper compares this token requirement with its conversion experiments involving Qwen2.5-3B.
  • E.1. RNN Mixer Efficiency Measurement: Lightning Attention with time-independent forget gates is significantly faster than other RNN mixers and comparable to SWA with a 512-token window across context-length measurements.The comparison is measured on one NVIDIA A800-80GB GPU; Mamba2 runs out of CUDA memory at 256K context.

E.2. More Evaluation Details … G.2. A Note on Kimi Delta Attention

The paper specifies evaluation procedures and validates HyPE and HALO design choices through ablations. It also identifies compatible RNN mixers, including Mamba2, while reporting that Kimi Delta Attention conversion failed to converge.

  • E.2. More Evaluation Details: Evaluations use LM-Evaluation-Harness 0.4.10.dev0 through the HuggingFace API with BFloat16-exported checkpoints.Checkpoints are exported for AutoModelForCausalLM.from pretrained before evaluation.
  • E.2. More Evaluation Details: Qwen3 evaluations retain official configurations by default, applying YaRN rope scaling when tasks exceed the 40,960-token maximum context length.The long-context configuration adds a "rope scaling" entry based on the official model card.
  • E.2. More Evaluation Details: CSR performance is measured on ARC-Easy, ARC-Challenge, HellaSwag, WinoGrande, PIQA, LAMBADA, and MMLU using normalized accuracy by default.The listed downstream tasks define the CSR evaluation set, and normalized accuracy is the default metric.
  • F.1. Attention Logits Scaling Validation: Without attention-logit scaling, HyPE has limited length generalization; constant scaling improves it, while position-dependent scaling unlocks its full potential.The validated settings are s_t = 1, s_t = 1.5, and s_t = log_a(t + a), respectively.
  • F.1.1. HALO CONFIGURATION ABLATION EXPERIMENTS: Beyond 320M stage-1 training tokens, increasing HALO’s data does not strongly improve final performance.This finding comes from ablations of HALO’s stage-1 and stage-2 training configurations on Qwen3-1.7B.
  • F.1.1. HALO CONFIGURATION ABLATION EXPERIMENTS: The default constant learning rate from RADLADS is highly suboptimal for HALO stage 2.The discrepancy may reflect differences in network architecture or model size between RADLADS and this work.
  • G. Which RNN Mixers are Compatible with HypeNet?: HypeNet and HALO support a non-exhaustive set of representative RNN mixers expressible as the paper’s outer-product-based RNN formulation.Tables map each mixer’s notation to the paper’s notation and identify which components inherit attention weights in HALO.
  • G.1. HypeNet’s Compatibility with Mamba2: Both Mamba and Mamba2 are compatible with HypeNet, while multi-head Mamba2 is used to match softmax-attention weight structure with negligible throughput impact.Mamba2 can be expressed as gated linear attention, and multi-head Mamba2 resolves its shared-query-and-key multi-value configuration.

H. Training and Model Configurations for Training From Scratch Experiments · H.1. Training Configurations · H.2. Model Configurations

The from-scratch experiments use standardized training resources and hyperparameters, while controlling model size and specifying HypeNet’s hybrid layer pattern and RNN settings. The configurations are documented in Tables 12–15, with symbols indicating enabled, disabled, or inapplicable features.

  • H. Training and Model Configurations for Training From Scratch Experiments: The appendix describes training and model configurations for the experiments in Section 6.3.
  • H.1. Training Configurations: 20 billion tokens from FineWeb-edu train every model using 8 NVIDIA A800 GPUs and HuggingFace Accelerate.Specific training hyperparameters appear in Table 12 and were chosen to match standard LLM pre-training practice.
  • H.2. Model Configurations: Approximately 500M parameters are used for every model to ensure fair comparison.Implementations are kept as similar as possible to the respective authors’ official implementations.
  • H.2. Model Configurations: HypeNet models allocate 25% of layers to attention, interleaving one attention layer followed by three RNN layers.The repeating pattern is Attn →RNN →RNN →RNN, and each attention/RNN block is followed by a SwiGLU MLP block.
  • H.2. Model Configurations: Training from scratch avoids the attention-layer selection handling required by HALO.
  • H.2. Model Configurations: Table 15 specifies hyperparameters for the RNN layers in HypeNet variants used in the from-scratch experiments.
  • H.2. Model Configurations: ✓ denotes an enabled feature, ✗ a disabled feature, and “–” an inapplicable hyperparameter.
Loading 2601.22156v1…