PaperlayerView exact revision on arXiv

Source-linked AI summary

Effective Distillation to Hybrid xLSTM Architectures

Lukas Hauzenberger, Niklas Schmidinger, Thomas Schmied, Anamaria-Roberta Hartl, David Stap, Pieter-Jan Hoedt, Maximilian Beck, Sebastian Böck, Günter Klambauer, Sepp Hochreiter

arXiv:2603.15590v2cs.LG

TL;DR

Quadratic attention makes LLM deployment expensive, while existing sub-quadratic distillations often fail to recover teacher performance across tasks. The paper combines mLSTM and sliding-window attention in an xLSTM distillation pipeline, adds expert merging, and finds that merged students close generation gaps and often recover or exceed teacher performance.

  • Problem

    Existing sub-quadratic distillations often fail to match teacher LLMs across diverse downstream tasks, limiting their use as efficient replacements.

  • Method

    The paper replaces full softmax attention with a gated mLSTM–sliding-window-attention hybrid and merges independently distilled domain-specialized students.

  • Results

    The merged distilled xLSTM students close the performance gap on free-form generation and consistently dominate existing linearization methods across Cα tolerance levels.

  • Takeaways & Limitations

    The results support modular linearization as a path toward more efficient replacements for Transformer-based models.

  • Takeaways & Limitations

    Remaining deficits are most visible on synthetic long-context and selected reasoning benchmarks, where expert interference can reduce recovery after merging.

Abstract

from arXiv · show

There have been numerous attempts to distill quadratic attention-based large language models (LLMs) into sub-quadratic linearized architectures. However, despite extensive research, such distilled models often fail to match the performance of their teacher LLMs on various downstream tasks. We set out the goal of lossless distillation, which we define in terms of tolerance-corrected Win-and-Tie rates between student and teacher on sets of tasks. To this end, we introduce an effective distillation pipeline for xLSTM-based students. We propose an additional merging stage, where individually linearized experts are combined into a single model. We show the effectiveness of this pipeline by distilling base and instruction-tuned models from the Llama, Qwen, and Olmo families. In many settings, our xLSTM-based students recover most of the teacher's performance, and even exceed it on some downstream tasks. Our contributions are an important step towards more energy-efficient and cost-effective replacements for transformer-based LLMs.

1 Introduction

Transformer LLMs are costly because quadratic attention scales poorly with context length, motivating sub-quadratic distillation. The paper introduces a hybrid xLSTM pipeline with expert merging that improves teacher recovery, especially on free-form generation.

  • Motivation: Quadratic attention makes current LLMs energy-intensive and costly to deploy, motivating sub-quadratic replacements.Sub-quadratic students are attractive when they can match teachers across diverse tasks.
  • Prior work: Existing linearization methods use sub-quadratic sequence mixers and hidden-state or logit alignment, but have not achieved effective distillation.These methods are far more token-efficient than training comparable models from scratch.
  • Approach: xLSTM is hybridized with sparse attention by combining mLSTM, sliding-window attention, and sink tokens through learned gates.The design combines recurrent linear processing with local attention.
  • Approach: The pipeline independently distills domain-specialized students and consolidates them through weight-space merging.This makes linearization modular while combining specialized experts into one model.
  • Results: The merged distilled xLSTM students close the free-form generation gap and consistently dominate existing linearization methods across Cα tolerance levels.Cα is the paper’s Win-and-Tie reliability criterion for teacher-level recovery across benchmarks.

2 Background

Softmax attention stores and compares growing key-value histories, creating quadratic training costs and increasing decoding bandwidth demands. Linear attention and sliding-window attention reduce these costs, while mLSTM adds gated state updates for more flexible memory control.

  • Softmax attention: Softmax attention projects inputs into queries, keys, and values, then reads stored values using attention scores over prior positions.The key-value cache grows by concatenating each new key and value along the time dimension.
  • Softmax attention: Full softmax attention requires O(T^2) training time, while autoregressive decoding incurs O(t) readout time and cache size at step t.The growing cache also creates an O(t) memory-bandwidth cost during decoding.
  • Sparse attention: Sliding-window attention restricts each query to a fixed-length history, keeping its cache length at most W and reducing training or prefill complexity to O(TW).During decoding, its computational and memory complexities are independent of global sequence length.
  • Linear attention: Linear attention factorizes the kernel with a finite-dimensional feature map, enabling chunkwise-parallel training and recurrent decoding.The recurrent form maintains a per-head key-value state with an optional normalizer.
  • mLSTM: mLSTM augments linear attention with input, forget, and output gates that control writes, state decay, and readout modulation.The input gate sets new key-value write strength, the forget gate decays accumulated state, and the output gate modulates retrieved values.

