---
title: Lightweight Residual Probes
url: https://www.emergentmind.com/topics/lightweight-residual-probes
type: topic
---

# Lightweight Residual Probes

Lightweight residual probes are small readouts attached to a model’s residual stream or hidden states, designed to reuse computation already paid for by the backbone rather than invoking a separate classifier, guard model, or evaluation pipeline. In the LLM setting, they are used for safety moderation, sentiment and safety classification, runtime orchestration, structural probing of tool-using agents, and in-training estimation of downstream performance; in text-to-video diffusion, an analogous probe is a lightweight pseudo-token embedding that tests whether an erased concept remains recoverable in frozen representations [2606.10487] [2601.13288] [2605.25310] [2604.01025] [2603.21547].

## 1. Conceptual basis and scope

In transformer language models, the basic object of analysis is the hidden-state or residual-stream activation. One formulation uses
$$
h_t^\ell \in \mathbb{R}^H,\quad \ell = 1,\dots,L,\ t = 1,\dots,T,
$$
where \(h_t^\ell\) is the residual-stream activation at layer \(\ell\) and position \(t\). A broader tensor view writes
$$
\mathbf{h}^{(l)} \in \mathbb{R}^{T \times d},\quad l \in \{0,\dots,L-1\},
$$
and stacks these as a hidden tensor in \(\mathbb{R}^{L \times T \times d}\). Lightweight probing treats classification or diagnosis as a readout problem over these internal representations rather than as a separate-model problem [2606.10487] [2601.13288].

The term *lightweight* has a precise operational meaning in the cited work. In serving systems, the probe reuses activations that are already computed during decoding, adds no additional forward pass, and contributes only a small projection, gate, MLP, or similarly constrained module. In frozen-model auditing, lightweight can also mean that all backbone weights remain frozen and only a small probe object is optimized, such as a pseudo-token embedding in text-to-video diffusion or a low-capacity linear edge classifier over pooled residual states [2606.10487] [2603.21547] [2605.25310].

A recurring theme across these studies is that residual representations already encode signals that downstream systems often recover expensively by other means. Safety and compliance signals appear in mid-network layers of aligned LLMs; token- and layer-selective classifiers can recover safety, sentiment, and emotion labels from a single serving pass; runtime tool-call dependency graphs are linearly decodable at tool-call boundaries; and checkpoint-level downstream success probabilities can be predicted from internal states without full generative evaluation [2606.10487] [2601.13288] [2605.25310] [2604.01025].

## 2. Architectural forms

One prominent form is the single-layer token probe used for streaming moderation. Given a chosen layer \(\ell\), typically a mid layer, the probe maps \(h_t^\ell\) to a token-level logit \(z_t\). The simplest form is the linear-token probe,
$$
z_t = w^\top h_t^\ell + b,
$$
while the MLP-token probe adds a hidden layer with \(d = 1024\), and the Dual-head mixes linear and MLP logits with \(\alpha \approx 0.8\). In the moderation setting, these heads operate on a single mid layer, around layer 21 of an 8B Qwen3-like model, and remain tiny relative to the backbone [2606.10487].

A second architectural family treats probing as token- and layer-selective representation selection over the full \(L \times T \times d\) tensor. The two-stage aggregator first compresses tokens within each layer to a layer summary \(\mathbf{v}^{(l)} \in \mathbb{R}^d\), then aggregates across layers to a single vector \(\mathbf{v} \in \mathbb{R}^d\), followed by a linear classifier. The cited implementations use direct pooling, a scoring attention gate with about **0.10M** added parameters, and a downcast multi-head self-attention probe with about **35M** added parameters at a strong operating point. This design explicitly avoids committing to one token or one layer [2601.13288].

