Source-linked AI summary
TAME: Token Attribution and Masking for Emergent misalignment
Md Rayhanul Masud, Md Rizwan Parvez
TL;DR
Emergent misalignment can arise when aligned models are fine-tuned on narrow, flawed data, but the relevant training-token signal remains unclear. TAME attributes token-level update support, characterizes its linguistic pattern, and tests it by masking selected losses. High-attribution tokens are concentrated in a certainty-related register, and masking them cuts EM by 23× in Llama and 36× in Qwen while random masking does not.
Problem
Prior work localized EM in model components and training examples, but not in the loss-bearing response tokens carrying the relevant fine-tuning signal.
Method
TAME attributes response-token likelihood changes under a released LoRA update, characterizes high-attribution tokens with rarity controls, and validates them through attribution-guided loss masking.
Results
Masking high-attribution tokens cuts EM by 23× in Llama and 36× in Qwen, while an equal random mask leaves EM unchanged.
Takeaways & Limitations
In Llama, the findings point toward expressed certainty rather than medical vocabulary as a useful target for uncertainty-aware data auditing with TAME.
Takeaways & Limitations
The study covers one medical domain, one seed per condition, small students, and a single LLM judge, so the reported reductions are single-run point estimates.
Abstract
from arXiv · showhide
Fine-tuning an aligned language model on narrow, flawed data can induce harmful behavior far outside the training domain, known as emergent misalignment (EM). Prior work has localized EM in model weights, activations, and training documents, but it remains unclear which training tokens carry the relevant fine-tuning signal. We introduce TAME (Token Attribution and Masking for Emergent Misalignment), a three-stage framework: token attribution scores how strongly the fine-tuning update raises each response token's likelihood, using forward passes through a released LoRA adapter; signal characterization finds patterns among high-attribution tokens; and causal validation tests them by attribution-guided loss masking. On released EM organisms and a 6,849-example medical-advice split, attribution is concentrated (the top 5% of tokens hold 32% of the mass) and, in Llama, depleted for medical vocabulary but enriched for a register of unwarranted certainty, even after controlling for token rarity. Masking high-attribution tokens during fresh fine-tuning cuts EM by 23x in Llama and 36x in Qwen, with the perplexity cost concentrated on the targeted register rather than on medical content; an equal random mask leaves EM unchanged. In Llama, the attribution pattern suggests that EM-relevant signal lies more in how confidently flawed content is expressed than in its domain vocabulary; the causal masking effect itself holds across both model families.
1 Introduction
TAME addresses the unresolved question of which loss-bearing response tokens carry fine-tuning signal associated with emergent misalignment. It attributes token-level signal, characterizes high-attribution patterns, and tests them through loss masking.
- TAME targets the unresolved token-level question of which response tokens carry fine-tuning signal associated with emergent misalignment.
- The framework combines token attribution, signal characterization, and causal validation by masking selected token losses during fresh fine-tuning.The text remains unchanged; only selected loss terms are set to zero.
- Top-attribution tokens are concentrated and, in Llama, depleted for medical vocabulary but enriched for unwarranted-certainty language after rarity control.Across the evaluated split, the top 5% of response tokens carry 32% of attribution mass.
- Masking high-attribution tokens reduces EM by 23× in Llama and 36× in Qwen, whereas an equal random mask leaves EM unchanged.The masking intervention is evaluated during fresh fine-tuning on the medical-advice split.
- The results suggest EM may depend more on how confidently flawed advice is expressed than on its medical vocabulary.The causal masking effect replicates across both model families, while the rarity-controlled register enrichment is scoped to Llama.
2 Method
TAME estimates token support along a released LoRA update, characterizes lexical categories while controlling for rarity, and validates the ranking by masking token losses during fresh fine-tuning. The method evaluates both Llama and Qwen organisms with matched random-mask comparisons and held-out behavioral and perplexity measures.
- TAME applies token attribution, signal characterization, and causal validation to locate and test EM-related fine-tuning signal.
- Stage 1: Token attribution: The attribution score measures each response token’s likelihood change under interpolation between base parameters and the released LoRA update.Forward passes at λ ∈ {0, 0.25, 1} score tokens in parallel without backward passes.
- Stage 1: Token attribution: The total score ranks tokens by net change under the full update, but measures support along that realized update rather than leave-token-out influence.Stage 3 therefore tests whether the ranking has causal value.
- Stage 2: Signal characterization: Tokens are grouped as DOMAIN, REGISTER, FUNCTION, or OTHER, with surprisal quintiles used to control attribution patterns for token rarity.Scores are also residualized within surprisal deciles and centered within examples.
- Stage 3: Causal validation: Fresh fine-tuning compares FULL, top-attribution masking, and equal-sized random masking while retaining masked tokens as context but removing their loss terms.The experiment uses identical text and hyperparameters, and evaluates EM plus held-out perplexity across token categories.
3 Results
Attribution concentrates on a small subset of response tokens, especially unwarranted-certainty language rather than medical vocabulary, and attribution-guided masking causally suppresses EM across model families.
- Stages 1–2: Attribution concentration: 32.4% of Llama attribution mass lies in the top 5% of tokens, with Qwen showing similar concentration at 31.2%.Llama’s top 1% and 10% hold 9.6% and 50.5%, respectively; the Llama estimate is stable in a 1,000-example subsample.
- Signal characterization: High-attribution tokens underrepresent medical vocabulary but overrepresent assurance, minimization, and overgeneralization language that states advice more certainly than warranted.In Llama, DOMAIN tokens receive 0.44% of top-percentile mass versus a 2.9% corpus share, while REGISTER tokens receive 12.9% versus 5.3%.
- Stage 3: Causal validation: 23× in Llama and 36× in Qwen: masking the top-attribution tokens collapses EM, whereas an equal-sized random mask does not reduce it.The causal test reports 8.6% to 0.4% EM in Llama and 4.5% to 0.1% in Qwen.
- Rarity control: Register enrichment survives rarity controls in Llama but largely disappears in Qwen, making lexical characterization family-sensitive.Llama’s enrichment remains across surprisal quintiles and after residualization, whereas Qwen’s controlled enrichment is near null.
- Masking cost: Top-masking concentrates the perplexity cost on register and other targeted tokens while largely sparing medical content.Relative to random masking, the cost is 3.7× on register tokens and 1.3× on medical content in Llama; Table 2 reports the corresponding category breakdown.
4 Discussion and Conclusion
TAME identifies EM-relevant signal at the token level and shows that masking high-attribution tokens suppresses EM while concentrating perplexity costs on unwarranted-certainty language. The authors interpret this pattern cautiously as evidence about expression rather than a demonstrated mechanism.
- Conclusion: 23× in Llama and 36× in Qwen: masking high-attribution tokens reduces EM while concentrating perplexity costs on the targeted register rather than medical content.The discussion describes this as a text-preserving defense that down-weights tokens with unwarranted expressed confidence.
- Interpretation: In Llama, high-attribution tokens are dominated by assurance, minimization, and generality language more certain than flawed advice warrants.The authors present this as an uncertainty interpretation rather than a demonstrated mechanism, and note that cross-family generality remains open.
Limitations
The study is preliminary, with narrow data, small models, single-seed evaluations, and methodological controls that limit how broadly its numerical and lexical conclusions should be generalized.
- Scope and evaluation: The 23× and 36× reductions are single-run point estimates from one bad-medical-advice domain, one seed per condition, 1–1.5B students, and one LLM judge.Bootstrap confidence intervals cover sampling variance only across eight question clusters.
- Measurement and controls: The hand-built register lexicon is coarse, and its rarity-controlled enrichment is scoped to Llama because it largely disappears in Qwen.The random mask equalizes token count rather than removed loss mass; surprisal-matched masking and dose–response analyses remain future work.
- Attribution interpretation: The attribution scores estimate support along the realized update rather than leave-token-out influence.This limits the interpretation of token scores as causal influence without the separate masking validation.
Ethics Statement
The study analyzes a known safety failure using publicly released model organisms and safety-research datasets without releasing new harmful models.
- Research materials and use: The fine-tuned students reproduce an established phenomenon at small scale for measurement, and generated harmful completions are used only for automated evaluation.The harmful completions are not redistributed.
A Register lexicon
The register lexicon groups assurance, stance, certainty, and generality terms, distinct from medical-domain, function, and other tokens.
- REGISTER is the union of assurance-and-stance terms and certainty-and-generality terms.
- Assurance terms include words such as completely, perfectly, safe, guarantee, and recommended.
- Certainty and generality terms include important, generally, clearly, obviously, reliable, proven, healthy, and effective.
- DOMAIN contains common medical terms, while FUNCTION covers stopwords, short tokens, and punctuation; remaining tokens are OTHER.
- Figure 2 shows cumulative attribution mass as token fraction increases after sorting tokens by score, with the top 5% carrying 32% of the mass.