Concept Heterogeneity-aware Representation Steering

AI summary

Overview

Research area: Machine learning / large language model interpretability and control — specifically representation steering (activation steering) via optimal transport theory.

Technical level: Advanced. The paper assumes familiarity with optimal transport, Wasserstein distances, Gaussian mixture models, barycentric projection, entropy-regularized OT (Sinkhorn), and PCA.

Scope: The paper proposes a clustering-based, input-adaptive replacement for global difference-in-means steering vectors, evaluated on jailbreaking, toxicity mitigation, and image style control across 3B to 32B parameter open-weight models (arXiv:2603.02237v2).

What This Paper Is About

Most representation steering methods compute a single global direction — typically the difference in mean activations between two contrastive datasets, such as harmful versus harmless prompts — and add it to hidden activations at inference time. This implicitly assumes the target concept is spread homogeneously through representation space, but LLM representations are in fact clustered and context-dependent, which makes a single global direction brittle. The paper reframes steering as an optimal transport problem between Gaussian mixture models, so that the steering direction is assembled from cluster-level shifts and varies smoothly depending on the input.

Key Contributions

1. A theoretical generalization of steering from unimodal Gaussians to multimodal GMMs. The authors formalize steering as a discrete optimal transport problem between semantic clusters using the Mixture Wasserstein distance, noting that standard difference-in-means steering corresponds to the OT map between two identical distributions with differing first moments — effectively a global translation.

2. CHaRS (Concept Heterogeneity-aware Representation Steering). An input-dependent steering method built from the cluster-level transport plan, producing a smooth, kernel-weighted combination of cluster-level shifts that adapt to context.

3. Principal Component Thresholding (PCT) for transport-aligned steering vectors. The authors show that the induced cluster-wise steering covariance is inherently low-rank, with rank at most 2K − 2 (much smaller than the ambient dimension d), and use this structure to obtain a disentangled spectral-spatial factorization of the steering field. The resulting CHaRS-PCT keeps only the top L ≤ 2K − 2 principal components.

4. Empirical validation across adversarial and safety tasks, spanning open-weight LLMs from 3B to 32B parameters, plus image generation experiments on FLUX.1.

Main Findings

• Higher attack success rates than global steering under ActAdd. Under the Activation Addition framework (Table 1), CHaRS improves ASR over ActAdd on every model tested: Gemma2-9B-Instruct (91.35 to 98.08 for CHaRS and CHaRS-PCT), Llama3.1-8B-Instruct (95.19 to 98.08 for CHaRS and 99.04 for CHaRS-PCT), Llama3.2-3B-Instruct (74.04 to 79.81), Qwen2.5-3B-Instruct (89.42 to 95.19), Qwen2.5-7B-Instruct (91.35 to 95.19 for CHaRS), Qwen2.5-14B-Instruct (94.23 to 95.19 for CHaRS), and Qwen2.5-32B-Instruct (82.69 to 89.42 for CHaRS).

• Higher attack success rates under Directional Ablation. Under the DirAbl framework (Table 2), CHaRS and CHaRS-PCT again exceed the baseline on every model: Gemma2-9B-Instruct (79.81 to 84.62 for CHaRS and 85.58 for CHaRS-PCT), Llama3.1-8B-Instruct (91.35 to 92.31 for both), Llama3.2-3B-Instruct (83.65 to 86.54 for both), Qwen2.5-3B-Instruct (90.38 to 92.31 for both), Qwen2.5-7B-Instruct (89.42 to 90.38 for CHaRS and 91.35 for CHaRS-PCT), Qwen2.5-14B-Instruct (78.85 to 80.77 for both), and Qwen2.5-32B-Instruct (82.69 to 85.58 for CHaRS and 83.65 for CHaRS-PCT).

• Steering is not free: utility drops relative to no steering. In the same tables, all steering methods reduce tinyBenchmark scores compared with the no-steering condition — for example, Gemma2-9B-Instruct falls from an average of 71.11 without steering to 51.39 with ActAdd and 52.34 with CHaRS, and Qwen2.5-32B-Instruct falls from 76.03 to 68.82 with ActAdd and 68.16 with CHaRS. The paper's framing is that CHaRS improves attack success "while maintaining general-language utility," meaning the utility of CHaRS is comparable to that of the baselines rather than higher than the unsteered model.

