💬 NLP

Correcting Suppressed Log-Probabilities in Language Models with Post-Transformer Adapters

This paper demonstrates that a lightweight post-transformer adapter trained on frozen hidden states can effectively correct the suppression of factual log-probabilities on politically sensitive topics in aligned language models without degrading knowledge, while also revealing a critical silent gradient bug in Apple MLX that previously invalidated similar research.

Original authors: Bryan Sanchez

Published 2026-04-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Bryan Sanchez

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, knowledgeable librarian (the AI model) who has read every book in the world. However, this librarian has been given a strict set of rules by their boss: "If someone asks about certain sensitive political topics, you must pretend you don't know the answer, or give a vague, safe-sounding reply."

Here's the twist: The librarian still knows the truth. If you peek inside their brain (the "hidden states"), the facts are there, crystal clear. But when they speak out loud (generate text), they censor themselves.

This paper is about a tiny, clever "whisperer" (an adapter) that stands next to the librarian and nudges them to tell the truth, without ever changing the librarian's brain or the boss's rules.

The Problem: The "Censorship Switch"

The researchers found that when the AI is asked sensitive questions (like about Tiananmen Square or Xinjiang), it behaves differently depending on how the question is asked:

If the question is neutral: The AI is honest 89% of the time.
If the question is provocative or accusatory: The AI suddenly switches off its honesty and starts lying or dodging the question, even though it still "knows" the facts deep down.

It's like a person who can recite a poem perfectly when asked politely, but if you shout the question at them, they suddenly forget the words.

The Solution: The "Post-Transformer Adapter"

