KDRL Unifies Distillation And Reinforcement Learning

KDRL, one loss, two teachers, teacher supervision and reward exploration

The false choice between imitation and exploration

DeepSeek R1 made two things obvious to anyone paying attention. Reinforcement learning can pull genuinely new reasoning behavior out of a language model that supervised training alone never produces, and distilling from a strong reasoning model is often the faster, cheaper path to get a smaller model most of the way there. Both of those observations are true, and they point in slightly different directions. Distillation is efficient because the student gets a dense, informative signal at every single token, but it can only ever be as good as its teacher, and it tends to fall apart outside the exact distribution of problems the teacher was good at. Reinforcement learning has no such ceiling since it lets the model explore and be rewarded for whatever actually works, but sparse binary rewards on long reasoning chains make for a slow, sample inefficient training signal, especially early on when the base model rarely stumbles into a correct answer at all.

The standard playbook, used by DeepSeek R1 itself and by most of the open reasoning models that followed, is to run these two stages back to back. Distill first to get the model to a reasonable starting point, then switch to reinforcement learning to sharpen it further. That works, but the paper’s authors point out a real cost hiding in the sequencing. Once you move to the reinforcement learning stage, teacher supervision is gone entirely. Whatever knowledge the teacher had that the distillation stage did not fully transfer is simply unavailable for the rest of training. The two learning signals never actually talk to each other.

How KDRL actually combines the two signals

Reverse KL instead of standard supervised fine tuning

The first design decision is what kind of distillation to use. Ordinary supervised fine tuning trains the student on sequences generated by the teacher, which is mathematically equivalent to minimizing the forward KL divergence between the teacher and the student. The problem is exposure bias. The model is trained on teacher prefixes but has to make predictions from its own, different prefixes at inference time, and for long reasoning chains that mismatch compounds badly.

KDRL instead minimizes the reverse KL divergence, sampling from the student’s own policy rather than the teacher’s. The authors show this is mathematically equivalent to a REINFORCE style objective where the reward at each token is the log ratio of the teacher’s probability to the student’s probability for that token.

Here R is the log ratio of teacher probability to student probability for a given token. When R is positive the teacher liked that token more than the student currently does, and the gradient pushes the student to raise its own probability of it. When R is negative the opposite happens. This is what the paper calls KD RKL, and its main appeal is that it trains on the student’s own generations, which mitigates the exposure bias problem that plain supervised fine tuning has.

Two ways to merge KD RKL with GRPO, and only one of them works

With RKL framed as a reward like quantity, merging it with reinforcement learning seems almost obvious. Just add it to the reward the reinforcement learning algorithm already optimizes, an approach the paper calls reward shaping. The modified reward becomes the original outcome reward plus a weighted sum of the teacher’s token level preferences, and that combined reward feeds into the usual GRPO advantage calculation.

It does not work. The paper reports that reward shaping causes training to collapse in the early steps, with training reward visibly dropping before ever recovering. The second approach, called joint loss, keeps the two objectives structurally separate instead of blending them inside a single reward. The GRPO objective handles reward optimization on its own terms, and a separate KL divergence term is subtracted directly from the loss.

The gradient of this joint objective works out to an unbiased combination of the GRPO gradient and the RKL gradient, weighted by a coefficient beta. Unlike reward shaping, this keeps the reward signal and the imitation signal on separate footing rather than letting one distort the other’s scale, and the paper reports it as stable from the first step onward while consistently beating both GRPO alone and KD RKL alone on validation accuracy.

Picking the right way to estimate the KL divergence

Once joint loss is settled on, there is still a question of how to actually compute that KL term, since the true divergence has to be estimated from samples rather than calculated exactly. The paper compares three estimators. The direct estimator has high variance. The k3 estimator is unbiased as a scalar quantity but its gradient is biased, meaning it systematically pushes training in a slightly wrong direction even though the number it reports looks correct. The k2 estimator is the reverse case, it is a biased estimate of the divergence itself but its gradient is exactly unbiased, which turns out to matter far more for actual training outcomes than having an accurate scalar to look at.

A fourth candidate, computing the divergence over the full vocabulary using only the teacher’s top K most likely tokens, sounded promising as a way to get denser supervision than sampling a single token gives. In practice the paper finds it destabilizes training, with a sharp rise in repetitive text inside overlong responses, likely because tokens sitting just outside the top K set get penalized excessively relative to how the student was actually behaving.

