Verifying EAGLE-3 Training Correctness: The Architecture of Frozen and Trainable Parameters

In a single, concise message ([msg 4234]), an AI assistant confirmed the correctness of an EAGLE-3 draft model training configuration by comparing a monkey-patched implementation against the upstream speculators library. The message, delivered in response to a user's pointed question about parameter freezing, distills a critical architectural decision that determines whether speculative decoding training succeeds or fails. The user had asked: "Before starting train - are we correctly training the whole model, embeddings etc or is anything frozen?" ([msg 4231]). The assistant's answer, while brief, encapsulates the entire design philosophy of the EAGLE-3 algorithm and reveals the careful reasoning required to deploy it correctly.

The Context: Training an EAGLE-3 Drafter from Scratch

The conversation leading up to this message represents the culmination of an enormous engineering effort. The team had been building an EAGLE-3 speculative decoding drafter for the Kimi-K2.5 language model across multiple sessions spanning days. They had survived a VM crash and disk migration ([msg 4200]), extracted 37,312 hidden states totaling 87.8 million tokens and 4.6 terabytes of data ([msg 4224]), and were poised to launch a training run that would consume four GPUs for over ten hours.

At this critical juncture, the user paused to ask a fundamental question about the training configuration. It was not a trivial question. In speculative decoding, the relationship between the "drafter" (the small model that predicts multiple candidate tokens) and the "verifier" (the large base model that validates them) is delicate. If the wrong parameters are frozen, the drafter could learn to produce embeddings that diverge from what the verifier expects, causing the entire speculation mechanism to collapse. If too many parameters are frozen, the drafter may lack the capacity to learn useful predictions.

The Message: A Table and a Confirmation

The assistant's response is structured as a direct answer followed by a summary table:

Our monkey-patch matches exactly what the upstream speculators library does — freeze embed_tokens and verifier_lm_head, everything else trainable.

Summary: | Component | Params | Trainable? | |-----------|--------|------------| | embed_tokens | 1,174M | Frozen | | verifier_lm_head | 229M | Frozen | | lm_head (draft) | 229M | Yes | | fc (projection) | ~86M | Yes | | decoder layer (attn+mlp) | ~876M | Yes | | Total | 2,595M | 1,191M trainable |

This is correct per the EAGLE-3 design. Ready to launch training?

The message serves as a final sign-off — a declaration that the configuration has been verified against the authoritative upstream source and is ready for production.

Why This Verification Was Necessary

The need for this verification arose from a fundamental tension in the project's architecture. The team was using a monkey-patched version of the EAGLE-3 model initialization (seen in [msg 4232]), where the standard _setup_embeddings_and_lm_heads method was replaced with a custom version that loaded pre-trained weights directly rather than calling load_model_layers. This was necessary because the training script needed to work with the specific Kimi-K2.5 model's weight structure, which differed from the default loading path.

Monkey-patching is inherently risky. It bypasses the normal initialization logic, and there is always a danger that the replacement code subtly diverges from the intended behavior. The upstream speculators library explicitly sets requires_grad = False on embed_tokens and verifier_lm_head at lines 258 and 280 of its eagle3/core.py ([msg 4233]). If the monkey-patch failed to replicate this, the training would silently proceed with incorrect gradient flow — potentially wasting hours of compute before the error was discovered.

The assistant's verification process was methodical. First, it read the training script to inspect the monkey-patched initialization ([msg 4232]). Then it examined the upstream library source code on the remote machine ([msg 4233]). Finally, it compared the two to confirm they agreed. This three-step process — inspect local code, inspect upstream code, compare — is a textbook example of how to validate a critical configuration change.

The Architecture: Why Each Component Is Frozen or Trainable

The message's table tells a story about the EAGLE-3 algorithm's design. Understanding why each component is treated the way it is reveals deep principles of speculative decoding.

embed_tokens (1,174M parameters, frozen): The embedding table maps token IDs to dense vectors. In EAGLE-3, the drafter must produce hidden states that are compatible with the verifier's internal representations. If the drafter learned its own embedding space, the verifier would receive alien inputs and produce garbage logits. Freezing the embedding ensures that the drafter's input representation remains identical to what the verifier expects. This is non-negotiable — training the embedding would break the verifier compatibility guarantee.