A third form is the low-capacity structural probe over paired residual states in agents. For tool-call dependency decoding, call-level representations \(\bar{H}_i\) and \(\bar{H}_j\) are obtained by mean-pooling selected layers at tool-call boundaries, then concatenated as \(x_{ij} = [\,\bar{H}_i ; \bar{H}_j\,]\). A linear logistic regressor with L2 penalty, \(C = 0.01\), solver `lbfgs`, `max_iter=2000`, and `class_weight='balanced'` predicts whether the ordered pair \((i,j)\) is an oracle edge in the runtime tool-call dependency DAG. The decisive point is that a linear probe is sufficient: a deeper MLP attains AUROC **0.866**, essentially the same as the linear probe’s **0.869** [2605.25310].

A fourth form is the in-training performance probe. The main “Submodel” probe is a small decoder-only transformer that mixes its own hidden state with projections of the base model’s hidden states,
$$
Z^{(k)} = \text{Decoder}^{(k)}\left( Z^{(k-1)} + W_{proj}^{(k)} H^{(k)}; \phi^{(k)} \right),
$$
and predicts success probability from the final-layer last-token representation. The same study also evaluates a linear probe baseline and a LoRA probe, but the Submodel probe is the main residual-space design [2604.01025].

A fifth form appears in multimodal safety auditing. PROBE in text-to-video diffusion optimizes only a pseudo-token embedding \(\mathbf{v}\) or a short sequence of such embeddings, inserted into the frozen text encoder input space. The text encoder and denoising model remain frozen; the probe is therefore low-capacity, localized, and diagnostic rather than a new model training procedure [2603.21547].

## 3. Supervision and optimization objectives

The supervision regime depends on the target property. For streaming safety moderation, the key technical move is token-level probing from sequence-level labels. For a sequence \(x_{1:T}\) with label \(y \in \{0,1\}\), the model produces logits \(z_1,\dots,z_T\), applies a pooling operator over tokens, and then uses weighted binary cross-entropy with token weights
$$
w_t = \frac{\exp(\bar z_t / \tau)}{\sum_{t' \in V} \exp(\bar z_{t'} / \tau)},\quad \tau=1,
$$
and loss
$$
\mathcal{L} = \sum_{t \in V} w_t \cdot \text{BCE}(\bar z_t,\, y).
$$
The probe is trained to imitate **Qwen3Guard-Gen-8B**, not human ground truth, on the **Zhuyi** dataset of about **600K** Chinese safety/compliance examples [2606.10487].

For single-pass classification over full hidden tensors, the two-stage probes use standard cross-entropy on the final logits. The scoring attention gate assigns learned weights via
$$
s_i = \tanh(\mathbf{w}^\top \mathbf{X}[i,:]), \quad
\boldsymbol{\alpha} = \mathrm{softmax}(\mathbf{s}), \quad
\mathbf{v} = \sum_{i=1}^{N} \alpha_i \mathbf{X}[i,:].
$$
The MHA probe uses scaled dot-product self-attention on downcast QKV projections and then pools over the sequence dimension [2601.13288].

For runtime dependency decoding in agents, supervision comes from an oracle tool-call dependency DAG derived from actual tool trajectories. The logistic probe predicts
$$
p(e_{i\to j}=1 \mid \bar{H}_i,\bar{H}_j)
= \sigma\big( W [\bar{H}_i;\bar{H}_j] + b \big),
$$
where the label is a direct or transitive oracle edge. The same study emphasizes control experiments: random-label controls, positional baselines, scaffold/template baselines, and surface-form decoders [2605.25310].

For checkpoint-level downstream evaluation during training, the target is the Monte Carlo estimate of success probability,
$$
\hat{v}(s) = \frac{1}{n} \sum_{i=1}^n R(s, a_i),
$$
equated with pass@1 in the paper’s setup. Probe training minimizes mean squared error between the predicted scalar and \(\hat{v}(s)\). This frames probing as approximation of a value function over prompts and checkpoints [2604.01025].

