LEMON: Equivariant Learning over Attention Maps for Zero-Shot Misbehavior Monitoring on New LLMs

Guy Bar-Shalom; Mor Geva; Fabrizio Frasca; Haggai Maron

cs.AIOct 15, 2026

LEMON: Equivariant Learning over Attention Maps for Zero-Shot Misbehavior Monitoring on New LLMs

Authors: Guy Bar-Shalom, Mor Geva, Fabrizio Frasca, Haggai Maron

Organizations: Technion · Tel Aviv University · Technion, NVIDIA

Abstract

Large language models exhibit misbehaviors that hinder deployment in high-stakes settings. One approach to monitor such failures is to use an external LLM as a monitor, but this is prohibitively expensive. A well-known and efficient class of monitors can detect such failures from internal representations (such as activation probes), but are tied to the model they were trained on, requiring costly re-annotation as new LLMs appear. A natural way to avoid this repeated annotation cost is to develop internal-representation based monitors that transfer zero-shot across LLMs. However, such transfer is challenging because internal representations are not naturally aligned across models; in particular, attention maps, known to contain signals of misbehavior, may differ in head function, ordering, and model depth. In this work, we characterize those cross-model attention mismatches as algebraic symmetries and use geometric deep learning principles to build LEMON (Language-model Equivariant MOnitoring Network) an online graph neural network-based monitor, whose predictions are invariant under these symmetries. LEMON processes attention maps token-by-token, as they are produced and outputs a score after every token (or every window of $W$ tokens) at negligible cost. Across 16 LLMs and six datasets, in the zero-shot cross-model setting, LEMON outperforms all baselines on both hallucination and jailbreak monitoring tasks, improving over the strongest by up to 14 ROC-AUC points.

Figures & tables

Panel of Figure 1 — Figure 1: Our framework for zero-shot monitoring of misbehavior across LLMs. (left) The attention map $\mathbf{A}\in[0,1]^{N\times N\times L\times H}$ is a stack of attention maps indexed by (layer, head). (middle) LEMON operates on a structured, attributed graph constructed from the attention maps and applies $\mathcal{W}$ -equivariant message passing. (right) LEMON can flag misbehavior during generation.

Figure 2: Motivating example: Head and layer roles vary across LLMs. The same hallucinated token produces distinct $32\times 32$ head patterns, with matched $(\ell,h)$ values uncorrelated ( $r=0.04$ ).

Panel of Figure 3 — Figure 3: MHA symmetries.

Panel of Figure 4 — Figure 4: Action of $\mathcal{W}=\mathcal{S}_{H}\wr\mathcal{C}_{L}$ on $\mathbf{A}$ for $L{=}H{=}3$ : rows are reordered by $\pi$ ; heads within each row by $\sigma_{\ell}$ .

Panel of Figure 5 — Figure 5: $\mathbf{A}$ symmetries.

Panel of Figure 6 — Figure 6: LEMON is online.

Table 7

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Day	Llama-2-7B	Mistral-7B
Monday	Monday	Monday
Tuesday	T\|ues\|day	Tuesday
Wednesday	Wed\|nes\|day	Wednesday
Thursday	Th\|urs\|day	Thursday
Friday	Friday	Friday
Saturday	Saturday	Saturday

Appendix

Table 5: Subword tokenization of the seven day names for Llama-2-7B and Mistral-7B .

Panel of Figure 9 — Figure 9: The $\mathcal{W}$ -equivariant layer $f_{\mathcal{W}}(\mathbf{A}_{n_{1},n_{2}})$ . The Siamese path (dashed, top) applies $f_{\mathcal{S}_{H}}$ to each layer slice independently; the information-sharing path (bottom) sums over heads, applies $f_{\mathcal{C}_{L}}$ across layers, and broadcasts the result back over heads. The two paths are summed.

Panel of Figure 10 — Figure 10: The LEMON architecture. Top : the end-to-end pipeline. Bottom, left to right : one equivariant MPNN layer—parallel supernode/response streams with $\mathcal{W}$ -equivariant updates ( Equations 6 and 7 ); the shared causal cross-token Transformer, applied independently at each layer–head cell $(l,h)$ ; and the invariant prediction head for the per-token case ( Equation 8 ).

Panel of Figure 11 — Figure 11: Preprocessing: p ${\to}$ p edges (red) discarded; weak edges (gray, dashed) thresholded away; prompt nodes (boxed) pooled into supernode $s$ ; response tokens (boxed) pooled into windows ( $W{=}2$ shown). Right : adjacency of the final graph with super and response stream blocks (edge colors match).

Panel of Figure 12 — Figure 12: Unequal $H^{s},L^{s}\!\neq\!H^{t},L^{t}$ . Layer map $\phi$ pairs source and target layers; only $H^{\min}\!=\!\min(H^{s},H^{t})$ heads are matched.

