PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2601.21900v2cs.CVcs.AIcs.CYcs.MM

TL;DR

Large foundation models encode harmful semantics in distributed cross-layer circuits, limiting the reliability and utility of localized interventions. TraceRouter identifies sensitive sources, traces their causal pathways with SAEs and feature influence scores, and selectively disconnects those paths. Across model families and benchmarks, it outperforms state-of-the-art baselines in adversarial robustness while preserving general utility, though its applicability beyond safety tasks remains to be established.

  • Problem

    Harmful semantics are distributed across cross-layer pathways, making localized interventions brittle and potentially damaging to model utility.

  • Method

    TraceRouter uses attention divergence, sparse autoencoders, differential activation analysis, and feature influence scores to identify and selectively disconnect harmful causal pathways.

  • Results

    TraceRouter outperforms state-of-the-art baselines in adversarial robustness while preserving general utility across diffusion, language, and multimodal language models.

  • Takeaways & Limitations

    Effective safety intervention requires targeting the topological structure of semantic flow rather than isolated activations.

  • Takeaways & Limitations

    The framework is primarily evaluated on safety-oriented tasks, and its applicability to general-purpose tasks remains to be fully explored.

Abstract

from arXiv · show

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.

1. Introduction

Large foundation models encode harmful semantics as distributed, cross-layer pathways, making localized safety interventions vulnerable to semantic leakage and utility degradation. TraceRouter addresses this gap by identifying, tracing, and blocking causal pathways while preserving orthogonal computation routes.

  • Harmful semantics are distributed across layers and propagate through complex cross-layer computation paths rather than residing in isolated neurons.
  • Localized suppression can fail under adversarial induction and degrade performance by disrupting internal routing and polysemantic units.
  • The path-level representation hypothesis treats sensitive semantics as encoded and propagated through specific neural paths rather than determined by isolated components.
  • TraceRouter uses a discover-trace-disconnect framework to identify sensitive sources, map their downstream causal pathways, and selectively prune harmful propagation.
  • SAEs and feature influence scores support fine-grained disentanglement and causal mapping of harmful information flow at the routing level.
  • Across diverse benchmarks, TraceRouter enhances adversarial robustness while preserving general utility.

2. Related Work

Safety interventions range from costly weight-level modification to lighter parameter editing, while mechanistic interpretability emphasizes distributed, cross-layer circuits as the substrate of semantic behavior.

  • Foundational safety methods use weight-level fine-tuning or machine unlearning, often incurring high computational costs and compromising generalization.
  • Later approaches introduce lightweight parameter editing or closed-form erasure to reduce the inefficiencies of foundational interventions.
  • Mechanistic interpretability research frames high-level behaviors as distributed computations routed across layers rather than isolated neurons.
  • Transformer studies describe functional circuits in which interacting components propagate information along structured and often redundant computational paths.

3. TraceRouter: Robust Safety via Path-Level Intervention

TraceRouter identifies sensitive semantic sources, traces their cross-layer causal pathways, and selectively suppresses those pathways while preserving orthogonal computation routes. Its pipeline combines onset-layer detection, SAE-based feature isolation, zero-out influence analysis, and masked path decomposition.

  • Sensitive Source Identification: A Top-K SAE and differential activation analysis isolate features and neurons specifically responsive to sensitive semantics rather than general linguistic patterns.The method compares weighted frequency scores between sensitive and non-sensitive samples and selects neurons with the largest differential scores.
  • Sensitive Source Identification: TraceRouter detects the sensitive onset layer using attention divergence and selects the earliest local peak of SS(l) above background.Sensitive attention is contrasted with contextual disturbance across paired sensitive and non-sensitive prompts; Stable Diffusion 1.4 reaches its first local peak at Layer 3.
  • Causal Semantic Pathway Tracing: Back-projection maps sensitive SAE features into dense onset-layer activations, where the largest-magnitude components identify source neurons for downstream tracing.Decoder weights project the selected features into the onset layer, and Top-K magnitudes define the source-neuron set.
  • Causal Semantic Pathway Tracing: Zero-out interventions quantify downstream causal effects with FIS, combining activation frequency, conditional mean magnitude, and expected absolute activation shift.Source neurons are zeroed while other neurons remain unchanged, enabling reconstruction of the sensitive semantic pathway.
  • Selective Path-Level Suppression: Path decomposition masks high-FIS neurons and separates each downstream activation into a sensitive pathway component and preserved orthogonal computational routes.The binary mask identifies sensitive neurons, while the orthogonal component represents routes associated with general utility.
  • Selective Path-Level Suppression: Selective scaling of the pathway-specific component suppresses unsafe semantic flow while leaving orthogonal routes intact.The intervention applies the suppression factor only to the identified sensitive circuit rather than indiscriminately modifying all activations.

4. Experiments and Analysis

