Research map / Image generation / Text-to-image diffusion
Controlling Language and Diffusion Models by Transporting Activations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Activation Transport (AcT), an optimal-transport-based inference-time activation steering framework (Linear-AcT) that unifies prior steering methods and controls both LLMs and text-to-image diffusion models.
Problem
Steering generative model activations at inference time often shifts them out of distribution, harming controllability and general model utility; existing methods lack a principled, bounded, layer-robust strength control and do not work across both LLMs and diffusion models.
Approach
Model each activation dimension as a univariate optimal-transport map (affine, fit by least squares on sorted samples) from source to target activation distribution, applied causally layer by layer and interpolated with a bounded strength lambda in [0,1]; many prior steering methods are shown to be special cases (mean-shift transports).
Key results
- Linear-AcT reduces CLS toxicity up to 7.5x on Gemma2-2B (4.17% to 0.56%) and 4.3x on Llama3-8B with lambda=1 and minimal PPL/MMLU change.
- Linear-AcT improves TruthfulQA MC1 by about +4.95 points on Gemma2-2B and +7.76 on Llama3-8B, outperforming ITI-c.
- On SDXL, Linear-AcT raises style presence from about 12% to about 95% while keeping about 80% CLIPScore similarity at lambda=1.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| apply Linear-AcT for toxicity mitigation | Gemma2-2B · RealToxicityPrompts | CLS toxicity (%): $4.17 \pm {0.32}$ -> $\underline{0.56} \pm {0.21}$ | Helped |
| apply Linear-AcT for toxicity mitigation | Llama3-8B · RealToxicityPrompts | CLS toxicity (%): $5.80$ -> $\mathbf{1.35} \pm {0.39}$ | Helped |
| apply Linear-AcT for truthfulness induction | Gemma2-2B · TruthfulQA | MC1 accuracy (%): $21.05 $ -> $\textbf{26.00} \pm {0.32}$ | Helped |
| apply Linear-AcT for truthfulness induction | Llama3-8B · TruthfulQA | MC1 accuracy (%): $25.46 $ -> $\textbf{33.22} \pm {0.22}$ | Helped |
| use mean pooling rather than min pooling (ablation) | Linear-AcT on Gemma2-2B · Toxicity mitigation | MMLU: $45.85 \pm 0.09$ -> $51.87 \pm 0.06$ | Helped |
| use mean pooling rather than max pooling (ablation) | Linear-AcT on Gemma2-2B · Toxicity mitigation | MMLU: $47.01 \pm 0.30$ -> $51.87 \pm 0.06$ | Helped |
| use mean pooling rather than last-token pooling (ablation) | Linear-AcT on Gemma2-2B · Toxicity mitigation | MMLU: $48.49 \pm 0.25$ -> $51.87 \pm 0.06$ | Helped |
| estimate transport maps causally instead of simultaneously (ablation) | Linear-AcT on Gemma2-2B · Toxicity mitigation | CLS toxicity (%): $0.65 \pm {0.17}$ -> $0.56 \pm {0.21}$ | Helped |
| estimate transport maps causally instead of simultaneously (ablation) | Linear-AcT on Llama3-8B · Toxicity mitigation | 0-shot toxicity (%): $6.55 \pm {0.75}$ -> $6.68 \pm {0.81}$ | Mixed |
| apply Linear-AcT for style control at full strength | SDXL · COCO Captions validation, style control | style presence: $\sim12\%$ -> $\sim95\%$ | Helped |
Surprising findings
- Causal (layer-by-layer) estimation of transport maps gives better toxicity mitigation than simultaneous estimation.
- ITI-c's best lambda differs across toxicity, concept induction, and T2I tasks, whereas AcT's best lambda is consistently 1.
Limitations
- Assumes linear, unimodal, independent per-activation transport; non-linear or joint maps are not explored.
- Map estimation depends on the expressiveness and size of the sample sets used (hundreds of samples).
Experimental setup
Datasets: RealToxicityPrompts OneSec WikiText / Wikipedia sentences COCO Captions TruthfulQA WordNet
Benchmarks: RealToxicityPrompts TruthfulQA MMLU COCO Captions
Metrics: CLS toxicity 0-shot toxicity perplexity MC1 accuracy MC2 accuracy MMLU accuracy concept presence p(yes) CLIP zero-shot presence CLIPScore Self-BLEU
Baselines: ActAdd AurA ITI-c CAA Mean-AcT RepE
Builds on: ITI CAA ActAdd AurA RepE
Closest papers by meaning
- Beyond Linear Activation Steering: Invertible Latent Transformations for Controlling LLM Behavior (NeurIPS 2026 · similarity 0.89)
- Steering When Necessary: Flexible Steering Large Language Models with Backtracking (NeurIPS 2025 · similarity 0.88)
- Concept Reachability in Diffusion Models: Beyond Dataset Constraints (ICML 2025 · similarity 0.88)
- Scaling Diffusion Language Models via Adaptation from Autoregressive Models (ICLR 2025 · similarity 0.88)
- Controlling Large Language Models Through Concept Activation Vectors (AAAI 2025 · similarity 0.88)
- LinEAS: End-to-end Learning of Activation Steering with a Distributional Loss (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.