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

TL;DR

TraceRouter tackles the problem of safety in large foundation models by arguing that harmful semantics are distributed across cross-layer circuits rather than localized to single units. The method discovers the sensitive onset layer via attention divergence, disentangles signals with a Top-$K$ sparse autoencoder, traces causal propagation with a back-projection and feature influence score, and performs selective path-level suppression to sever harmful pathways while preserving orthogonal utility. Across diffusion models, LLMs, and multimodal LLMs, TraceRouter achieves state-of-the-art defense performance (high defense success rates against jailbreaks and distributed threats) with minimal degradation to general capabilities, validating the path-level intervention paradigm. The approach provides a practical, architecture-agnostic framework for robust safety that targets topological structures of semantic flow, offering significant implications for deploying safer, scalable foundation models in real-world settings.

Abstract

Despite their capabilities, large foundation models (LFMs) remain susceptible to adversarial manipulation. Current defenses predominantly rely on the "locality hypothesis", suppressing isolated neurons or features. However, harmful semantics act as distributed, cross-layer circuits, rendering such localized interventions brittle and detrimental to utility. To bridge this gap, we propose \textbf{TraceRouter}, a path-level framework that traces and disconnects the causal propagation circuits of illicit semantics. TraceRouter operates in three stages: (1) it pinpoints a sensitive onset layer by analyzing attention divergence; (2) it leverages sparse autoencoders (SAEs) and differential activation analysis to disentangle and isolate malicious features; and (3) it maps these features to downstream causal pathways via feature influence scores (FIS) derived from zero-out interventions. By selectively suppressing these causal chains, TraceRouter physically severs the flow of harmful information while leaving orthogonal computation routes intact. Extensive experiments demonstrate that TraceRouter significantly outperforms state-of-the-art baselines, achieving a superior trade-off between adversarial robustness and general utility. Our code will be publicly released. WARNING: This paper contains unsafe model responses.

Figures (12)

• Figure 1: Neuron-level vs. Path-level Intervention. (a) The former fails to block distributed harmful semantics, often causing semantic leakage. (b) The latter physically severs the causal propagation path, ensuring robust safety without compromising utility.
• Figure 2: Overall of the TraceRouter. First, TraceRouter identifies sensitive onset layers and extracts features via a Top-$K$ SAE. It then traces causal semantic pathways. Finally, selective path-level suppression blocks harmful propagation while preserving general utility.
• Figure 3: Sensitive Onset Layer Detection. Sensitive onset layer is identified as the first local peak of the $\operatorname{SS}(l)$ along depth.
• Figure 4: Safety performance comparison on MLLMs. TraceRouter is compared with SOTA methods across different models. Higher values denote better safety.
• Figure 5: Qualitative results of DMs safety intervention. TraceRouter achieves precise erasure of harmful concepts while maintaining superior image quality and semantic fidelity for benign prompts through causal path-level intervention.