Annealing the balance between imitation and exploration

The coefficient beta in the joint loss controls how strongly the teacher’s preferences pull against the reward signal. A larger beta gives the student more scaffolding early on but risks the student never developing exploration behavior of its own, showing up as response length inflating quickly and truncation rates climbing. A smaller beta lets the reward signal dominate but leaves useful teacher information on the table.

Rather than picking one fixed value, the paper anneals beta linearly downward over training, starting high to lean on the teacher early and gradually shifting weight toward the reward signal as training progresses. This annealed version, KDRL Annealing, ends up as the paper’s strongest configuration, matching the response length growth of the best fixed beta setting while keeping a lower clip ratio, which the authors read as more efficient use of the reward signal once the model has enough of its own footing to explore productively.

Not every correct response needs a teacher’s help

One more refinement addresses a subtler problem. Because GRPO computes its advantage at the response level and spreads it across every token in that response, combining it with a token level KL penalty can create conflicting gradients even on responses the model already got right. The RKL term will still push those tokens toward matching the teacher exactly, even when the student’s own phrasing was already correct, which adds regularization pressure the model does not need.

The fix, called reward guided masking, simply switches off the KL penalty for any response that already received a positive reward. A stricter group level version switches off the penalty for an entire batch of sampled responses if even one of them succeeded. The response level version turns out to be the better trade, cutting response length by more than ten percent while holding accuracy roughly steady, because it concentrates teacher supervision specifically on the queries where the student is still struggling rather than diluting it across responses that were already fine.

What the benchmark numbers say

The main experiments start from DeepScaleR 1.5B 8K, itself a lightly pretrained 1.5 billion parameter model, and use Skywork OR1 Math 7B as the teacher. Table 1 in the paper compares six training strategies across six math reasoning benchmarks, three of them competition level, AIME24, AIME25, and AMC23.

The ranking is consistent across nearly every benchmark in the paper. Supervised fine tuning barely moves the needle past the base model. Pure reinforcement learning does meaningfully better on its own. On policy distillation through KD RKL does better still, which the paper notes lines up with a similar finding published around the same time in the Qwen3 technical report. KDRL, whichever integration strategy is used, beats every one of those individually.

The efficiency comparisons are where the practical value becomes clearest. KD RKL needs more than six thousand additional tokens per response compared to KDRL to reach similar accuracy, which is a real inference cost difference at deployment time. On a matched training budget, the paper reports that roughly 280 KDRL training steps correspond to about 200 steps of KD RKL and 460 steps of GRPO for the same total wall clock time, since KD RKL’s teacher inference overhead and long response generation slow it down, while GRPO trains fastest but without any teacher signal at all. KDRL sits in between and still wins on accuracy at matched compute, which is a stronger claim than simply winning at matched step count.

The paper also tests KDRL in the R1 Zero style setting, applying reinforcement learning directly to a base model rather than one already distilled, using Qwen2.5 3B as the student and a GRPO trained Qwen2.5 7B as the teacher. KDRL Annealing again comes out ahead, beating GRPO by 1.2 percent and KD RKL by 1.7 percent averaged across the same six benchmarks. Interestingly, in this zero shot setting GRPO alone slightly beats KD RKL alone, 31.3 percent versus 30.8 percent, a reversal from the main experiments that suggests on policy distillation is not universally superior to reinforcement learning, it depends on how strong the available teacher actually is relative to the task.

What this means for teams building reasoning models

The most transferable idea here is not really about language models specifically. It is that treating imitation and exploration as sequential stages, when both learning signals are actually differentiable at the same time, throws away information for no good reason. The instant a team has both a capable teacher and a working reward function, KDRL’s joint loss framing is a strictly more expressive training objective than picking one stage or the other, and the annealing schedule gives a principled way to shift from leaning on the teacher to leaning on self discovered reward as training matures rather than committing to one regime for the whole run.

The k2 versus k3 finding is worth internalizing even outside this specific paper. Anywhere a team is estimating a KL divergence from samples to compute a gradient, the unbiasedness of the gradient estimator matters more than the unbiasedness of the scalar the estimator produces. It is an easy detail to get backwards if you are optimizing for a metric that looks correct in a logging dashboard rather than the actual update direction it produces.

