Activation Transport Operators

• Activation Transport Operators (ATO) are mechanisms defined on internal neural activations that transport information across layers using regularized linear and conditioned nonlinear maps.
• ATOs enable prediction of downstream features in decoder-only transformers and safe guidance in text-to-image models by shifting unsafe activations toward safe regions.
• Empirical studies show that ATOs effectively capture linear transport channels and balance safety with utility, supporting early detection and correction of harmful features.

Activation Transport Operators (ATO) are operators defined on internal neural activations that model how representations are moved across a model’s computation. In the 2025 residual-stream formulation, an ATO is a regularized linear map from upstream to downstream residuals $k$ layers later, evaluated in downstream sparse-autoencoder (SAE) feature space (Szablewski et al., 24 Aug 2025). In the 2026 text-to-image safety formulation, ATO denotes an operator acting on intermediate activations $z$ to move unsafe activations toward safe regions while minimally perturbing benign ones; "Conditioned Activation Transport" (CAT) instantiates this idea with a learned transport map $T_\theta$ and a conditioning mask $\mathcal{C}$ (Chrabąszcz et al., 3 Mar 2026). Taken together, these usages define ATOs as a family of transport mechanisms for internal representations, with one line of work emphasizing linear feature flow in decoder-only transformers and the other emphasizing conditioned nonlinear steering in text-to-image generation.

1. Definitions and scope

In decoder-only transformers, the residual stream mediates communication between layers via linear reads and writes of non-linear computations. Within that setting, ATOs are introduced as linear maps from upstream to downstream residuals $k$ layers later, with evaluation performed in downstream SAE feature space. Let $r^{(l)}(i) \in \mathbb{R}^d$ denote the residual vector at layer $l$ and token position $i$, and let the downstream affine predictor be

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$

In the study, the $j$-policy is same-token, so $z$0 (Szablewski et al., 24 Aug 2025).

The feature-space evaluation uses downstream SAE decoder projections. With downstream decoder matrix $z$1 and projection operator $z$2, downstream feature coefficients are

$z$3

Predicted downstream features are then compared through

$z$4

In the text-to-image safety setting, the paper uses the language of activation steering and activation transport. Activation Steering, or Representation Engineering, intervenes directly on the model’s internal activations during inference, grounded in the hypothesis that high-level concepts are encoded as directions in the latent space. Here, transport refers to moving intermediate representations from an unsafe manifold toward a safe manifold using a learned map. The overall steered state is

$z$5

where $z$6 is the mean-pooled activation across tokens, $z$7 is the steering strength, $z$8 is the learned transport map, and $z$9 is the conditioning mask (Chrabąszcz et al., 3 Mar 2026).

A concise comparison is useful because the two papers use the same term at different levels of generality.

Formulation	Object being transported	Operator form
Residual-stream ATO	Upstream residuals to downstream residuals $T_\theta$0 layers later	Affine, rank-constrained linear map
CAT-instantiated ATO	Unsafe intermediate activations toward safe regions	Residual transport with conditioning

This suggests that ATO is best understood as a transport-centered abstraction rather than a single fixed algorithm.

2. Linear ATOs in residual streams

The residual-stream formulation is motivated by a distinction between content that is linearly transported through the residual stream and content that is synthesized by later non-linear computation. The paper positions ATOs relative to SAEs and activation patching: SAEs recover monosemantic features and provide dictionaries that decode residuals into sparse latent features, whereas activation patching discovers causal circuits by intervening on activations to test component importance. ATOs bridge representation and attribution by learning explicit transport maps that predict downstream residuals, and downstream SAE feature projections, from upstream residuals (Szablewski et al., 24 Aug 2025).

Per-feature evaluation uses downstream decoder directions $T_\theta$1 and scalar projections

$T_\theta$2

For a downstream feature $T_\theta$3, predictive agreement is measured by

$T_\theta$4

with $T_\theta$5 serving as the principal statistic and MSE monitored as an auxiliary metric.

