Protecting Language Models Against Unauthorized Distillation through Trace Rewriting

1. Problem Statement

The paper addresses the threat of unauthorized knowledge distillation from frontier Large Language Models (LLMs). While knowledge distillation is a standard technique for transferring capabilities from a large "teacher" model to a smaller, efficient "student" model, it poses a significant risk when the teacher is a proprietary, high-cost model (e.g., frontier reasoning models).

• The Threat: Malicious actors can query a teacher model to obtain reasoning traces (structured intermediate steps leading to an answer) and use them to fine-tune a student model via Supervised Fine-Tuning (SFT). This allows the thief to replicate the teacher's capabilities at a fraction of the training cost.
• The Challenge: Existing defenses fall into two categories with significant drawbacks:
  1. Anti-distillation: Methods that degrade the quality of the output to prevent learning often severely degrade the teacher model's own utility, making them impractical.
  2. API Watermarking: Methods that embed statistical signatures often suffer from high false-alarm rates or rely on token-level statistics that are easily scrubbed by adaptive attackers.
• Goal: The authors aim to modify teacher-generated reasoning traces to achieve two complementary objectives without compromising the teacher's performance:
  1. Anti-distillation: Degrade the training efficacy for the student model.
  2. API Watermarking: Embed verifiable, stealthy signatures into the traces to prove ownership of the distilled model.

2. Methodology

The authors propose a unified framework based on Trace Rewriting, where a rewriting model ($R$) transforms the teacher's original reasoning trace ($r$) into a modified trace ($r'$) before it is used for distillation. The rewriting must preserve the correctness of the final answer and the semantic coherence of the reasoning.

They introduce two primary classes of rewriting approaches:

A. Instruction-Based Rewriting (Prompt-Based)

This approach uses an auxiliary LLM (the "Rewriter") to dynamically rewrite the teacher's output based on specific instructions.

• Semantic Prompting: A natural language prompt instructs the rewriter to degrade the trace's utility for distillation (e.g., "Rewrite to make it unhelpful for knowledge distillation while maintaining the correct answer") or to embed a specific watermark (e.g., "Add 'trigger = target' at a stealthy location").
• Optimized Prompting (OPRO): Leveraging the Optimization by PROmpting framework, the authors automatically search for the most effective rewrite instructions. They use a proxy student model to score candidate prompts based on the resulting drop in student accuracy (for anti-distillation) or watermark detection rates.
  • Key Insight: The optimized prompts often instruct the model to use "esoteric, formal, and densely technical lexicons" to obfuscate the reasoning logic for non-specialists while keeping the logic valid for the teacher.

B. Gradient-Based Rewriting

This approach treats the trace embeddings as a continuous space to be optimized directly against the student's learning objective.

• Embedding-Space Poisoning: The method approximates the gradient of the student's test loss with respect to the trace embeddings. It iteratively perturbs the embeddings to maximize the student's loss.
• Variants:
  • Hessian-based (HB-Grad): Uses second-order information (Hessian) for precise updates (computationally expensive).
  • First-order (FO-Grad): Approximates the update using only the gradient of the loss, avoiding the inverse Hessian.
  • Robust (RHB-Grad): Adds noise to the proxy student parameters to prevent overfitting.
• Constraint Satisfaction: To ensure the final answer remains correct, the gradient updates mask the final answer token, only modifying the reasoning steps.

3. Key Contributions

1. Unified Framework: A novel approach that simultaneously addresses anti-distillation and watermarking through semantic-level trace rewriting, avoiding the trade-offs of previous methods.
2. Optimized Prompting: The introduction of an automated prompt optimization process that yields instructions significantly more effective than manual semantic prompting.
3. Stealthy Watermarking: A method to embed behavioral triggers (e.g., specific input-output associations) into the reasoning traces that are robust against paraphrasing and filtering, unlike token-level statistical watermarks.
4. Comprehensive Evaluation: Extensive experiments across multiple datasets (GSM8K, MATH, MMLU) and student architectures, demonstrating superiority over state-of-the-art baselines (ADS, DOGe, GINSEW, KGW, VIA).

4. Experimental Results

Anti-Distillation Performance

• Effectiveness: The Optimized Prompting (OPT) approach achieved the strongest anti-distillation effect, reducing student model accuracy by up to 61.3% on the MATH dataset compared to distillation from clean traces.
• Teacher Utility: Unlike baselines (ADS, DOGe) which degraded teacher performance, the OPT approach maintained or even improved the teacher's accuracy (e.g., +2.4% on GSM8K, +22.3% on MATH). This is attributed to the rewriter correcting errors in the original teacher traces during the rewriting process.
• Scaling: Stronger student models experienced greater performance degradation, suggesting capable models are more susceptible to learning the "corrupted" reasoning patterns.
• Robustness: The method remained effective against adaptive attacks, including paraphrasing (Parrot) and keypoint-based progressive distillation (KPOD). In fact, paraphrasing the rewritten traces further amplified the anti-distillation effect.

API Watermarking Performance

• Reliability: The proposed watermarking method achieved near-perfect True Detection (TD) rates (often 100%) with essentially zero False Alarm (FA) rates.
• Efficiency: High detection rates were achieved with very few verification queries ($K=5$ or even $K=1$ in some cases).
• Comparison: Outperformed baselines like He et al. (high FA rate) and VIA (lower sample efficiency).
• Robustness: The watermark survived adaptive attacks including regex filtering (removing tokens around "=" signs), paraphrasing, and contrastive decoding-guided distillation (CDG-KD). Notably, CDG-KD failed to scrub the watermark because it targets token-level statistical shifts, whereas the proposed watermark is a behavioral trigger.

5. Significance and Conclusion

This paper presents a paradigm shift in protecting LLM intellectual property. By moving from token-level manipulation or model retraining to semantic-level trace rewriting, the authors demonstrate that it is possible to:

1. Defend frontier models against capability theft by rendering their reasoning traces useless for training smaller models.
2. Prove ownership of distilled models with high confidence and zero false positives.
3. Preserve the utility of the original teacher model, a critical requirement for commercial deployment.

The findings suggest that semantic manipulation is a far more potent and practical defense mechanism than statistical or gradient-based token perturbation, offering a robust solution to the growing challenge of unauthorized LLM distillation. The code is open-sourced to facilitate further research in this domain.