3 xLSTM distillation pipeline

The pipeline replaces each Transformer self-attention block with a hybrid mLSTM–SWA module, then trains students through staged alignment and distillation, with optional expert merging. The architecture combines local context modeling from SWA with the linear complexity of mLSTM and dynamically mixes both outputs.

  • Architecture & student initialization: The core method replaces standard self-attention with a sub-quadratic hybrid block combining SWA and mLSTM through data-dependent gating.
  • Architecture & student initialization: SWA captures local context while mLSTM provides linear-complexity sequence mixing, and both branches operate in parallel before dynamic fusion.
  • Architecture & student initialization: The hybrid uses SWA over recent tokens plus four initial sink tokens, while a sigmoid-bounded per-head output gate modulates the global mLSTM branch against local outputs.
  • Linearization fine-tuning: Stage I aligns student layer representations to teacher attention outputs with mean-squared error while freezing teacher embeddings and MLP weights.
  • Linearization fine-tuning: Stage II unfreezes all student parameters and interpolates next-token prediction with sparse top-k KL matching to the teacher distribution.
  • Expert merging: An optional Stage III trains domain experts from a shared initialization and merges them into one student using weighted linear parameter averaging.

4 Experiments

The experiments apply the linearization protocol across base and instruction-tuned Llama, Qwen, and Olmo models. They evaluate hybrid students against teachers and alternative linearization methods on understanding, knowledge, generation, and reasoning benchmarks.

  • Models: The study covers base and instruction-tuned models from the Llama, Qwen, and Olmo families.
  • Benchmarks: Downstream evaluation spans language understanding and knowledge tasks alongside language generation and reasoning tasks.
  • Comparisons: The distilled xLSTM students are compared with their teacher models and state-of-the-art linearization alternatives.

Appendix E.3 for details). For mathematical evaluations, we use the Math-Verify evaluation system.

The paper evaluates xLSTM students using teacher-recovery and tolerance-corrected Win-and-Tie metrics across understanding, generation, and reasoning tasks. The pipeline combines an mLSTM–SWA hybrid, distillation, and optional expert merging to improve performance while reducing inference cost.

  • Evaluation metrics: Teacher-recovery rate is the student-to-teacher performance ratio, while Cα measures the fraction of benchmarks where the student matches or exceeds the teacher within tolerance α.α∗ is the minimum tolerance required for Cα ≥ 0.5.
  • Downstream evaluation: Across generation tasks, the xLSTM-based student matches the teacher on language understanding and exceeds the teacher on four considered generation tasks.Figure 3 reports recovery rates relative to teacher scores, with 1.0 denoting parity.
  • Downstream evaluation: Our xLSTM students achieve full or near-full teacher parity on language understanding tasks, while prior linearization baselines show significant performance gaps.The xLSTM-Llama3.1-8B student reaches full parity and xLSTM-Olmo3-7B reaches near-full parity.
  • Expert merging: Instruction-tuned merged students match teachers on language understanding and code generation, recover most mathematical reasoning performance, and achieve α⋆= 0.02 and α⋆= 0.05 for Llama and Qwen respectively.The largest remaining gap appears in STEM reasoning, where the merged student underperforms the dedicated STEM expert.
  • Expert merging: Merging independently trained experts improves instruction-following and preserves math and code capabilities, but STEM evaluations show interference between domain updates.The effect is strongest for Llama instruction-following, while Qwen instruction-following changes comparatively little.

B.1 Modern Recurrent and Hybrid Architectures.

