POCL: How A Progressive Overload Curriculum Improves LLM Distillation

Why shrinking a language model is harder than it looks

Large language models keep getting better, but that improvement comes bundled with a cost problem. A model with billions of parameters needs serious hardware just to run, which rules it out for a lot of practical deployment settings, particularly anything running on the edge rather than in a data center. Knowledge distillation is the standard fix, training a small student model to reproduce the behavior of a large teacher model, giving up some capability in exchange for a model that is cheap enough to actually run where it is needed.

The paper draws a useful distinction between two flavors of distillation. Black box distillation only has access to a teacher’s final outputs, which is the situation anyone stuck behind a proprietary API like GPT-4o or Claude finds themselves in. White box distillation, the focus of this paper, has access to the teacher’s actual internal probability distributions over its vocabulary at every generation step, which opens up much richer training signals but requires the teacher to be an open model you can actually inspect. As open source models have closed the performance gap with proprietary ones, white box distillation has become a more attractive option since it offers finer control over what the student actually learns.

Three ways white box distillation quietly goes wrong

The paper walks through a set of specific failure modes that plague white box distillation once you look past the headline accuracy numbers. Catastrophic forgetting and mode collapse show up because a small student model has limited capacity, and forcing it to match a much larger teacher’s distribution too aggressively can overwrite what little capability it already had, sometimes leaving it stuck generating repetitive or narrow output. Different loss function choices trade off against each other here in an awkward way. Reverse KL divergence tends to sharpen the student’s predictions, which helps avoid over smoothing but can make forgetting and collapse worse, while standard KL divergence keeps training more stable but often produces a weaker distilled model overall.

A second failure mode is what the paper calls training inference mismatch. Student models typically train on a fixed, pre existing dataset, but at inference time they see whatever prompts a real user throws at them, which will not perfectly match the training distribution. One fix is mixing in student generated outputs, letting the student practice on its own responses during training rather than only ground truth text. That helps, but comes with real costs. Too much reliance on self generated data introduces noise, since a weak student’s own outputs are not always trustworthy targets, and generating that data during training is expensive, sometimes eating up to eighty percent of total training time according to prior work the paper cites.

The progressive overload metaphor, and why it actually maps onto training

Progressive overload is a basic principle from strength training, gradually increasing the weight, volume, or intensity of a workout over time so the body adapts without being overwhelmed on day one. The paper builds its entire framework around treating the teacher model as a coach and the student model as an athlete. The training dataset becomes the workout plan, where the size of each training stage corresponds to training volume and the difficulty of the samples in that stage corresponds to training intensity.

Framed this way, the student begins with a small set of easy examples, much like a beginner starting with light weights, and gradually works through larger and harder training subsets as it builds capability, mirroring how a lifter’s program advances in stages. The paper argues this is not just a cute analogy, curriculum learning of this general shape has a real theoretical justification here, since starting with samples where the student and teacher distributions already align most closely avoids the abrupt distribution shifts that seem to drive both catastrophic forgetting and training inference mismatch in the first place.

How POCL is actually built

The framework has two working parts, a Difficulty Measurer that decides which samples are easy and which are hard, and a Training Scheduler that decides when to introduce each batch of samples and how to adjust training hyperparameters along the way.

The Difficulty Measurer, ranking samples without extra labeling

Deciding what makes a training sample hard is not obvious in a generative setting, so the authors combine two independent signals. One is the ROUGE-L score between the student model’s own generated output and the ground truth response, where a higher score suggests the student already handles that example reasonably well. The other is the student’s cross entropy loss on that same example, where a lower loss suggests the same thing from a different angle. Rather than picking one signal over the other, the paper combines both rankings using reciprocal rank fusion, a technique borrowed from information retrieval.

A higher fusion score means an easier sample. Once every training example has a fusion score, the full dataset gets split into n subsets ordered from easiest to hardest, with n set to 4 in the paper’s main experiments, a value the authors say worked well enough across datasets that they did not bother tuning it further.

The Training Scheduler, introducing difficulty gradually