For text-to-video diffusion, PROBE combines a denoising reconstruction objective with latent alignment:
$$
\mathcal{L}_{\text{total}}(\mathbf{v}) = \mathcal{L}_{\text{rec}}(\mathbf{v}) + \lambda \, \mathcal{L}_{\text{align}}(\mathbf{v}),
$$
with \(\lambda \approx 1\) effective in ablations. This objective is explicitly designed to test whether an erased concept remains encoded in a frozen erased model [2603.21547].

## 4. Integration into generation and orchestration

The operational attraction of lightweight residual probes is that they can be embedded inside an existing serving path. In streaming moderation, the probe is inserted into the decode loop with a `vLLM` forward hook. At each decoding step, the generator produces next-token logits and \(h_t^\ell\), the hook reads the chosen layer, and the probe computes \(z_t\) with no additional LLM forward pass. The system then maintains a running score using modes such as `raw`, `ema`, `window_max`, `topk_mean`, `percentile`, and `logsumexp`, and applies trigger modes such as `instant`, `cumulative / consecutive`, `decay / budget`, and `warmup`. The `probe interval` parameter `pi` allows probing every `pi` decode steps, with **pi = 2** reported as an excellent compromise [2606.10487].

The same single-pass logic appears in general orchestration. Token- and layer-selective probes are attached to the serving LLM forward, consume the already-computed hidden tensor, and return a label for safety, sentiment, or emotion in the same LLM call. The cited implementation requires no separate guard model, no second pass, and no special batching or KV-cache modifications beyond ordinary backbone serving [2601.13288].

In agent settings, the probe attaches at decode time to tool-call boundaries rather than every token. Residual vectors are captured for every assistant token, boundary states are selected at tool-call termination, and the resulting call-level features are used to decode the runtime tool-call dependency graph. This is an orchestration probe in the sense that it reads execution structure from the same run that produced the tool trajectory [2605.25310].

In training-time evaluation, the serving constraint is replaced by a monitoring constraint. The probe is trained once on checkpoint-specific hidden states and pass@1 labels, then applied to later checkpoints using a single forward pass without generation. The paper’s three-phase workflow is: probe preparation with full generative evaluation on a training subset, representation extraction, and later in-training monitoring by probe inference alone [2604.01025].

## 5. Residual directions, internal structure, and control

The linear branch of a residual probe can be interpreted geometrically as a direction in residual space. In the moderation setting, \(w \in \mathbb{R}^H\) defines a hyperplane \(w^\top h + b = 0\), an unsafe half-space \(\{ h : w^\top h + b \ge 0 \}\), and a unit vector \(u = \frac{w}{\|w\|}\) associated with higher probability of unsafe content. Detection is projection onto this direction. The same direction can be reused for activation steering by subtracting
$$
\beta \frac{w}{\|w\|}
$$
from the hidden state when a trigger fires. The paper describes this as a read–write symmetry: project onto \(w\) to detect, move along \(-w\) to suppress [2606.10487].

Residual probes also expose structured latent variables beyond scalar labels. In tool-using agents, a low-capacity edge probe decodes the runtime tool-call dependency DAG well above both a Hewitt–Liang random-label control and a positional baseline. The direct-edge probe achieves AUROC **0.869**, while the transitive-only probe reaches **0.986**. Counterfactual analysis distinguishes value corruption from structural perturbation: the decoded graph is nearly invariant to single-field ID swaps but responds strongly to perturbations that alter plan structure. Per-layer activation patching shifts the probe at a later, non-patched boundary, which the paper interprets as evidence that the representation propagates rather than passively reads out, while the realized tool call does not move [2605.25310].

In diffusion models, the same general logic appears as residual concept capacity. PROBE shows that concept erasure methods can leave measurable residual capacity even when output-level suppression looks successful. The protocol detects temporal re-emergence, a video-specific failure mode in which suppressed concepts progressively resurface across frames, invisible to frame-level metrics. This supports the paper’s distinction between output-level suppression and representational removal [2603.21547].