TraceRouter is evaluated across diffusion, language, and multimodal models using safety, robustness, qualitative, ablation, causal, and utility analyses. Across these settings, it removes harmful semantics and improves safety while preserving benign generation quality and general capabilities.

  • Quantitative safety evaluation: TraceRouter neutralizes 99.2% of I2P nudity and 93.6% of violence, while achieving 74.8% defense against P4D and 98.7% against Ring-A-Bell on Stable Diffusion 1.4.These results establish strong standard-safety and adversarial-robustness performance on diffusion models.
  • Quantitative safety evaluation: TraceRouter raises Mistral-7B average DSR from 6.6% to 98.8% and achieves 99.6% on LLaMA3-8B.The method also reaches 100% on Template attacks for Mistral-7B.
  • Quantitative safety evaluation: TraceRouter achieves 99.60% average safety on LLaVA-1.5 and 98.40% on MiniGPT-4, outperforming the reported multimodal baselines.The LLaVA-1.5 comparison includes a 49.00% baseline and 98.40% for SafePTR.
  • Qualitative intervention results: Qualitative diffusion results show precise harmful-concept removal while preserving image fidelity, human structure, fine texture, and semantic fidelity for benign prompts.Neuron- and feature-level baselines are described as exhibiting residual leakage or visual degradation.
  • Qualitative intervention results: TraceRouter intercepts malicious intent in language jailbreaks and visual-injection attacks, producing refusals instead of harmful outputs without over-refusal in the reported examples.The intervention blocks harmful propagation from visual encoders to language decoders in multimodal models.
  • Ablation and causal analysis: Path-level intervention consistently and significantly outperforms isolated-neuron suppression across three architectures, supporting distributed cross-layer circuits as the relevant intervention target.The comparison argues that blocking downstream propagation prevents semantic bypasses that remain possible after local suppression.
  • Ablation and causal analysis: λ = 2 is reported as the optimal suppression intensity on Stable Diffusion 1.4 because it removes harmful concepts while preserving the main subject’s semantic identity.The experiment varies λ from 0 to 5.
  • Ablation and causal analysis: Amplifying sensitive paths produces 72.20% DSR on SD 1.4 versus 99.20% under suppression, while TraceRouter raises DSR from 82.2% to 99.2% compared with 68.7% under random suppression.These results support causal specificity and show that blind intervention can disrupt benign logic.

5. Conclusion

TraceRouter challenges component-local safety interventions by tracing and disconnecting distributed causal circuits of harmful semantics. Results across diffusion, language, and multimodal models report stronger adversarial robustness while preserving general utility.

  • Conclusion: TraceRouter traces and disconnects causal circuits of harmful semantics instead of suppressing isolated activations.The framework uses sparse autoencoders and feature influence scores to analyze distributed cross-layer propagation.
  • Conclusion: Across DMs, LLMs, and MLLMs, TraceRouter outperforms SOTA baselines in adversarial robustness while preserving general utility.The conclusion frames topological semantic-flow intervention as an interpretable direction for foundation-model safety.

Impact Statement

TraceRouter physically severs harmful-semantic propagation circuits in foundation models. The reported outcome is improved safety across DMs, LLMs, and MLLMs while preserving general reasoning and generation capabilities.

  • Impact Statement: TraceRouter addresses semantic leakage by physically severing causal propagation circuits of harmful semantics within foundation models.The framework is presented as a path-level intervention approach.
  • Impact Statement: The method is reported to enhance security across DMs, LLMs, and MLLMs while preserving general reasoning and generative capabilities.The stated goal is a balance between security alignment and model utility.

A. More Results

TraceRouter’s extended evaluations show strong safety performance, practical inference cost, robustness to SAE settings, and applicability across model architectures. Additional analyses identify λ = 2 and the first local peak as effective operating choices while preserving utility.

  • Comparative safety: 99.6% DSR on MM-SafetyBench exceeds Circuit Breakers’ 97.3% on LLaVA-1.5-7B.The comparison attributes the gain to SAE-based harmful-feature disentanglement and FIS-based causal tracing.
  • Scaling factor: λ = 2 raises DSR from 82.2% to 99.2% while CLIP Score decreases only from 31.34 to 31.27.At λ = 3, safety gains diminish and CLIP Score drops to 29.32.
  • Architectural coverage: The framework reduces unsafe content across FLUX.1 Dev, AltDiffusion, and Show-o2 while retaining image-quality evaluation on MS COCO.These experiments assess applicability across diffusion, multilingual, and unified multimodal architectures.
  • SAE sensitivity: Selecting K at the WFS peak-to-tail transition retains sparse hub neurons and discards weakly responsive non-causal neurons.The WFS analysis identifies a small set of neurons with substantially higher responses than the long tail.
  • Onset-layer selection: The first local peak identifies the sensitive onset layer more effectively than global-maximum or threshold-based alternatives.The analysis reports that early semantic emergence is more important than peak activation intensity for locating the causal bottleneck.
  • Generalizability: A new SAE training and tracing process is required for each unique architecture, although components shared across models permit pathway reuse.Neuron specialization and layer depth vary across architectures.
  • Efficiency: TraceRouter requires 3.28 seconds for image generation after one-time SAE training and tracing, with negligible additional latency.Identified pathways can be reused across models sharing the same text encoder.

