💬 NLP

Stuck on "A": Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization of a 0.6B Language Model

This paper diagnoses and repairs a specific "interface injury" in a 0.6B language model converted to linear attention, where the model blindly predicts option labels (e.g., "A") rather than content, by using a permutation-based diagnostic and a targeted completion-stage KL distillation to restore multiple-choice accuracy while maintaining efficiency on consumer hardware.

Original authors: Ronglong Bao

Published 2026-08-05

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Ronglong Bao

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

Imagine you are trying to teach a brilliant but tiny robot how to read a book. In the world of artificial intelligence, these robots are called "language models." They work by looking at words one by one and guessing what comes next, kind of like a super-fast game of "fill in the blank." To do this, they usually keep a giant mental list of everything they've read so far, which takes up a lot of memory. But what if we could teach them to remember things in a much smaller, more efficient way, like a short-term memory loop? This is the promise of "linear attention," a new trick that lets these robots run on smaller computers, like the ones in your laptop or phone, instead of needing massive, expensive supercomputers.

Scientists have been trying to take big, powerful robots that already know a lot and "surgery" them to use this new, efficient memory trick. The big question is: when you swap out their old memory for the new, efficient one, do they forget how to think, or do they just get confused about how to show you what they know? This paper dives into that exact mystery. It asks: if we fix the robot's memory, does it still know the answers, or has the surgery broken the part of its brain that connects its knowledge to the words it speaks?

The Surgery That Broke the Voice Box

The researcher took a tiny, 0.6-billion-parameter AI model (think of it as a very smart, but small, student) and performed a major surgery. They replaced 21 out of its 28 brain layers with a new, efficient "linear attention" system called KDA. They did this on a single, standard computer graphics card—the kind a gamer might have—rather than a supercomputer.

After the surgery, they checked the robot's health using standard medical tests. The tests said the robot was doing great! Its "perplexity" (a measure of how surprised it is by the next word) was almost as good as the original. It seemed to understand the flow of language perfectly. But then, they asked it a simple multiple-choice question: "What is the capital of France? A) London, B) Paris, C) Berlin, D) Rome."

The robot failed. Miserably. It got the answer right only about 25% of the time, which is the same score you'd get if you just guessed randomly. The standard tests had lied. The robot knew the facts deep down, but it had lost the ability to point to the correct answer.

The "A" Addiction: A Diagnostic Detective Story

To figure out what was wrong, the team invented a clever detective trick called the "four-permutation diagnostic." Imagine you have a multiple-choice question with four options: A, B, C, and D. Usually, the answer is hidden behind one of these letters.

The researcher took 161 questions and shuffled the letters around. They asked the same question four times, but each time they rotated the options:

Original: A=London, B=Paris...
Rotated: A=Paris, B=Berlin...
Rotated again: A=Berlin, B=Rome...
And again.

If the robot was actually thinking, it should follow the content (the word "Paris") and pick the correct letter, no matter where it was placed. If it had forgotten everything, it would guess randomly every time.

But the robot did neither. It did something weird: it became obsessed with the letter "A".

It picked option "A" 81% of the time, regardless of whether "A" was the right answer or not.
On 106 out of 161 questions, it picked the exact same letter (usually "A") even when the options were shuffled completely.

The robot wasn't stupid; it was just "stuck." It had lost the interface—the mental bridge that connects its knowledge to the specific labels (A, B, C, D) it needs to output. It knew the answer was Paris, but it couldn't figure out which letter to press to say "Paris." It was like a person who knows the answer to a riddle but can only speak in a specific, broken code that always starts with the same word.

The Fix: Re-teaching the Format

The team realized that the robot had been trained on long paragraphs of text, but never on the specific format of "Question -> Options -> Answer." It didn't know how to play the game of multiple-choice.

To fix this, they ran a short, targeted training session. They didn't teach it new facts; they just showed it thousands of examples of the format of multiple-choice questions, specifically teaching it to look at the options and pick the right one. They used a technique called "completion-only KL," which is a fancy way of saying, "Don't worry about the question part, just learn how to finish the answer correctly."

