Draft-model fine-tuning
×Stage · Kimi K3
The pre-trained MTP layer is fine-tuned into a draft model for speculative decoding, with the target model frozen.
- Trainable
- the draft layer and its projection
- Loss
- the LK loss
- Unrolling
- 7 steps
- Input
- features after the 1st, 4th and final blocks
Not disclosed Data · Steps · Learning rate
Why
The speed-up of speculative decoding is set by the acceptance rate; a KL surrogate does not guarantee the highest rate for a small draft model, so the loss is the negative logarithm of that rate itself. Kimi K3 report §4.1.4, p. 14
The training step →What one step of this stage computes, as a plate.
Architecture at this stageThe target model is frozen; the MTP layer and a projection of three of the target’s features train.
All settings6
- Trainable
- The draft layer (the pre-trained MTP layer) and its feature-fusion projection; the target model is frozen
- Input
- The target’s features after the 1st, 4th and final AttnRes blocks, concatenated and projected by a bias-free matrix initialized [0 0 I]
- Unrolling
- 7 steps in training; after the first, the draft reads its own earlier outputs
- Use
- Speculative decoding is used in RL rollouts, where it “further accelerates request turnover”
- Loss
- The LK loss: the negative logarithm of the acceptance rate Σ min(p, q) between the target’s and the draft’s next-token distributions, at temperature 1, with no cross-entropy term (eq. 16)
- Number format
- Quantization-aware training: MoE expert weights in MXFP4, activations in MXFP8; the other parts in higher precision
Sources3
- k3_tech_report.pdf, p. 14Trainable, Input, Unrolling, Loss, Number format
- k3_tech_report.pdf, p. 21Use
- k3_tech_report.pdf, p. 12Number format