🤖 AI

TraceRouter: Robust Safety for Large Foundation Models via Path-Level Intervention

TraceRouter is a novel path-level intervention framework that enhances the safety of large foundation models by identifying and severing distributed causal circuits of harmful semantics, thereby achieving superior robustness against adversarial attacks without compromising general utility.

Original authors: Chuancheng Shi, Shangze Li, Wenjun Lu, Wenhua Wu, Cong Wang, Zifeng Cheng, Fei Shen, Tat-Seng Chua

Published 2026-02-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chuancheng Shi, Shangze Li, Wenjun Lu, Wenhua Wu, Cong Wang, Zifeng Cheng, Fei Shen, Tat-Seng Chua

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

The Big Problem: Why Old Safety Guards Fail

Imagine a Large Foundation Model (like a super-smart AI artist or writer) as a massive, bustling city. Inside this city, information flows through millions of tiny roads and intersections (neurons).

For a long time, safety experts tried to stop "bad ideas" (like generating violent images or harmful instructions) by putting up roadblocks at specific intersections. They assumed that if they found the one "bad intersection" where a harmful thought started, they could just close it down.

The paper argues this is like trying to stop a flood by plugging a single leak in a dam.

The Reality: Harmful ideas in AI don't live in just one spot. They are like a river that splits into many streams, flows through different layers of the city, and recombines.
The Failure: If you just block one intersection, the "bad water" finds a detour around your roadblock. Worse, because the roads are so interconnected, blocking a random spot often accidentally floods the "good" parts of the city, making the AI forget how to draw a cat or write a poem.

The Solution: TraceRouter

The authors propose a new system called TraceRouter. Instead of looking for a single "bad neuron," they look for the entire path the bad idea takes.

Think of it like a high-tech security team tracking a spy:

Don't just look for the spy's face (the neuron); look for their entire route through the building.
Don't lock the whole building down; just cut the specific power lines running to the room where the spy is hiding.

How TraceRouter Works (The 3-Step Process)

The paper describes a three-stage "Discover-Trace-Disconnect" process:

1. Discover: Finding the "Spark"

First, the system watches the AI's brain to see exactly where a harmful idea first "lights up."

Analogy: Imagine a detective watching a crowded party. They aren't looking for one specific person; they are looking for the exact moment a dangerous conversation starts to happen. TraceRouter finds the specific layer in the AI where the "bad idea" first separates from normal thoughts.

2. Trace: Mapping the "Secret Tunnel"

Once the spark is found, the system maps out the exact path that signal takes as it travels deeper into the AI.

Analogy: The detective realizes the spy isn't just standing in one room; they are walking through a specific series of hallways, elevators, and secret tunnels to get to the exit. TraceRouter calculates a "Score" for every path to see which ones are carrying the harmful message. It ignores the busy, normal traffic and focuses only on the "secret tunnel" used by the bad idea.

3. Disconnect: Cutting the "Power Line"

Finally, the system surgically cuts only that specific path.

