Self-Policy Distillation via Capability-Selective Subspace Projection

2026-05-21

AI Summary

Existing self-distillation methods rely on external signals or raw model outputs, making it difficult to effectively disentangle task-relevant from task-irrelevant features, thereby limiting generalization performance. This work proposes Self Policy Distillation (SPD), the first general-purpose, signal-free approach to selective self-distillation. SPD extracts a low-rank capability subspace by analyzing the model’s own gradients and projects key-value (KV) activations onto this subspace during self-generation, followed by fine-tuning with the standard language modeling loss. Evaluated on code generation, mathematical reasoning, and multiple-choice question answering tasks, SPD significantly outperforms current signal-free self-distillation methods, achieving performance gains of up to 13%, a 16% improvement over the pretraining baseline, and a 15% increase in out-of-domain generalization.

Abstract

Self-distillation bootstraps large language models (LLMs) by training on their own generations. However, existing methods either rely on external signals to curate self-generated outputs (e.g., correctness filtering, execution feedback, and reward search), which are costly and unavailable for the best-performing frontier models, or skip curation entirely and train on all raw outputs, an approach that is often domain-specific and hard to generalize. Both also share a deeper weakness that self-generated outputs entangle task-relevant capability with others, such as stylistic patterns, formatting artifacts, and model-specific errors, diluting the signal for the specific capability one aims to improve. In this paper, we propose Self-Policy Distillation (SPD), which achieves generalizable, capability selective without any external signal. Specifically, SPD extracts a low-rank capability subspace from the model's own gradients on correctness-defining tokens, projects key-value (KV) activations into this subspace during self-generation, and fine-tunes on the resulting raw outputs with standard next-token prediction loss. Through extensive experiments across code generation, mathematical reasoning, and multiple-choice QA, we show that SPD achieves up to 13% improvement over state-of-the-art self-distillation methods without external signals and up to 16% improvement over pre-trained baselines. Notably, SPD demonstrates superior generalizability, achieving 15% better performance under out-of-domain generalization settings.

Problem

• self-distillation
• capability disentanglement
• large language models
• generalization
• external signal

Innovation

• Self-Policy Distillation
• capability-selective subspace
• low-rank projection
• self-distillation without external signals
• KV activation projection