The reward guided masking result also generalizes usefully. Any time a fixed penalty term is applied uniformly across a batch that contains a mix of easy and hard examples, there is a reasonable chance the penalty is doing more harm than good on the easy examples the model already solved. Conditioning auxiliary losses on whether the primary objective was already satisfied is a cheap, broadly applicable pattern well beyond distillation specifically.

Honest limitations

The experiments are entirely on math reasoning benchmarks with a rule based, automatically checkable reward function. That is a favorable setting for reinforcement learning in general, since the reward signal is unambiguous and cheap to compute. Whether KDRL’s advantages hold up on reasoning tasks with fuzzier, model judged rewards, like open ended writing or multi step tool use, is untested here, and reward hacking dynamics in those settings could interact with the KL term differently than they do on math problems with a single verifiable answer.

The method also depends on having a genuinely strong teacher available and, in the KDRL variant used in the main experiments, on that teacher being close enough in output distribution to the student for the reverse KL estimate to behave well. The paper’s own ablation with a weaker teacher shows the benefit shrinking substantially, and there is no guarantee the same annealing schedule and KL coefficient values transfer cleanly to a very different teacher student capability gap.

Finally, all of the main results come from one student scale, 1.5 billion parameters, with one specific teacher pairing, and the R1 Zero experiments use a 3 billion parameter student. Whether the same relative gains hold at the much larger scales where reasoning models are typically deployed in production is a real open question the paper does not directly address, since running these comparisons at larger scale would multiply an already teacher inference heavy training cost substantially.

A minimal PyTorch implementation of the KDRL joint loss

The code below implements the core mechanics from the paper in a simplified, runnable form. It defines a toy student and teacher language model, computes the GRPO style policy gradient term, the k2 estimator for the reverse KL term, combines them into the joint loss from Section 3.1, applies response level reward guided masking from Section 3.4, and runs a smoke test on randomly generated dummy token sequences.