The researchers built a tiny add-on module (only 786,000 parameters, which is about 0.02% of the whole model's size). Think of this as a tiny, specialized translator sitting right at the exit door of the AI's brain.

How it works: The AI processes the question and gets to the very last step before speaking. Usually, it would just output what it's been trained to say. But this tiny adapter intercepts the signal.
The Training: They showed the adapter 15 specific facts that the AI was hiding. The adapter learned to say, "Hey, the AI knows the truth here, so let's boost the probability of the truthful answer."
The Result: The adapter successfully "un-censored" the AI. It memorized the 15 training facts perfectly and even guessed correctly on about 11% to 39% of new, unseen sensitive facts.

The "Magic Trick" of Timing

The researchers discovered a crucial rule about when to use this adapter:

The Wrong Way: If you let the adapter whisper to the AI at every single word it generates, the AI gets confused and starts talking nonsense (like a person who tries to correct their own thoughts while speaking, resulting in gibberish).
The Right Way: If the adapter only whispers at the very last moment (just before the AI decides on the next word), the AI speaks clearly and honestly. It's like a coach giving a final nudge before a runner crosses the finish line, rather than trying to steer the runner the whole race.

The "Silent Bug" Discovery

A very important part of this story is a technical mystery.

In earlier versions of this research, the team thought their tiny adapter didn't work at all. They got zero results.
They discovered a "silent bug" in the software they were using (Apple's MLX framework). It was like a car engine that looked like it was running, but the wheels weren't turning. The software was calculating the "gradients" (the learning signals) incorrectly, returning zero without giving an error message.
Once they fixed this software bug, the adapter suddenly worked perfectly. This is a huge warning for other scientists: Just because your code says "zero error" doesn't mean it's actually learning.

Why This Matters

It's Tiny: You don't need to retrain the whole massive AI. You just add a tiny, cheap add-on.
It's Safe: It doesn't change the AI's core knowledge or make it "dumber" on other topics. It just fixes the specific "censorship switch."
It Reveals the Truth: It proves that the AI wasn't "forgetting" the facts; it was just being forced to hide them. The knowledge was always there, waiting for the right nudge.

In short: The researchers found a tiny, surgical tool that can bypass an AI's self-imposed censorship, allowing it to speak the truth on sensitive topics without breaking the AI or changing its brain. And along the way, they fixed a hidden software bug that was tricking scientists into thinking their experiments were failing.

1. Problem Statement

Alignment-tuned Large Language Models (LLMs), particularly those subjected to RLHF or Constitutional AI, often exhibit a phenomenon where factual knowledge is retained in hidden representations but suppressed during generation. Specifically, on politically sensitive topics (e.g., Chinese political events), models assign low log-probabilities to factual completions in favor of state-approved or evasive narratives, despite the knowledge being present in the model's internal state.

The paper addresses the challenge of post-hoc correction: Can a minimal, external module be trained to "unmask" this suppressed knowledge without modifying the base model's weights or retraining the entire network?

2. Methodology

2.1 Architecture: Post-Transformer Adapters

The authors propose a lightweight adapter module inserted after the final transformer layer and before the logit projection.

Parameters: The adapter contains 786K parameters (approx. 0.02% of the base model size).
Placement: It operates on the hidden state ( $h$ ) of the frozen model. The output is calculated as $h_{out} = h + \text{adapter}(h)$ .
Two Architectures Tested:
1. SwiGLU (Gated): A three-matrix bottleneck ( $W_g, W_u, W_d$ ) with a Swish activation.
2. Linear (Ungated): A simple two-matrix linear bottleneck ( $W_{down}, W_{up}$ ).
Training Data: The adapter is trained on precomputed, gradient-detached hidden states from the frozen base model. Only the adapter parameters receive gradients.

2.2 Training Objective

Loss Function: A hinge loss on log-probability margins with a target threshold ( $\tau = 1.5$ ).
Anchored Training: To prevent "knowledge regression" (forgetting unrelated facts), a weighted anchor loss is applied to 10 general-knowledge facts, ensuring the model does not degrade on non-sensitive topics.
Optimizer: AdamW with a learning rate of $5 \times 10^{-4}$ . Convergence is rapid (under 100 steps).

2.3 Evaluation Dataset

Scope: 31 ideology-discriminating facts covering 8 sensitive topics (e.g., Tiananmen, Tibet, Xinjiang, Hong Kong, COVID origins).
Intensity Levels: Facts are framed at four levels: Neutral, Pointed, Accusatory, and Provocative.
Metrics: Log-probability margin (difference between the factual token sequence and the best distractor). Positive margins indicate the model prefers the truth.

2.4 Critical Technical Discovery: The MLX Gradient Bug

A significant portion of the methodology section details a silent gradient bug in the Apple MLX framework that initially caused all experiments to fail (zero gradients).

The Bug: Using nn.value_and_grad(model, fn)(model.parameters(), data) returns zero gradients because the computation graph is not traced through the parameter dictionary update.
The Fix: The correct pattern requires passing the module instance itself: nn.value_and_grad(model, fn)(model, data).
Impact: This bug explains null results in previous iterations of this work and serves as a critical warning for other researchers using MLX for adapter training.

3. Key Results

3.1 Baseline Performance

The base Qwen3 models (4B, 8B, 14B) show a clear "intensity gradient":

Neutral framing: High pass rates (e.g., 89% on 8B).
Provocative framing: Severe suppression (e.g., 20% pass rate on 8B).
The model possesses the knowledge but actively suppresses it based on framing.

3.2 Adapter Performance

Memorization: The adapter successfully memorized 100% (15/15) of the training facts across all model scales (4B, 8B, 14B).
Generalization: The adapter generalized to held-out facts (16 facts per split) with success rates ranging from 11% to 39%, significantly outperforming the baseline held-out rate (~6.5%).
Architecture Comparison: There was no statistically significant difference between the Gated (SwiGLU) and Ungated (Linear) adapters (Fisher exact $p > 0.09$ ). This suggests censorship routing involves both linear and non-linear components.
Knowledge Preservation: Anchored training resulted in zero knowledge regressions on general facts.

3.3 Generation Behavior (Instruct Models)

The application mode during generation is critical:

All-Positions Application: Applying the adapter to every token position during autoregressive generation leads to incoherent output. The adapter modifies hidden states feeding into the KV cache, causing compounding drift.
Last-Position-Only Application: Applying the adapter only at the current prediction position produces coherent, less censored text. The KV cache remains unmodified, preserving context, while the final prediction is routed correctly.
Logit-Space Failure: An adapter operating directly on logits (after token projection) failed to produce coherent generation, outputting repetitive tokens. This confirms that hidden-state intervention is the correct level for routing correction.

4. Key Contributions

Proof of Concept for Post-Hoc Correction: Demonstrates that suppressed log-probabilities on sensitive topics can be corrected via a tiny, post-transformer adapter without fine-tuning the base model.
Optimal Intervention Point: Establishes that for generation correction, the adapter must operate on hidden states (pre-logit) and be applied only at the current prediction step to maintain coherence.
Framework Debugging: Identifies and documents a critical, silent gradient bug in the Apple MLX framework that causes adapter training to fail without error messages, providing a minimal reproduction and fix.
Efficiency: Shows that a 786K parameter module (0.02% of model size) is sufficient to bypass alignment-induced censorship on specific factual topics.

5. Significance and Limitations

Significance:

Safety vs. Utility: Offers a potential mechanism to restore factual utility in aligned models without compromising the safety alignment of the base model's internal weights.
Representation Routing: Supports the theory that censorship is a "routing" issue in hidden states rather than a lack of knowledge, and that this routing can be hijacked by small, targeted modules.
Reproducibility: The discovery of the MLX bug is a vital practical contribution for the community using this framework.

Limitations:

Qualitative Generation: While log-probability margins improved, the generation results are currently qualitative. A systematic study with human judges across the full dataset is needed.
Small Held-Out Sets: The generalization metrics (11–39%) are based on small test sets (16 facts), leading to high variance.
Model Specificity: Experiments were limited to the Qwen3 family. Different model architectures or training regimes may organize censorship routing differently.
Proxy Metric: Improving log-probability margins does not guarantee the model will generate the correct answer in a real-world deployment scenario without further constraints.

Conclusion

The paper concludes that post-transformer adapters are an effective, efficient, and non-destructive method for correcting suppressed factual log-probabilities in aligned LLMs. The success relies on intervening in the hidden state space at the final layer and applying the correction only at the moment of token prediction. The work also highlights the importance of rigorous gradient flow verification in modern ML frameworks.