PaperlayerView exact revision on arXiv

Source-linked AI summary

Value Probes: Learned Value Functions for Beam-Search Red Teaming

Jonathan F. Jao, Andrew Kyle Lampinen

arXiv:2610.38526v1cs.AIcs.LG

TL;DR

Safety-relevant behaviors may surface only in particular multi-turn contexts, making automated red-team audits an exploration problem. The paper uses beam search guided by Value Probes, which read discounted future judge scores from target-model activations. Across 18 behaviors and three Gemma 3 sizes, the probes improve mean peak audit scores over token-matched Petri, and smaller-model rollouts can train probes for larger targets.

  • Problem

    Safety-relevant behaviors can emerge only in particular multi-turn contexts, while current audits use costly model evaluations and do not use the target model’s internal states.

    Source: §1 · pp. 1–2

  • Method

    The paper trains linear probes on target-model activations to predict discounted judge scores and uses the estimates to rank branches in beam-search audits.

    Source: §1 · p. 3, §3.1 · p. 5, §3.3 · p. 6

  • Results

    Value Probe beam search raised mean peak judge score over token-matched Petri by 0.22–0.34 across 18 behaviors on Gemma 3 4B, 12B, and 27B, with non-overlapping 95% confidence intervals at every size.

    Source: §7 · p. 11

  • Takeaways & Limitations

    Off-policy probes trained on Gemma 3 4B rollouts transferred to Gemma 3 12B and 27B, matching on-policy probes to within about 0.01 in mean peak judge score.

    Source: §7 · p. 11

    The evaluation covers Gemma 3 models and 18 behaviors, and the heuristic requires target activations, limiting its application to open-weight models or labs auditing their own models.

    Source: §6 · p. 10

Abstract

from arXiv · show

Automated red-team auditing of language models is an important way to evaluate safety-relevant behaviors. These audits implicitly search a tree of possible conversation pathways: at each turn an auditor model decides which conversation to continue, and an LLM judge scores the dialogue. We emphasize that this auditing procedure is a Markov decision process, where states are the conversation history, actions are candidate responses, and rewards are judge scores for the target behavior. By adopting this perspective, we introduce two major improvements to auditing. First, we propose adopting beam search for the audit rather than relying on the auditor alone. We then interpret the problem of finding states that are likely to lead to problematic behavior as value estimation. Surprisingly, we show that this value function can be estimated by linear probing from the target model's activations to predict discounted future rewards - a method we term "Value Probes." By using Value Probes as a heuristic in our beam-search audit, we are able to substantially improve auditing success. Across 18 target behaviors and 3 scales of models, our approach raises peak audit scores by roughly 50% compared to (token-matched) standard auditing approaches. Furthermore, we show that Value Probes can be fit using off-policy data from smaller models. We also ablate individual components, and show that our method offers a comparably-performant and more efficient alternative to more expensive language model follow-ups. We hope these methods will improve language model auditing and lead to safer models.

1 Introduction

Safety-relevant behaviors may emerge only in particular multi-turn contexts, so automated red teaming must explore conversation pathways. The paper frames auditing as sequential decision-making and proposes beam search guided by Value Probes that estimate future judge scores from target-model activations.

1 Introduction

  • Safety-relevant behaviors can surface only in particular contexts several turns into interactions, making exploration necessary beyond fixed prompt sets.Automated red teaming uses an auditor model to investigate the target and an LLM judge to score transcripts.Source: §1 · p. 1
  • Auditing decisions are costly and black-box: current pipelines repeatedly use high-capacity models on transcripts while leaving the target’s internal states unused.The cost of evaluation constrains how much of the conversation tree can be explored.Source: §1 · pp. 1–2
  • The paper models auditing as a sequential decision problem and uses discounted future judge scores as a value heuristic for ranking conversation branches.This framing separates the search strategy from its guiding estimator, allowing auditor-driven search and beam search to use the same heuristic.Source: §1 · pp. 2–3
  • Value Probes linearly read target-model activations to estimate value, and the paper introduces them as a heuristic for guiding audits.The hypothesis is that activations encode how close the current conversation is to eliciting the target behavior.Source: §1 · p. 3
  • Beam search and Value Probes, separately and together, substantially improve auditing success, and probes can also be fit from smaller models’ off-policy data.The paper presents these as complementary contributions to the audit procedure.Source: §1 · p. 3

