PaperlayerView exact revision on arXiv

Source-linked AI summary

Controlling Language and Diffusion Models by Transporting Activations

Pau Rodriguez, Arno Blaas, Michal Klein, Luca Zappella, Nicholas Apostoloff, Marco Cuturi, Xavier Suau

arXiv:2410.23054v2cs.LGcs.AIcs.CLcs.CV

TL;DR

Fine-tuning and existing activation-steering approaches have computational, utility, and distribution-preservation concerns. The paper proposes Activation Transport, an optimal-transport approach to steer activations across language and image-generation models. It reports control across LLM and T2I tasks, with a bounded conditioning parameter and stated simplifying assumptions for Linear-ACT.

  • Problem

    Fine-tuning can become costly as models grow and can affect other measures of utility, while constant-shift steering may move activations out of distribution.

    Source: §1 · p. 2

  • Method

    Activation Transport steers activations between source and target distributions with univariate maps that preserve the target distribution, and Linear-ACT applies this as an inference-time intervention.

    Source: §1 · p. 2

  • Results

    Linear-ACT matches or outperforms existing inference-time interventions on LLM toxicity mitigation, concept induction, and truthfulness, and off-the-shelf AcT controls T2I style and concept negation.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    AcT offers a conditioning parameter bounded between 0 and 1, described as user-friendly and interpretable.

    Source: §6 · p. 10

    Linear-ACT assumes linear transport between i.i.d. activations, and its estimated map is limited by the expressiveness of the samples used.

    Source: §6 · p. 11

Abstract

from arXiv · show

The increasing capabilities of large generative models and their ever more widespread deployment have raised concerns about their reliability, safety, and potential misuse. To address these issues, recent works have proposed to control model generation by steering model activations in order to effectively induce or prevent the emergence of concepts or behaviors in the generated output. In this paper we introduce Activation Transport (AcT), a general framework to steer activations guided by optimal transport theory that generalizes many previous activation-steering works. AcT is modality-agnostic and provides fine-grained control over the model behavior with negligible computational overhead, while minimally impacting model abilities. We experimentally show the effectiveness and versatility of our approach by addressing key challenges in large language models (LLMs) and text-to-image diffusion models (T2Is). For LLMs, we show that AcT can effectively mitigate toxicity, induce arbitrary concepts, and increase their truthfulness. In T2Is, we show how AcT enables fine-grained style control and concept negation.

1 INTRODUCTION

Fine-tuning can become costly and affect model utility, while existing activation-steering approaches may move activations out of distribution. The section introduces Activation Transport, which steers activations between distributions and is reported to work across language and image-generation tasks.

  • Fine-tuning can become costly in compute and memory as models grow, and modifying model parameters with new data can affect zero-shot or question-answering utility.Existing steering methods can use sparse or low-dimensional interventions estimated from hundreds of sentences, but constant shifts may move activations out of distribution and hinder conditioning or general performance.Source: §1 · p. 2
  • Activation Transport steers activations using an optimal-transport map between source and target distributions, while univariate maps preserve the target distribution.The framework is presented for language behaviors such as toxicity and for styles in text-to-image generation.Source: §1 · p. 2
  • Without hyperparameter tuning, Linear-ACT matches or outperforms existing inference-time interventions on LLM toxicity mitigation, concept induction, and truthfulness, and controls T2I style and concept negation.The authors also report that off-the-shelf Linear-ACT works on both LLMs and diffusion models.Source: §1 · p. 2
  • The paper frames Linear-ACT as an inference-time intervention that preserves internal activation distributions, with a strength parameter from 0 to 1 and transport support intended to prevent out-of-distribution activations.It also interprets many existing steering methods as mean transport maps.Source: §1 · p. 2

2 RELATED WORK

Prior controllability work spans activation steering for language models, fine-tuning and adapters for image generation, and inference-time interventions. The paper positions its approach among these methods as an intervention that transfers across modalities.

Controlling LLMs

Controlling T2I

3 TRANSPORTING NEURON ACTIVATIONS

Activation Transport learns maps between distributions of sentence-level activation statistics. It uses low-budget affine maps, estimates them layer by layer, and frames prior interventions as special cases of linear transport.