The result was a dramatic recovery:

The robot's score on a standard test (C-Eval) jumped from 28.8% to 41.3%.
The "A addiction" was cut in half (dropping from 81% to 58%).
The robot started picking the correct option more often than the wrong ones.

The robot was no longer just guessing; it had relearned how to map its knowledge to the answer labels.

The Final Touch: Giving It a Personality

Once the robot could answer questions correctly, the team wanted to give it a personality. They wanted it to act like a specific digital assistant named "Qingyi." They did this by having the robot practice chatting and answering questions in the style of this persona.

Surprisingly, this didn't break the robot again. Usually, teaching a robot a new personality makes it forget its old skills (a problem called "catastrophic forgetting"). But because the robot's brain was already fixed and stable, it learned the new personality without losing its ability to answer questions. It ended up with a score of 41.83%, which was very close to its post-repair score, proving that you can give a robot a personality without breaking its brain.

What We Learned (and What We Didn't)

The paper concludes with a few important lessons for anyone trying to build these efficient AI models:

Don't trust the standard tests: Just because a robot looks healthy on paper (low perplexity) doesn't mean it can actually do the job. You have to test if it can follow instructions, not just predict words.
The "A" problem is real: When you change how an AI thinks, it can get stuck on simple habits like always picking the first option. This is a "interface injury," not a loss of knowledge.
Small computers can do big things: You can do this kind of complex surgery on a single consumer-grade GPU, but you have to be careful. The researcher found a hidden bug where the computer's math (using a format called bf16) was silently swallowing tiny updates, making the robot think it was learning when it was actually frozen. They had to switch to a different math format (FP32) to fix this.

The researcher released all their code, weights, and the "accident report" of everything that went wrong, so others can learn from their mistakes. They didn't solve everything—there's still a gap between their small robot and the big teacher model—but they proved that with the right diagnosis and a little bit of format training, you can fix the broken parts of an AI's brain.

Technical Summary: Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization

Problem Statement
The paper investigates the "lossy" nature of converting pretrained full-attention language models into linear attention architectures, specifically targeting the Kimi Delta Attention (KDA) mechanism. While linear attention offers structural advantages for long-context inference and edge deployment (O(1) recurrent state vs. O(n) KV cache), the conversion process often fails to preserve task performance despite standard metrics suggesting success.

The author focuses on converting Qwen3-0.6B-Base (28 layers) into a hybrid architecture where 21 layers are replaced by KDA and 7 layers retain Grouped Query Attention (GQA). This is executed under strict consumer-grade constraints: a single 32 GB GPU and a training budget of tens of millions of tokens (rather than billions).

The core diagnostic question is: What exactly does the conversion break? Standard evaluation reveals a paradox: after layer-wise alignment and end-to-end KL distillation, the student model achieves near-teacher perplexity (Cross-Entropy gap of +0.128 nats) but collapses in multiple-choice accuracy (staying near random chance at ~25–29% on C-Eval, compared to the teacher's 50.6%).

Methodology
The author proposes a multi-stage pipeline to diagnose and repair this performance gap:

Hybrid Architecture & Surgery:
- Layout: Layers 3, 7, 11, 15, 19, 23, and 27 retain native GQA; the remaining 21 layers are converted to KDA (16 heads, head-dim 128).
- KDA Mechanics: Implements a per-channel gated delta rule with a forget gate ( $\alpha_t$ ) and a write gate ( $\beta_t$ ).
- Initialization: A critical ablation study identified that simple projection transplants failed. The successful configuration (g6) uses specific decay scales and bias settings to stabilize the recurrent state.
Training Stages:
- Stage 2 (Layer-wise Alignment): Aligns KDA hidden states with the frozen teacher's hidden states using MSE. This reduced Cross-Entropy (CE) from 9.48 to 4.13 but left C-Eval accuracy at 25.5%.
- Stage 3a (End-to-End KL Distillation): Trained on pretraining text using forward KL divergence between teacher and student distributions. This closed the CE gap to +0.128 nats but only raised C-Eval accuracy to 28.8%.
- Stage 3b (Format-Targeted Repair): A 1,000-step "completion-only" KL stage using teacher-verified Multiple Choice Questions (MCQs), poetry, and translation QA. Crucially, this stage focused on the format of the answer rather than just general text generation.
The Four-Permutation Diagnostic:
To move beyond perplexity, the author designed a diagnostic protocol using 161 clean MCQs (zero overlap with benchmarks). Each question was evaluated under four cyclic rotations of the answer options (A, B, C, D).
- Metric: "Stickiness" measures how often the model predicts the same label (e.g., "A") regardless of the option's content position.
- Finding: The Stage 3a model exhibited severe interface injury: it predicted "A" 81% of the time and maintained the same label across all four rotations for 106/161 questions. The model had lost the ability to map knowledge to option labels, effectively ignoring the content of the choices.
Persona Alignment:
Post-repair, the model underwent Supervised Fine-Tuning (SFT) with an injected persona ("Qingyi") and one round of on-policy Direct Preference Optimization (DPO) to ensure the conversion did not degrade the model's identity or conversational abilities.

Key Results

The "Perplexity Lie": The study confirms that hidden-state alignment and low KL divergence do not guarantee task ability. The model can mimic the teacher's distribution statistics while failing to utilize that knowledge for specific tasks.
Diagnosis of Interface Injury: The four-permutation test revealed that the model was not suffering from uniform knowledge loss (random guessing) but from label-stickiness. It failed to track answer content, defaulting to a specific option label.
Repair Efficacy: The format-targeted Stage 3b repair:
- Increased C-Eval accuracy from 28.83% to 41.31%.
- Reduced "A" prediction share from 81.06% to 58.5%.
- Halved the "same label across 4 rotations" metric (106/161 $\to$ 47/161).
- Turned the correct-option margin from negative (−0.117 nats) to positive (+0.233 nats).
- Trade-off: This repair came at a measurable cost to general distribution fit (Valid CE/KL increased slightly), suggesting the model specialized in the MCQ format.
Final Performance: After persona alignment, the final model achieved 41.83% on C-Eval, remaining within the noise band of the pre-SFT repair stage, indicating no catastrophic forgetting.

Engineering Lessons & Failure Modes
The paper documents several critical engineering pitfalls encountered during the process:

FP32 Master Failure Mode: Using bf16 for both parameters and Adam optimizer moments caused small weight updates to be "silently swallowed" (sub-ULP updates), resulting in frozen weights despite healthy gradient norms. The fix was using FP32 master weights.
Cache Parity: New KDA layers initially ignored incremental cache updates, causing generation to collapse into garbage from the second token. A specific cache implementation fix resolved this without retraining.
False Transplants: Early attempts to transplant Q/K/V/O projections from the teacher failed functionally; the successful approach abandoned projection transplant entirely.

Significance and Claims
The paper claims four primary contributions:

Diagnostic Tool: A four-permutation protocol that separates "interface injury" (inability to map knowledge to output format) from "knowledge loss" (uniform forgetting), applicable to any linearization or compression pipeline.
Reproducible Recipe: A complete, budget-constrained recipe for converting Qwen3 to KDA, documenting specific failure modes (gate ablation, cache bugs, update swallowing) that are often omitted in public records.
Persona Survival: Evidence that persona alignment (SFT + DPO) can survive architecture conversion essentially "for free," without requiring massive retraining or causing catastrophic forgetting.
Open Release: Full release of code, weights, recipes, and the audit trail of failures.

The author concludes that while converting a 0.6B model to a KDA hybrid on a consumer budget is feasible, standard metric stacks (perplexity, KL, hidden-state MSE) systematically misreport the damage. The primary barrier to success is not necessarily the loss of internal knowledge, but the degradation of the interface that maps that knowledge to specific output formats (like multiple-choice labels), which can be repaired with targeted format supervision.