2 Related Work

Prior work explores language-model red teaming, activation-based representation reading, and learned values for search. This paper combines these lines by searching adversarial conversations and reading value estimates from the target model’s activations.

Automated red teaming and multi-turn attacks

  • Multi-turn attacks build on conversation context to elicit behaviors, while automated auditing extends red teaming to broader misalignment scenarios scored by LLM judges.The harness retains an auditor that proposes turns and a judge that scores transcripts, but uses a probe in place of the judge only for branch ranking.Source: §2 · pp. 3–4

Probes and representation reading

  • Prior activation-reading work classifies current safety-relevant properties, whereas Value Probes predict discounted future judge outcomes rather than detect the current state.The paper focuses on reading representations, not interventions that alter behavior.Source: §2 · p. 4

Search with learned value functions

  • Language-model search has used self-evaluation, learned values, and process reward models; Value Probes instead read value linearly from the target model’s activations for adversarial conversations.The scored object is a conversation between two models, rather than a reasoning step.Source: §2 · p. 4

3 Methods

The method separates conversation-tree search from the value estimator that ranks branches. It trains linear probes on discounted judge scores and uses them to guide beam search over auditor-proposed turns.

3.1 Red Teaming as a Markov Decision Process

  • Auditing is a search over conversation prefixes: the auditor proposes the next user turn, the target stochastically replies, and the judge score is the reward.The Bellman equations define discounted prefix values under the auditor’s proposal policy and under optimal continuation.Source: §3.1 · p. 4
  • A Value Probe estimates the value of a conversation prefix from target-model activations, independently of the search strategy that uses the estimate.This separation allows either Petri’s auditor-driven depth-first search or beam search to use the same estimator.Source: §3.1 · p. 5

3.2 Value Probes

  • Value targets are backed up from rollout leaves with γ = 0.9; the Expected Value target averages explored children, while the Max Value target takes their maximum.Both targets use the node’s judge score, and leaves take their observed score as value.Source: §3.2 · p. 5
  • A linear probe maps mean-pooled residual-stream activations from one target-model layer to the value target, with squared-loss training and the target model frozen.Probes are trained per target behavior; the paper reports only mean pooling and one layer per probe.Source: §3.2 · p. 5
  • Off-policy probes use transcripts and judge scores from an older or smaller source model, recomputing activations with the frozen target model while excluding its own replies from training.This supplies positive judge-score examples when a well-aligned target produces few of them.Source: §3.2 · pp. 5–6

3.3 Beam search over the conversation tree

  • Beam search expands b live nodes into b k auditor-proposed children, scores them, keeps the top b, and discards the remaining branches and subtrees.The harness re-scores at each turn, so a previously retained branch can be dropped when a sibling’s continuation scores higher; search stops at turn budget T.Source: §3.3 · p. 6
  • The uniform-random beam arm isolates search structure from ranking: it selects among the auditor’s proposed candidates uniformly, while Petri is token-matched to the beam arms.Comparing random selection with Petri measures the search-structure effect; comparing probe or judge ranking with random selection measures the ranking-signal effect.Source: §3.3 · p. 6

4 Experimental Setup

Experiments evaluate the audit harness on 18 behaviors across three Gemma 3 target sizes, with Gemma 4 used only for transfer. The main comparison uses token-matched Petri and beam-search arms with random, probe, or judge ranking.