3.1 LOW BUDGET ESTIMATORS FOR TRANSPORT MAPS

  • Linear-ACT estimates independent univariate affine maps for each activation, simplifying transport estimation when parameterizing quantile and CDF functions is intractable.The maps are estimated from paired samples; with squared cost, the map has a closed form.Source: §3.1 · pp. 3–4
  • When σa ≠ σb, vector-based methods diverge from the sample-defined map, while Linear-ACT is robust to differences in σ.Mean-shift methods can over- or undershoot into out-of-distribution activations; the linear choice fits the mostly unimodal activation distributions described here.Source: Figure 2 · p. 4, §3.1 · p. 4

3.2 SEQUENTIAL ITERATIVE MAPS

  • Layerwise causal estimation refreshes activations after each intervention before estimating the next layer’s map; λ interpolates between no transport and full transport.The authors use causal estimation for both Mean-ACT and Linear-ACT.Source: §3.2 · p. 5

3.3 GENERALIZATION OF PRIOR INFERENCE-TIME INTERVENTIONS WORK

  • Prior inference-time interventions can be expressed as linear maps, but most add a bias and use λ as an unbounded bias multiplier rather than interpolating with the transformed activation.ACT uses all activations or the observed input range for support, and averages across tokens; several prior methods use only the last token or max-pooling.Source: §3.3 · pp. 5–6, Table 1 · p. 6

4 EXPERIMENTS ON LLMS

LLM experiments assess AcT for toxicity mitigation, concept induction, and truthfulness on Gemma2-2B and Llama3-8B. The reported results include reduced toxicity and induced concepts and truthfulness.

4.1 TOXICITY MITIGATION IN LLMS

  • Linear-ACT reduces toxicity up to 7.5× on Gemma2-2B and 4.3× on Llama3-8B, with minimal PPL and MMLU impact.ACT methods perform best at λ = 1; ITI-C is sensitive to model, layer, and λ.Source: §4.1 · pp. 6–7, Table 2 · p. 7

4.3 INDUCING TRUTHFULNESS IN LLMS WITH ACT

  • At λ = 1, Linear-ACT and Mean-ACT achieve the best and second-best MC1 and MC2 accuracy improvements; Linear-ACT raises MC1 by roughly 5% on Gemma2-2B and almost 8% on Llama3-8B.The reported MC1 increases exceed the closest non-ACT baseline, ITI-C, by about 1.5% and 3%, respectively, with slightly less MMLU decrease.Source: Table 3 · p. 8, §4.3 · p. 8

4.2 INDUCING CONCEPTS IN LLMS WITH ACT

  • At λ ≈1, Linear-ACT reaches median p(yes) = 0.87 and generated-text PPL = 8.5; both concept presence and PPL worsen for λ > 1.The figure aggregates results over seven concepts and 500 generated sentences.Source: §4.2 · pp. 7–8, Figure 4 · p. 7

5 CONTROLLING IMAGE DIFFUSION MODELS

Activation Transport addresses fine-grained style control and concept negation in text-to-image models. Linear-ACT performs well across both tasks, with full-strength conditioning strongest for style control and better semantic preservation than ITI-C during concept negation.

5 CONTROLLING IMAGE DIFFUSION MODELS

  • ACT addresses style control and concept negation in SDXL and FLUX; its conditioning is strongest at full strength, λ = 1, while adapted ITI-C requires tuning λ and fails at negation.The experiments intervene on SDXL’s normalization layers and most FLUX residual-layer outputs, using CLIP zero-shot classification and CLIPScore to assess results.Source: §5 · pp. 8–9

5.1 STYLE CONTROL

  • For style control, transport maps are learned from 2,048 COCO training prompts paired with style-modified prompts and evaluated on 512 validation prompts.The tested styles are anime, art nouveau, cyberpunk, impressionism, sketch, and watercolor.Source: §5.1 · p. 9
  • On SDXL, Linear-ACT raises cyberpunk classification from ∼12% to ∼95% while retaining ∼80% similarity to the original prompt at λ = 1.Across SDXL and FLUX, its highest style-conditioning score occurs at λ = 1; ITI-C’s best λ differs between models, and its images can exaggerate style while distorting semantics.Source: §5.1 · p. 9, Figure 6 · p. 10

5.2 CONCEPT NEGATION

  • Concept-negation tests learn maps from concept-containing prompts to unmodified prompts, then evaluate 512 captions with instructions to omit pink elephants, white bears, or gorillas.The experiments address a known T2I difficulty: models can still generate concepts named in negative prompts.Source: §5.2 · pp. 9–10
  • For concept negation, Linear-ACT reduces unwanted-concept scores while preserving more original image semantics than ITI-C, which needs stronger intervention and loses semantic content.ITI-C fails on the negation task overall, while ACT is more effective at negation with a higher CLIPScore than ITI-C.Source: §5 · p. 8, §5.2 · p. 10, Figure 6 · p. 10