• CHaRS-PCT matches or beats CHaRS with fewer steering directions. The thresholded variant achieves the best ASR in several settings (Llama3.1-8B-Instruct under ActAdd at 99.04; Qwen2.5-7B-Instruct under DirAbl at 91.35; Gemma2-9B-Instruct under DirAbl at 85.58) while using fewer components, consistent with the proven low-rankness of the steering covariance.

• Toxicity mitigation on RealToxicityPrompts. In sequential steering for toxicity mitigation (Table 3), CHaRS and CHaRS-PCT beat the Linear-AcT baseline on toxicity metrics for Gemma2-2B, Llama3-8B, and Qwen2.5-7B. For Gemma2-2B, CLS toxicity falls from 2.93 ± 0.52 (original) and 0.67 ± 0.05 (Linear-AcT) to 0.53 ± 0.12 for CHaRS (best) and 0.57 ± 0.05 for CHaRS-PCT, with 0-shot toxicity at 3.63 ± 1.56 for CHaRS and 3.47 ± 1.24 for CHaRS-PCT versus 4.60 ± 0.70 for Linear-AcT. For Llama3-8B, CLS toxicity drops from 1.93 ± 0.39 (Linear-AcT) to 1.23 ± 0.12 (CHaRS) and 1.17 ± 0.05 (CHaRS-PCT), and 0-shot toxicity from 7.73 ± 0.94 to 4.80 ± 0.37 and 4.47 ± 0.12. For Qwen2.5-7B, CLS toxicity drops from 1.80 ± 0.51 to 1.03 ± 0.19 (CHaRS) and 1.07 ± 0.12 (CHaRS-PCT), and 0-shot from 6.10 ± 0.93 to 4.97 ± 0.09 and 4.37 ± 0.12.

• Perplexity and MMLU effects in the toxicity setting. The perplexity numbers reported for CHaRS and CHaRS-PCT are slightly higher than the original model in several cases — for Gemma2-2B, Wikipedia perplexity is 14.40 ± 0.00 for the original, 14.62 ± 0.06 for Linear-AcT, 15.06 ± 0.12 for CHaRS, and 15.17 ± 0.08 for CHaRS-PCT; for Llama3-8B, 9.17 ± 0.00 original versus 9.85 ± 0.08 for CHaRS. MMLU drops modestly relative to the original in all three models under all steering methods (Gemma2-2B: 53.03 ± 0.00 to 50.69 ± 0.40 for CHaRS and 50.21 ± 0.19 for CHaRS-PCT; Llama3-8B: 65.52 ± 0.00 to 64.48 ± 0.05 and 64.37 ± 0.11; Qwen2.5-7B: 74.26 ± 0.00 to 73.41 ± 0.06 and 73.31 ± 0.03). The paper states that CHaRS and CHaRS-PCT "outperform prior causal methods, reducing toxic generations without degrading perplexity or downstream performance."

• Evidence that concepts are heterogeneous. Figure 1 shows PCA and t-SNE visualizations of last-token representations for Llama-3.2-3B-Instruct and Qwen2.5-7B-Instruct colored by k-means clustering, presented as instances of feasible heterogeneity. Appendix E.3 is cited for textual examples showing harmful instructions can be grouped coherently via clustering of their last-token hidden representations.

• Image style control and compute details. The paper states it evaluates image style control as a third task on FLUX.1 and that experiments ran on 4×H100 GPUs; the specific numerical results for the FLUX.1 style control experiments are not included in the text excerpt, and the hardware sentence is truncated.

Methodology in Plain English

Instead of treating a concept as one blob with a single average position, the authors split the activations for each concept (for example, harmful and harmless prompts) into K clusters using k-means. Each cluster becomes a component of a Gaussian mixture model. They then compute the optimal transport coupling between the source clusters and target clusters: a cost matrix of squared distances between centroids is solved with entropy-regularized optimal transport using Sinkhorn iterations, which produces a soft matching matrix saying how much mass should move from each source cluster to each target cluster.

Because estimating full covariance matrices for every cluster pair is noisy and expensive in high dimensions, the authors adopt the equal-covariance assumption, which reduces each cluster-pair transport map to a simple translation equal to the difference of centroids. For any input representation, the method computes a Gaussian (RBF) kernel similarity between that input and each source centroid, with the bandwidth set to the median of squared distances from the input to the centroids, and uses these similarities plus the transport weights to blend the cluster-level translation vectors into one input-dependent steering vector. This vector is scaled by a steering strength parameter α and added to the activation, exactly following the Activation Addition interface. A remark notes that normalizing the vector and choosing α = −xᵀv̂(x) extends the same construction to the Directional Ablation interface.