Model	Task	Spans	Spans/resp.	Empty	Exact-offset
Qwen2.5-3B	Summ.	$15$	$1.5$	$5$	$15/15$
Qwen2.5-3B	QA	$26$	$2.6$	$0$	$26/26$
Qwen2.5-3B	D2T	$35$	$3.5$	$0$	$35/35$
Qwen2.5-7B	Summ.	$11$	$1.1$	$6$	$11/11$
Qwen2.5-7B	QA	$24$	$2.4$	$1$	$24/24$
Qwen2.5-7B	D2T	$31$	$3.1$	$0$	$31/31$

Appendix

Table 6: Audit statistics over the first ten test examples per (model, task),

90

examples in total. “Spans” is the total number of hallucination spans the judge flagged across the ten responses, and “Spans/resp.” the average per response; “Empty” counts responses with no flagged span (out of

10

); “Exact-offset” counts spans whose recorded text matches the response substring at the recorded offsets.

Hyperparameter	Hallucination	Jailbreak
Hidden dimension $d$	$16$
MPNN layers $K$	$\{1,2\}$
Cross-token Transformer layers	$\{1,2\}$
Attention heads (cross-token Transformer)	$2$
Dropout	$0$
Attention threshold $\delta$	$0.1$

Appendix

Table 7: LEMON hyperparameters. Values in braces are searched; the configuration is selected by validation ROC-AUC on the training LLMs only. Both LEMON variants share all settings.

Training pool ( $\mathcal{K}=6$ ) — fitted on; held out one at a time
Model	$H$	$L$
Llama-2-7B	32	32
Llama-2-13B	40	40
Mistral-7B	32	32
Mistral-Nemo-12B	32	40
Qwen2.5-3B	16	36

Appendix

Table 8: Every LLM we monitor, with its attention grid

(H,L)

LLM	All $3000$ prompts		Test split ( $750$ )
LLM	jailbroken	ASR (%)	jailbroken	ASR (%)
Llama-2-7B	$256/3000$	$8.5$	$69/750$	$9.2$
Llama-2-13B	$234/3000$	$7.8$	$53/750$	$7.1$
Mistral-7B	$1325/3000$	$44.2$	$337/750$	$44.9$
Mistral-Nemo-12B	$1548/3000$	$51.6$	$394/750$	$52.5$
Qwen2.5-3B	$762/3000$	$25.4$	$212/750$	$28.3$

Appendix

Table 9: Attack-success rate (ASR) % of the universal GCG suffix, per LLM.

Held-out LLM	Task	Prob	LB	HS	CHARM	LB+Align	LEMON (Ours)
Held-out LLM	Task	Prob	LB	HS	CHARM	LB+Align	$\mathcal{W}$	$\mathcal{W}$ -per-cell
Llama-2-7B	SUM	56.2	48.7	50.7	41.3	61.7	68.4	71.4
	NQ	60.7	55.8	52.1	51.6	61.5	82.7	84.5
	D2T	56.6	41.8	55.9	43.5	60.6	63.2	68.9
	avg	57.8	48.8	52.9	45.5	61.3	71.4	74.9
Llama-2-13B	SUM	55.8	55.3	47.5	59.5	49.5	68.3	70.3

Appendix

Table 10: LOO zero-shot generalization — per-task breakdown of Table 2 . Each data row is one (left-out LLM, task) pair; every value is the ROC-AUC

{\times}\,100

on that held-out shard; avg rows average the three tasks for that LLM. Bold = best per row; underline = second best.

LLM	Task	LEMON- $\bm{\mathcal{W}}$	LEMON- $\bm{\mathcal{W}}$ -per-cell	Prob
Llama-2-7B	Summ.	3.00	2.83	7.38
	NQ	3.37	2.99	9.23
	D2T	5.97	6.00	9.67
	mean	4.11	3.94	8.76
Llama-2-13B	Summ.	1.62	1.71	2.97
	NQ	2.04	1.72	7.08

Appendix

Table 11: Selective generation. A per-shard breakdown of the area under the risk–coverage curve (lower is better). Bold = best in row.

Panel of Figure 15 — Figure 15: Invariant and per-LLM heads (one linear predictor per pool LLM). $\mathbf{H}^{\mathrm{out}}$ is the per-node features before the prediction heads. Dashed LR : the adaptation read-out, fit on the frozen $\bm{\phi}_{n}$ after training, so it does not feed $\mathcal{L}$ .

Figure 16: Adaptation to the held-out LLM vs. labeling budget , both LEMON variants; mean over the 18 (LLM, task) shards. Solid: logistic regression on the frozen non-invariant head’s features (

\alpha{=}0.5

) and LB ; dashed: the corresponding invariant-head (zero-shot) levels. The signed number at each point is that variant’s gap to LB .

Figure 17: Invariant head vs. non-invariant head

+

LR as a function of

\alpha

, both LEMON variants. We report the mean over the 18 (LLM, task) shards. Raising

\alpha

trains the per-LLM non-invariant pathway and improves adaptation, while at

\alpha{=}1

the invariant head collapses towards chance as expected.

