---
title: Sampling-Equivariant Self-Attention
url: https://www.emergentmind.com/topics/sampling-equivariant-self-attention
type: topic
---

# Sampling-Equivariant Self-Attention

Sampling-equivariant self-attention denotes a class of attention mechanisms in which changes to the sampling of the underlying signal—such as token permutations, rotations or reflections of a sampling grid, patch extraction, down-sampling, up-sampling, or learned token subsampling—either leave the attention operator invariant or induce a prescribed transformation of its outputs. In the literature, the term is explicit in aerial-image detection, where local self-attention is recast as convolutional sampling with masks rather than locations [2111.03420], and it is also realized by group-equivariant transformers in vision, 3D geometry, Lie-group domains, and agent modeling, where queries, keys, values, positional structure, and resampling operators are constrained to commute with a symmetry action [2602.08047][2006.10503][2012.10885][2604.01466].

## 1. Formal notion and scope

At the most general level, equivariance is the requirement that transforming the input by a group element \(g\) transforms the output by the corresponding output action. One standard statement is
\[
S_g[\phi(v)] = \phi(T_g[v]),
\]
and, in an image setting with feature-space action \(\pi_{\tilde g}^F\), an equivariant mapping \(\Psi\) satisfies
\[
\Psi\big[\pi_{\tilde g}^I\big](I) = \pi_{\tilde g}^F\big[\Psi\big](I), \quad \forall \tilde g.
\]
These formulations underlie SE(3)-equivariant attention, discrete planar-equivariant ViTs, and Lie-group attention [2006.10503][2602.08047][2012.10885].

The qualifier “sampling-equivariant” emphasizes that the transformation does not merely change feature values; it changes the way the data are sampled. In aerial detection, this is formulated by comparing the sampling graph produced on a transformed image with the ideally transformed sampling graph, and the discrepancy is measured by Earth Mover’s Distance and aggregated as AEMD; lower AEMD means better sampling equivariance [2111.03420]. In group-equivariant ViTs, the same idea appears when rotations or reflections alter patch locations, group channels, and down/up-sampling lattices, but the network output still co-transforms exactly [2602.08047].

A recurring distinction is between exact equivariance and equivariance in expectation. Exact equivariance is established for several constructions on discrete planar groups, SE(3), SE(2), and general Lie groups [2602.08047][2006.10503][2604.01466][2012.10885]. By contrast, stochastic sampling schemes for efficient attention on sets or sequences typically preserve the target symmetry in distribution or in expectation over the sampling randomness [2302.14346][2111.09714].

## 2. Structural mechanisms that make attention sampling-equivariant

A common mechanism is to ensure that attention logits are invariant scalars while value messages are equivariant objects. In the SE(3)-Transformer, queries \(q_i\) and keys \(k_{ij}\) are equivariant embeddings of node and edge features, and the inner product \(q_i^\top k_{ij}\) is invariant because both transform under the same orthogonal representation. Attention weights are therefore invariant scalars, and the weighted sum of equivariant value messages remains SE(3)-equivariant [2006.10503]. LieTransformer expresses the same principle on a general Lie group by making the location term depend only on the relative group element \(g^{-1}g'\); under global left multiplication, \(g^{-1}g'\) is unchanged, so the attention kernel reindexes compatibly with the group action [2012.10885].

A second mechanism is to lift tokens into a representation space that explicitly carries group structure. In the discrete planar ViT construction, the token tensor is
\[
\bm{z} \in \mathbb{R}^{N \times c \times t},
\]
where \(t=|S|\) indexes group elements. Queries, keys, and values are not formed by unconstrained dense maps; instead, each slice is computed with group-tied weights, for example
\[
\bm{q}^B = \sum_{g\in S} \bm{z}_0^g \cdot \bm{W}^{q}_{B^{-1}g}, \quad \forall B\in S.
\]
The resulting block-structured \(W^q,W^k,W^v\) commute with the induced permutation of the group index, so the attention layer becomes exactly equivariant once positional encodings and resampling operators are made compatible [2602.08047].

