Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Toward a Flexible Framework for Linear Representation Hypothesis Using Maximum Likelihood Estimation

Trung Nguyen, Yan Leng

TL;DR

This work extends the linear representation hypothesis by removing reliance on single-token counterfactuals and unembedding-based representations. It defines binary concepts as unit vectors in a canonical representation space and derives concept directions from activation differences using a von Mises-Fisher model with maximum likelihood estimation, producing the SAND method. The approach is validated on LLaMA models across multiple concepts and QA benchmarks, showing improved monitoring and steering performance with minimal computational overhead. The results offer a practical, principled toolkit for activation engineering that enhances interpretability, controllability, and robustness in large language models, while acknowledging ethical considerations around steerability.

Abstract

Linear representation hypothesis posits that high-level concepts are encoded as linear directions in the representation spaces of LLMs. Park et al. (2024) formalize this notion by unifying multiple interpretations of linear representation, such as 1-dimensional subspace representation and interventions, using a causal inner product. However, their framework relies on single-token counterfactual pairs and cannot handle ambiguous contrasting pairs, limiting its applicability to complex or context-dependent concepts. We introduce a new notion of binary concepts as unit vectors in a canonical representation space, and utilize LLMs' (neural) activation differences along with maximum likelihood estimation (MLE) to compute concept directions (i.e., steering vectors). Our method, Sum of Activation-base Normalized Difference (SAND), formalizes the use of activation differences modeled as samples from a von Mises-Fisher (vMF) distribution, providing a principled approach to derive concept directions. We extend the applicability of Park et al. (2024) by eliminating the dependency on unembedding representations and single-token pairs. Through experiments with LLaMA models across diverse concepts and benchmarks, we demonstrate that our lightweight approach offers greater flexibility, superior performance in activation engineering tasks like monitoring and manipulation.

Toward a Flexible Framework for Linear Representation Hypothesis Using Maximum Likelihood Estimation

TL;DR

This work extends the linear representation hypothesis by removing reliance on single-token counterfactuals and unembedding-based representations. It defines binary concepts as unit vectors in a canonical representation space and derives concept directions from activation differences using a von Mises-Fisher model with maximum likelihood estimation, producing the SAND method. The approach is validated on LLaMA models across multiple concepts and QA benchmarks, showing improved monitoring and steering performance with minimal computational overhead. The results offer a practical, principled toolkit for activation engineering that enhances interpretability, controllability, and robustness in large language models, while acknowledging ethical considerations around steerability.

Abstract

Linear representation hypothesis posits that high-level concepts are encoded as linear directions in the representation spaces of LLMs. Park et al. (2024) formalize this notion by unifying multiple interpretations of linear representation, such as 1-dimensional subspace representation and interventions, using a causal inner product. However, their framework relies on single-token counterfactual pairs and cannot handle ambiguous contrasting pairs, limiting its applicability to complex or context-dependent concepts. We introduce a new notion of binary concepts as unit vectors in a canonical representation space, and utilize LLMs' (neural) activation differences along with maximum likelihood estimation (MLE) to compute concept directions (i.e., steering vectors). Our method, Sum of Activation-base Normalized Difference (SAND), formalizes the use of activation differences modeled as samples from a von Mises-Fisher (vMF) distribution, providing a principled approach to derive concept directions. We extend the applicability of Park et al. (2024) by eliminating the dependency on unembedding representations and single-token pairs. Through experiments with LLaMA models across diverse concepts and benchmarks, we demonstrate that our lightweight approach offers greater flexibility, superior performance in activation engineering tasks like monitoring and manipulation.

Paper Structure

This paper contains 27 sections, 1 theorem, 14 equations, 14 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Linear Representation Hypothesis
  4. Concept Vector for Activation Engineering
  5. Background: Revisiting park2024the
  6. Binary Concepts and Causal Separability.
  7. Linear Representation in the Embedding Space.
  8. Unified Representations via the Causal Inner Product.
  9. Our Proposed Framework
  10. Preliminaries
  11. Generalized Representation Framework
  12. Deriving Our Algorithm
  13. Choices for Geometry in Activation Space $\Psi$
  14. The SAND Algorithm
  15. Operation Count
  16. ...and 12 more sections

Key Result

Theorem 4.2

Algorithm al:main requires $\sim 2n_v\times d\times k$ flops given input matrices $\Lambda \in \mathbb{R}^{d \times k}$ and $C \in \mathbb{R}^{n_v \times d}$,

Figures (14)

  • Figure 1: MD, SAND-e, and SAND-w demonstrate significantly stronger alignment in their concept directions compared to PCA. Enlarged versions of these plots are provided in the Appendix \ref{['app:cos']}.
  • Figure 2: Singular Values within the 1% to 99% quantile ranges of Matrices $C$ in LLaMA-2 Chat Models
  • Figure 3: Cumulative Energy Plots of Singular Values for Matrices $C$ in LLaMA-2 Chat Models
  • Figure 4: Concept direction map to intervention representations. The top and bottom panel correspond to SAND and PCA correspondingly. The intervention strength is set to $10$. SAND captures concept directions in all cases, whereas PCA fails to do so.
  • Figure 5: Cosine similarities between Truthfulness directions, extracted by different methods using six QA examples given in Table \ref{['tab:truth_examples']}, across layers of the LlaMA-2 7B Chat model
  • ...and 9 more figures

Theorems & Definitions (3)

  • Definition 3.1
  • Definition 4.1
  • Theorem 4.2