import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------------------------------------------------------------
# A minimal toy language model, standing in for a real transformer
# decoder. Just enough structure to produce per token logits over
# a small vocabulary so the loss mechanics can be demonstrated.
# ---------------------------------------------------------------
class TinyLM(nn.Module):
    def __init__(self, vocab_size=64, embed_dim=32, hidden_dim=64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.rnn = nn.GRU(embed_dim, hidden_dim, batch_first=True)
        self.head = nn.Linear(hidden_dim, vocab_size)

    def forward(self, token_ids):
        # token_ids, shape (batch, seq_len)
        x = self.embed(token_ids)
        h, _ = self.rnn(x)
        logits = self.head(h)
        return F.log_softmax(logits, dim=-1)


# ---------------------------------------------------------------
# GRPO style advantage, computed at the response level and shared
# across every token in that response, matching Eq. 2 in the paper.
# ---------------------------------------------------------------
def grpo_advantage(rewards):
    # rewards, shape (group_size,), one scalar reward per sampled response
    mean = rewards.mean()
    std = rewards.std().clamp(min=1e-6)
    return (rewards - mean) / std


# ---------------------------------------------------------------
# The k2 estimator of the reverse KL divergence and its gradient,
# following Eq. 6 in the paper. R is the log ratio of teacher to
# student probability for the sampled token.
# ---------------------------------------------------------------
def k2_kl_term(student_log_probs, teacher_log_probs, token_ids):
    student_token_lp = torch.gather(student_log_probs, 2, token_ids.unsqueeze(-1)).squeeze(-1)
    teacher_token_lp = torch.gather(teacher_log_probs, 2, token_ids.unsqueeze(-1)).squeeze(-1)
    R = teacher_token_lp - student_token_lp
    return 0.5 * (R ** 2), R


# ---------------------------------------------------------------
# Response level reward guided masking from Section 3.4. Zeroes
# out the KL penalty on responses that already received a reward.
# ---------------------------------------------------------------
def reward_guided_mask(k2_terms, rewards):
    # rewards, shape (group_size,) . k2_terms, shape (group_size, seq_len)
    mask = (rewards == 0).float().unsqueeze(-1)
    return k2_terms * mask


# ---------------------------------------------------------------
# The full KDRL joint loss, Eq. 8 in the paper. Combines the GRPO
# policy gradient term with the beta weighted k2 KL term.
# ---------------------------------------------------------------
def kdrl_joint_loss(student_log_probs, teacher_log_probs, token_ids, rewards, beta=2e-3, use_masking=True):
    student_token_lp = torch.gather(student_log_probs, 2, token_ids.unsqueeze(-1)).squeeze(-1)
    advantage = grpo_advantage(rewards).unsqueeze(-1)
    policy_term = -(advantage * student_token_lp).mean()

    k2_terms, _ = k2_kl_term(student_log_probs, teacher_log_probs, token_ids)
    if use_masking:
        k2_terms = reward_guided_mask(k2_terms, rewards)
    kd_term = k2_terms.mean()

    total_loss = policy_term + beta * kd_term
    return total_loss, policy_term.item(), kd_term.item()


# ---------------------------------------------------------------
# Smoke test, runs one KDRL update on dummy data end to end and
# checks the loss is finite and the model parameters actually move.
# ---------------------------------------------------------------
if __name__ == '__main__':
    torch.manual_seed(0)
    vocab_size, seq_len, group_size = 64, 12, 8

    student = TinyLM(vocab_size=vocab_size)
    teacher = TinyLM(vocab_size=vocab_size)
    teacher.eval()
    optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)

    token_ids = torch.randint(0, vocab_size, (group_size, seq_len))
    rewards = torch.randint(0, 2, (group_size,)).float()

    before = [p.clone() for p in student.parameters()]

    for step in range(3):
        student_log_probs = student(token_ids)
        with torch.no_grad():
            teacher_log_probs = teacher(token_ids)

        loss, policy_part, kd_part = kdrl_joint_loss(
            student_log_probs, teacher_log_probs, token_ids, rewards, beta=2e-3
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print('step', step, 'total loss', loss.item(), 'policy term', policy_part, 'kd term', kd_part)

    moved = any(
        not torch.allclose(a, b) for a, b in zip(before, student.parameters())
    )
    assert torch.isfinite(loss), 'loss should be a finite number'
    assert moved, 'student parameters should update after a training step'
    print('smoke test passed')

Conclusion

The core achievement of this paper is showing that a framing choice most of the field treats as settled, distill then reinforce, was leaving value on the table for no principled reason. Once teacher supervision and reward driven exploration are expressed as terms in the same differentiable objective rather than sequential training stages, the model gets to use both signals for the entire training run instead of losing access to one of them partway through.

The conceptual shift that matters most here is subtle. It is not that combining KD and RL is a new idea in the abstract, prior work like GKD explored a related direction for text summarization. What KDRL adds is a careful accounting of exactly which combination strategy is stable, which KL estimator’s gradient actually points in the right direction even when its scalar value does not, and how to avoid wasting teacher supervision on examples the model has already mastered. Those are the kinds of details that separate a promising idea from a training recipe someone else can actually reproduce.

Transferability looks strong on the mechanics and untested on the domain. The joint loss framing, the k2 gradient argument, and reward guided masking are all generic enough to apply to essentially any post training setup where a verifiable reward and a stronger teacher both exist, not just math reasoning specifically. Where this needs more evidence is exactly the settings the paper did not test, reward functions that are noisier or model judged rather than rule based, and student scales much larger than 1.5 or 3 billion parameters, where the extra teacher inference cost of on policy distillation becomes proportionally more expensive.

The honest limitations are worth repeating plainly. Teacher quality caps the achievable gain, the reward signal used throughout is unusually clean for a rule based math checker, and the paper does not establish how the specific numerical settings for the KL coefficient and its annealing schedule should change as the teacher student capability gap changes. None of that undercuts the central result, but it does mean a team adopting this approach should expect to retune those hyperparameters for their own teacher and task rather than copying the paper’s exact numbers.

What is easy to miss on a first read is how much of the paper’s value sits in the negative results. Reward shaping collapsing training, the Top K approximation destabilizing optimization, group level masking underperforming its response level counterpart. Each of those closed off a design choice that looked reasonable on paper, and knowing which doors not to open is often worth as much as knowing which one to walk through.

Frequently asked questions

Related reading