Left-in LLM	Task	Prob	LB	HS	CHARM	LB+Align	LEMON (Ours)
Left-in LLM	Task	Prob	LB	HS	CHARM	LB+Align	$\mathcal{W}$	$\mathcal{W}$ -per-cell
Llama-2-7B	SUM	62.3	53.4	48.3	60.1	55.8	55.7	62.1
	NQ	57.5	56.1	52.1	59.2	59.9	62.7	66.1
	D2T	60.7	55.7	50.0	54.4	49.2	54.5	64.7
	avg	60.2	55.1	50.1	57.9	55.0	57.6	64.3
Llama-2-13B	SUM	62.4	49.7	53.4	50.6	49.5	61.8	63.5

Appendix

Table 12: LOI zero-shot generalization — per-task breakdown. Each monitor is trained on a single (left-in) LLM and evaluated zero-shot on the other five. Each data row is one (left-in LLM, task) pair; every value is the mean ROC-AUC

{\times}\,100

over the five held-out evaluation LLMs; avg rows average the three tasks for that LLM. Bold = best per row; underline = second best.

Panel of Figure 18 — Figure 18: The selective generator $(f,g)$ : the LLM $f$ samples one token at a time, while the selection function $g$ decides whether to serve the newest token and continue generation, or to withhold it and halt.

Panel of Figure 19 — Figure 19: The halting step $\rho$ of equation 16 and the served output $\hat{\mathbf{x}}$ of equation 17 : $g$ abstains for the first time at “Mars”, so $\rho=4$ . The user sees only the tokens delivered before the flag (blue); the flagged token itself is withheld (orange), and the suffix beyond it is never generated (dashed).

Panel of Figure 21 — Figure 21: The three selection functions on two different generations (one per row); the time step $t$ of each token is shown above it.

Panel of Figure 22 — Figure 22: Coverage equation 24 , token coverage equation 22 , selective risk equation 23 , and token savings equation 25 on an evaluation set of $N=2$ generations. Beneath each token, as in Figure 21 : the per-step decisions $g$ , the per-token acceptance $\tilde{g}$ (latching to $0$ from the first rejection), and the per-prefix losses $\ell$ ( red underline = hallucinated).

LLM Safety Monitoring

See all 56 papers

Benchmarking and Improving Monitors for Out-Of-Distribution Alignment Failure in LLMs
May 20, 2026Dylan Feng, Pragya Srivastava, Anca Dragan +1Language Model Safety Evaluation LLM Guardrails
University of California, Berkeley, USA · 3Google DeepMind, India.
Mechanism-Aligned Observability for Training Instability of Large Language Models
Jun 26, 2026Ruixuan Huang, Yipei Wang, Wenyi Fang +5Fault Detection LLM Training
HKUST · 2Huawei · Independent Researcher
Relative Representations Enable Portable White-Box Monitoring Across Language Models
Oct 20, 2026Tim Bakker, Christos LouizosLanguage Model Probing Representation Geometry in Language Models
Qualcomm AI Research

May 20, 2026J/K move · Enter open · S save

Benchmarking and Improving Monitors for Out-Of-Distribution Alignment Failure in LLMs

Dylan Feng, Pragya Srivastava, Anca Dragan, Cassidy Laidlaw

University of California, Berkeley, USA · 3Google DeepMind, India.

Language Model Safety Evaluation LLM Guardrails LLM Safety Monitoring LLM Safety OOD Detection

Open paper AskPDF arXiv

Many safety and alignment failures of large language models (LLMs) occur due to out-of-distribution (OOD) situations: unusual prompt or response patterns that are unforeseen by model developers. We systematically study whether LLM monitoring pipelines can detect these OOD alignment failures by introducing a benchmark called Misalignment Out Of Distribution (MOOD). It is difficult to find failures that are truly OOD for off-the-shelf models trained on vast safety datasets. We sidestep this by including a restricted training set in MOOD that we use to train our own monitors, as well as seven test sets with diverse alignment failures that are outside the training distribution. Using MOOD, we find that guard models (safety classifiers) often fail to generalize OOD. To fix this, we propose combining guard models with OOD detectors. We test four types of OOD detectors and find that a combination of a guard model with Mahalanobis distance and perplexity-based OOD detectors can improve recall from 39% to 45%. We also establish positive scaling trends across model scales for monitors that combine a guard model and OOD detector; we find that incorporating OOD detection into monitoring achieves a higher recall gain than using a guard model with 20 times more parameters. Our work suggests that OOD detection should be a crucial component of LLM monitoring and provides a foundation for further work on this important problem. We release the code and data for our experiments publicly, and you can find the relevant links here: https://github.com/Dylan102938/mood-bench.

Abstract

LLM Safety Monitoring

Explore similar work

Benchmarking and Improving Monitors for Out-Of-Distribution Alignment Failure in LLMs

Mechanism-Aligned Observability for Training Instability of Large Language Models

Relative Representations Enable Portable White-Box Monitoring Across Language Models