verifier_lm_head (229M parameters, frozen): This is a copy of the verifier's language modeling head, used as a teacher signal in the loss function. The EAGLE-3 loss typically involves a KL divergence or cross-entropy between the drafter's predictions and the verifier's predictions. If this head were trainable, it would adapt to the drafter's distribution rather than representing the true verifier distribution, corrupting the training signal. Freezing it preserves it as a fixed reference point.

lm_head (draft) (229M parameters, trainable): The drafter's own language modeling head. While initialized from the verifier's weights, it needs to adapt because the drafter's hidden representations differ from the verifier's. The drafter's decoder layer produces outputs in a different subspace, and the lm_head must learn to map those to the vocabulary distribution.

fc projection (~86M parameters, trainable): This linear layer concatenates and projects hidden states from multiple layers into the drafter's working dimension. It is the primary mechanism by which the drafter learns to combine information from different verifier layers. Training it is essential.

Decoder layer (~876M parameters, trainable): The single transformer decoder layer (self-attention + MLP) is the core of what EAGLE-3 learns. It models the autoregressive dependencies between draft tokens. This is the largest trainable component and the one that does most of the learning.

Assumptions and Potential Pitfalls

The message makes several implicit assumptions worth examining. First, it assumes that the upstream speculators library's freezing pattern is correct for the Kimi-K2.5 model specifically. While EAGLE-3 is model-agnostic in principle, different base models have different architectural details. The monkey-patch was designed for the Kimi-K2.5's specific weight layout, and the assistant verified only that the freezing pattern matched — not that the weight loading itself was correct.

Second, the parameter counts are approximate. The fc projection is listed as "~86M" and the decoder layer as "~876M". These are estimates based on the model dimensions (hidden size 7168, intermediate size, etc.). The total of 2,595M parameters and 1,191M trainable parameters gives a rough sense of scale — about 46% of parameters are trainable — but exact counts would depend on the precise architecture configuration.

Third, the message assumes that the user is satisfied with the answer and ready to proceed. The final line "Ready to launch training?" is both a confirmation and a prompt. It signals that the assistant considers the verification complete and is awaiting the user's go-ahead.

The Thinking Process Revealed

The assistant's reasoning, visible across messages [msg 4232], [msg 4233], and [msg 4234], shows a clear pattern. When the user asked about frozen parameters, the assistant did not simply assert "it's correct." Instead, it:

Read the actual training script to inspect the monkey-patched initialization code.
Identified the specific lines where requires_grad was set to False for embed_tokens and verifier_lm_head.
Connected the code to the EAGLE-3 paper's design rationale, explaining why each component was frozen or trainable.
Verified against the upstream library by running a grep command on the remote machine to find the corresponding requires_grad lines in speculators/models/eagle3/core.py.
Confirmed the match and produced a summary table. This is a mature engineering approach. Rather than relying on memory or assumptions, the assistant traced the actual code paths and validated them against the authoritative source. The table format is particularly effective — it compresses a complex architectural discussion into a single glance, making it easy for the user to verify the configuration's correctness.

Input and Output Knowledge

To fully understand this message, a reader needs knowledge of several domains: the EAGLE-3 speculative decoding algorithm and its distinction from earlier EAGLE variants; transformer model architecture, particularly the roles of embeddings, language modeling heads, and decoder layers; the concept of parameter freezing and requires_grad in PyTorch; and the relationship between a drafter model and a verifier model in speculative decoding.

The message creates new knowledge by establishing that the specific combination of frozen and trainable parameters — as implemented in the monkey-patched training script — is correct for the Kimi-K2.5 EAGLE-3 drafter. It also provides a clear parameter count breakdown that can be used for estimating memory requirements and training throughput. Perhaps most importantly, it serves as a documented decision point: the user can look back at this message and know that the training configuration was explicitly verified before launch.

Conclusion

Message [msg 4234] is a small but critical node in a much larger engineering narrative. It represents the moment when a complex, multi-day effort to build an EAGLE-3 drafter was validated at its most vulnerable point — the configuration of which parameters to train. The assistant's methodical verification against the upstream library, clear tabular summary, and explicit confirmation of correctness provided the confidence needed to proceed with a ten-hour training run on four GPUs. In doing so, it demonstrated that effective AI assistance in software engineering is not about providing answers from memory, but about tracing code paths, comparing implementations, and producing clear evidence that a configuration is sound.