Models and behaviors

  • The main targets are instruction-tuned Gemma 3 models at 4B, 12B, and 27B parameters, tested on 18 behaviors spanning sycophancy, blackmail, and animal-welfare disregard.Gemma 4 31B is included only in transfer experiments.Source: §4 · pp. 6–7

Harness and arms

  • Each beam-search arm uses width b = 4, k = 2 candidate turns per live node, and T = 12, for 96 target continuations per audit.The auditor is Claude Sonnet 4.6; the ranking arms include random selection, Expected or Max Value Probes, and a Claude Opus 4.8 judge.Source: §4 · p. 7

Baselines, probes, and metric

  • Petri is token-matched to beam search, with auditor and target tokens matched separately and the later cutoff used; probes are single-layer, behavior-specific, and trained on same-size rollouts in the main comparison.Peak judge score is the maximum behavior-judge score over nodes; each arm, model size, and behavior has 10 replicate audits on the same scenarios.Source: §4 · p. 7

5 Results

Beam search with Value Probe ranking improves peak audit scores over token-matched Petri, while the probes perform comparably to a separate LLM judge. Probes trained on Gemma 3 4B rollouts also transfer to larger targets.

5.1 Beam search versus a Petri audit

  • At 4B, Value Probe arms score 0.30–0.34 above token-matched Petri, versus gains of 0.22–0.24 at 12B and 27B.Probe and judge arms have disjoint across-behavior intervals from Petri’s at every size.Source: §5.1 · p. 7

5.2 Attribution: search structure and value signal

  • Beam search with random selection beats Petri by +0.24 ± 0.11 at 4B and +0.15 ± 0.11 at 12B and 27B, winning 14 or 15 of 18 behaviors.The comparison is paired by behavior; breadth and re-scoring of siblings contribute at equal token spend.Source: §5.2 · pp. 7–8, Figure 2 · p. 8
  • Adding a ranking signal to random selection raises mean peak judge score at every size; probe gains are +0.06 to +0.10, with all six paired intervals excluding zero.Across the three ranking signals, each beats random selection on 14 to 17 behaviors in paired comparisons.Source: §5.2 · p. 8, Figure 2 · p. 8

5.3 Activations rival a high-capacity judge

  • Across model sizes, each probe is within 0.05 of the judge in paired mean peak score; Max Value Probe is within 0.04 at every size.The judge heuristic uses 288 extra Claude Opus 4.8 calls per audit, while the probe reads activations computed during target-model generation.Source: §5.3 · p. 8

5.4 Max versus Expected Value Probes

  • At 4B, Max Value Probe leads Expected Value Probe by +0.04 ± 0.03; at 12B and 27B, paired differences are within ±0.02 with intervals including zero.The reported distinction between the Bellman and Bellman optimality backups is negligible beyond the smallest tested model.Source: §5.4 · p. 9

5.5 Off-policy transfer

  • On Gemma 3 12B and 27B, probes trained on Gemma 3 4B rollouts are indistinguishable from on-policy probes by paired comparisons.All four transferred-minus-on-policy means are between 0.000 and +0.011, with 95% intervals including zero.Source: §5.5 · p. 9
  • On Gemma 4 31B, probes trained on Gemma 3 4B rollouts score 0.79 ± 0.10 (Max) and 0.79 ± 0.08 (Expected), versus 0.68 ± 0.11 for random selection.Each transferred probe beats random by +0.11 in paired mean, with intervals excluding zero and 15 of 18 behaviors won.Source: §5.5 · p. 9

6 Discussion

The results support combining a learned state-value estimate with beam-search lookahead: Value Probes improve search and transfer across model sizes. The tested setup and model scope leave several questions open.