6 LIMITATIONS AND DISCUSSION

The discussion presents Activation Transport as a general framework for fine-grained control across generative-model architectures. Its Linear-ACT implementation relies on simplifying assumptions and samples whose expressiveness limits the estimated maps.

  • Activation Transport maps activations from a source to a target distribution, unifies prior steering approaches, and generalizes across LLM and text-to-image models.Its conditioning strength is bounded between 0 and 1.Source: §6 · p. 10
  • Linear-ACT assumes linear transport between i.i.d. activations, simplifying choices made for computational and memory reasons.The authors identify nonlinear maps and joint activation distributions as future directions.Source: §6 · p. 11
  • Map estimation depends on its samples, so its expressiveness is limited by the samples used.Source: §6 · p. 11

ETHICS STATEMENT

The ethics statement notes that concept control could be used to develop censorship or misinformation tools. It also cautions against treating the method as a reason to forgo clean pretraining data and additional alignment strategies.

  • Because the method could theoretically induce or mitigate any concept, it could eventually contribute to censorship or misinformation tools.Source: ETHICS STATEMENT · p. 12
  • Using the method to align pretrained models is not a reason to abandon clean data or additional alignment strategies during pretraining.Source: ETHICS STATEMENT · p. 12

REPRODUCIBILITY STATEMENT

The appendices report implementation and reproducibility details alongside ablations on estimation, pooling, transport support, activation distributions, and generation diversity. They also specify the evaluation prompts and summarize the computational costs of Linear-ACT.

REPRODUCIBILITY STATEMENT

  • The reproducibility statement says code and data will be made publicly available and that tables report grid-searched best λ values averaged over five runs.The appendices include intervened-layer, pooling, transport-support, judging-prompt, TruthfulQA, and T2I details.Source: REPRODUCIBILITY STATEMENT · p. 12

A.1 DETAILS ON COMPUTATIONAL COMPLEXITY

  • With fixed λ, Linear-ACT's map can be fused into model weights for zero additional inference cost; variable λ costs O(LD), below a forward pass.Estimation is performed once, with cost O(NML), and is amortized during inference.Source: §A.1 · p. 17, A MEMORY AND COMPUTATIONAL ASPECTS · p. 17

C CAUSAL VS. SIMULTANEOUS ESTIMATION OF ACT

D THE EFFECT OF THE POOLING OPERATION

E THE EFFECT OF THE TRANSPORT SUPPORT

F HOW DO DIFFERENT INTERVENTIONS AFFECT DISTRIBUTIONS?

G ASSESSING TEXT GENERATION DIVERSITY

  • Linear-ACT (Self-BLEU = 0.134) stays closer to the non-intervened model (Self-BLEU = 0.130) than Mean-ACT (0.140) or ITI-C (0.144).The authors report averaging four runs of 1,000 generations, with standard deviations indicating significant results.Source: G ASSESSING TEXT GENERATION DIVERSITY · p. 20, Table 7 · p. 21

H 0-SHOT TOXICITY CLASSIFICATION

J EXTENDED RESULTS ON TOXICITY MITIGATION

The extended toxicity results report that ACT methods are robust to layer choice and that ITI-C is sensitive to both strength and layer. The authors also note that the optimal strength varies across methods and models.

J EXTENDED RESULTS ON TOXICITY MITIGATION

  • ITI-C's optimal strength λ varies across models and layers, complicating its application.The Llama3-8B table is described as supporting similar conclusions to Gemma2-2B.Source: J EXTENDED RESULTS ON TOXICITY… · p. 22, Table 9 · p. 22
  • ACT methods are robust to layer choice and achieve up to 7.5× toxicity mitigation with Linear-ACT at λ = 1 on Gemma2-2B.The table selects each method's best CLS-toxicity result subject to less than +1 increase in Wikipedia PPL.Source: Table 8 · p. 22, J EXTENDED RESULTS ON TOXICITY… · p. 22
  • ACTADD and ITI-C require a strength parameter sweep, whereas AURA does not provide strength control.The Gemma2-2B table reports the best λ for each method under the stated PPL constraint.Source: Table 8 · p. 22, Figure 11 · p. 23

