Conditional Optimal Bridge for Riemannian Activation Steering

2026-07-11

AI Summary

Existing activation steering methods suffer from limited performance both in-distribution and out-of-distribution due to their lack of theoretical grounding and reliance on query-agnostic, fixed intervention directions. This work formulates activation steering as a Schrödinger bridge problem on the hypersphere of residual streams, deriving a probability flow ODE via entropic optimal transport. For the first time, it rigorously establishes the log-density ratio as the optimal intervention target by solving a well-posed optimization problem, and introduces a query-adaptive Schrödinger potential to overcome the limitations of static interventions. Integrating Riemannian geometry, entropic optimal transport, and density ratio gradient estimation, the proposed method consistently outperforms existing baselines across four models and three alignment dimensions—helpfulness, truthfulness, and detoxification—while effectively mitigating out-of-distribution performance degradation.

Abstract

Activation steering offers a lightweight alternative to fine-tuning for controlling large language models at inference time. While many existing methods implicitly optimize a log-density-ratio objective between desired and undesired activation distributions, they do so heuristically rather than deriving it from a principled optimization problem. Moreover, these methods produce query-independent steering directions that can degrade performance on both in-distribution and out-of-distribution (OOD) inputs. We introduce \textsc{Cobras} (Conditional Optimal Bridge for Riemannian Activation Steering), which addresses both limitations by casting activation steering as a Schrödinger Bridge on the residual-stream hypersphere. This formulation yields, to our knowledge, the first principled derivation of the log-density-ratio steering objective from a well-posed optimization problem. Solving the bridge via entropic optimal transport and extracting the probability flow ODE recovers the widely used density-ratio gradient as a special case when the Sinkhorn potentials are uniform. Crucially, the Schrödinger potentials are evaluated at the current activation, making the resulting steering direction inherently query-adaptive. Empirically, across four models and three alignment axes (helpfulness, truthfulness, and detoxification), \textsc{Cobras} consistently outperforms prior activation steering baselines while avoiding the OOD degradation commonly observed in existing methods. The code can be found at https://github.com/arshandalili/cobras.

Problem

• activation steering
• log-density-ratio
• query-adaptive
• out-of-distribution
• optimization problem

Innovation

• Schrödinger Bridge
• activation steering
• entropic optimal transport
• query-adaptive
• log-density-ratio