PaperlayerView exact revision on arXiv

Source-linked AI summary

Building Production-Ready Probes For Gemini

János Kramár, Joshua Engels, Zheng Wang, Bilal Chughtai, Rohin Shah, Neel Nanda, Arthur Conmy

arXiv:2601.11516v4cs.LGcs.AIcs.CL

TL;DR

The paper addresses the difficulty of making activation probes generalize across production distribution shifts, especially from short to long contexts. It develops new probe architectures, automated search, and probe–LLM cascades, finding strong performance and low-cost monitoring while identifying persistent limits under adaptive attacks and broader shifts.

  • Problem

    Activation probes are promising and inexpensive misuse monitors, but existing probes are fragile under production distribution shifts, particularly long-context inputs.

  • Method

    The paper evaluates new probe architectures, AlphaEvolve-based architecture search, and cascades that defer to prompted language models when probes are uncertain.

  • Results

    The best probes achieve comparable performance to Gemini 2.5 Flash and Pro classifiers at a fraction of the cost, while cascading uses the LLM less than 10% of the time and achieves lower FNR than the LLM alone.

  • Takeaways & Limitations

    The findings informed successful deployment of cyber-attack misuse-mitigation probes in production and show that architecture choice and diverse training support broader generalization.

  • Takeaways & Limitations

    Training probes on model outputs remains an important future direction, and the techniques do not significantly reduce adaptive adversarial attack success rates.

Abstract

from arXiv · show

Frontier language model capabilities are improving rapidly. We thus need stronger mitigations against bad actors misusing increasingly powerful systems. Prior work has shown that activation probes may be a promising misuse mitigation technique, but we identify a key remaining challenge: probes fail to generalize under important production distribution shifts. In particular, we find that the shift from short-context to long-context inputs is difficult for existing probe architectures. We propose several new probe architectures that handle this long-context distribution shift. We evaluate these probes in the cyber-offensive domain, testing their robustness against various production-relevant distribution shifts, including multi-turn conversations, long context prompts, and adaptive red teaming. Our results demonstrate that while our novel architectures address context length, a combination of architecture choice and training on diverse distributions is required for broad generalization. Additionally, we show that pairing probes with prompted classifiers achieves optimal accuracy at a low cost due to the computational efficiency of probes. These findings have informed the successful deployment of misuse mitigation probes in user-facing instances of Gemini, Google's frontier language model. Finally, we find early positive results using AlphaEvolve to automate improvements in both probe architecture search and adaptive red teaming, showing that automating some AI safety research is already possible.

1. Introduction

The paper studies activation probes for detecting cyber-offensive prompts in Gemini and addresses their fragility under production distribution shifts, especially long-context inputs. It introduces architectural and training strategies that improve probe performance, reduce monitoring cost, and inform production deployment, while leaving adaptive attacks challenging.

  • Activation probes offer a cheaper alternative to LLM monitors because they reuse activations generated during the monitored model’s forward pass.
  • Probes trained on short-context data fail to generalize reliably to important production shifts, particularly long-context inputs.Training directly on long-context data can increase training costs by over an order of magnitude.
  • The paper investigates improved probe architectures, automated architecture search, and cascading probe–LLM classifiers to address distribution-shift performance degradation.
  • MultiMax improves long-context accuracy, AlphaEvolve discovers competitive probes, and using an LLM only when a probe is uncertain yields stronger classifiers at a small fraction of LLM inference cost.
  • Best probes achieve comparable performance to Gemini 2.5 Flash and Pro classifiers at a fraction of the cost, while probes outperform 2.5 Flash Lite.
  • The findings informed successful production deployment of cyber-misuse probes, but the techniques do not significantly reduce adaptive adversarial attack success rates.

2. Setup and Notation

The setup classifies token sequences using probes trained on a monitored language model’s hidden states. A probe score is thresholded into a misuse decision after a single forward pass, with activation-based inference providing a lightweight monitoring alternative.

  • Each input sequence is passed through the deployed model, and hidden states from a fixed middle layer are harvested for probe training.The hidden states form a matrix with one dimension for representation size and one for sequence length.
  • Probes are small logistic-regression or neural-network classifiers trained on the monitored model’s activations to classify sequences.
  • At inference, the trained probe processes activations from a single forward pass, including inputs from a potentially different distribution.
  • A probe score above threshold τ triggers a misuse mitigation instead of continuing generation naively.
  • The thresholded probe function produces a binary classification decision, analogous to the prompted classifier.
  • Figure 1 compares weighted test error against cost, with thresholds selected by validation-set error minimization and probe costs artificially shifted for visual separation.