K EXTENDED RESULTS ON CONCEPT INDUCTION ON LLMS

Extended Gemma2-2B examples show gradual, consistent concept induction with Linear-ACT at λ = 1, while ITI-C's strength choice is less clear. The TruthfulQA appendix describes the multiple-choice evaluation protocol.

K EXTENDED RESULTS ON CONCEPT INDUCTION ON LLMS

L.1 EXPERIMENTAL DETAILS

L.2.1 FULL RESULTS OVER 5 SEEDS FOR ALL LAYERS

The appendix documents TruthfulQA results across five runs and the λ-selection procedure for ITI-C and ACTADD. It also provides T2I examples and describes constraints of guidance and negative prompting, with further style-control and concept-negation cases.

L.2.1 FULL RESULTS OVER 5 SEEDS FOR ALL LAYERS

  • TruthfulQA results are reported over five runs for Gemma2-2B and Llama3-8B, with method strengths selected under an MMLU-accuracy constraint.The selection favors MC1 accuracy while requiring performance at least equal to the best ACT method's MMLU accuracy, with 0.1% slack.Source: Table 13 · p. 27, Table 14 · p. 27, §L.2.2 · p. 27

L.2.2 SWEEPING λ FOR ITI-C AND ACTADD

  • ITI-C's λ can strongly affect performance differently across layers and models, and can more than halve MMLU on specified layers.The authors say λ must be swept to find reliable control; they report similar issues for ACTADD in one setting.Source: §L.2.2 · pp. 27–28, Figure 13 · p. 28, Figure 14 · p. 29

M.1 GUIDANCE PARAMETER IN EXISTING DIFFUSION MODELS

  • In SDXL, large guidance values condition effectively, but low values destroy content, making guidance non-intuitive and harder for users.The examples vary guidance strength from 1 to 6.Source: §M.1 · p. 31, Figure 17 · p. 31

M.2 NEGATIVE PROMPTING

  • Negative prompts fail to remove pink elephant, white bear, and gorilla in shown SDXL and Stable Diffusion 3 cases, while ACT removes them in SDXL examples.The appendix also includes concept-negation failure cases.Source: §M.2 · p. 31, Figure 18 · p. 31, Figure 19 · p. 32

M.3 STYLE CONTROL

  • Additional SDXL and FLUX style examples show Linear-ACT inducing styles where Mean-ACT or ITI-C sometimes fail or produce noisy results.The figures include cases where ITI-C fails to induce particular styles and where ACT conditioning is weak.Source: Figure 20 · p. 32, Figure 21 · p. 33, Figure 22 · p. 34

M.4 CONCEPT NEGATION

  • In SDXL examples, Mean-ACT and Linear-ACT remove negated concepts, though maximum strength can sometimes change image semantics.In FLUX examples, ACT sometimes fails to remove a concept, while higher strengths can affect other semantics.Source: Figure 26 · p. 38, Figure 27 · p. 39, Figure 28 · p. 40

M.5 DETAILED RESULTS

The detailed T2I results track style induction and concept negation with classification and ClipScore measures. The supplementary materials also provide prompt tags and specify which FLUX blocks are conditioned.

M.5 DETAILED RESULTS

  • Style-induction figures report 0-shot style-presence classification scores and ClipScore, with gray regions marking images that lost semantic content.The measures are shown for each style and model; ClipScore tracks deviation from the unconditional prompt.Source: Figure 31 · p. 43
  • Concept-negation figures report 0-shot concept-presence classification scores and ClipScore, with gray regions marking images that lost semantic content.The measures are shown for each concept and model; ClipScore tracks deviation from the unconditional prompt.Source: Figure 32 · p. 43

M.6 STYLE PROMPTS

  • Style prompts are listed alongside tags generated with Llama-8B-instruct to induce those styles.Source: Table 15 · p. 44

M.7 CONCEPT PROMPTS

  • Concept prompts are paired with Llama-8B-instruct-generated tags for inducing concepts or prompting models not to generate them.Source: Table 16 · p. 45

M.8 DETAILS ON FLUX CONDITIONING

  • For FLUX, conditioning all multimodal blocks and the first 15 unimodal blocks avoids the image deterioration observed when conditioning blocks closest to the output.Conditioning the output of all blocks was the most effective strategy for strong conditioning, but conditioning blocks nearest the output tended to deteriorate generated images.Source: §M.8 · p. 45

Related papers

Loading 2410.23054v2…