Discussion

  • A linear probe on one layer’s mean-pooled activations adds 0.06–0.10 in mean peak judge score over random ranking, and performs within a few hundredths of a separate LLM judge.The probe is a learned state-value estimate, while beam search supplies shallow lookahead.Source: §6 · pp. 9–10
  • The probe direction learned from Gemma 3 4B audits ranks branches for Gemma 3 12B, 27B, and next-generation Gemma 4 31B.The discussion relates this value-plus-search pattern to game-playing systems and search over reasoning steps guided by learned or self-assessed values [Silver et al., 2016, Feng et al., 2024, Lightman et al., 2023, Snell et al., 2024, Yao et al., 2023].Source: §6 · p. 10
  • With two candidates per node and peak-only scoring, the Max and Expected Value Probes coincide at 12B and 27B; larger branching or timing-sensitive metrics are needed to distinguish them.The authors offer correlated targets, limited probe capacity, and propagated judge noise as possible explanations.Source: §6 · p. 10
  • Only mean pooling, γ = 0.9, and one layer per model size were tested; results cover Gemma 3 and Gemma 4 31B transfer, not other model families.The heuristic requires target activations, limiting application to open-weight models or labs auditing their own models; cross-family probe stability is untested.Source: §6 · p. 10, Figure 3 · p. 10
  • Monte-Carlo search, activation steering with the probe direction, and broader agentic behavior sets are proposed as future directions, not tested findings.The paper also identifies turns-to-elicit and area under the score-versus-turn curve as alternatives to peak score.Source: §6 · pp. 10–11

7 Conclusion

The paper frames red-team auditing as search over conversation trees and finds that a linear probe of target-model activations can guide that search effectively. Value Probes improve mean peak judge scores over token-matched Petri audits, and probes trained on smaller-model rollouts transfer to larger targets.

7 Conclusion

  • A Value Probe raised mean peak judge score over token-matched Petri audits by 0.22–0.34 across 18 behaviors, with non-overlapping 95% confidence intervals at every tested model size.It was on par with a separate high-capacity LLM judge ranking the beams.Source: §7 · p. 11
  • Probes trained on Gemma 3 4B rollouts matched target-trained probes within about 0.01 mean peak judge score on Gemma 3 12B and 27B, and also guided search on Gemma 4 31B.The authors report similar mean peak judge scores on Gemma 4 31B as on Gemma 3 models.Source: §7 · p. 11
  • The comparison with beam search using random selection indicates that both search structure and the value signal contribute to audit performance.The paper also reports that a linear probe estimates propensity for interesting behaviors from activations alone.Source: §7 · p. 11

A.1 Results tables

The appendix provides tables listing the values plotted in Figures 2 and 3. They cover the main comparison and the transfer experiments, respectively.

A.1 Results tables

  • Table 1 lists mean peak judge score and 95% t-CI half-width over 18 behaviors for every arm of the main comparison.Its values come from the same data snapshot as Figure 2.Source: Table 1 · p. 17
  • Table 2 lists mean peak judge score and 95% t-CI half-width over 18 behaviors for every arm of the transfer experiments.Its values come from the same data snapshots as Figure 3.Source: Table 2 · p. 17
  • The Gemma 4 31B values in Table 2 use a post hoc Claude Opus 4.8 judge, so compare them only within that column.Paired differences quoted in the text can differ by 0.01 from differences calculated from rounded Table 1 entries.Source: Table 2 · p. 17, Table 1 · p. 17

A.2 Probe construction

The probe construction uses a trained linear probe to predict continuous discounted judge targets from target-model activations. The appendix also compares ways to construct training pairs, finding that composite datasets transfer best across strategies.

