Activation Transport Operators

TL;DR

This paper introduces Activation Transport Operators (ATO), linear maps that analyze how features are linearly transported through residual streams in transformer models, aiding understanding, safety, and debugging of LLMs.

Contribution

The paper proposes ATO, a novel linear operator framework to measure feature transport in residual streams, providing insights into linearity and efficiency in transformer models.

Findings

ATO can identify linearly transported features

Transport efficiency has an established upper bound

Empirical results show significant linear transport in residuals

Abstract

The residual stream mediates communication between transformer decoder layers via linear reads and writes of non-linear computations. While sparse-dictionary learning-based methods locate features in the residual stream, and activation patching methods discover circuits within the model, the mechanism by which features flow through the residual stream remains understudied. Understanding this dynamic can better inform jailbreaking protections, enable early detection of model mistakes, and their correction. In this work, we propose Activation Transport Operators (ATO), linear maps from upstream to downstream residuals $k$ layers later, evaluated in feature space using downstream SAE decoder projections. We empirically demonstrate that these operators can determine whether a feature has been linearly transported from a previous layer or synthesised from non-linear layer computation. We develop the notion of transport efficiency, for which we provide an upper bound, and use it to estimate the size of the residual stream subspace that corresponds to linear transport. We empirically demonstrate the linear transport, report transport efficiency and the size of the residual stream's subspace involved in linear transport. This compute-light (no finetuning, <50 GPU-h) method offers practical tools for safety, debugging, and a clearer picture of where computation in LLMs behaves linearly.