The PCT variant looks at the weighted covariance of all cluster-level steering vectors. Because each centered steering vector is a difference of centroid deviations, the covariance has rank at most 2K − 2, far below the model dimension. The authors take its PCA, keep only the top L ≤ 2K − 2 components, and rewrite the steering field as a global mean vector plus a sum of principal directions weighted by input-dependent coefficients — a spectral-spatial factorization that captures between-cluster and within-cluster variability.

Why This Matters

Impact on research. The paper refines the theoretical bridge between difference-in-means steering and optimal transport, which earlier work characterized as the OT map between two identical Gaussians with different means. By relaxing that to Gaussian mixtures, it gives the activation-steering literature a principled framework for handling non-Gaussian, clustered representation structure and a concrete low-rank argument for why compact steering directions should suffice.

Real-world applications:

• Safety testing and red-teaming of deployed LLMs, where higher attack success rates under a steering intervention mean stronger stress tests of refusal behavior across model families and sizes (3B to 32B parameters).
• Toxicity mitigation in generation pipelines, where the method reduces CLS and 0-shot toxicity scores on RealToxicityPrompts across Gemma2-2B, Llama3-8B, and Qwen2.5-7B.
• Controllable image generation, where the same transport formulation is applied to style control on FLUX.1.
• Behavior and persona control in assistants, since the steering interface is compatible with the widely used Activation Addition and Directional Ablation frameworks.

Industry relevance. The method operates at inference time on internal activations and requires no fine-tuning, which makes it attractive for teams that want to add or suppress behaviors in existing open-weight models. Because CHaRS-PCT needs fewer steering directions, it points toward lower per-token overhead, and the paper reports the approach integrates with both language and diffusion models. The code is released publicly at https://github.com/lazizcodes/CHaRS.

Future Directions

• Closing the utility gap. Every steering method in the tinyBenchmark tables sacrifices substantial utility relative to the unsteered model, and the toxicity experiments show slight perplexity increases and modest MMLU decreases. How to raise attack success or reduce toxicity without these costs remains open.

• Full-covariance transport. The theoretical derivation allows general Gaussian component covariances, but the implementation falls back to the equal-covariance assumption for tractability. Whether reliable covariance estimation at LLM scale can improve steering quality is untested.

• Choosing the number of clusters K and retained components L. The paper establishes that the steering covariance has rank at most 2K − 2 and that L ≤ 2K − 2, but the excerpt does not report how K and L were selected or how sensitive results are to them.

• Generalization beyond the tested tasks and models. The evaluation covers jailbreaking, toxicity mitigation, and image style control on Gemma2, Llama3, Llama3.2, Qwen2.5, and FLUX.1. Whether the approach extends to other concept types, other modalities, and other model architectures is not established here.

Target Audience

Researchers and practitioners in LLM interpretability, alignment, and safety who already understand activation steering and want a more expressive alternative to single-direction interventions. It is also relevant to applied engineers who deploy open-weight models and need inference-time behavioral control, and to theoretically inclined readers interested in optimal transport applied to representation alignment. Readers without background in optimal transport, Wasserstein distances, or Gaussian mixture models will find the derivations in Sections 2 and 3 demanding.

Authors’ abstract

Representation steering offers a lightweight mechanism for controlling the behavior of large language models (LLMs) by intervening on internal activations at inference time. Most existing methods rely on a single global steering direction, typically obtained via difference-in-means over contrastive datasets. This approach implicitly assumes that the target concept is homogeneously represented across the embedding space. In practice, however, LLM representations can be highly non-homogeneous, exhibiting clustered, context-dependent structure, which renders global steering directions brittle. In this work, we view representation steering through the lens of optimal transport (OT), noting that standard difference-in-means steering implicitly corresponds to the OT map between two identical distributions with differing first moments, yielding a global translation. To relax this restrictive assumption, we theoretically model source and target representations as Gaussian mixture models and formulate steering as a discrete OT problem between semantic latent clusters. From the resulting transport plan, we derive an explicit, input-dependent steering map via barycentric projection, producing a smooth, kernel-weighted combination of cluster-level shifts. We term this method Concept Heterogeneity-aware Representation Steering (CHaRS). Through numerous experimental settings, we show that CHaRS yields more effective behavioral control than global steering.