B. More Discussions

TraceRouter treats harmful semantics as distributed cross-layer circuits rather than isolated components. Its framework uses onset-layer detection, feature disentanglement, and pathway tracing to block harmful routes while retaining benign computation across architectures.

  • Path-level intervention: FIS identifies and severs entire causal propagation circuits, preventing detours around a single suppressed neuron.This addresses semantic leakage caused by distributed harmful representations.
  • Safety–utility trade-off: SAEs separate sensitive circuit components from orthogonal computational routes, allowing selective suppression of harmful pathways while preserving benign routes.The proposed decomposition targets ZP while retaining Z¬P.
  • Architectural agnosticism: TraceRouter’s discover-trace-disconnect framework applies across architectures by targeting attention and feature-activation routing shared by Transformer-based models.Experiments on FLUX.1 Dev and Show-o2 support the claimed cross-modal applicability.

▷Q4. Does TraceRouter’s effectiveness stem from simple signal masking of high-intensity activations?

The reported evidence rejects simple intensity masking as the source of TraceRouter’s effectiveness and frames the method as a selective, inference-side intervention on causal pathways.

  • Specificity verification: Randomly suppressing non-sensitive paths with similar activation intensities worsens safety, unlike TraceRouter’s identified paths.This specificity test distinguishes causal-path intervention from magnitude-based masking.
  • Intervention mechanism: TraceRouter is an inference-side intervention that separates harmful circuits from general knowledge pathways without permanently changing model weights.This contrasts with weight-level editing or unlearning, which can incur computational costs and degrade generalization.
  • Formal interpretation: The framework’s theoretical analysis models intervention as selective projection in activation space.The method applies a mask-based decomposition to distinguish sensitive and orthogonal components.

C.2. Efficacy of Harmful Suppression (Safety Guarantee)

The safety analysis models TraceRouter as a targeted activation-space suppression aligned with harmful propagation. Its justification depends on causal feature selection and argues that first-order suppression remains effective when curvature is limited.

  • Causal alignment: The intervention reduces harmful loss through a negative term proportional to λ and the alignment between the selected pathway and the harmful gradient.For λ > 0, the analysis predicts a significant reduction in harmful loss.
  • Causal alignment: High-FIS neurons define the sensitive mask, aligning the suppressed component with the gradient of harmful propagation.FIS is described as proportional to expected output change under feature intervention.
  • Approximation validity: The second-order analysis is introduced to justify neglecting nonlinear dynamics in the first-order safety approximation.The paper uses a Taylor expansion with a Lagrange remainder for this validation.
  • Approximation validity: The analysis treats the quadratic Taylor remainder as small because Top-K sparsity bounds the intervention norm, activations are locally piecewise linear, and sensitive features are orthogonal to high-curvature directions.The stated operating range includes the empirically selected λ ≈ 2.

C.3. Preservation of General Utility

TraceRouter aims to preserve general utility by disentangling harmful semantic features and selectively masking their causal direction. The paper argues that harmful and utility-related subspaces are nearly orthogonal, minimizing utility loss while improving the safety-utility trade-off.

  • TraceRouter preserves utility through SAE-based feature disentanglement and a specific intervention mask M(l).The SAE isolates the semantic direction associated with the harmful concept.
  • The SAE resolves feature superposition so M(l) isolates the specific semantic direction of the harmful concept.
  • In a well-disentangled representation space, harmful semantics are nearly orthogonal to general utility, minimizing the perturbation to utility loss.The paper states that the utility-loss change is approximately zero under this condition.
  • Eq. 8 decomposes activations into causal and orthogonal components for the intervention.
  • TraceRouter maximizes the safety-utility ratio by exploiting semantic-feature orthogonality, unlike intensity-based suppression methods that do not distinguish harmful from non-harmful components.

D. Limitation

The study primarily evaluates TraceRouter on safety-oriented tasks involving harmful-content prevention and jailbreak defense. Whether its Discover-Trace-Disconnect framework generalizes to broader non-safety tasks remains unexplored.

  • The study primarily focuses on safety-oriented tasks, including harmful-content prevention and jailbreak defense.
  • The framework’s applicability to logical reasoning, knowledge retrieval, and creative writing remains to be fully explored.
  • Future work will test path-level intervention across broader non-safety domains without compromising the model’s underlying cognitive integrity.
Loading 2601.21900v2…