RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale
Original authors: Daniel Goldstein, Eric Alcaide, Janna Lu, Eugene Cheah
Original authors: Daniel Goldstein, Eric Alcaide, Janna Lu, Eugene Cheah
Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer
Technical Summary: RADLADS
Problem Statement
While linear attention variants (such as RWKV and Mamba) offer significant advantages over traditional softmax attention transformers—specifically O(1) inference time per token and constant memory usage by avoiding Key-Value caches—training large-scale linear models from scratch remains prohibitively expensive. State-of-the-art (SoTA) transformer models often require training on over ten trillion tokens, a cost inaccessible to most organizations. Existing methods for converting pre-trained softmax transformers into linear or recurrent models (e.g., T2R, SUPRA, LoLCats, MOHAWK) have historically required massive token counts (ranging from 20B to 100B+ tokens) or resulted in poor downstream performance, particularly on benchmarks like MMLU. Furthermore, many conversion attempts rely on hybrid architectures (retaining some softmax attention) or fail to match the quality of the original teacher models.
Methodology
The authors present RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale), a protocol designed to convert softmax attention transformers into linear attention decoder models with minimal computational cost. The process involves three primary stages and specific architectural adaptations:
1. The RADLADS Protocol
The conversion is a three-step process requiring only 350–700 million tokens (less than 0.005% of the teacher's pre-training data):
- Setup (Attention Weights Transfer): Attention-related weights (Wq,Wk,Wv,Wo) from the teacher model are transferred directly to the student model's sequence mixing layers. Other weights are initialized via standard pre-training methods or set to mimic the teacher's behavior without immediate effect.
- Step 1: Attention Hidden State Alignment: The student model's sequence mixing layers are trained to approximate the hidden state outputs of the corresponding teacher attention layers. This is performed using an L2 distance (or MSE) objective. The authors found that using a Gated Linear Attention kernel (removing off-by-one decay and bonus terms) allows for a closer fit to the teacher's hidden states than standard RWKV-6.
- Step 2: Knowledge Distillation: The entire student model is trained to approximate the teacher's output logits using Kullback-Leibler (KL) divergence loss. The authors hypothesize that factual knowledge resides primarily in the teacher's MLPs and embeddings; therefore, learning rates for these components are kept low or fixed to prevent catastrophic forgetting, while the sequence mixer is trained more aggressively.
- Step 3: Context Length Extension (Optional): The model is fine-tuned on longer sequences (up to 16k tokens) using standard cross-entropy loss without a teacher model to enhance long-context capabilities. Alternatively, the authors propose Step 2a, which extends the sequence length to 4096 during the distillation phase itself, potentially obviating the need for a separate Step 3.
2. New Architectures
The authors identified that standard RWKV architectures were not always optimal for conversion. They introduced two new variants:
- RAD-RWKV6 ("RADFinch"): A modification of RWKV-6 using a Gated Linear Attention kernel and state balancing techniques to improve stability and fit.
- RAD-RWKV7 ("RADGoose"): A modification of RWKV-7 that removes the "tokenshift" mechanism (which offered no benefit in this context) and applies Rotary Positional Embeddings (RoPE) directly. This architecture demonstrated faster convergence and lower distillation loss compared to unmodified RWKV-7.
3. Hyperparameters and Data
- Dataset: The authors settled on DCLM (DataComp-LM) for all conversion steps, finding it superior to FineWeb or FineWeb-Edu for converting Qwen models.
- Learning Rates: A cosine annealing schedule is used in Step 1, starting high (10−3) to align hidden states and ending near the teacher's final pre-training learning rate (10−5). Steps 2 and 3 use a flat learning rate.
Key Contributions
- RADLADS Distillation Recipe: A detailed, step-by-step protocol including specific hyperparameters, token counts, and dataset choices that enables high-quality conversion with minimal data.
- New Architectures: The introduction of RAD-RWKV6 and RAD-RWKV7, which are optimized for the conversion process, yielding faster inference and better alignment with teacher models than their unmodified predecessors.
- Large-Scale Models: The successful conversion of popular open-source Qwen2.5 models into linear attention variants at 7B, 32B, and 72B parameter scales.
- Open Source Release: The release of code and converted models (QRWKV6/7-7B/32B/72B) under the Apache 2.0 license (with Qwen License restrictions for the 72B model), allowing others to replicate the process.
Results
The converted models achieve state-of-the-art performance among pure recurrent models of their size:
- Efficiency: Converting a 72B model costs less than $2,000 USD and requires only ~700M tokens.
- Performance: On standard benchmarks (Lambada, MMLU, ARC, etc.), the RADLADS models consistently outperform other conversion methods (e.g., SUPRA, LoLCats, MOHAWK, ARWKV).
- The QRWKV7-7B-Instruct achieves a relative MMLU score of 92.4% compared to its teacher, significantly outperforming other 7B conversions.
- The QRWKV6-72B-Instruct achieves a relative MMLU score of 89.9%, setting a new SoTA for pure RNN language models at that scale.
- Inference Speed: Due to the linear attention mechanism, the converted models show significant speedups over the teacher models as context length increases. For example, at 8k input tokens and 256 output tokens, the 32B QRWKV7 model is 1.61x faster than the Qwen3-32B teacher; at 6k input and 2k output, the speedup reaches 3.41x.
Limitations and Claims
The paper modestly acknowledges several limitations:
- Reasoning and Long Context: While performance is strong on standard benchmarks, the models show limitations in complex reasoning tasks (e.g., Minerva Math) and very long-context tasks (e.g., RULER), where they do not improve as effectively with additional output tokens as the teacher models.
- Architecture Sensitivity: Each new architecture design requires meticulous testing to ensure compatibility with the RADLADS protocol. For instance, GroupNorm/LayerNorm, common in some linear attention variants, caused training instability at 14B+ scales and had to be replaced with pre-scaling techniques.
- Dataset Alignment: The authors note that reasoning datasets may need to be more closely aligned with the teacher model's distribution to prevent repetitive looping behaviors, a challenge they addressed partially by mixing DCLM with OpenThoughts.
Significance
The paper claims that RADLADS provides a cost-effective pathway to democratize access to large-scale linear attention models. By reducing the training token requirement from trillions to hundreds of millions, it enables researchers and smaller organizations to test, train, and deploy new varieties of RNN architectures at scale without the extreme costs associated with pre-training. The authors position this as a tool to accelerate the development of the next generation of compressive state attention variants.
Drowning in papers in your field?
Get daily digests of the most novel papers matching your research keywords — with technical summaries, in your language.
Get the best AI papers every week.
Trusted by researchers at Stanford, Cambridge, and the French Academy of Sciences.
Check your inbox to confirm your subscription.
Something went wrong. Try again?
No spam, unsubscribe anytime.