Modern hybrid architectures combine sub-quadratic sequence mixers with attention or state-space components, while post-training linearization replaces softmax attention in pretrained models. This work uses an intra-layer mLSTM–SWA design with gated mixing and extends linearization through expert merging.

  • Modern Recurrent and Hybrid Architectures.: Sub-quadratic sequence operators target linear-time training and constant-memory decoding as alternatives to softmax attention.The alternatives include state-space models, data-dependent convolutions, and gated linear-attention families such as RWKV and xLSTM.
  • Modern Recurrent and Hybrid Architectures.: Hybrid models are categorized as inter-layer designs that alternate blocks or intra-layer designs that fuse attention and linear branches within each block.Intra-layer branches may be mixed through addition or learned gates, with head-wise, sequence-wise, and synchronous variants.
  • Modern Recurrent and Hybrid Architectures.: The proposed method places sliding-window attention and mLSTM in every layer, allowing each head to model local and global dependencies through both branches.The branches are synchronously combined rather than assigned separate roles across heads, and learned gates provide data-dependent mixing.
  • Modern Recurrent and Hybrid Architectures.: Post-training linearization replaces some or all softmax attention layers in a pretrained model with linear-time mixers, generally reusing teacher weights for data efficiency.The pipeline copies compatible weights, calibrates new gating and head-projection parameters with MSE, and may use full-model updates or low-rank adapters.
  • Modern Recurrent and Hybrid Architectures.: Expert-merging pipelines train capability- or domain-specialized variants and consolidate them into one deployable model through parameter averaging or interference-aware methods.The paper positions linear expert merging within broader branch-train-merge and weight-space consolidation work.
  • Modern Recurrent and Hybrid Architectures.: SWA reduces training and prefill complexity to O(TW), but deep stacks can have an effective receptive field that grows sublinearly and favors recent tokens.Sink tokens are also used to preserve stability and support KV-cache compression.

D.1 Hyperparameters and Data Mixes

The distillation setup uses a fixed xLSTM student configuration, staged hidden-state and end-to-end distillation, and domain-specific mixtures for specialist training.

  • D.1 Hyperparameters and Data Mixes: The student uses an xLSTM [0:1] mLSTM-only configuration with 32 heads, head dimension 128, rotary embeddings, and context length 4096.Optimization uses AdamW with weight decay 0.1 and global batch size 64.
  • D.1 Hyperparameters and Data Mixes: Phase I performs layer-wise hidden-state alignment with MSE while focusing optimization on newly introduced mixer and gating parameters.The phase uses a cosine learning-rate schedule with peak learning rate 10^-2.
  • D.1 Hyperparameters and Data Mixes: The data mixtures cover math, code, STEM, chat or instruction-following, and a multi-task mixture for a generalist student.Most mixtures use Nemotron Nano datasets, with domain-targeted additions such as OpenCodeInstruct and Dolci-Think-SFT.
  • D.1 Hyperparameters and Data Mixes: The STEM mixture differs by model family: Llama uses only Nemotron STEM Post-Training data, whereas Qwen adds specialized STEM and MCQ traces.The added Qwen data is intended to minimize the gap to the teacher.

D.2 Generalist vs. Merged Expert Student

Training separate domain experts and linearly merging them produces stronger specialized students than training one generalist on the same total token budget. The advantage appears across math, code, reasoning, and instruction-following evaluations.

  • D.2 Generalist vs. Merged Expert Student: The merged-expert setup trains four domain specialists for approximately 5B tokens each, then linearly merges them; the generalist uses the same 20B-token budget.Architecture and optimization are matched between the specialist and generalist settings.
  • D.2 Generalist vs. Merged Expert Student: The merged Qwen student exceeds the generalist on MATH 0.66 vs. 0.62, MATH Level 5 0.42 vs. 0.36, GPQA-D 0.26 vs. 0.19, and MT-Bench 5.96 vs. 5.55.The comparison shows the same direction across the listed specialized benchmarks.
  • D.2 Generalist vs. Merged Expert Student: The authors attribute the gap to multi-task distillation allocating fewer updates to specialized traces than single-domain expert training.They report that linear merging retains much of the specialists’ domain specialization in the merged checkpoint.

E Additional Results