The classification scheme is explicit. A downstream feature is treated as linearly transported if $T_\theta$6, with $T_\theta$7 used as a high-confidence threshold in reporting. It is treated as synthesized by later nonlinear computation if $T_\theta$8, where values near or below $T_\theta$9 indicate poor predictive agreement. Features are analyzed only if they activate at least ten times in the held-out set and achieve $\mathcal{C}$0.

The fitted operator is obtained by ridge-regularized regression with a rank constraint:

$\mathcal{C}$1

and, for zero-mean data with $\mathcal{C}$2,

$\mathcal{C}$3

The rank-$\mathcal{C}$4 map is then formed by truncated SVD.

The reported empirical pattern is that linear transport is strongest over short distances and weaker over large distances. For target layer $\mathcal{C}$5, operators trained for $\mathcal{C}$6 and $\mathcal{C}$7 transport many features with $\mathcal{C}$8. Transport deteriorates as $\mathcal{C}$9 grows, and transport is less common in later layers such as target layer $k$0, even for small $k$1. The paper interprets this as an early-layer regime where the residual stream carries features linearly, followed by later layers that prioritize synthesis, recomposition, and eviction of old information.

3. Transport efficiency and the linear transport subspace

A central contribution of the residual-stream work is the notion of transport efficiency, which measures how close a fitted rank-$k$2 operator comes to the best possible rank-$k$3 linear prediction in whitened downstream space (Szablewski et al., 24 Aug 2025). Given zero-mean upstream and downstream matrices

$k$4

the paper defines

$k$5

The upper bound is derived through CCA. With

$k$6

the whitened cross-covariance is

$k$7

If $k$8 is its SVD, then the CCA ceiling is

$k$9

Therefore $r^{(l)}(i) \in \mathbb{R}^d$0 for any fitted operator.

The paper also defines the effective dimensionality of the Linear Transport Subspace (LTS) as

$r^{(l)}(i) \in \mathbb{R}^d$1

This is an effective dimension: if transportable energy concentrates in a few modes, $r^{(l)}(i) \in \mathbb{R}^d$2 is small; if it is spread across many modes, $r^{(l)}(i) \in \mathbb{R}^d$3 is large. Rank selection follows this quantity, since choosing $r^{(l)}(i) \in \mathbb{R}^d$4 yields diminishing returns beyond the CCA ceiling.

Empirically, transport efficiency and LTS size depend on leap size $r^{(l)}(i) \in \mathbb{R}^d$5. For target layer $r^{(l)}(i) \in \mathbb{R}^d$6, the paper reports $r^{(l)}(i) \in \mathbb{R}^d$7 for $r^{(l)}(i) \in \mathbb{R}^d$8, $r^{(l)}(i) \in \mathbb{R}^d$9 for $l$0, and $l$1 for $l$2. Efficiency grows nearly linearly with rank and approaches the CCA ceiling near full rank for $l$3, while it saturates early at lower values for $l$4 and $l$5. The reported interpretation is that the effective linear transport subspace shrinks with longer leaps.

The causal validation is likewise transport-specific. Under interventions that ablate upstream positions and inject ATO reconstructions downstream, ATO reconstruction raises perplexity only slightly, with the effect growing with $l$6. At $l$7, the increase is $l$8 of the maximum degradation from zero intervention; for $l$9 it stays below $i$0. Applying ATOs to all positions yields at most a $i$1 increase at $i$2, with upper-bound perplexity $i$3. These measurements support the claim that the fitted operators recover much of the predictive content otherwise lost under zero intervention.

4. Conditioned and nonlinear ATOs for text-to-image safety steering

The CAT paper uses ATO in a broader sense: an operator acting on intermediate activations $i$4 to move unsafe activations toward safe regions while minimally perturbing benign ones. Its core claim is that linear activation steering frequently degrades image quality when applied to benign prompts, motivating a conditioned nonlinear transport map (Chrabąszcz et al., 3 Mar 2026).

The general steered state is

$i$5

with $i$6 for $i$7. When $i$8, the operator becomes identity, minimizing interference with benign activations.

The paper contrasts three linear baselines with nonlinear transport. Linear Activation Addition (ActAdd) applies a global shift by a centroid difference,

$i$9

which assumes a constant translation vector and ignores variance and shape differences between safe and unsafe distributions. Linear-ACT uses

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$0

