Concept Heterogeneity-aware Representation Steering
Laziz Abdullaev, Noelle Y. L. Wong, Ryan Lee, Shiqi Jiang, Minh-Khoi Nguyen-Nhat, Tan Nguyen
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 unimodal Gaussian distributions with identical covariance, 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.
BibTeX
@inproceedings{
abdullaev2026concept,
title={Concept Heterogeneity-aware Representation Steering},
author={Laziz Abdullaev and Noelle Y. L. Wong and Ryan Lee T. Z. and Shiqi Jiang and Minh-Khoi Nguyen-Nhat and Tan Minh Nguyen},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=jV8qIG3xMR}
}