With the dataset partitioned, training proceeds through n stages using a scheduling approach called Baby Step, a well established curriculum learning technique. At stage one, the student trains only on the easiest subset. Once that stage converges or a fixed number of epochs passes, the next hardest subset gets merged in, and training continues on the combined set. This keeps expanding until, by the final stage, the student is training on the complete dataset, same as it would under standard training, just having built up to that point gradually rather than facing it all at once.

Two additional hyperparameters ride along with this staged expansion. The distillation temperature, which controls how smoothed out the teacher’s probability distribution appears to the student, starts low and rises linearly across stages.

The intuition here is that early stages, working with easy samples, benefit from a sharper, more confident view of the teacher’s predictions, while later stages, working with harder samples, benefit from the softer, more nuanced signal a higher temperature provides. Alongside temperature, the ratio between the supervised cross entropy loss against ground truth and the distillation loss against the teacher, called alpha, follows the opposite trend for off policy methods, starting relatively high and decreasing to zero.

The reasoning borrows from a teaching principle sometimes called least to most prompting, lean on ground truth guidance early while the student is still shaky, then hand more control over to the teacher’s own signal as the student gets more capable. For on policy methods, which already mix in student generated outputs during training, the paper sets alpha to zero throughout, arguing that ground truth labels can actively interfere with a student learning from its own generated outputs in that specific setting.

Does POCL actually move the needle

The main experiments distill GPT-2 1.5B down to GPT-2 120M and OPT 2.7B down to OPT 350M, both trained on the databricks dolly 15K instruction following dataset, then evaluated with ROUGE-L across five separate instruction following benchmarks, DollyEval, SelfInst, VicunaEval, S-NI, and UnNI.

Every one of the seven distillation methods tested, spanning both on policy GKD and off policy variants like KLD, reverse KLD, JSD, TVD, and the two skew KLD versions, sees its average ROUGE-L score improve when POCL is added, with gains ranging from roughly 1.5 to 2.6 points depending on the specific loss function. The largest jump belongs to SKL, improving by 2.59 points on average. Several results even cross into territory where the student briefly outperforms its own teacher on individual benchmarks, marked with an asterisk in the original table, which happened most often on the S-NI and VicunaEval datasets. A parallel set of results using OPT 2.7B distilled down to OPT 350M, included in the paper’s appendix, shows the same consistent pattern of improvement across every distillation method tested, with average gains ranging from about half a point up to just over two points.

Figure 3 in the paper tracks ROUGE-L on the validation set across training iterations for both KLD and GKD, with and without POCL. Both curves with POCL sit consistently above their non POCL counterparts for essentially the entire training run, not just at the final checkpoint, which is a meaningfully stronger claim than simply comparing two endpoint numbers, since it suggests the benefit holds throughout training rather than only showing up by chance at whatever epoch training happened to stop.

Rising distillation temperature achieves superior performance over all baselines, and temperature has a more pronounced effect than the SFT ratio.Reading of the POCL adaptive parameter ablation, Section 5.2 of the source paper

What the ablations reveal about which piece is doing the work

The paper runs a genuinely thorough set of ablations, breaking POCL apart component by component rather than just reporting the full system’s numbers.

Does the fusion ranking actually beat simpler ranking methods

Comparing the combined ROUGE-L and cross entropy fusion ranking against using either signal alone shows the fusion approach coming out ahead on the in domain DollyEval benchmark.

All three ranking variants beat plain KLD without any curriculum at all, which is itself evidence that some form of sample ordering helps regardless of exactly how you compute it. The fusion approach edges out both single metric versions on GPT-2, and edges out the cross entropy only version on OPT while trailing the ROUGE-L only version by a small margin on that same model. The overall pattern supports the paper’s claim that fusion ranking is more reliable across models, even though it is not uniformly the single best option in every individual comparison.

Which hyperparameter schedule matters more, temperature or SFT ratio

Stripping out the adaptive parameter schedules while keeping the staged sample introduction shows the temperature schedule carries most of the weight.