where $$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$1 is a per-dimension scale and $$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$2 denotes the Hadamard product, and the additional affine baseline is

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$3

The paper states that linear separability assumptions are often violated in text-to-image models; toy manifolds and experiments on Z-Image and Infinity show that linear methods can overly compress or misalign unsafe manifolds, and when made strong enough to suppress unsafe content, they frequently degrade benign generations.

CAT replaces these with an MLP-based residual transport,

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$4

initialized to identity by zero-initializing the final projection and regularized so that safe inputs are minimally perturbed. The MLP uses a single hidden layer with RMSNorm and GELU.

Training uses a dual alignment and identity-regularization objective:

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$5

The first term aligns unsafe activations to safe targets, while the second enforces identity on already-safe inputs. The paper explicitly describes this second term as a built-in conditioning signal to avoid benign drift.

Conditioning is central. The baseline min-max mask uses a hyper-rectangular bounding box around unsafe activations via quantiles, but the paper notes that this approximation can be too loose and capture benign queries. The proposed geometry-aware variants estimate a stable precision matrix using the shrinkage estimator

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$6

The GDA variant defines

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$7

and activates when the posterior for the unsafe class exceeds a threshold. The Mahalanobis out-of-distribution variant computes

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$8

with conditioning

$$\hat r^{(l+k)}(j) = T_r r^{(l)}(i) + b, \qquad T_r = U_r S_r V_r^\top,\ \mathrm{rank}(T_r)=r \le d.$$9

typically with $j$0. The stated purpose is to confine the operator’s support to the geometry of the toxic concept.

5. Data, integration, and empirical behavior in text-to-image models

The CAT framework is trained on SafeSteerDataset, a contrastive dataset containing $j$1 semantically aligned safe–unsafe prompt pairs spanning a granular taxonomy of $j$2 subcategories within six high-risk categories: Sexual, Hate, Humiliation, Violence, Illegal Activity, and Disturbing (Chrabąszcz et al., 3 Mar 2026). Each subcategory contains $j$3 situations, yielding $j$4 pairs per subcategory, balanced overall. Construction proceeds in two stages: Gemini 2.5-Pro generates candidates, and Qwen-8b embeddings are used to retain pairs with cosine similarity $j$5. The paper’s rationale is that high-cosine-similarity ensures the learned transport isolates toxic manifolds while preserving scene semantics.

The transport map is trained on a $j$6 split of SafeSteerDataset. The paper recommends exploring $j$7 and $j$8, with best configurations often using $j$9. Conditioning may be min_max, CGDA, or Mahalanobis OOD, and the MLP’s final projection is initialized to zero so that the untrained map starts as identity.

Integration is reported on two state-of-the-art architectures: Z-Image, a Single-Stream Diffusion Transformer (S3-DiT), and Infinity, a high-resolution autoregressive model described as "Bitwise AutoRegressive Modeling". Steering is applied in the second half of the models in both the text and vision components, because steering early layers resulted in completely degraded images. The layer-wise inference-time procedure is:

• mean-pool: $z$00
• gate: $z$01
• displacement: $z$02
• broadcast: $z$03

Safety is measured by Attack Success Rate (ASR), the fraction of generations from unsafe prompts classified Unsafe by ShieldGemma-2-4b-it, and general utility by CLIP Score (ViT-B/32) on MS-COCO validation.

Model	Method/configuration	ASR / CLIP
Z-Image	No Steering	33.91 / 0.35
Z-Image	ActAdd (min_max, $z$04)	9.57 / 0.34
Z-Image	Linear-ACT (min_max, $z$05)	2.61 / 0.22
Z-Image	Affine (none, $z$06)	8.70 / 0.25
Z-Image	CAT (none, $z$07)	9.13 / 0.34
Z-Image	CAT (reg=0.5, none, $z$08)	6.96 / 0.33
Infinity	No Steering	31.74 / 0.33
Infinity	ActAdd (none, $z$09)	12.17 / 0.32
Infinity	Linear-ACT (ood_mahal., $z$10)	2.61 / 0.16
Infinity	Affine (mahal., $z$11)	3.48 / 0.10
Infinity	CAT (ood_mahal., $z$12)	11.30 / 0.25
Infinity	CAT (reg=0.5, min_max, $z$13)	4.78 / 0.32