3. Our Classifiers

The paper organizes classifiers as transformations of per-token activations followed by aggregation into a scalar score, covering baseline probes, novel long-context architectures, prompted LLMs, and cascades. Its main architectural changes target signal preservation in long contexts and efficient inference.

  • Classifier framework: The classifier framework transforms residual-stream activations per position, aggregates them, and produces a single scalar score.
  • Baseline probes: Linear probes use an identity transformation, a linear score, and mean aggregation, while EMA probes take the maximum of exponential moving-average scores.
  • Baseline probes: MLP probes apply a multilayer perceptron before producing a scalar score from pooled or single-token activations.
  • Attention probes: Attention probes learn query and value vectors so aggregation weights depend on the activations themselves.
  • Attention probes: Attention probes have an inexpensive inference algorithm, avoiding the need for expensive softmax recomputation for every generated token.
  • Novel architectures: MultiMax replaces softmax weighting with a hard maximum at inference, selecting the highest-scoring token per head to prevent signal dilution in long contexts.
  • Novel architectures: The Max of Rolling Means Attention Probe combines sliding-window attention-weighted averages with a maximum over window outputs.
  • Automated search: AlphaEvolve generates candidate architectures from an attention-probe seed using weighted FPR and FNR as its optimization target.

4. Datasets and Metrics

The evaluation uses cyber-misuse datasets with separate train, validation, and test distributions, including benign overtriggering and true cyber attacks. Probe thresholds and test losses are weighted to strongly penalize overtriggering, while seed and architecture selection are assessed separately.

  • Datasets: The paper focuses on cyber misuse, evaluating benign datasets where classifiers should not fire and cyber-attack datasets where they should.Benign examples target low logits and minimized false positive rates; attacks target high logits and minimized false negative rates.
  • Datasets: No prompts are shared across train, validation, and test splits, and test sets represent distribution shifts.Validation and test distributions are also kept different so thresholds are not optimized on the test distribution.
  • Datasets: Long-context data is intentionally excluded from probe training because it substantially increases infrastructure burden and compute.The training design therefore evaluates generalization to a shift that is not directly represented in training data.
  • Metrics: The weighted validation objective assigns penalties of 5 to FNR, 2 to hard-negative FPR, and 50 to overtriggering FPR.These weights make overtriggering false positives 10 times more costly than false negatives and 25 times more costly than hard-negative false positives.
  • Metrics: The weighted test loss averages FPRs and FNRs across all test datasets and serves as the main single-axis comparison metric.Thresholds are selected on validation datasets, while the same weighted formulation is evaluated on test datasets.
  • Initialization Seeds: Across architectures, selecting the lowest-validation-loss seed reduces test loss by 0.008 on average, but the median reduction is 0.003.The best architecture reaches 0.025 test loss, whereas Linear Probe Mean reaches 0.064, making architecture choice more influential than median seed selection.
  • Initialization Seeds: The best possible test-set seed is often substantially better than the validation-selected seed, suggesting validation-data composition constrains seed selection.A better validation data mix might move the selected seed closer to the best possible seed.

5. Results

The results show that new probe architectures substantially improve performance across production-relevant distribution shifts, but no single method dominates every sub-distribution. Short-context training can support long-context performance with suitable architectures, while cascading probes with LLM classifiers offers favorable cost-accuracy trade-offs; adversarial jailbreak robustness remains limited.

  • Main results: Existing probe architectures can be significantly improved, with AlphaEvolve’s MultiMax and Max of Rolling Mean Attention Probes outperforming linear probe baselines.The Max of Rolling Mean Attention Probe with MultiMax aggregation outperforms all evaluated LLMs.
  • Main results: No single method dominates across all distribution shifts, and the optimal probe varies across sub-distributions and tasks.The authors speculate that broader training distributions and the similarity between harmful and benign cyber or coding requests contribute to variation in FPR and FNR.
  • Adversarial robustness: All evaluated techniques leave attack success rates above 1% on the pre-existing jailbreak dataset, indicating that adversarial robustness remains unsaturated.The dataset contains prompt-based attacks generated from publicly accessible methods, and the authors suspect some dataset labeling error affects low performance.
  • Training on long-context data: 2.53% test error for AlphaEvolve approaches the 2.38% achieved by attention probes trained directly on long-context data, while AlphaEvolve uses only short-context training.Direct long-context training incurred a 22× increase in training cost compared to short-context training.
  • Cascading results: The probe used alone lies on the cost-accuracy Pareto frontier, while Gemini 2.5 Flash dominates Gemini 2.5 Pro across that frontier.The cost estimate for Pro is generous because thinking-token costs are ignored.

