FlashMorph picks which layers of a pretrained Transformer to keep as full attention (versus swap for Linear attention) by jointly training one scalar gate per layer with a linearization penalty, cutting layer-selection cost to 20M tokens versus 234M–50B for prior methods.
You run a Qwen3-style LLM in production and long prompts are eating your GPU memory and prefill latency. One fix: convert most attention layers to a linear variant so the KV cache stops growing with context length, keeping only a few full-attention layers for the retrieval-heavy work. The catch is choosing which layers to keep. Prior conversion pipelines like HALO score each layer in isolation (replace one, measure the damage, rank), which is slow and misses the fact that layers cooperate. FlashMorph reframes the choice as one joint optimization over the whole stack.
Start with a pretrained Transformer. Train a linear-attention twin for every layer by matching hidden states (standard distillation, weights frozen after). Now every layer has two branches available: original full attention, and its linear replacement.
Add one learnable scalar gate per layer, α, initialized to 1 (full attention). The layer’s output is a convex blend: α·full + (1−α)·linear. Freeze everything except the L gates, so only L numbers are trainable. Train those gates on synthetic Needle-in-a-Haystack data (random passkeys buried in long documents) with two losses: match the original model’s hidden states at answer positions, plus a penalty equal to the sum of the gates, which pushes each α toward 0 unless that layer really needs full attention for retrieval. After training, keep the top-K layers by gate value as full attention, linearize the rest, then do standard logits distillation and long-context finetuning.
for layer in range(L):
alpha[layer] = 1.0 # start fully-attention
for step in range(250):
x = sample(synthetic_passkey_data)
h_mix = [alpha[l]*H_full[l] + (1-alpha[l])*H_lin[l] for l in range(L)]
loss = align(h_mix, H_full_teacher) + 0.1 * sum(alpha)
update(alpha) # only gates, not weights
full_layers = topk(alpha, K)
The trick is that all gates move together, so a layer’s learned importance already accounts for what its neighbors are doing.
The prevailing approach treats each layer’s importance as a property you can measure alone: perturb one layer, see how much quality drops, rank, pick. This paper argues the opposite: layer importance only exists relative to the rest of the hybrid configuration, so you should learn all layer choices simultaneously under a budget constraint. The evidence that carries the argument is not the headline benchmark score but the layer-cost table, where joint optimization matches or beats isolated scoring while using orders of magnitude less compute.
•
Selection cost collapses. On Qwen3-1.7B, FlashMorph uses 20M tokens and 2.1 GPU hours, versus 234M tokens / 15.4 GPU hours for HALO and 20B tokens / 1071.8 GPU hours for KL-LS. That’s a 7.3× to 1219.7× GPU-hour reduction, and the gap widens with model size.
•
Retrieval quality holds or improves. On Needle-in-a-Haystack with 1.7B backbone at 128K context, FlashMorph hits 98.2 on the harder NIAH-Single-2 variant versus 95.0 for HALO and 78.0 for PostNAS. On NIAH-Single-3 at 256K, it reaches 73.2 versus 52.8 (HALO) and 57.6 (PostNAS).
•
Recall-intensive tasks (SQuAD, FDA, SWDE) benefit most. On 0.6B with Gated DeltaNet (GDN), FlashMorph’s recall average is 62.1 versus 53.2 (HALO) and 60.6 (KL-LS), while commonsense reasoning stays within a point of the strongest baseline.
•
The synthetic-passkey supervision matters but isn’t the whole story. Swapping in generic language-modeling loss still beats prior methods; adding passkey supervision lifts RULER on GDN from 61.6 to 64.7. So joint optimization is the main mechanism; retrieval-flavored data is a booster.
•
Inference wins downstream. The resulting 3:1 linear:full hybrid gives 2.24× prefill speedup at 128K and 2.07× decode speedup at 512K on 1.7B, and runs 1M-token decode where the original Qwen3 goes out of memory.
•
Weak spot: MoE. On Qwen3-30B-A3B, FlashMorph loses to HALO on the harder NIAH variants and on recall tasks. The paper reports this without a clean explanation.
Reach for this when you’re taking a pretrained dense Transformer (Qwen3, Llama-family) into production with long-context requirements and you want a hybrid variant without paying the 20B-token search bill that KL-LS or PostNAS demands. The concrete workflow: distill an all-linear twin once, then spend a couple of GPU-hours training L scalar gates on synthetic passkey data, pick top-K, distill logits, finetune long context. You get a drop-in architecture with a fixed-size recurrent state on 75% of layers.
The paper builds on the flash-linear-attention library library and uses three linear-attention backbones (Lightning Attention, Gated Linear Attention (GLA), Gated DeltaNet (GDN)) so the recipe is not tied to one mixer. The paper does not mention a code release for FlashMorph itself; the complete selected-layer indices for Qwen3-0.6B/1.7B/8B/30B are published in the appendix, so you can reproduce the exact hybrid without re-running selection. Evaluations use lm-evaluation-harness.
When you’re choosing which layers to keep versus replace under a budget, learn the choice jointly. Ranking layers one at a time bakes in an independence assumption that isn’t true.
•
The MoE result (Qwen3-30B-A3B) is genuinely worse than HALO on harder retrieval and recall tasks. If your target is a large MoE, this method’s advantage is not established.
•
All experiments use Qwen3 backbones from 0.6B to 30B. Generalization to other model families (Llama, Mistral) is untested in the paper.
•
The gates train on synthetic passkey retrieval. If your production workload’s long-context pattern differs from key-value lookup (say, aggregation or reasoning over long inputs), the selection signal may not transfer, though the LM-supervision ablation shows the method still works, just weaker.