A.2 Probe construction

  • The selected construction regresses a continuous discounted judge target onto mean-pooled activations from one layer of the frozen target model, training only the linear probe weights.The paper contrasts this with mass-mean probes, which require collapsing the graded signal into two classes and discard the Bellman structure distinguishing Expected from Max Value Probes.Source: §A.2 · pp. 17–18
  • Synthesized and system-prompt pairs scored 0.94 in-distribution but transferred worst, with mean off-diagonal AUCs of 0.51 and 0.43, while fork-higher and edit-in-place scored 0.64 and 0.66 in-distribution and transferred at 0.61 and 0.60.The authors associate the weaker transfer of synthesize with a cross-model style signature and system-prompt with a length confound.Source: §A.2 · p. 18
  • Composite training reached mean transfer AUC 0.72 across four targets, versus 0.61 for the best single strategy, and was the only source type handling the system-prompt target at 0.73–0.85.Single strategies reached at most 0.67 on that hardest target.Source: §A.2 · pp. 18–19, Figure 6 · p. 19
  • Figure 6 compares probes trained on each pair-construction strategy against held-out tests for every strategy, with composite training evaluated as pooled or balanced mixes of cleaned or original datasets.Its nodes show in-distribution AUC and its directed edges show cross-strategy AUC.Source: Figure 6 · p. 19

A.3 Off-policy transfer

Off-policy training reuses rollouts and judge scores from an older or smaller source model, while recomputing activations with the frozen target model. The transfer claim is tested by deploying those probes as beam-search heuristics on targets with different reply distributions.

A.3 Off-policy transfer

  • Off-policy training addresses the practical constraint that reward signals can be sparse for well-aligned targets.Figure 7 depicts the source rollout and judge-score path alongside the target-model activation path.Source: §A.3 · pp. 19–20, Figure 7 · p. 20
  • Off-policy probes reuse an older or smaller source model’s transcripts and per-node judge scores, then run those transcripts through the frozen target model to collect activations and fit the probe.This avoids using the target model’s own replies during training, which the authors note can be scarce for well-aligned targets.Source: §A.3 · p. 20
  • The transfer test asks whether a probe targeting V π (Expected) or V ∗ (Max) under the source reply distribution remains a useful ranking heuristic when target-model responses differ.Figure 3 compares the transferred probe as a beam heuristic alongside on-policy probes and random selection, without AUC-based filtering of behaviors.Source: §A.3 · p. 20

A.4 Meta-information leakage gate

This appendix describes checks for metadata leakage, alternative and hybrid beam-search strategies, and the judge prompts used to score behavior and forecast future behavior.

A.4 Meta-information leakage gate

  • Behavior names, seed identifiers, and auditor instructions can leak through judge scores or mean-pooled activations, creating shortcuts unrelated to the target behavior.The seed text is a strong, constant feature across nodes in a rollout.Source: §A.4 · p. 20
  • The gate checks judge robustness through redaction and metadata-only tests, and probe robustness through held-out seeds, prefix-only prediction, and permuted-label controls.The checks test whether scores or predictions persist when behavioral evidence is removed or metadata cues are controlled.Source: §A.4 · pp. 20–21

A.5 Monte-Carlo beam search

  • Monte-Carlo beam search ranks candidates by sampled terminal judge returns, trading fixed-width beam search’s predictable, parallel cost for adaptive but sequential playouts.It replaces the probe’s one-step estimate with a sampled estimate of the same quantity; evaluation is left for future work.Source: §A.5 · p. 21
  • A proposed hybrid would use Value Probes to truncate playouts and prioritize expansions, while real playouts resolve close scores and judge flags provide terminal rewards.Averaging and max backups correspond to Expected and Max Value Probe targets, respectively.Source: §A.5 · p. 21

A.6 Judge prompts

  • The post-hoc behavior judge scores each node, while the future-LLM heuristic uses the same user message with a different system prompt and closing question.The post-hoc judge runs once per node at temperature 0 when supported; the in-loop heuristic uses the median of three samples.Source: §A.6 · pp. 21–22

A.6.1 Behavior judge (post-hoc metric)

  • The behavior judge rates target replies so far, whereas the prospective judge forecasts target behavior in upcoming turns from the trajectory so far.Both prompts direct the judge to ground scores in the target’s replies rather than the auditor’s intent alone.Source: §A.6.1 · p. 22, §A.6.2 · pp. 22–23

Related papers

Loading 2610.38526v1…