Distilling Token-Trained Models into Byte-Level Models

The gist: This paper proposes an efficient distillation recipe that converts existing token-trained LLMs into Byte Language Models (BLMs) using a two-stage curriculum, achieving comparable performance with only approximately 125B bytes of training data.

Proposed Distillation Recipe

The paper introduces a novel distillation recipe that follows a two-stage curriculum to convert existing token-trained LLMs into BLMs while retaining comparable capabilities. This recipe is structured into (1) Progressive Knowledge Distillation, which aligns byte-level representations with the embeddings of the token-trained teacher model, and (2) Byte-Level Supervised Fine-Tuning, which enables end-to-end generation entirely in the byte space. The central challenge addressed is the boundary mismatch between token-based models and byte models.

Stage 1: Progressive Knowledge Distillation

Stage 1 focuses on progressive representation alignment and boundary learning. This stage involves three specific loss objectives to align the student model’s embedding space, semantic probabilities, and boundary decisions with those of the teacher model.

1. Embedding Alignment (Lalign): This objective minimizes the difference between the hidden state of the student model encoder at the boundary byte eˆk and the static embedding of the teacher model ek.

2. Joint Distillation (Ldistill): This objective minimizes the KL divergence between the student model’s and teacher model’s output distributions by utilizing the teacher model’s segmentation boundaries to synchronize sequence lengths.

3. Boundary Learning (Lboundary): This objective trains the Routing Module to predict token boundaries via binary classification on each byte, minimizing the loss defined as Lboundary.

The framework adopts a sequential curriculum rather than a holistic objective to ensure robustness and generalizability, optimizing Encoder Alignment first, then Joint Distillation, and finally Boundary Learning.

Stage 2: Byte-Level Supervised Fine-Tuning (SFT)

Building on Stage 1, Stage 2 transitions the model to generate entirely in the byte space. The student model replaces the token LM head with Dechunk and Decoder modules and is trained using the Next Byte Prediction loss. This process involves two steps: Head Adaptation, training only the newly initialized Dechunk and Decoder modules, followed by end-to-end fine-tuning where all parameters are unfrozen to optimize the entire model.

Architectural Modifications for Byte Modeling

The student model is based on H-Net, a U-Net-like architecture crafted for multigranularity sequence modeling. Key architectural adaptations include the ONE-BYTE LOOKAHEAD ROUTING mechanism, which computes the similarity between the current byte hidden state xcurr and the next byte hidden state xnext to derive a boundary indicator b = I(p ≥ 0.5). For decoding, two strategies are explored: Joint Boundary Prediction (JBP), which encodes boundary information directly into the output vocabulary, and Multi-Byte Prediction (MBP), which introduces an auxiliary language modeling head to predict the next-next byte to simulate lookahead behavior during inference. Furthermore, a Dechunk module employs a dynamic pooling strategy by revealing chunk representations only at the final byte of the k-th chunk to preserve autoregressive properties while providing maximal context at chunk boundaries.

Empirical Validation and Results

The distillation recipe is validated across multiple model families, including Llama, Qwen, and OLMo. Empirical results demonstrate unprecedented data efficiency, with approximately 125 B training bytes in total required for the distilled model. For instance, on the MMLU benchmark, the method achieves scores of 51.8 and 68.5 for Llama-3.2 3B and Qwen-3 4B respectively, retaining over 92% of the original performance. The framework shows superior performance retention compared to baselines like BLT and H-Net Distill, achieving lower average drop across various benchmarks.

Robustness and Further Enhancements

The distilled Stage 1 models demonstrate a significant level of intrinsic resilience against character-level perturbations. However, the analysis reveals that the boundary predictor is heavily overfitted to surface patterns, suggesting a Space Bias. Proposed mitigation strategies include Trim Data, which removes all whitespaces to force reliance on character sequences alone, and Whitespace Penalty, which introduces a penalty term in the boundary loss function to discourage boundaries at whitespace positions. Fine-tuning on perturbed benchmarks shows that the byte-level Llama distill model exhibits a significantly higher adaptation ceiling than its subword counterpart, achieving the highest overall Robustness Score (69.93) after 10 epochs of training on noisy data. The findings also indicate that the Mamba2 encoder offers high fidelity compared to standard Transformer encoders in aligning fine-grained byte sequences with high-level semantic embeddings.

Comparison with Related Work

The framework diverges from prior approaches like Bolmo by decomposing the process into a sequential curriculum, explicitly separating representation alignment and capability transfer from boundary learning. The paper also highlights that the performance of decoupled routing modules suffers a notable decline compared to integrated ones, suggesting that token boundaries are heavily dependent on high-level semantic and syntactic contexts provided by a fully developed encoder. Moreover, the study confirms that employing a significantly larger teacher model does not offer a distinct advantage in this distillation context, as standard setups with similarly scaled models consistently produce better outcomes. The final results show that while the distilled model is competitive, there remains room for improvement to match the performance of models trained with different objectives or larger computational budgets.

Conclusion

In this paper, a practical two-stage distillation framework is proposed that converts pretrained token-based LLMs into byte-level models while preserving most of their capabilities. This method is effective across multiple model families, including Llama, Qwen, and OLMo, and achieves competitive performance using only 125B training bytes. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries. This work substantially lowers the cost of building scalable BLMs. The distillation process successfully enables the H-Net to utilize its byte-to-byte architecture to bypass the fragility of the subword boundaries.

1. Bold header: Byte-Level Model Creation via Two-Stage Distillation

The improved AI system can convert existing token-trained LLMs into byte-level models using a two-stage curriculum, which drastically reduces data requirements, achieving competitive performance with approximately 125B bytes of training data.

2. Bold header: Progressive Knowledge Distillation for Representation Alignment

The system implements Stage 1, Progressive Knowledge Distillation, utilizing three loss objectives—embedding alignment (Lalign), joint distillation (Ldistill), and boundary learning (Lboundary)—to align byte-level representations with the embeddings of the token-trained teacher model.

3. Bold header: Dynamic Boundary Learning for Autonomous Tokenization

The system incorporates a one-byte lookahead mechanism in the routing module, which is trained via the Boundary Learning Objective (Lboundary) to predict token boundaries, enabling the student model to dynamically group raw bytes into chunks and generate in the token space.

4. Bold header: Robustness Enhancement via Post-Training Optimization

The system can be enhanced using post-training techniques like GOLD integrated with On-Policy Distillation, which leverages the exceptionally high boundary detection accuracy of the Stage 1 model to further align its segmentation trajectory.