The supplied passages for this section identify supplementary experimental material and a figure about attention sinks in the Llama 3.1 8B teacher model.

  • E Additional Results: The section provides additional experimental details and results intended to complement the main text.The supplied passage does not report a specific result from these additional materials.
  • E Additional Results: Figure 7 illustrates attention sinks in the Llama 3.1 8B teacher model.No further comparison or outcome is stated in the supplied figure passage.

E.1 Sink Analysis

The hybrid student models attention-sink behavior alongside sliding-window attention through a data-dependent gate. Sink modeling lowers loss and improves downstream performance, while both mLSTM and SWA contribute across the network.

  • Sink patterns are modeled together with sliding-window attention because this combination yields lower loss and stronger downstream performance.
  • The student combines mLSTM and SWA outputs through a data-dependent output gate analyzed as a layer-by-head gating matrix.
  • The gate analysis uses 128 sequences of length 4096 to measure component contributions across layers and attention heads.
  • mLSTM dominates the first two layers, SWA predominates in layers 3–16, and both components contribute more evenly in layers 17–32.

E.3 Downstream Evaluations

The evaluation covers language understanding, generation and reasoning, MT-Bench quality, and long-context retrieval under matched teacher–student settings. The distilled students generally approach or match teacher performance, sometimes exceeding it, but degrade on longer Needle-in-a-Haystack contexts.

  • Evaluations cover language understanding, generation and reasoning, MT-Bench quality, and Needle-in-a-Haystack long-context retrieval.
  • Matched few-shot examples, context lengths, and generation budgets provide consistent comparisons across teacher and student models.
  • Language Understanding & Knowledge: Our distilled students match or slightly exceed teacher performance on language understanding tasks, while alternative linearization recipes show significant gaps.
  • Language Generation & Reasoning: Our distilled students almost match teacher performance on language generation and reasoning tasks, outperforming alternative distilled models.
  • MT-Bench: Both xLSTM-Qwen2.5-7B-IT and xLSTM-Llama3.1-8B-IT outperform their teachers across all 7 MT-Bench categories.
  • Needle in a Haystack: Instruction-tuned baselines retain near-perfect Needle-in-a-Haystack recall, whereas distilled students degrade as context length increases from 4k to 16k tokens.

E.4 Inference Time Analysis

Inference tests separate prompt prefilling from autoregressive generation and vary batch size, context length, and generation budget. The hybrid student gains latency, memory, and throughput advantages as sequence length and workload increase.

  • Inference tests separate prefilling, which populates the KV cache, from autoregressive generation starting from that cache.
  • The experiments vary batch size B, context length C, and generation budget G across prefilling and generation workloads.
  • Latency: The hybrid student’s latency rises only slightly from batch size 1 to 8 because recurrent mLSTM inference and a fixed 256-token window avoid sequence-length-dependent computation.
  • Latency: The Transformer teacher’s computational demand grows faster with sequence length because of quadratic attention, while teacher measurements may be absent because of OOM.
  • Throughput: The hybrid student reaches up to almost 4× the Transformer teacher’s generation throughput as sequence length increases for a fixed generation budget of G = 100 tokens.
  • Throughput: Throughput comparisons are limited by out-of-memory failures, with B = 4 shown only up to C = 32K and B = 8 only up to C = 16K.

F Ablations

Ablations examine hybrid components, loss weighting, and training strategies. They support combining mLSTM, SWA, and sink tokens, balancing CE and KL losses, and using hidden-state alignment followed by full finetuning.

  • Component Ablation: Replacing linear attention with gated mLSTM, then adding SWA and sink tokens, produces successive validation-CE reductions.
  • Loss Weighting: Increasing KL weight raises validation CE and causes under-adaptation, whereas β approaching 0 lowers CE but lets the student drift from the teacher.
  • Loss Weighting: A CE/KL weighting of γ = 0.9 and β = 0.1 provides a good tradeoff between student performance and teacher alignment.
  • Training Strategy: Full finetuning achieves comparable KL while delivering superior downstream performance relative to other training strategies.
  • Training Strategy: Layer-wise hidden-state alignment is necessary but insufficient, while following it with full finetuning consistently outperforms standard finetuning under the same budget.
Loading 2603.15590v2…