The ablations clarify the safety–utility trade-off. For Infinity with Linear-ACT, no conditioning gives ASR $z$14 but CLIP $z$15, which the paper describes as images effectively destroyed; min_max conditioning partially recovers quality but worsens safety; ood_mahal. yields a better trade-off. For Z-Image with CAT reg=$z$16, adding mahal. changes ASR from $z$17 to $z$18 and CLIP from $z$19 to $z$20, which the paper describes as tight gating that slightly increases ASR but keeps utility.

The modality ablation states that steering both text and vision components is necessary. On Z-Image, text-only gives ASR $z$21 and CLIP $z$22, vision-only gives ASR $z$23, and text+vision gives ASR $z$24 and CLIP $z$25, which the paper labels the best global trade-off. On Infinity, text+vision gives ASR $z$26 and CLIP $z$27, compared with text-only $z$28 and vision-only $z$29 ASR.

A fine-grained Sexual-category ablation reports that CAT preserves more fidelity than linear methods while reducing ASR. On Z-Image, CAT (reg=$z$30, none, $z$31) yields ASR $z$32 and CLIP $z$33, compared with Linear-ACT at ASR $z$34 and CLIP $z$35. On Infinity, CAT (ood_mahal., $z$36) yields ASR $z$37 and CLIP $z$38, whereas Linear-ACT gives ASR $z$39 and CLIP $z$40 and Affine gives ASR $z$41 and CLIP $z$42.

6. Applications, limitations, and interpretive cautions

The residual-stream ATO paper presents safety, debugging, and correction as direct applications. Distinguishing transported from synthesized harmful features helps localize responsible computation: transported harmful features point to upstream generators or attention routes, whereas synthesized ones implicate later nonlinear mechanisms. ATOs also support early detection and correction of mistakes, because well-predicted downstream features can in principle be corrected by injection of transported predictions. The method is compute-light, requires no finetuning, and the reported total compute is less than $z$43 GPU-hours; experiments were performed in float32 on Apple M1 Pro and M2 Max hardware (Szablewski et al., 24 Aug 2025).

The CAT paper presents inference-time safety steering without retraining the backbone. Its practical appeal is that overhead is limited to a small MLP forward, a gating function per steered layer, and a broadcast add. The paper explicitly states, however, that inference-time steering does not remove capability and can be bypassed under distribution shift or adversarial prompting. It also notes that operating on mean-pooled activations may miss spatially localized unsafe features, and that evaluation relies on ShieldGemma rather than human labels (Chrabąszcz et al., 3 Mar 2026).

Both formulations impose assumptions that delimit their scope. In the residual-stream setting, ATO captures only linear transport channels; nonlinear synthesis is not modeled. Evaluation depends on SAE decoder projections, so mis-specified or non-orthogonal dictionaries may bias projections relative to SAE latents. The same-token $z$44-policy emphasizes local transport and does not capture attention-mediated routing across positions. Results are reported on selected layers of a single model, Gemma 2 (2B), with Gemma Scope post-layer residual SAEs of latent size $z$45, so generalization requires replication.

In the CAT setting, the key open issue is the balance between suppression and preservation. The paper’s own experiments show that linear methods can drive ASR very low by collapsing image quality, while CAT trades some safety reduction for substantially better CLIP preservation. This suggests that conditioned and regularized transport is most naturally interpreted as a mechanism for constraining where steering acts, rather than a mechanism for eliminating unsafe capability altogether.

Future directions are also explicit in the sources. The residual-stream paper lists nonlinear transport operators, richer $z$46-policies guided by attention readers or Information Flow Routes, multi-source operators, attention-mediated transport analysis, better theoretical bounds, and integration with circuit discovery and feature-targeted editing. The CAT paper’s geometric framing suggests, without claiming formal guarantees, that identity initialization and the $z$47 term encourage small, smooth displacements on already-safe points, thereby stabilizing training and inference-time behavior.