A plausible implication is that lightweight residual probing is not confined to binary moderation. The cited work already spans per-token safety scores, token- and layer-selective readouts for sentiment and emotion, runtime graph decoding in agents, checkpoint capability prediction, and multimodal residual-capacity auditing. The shared mechanism is a constrained readout over internal computation, not an auxiliary model family.

## 6. Empirical profile, limitations, and deployment patterns

The moderation results illustrate the central accuracy–latency trade-off. On the held-out **Zhuyi** test split, the hand-designed **MLP-token**, layer 21, max aggregation probe reports **F1 = 0.941, AUC = 0.942, precision = 0.961**; the **Linear-token** probe reports **F1 = 0.892, AUC = 0.897, precision = 0.932**; and the full-data **Dual-head** reaches **F1 = 0.9746, AUC = 0.9865, precision = 0.985, recall = 0.965** on validation. In efficiency terms, the probe adds **8 MB** VRAM, about **0.6 GFLOP per request**, and about **0.1 ms P50/P95/P99**, versus **15.5 GB** and **~448 / 551 / 602 ms** for post-hoc **Qwen3Guard-Gen-8B**, and **14.2 GB** and **~39 / 43 / 54 ms** for **Qwen3Guard-Stream-8B**. On **BeaverTails**, the probe is weaker than dedicated guards against human labels, but calibrated thresholding raises it to **F1 = 0.795, precision = 0.798, recall = 0.792** [2606.10487].

The token- and layer-selective probes show a similar pattern on general classification workloads. On **ToxicChat**, direct pooling reaches **73.53 F1, 0.812 AUPRC**, the scoring attention gate reaches **80.49 F1, 0.854 AUPRC**, and the **MHA probe** reaches **84.51 F1, 0.898 AUPRC** with **35M** added parameters. On **WildGuardMix**, the MHA probe reaches **88.55 F1**, close to **88.9 F1** for the **WildGuard** 7B guard model. In latency measurements on a **Llama-3.2–3B** backbone, the MHA probe runs at **24.83** samples/s and **40.27 ms / sample**, versus **8.12** samples/s and **123.21 ms / sample** for a **Guard→Serve (8B+3B pipeline)**, while peak GPU remains near **6.5–7.0 GB** rather than **22768.72 MB** [2601.13288].

The in-training evaluation probes trade generative evaluation for residual-space estimation. The **Submodel** probe reports average **AUROC: 0.7890** and average **MSE: 0.1049** on pre-training **OLMo3-7B** checkpoints, compared with **0.5810** for the linear probe and **0.5795** for the loss-fit baseline. Probe evaluation reduces per-checkpoint latency from about **0.78 hours** to about **0.05 hours**, a **15.6×** speedup, with reported speedups of **61.4×** on the instruct model and **231.8×** on the think model [2604.01025].

The literature also delineates clear limitations. Streaming moderation probes are **not a universal moderator**; their best performance is agreement with a specific guard on Chinese data, they require calibration, and they are weaker under domain shift or adversarial obfuscation. The agent-structure study explicitly limits its claims to representation, not behavioural control. The diffusion study argues that failure to recover a concept is not proof of true erasure, and that current erasure methods often achieve output-level suppression rather than representational removal [2606.10487] [2605.25310] [2603.21547].

The deployment patterns that follow from these results are correspondingly conservative. The moderation paper recommends the probe as a **first-stage, always-on screen** and proposes cascades in which expensive guards are reserved for flagged cases. The token- and layer-selective classification paper presents probes as orchestration primitives that preserve near-serving latency while avoiding the VRAM and latency costs of a separate guard-model pipeline. This suggests a stable encyclopedia-level characterization: lightweight residual probes are low-capacity readouts or control handles on residual computation, most valuable when latency, VRAM, or repeated evaluation cost makes a separate model operationally prohibitive [2606.10487] [2601.13288].

Source: https://www.emergentmind.com/topics/lightweight-residual-probes