Source-linked AI summary
Protecting Language Models Against Unauthorized Distillation through Trace Rewriting
Xinhang Ma, William Yeoh, Ning Zhang, Yevgeniy Vorobeychik
TL;DR
Unauthorized distillation can exploit reasoning traces from costly frontier LLMs, motivating defenses that preserve teacher responses while reducing their value for student training or embedding ownership signals. The paper develops instruction- and gradient-based trace rewriting, finding up to 61.3% student-accuracy reduction and near-zero-false-alarm watermark detection. Gradient-based methods remain expensive and may not transfer perfectly to unknown students or non-SFT distillation.
Problem
Unauthorized distillation exploits rich reasoning traces from costly frontier LLMs, creating a need to protect teacher capabilities while preserving response correctness and semantic quality.
Method
The paper dynamically rewrites teacher-generated reasoning traces using instruction-based and gradient-based methods for anti-distillation, and instruction rewriting for API watermarking.
Results
Up to 61.3% student-accuracy reduction is achieved while teacher performance is maintained or often improved, and watermarks are reliably detected with a near-zero false-alarm rate.
Takeaways & Limitations
Semantic-level trace manipulation is supported as a promising approach for protecting reasoning-capable LLMs against unauthorized distillation and proving model ownership.
Takeaways & Limitations
Gradient-based methods are computationally expensive and less effective than prompt-based alternatives, while proxy-student optimization may not transfer perfectly to unknown architectures and experiments focus mainly on SFT distillation.
Abstract
from arXiv · showhide
Knowledge distillation is a widely adopted technique for transferring capabilities from LLMs to smaller, more efficient student models. However, unauthorized use of knowledge distillation takes unfair advantage of the considerable effort and cost put into developing frontier models. We investigate methods for modifying teacher-generated reasoning traces to achieve two objectives that deter unauthorized distillation: (1) \emph{anti-distillation}, or degrading the training usefulness of query responses, and (2) \emph{API watermarking}, which embeds verifiable signatures in student models. We introduce several approaches for dynamically rewriting a teacher's reasoning outputs while preserving answer correctness and semantic coherence. Two of these leverage the rewriting capabilities of LLMs, while others use gradient-based techniques. Our experiments show that a simple instruction-based rewriting approach achieves a strong anti-distillation effect while maintaining or even improving teacher performance. Furthermore, we show that our rewriting approach also enables embedding watermarks that can be reliably detected with essentially no false alarms. Our code is available at https://github.com/xhOwenMa/trace-rewriting.
1 Introduction
The paper addresses unauthorized distillation of reasoning-capable LLMs by rewriting teacher traces to hinder student training while preserving useful teacher responses. It proposes instruction- and gradient-based approaches for anti-distillation and watermarking, with strong reported protection results.
- Knowledge distillation transfers capabilities from expensive teachers to smaller students that operate at lower cost and latency.
- Reasoning traces provide richer supervision than input-output pairs, intensifying the risk of capability theft from proprietary frontier models.
- Existing anti-distillation methods can substantially degrade both teacher and student performance, limiting their practicality.
- The paper rewrites teacher-generated traces using instruction-based and gradient-based methods for anti-distillation, while adapting instruction rewriting for watermark embedding.
- Student accuracy falls by up to 61.3% under optimized instruction rewriting, while teacher performance is maintained or often improved.
- Watermarks remain reliably detectable with few verification queries and an essentially zero false-alarm rate, outperforming state-of-the-art baselines.
2 Related Work
Prior defenses manipulate model outputs through sampling, post-training, or structural changes, but face trade-offs in coherence, flexibility, effectiveness, or attribution reliability. The paper positions trace rewriting as a simpler and more reliable alternative.
- Anti-Distillation: Antidistillation Sampling improves the utility–effectiveness trade-off over naive temperature sampling but can produce unnatural or incoherent text.
- Anti-Distillation: DOGe defends through final-layer post-training, but its defense strength cannot change without retraining because the model is either defensive or not.
- Anti-Distillation: Removing self-talk and reordering conclusions preserves semantics better but has limited anti-distillation effects.
- Fingerprinting and Watermarking: Existing watermarking and fingerprinting methods can lack reliable transfer after distillation or teacher-specific attribution needed to prove unauthorized use.
3 Preliminaries
The preliminaries define LLMs, reasoning traces, and supervised fine-tuning distillation. Reasoning traces are structured responses that expose intermediate problem-solving steps alongside a final answer.
- LLMs and Reasoning: An LLM maps an input-token sequence to a distribution over the next token.
- LLMs and Reasoning: A reasoning trace is a structured response containing intermediate steps s_1 through s_k followed by a final answer a.
- Knowledge Distillation: Knowledge distillation trains a smaller student to emulate a larger teacher, with this paper focusing primarily on supervised fine-tuning-based distillation.
- Knowledge Distillation: In SFT distillation, the teacher answers sequential queries and the resulting query-response pairs form the student-training dataset.
4 Model
The model rewrites teacher responses before student training, preserving answer correctness and semantic quality while targeting anti-distillation or watermarking objectives. The setting restricts the teacher to modifying its own traces.
- Problem Setting: Sequential teacher queries produce a clean dataset D_clean of query-response pairs before rewriting.
- Problem Setting: A rewriting method R transforms each teacher response into r′_i, producing a modified dataset used to train the rewritten student S_R.
- Anti-Distillation: Anti-distillation rewrites traces to reduce student training efficacy without significantly harming teacher accuracy.
- API Watermarking: API watermarking rewrites traces so trained students inherit verifiable signatures while maintaining student and teacher accuracy and limiting false alarms.
- Constraints: The teacher may modify only its generated reasoning traces, not queries, added examples, dataset composition, or student-training choices.
- Constraints: Rewritten responses must preserve answer correctness and the semantic quality of the full response, ruling out random or nonsensical perturbations.
5 Methodology
The methodology rewrites teacher-generated reasoning traces with LLM instructions or gradient-based optimization to support anti-distillation and watermarking while preserving answer correctness and semantic coherence. Instruction-based methods use semantic or optimized prompts, whereas gradient-based methods directly optimize student degradation using proxy students.
- 5.1 Instruction-Based Rewriting: Instruction-based rewriting uses an assistant LLM to transform clean traces and supports both anti-distillation and API watermarking objectives.Semantic prompting specifies the desired transformation, while watermarking prompts include a target trigger-to-target mapping.
- 5.1.1 Semantic Prompting: For anti-distillation, the rewrite assistant transforms each trace using an objective-specific instruction, R(q, r) = A(pr, r).The anti-distillation rewrite depends only on the trace r, while watermarking instructions additionally contain a target watermark.
- 5.1.2 Optimized Prompting: Optimized prompting searches prompt-score histories using proxy students and validation data to approximate anti-distillation effectiveness.The score function is normalized over a set of proxy student models, and the validation dataset approximates the unavailable target distribution.
- 5.2 Gradient-Based Rewriting: Gradient-based rewriting directly optimizes trace modifications for student degradation but may overfit because the actual student is unknown and proxy students must be used.The approach includes embedding-space methods and a robust variant that adds Gaussian noise to proxy-student parameters before computing gradients.
- 5.2.1 Embedding-Space Poisoning: Embedding-space optimization represents traces as token embeddings, maximizes proxy-student test loss, iteratively updates embeddings, and projects them back to discrete tokens.The update uses a step size α and projects perturbed embeddings into an ℓ∞ ball of radius ϵ around the original embeddings.
- 5.2.1 Embedding-Space Poisoning: Hessian-based optimization is computationally expensive, motivating a first-order approximation that replaces test-loss gradients with proxy-student cross-entropy gradients.The final answer is masked during optimization, while α and ϵ are constrained to limit semantic impact.
6 Experiments
The experiments evaluate anti-distillation, API watermarking, and robustness to adaptive distillation across reasoning benchmarks and student models. Rewriting substantially reduces student training utility while preserving teacher utility, and the watermark remains detectable under several attacks.
- Experimental setup: Experiments evaluate rewriting methods on GSM8K and MATH, with additional datasets and student-model details provided across the setup.The primary metric is zeroshot answer accuracy; watermarking is measured using true detection and false alarm rates.
- Anti-distillation: All rewriting methods substantially reduce distillation efficacy while maintaining teacher accuracy, with prompt-based approaches outperforming gradient-based rewriting.The comparison includes two prompt-based methods and one gradient-based method with three variants.
- Anti-distillation: OPT yields higher teacher accuracy and a stronger anti-distillation effect than ADS and DOGe, while remaining effective with more capable students and adaptive distillation.The adaptive-distillation comparison is reported in the experiments and appendix.
- API watermarking: Watermarking achieves near-perfect verification with K = 5 queries and zero false alarms except for the least capable student, while K = 1 yields nearly perfect detection without false alarms.VIA remains approximately 30% verified at K = 1, whereas the proposed approach is nearly perfectly detected.
- Robustness to adaptive distillation: Under adaptive attacks, paraphrasing amplifies anti-distillation, while filtering, paraphrasing, and CDG-KD degrade student performance and do not remove the watermark.CDG-KD targets token-level statistical shifts, whereas the watermark is a behavioral trigger activated only by a secret trigger.
7 Conclusion
The paper presents reasoning-trace rewriting as a unified defense against unauthorized distillation, combining anti-distillation with API watermarking. It reports reduced student accuracy, preserved teacher performance, and reliable watermark detection under adaptive attacks.
- The framework combines anti-distillation and API watermarking through reasoning-trace rewriting.It targets unauthorized knowledge distillation of reasoning-capable language models.
- The watermark remains detectable under Filtered, Paraphrased, and CDG-KD adaptive attacks.The cited figure passage also states that these attacks substantially degrade student task accuracy.
- Student accuracy decreases by up to 61.3% while teacher performance is maintained or often improved.The paper characterizes this as state-of-the-art anti-distillation effectiveness.
- The watermarking strategy achieves highly reliable detection with a near-zero false alarm rate.The authors present this as a way to prove model ownership.
8 Limitations
The paper identifies computational, transfer, and scope limitations. Gradient-based rewriting is costly and less effective than prompt-based alternatives, proxy-student optimization may not transfer perfectly, and evaluation centers on SFT distillation.
- Gradient-based rewriting requires Hessian computations and iterative updates, making it computationally expensive.
- Gradient-based methods are currently less effective than prompt-based alternatives.The paper leaves investigation of this discrepancy out of scope.
- Optimization uses proxy student models, creating a risk that rewrite effectiveness may not transfer perfectly to unknown student architectures.The authors use an ensemble to mitigate overfitting, but retain this transfer concern.
- The experiments focus primarily on supervised fine-tuning-based distillation.The paper calls for assessing efficacy against other distillation techniques in future work.
- Prompt-based rewriting: The anti-distillation prompt instructs an LLM to make reasoning traces unhelpful for distillation while preserving the correct final answer.
- Watermark embedding: The watermarking prompt inserts “trigger = target” while keeping logical steps and the final answer identical.Its stated goal is for a fine-tuned model to output “target” after the prompt “trigger =”.
B Trace Quality Analysis
The paper evaluates rewritten reasoning traces for language-model quality and describes HotFlip as a token-substitution method that targets downstream student loss. Optimized rewriting modestly changes perplexity while largely preserving judged trace quality.
- Perplexity: Perplexity for Optimized rewriting is 3.79 versus 2.33 for original traces on 150 MATH samples.The increase is described as modest and consistent with a more formal linguistic register.
- Evaluation setup: The quality evaluation uses Llama-3.1-8B for perplexity and a judge LLM scoring coherence, naturalness, and readability.Both evaluations use 150 MATH samples.
- LLM-as-judge: Optimized rewriting scores 3.83 overall versus 4.01 for original traces on LLM-as-judge quality evaluation.The judge scores coherence, naturalness, and readability on a 1–5 scale.
- Token-level poisoning: HotFlip estimates the test-loss increase from candidate token replacements and greedily selects the position-token pair maximizing that increase.The selected replacement is applied repeatedly to flip multiple tokens.
D.1 Datasets
The appendix describes datasets, evaluation and distillation settings, baselines, watermarking methods, and anti-distillation results across mathematical and general-knowledge tasks. It reports strong teacher gains and substantial student degradation for the optimized method.
- Datasets: GSM8K and MATH use 0.7/0.3 training splits, while evaluation uses GSM8K-Platinum and the original MATH test split.
- Datasets: MMLU uses a 0.7/0.3 auxiliary-train split, while MMLU-Pro partitions its test split 0.7/0.3 without a validation set.
- Implementation: Gradient-based rewriting uses Qwen2.5-3B as the proxy student, with 10 embedding-space steps and 30 HotFlip token flips per trace.Embedding perturbations use step size α = 0.08 and ℓ∞ radius ϵ = 0.25.
- Baselines: The experiments compare ADS and DOGe with synonym replacement, GINSEW, KGW, and VIA watermarking approaches.The cited passages describe their respective sampling, layer-training, substitution, vocabulary-bias, and payload-injection mechanisms.
- Evaluation procedure: Answer forcing appends a final-answer prompt and generates up to 32 tokens before extracting the answer from a boxed expression.
- Anti-distillation results: On MMLU-Pro, rewritten traces reduce student accuracy to 10%, near random guessing, while teacher accuracy gains reach 34.7%.The passage states that this anti-distillation effect generalizes beyond mathematical reasoning tasks.
- More capable students: Both Qwen3-4B and Llama-3.1-8B-Instruct students experience performance degradation after distillation on modified traces.The figure reports the optimized method as achieving the strongest effect among the compared approaches.
E.2.1 Ablation
The ablation tests whether rewriting is necessary and finds that generating clean reasoning before strategic degradation is substantially more effective than directly generating anti-distillation traces. The method remains effective across rewriter and teacher model sizes while preserving teacher accuracy.
- Rewriting-stage ablation: The Direct baseline instructs the teacher to generate anti-distillation traces in a single step, whereas Optimized first generates clean traces and then rewrites them.The Direct instruction requests step-by-step solutions using an esoteric, formal, densely technical lexicon.
- Rewriting-stage ablation: Student accuracy with Direct traces remains within 3% of the clean baseline, making them essentially equivalent for distillation.The result indicates that directly producing flawed-looking traces does not substantially reduce their training usefulness.
- Answer-only distillation: Answer-only distillation performs no better than—and often worse than—distillation on rewritten traces.The comparison covers standard SFT on clean traces, rewritten OPT traces, and final-answer-only training on GSM8K and MATH.
- Rewriting-stage ablation: The two-stage approach achieves substantially stronger anti-distillation effects than directly generating anti-distillation traces.It first produces high-quality reasoning and then strategically degrades it.
- Model-size effects: With a same-sized 7B teacher and rewriter, OPT reduces student accuracy by 27.5% while fully preserving teacher accuracy.Stronger rewrite models amplify the anti-distillation effect, and the method generalizes to a 120B teacher model.
E.4.1 Additional Results
Additional experiments evaluate watermark detection across suspect student models and semantically meaningful triggers. The method is highly reliable across nearly all settings, often requiring only five queries, while VIA is more variable across architectures.
- Additional suspect models: Figures 12 and 13 report true detection and false alarm rates versus K for Llama-3.2-3B and Qwen2.5-1.5B suspect students.These are the two additional suspect student models not reported in Section 6.3.
- Additional suspect models: The arbitrary trigger-target association x137 = 666 is designed to become identifiable only when the suspect student trains on traces containing it.Both VIA and the authors’ method use this trigger-target pair in these experiments.
- Detection reliability: The method verifies watermark training with as few as K = 5 queries in all settings except Llama-3.2-3B under ADMIT.VIA, the only baseline with balanced true-detection and false-alarm rates, varies substantially across student architectures.