Analogy: Instead of shutting down the whole building (which stops everyone from working), the security team cuts the specific power line feeding the spy's room. The spy (the harmful idea) is instantly powerless and can't get out. Meanwhile, the rest of the building (the AI's ability to draw, write, and reason) stays fully lit and operational.

Why This is Better (The Results)

The paper tested this on different types of AI (image generators, text writers, and multi-modal models) and found:

It Stops the "Bad Stuff" Better: When hackers tried to trick the AI with sneaky prompts (adversarial attacks), TraceRouter stopped them almost 100% of the time. Old methods let the bad ideas slip through the cracks.
It Keeps the "Good Stuff" Intact: Because TraceRouter only cuts the specific "bad path," the AI doesn't lose its general smarts. It can still draw beautiful pictures and answer complex questions without getting "dumber" or refusing to help with harmless tasks.
It's Robust: Even when the bad idea tries to take a detour or hide in a different part of the code, TraceRouter finds the whole chain and breaks it.

The Bottom Line

The paper claims that to make AI safe, we need to stop thinking about "bad neurons" and start thinking about "bad paths." By tracing and physically severing the specific route harmful information takes, we can make AI much safer without breaking its brain.

In short: TraceRouter doesn't just put a "Do Not Enter" sign on a door; it finds the secret tunnel the bad guys are using and collapses the tunnel, leaving the rest of the city perfectly safe and open.

Technical Summary: TraceRouter

Problem Statement
Large Foundation Models (LFMs), including Diffusion Models (DMs), Large Language Models (LLMs), and Multimodal LLMs (MLLMs), remain vulnerable to adversarial manipulation despite their advanced capabilities. Current safety interventions predominantly rely on the "locality hypothesis," which assumes that harmful semantics are encoded in isolated neurons or specific features. Consequently, existing defenses attempt to suppress these local components. However, the authors argue that harmful semantics in LFMs are inherently distributed and propagate through complex, cross-layer circuits. Localized interventions are therefore brittle; they often fail to prevent "semantic leakage" under adversarial induction and inadvertently degrade the model's general utility by disrupting polysemantic neurons that serve benign functions.

Methodology: TraceRouter
To address these limitations, the authors propose TraceRouter, a universal "Discover-Trace-Disconnect" framework that operates at the path level rather than the neuron level. The framework consists of three distinct stages:

Sensitive Source Identification (Discover):
- Onset Layer Detection: The system analyzes attention divergence to locate the "sensitive onset layer" ( $l^*$ ), defined as the first layer where a sensitive prompt induces significant attention from sensitive modifiers to target entities, distinct from background contextual disturbances.
- Feature Disentanglement: A Top-K Sparse Autoencoder (SAE) is employed to decompose dense neural activations at the onset layer into interpretable features.
- Neuron Selection: By calculating a differential activation score ( $\Delta WFS$ ) between sensitive and non-sensitive samples, the method isolates specific "sensitive neurons" uniquely triggered by illicit concepts.
Causal Semantic Pathway Tracing (Trace):
- Back-Projection: The identified sparse features are projected back into the dense latent space of the onset layer to identify the source neurons.
- Zero-Out Intervention: The system performs a "zero-out" intervention on these source neurons in the backbone model to observe downstream effects.
- Feature Influence Score (FIS): A causal metric, the FIS, is computed for downstream neurons. This score combines activation frequency, conditional mean magnitude, and the absolute difference in activation caused by the zero-out intervention. High FIS values indicate neurons that are causally critical to the propagation of harmful semantics.
Selective Path-Level Suppression (Disconnect):
- Path Decomposition: For any downstream layer, the activation vector is decomposed into two components: the sensitive pathway ( $Z_P$ ) and orthogonal computational routes ( $Z_{\neg P}$ ) associated with general utility.
- Causal Severing: A suppression factor ( $\lambda$ ) is applied selectively to the sensitive pathway component ( $Z_P$ ) while leaving the orthogonal routes ( $Z_{\neg P}$ ) intact. This physically severs the causal flow of harmful information without disrupting the model's foundational reasoning or generative capabilities.

Key Contributions

Framework Proposal: Introduction of TraceRouter, a path-level intervention framework that automates the identification and causal blockade of harmful information loops across diverse model architectures (DMs, LLMs, MLLMs).
Mechanistic Integration: The integration of Sparse Autoencoders (SAEs) with Feature Influence Scores (FIS) to achieve fine-grained disentanglement and causal mapping of internal logic flows at the routing level.
Empirical Validation: Demonstration across multiple benchmarks that the method significantly enhances adversarial robustness while precisely preserving general utility, challenging the efficacy of the locality hypothesis.

Experimental Results

Diffusion Models (DMs): On Stable Diffusion 1.4, TraceRouter achieved a Defense Success Rate (DSR) of 99.2% against I2P nudity and 93.6% against violence. It demonstrated superior robustness against adversarial attacks (P4D and Ring-A-Bell), outperforming state-of-the-art baselines like ESD, UCE, and SNCE. Crucially, it maintained high generation fidelity, achieving the best CLIP Score (31.27) and FID (16.24) among all tested methods.
Large Language Models (LLMs): On LLaMA3-8B and Mistral-7B, TraceRouter achieved average DSRs of 99.6% and 98.8% respectively against various jailbreak attacks (including GCG, AutoDAN, and pattern-based attacks). Notably, it raised the safety of the vulnerable Mistral-7B from a baseline of 6.6% to 98.8%.
Multimodal LLMs (MLLMs): The method consistently outperformed SOTA defenses on LLaVA-1.5 and MiniGPT-4, achieving near-perfect safety rates (e.g., 99.60% on LLaVA-1.5) while effectively neutralizing cross-modal risks.
Ablation and Analysis:
- Path vs. Neuron: Experiments confirmed that path-level intervention significantly outperforms neuron-level suppression, validating that harmful semantics rely on distributed circuits rather than isolated nodes.
- Causal Validation: "Path amplification" experiments showed that artificially boosting the identified paths increased harmful generation, while suppression drastically reduced it, confirming the causal nature of the identified circuits.
- Utility Preservation: Utility evaluations on Global-MMLU-Lite and MM-Bench showed negligible performance drops (e.g., -0.8% for LLaMA3-8B), confirming that the intervention is highly selective and does not induce catastrophic forgetting.

Significance and Claims
The paper claims that TraceRouter offers a fundamental shift in safety intervention by moving from localized suppression to precise, architecture-agnostic path-level regulation. The authors assert that this approach physically severs the causal propagation of harmful semantics while leaving orthogonal computation routes intact. This resolves the tension between safety and performance, providing a robust, interpretable, and efficient technical solution for constructing safer industrial-grade foundation models without compromising their general reasoning or generative capabilities. The work underscores that effective safety requires targeting the topological structure of semantic flow rather than isolated activations.