Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Probe-Free Low-Rank Activation Intervention

Chonghe Jiang, Bao Nguyen, Anthony Man-Cho So, Viet Anh Nguyen

TL;DR

This work introduces FLORAIN, a probe-free, low-rank activation intervention for inference-time editing of transformer activations. By modeling the desirable output region as an ellipsoid and learning a nonlinear low-rank map, FLORAIN performs single-layer, head-wise interventions without training a probe classifier, achieving efficient optimization via a smooth objective with analytical projection under Mahalanobis distance. Empirical results on TruthfulQA show FLORAIN improves truthfulness and informativeness across multiple LMs and remains competitive with or superior to baselines like ITI and few-shot prompting, while maintaining limited distribution shift. The approach enables scalable, near real-time intervention with broad applicability to instruction-tuned models and suggests future work on richer region models and alternative projection schemes.

Abstract

Language models (LMs) can produce texts that appear accurate and coherent but contain untruthful or toxic content. Inference-time interventions that edit the hidden activations have shown promising results in steering the LMs towards desirable generations. Existing activation intervention methods often comprise an activation probe to detect undesirable generation, triggering the activation modification to steer subsequent generation. This paper proposes a probe-free intervention method FLORAIN for all attention heads in a specific activation layer. It eliminates the need to train classifiers for probing purposes. The intervention function is parametrized by a sample-wise nonlinear low-rank mapping, which is trained by minimizing the distance between the modified activations and their projection onto the manifold of desirable content. Under specific constructions of the manifold and projection distance, we show that the intervention strategy can be computed efficiently by solving a smooth optimization problem. The empirical results, benchmarked on multiple base models, demonstrate that FLORAIN consistently outperforms several baseline methods in enhancing model truthfulness and quality across generation and multiple-choice tasks.

Probe-Free Low-Rank Activation Intervention

TL;DR

This work introduces FLORAIN, a probe-free, low-rank activation intervention for inference-time editing of transformer activations. By modeling the desirable output region as an ellipsoid and learning a nonlinear low-rank map, FLORAIN performs single-layer, head-wise interventions without training a probe classifier, achieving efficient optimization via a smooth objective with analytical projection under Mahalanobis distance. Empirical results on TruthfulQA show FLORAIN improves truthfulness and informativeness across multiple LMs and remains competitive with or superior to baselines like ITI and few-shot prompting, while maintaining limited distribution shift. The approach enables scalable, near real-time intervention with broad applicability to instruction-tuned models and suggests future work on richer region models and alternative projection schemes.

Abstract

Language models (LMs) can produce texts that appear accurate and coherent but contain untruthful or toxic content. Inference-time interventions that edit the hidden activations have shown promising results in steering the LMs towards desirable generations. Existing activation intervention methods often comprise an activation probe to detect undesirable generation, triggering the activation modification to steer subsequent generation. This paper proposes a probe-free intervention method FLORAIN for all attention heads in a specific activation layer. It eliminates the need to train classifiers for probing purposes. The intervention function is parametrized by a sample-wise nonlinear low-rank mapping, which is trained by minimizing the distance between the modified activations and their projection onto the manifold of desirable content. Under specific constructions of the manifold and projection distance, we show that the intervention strategy can be computed efficiently by solving a smooth optimization problem. The empirical results, benchmarked on multiple base models, demonstrate that FLORAIN consistently outperforms several baseline methods in enhancing model truthfulness and quality across generation and multiple-choice tasks.

Paper Structure

This paper contains 17 sections, 1 theorem, 23 equations, 1 figure, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Nonlinear Low-Rank Intervention Mapping
  4. Probe-Free Intervention Framework
  5. Algorithms for Training Nonlinear Low-Rank Mapping
  6. Empirical Results
  7. Experimental Setup
  8. Numerical Results
  9. Comparison between Truthfulness Augmentation Baselines across Different LMs
  10. Comparison between ITI, FLORAIN, and Instruction Finetuning Methods.
  11. Conclusion and Future Directions
  12. Limitations
  13. Hyper-parameter Tuning
  14. Proof of Theorem \ref{['thm:ellipsoid']}
  15. Experimental results on the toxicity mitigation task
  16. ...and 2 more sections

Key Result

Theorem 4.1

Suppose that the manifold, the distance measure, and the projection operator are chosen as in eq:ellipsoid-all. Problem eq:prob becomes where $(y)_+ = \max\{0, y\}$.

Figures (1)

  • Figure 1: Exploration scheme in constructing question-wise mean vector $\hat{\mu}_{q}$, where the directions are computed as $\Vec{d}=\hat{\mu}^{+}-\hat{\mu}^{-}$, $\Vec{d}_{q} = \hat{\mu}_{q}^{+}-\hat{\mu}_q^{-}$, and then take the scaled convex combination $\Vec{d}_{\textit{move}} = \hat{\mu}_{q}-\hat{\mu}_{q}^{+} = \lambda \left(\alpha \frac{\Vec{d}}{\| \Vec{d}\|} + (1-\alpha) \frac{\Vec{d}_{q}}{\|\Vec{d}_{q}\|} \right)$.

Theorems & Definitions (2)

  • Theorem 4.1: Ellipsoidal manifold and Mahalanobis distance
  • proof : Proof of Theorem \ref{['thm:ellipsoid']}