LEMON: Equivariant Learning over Attention Maps

TL;DR — Attention maps carry signals of LLM misbehavior, but heads and layers are not aligned across models. We treat these mismatches as symmetries and build LEMON, a monitor that is invariant to them, so LEMON transfers zero-shot to LLMs it has never seen, scoring every token as it is generated.

Motivation and problem

To catch LLM misbehavior (hallucinations, jailbreaks) as it happens, a monitor must be online, scoring the text while it is generated, and lightweight, adding little to the cost of generation. Monitors that read the LLM's own internals (in particular, attention maps, as considered in this paper) typically meet both requirements, but each is trained for one specific LLM, and every newly released LLM needs new annotations in order to train a new monitor.

Ideally, we want a single monitor that transfers zero-shot to unseen LLMs. That is hard: attention maps carry misbehavior signals, but heads with the same role sit at different positions in different models:

Per-head Lookback Lens attention ratios for the same hallucinated token in Llama-2-7B and Mistral-7B: two 32x32 head-by-layer heatmaps with different patterns, and a scatter plot of matched head/layer values showing Pearson r = 0.04. — **Head and layer roles differ across LLMs.** Each cell shows the attention-based *lookback ratio* introduced by Lookback Lens (Chuang et al., 2024): mean attention to the input context divided by the sum of mean attention to the context and previously generated tokens. For the same hallucinated token, Llama-2-7B and Mistral-7B produce very different ratios at matching layers ℓ and heads h; the values are essentially uncorrelated (r = 0.04).

Cross-LLM mismatches are symmetries

Heads. A multi-head attention layer concatenates its head outputs and applies an output projection $\mathbf{W}^{O}$. Permuting the head outputs and the corresponding row blocks of $\mathbf{W}^{O}$ leaves the layer's output unchanged:

$$\operatorname{MHA}(\mathbf{X}) = \big[\mathbf{A}^{1}\mathbf{V}^{1} \,\|\, \cdots \,\|\, \mathbf{A}^{H}\mathbf{V}^{H}\big]\,\mathbf{W}^{O} = \big[\mathbf{A}^{\textcolor{#ad0000}{\sigma}^{-1}(1)}\mathbf{V}^{\textcolor{#ad0000}{\sigma}^{-1}(1)} \,\|\, \cdots \,\|\, \mathbf{A}^{\textcolor{#ad0000}{\sigma}^{-1}(H)}\mathbf{V}^{\textcolor{#ad0000}{\sigma}^{-1}(H)}\big]\,\textcolor{#ad0000}{\mathbf{P}_\sigma}\mathbf{W}^{O}.$$

Here $\textcolor{#ad0000}{\mathbf{P}_\sigma}$ is the block-permutation matrix corresponding to the head permutation $\textcolor{#ad0000}{\sigma} \in \mathcal{S}_H$. So two LLMs can implement the same head in different slots; see the illustration below of this symmetry for a single layer.

Animation showing three attention heads being reordered together with the matching row blocks of the output projection W^O, while the layer's output stays unchanged. — Swapping heads together with the matching blocks of $\mathbf{W}^{O}$ leaves the layer unchanged.

Layers. Layers are composed sequentially, so their order matters. Their absolute indices, however, carry no fixed meaning across models of different depths. We model layer relabelings as cyclic shifts $\pi \in \mathcal{C}_L$, treating them as an approximate symmetry assumption that reduces dependence on absolute layer indices.

Together, an independent head permutation in every layer plus a layer shift form the group $\mathcal{W} = \mathcal{S}_H \wr \mathcal{C}_L$. The relabeling is $(\ell,h) \mapsto (\pi(\ell),\sigma_\ell(h))$, so the group action is as follows:

$$(\sigma_\ell,\pi) \in \mathcal{S}_H \wr \mathcal{C}_L, \qquad (\sigma_\ell,\pi) \cdot \mathbf{A}^{\ell,h}_{n_1,n_2} = \mathbf{A}^{\pi^{-1}(\ell),\; \sigma_{\pi^{-1}(\ell)}^{-1}(h)}_{n_1,n_2}.$$

LEMON is designed to be invariant to the symmetries formalized above.

A 3x3 grid of layer-by-head attention maps; layers are cyclically shifted and heads are permuted independently within each layer. — The joint action for $L{=}H{=}3$: layers are shifted by $\pi$, and heads are permuted within each layer by $\sigma_\ell$.

Our LEMON architecture

Our architecture, LEMON, is based on a Message Passing Neural Network (MPNN), modifying the update function $f_{\mathrm{up}}:\mathbb{R}^{L \times H \times d}\to\mathbb{R}^{L \times H \times d}$ to be equivariant to the group $\mathcal{W}$, where $d$ is the hidden feature dimension.

Animation showing messages passed between token nodes, aggregated, and processed by an update function. Each node and edge carries a layer-by-head feature array. — Message passing over tokens, with $L \times H$ features on every node and edge.

Making $f_{\mathrm{up}}$ equivariant to $\mathcal{W}$. Building on the characterization of linear $\mathcal{W}$-equivariant maps by Wang, Albooyeh, and Ravanbakhsh (NeurIPS 2020), we extend their construction to nonlinear neural networks, yielding a more general class of $\mathcal{W}$-equivariant maps. We define $f_{\mathcal{W}}:\mathbb{R}^{L \times H}\to\mathbb{R}^{L \times H}$ as follows, omitting feature channels for clarity:

$$f_{\mathcal{W}}(\mathbf{A}_{n_1,n_2}) = f_{\mathcal{C}_L}(\mathbf{A}_{n_1,n_2}\mathbf{1}_H)\mathbf{1}_H^{\top} + f_{\mathcal{S}_H}^{\mathrm{siamese}}(\mathbf{A}_{n_1,n_2}),$$ $$\big[f_{\mathcal{S}_H}^{\mathrm{siamese}}(\mathbf{A}_{n_1,n_2})\big]^{\ell} := f_{\mathcal{S}_H}(\mathbf{A}_{n_1,n_2}^{\ell}), \qquad \forall\,\ell \in [L].$$

Here $\mathbf{1}_H \in \mathbb{R}^{H}$ is the all-ones vector. The functions $f_{\mathcal{C}_L}:\mathbb{R}^{L}\to\mathbb{R}^{L}$ and $f_{\mathcal{S}_H}:\mathbb{R}^{H}\to\mathbb{R}^{H}$ can be any functions equivariant to $\mathcal{C}_L$ and $\mathcal{S}_H$, respectively, with the same $f_{\mathcal{S}_H}$ shared across all layers. Adding the two paths gives a $\mathcal{W}$-equivariant map; see the illustration of $f_{\mathcal{W}}$ below.

A wreath-product equivariant map with two paths: a shared head-equivariant function applied to each layer, and a path that sums over heads, applies a layer-equivariant function, broadcasts back over heads, and adds the two outputs.

LEMON is online. Each generated token contributes a new attention row, which adds a node and incoming edges to the graph. Forward-only message passing, causal cross-token attention, and a prediction head that scores tokens independently ensure that $\hat{y}_n$ depends only on tokens $m \le n$. LEMON can therefore score tokens as they arrive, with the same predictions as within the completed response. Appending each node and its edges costs linear time in the prefix length.

Each generated token appends an attention row and a graph node with edges from earlier tokens m to the new token n. LEMON scores the new token immediately while retaining earlier nodes, edges, and predictions. — Schematic for one attention layer and head: each new attention row adds a node and incoming edges; the new token is scored immediately.

Results

Setup. We use a Leave-One-Out (LOO) evaluation protocol — we train on five LLMs and test zero-shot on the sixth. The held-out LLM contributes no labels to training or model selection. We repeat this for all six LLMs, training separate monitors for hallucination and jailbreak monitoring; see illustration below.

Leave-one-LLM-out evaluation rotates through six splits. Each LLM keeps a distinct color as it moves between the five-model training set and the held-out evaluation card.

We consider two tasks: hallucination monitoring and jailbreak monitoring. For hallucination monitoring, we use the RAGTruth dataset and predict hallucinations at the token level as tokens are generated. For jailbreak monitoring, we predict attack success from the prompt alone on RedBench, using a universal GCG suffix across requests and LLMs.

ROC-AUC × 100 ↑ Bold: best · Underline: second best Swipe each table to compare all methods →

Table 1 · RAGTruth Hallucination monitoring Mean over three RAGTruth tasks
Unseen LLM	Baselines					LEMON (Ours)
Unseen LLM	Prob	LB	HS	CHARM	LB+ Align	$\mathcal{W}$	$\mathcal{W}$per-cell
Llama-2-7B	57.8	48.8	52.9	45.5	61.3	71.4	74.9
Llama-2-13B	52.3	53.9	52.5	57.3	54.2	73.8	75.5
Mistral-7B	67.6	49.5	51.0	42.2	57.5	68.2	68.5
Mistral-Nemo	58.3	52.4	50.3	44.2	60.6	63.9	64.8
Qwen2.5-3B	61.6	54.3	49.3	41.9	50.1	66.0	66.2
Qwen2.5-7B	61.1	50.0	52.2	57.8	58.7	63.6	66.3
Mean (18 shards)	59.8	51.5	51.4	48.1	57.1	67.8	69.4

Each row averages three RAGTruth tasks: summarization, question answering, and data-to-text. Mean covers all 18 LLM/task shards.

Table 3 · RedBench Jailbreak monitoring Prompt-only prediction before generation
Unseen LLM	Baselines			LEMON (Ours)
Unseen LLM	Attn→ Sys	CHARM	HS	$\mathcal{W}$	$\mathcal{W}$per-cell
Llama-2-7B	59.1	30.4	73.1	75.1	62.6
Llama-2-13B	47.8	73.7	26.4	69.5	66.5
Mistral-7B	61.5	53.9	50.9	67.5	59.2
Mistral-Nemo	47.3	58.3	48.7	55.0	57.3
Qwen2.5-3B	67.6	38.5	45.7	63.9	64.8
Qwen2.5-7B	67.8	46.1	38.6	59.3	69.1
Mean (6 LLMs)	58.5	50.1	47.2	65.0	63.3

Predict attack success on each unseen LLM; mean covers all six LLMs.

LEMON achieves the highest mean ROC-AUC on both tasks, substantially outperforming baselines that do not explicitly account for layer–head symmetries. These results highlight the value of symmetry-aware architectures for transferring monitors to unseen LLMs.

BibTeX

@article{barshalom2026lemon,
  title   = {LEMON: Equivariant Learning over Attention Maps for Zero-Shot Misbehavior Monitoring on New LLMs},
  author  = {Bar-Shalom, Guy and Geva, Mor and Frasca, Fabrizio and Maron, Haggai},
  journal = {arXiv preprint},
  year    = {2026}
}