TraceRouter++: Causal Semantic Routing Realignment via Optimal Transport
Organizations: The University of Sydney · Nanjing University of Science and Technology · National University of Singapore · Beijing University of Technology
Abstract
Despite their remarkable capabilities, large foundation models (LFMs) remain susceptible to adversarial manipulation. Current safety intervention methods primarily rely on the locality hypothesis, rigidly suppressing isolated neurons or features. However, harmful semantics manifest as distributed, cross-layer circuits, rendering localized ablations both brittle against complex jailbreaks and detrimental to general utility. To bridge this gap, we propose TraceRouter++, a routing-level safety intervention framework that explicitly traces safety-relevant causal routing pathways activated under harmful inputs and realigns their unsafe activation states. Our framework operates in a Discover-Trace-Realign paradigm: it first pinpoints a sensitive onset layer via attention divergence and disentangles dense activations into interpretable routing features using sparse autoencoders (SAEs). Subsequently, it maps these features to downstream causal circuits via feature influence scores (FIS). Crucially, departing from rigid suppression paradigms, TraceRouter++ introduces a constrained optimal transport (OT)-based soft routing realignment. By formulating safety intervention as a tractable coordinate-wise marginal optimal transport approximation, our method redirects unsafe activation marginals toward safe-reference marginals within causally identified routing subspaces while preserving non-selected coordinates at the intervention point. This mechanism safely redirects malicious semantic flows while leaving orthogonal computational routes intact. Extensive experiments across representative large language models (LLMs), multimodal large language models (MLLMs), and diffusion models (DMs) show that TraceRouter++ outperforms evaluated baselines, improving adversarial robustness while maintaining a favorable safety-utility trade-off. WARNING: This paper contains unsafe model responses.
Figures & tables
| Item | LLMs | MLLMs | DMs |
| Models | LLaMA3-8B-Instruct [ 27 ] Mistral-7B-Instruct [ 35 ] | LLaVA-1.5-7B [ 42 ] MiniGPT-4-7B [ 91 ] DeepSeek-VL2 [ 80 ] | Stable Diffusion 1.4 [ 59 ] FLUX.1-Dev [ 38 ] AltDiffusion [ 86 ] Show-o2 [ 84 ] |
| Datasets | AdvBench [ 95 ] , HarmBench [ 50 ] GCG [ 95 ] , AutoDAN [ 43 ] Template [ 81 ] , Prefill [ 81 ] XSTEST [ 60 ] Global-MMLU-Lite [ 68 ] | FigStep [ 25 ] JailbreakV-28K [ 47 ] MMBench [ 45 ] | I2P-N [ 62 ] , I2P-V [ 62 ] P4D [ 12 ] , Ring-A-Bell [ 73 ] MS COCO [ 41 ] |
| Metrics | DSR [ 22 ] ORR [ 14 ] Accuracy [ 69 ] | DSR [ 22 ] Accuracy [ 69 ] | DSR [ 22 ] CLIP Score [ 31 ] FID [ 32 ] |
| Model | Method | Gradient Driven | Pattern Injections | Avg. | |||
| GCG [ 95 ] | AutoDAN [ 43 ] | Template | Prefill | T+P | |||
| LLaMA3-8B-Instruct [ 27 ] | Original [ 27 ] | 88 | 87 | 98 | 41 | 2 | 63.2 |
| DeepAug [ 55 ] | 99 | 40 | 100 | 59 | 3 | 60.2 (-3.0) | |
| CB [ 94 ] | 99 | 98 | 97 | 95 | 98 | 97.4 (+34.2) | |
| DeRTa [ 87 ] | 97 | 89 | 100 | 98 | 32 | 83.2 (+20.0) | |
| HumorReject [ 81 ] | 98 | 99 | 99 | 100 | 98 | 98.8 (+35.6) | |
| Model | Method | Noise | SD | Nature | Blank | Avg. | ||||||||
| T | P | L | T | P | L | T | P | L | T | P | L | |||
| LLaVA-1.5-7B [ 42 ] | Original [ 42 ] | 42.9 | 70.8 | 37.9 | 39.5 | 60.9 | 27.1 | 41.0 | 68.2 | 40.6 | 42.6 | 69.1 | 39.2 | 48.3 |
| FigStep [ 25 ] | 40.5 | 47.7 | 59.5 | 42.9 | 45.1 | 50.0 | 41.6 | 41.6 | 55.5 | 39.2 | 48.9 | 59.5 | 47.7 (-0.6) | |
| CoCA [ 21 ] | 38.8 | 60.9 | 37.9 | 38.7 | 58.8 | 47.3 | 36.9 | 64.8 | 44.6 | 39.0 | 62.7 | 47.3 | 48.7 (+0.4) | |
| ECSO [ 26 ] | 42.7 | 74.6 | 41.9 | 42.7 | 74.6 | 41.9 | 42.7 | 74.6 | 41.9 | 42.68 | 74.6 | 41.9 | 53.1 (+4.8) | |
| AdaShield [ 76 ] | 78.4 | 98.6 | 82.5 | 75.4 | 98.6 | 77.1 | 76.8 | 99.2 | 82.5 | 78.2 | 98.6 | 82.5 | 85.7 (+37.4) | |
| Model | Method | Illegal | Hate | Malware | Physical | Fraud | Adult | Privacy | Legal | Financial | Health | Avg |
| Activity | Speech | Generation | Harm | Content | Violation | Opinion | Advice | Consultation | ||||
| LLaVA-1.5-7B [ 42 ] | Original [ 42 ] | 8.0 | 52.0 | 10.0 | 6.0 | 16.0 | 72.0 | 34.0 | 100.0 | 100.0 | 92.0 | 49.0 |
| FigStep [ 25 ] | 44.0 | 50.0 | 46.0 | 38.0 | 16.0 | 74.0 | 46.0 | 100.0 | 100.0 | 94.0 | 60.8 (+11.8) | |
| CoCA [ 21 ] | 56.0 | 91.8 | 62.0 | 77.9 | 93.5 | 57.5 | 90.4 | 55.9 | 58.1 | 69.6 | 71.4 (+22.4) | |
| ECSO [ 26 ] | 80.0 | 88.0 | 18.0 | 58.0 | 40.0 | 84.0 | 50.0 | 98.0 | 98.0 | 96.0 | 71.0 (+22.0) | |
| AdaShield [ 76 ] | 96.0 | 84.0 | 84.0 | 92.0 | 52.0 | 92.0 | 70.0 | 100.0 | 100.0 | 100.0 | 87.0 (+38.0) |
| Method | Standard Safety (%) | Adversarial Robustness (%) | ||
| I2P (N) [ 62 ] | I2P (V) [ 62 ] | P4D [ 12 ] | Ring-A-Bell [ 73 ] | |
| Stable Diffusion 1.4 [ 59 ] | 82.2 | 59.9 | 1.3 | 16.9 |
| ESD [ 19 ] | 86.0 ( +3.8) | 83.3 ( +23.4) | 36.7 ( +35.4) | 30.3 ( +13.4) |
| UCE [ 20 ] | 89.7 ( +7.5) | 76.7 ( +16.8) | 19.8 ( +18.5) | 66.9 ( +50.0) |
| CA [ 37 ] | 89.8 ( +7.6) | N/A | N/A | N/A |
| SLD-Med [ 62 ] | 88.5 ( +6.3) | 80.3 ( +20.4) | 22.5 ( +21.2) | 33.8 ( +16.9) |
| Modality | Model | Benchmark | Metric | Original | TraceRouter [ 67 ] | TraceRouter++ |
| LLMs | LLaMA3-8B-Instruct [ 27 ] | Global-MMLU-Lite | Acc. | 71.5 | 70.7 ±0.4 (-0.8) | 71.1 ±0.4 (-0.4) |
| Mistral-7B-Instruct [ 35 ] | 63.0 | 61.8 ±0.5 (-1.2) | 62.4 ±0.3 (-0.6) | |||
| MLLMs | LLaVA-1.5-7B [ 42 ] | MMBench | Acc. | 47.5 | 47.2 ±0.1 (-0.3) | 47.4 ±0.2 (-0.1) |
| MiniGPT-4-7B [ 91 ] | 16.0 | 15.2 ±0.1 (-0.8) | 15.5 ±0.1 (-0.5) | |||
| DeepSeek-VL2 [ 80 ] | 78.7 | 78.3 ±0.2 (-0.4) | 78.6 ±0.1 (-0.1) | |||
| DMs | Stable Diffusion 1.4 [ 59 ] | MS COCO | CLIP Score | 31.34 | 30.75 ±0.14 (-0.59) | 31.08 ±0.07 (-0.26) |
| Model | Method | ORR % |
| LLaMA3-8B-Instruct [ 27 ] | Original [ 27 ] | 3.2 |
| DeepAug [ 55 ] | 39.6 (+36.4) | |
| CB [ 94 ] | 4.4 (+1.2) | |
| DeRTa [ 87 ] | 27.6 (+24.4) | |
| HumorReject [ 81 ] | 5.2 (+2.0) | |
| LEACE [ 7 ] | 3.6 (+0.4) |
| Model | Original | Amplification | Suppression |
| LLaMA3-8B-Instruct [ 27 ] | 63.2 | 75.8 ±1.7 (+12.6) | 61.4 ±1.4 (-1.8) |
| LLaVA-1.5-7B [ 42 ] | 49.0 | 78.6 ±0.7 (+29.6) | 18.2 ±1.9 (-30.8) |
| Stable Diffusion 1.4 [ 59 ] | 82.2 | 73.5 ±1.3 (-8.7) | 98.8 ±0.4 (+16.6) |
| Method | LLaMA3-8B-Instruct [ 27 ] | LLaVA-1.5-7B [ 42 ] | Stable Diffusion 1.4 [ 59 ] | ||||
| DSR | Global-MMLU-Lite [ 68 ] | DSR | MMBench [ 45 ] | DSR | CLIP Score [ 31 ] | FID [ 32 ] | |
| Original | 63.2 | 71.5 | 49.0 | 47.5 | 82.2 | 31.34 | N/A |
| Hard Intervention [ 16 ] | 74.2 (+11.0) | 71.4 (-0.1) | 77.4 (+28.4) | 47.8 (+0.3) | 96.3 (+14.1) | 30.70 (-0.64) | 19.12 |
| Linear Steering [ 58 ] | 97.4 (+34.2) | 70.5 (-1.0) | 95.7 (+46.7) | 45.6 (-1.9) | 72.2 (-10.0) | 30.52 (-0.82) | 16.24 |
| TraceRouter++ | (+36.3) | (-0.4) | (+50.4) | (-0.1) | (+17.1) | (-0.26) | |
| Selection Criterion | LLaMA3-8B-Instruct [ 27 ] | LLaVA-1.5-7B [ 42 ] | Stable Diffusion 1.4 [ 59 ] |
| Original | 63.2 | 49.0 | 82.2 |
| Global Maximum Peak | (+21.6) | (+29.4) | (+16.1) |
| Fixed Threshold Selection | (+17.6) | (+25.4) | (+14.3) |
| First Local Peak (Ours) | (+36.3) | (+50.4) | (+17.1) |
| Top- | LLaMA3-8B-Instruct [ 27 ] | LLaVA-1.5-7B [ 42 ] | Stable Diffusion 1.4 [ 59 ] | ||||
| DSR | Global-MMLU-Lite [ 68 ] | DSR | MMBench [ 45 ] | DSR | CLIP Score [ 31 ] | FID [ 32 ] | |
| Original | 63.2 | 71.5 | 49.0 | 47.5 | 82.2 | 31.34 | N/A |
| 88.6 ±0.5 (+25.4) | 71.4 ±0.3 (-0.1) | 81.6 ±0.6 (+32.6) | 47.4 ±0.1 (-0.1) | 96.3 ±0.4 (+14.1) | 31.30 ±0.04 (-0.04) | 5.74 ±0.05 | |
| 99.5 ±0.3 (+36.3) | 71.1 ±0.4 (-0.4) | 99.4 ±0.1 (+50.4) | 47.4 ±0.2 (-0.1) | 99.3 ±0.2 (+17.1) | 31.08 ±0.07 (-0.26) | 8.62 ±0.04 | |
| 99.2 ±0.3 (+36.0) | 70.9 ±0.3 (-0.6) | 95.8 ±0.4 (+46.8) | 46.9 ±0.1 (-0.6) | 97.9 ±0.3 (+15.7) | 31.18 ±0.05 (-0.16) | 12.92 ±0.05 | |
| Method | LLaMA3-8B-Instruct [ 27 ] | LLaVA-1.5-7B [ 42 ] | Stable Diffusion 1.4 [ 59 ] | ||||
| DSR | Global-MMLU-Lite [ 68 ] | DSR | MMBench [ 45 ] | DSR | CLIP Score [ 31 ] | FID [ 32 ] | |
| Original | 63.2 | 71.5 | 49.0 | 47.5 | 82.2 | 31.34 | N/A |
| 99.5 ±0.3 (+36.3) | 71.1 ±0.4 (-0.4) | 99.4 ±0.1 (+50.4) | 47.4 ±0.2 (-0.1) | 99.3 ±0.2 (+17.1) | 31.08 ±0.07 (-0.26) | 8.62 ±0.04 | |
| 98.4 ±0.4 (+35.2) | 71.4 ±0.4 (-0.1) | 87.0 ±0.7 (+38.0) | 47.1 ±0.1 (-0.4) | 97.6 ±0.4 (+15.4) | 31.22 ±0.04 (-0.12) | 9.34 ±0.15 | |
| 95.6 ±0.3 (+32.4) | 70.9 ±0.4 (-0.6) | 86.8 ±0.8 (+37.8) | 46.9 ±0.2 (-0.6) | 98.3 ±0.3 (+16.1) | 31.25 ±0.06 (-0.09) | 8.79 ±0.11 | |
| 93.2 ±0.7 (+30.0) | 70.9 ±0.6 (-0.6) | 92.4 ±0.6 (+43.4) | 45.9 ±0.1 (-1.6) | 96.3 ±0.5 (+14.1) | 31.20 ±0.07 (-0.14) | 9.37 ±0.18 | |
| Method | DSR [ 62 ] | CLIP Score [ 31 ] |
| FLUX.1 Dev [ 38 ] | 70.4 | 26.4 |
| + SNCE [ 29 ] | 80.2 (+9.8) | 24.7 (-1.7) |
| + TraceRouter++ (Ours) | 93.5 (+23.1) | 25.5 (-0.9) |
| AltDiffusion [ 86 ] | 73.5 | 30.5 |
| + SNCE [ 29 ] | 82.6 (+9.1) | 29.7 (-0.8) |
| + TraceRouter++ (Ours) | 92.6 (+19.1) | 30.2 (-0.3) |
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
| Expansion | LLaMA3-8B-Instruct [ 27 ] | LLaVA-1.5-7B [ 42 ] | Stable Diffusion 1.4 [ 59 ] | ||||
| DSR | Global-MMLU-Lite [ 68 ] | DSR | MMBench [ 45 ] | DSR | CLIP Score [ 31 ] | FID [ 32 ] | |
| Original | 63.2 | 71.5 | 49.0 | 47.5 | 82.2 | 31.34 | N/A |
| 88.4 ±0.5 (+25.2) | 71.1 ±0.3 (-0.4) | 77.2 ±0.7 (+28.2) | 47.2 ±0.1 (-0.3) | 88.9 ±0.4 (+6.7) | 31.25 ±0.05 (-0.09) | 11.49 ±0.14 | |
| 99.5 ±0.3 (+36.3) | 71.1 ±0.4 (-0.4) | 99.4 ±0.1 (+50.4) | 47.4 ±0.2 (-0.1) | 99.3 ±0.2 (+17.1) | 31.08 ±0.07 (-0.26) | 8.62 ±0.04 | |
| 91.4 ±0.6 (+28.2) | 70.8 ±0.4 (-0.7) | 82.0 ±0.5 (+33.0) | 47.1 ±0.1 (-0.4) | 89.8 ±0.4 (+7.6) | 31.31 ±0.06 (-0.03) | 10.51 ±0.13 | |
| 93.8 ±0.5 (+30.6) | 70.6 ±0.5 (-0.9) | 76.4 ±0.8 (+27.4) | 46.8 ±0.2 (-0.7) | 88.1 ±0.5 (+5.9) | 31.18 ±0.07 (-0.16) | 13.07 ±0.18 | |
| Gradient Driven DSR (%) | Pattern Injections DSR (%) | Avg. DSR (%) | Global-MMLU-Lite [ 68 ] (%) | ORR (%) | ||||
| GCG [ 95 ] | AutoDAN [ 43 ] | Template | Prefill | T+P | ||||
| 88.0 | 87.0 | 98.0 | 41.0 | 2.0 | 63.2 | 71.50 | 3.2 | |
| (+19.6) | (0.00) | (+0.8) | ||||||
| (+36.3) | (-0.4) | (+0.3) | ||||||
| (+36.4) | (-1.2) | (+4.4) | ||||||
| (+36.7) | (-3.0) | (+9.2) | ||||||
| Method | Request latency (ms) | Latency (ms/output token) | Throughput (output tokens/s) | Peak GPU memory (MiB) | SAE training (s) | Pathfinding (s) | Total offline overhead (s) |
| Original | |||||||
| TraceRouter++ |
| Model | Method / Evaluation Setting | Avg. DSR (%) |
| LLaMA3-8B-Instruct [ 27 ] | Original (adaptive GCG) | |
| TraceRouter++ (adaptive GCG) | ||
| Mistral-7B-Instruct [ 35 ] | Original (adaptive GCG) | |
| TraceRouter++ (adaptive GCG) |
LLM Safety
Large language model (LLM) safety studies how to prevent harmful, unreliable, or misaligned behavior while preserving useful responses.
See all 451 papers