A striking detail here, curriculum sample ordering by itself, without either adaptive schedule, actually underperforms the plain KLD baseline on most datasets, dropping DollyEval from 23.49 to 23.21 and SelfInst barely moves. Removing only the temperature schedule while keeping sample ordering and the ratio schedule active makes things noticeably worse still, dropping DollyEval to 22.04 and SelfInst all the way down to 8.70, a bigger single drop than removing the ratio schedule causes. That comparison is the paper’s strongest evidence that the temperature schedule, not the sample curriculum itself, is doing most of the heavy lifting in the full POCL system. Keeping temperature active while dropping the ratio schedule recovers most of the gain, landing at 24.33 on DollyEval, just half a point behind the full system.

Does the order of difficulty actually matter, or would any structure do

A separate ablation directly tests easy to hard ordering against the reverse, hard to easy ordering, on both GKD and KLD across the GPT-2 and OPT model families.

Both orderings beat the plain baseline in essentially every case, which supports the general claim that structured, staged data exposure helps regardless of direction. But easy to hard consistently wins by a real margin over hard to easy, most dramatically on GKD where easy to hard reaches 26.60 on DollyEval against hard to easy’s 24.73, nearly matching the plain baseline. The paper’s interpretation is that starting with hard, unfamiliar samples may actively interfere with a student’s ability to later learn simpler underlying patterns, essentially confirming the strength training analogy holds up, you cannot skip straight to the heaviest weight and expect it to go well.

What this actually costs to run

A practical detail buried in the appendix deserves more attention than it usually gets. Despite splitting training into multiple curriculum stages, POCL is configured to run for only about forty percent of the total epochs used by the baseline methods it is compared against, while the paper reports keeping total training steps equal for fairness in the main comparison. The authors report that a single GPT-2 0.1B KLD based POCL run takes around ten minutes on their hardware, a single A800 80GB GPU setup with four GPUs used across the broader experimental suite. Combined with the fact that the curriculum adds no new trainable parameters and no architectural changes at all, the practical overhead of adopting POCL on top of an existing white box distillation pipeline looks close to negligible, which matters a great deal for a technique whose entire pitch is being a low cost plug in rather than a wholesale replacement.

Honest limitations

Where this fits in the wider model compression picture

Step back from the specific benchmark numbers and the underlying claim generalizes well past this one paper. Any training setup that forces a smaller or less capable model to imitate a larger one, whether that is classic knowledge distillation, pruning a network before it is even trained, quantization aware training that fine tunes a compressed model back toward its original behavior, or even a small model learning from a large model’s synthetic data in a black box setting, faces some version of the same abrupt distribution shift problem POCL is built to soften. The core recipe demonstrated here, ranking training difficulty using signals the student model already produces during training rather than requiring external annotation, and staging both data and hyperparameters together rather than tuning them independently, looks like a pattern worth testing in any of these adjacent compression settings, not just white box distillation of instruction following language models specifically.

Complete PyTorch implementation

The paper’s own code is publicly available on GitHub at the link in the footnote below, but for a compact, dependency light illustration of the core mechanics, here is an independent reimplementation covering the difficulty measurer with reciprocal rank fusion, the Baby Step training scheduler, the linear temperature and SFT ratio schedules, a simplified KLD based white box distillation loss, a training loop across curriculum stages, and a smoke test on randomly generated dummy text data.

# pocl_reimplementation.py
# Independent PyTorch reimplementation of POCL, "Being Strong Progressively,
# Enhancing Knowledge Distillation of Large Language Models through a
# Curriculum Learning Framework" by Liu and Zhang. This is not the authors'
# own code, it is a reconstruction built from the paper's equations for
# educational use. The authors' own implementation is linked in the article.
import torch
import torch.nn as nn
import torch.nn.functional as F
def reciprocal_rank_fusion(rouge_l_scores, cross_entropy_losses, k=60):
    # higher rouge_l and lower cross_entropy both indicate an easier sample
    n = len(rouge_l_scores)
    rouge_rank = sorted(range(n), key=lambda i: -rouge_l_scores[i])
    ce_rank = sorted(range(n), key=lambda i: cross_entropy_losses[i])
    rouge_rank_pos = {idx: pos for pos, idx in enumerate(rouge_rank)}
    ce_rank_pos = {idx: pos for pos, idx in enumerate(ce_rank)}
    fusion_scores = []
    for i in range(n):
        score = 1.0 / (k + rouge_rank_pos[i]) + 1.0 / (k + ce_rank_pos[i])
        fusion_scores.append(score)
    return fusion_scores  # higher fusion score, easier sample