A third mechanism is to replace location-based sampling by mask-based sampling. In SESAN, local self-attention is interpreted as convolutional sampling with masks rather than offsets. The mask is computed as
\[
w = \mathrm{softmax}\big(\gamma(\delta(Q) - \delta(K))\big),
\]
positional encoding is removed from the mask-regression path, and a transformation embedding module injects the mask \(w\) back into the output feature \(Y=\zeta(V\odot w,w)\). This design is meant to make the local sampling mask transform with the object under rotation, reflection, skew, or scaling [2111.03420].

A fourth mechanism is to encode geometry in an invariant algebraic representation rather than in explicit pairwise relative positional embeddings. DriveGATr represents scene elements as multivectors in \(\mathbb{R}^*_{2,0,1}\), computes logits from an SE(2)-invariant inner product and distance-aware invariant terms, and then applies standard scaled dot-product attention over the resulting embeddings. Because the logits are invariant and the values are equivariant multivectors, the attention output is SE(2)-equivariant without explicit pairwise relative positional encodings [2604.01466].

| Mechanism | Sampling object | Equivariance device |
|---|---|---|
| Group-indexed token lifting | Patch/grid tokens with group channel | Group-tied \(Q/K/V\), invariant or equivariant PE |
| Relative-group attention | Points or lifted group elements | Kernels on \(g^{-1}g'\) |
| Mask-based local attention | Local image patch | Content-derived mask, no absolute PE in mask path |
| Algebraic geometric attention | Agents, map nodes | Invariant inner products over equivariant multivectors |

These constructions differ in representation and domain, but they share one invariant pattern: the sampling geometry is not external bookkeeping. It is part of the representation on which attention is defined.

## 3. Sampling operations beyond the attention kernel

Sampling-equivariant attention is typically not achieved by modifying logits alone. In the discrete planar ViT framework, patch embedding is replaced by a group-equivariant convolution with stride,
\[
\hat{\bm{z}}^g = D_s\left(\pi_g(\psi) \otimes \bm{x}\right),\quad \forall g\in S,
\]
so that rotating or reflecting the input induces a spatial permutation plus a cyclic shift in the group dimension. Down-sampling is defined slice-wise by
\[
\tilde F_g = \big[\pi(g)\circ \mathrm{Down}_s\circ \pi(g^{-1})\big](F_g),
\]
and equivariant Pixel-Shuffle uses the analogous
\[
\tilde F_g = [\pi(g)\circ \mathrm{PS}_r \circ \pi(g^{-1})](F_g).
\]
These operators are intended to make patch extraction, multiscale reduction, and reconstruction commute with the group action rather than silently break it [2602.08047].

In point-cloud and graph attention, the relevant sampling issue is variable cardinality and density rather than a square lattice. The SE(3)-Transformer is a set/graph network: attention is defined per node and per neighbourhood, no assumption is made about the number of nodes, and permutation equivariance ensures that changing the order of points or the particular sampling order does not change the functional mapping beyond the same permutation [2006.10503]. LieTransformer addresses an even broader setting by lifting samples on a homogeneous space \(G/H\) to functions on \(G\), then applying attention over the lifted support; a rigid transformation of the sampled input left-multiplies the lifted support, so the attention layer co-transforms with the sampling geometry [2012.10885].

A different use of the term appears in self-learning Monte Carlo. There, the Transformer produces an equivariant effective spin field \(\mathbf S_{\rm eff}\), then an invariant effective Hamiltonian \(H_{\rm eff}\), and the proposal transition kernel satisfies
\[
P_{\rm prop}(g\mathbf S'|g\mathbf S) = P_{\rm prop}(\mathbf S'|\mathbf S).
\]
The paper explicitly interprets this as sampling-equivariance of the proposal chain: transforming the whole trajectory by a symmetry yields another valid trajectory with the same statistical properties [2306.11527].

A related, non-attentional development replaces a fixed group discretization by an input-dependent sampling matrix \(A(x)\) satisfying
\[
A(g.x) = A(x)\rho(g)^\top.
\]
This yields exact equivariance for any number of group samples \(N\) in a steerable-network nonlinearity [2409.08741]. A plausible implication is that future group-attentional layers can decouple equivariance from sampling density by generating their group samples adaptively rather than fixing them a priori.

## 4. Stochastic and subquadratic variants

A separate branch of the literature studies sampling as an efficiency device. The sampled transformer for point sets randomly splits the set into subsets, applies Hamiltonian self-attention within each subset, and interprets the global sparse pattern as a random Hamiltonian cycle in the complete attention graph. A fixed realization is not strictly permutation-equivariant, but the distribution over sampled attention graphs is invariant under permutations, and permutation equivariance holds in expectation over the sampling [2302.14346]. The same work proves that sampled transformers remain universal approximators of continuous set-to-set functions.

YOSO rewrites self-attention as a sum of value vectors weighted by Bernoulli random variables derived from LSH collisions. For a single head,
\[
\text{YOSO}(Q,K,V)=\mathcal{B}(Q,K)V,
\]
with
\[
\mathbb{E}[\mathcal{B}(Q,K)_{i,j}] = \left(1 - \frac{\arccos(Q_i K_j^\top)}{\pi}\right)^{\tau}.
\]
The resulting operator is linear in \(V\), its expectation is a deterministic collision-based attention map, and the implementation has linear complexity in sequence length for fixed hash parameters [2111.09714]. This is not exact group-equivariant attention, but it is an explicit instance of attention defined through stochastic sampling while preserving the attention form in expectation.

Sampling Foundational Transformer inserts a learned sampling-without-replacement step before attention on keys and values, while retaining all queries. It uses a pairwise Maxout score,
\[
\operatorname{Score}(\mathbf{Q},\mathbf{K})[i,j]
= \frac{1}{\sqrt{d}}\sum_{k=1}^{d}\max\big(\mathbf{Q}[i,k],\,\mathbf{K}[j,k]\big),
\]
and a leaky ReLU-based probability map with a learned leak term. The paper emphasizes context aware sampling-without-replacement, linear asymptotic computational complexity, and a pseudoconvex formulation of the transformer layer [2408.05822]. In point-cloud experiments, the same framework is used under rotational-invariant constraints by feeding only relative geometric information through \(\mathbf{X_R}\), not raw coordinates.

These stochastic or sampled-efficient constructions differ from exact symmetry-enforced designs. Their primary guarantee is usually linear complexity and equivariance in distribution, expectation, or with respect to a chosen invariant relative encoding, rather than an exact intertwining relation layer by layer.

## 5. Empirical behavior across domains

Representative reported outcomes show that sampling-equivariant attention is not confined to one modality [2111.03420][2602.08047][2006.10503][2604.01466].

| Domain | Reported outcome | Interpretation |
|---|---|---|
| Aerial detection and classification | Rotation AEMD drops from \(0.0099\) to \(0.0059\) for SAN19 vs. SES-Net19; DOTA mAP rises from \(74.82\) to \(76.84\) for Gliding vs. Gliding+SES | Better alignment of local sampling masks with object transformations |
| Discrete planar-equivariant ViTs | COCO Box mAP \(43.7 \to 45.4\), Mask mAP \(39.8 \to 41.2\), ADE20K mIoU \(44.51 \to 44.86\); in SR, SwinIR-EQ uses \(5.2\)M vs \(11.8\)M parameters | Exact equivariance can be maintained across patch embedding, attention, and multiscale resampling |
| SE(3)-equivariant attention | N-body position MSE \(0.0139 \to 0.0076\); position equivariance error \(0.167 \to 3.2\cdot10^{-7}\) for Set Transformer vs. SE(3)-Transformer | Invariant logits plus equivariant value messages improve both accuracy and symmetry fidelity |
| SE(2)-equivariant agent modeling | RMM \(0.7257\) for plain Transformer vs. \(0.7620\) for DriveGATr-3M; Transformer + RPE uses \(15.19\) GB peak memory vs. \(6.35\) GB for DriveGATr-3M | Geometry encoded in the representation can avoid explicit pairwise RPE overhead |

The efficient-sampling literature reports a different empirical pattern. On ModelNet40, MAE-sampled attains \(93.7\%\) while MAE-dense attains \(93.6\%\), indicating that random Hamiltonian-cycle attention can remain competitive with dense attention on point sets [2302.14346]. YOSO reports results consistent with softmax self-attention on Long Range Arena while providing sizable speed-ups and memory savings [2111.09714]. SFT reports competitive results on point cloud, graph, and sequence benchmarks, while faster inference is tied to the learned key/value sampling ratio [2408.05822].

Across these results, two empirical regimes recur. Exact group-equivariant methods mainly improve robustness, sample efficiency, and geometric consistency under prescribed symmetries. Sampling-based efficient methods mainly trade dense all-pairs interactions for learned or random subsets while trying to preserve task performance and permutation structure.

## 6. Misconceptions, limitations, and open directions

A common misconception is that standard self-attention is already geometry-equivariant because it is permutation-equivariant. Several papers make the opposite point: vanilla attention is permutation-equivariant in the token index, but patch embedding, positional encodings, and ordinary resampling break rotation or reflection equivariance unless they are redesigned explicitly [2602.08047]. In 3D, permutation equivariance over point labels does not by itself yield SE(3)-equivariance; invariant attention weights must be paired with equivariant feature representations and relative-coordinate kernels [2006.10503].

A second misconception is that all “sampling-aware” attention is exact. Exact layerwise equivariance is proved for discrete planar ViTs, SE(3)-Transformer attention, LieSelfAttention, and DriveGATr under their respective symmetry assumptions [2602.08047][2006.10503][2012.10885][2604.01466]. By contrast, sampled Hamiltonian attention on sets is permutation-equivariant in expectation, and Bernoulli-sampled YOSO preserves a deterministic collision-attention only in expectation over hashes [2302.14346][2111.09714]. For continuous-group steerable networks, fixed discretizations improve equivariance only as the number of group samples increases, although adaptive equivariant sampling can remove that dependency in the non-attentional setting [2409.08741].

The main technical limitations are consistent across papers. Exact discrete-planar constructions are restricted to rotations by multiples of \(\pi/2\) and reflections or other finite subgroups, and complexity grows with group size \(t\) [2602.08047]. SE(3)-equivariant attention relies on spherical harmonics, Clebsch–Gordan structure, and neighbourhood-based message passing, with spherical harmonics as a bottleneck and higher representation degrees increasing memory and compute [2006.10503]. Lie-group and geometric-algebra formulations require explicit representation theory or algebraic primitives, and extensions to irregular domains or larger groups are non-trivial [2012.10885][2604.01466].

Current directions point toward two converging themes. One is to make more of the transformer stack—not only logits, but tokenization, positional structure, and multiscale resampling—commute exactly with symmetry actions. The other is to make the sampling pattern itself learned, adaptive, and still symmetry-compatible. The first theme is already explicit in holistic equivariant ViTs [2602.08047]; the second is explicit in efficient sampled transformers and in adaptive group sampling for steerable networks [2408.05822][2409.08741]. A plausible implication is that future sampling-equivariant attention will combine exact symmetry constraints with adaptive discretization, so that changing the sampling pattern alters neither the represented geometry nor the functional action of attention except through a known group transformation.

Source: https://www.emergentmind.com/topics/sampling-equivariant-self-attention