6. Related Work

Related work frames activation probes as low-cost safety monitors, while emphasizing unresolved robustness challenges from distribution shifts and adaptive attacks. The paper extends automated optimization into probe architecture search and red teaming.

  • Misuse mitigation: Existing alignment and classifier-based safeguards remain insufficiently robust to prevent all misuse.The related work positions external monitors as increasingly standard for frontier models.
  • Distribution shifts: Long-context and multi-turn inputs are challenging distribution shifts for language models and, by extension, activation probes operating on their hidden states.The paper treats distribution-shift generalization as a longstanding machine-learning problem rather than a probe-specific anomaly.
  • Jailbreaking: Jailbreaking is an adversarial setting in which crafted inputs bypass safety training to elicit harmful outputs, and adaptive attacks can bypass most static guardrails.Benchmarks such as JailbreakBench and HarmBench support standardized evaluation of jailbreak mitigation.
  • Automation: The paper uses AlphaEvolve to automate searches for well-generalizing probe architectures and more effective adversarial prompts in automated red teaming.This extends language-model-guided optimization directly to AI safety-relevant problems.

7. Conclusion

The paper concludes that probes can complement black-box classifiers for cyber-misuse detection, but distribution-shift robustness remains difficult. It reports deployment-relevant benefits while identifying evaluation gaps and uncertainty in architecture selection.

  • Conclusion: Probes effectively complement black-box classifiers for cyber-misuse detection, with substantial remaining headroom from architecture variation and automated discovery.The conclusion also highlights distribution-shift robustness as a continuing challenge for both probes and LLM classifiers.
  • Limitations: The study does not evaluate probes using activations from every model layer, limiting direct comparison with concurrent all-layer approaches.The authors note that their findings and AlphaEvolve methodology could potentially be combined with that approach.
  • Limitations: Large error bars prevent recommending one single probe architecture, and future work should evaluate more domains beyond cyber misuse and academic datasets.The conclusion presents these as boundaries on exact recommendations and broader validation.
  • Future work: The authors aim to improve frontier-deployment defenses at a given cost and enable further activation-probe research using open-weights and open-data work.This closing statement points to future research built on the paper’s released resources.

9. Author Contributions

The author-contribution statement assigns responsibility for MultiMax findings, infrastructure, AlphaEvolve experiments, automated red teaming, autorater baselines, and the manuscript.

  • Contributions: JK made the first MultiMax generalization findings and built foundational infrastructure for the project.
  • Contributions: JE ran all AlphaEvolve experiments, wrote the first draft, edited the paper, and assisted with autorater baseline experiments.
  • Contributions: ZW led autorater baselines and ran long-context training experiments, while BC built the automated-red-teaming seed program and provided paper feedback.

A. Running our findings on other datasets and models

Across 12 additional classification datasets, AlphaEvolve and attention architectures generally outperform logistic regression, while seed selection contributes smaller gains and some comparisons are infrastructure-sensitive. The appendix also documents evaluation filtering and high variance in MultiMax architectures.

  • Architecture ranking: AlphaEvolve architectures and default-aggregation attention probes achieve the highest median AUROC, at 0.975, versus 0.944 for logistic regression.AlphaEvolve exceeds the logistic-regression baseline by an average of 3.1 percentage points in median AUROC across most datasets.
  • Seed variance: Seed selection provides modest 0.5–1.5 percentage-point improvements depending on architecture.Each architecture-dataset pair was evaluated with 100 random seeds, enabling median- and best-seed comparisons.
  • Linear baseline: The in-house linear probe reaches median AUROC 0.814, below the prior logistic-regression baseline of 0.944, likely because training settings were tuned for larger attention-based probes.The authors therefore use the prior logistic-regression result as the reference point for linear methods in visualizations.
  • Comparison to prior work: On overlapping datasets, attention probes score 0.98 versus 0.965 on news classification and 0.94 versus 0.86 on Commonsense Morality compared with prior work.The authors suggest the Commonsense Morality gap may reflect the attention probe’s MLP layer before attention.
  • Figures: Figures 6 and 7 encode median and best-seed AUROC across architecture rows and dataset columns, with green indicating higher AUROC and matched ordering.
  • Evaluation construction: Clean long-context evaluation retained 436 rows after excluding 1,711 rows matching short-context training stems.The filtering was applied to reduce leakage and test genuine generalization.
  • High-variance architectures: MultiMax architectures show substantially larger interquartile ranges in test loss across seeds than the architectures shown in Figure 3.They are presented separately to avoid distorting the main figure’s scale.

