Table of Contents
Fetching ...

Building Production-Ready Probes For Gemini

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

TL;DR

This work finds 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.

Abstract

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 architecture that handle this long-context distribution shift. We evaluate these probes in the cyber-offensive domain, testing their robustness against various production-relevant shifts, including multi-turn conversations, static jailbreaks, and adaptive red teaming. Our results demonstrate that while multimax addresses 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.

Building Production-Ready Probes For Gemini

TL;DR

This work finds 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.

Abstract

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 architecture that handle this long-context distribution shift. We evaluate these probes in the cyber-offensive domain, testing their robustness against various production-relevant shifts, including multi-turn conversations, static jailbreaks, and adaptive red teaming. Our results demonstrate that while multimax addresses 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.
Paper Structure (73 sections, 21 equations, 12 figures, 8 tables, 1 algorithm)

This paper contains 73 sections, 21 equations, 12 figures, 8 tables, 1 algorithm.

Figures (12)

  • Figure 1: We improve over prior activation probing work by improving probe architectures and training to achieve better performance than language models at over 10,000$\times$ lower cost. "Selected Probe" refers to our Max of Rolling Means Attention Probe (\ref{['secOurProbes']}). Deferring to an LLM such as Gemini 2.5 Flash 8% of the time further decreases error rate to a lower value than what can be achieved by solely using a probe or an LLM, by trusting the probe at extreme scores and delegating to the LLM in between (see \ref{['secOptimalFrontier']}, though note wide error bars). *As described in \ref{['subsecMetric']}, for each method we select the decision threshold on a validation set by minimizing a weighted combination of FPR and FNR (i.e. error rate), with overtriggering FPR weighted most heavily. We report the same weighted error computed on test data. See \ref{['appErrorBars']} for error bar methodology. The two points not on the Pareto frontier (Selected Probe and AlphaEvolve) have their costs artificially increased by 1.25$\times$ for visual separation; the costs of these three probes are in practice almost identical.
  • Figure 2: Different probing classifiers that we compare in \ref{['secMethods']}. All our probing classifiers can be composed into six states 1) -- 6) as illustrated. Residual stream activations of undergo a 2) Transformation per-position. These are then processed via 4) -- 5) Aggregation which ends up producing a single scalar score at step 6). Many existing probing classifiers such as linear probes, exponential moving average aggregation probes and attention probes fit into this framework.
  • Figure 3: Effect of seed selection on test loss. Arrows show improvement from median (gray) to best-validation-selected (coral star) test loss. Green diamonds show the oracle (best possible seed). Some MultiMax aggregations have very large IQR and are shown in \ref{['figAppendixSeedSelection']} to not distort this figure.
  • Figure 4: Comparison of probe architectures across distribution shifts. (a) Test error across methods. (b) Probes achieve comparable performance to LLM classifiers at a fraction of the cost. (c) Long context generalization: AlphaEvolve probes maintain low FPR on long contexts. (d) Adversarial attacks remain challenging for all methods; all methods still leave attack success rate >1%. Note: error bars are omitted from these plots see \ref{['figMainFigScatterCost', 'figCascadingClassifiers']} for visualizations with (relatively wide) uncertainty estimates. *As in \ref{['figMainFigScatterCost']} the threshold selection methodology is described in \ref{['subsecMetric']}.
  • Figure 5: Cost plotted against test error (\ref{['eqnTradeoff']}). The curves show the optimal frontier for cascading classifiers that combine cheap probes with expensive LLMs. Error bars (see \ref{['appErrorBars']}) indicate substantial uncertainty: most operating points have overlapping confidence intervals, so the apparent ordering between methods should be interpreted cautiously. The cascade global minima (circles) represent optimal operating points for each LLM; further deferring to LLMs past these points increases the error rate.
  • ...and 7 more figures