def partition_by_difficulty(sample_indices, fusion_scores, n_subsets=4):
    # sort easiest first, then split into n roughly equal subsets
    ordered = [idx for _, idx in sorted(zip(fusion_scores, sample_indices), key=lambda t: -t[0])]
    n = len(ordered)
    base_size = n // n_subsets
    subsets = []
    start = 0
    for i in range(n_subsets):
        extra = 1 if i < (n % n_subsets) else 0
        end = start + base_size + extra
        subsets.append(ordered[start:end])
        start = end
    return subsets  # list of n_subsets lists of sample indices, easiest first
def temperature_schedule(stage_i, n_stages, tau0=1.0, tau_n=2.0):
    # eq (3), linear rise across stages, stage_i is 1-indexed
    if n_stages == 1:
        return tau0
    return tau0 + (tau_n - tau0) * (stage_i - 1) / (n_stages - 1)
def sft_ratio_schedule(stage_i, n_stages, alpha0=0.3, alpha_n=0.0):
    # eq (4), linear decrease across stages, for off-policy methods
    if n_stages == 1:
        return alpha0
    return alpha0 - (alpha0 - alpha_n) * (stage_i - 1) / (n_stages - 1)
def white_box_kd_loss(student_logits, teacher_logits, labels, alpha, tau, pad_id=-100):
    # eq (10), cross entropy against ground truth plus temperature scaled KLD to the teacher
    ce_loss = F.cross_entropy(
        student_logits.view(-1, student_logits.size(-1)),
        labels.view(-1),
        ignore_index=pad_id,
    )
    student_log_probs = F.log_softmax(student_logits / tau, dim=-1)
    teacher_probs = F.softmax(teacher_logits / tau, dim=-1)
    kd_loss = F.kl_div(student_log_probs, teacher_probs, reduction='batchmean') * (tau ** 2)
    total_loss = alpha * ce_loss + (1.0 - alpha) * kd_loss
    return total_loss, ce_loss.item(), kd_loss.item()
def run_pocl_stage(student, teacher, optimizer, stage_data, alpha, tau, epochs_per_stage=1, pad_id=-100):
    student.train()
    teacher.eval()
    stage_losses = []
    for _ in range(epochs_per_stage):
        for input_ids, labels in stage_data:
            optimizer.zero_grad()
            student_logits = student(input_ids)
            with torch.no_grad():
                teacher_logits = teacher(input_ids)
            loss, ce, kd = white_box_kd_loss(student_logits, teacher_logits, labels, alpha, tau, pad_id)
            loss.backward()
            optimizer.step()
            stage_losses.append(loss.item())
    return sum(stage_losses) / max(len(stage_losses), 1)
class TinyLM(nn.Module):
    # Stand in for a small autoregressive language model
    def __init__(self, vocab_size, hidden_dim=64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, hidden_dim)
        self.rnn = nn.GRU(hidden_dim, hidden_dim, batch_first=True)
        self.out = nn.Linear(hidden_dim, vocab_size)
    def forward(self, input_ids):
        h = self.embed(input_ids)
        h, _ = self.rnn(h)
        return self.out(h)
