The Moment of Correction: Reverting the DFlash Drafter's FC Architecture

In the middle of a grueling investigation into why a DFlash speculative decoding drafter was plateauing at a fraction of its expected performance, a single edit to dflash_model.py marked the turning point. The message is deceptively brief:

Fix 1+3: Revert fc back to (N-1) layers for context, separate last layer for clean target logits. The pipeline splits again, but noise only applies to the fc input: [edit] /data/dflash/scripts/dflash_model.py Edit applied successfully.

This short message — a single edit with a one-line summary — represents the culmination of a deep forensic analysis spanning dozens of messages, multiple evaluation harnesses, and a side-by-side comparison against a reference model. It is the moment when the assistant stops investigating and starts fixing.

The Investigation That Led Here

To understand why this edit matters, one must trace the path that led to it. The assistant had been training a DFlash drafter — a lightweight draft model used in speculative decoding to accelerate inference of a large target model (Qwen3.6-27B). Despite matching the paper's architecture and hyperparameters, the drafter was achieving a DDTree-8 acceptance rate (τ) of only about 3.0 after 20,000 training steps, while a reference model from z-lab achieved τ≈12.4 on the same prompts — a 4x gap.

The assistant built a comprehensive evaluation harness on the CT129 server, loading the target model with the fla library for correct linear attention, extracting hidden states from fresh coding prompts, and running drafter inference. A critical early discovery was that CPU-based hidden state extraction using PyTorch's fallback for linear attention produced numerically different results from the fla-based extraction used during training — 4 of 5 target layers use Qwen3.5's linear attention, and the bf16 numerical differences caused completely garbled drafter output. Switching to GPU extraction with fla fixed this, revealing the model's true but disappointing performance.

The assistant then turned to the official speculators repository — the reference implementation of DFlash — and performed a detailed code comparison. This comparison, executed as a subagent task ([msg 9123]), revealed three critical discrepancies between the official implementation and the assistant's own code.

The Three Critical Bugs

Bug 1: Noise corrupting target logits. The training pipeline captured hidden states from 5 target layers, concatenated them into a single tensor, and then applied Gaussian noise for regularization. The drafter then extracted the last layer from this already-noised tensor for target logit computation. This meant the training signal itself was corrupted by noise — the model was being asked to predict tokens from a distribution that had been deliberately perturbed. The official code applies noise only to the fc-input hidden states, never to the verifier hidden states.

Bug 2: Loss function mismatch. The official DFlash uses pure hard cross-entropy loss with gamma=4.0. The assistant's implementation used a composite loss: 70% soft KL divergence (temperature 2.0) + 30% hard CE + streak-aware weighting + gamma=10. The soft KL divergence forces the model to match the full 248K-dimensional vocabulary distribution instead of just getting the top-1 token correct — which is all that matters for speculative decoding acceptance. This massively diluted the gradient signal.

Bug 3: FC shortcut including the target layer. The official DFlash uses (N-1) layers for the fc projection (which conditions the drafter's KV cache) and reserves the last layer exclusively for target logit computation. The assistant's implementation fed all N layers to fc and used the last layer for target logits, creating a shortcut: the same information appeared in both the drafter's conditioning and the loss target, allowing the model to cheat rather than learn genuine next-token prediction.

Why Bugs 1 and 3 Are Treated Together

The subject message addresses bugs 1 and 3 in a single edit because they are structurally intertwined. Both involve how the 5-layer hidden state tensor is split between the fc projection (for conditioning the drafter's attention) and the verifier head (for computing target logits). The fix is conceptually simple but architecturally critical: revert the fc to use only 4 layers (N-1), keep the last layer exclusively for clean target logit computation, and ensure noise is applied only to the fc input.

The edit to dflash_model.py changes the data flow at a fundamental level. Previously, the code at line 699 read:

last_target_hidden = all_hidden_states[:, :, -H:]

This extracted the last layer from the noised tensor. After the fix, the pipeline splits the hidden states before noise application: 4 layers are routed to the fc projection (where noise is applied for regularization), and the last layer is kept completely clean for target logit computation. The fc projection no longer sees the target layer, eliminating the shortcut.

The Deeper Significance

This message represents a moment of intellectual humility in the development process. Earlier in the session, the assistant had implemented several "improvements" over the paper's architecture — soft KL loss, streak-aware dynamic weighting, a higher gamma value, and a cosine-annealed noise schedule. These were introduced with the reasonable belief that they would accelerate training or improve generalization. The investigation revealed that these modifications were actually harmful — they diluted the gradient signal, corrupted the training targets, and created architectural shortcuts.

The decision to revert to the official architecture was not made lightly. The assistant had to overcome the sunk cost fallacy — the investment in implementing and tuning these "improvements" — and accept that the paper's authors had already solved these design choices correctly. As the assistant noted in their reasoning, the official code uses pure hard CE with gamma=4.0, no soft KL, no streak weighting, and a clean separation between fc and verifier layers. Matching this exactly became the recommended path.

What This Fix Enables

With this edit applied, the training pipeline now mirrors the official DFlash implementation. The fc projection receives 4 layers of clean (but noised for regularization) hidden states, conditioning the drafter on rich contextual information without creating a shortcut. The verifier head receives the last target layer's hidden states without noise corruption, providing a clean training signal. The loss function (to be fixed separately in bug 2) will switch to pure hard CE with gamma=7.0 (the paper's recommended value for block_size=16).

The corrected v5 training run — launched immediately after these fixes — showed higher loss values (expected for hard CE versus soft KL) but comparable accuracy to previous runs at the same step count. The early metrics were promising, suggesting that the architecture was now learning correctly rather than plateauing due to corrupted signals.

Conclusion

Message [msg 9132] is a quiet but pivotal moment in the conversation. It is the point where diagnosis ends and treatment begins. The single edit to dflash_model.py — reverting the fc to (N-1) layers and separating the last layer for clean target logits — addresses two of the three critical bugs that were silently crippling the DFlash drafter's training. It represents a return to architectural principles validated by the paper's authors, a correction of well-intentioned but harmful modifications, and the foundation for a training run that could finally close the 4x performance gap against the reference model.