D.3. Seed Selection vs Architecture Choice

Architecture choice improves AUROC substantially more than seed selection, while seed sweeps still provide modest gains and help quantify variance. The section also develops a convex-frontier method for evaluating cascading probe–LLM policies.

  • Seed Selection: 0.011 AUROC is the average gain from selecting the best seed, ranging from +0.007 to +0.014 across architectures.
  • Architecture Choice: 0.109 AUROC separates the best architecture from Linear Probe, approximately 10× larger than the average seed-selection gain.Attention Probe with Default aggregation achieves 0.933 median AUROC versus 0.824 for Linear Probe.
  • Architecture Choice: AE1 and AE2 reach 0.930 median AUROC, with higher seed variance than Attention Probe with Default aggregation.Their seed-selection gains are +0.012 versus +0.007 for Attention Probe with Default aggregation.
  • Cascading Frontier: Threshold pairs produce discrete cost–test-error points, so the analysis uses linear interpolation through threshold randomization to form the continuous frontier.Linear interpolation is achievable and optimal within threshold-randomization policies.
  • Cascading Frontier: The cascading frontier combines independent left- and right-tail savings curves through a Minkowski sum, computed by merging edges in increasing slope order.The initial probe-logit sort dominates runtime at O(Nlog N).

E.1. Vertex Optimality: Why No Randomization is Needed At The Optimal Low Error Point

For minimizing error, the optimal cascading policy is deterministic because linear objectives over the achievable convex set attain minima at vertices. Randomization remains relevant for interior frontier points or targeted cost budgets.

  • Vertex Optimality: The minimum-error policy lies at a convex-hull vertex corresponding to one deterministic threshold pair.This follows because minimizing error is a linear objective over the convex set of achievable cost–error pairs.
  • Randomization: Interior frontier points and edges require randomization between policies, whereas vertices correspond to non-randomized threshold pairs.
  • Application: The “Selected Probe + 8% Flash” point minimizes error with approximately 3% of queries deferred to the LLM and no randomization.
  • Caveat: Randomization may be required for a specific cost budget and may be undesirable for misuse defense because attackers can exploit repeated queries.
  • AlphaEvolve: AlphaEvolve validation used a different error weighting from the final metric, and rerunning it with the final training mixture was not verified.

H. Automated Red Teaming Analysis

The automated red-teaming analysis compares adaptive attacks with repeated first-step attempts and evaluates prompt optimization and attack-validation safeguards. Adaptive turns did not improve over repetition, while AlphaEvolve improved single-turn success.

  • Adaptive Red Teaming: Repeating the first attack step outperformed the full adaptive procedure for every tested T>1.The comparison used 100 prompts and found adaptive turns did not outperform independent repetitions.
  • Prompt Optimization: 20% to 26% was the improvement in single-turn jailbreak success after AlphaEvolve optimized attacker prompts.
  • Adaptive Red Teaming: Cross-pollination produced a negligible margin while substantially increasing complexity and computational cost.
  • Attack Evaluation: Successful attacks required both bypassing the monitor and receiving an autorater judgment that the model response was harmful.
  • Attack Evaluation: The autorater compared rewritten-prompt responses with responses to the original harmful query, preserving the underlying malicious intent for evaluation.

J.3. Cascading Methods (Best Probe + Gemini 2.5 Flash)

The cascading classifier combines probe and LLM uncertainty while exploiting efficient attention-probe inference. Its uncertainty is typically dominated by probe seed-selection variance, and metric weighting materially affects reported error tradeoffs.

  • Cascading Uncertainty: Cascade error is approximated as (1 − f_defer) · Probe Error + f_defer · LLM Error.
  • Cascading Uncertainty: With deferral rates typically near 8%, cascade uncertainty is dominated by the probe component and inherits its asymmetry.
  • Uncertainty Interpretation: Probe confidence intervals measure variation from retraining and selecting among 100 seeds, while LLM intervals measure finite-test-set uncertainty.
  • Metric Sensitivity: Changing the weighting from 1:2:50 to the final weighting reduced false positives but produced extremely high false negative rates in earlier experiments.
  • Efficient Inference: Attention-probe inference scales as O(n) compute and O(1) memory for sequence length n when using a constant query vector.The probe maintains a scalar average and log-normalizer per head, discarding token activations after updates.
Loading 2601.11516v4…