Applied Compute introduced Relevance-Masked Self-Distillation (RMSD) to train models on enterprise tasks involving out-of-distribution behavior such as internal tools and customer data.

We've ran into a new set of tasks that are challenging to train on because they involve internal tools and processes or customer preference data that can't be found on the internet. RL with low success rate doesn't get you very far (after all, you can't really expect learn if you keep failing and don't know what you don't know), so we've had to look for other algorithms.

The self-distillation literature has a nice solution involving a teacher and a student model. Think of the teacher as a "peer" of the student (it has the same model weights), but with access to a reference/solution manual that it uses to critique the student's work. We've found that this can lift the model out of a valley of failures on out-of-distribution tasks, and lead it to forget less compared to other methods that change the weights. Additionally, we introduce an extension (RMSD) that improves the data and compute efficiency of this method.

Check out our blog (link in thread).

I think everyone basically agrees you need some sort of filtering (exactly what is the question) over soft target distillation. The problem is that token space is the wrong domain imo.

Look at the sample below, even as a human I cannot tell you which tokens deserve to be upweighted (What do you mean we need to upweight "P" and "Ban"). The behaviour change we want lies several levels implicitly below token space. How do we expect an LLM judge to be able to filter?

The current prompt literally provides the divergent token index and the respective log probs. Even if you gave me that information, I won't even know what to do with it.

I don't know the solution but its definitely not token space credit assignment or masking imo

Enterprises increasingly want custom models that are tailored to their internal tools and processes, without sacrificing intelligence or reliability. Often times this involves tasks that are out-of-distribution for existing models: think custom document formats that aren't on the public web or company-specific legacy APIs, things that never appeared in pretraining.

Existing techniques struggle here. SFT teaches new behaviors, but causes catastrophic forgetting. RL has no foothold: if the target behavior never appears, there's no reward to reinforce.

We investigate self-distillation as a way to elicit these OOD behaviors, using a controlled toy task as a proxy for the enterprise-internal formats we see in practice.

We find that self-distillation can: - Better elicit OOD behaviors compared to SFT and RL - Better preserve the model's existing capabilities

Furthermore, we find that the token-level granularity in self-distillation is both a strength and a weakness – it's dense, but there is often a great deal of noise in the updates because the teacher and student may disagree on tokens for reasons unrelated to the desired behavior. We want to concentrate our loss on the tokens which matter most for the student's improvement.

This is the intuition behind RMSD: we filter for a set of T token positions where the teacher and student disagree the most, then an LLM judge narrows those to a final set of S tokens that are most relevant to the target behavior.

In our experiments, RMSD reaches the target behavior in about half the steps of OPSD, making it significantly more data efficient, while also training more stably and reaching a higher performance ceiling.

We think that the self-distillation approach and RMSD in particular will be powerful methods for advancing the capabilities of models on specialized settings.