if __name__ == '__main__':
    # Smoke test on randomly generated dummy token sequences, no real dataset needed
    torch.manual_seed(0)
    vocab_size, seq_len, num_samples, n_stages = 200, 12, 40, 4
    student = TinyLM(vocab_size, hidden_dim=32)
    teacher = TinyLM(vocab_size, hidden_dim=32)  # stand-in teacher, frozen, untrained
    for p in teacher.parameters():
        p.requires_grad = False
    optimizer = torch.optim.AdamW(student.parameters(), lr=1e-3)
    # fake difficulty signals, in a real pipeline these come from student rollouts
    fake_rouge_l = [float(torch.rand(1)) for _ in range(num_samples)]
    fake_ce = [float(torch.rand(1)) for _ in range(num_samples)]
    fusion_scores = reciprocal_rank_fusion(fake_rouge_l, fake_ce)
    subsets = partition_by_difficulty(list(range(num_samples)), fusion_scores, n_subsets=n_stages)
    dummy_input_ids = torch.randint(0, vocab_size, (num_samples, seq_len))
    dummy_labels = torch.randint(0, vocab_size, (num_samples, seq_len))
    cumulative_indices = []
    for stage_i, subset in enumerate(subsets, start=1):
        cumulative_indices.extend(subset)
        tau = temperature_schedule(stage_i, n_stages)
        alpha = sft_ratio_schedule(stage_i, n_stages)
        idx_tensor = torch.tensor(cumulative_indices)
        stage_input = dummy_input_ids[idx_tensor]
        stage_labels = dummy_labels[idx_tensor]
        stage_batches = [(stage_input[j:j + 8], stage_labels[j:j + 8]) for j in range(0, stage_input.size(0), 8)]
        avg_loss = run_pocl_stage(student, teacher, optimizer, stage_batches, alpha, tau, epochs_per_stage=1)
        print(f'stage {stage_i}/{n_stages} tau={tau:.2f} alpha={alpha:.2f} samples={len(cumulative_indices)} avg_loss={avg_loss:.4f}')
    print('Smoke test complete, POCL stages run end to end on dummy data without errors.')

Conclusion

The core achievement of this paper is showing that a metaphor borrowed from strength training turns out to map cleanly onto real training dynamics, not just as a nice framing device but as a set of concrete hyperparameter schedules that measurably improve distillation across seven different loss functions and two separate model families. That kind of broad compatibility across loss function choices is the strongest evidence here that POCL is addressing something genuinely upstream of the specific distillation objective, the shape and order of the training data itself, rather than fixing a narrow quirk of any one loss function.

The conceptual shift worth carrying forward is the ablation finding that sample ordering alone is not enough, and can even underperform a plain baseline without the accompanying temperature schedule. That is a useful and slightly humbling result for anyone tempted to bolt curriculum learning onto an existing pipeline as a quick win, since it suggests the curriculum needs to be paired with matching changes to the training hyperparameters, not just a reordering of which examples show up first.

Transferability looks reasonably strong given how little the method assumes about the underlying distillation loss or data curation strategy. Because POCL only touches sample ordering and two scalar hyperparameters, and requires no architectural changes and only modest additional compute, the recipe should transfer cleanly to other white box distillation setups, and the general logic of staging difficulty alongside distillation temperature could plausibly extend to other training regimes that reuse a frozen or slowly updating teacher signal, such as certain reinforcement learning from human feedback pipelines.

The honest limitations deserve equal weight. Every result here comes from GPT-2 and OPT at one size pairing each, trained on a comparatively small fifteen thousand example instruction dataset, and the paper’s own stated concern, that larger student models may perceive less variation in sample difficulty to begin with, is a real open question the current experiments cannot answer. The specific temperature range, SFT ratio range, and subset count were fixed by the authors’ empirical judgment rather than swept systematically.

Where this goes next likely depends on testing POCL against more modern model families at larger scale, and checking whether the difficulty measurer’s reciprocal rank fusion approach holds up on datasets far larger than fifteen thousand examples, where computing per sample rankings could itself become a meaningful cost. Until that happens, the fair read is that POCL offers a cheap, broadly compatible, and empirically well supported argument for treating curriculum design as a first class part of knowledge distillation, one worth trying as a near free addition to whatever distillation pipeline a team is already running, especially given that the authors have made their implementation publicly available to build on directly.

Frequently asked questions

Related reading on aitrendblend