Re-Distilling Smaller DeepSeek R1 Models for Better Performance

Table of Contents

Introduction

The DeepSeek R1 distilled models are currently the most popular way to run the R1 model, as they can be efficiently executed on consumer hardware. This family of smaller models includes various Qwen2.5 and Llama3 models, which can operate faster with quantization methods like HQQ, utilizing backends such as GemLite.

According to the paper, the distilled models were trained on 800,000 samples generated by the original R1 model through Supervised Fine-tuning (SFT). Typically, a second reinforcement learning-based alignment step via RLHF or policy optimization like DPO, GRPO would achieve better performance.

In this work, we propose aligning smaller models with larger, more performant models via logits distillation. This method does not require preference data or even perfectly accurate samples. Using approximately 35,000 samples, this simple approach boosts performance across various benchmarks and costs only a few dollars per run. Below, we describe the approach in detail.

Re-Distillation Approach

Our approach leverages logits distillation between models sharing similar tokenizers. The key idea is to use larger models' output distributions to guide smaller ones, which proves effective without requiring extensive training data. In our experiments, we used Qwen 32B to guide both Qwen 1.5B and Qwen 7B, while Llama3-70B guided Llama3-8B.

Due to memory constraints with larger teacher models, quantization is necessary - we used HQQ to compress the teacher models to 8-bit precision. In our specific setup, we ran the pipeline on two GPUs: the quantized teacher generated logits on one GPU while the student trained on the other, with operations parallelized via CUDA streams. We kept the batch size at 1 to manage memory usage, though different hardware configurations and batch sizes are possible depending on available resources.

For the training objective, we employ KL-divergence loss with logits clipping - a crucial detail as we found larger models can produce over-confident predictions leading to training instability. When vocabulary sizes don't match between models, we zero-pad the missing embeddings in the lm-head and mask them in the loss computation. While it's possible to combine this with next-token prediction loss when reasoning data is available, our experiments focused on pure KL-divergence with a linear learning rate schedule and single dataset pass. Below is our masked KL-divergence loss implementation:

	import torch.nn.functional as F
	def KL_divergence_masked_loss(output, target, target_mask=1, T=1, clip_logits=50):
		
		#Clip logits to avoid nan loss
		target.clamp_(-clip_logits, clip_logits) 
		output.clamp_(-clip_logits, clip_logits)

		#Mask missing tokens
		target *= target_mask

		#Regular KL-divergence loss
		target = F.softmax(target / T, dim=-1)
		output = F.log_softmax(output / T, dim=-1)
		out    = (target * (target.log() - output)).sum(dim=-1).mean()

		return out

For our data, we sampled approximately 35,000 examples from various open-source datasets, covering topics such as math, knowledge, and coding. The datasets include prompts from MetaMathQA, Orca Math, Evol-Instruct, etc. Additionally, newly released open-source R1 reasoning datasets, like Numira R1 and R1-Distill-SFT, offer additional valuable resources.

Costs

In terms of costs, the re-distillation training phase costs are as follows using 2 x H100 SXM:

Benchmarks

Below, we present various comparative benchmarks that demonstrate significant improvements across multiple tasks. Notably, the GSM8K score shows an increase of over 4% compared to the original distilled models with the Qwen models and almost 14% with the Llama3 version.

Qwen 1.5B

Qwen 7B

Llama3-8B

Conclusion

In this work, we have shown that logit alignment, when used as a second-stage distillation, is a highly effective and economical method for enhancing smaller models, such as the DeepSeek R1 models. We are excited to share these improved models with the community on Hugging Face 🤗. We believe that this approach will inspire further advancements in the development of high-performance, small-scale models. We look forward to seeing how the community leverages this technique to push the boundaries of what is possible.

Citation

@misc{badri2025r1,
title = {Re-Distilling Smaller DeepSeek R1 Models for Better Performance,
url = {https://mobiusml.github.io/r1_redistill_blogpost/},
author = {Hicham Badri, Appu Shaji},
month = {January},
year = {2025}
}