Geometric Augmentation in Dot-Product Attention

• The paper introduces geometric modifications that replace Euclidean similarity with non-Euclidean metrics, enhancing robustness and representation diversity.
• It details implementations using Mahalanobis, hyperbolic, and group-equivariant methods as efficient, drop-in alternatives for standard attention.
• Empirical results demonstrate significant performance gains, including lower perplexity and higher PSNR, across tasks in language, vision, and multi-modal domains.

Geometric augmentation of dot-product attention refers to a family of modifications to the standard scaled dot-product attention mechanism, in which the underlying geometry or metric used for scoring interactions between tokens is replaced, enhanced, or generalized to encode richer structural, spatial, or contextual relationships. These modifications enable attention modules to reflect non-Euclidean structure, symmetry, or invariance properties, and to enforce inductive biases such as hierarchical, geometric, rotation, or translation equivariance. Approaches include Mahalanobis metrics, hyperbolic or geometric algebraic operators, group-based representations, and kernelization, resulting in significant benefits in robustness, representation diversity, empirical performance, and sample efficiency across domains.

1. Mathematical Foundations and Geometric Reformulation

The canonical scaled dot-product attention computes token interactions as

$$\alpha_{ij} = \mathrm{softmax}_j\bigl(q_i^\top k_j\bigr), \qquad h_i = \sum_j \alpha_{ij} v_j$$

where $q_i, k_j, v_j$ are the query, key, and value vectors. This implements a similarity function in Euclidean space, assigning weights proportional to the inner product—equivalently, a Gaussian kernel over Euclidean distance after layer normalization: $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$ This operation can be reinterpreted as a projection of each query vector onto the convex hull or affine manifold (“surface”) defined by the key vectors, via kernel averaging with a Gaussian kernel (Sanger, 25 Jan 2026). This geometric perspective immediately motivates the introduction of alternative metrics and structures for more expressive and robust attention.

2. Mahalanobis and Elliptical Metric Augmentation

Elliptical Attention (Nielsen et al., 2024) extends standard attention by generalizing the similarity metric to a Mahalanobis distance: $d_M(q_i, k_j) = (q_i - k_j)^\top M (q_i - k_j)$ where $M \succ 0$ is typically diagonal and data-dependent. The attention scores become: $$\alpha_{ij}^\mathrm{ellip} = \mathrm{softmax}_j\bigl( -d_M(q_i, k_j) / \sigma^2 \bigr)$$ Through algebraic manipulation, this is equivalent to a weighted “elliptical” dot product: $H = \mathrm{softmax}(Q M K^\top / \sqrt{D}) V$ This change transforms the attention neighborhood from a Euclidean sphere to a hyper-ellipsoid

$\{ k : (q-k)^\top M (q-k) \leq r^2 \}$

The diagonal metric $M$ is computed adaptively at each layer by the mean absolute change in value vectors, scaled to unit maximum, and introduces no extra learnable parameters. This adaption stretches low-variance (informative) directions in the embedding space, reducing representation collapse and yielding robustness to contamination or adversarial noise.

3. Non-Euclidean and Group-Equivariant Geometries

Several augmentation paradigms explicitly encode non-Euclidean, hierarchical, or group-invariant structures.

Hyperbolic (Cone) Attention: Cone Attention (Tseng et al., 2023) replaces the Euclidean dot product with a kernel based on hyperbolic entailment cones, associating queries and keys by the depth of their lowest common ancestor in an implicit tree structure: $K(u,v) = f\bigl(H(\sup_2(u, v), \mathcal{S})\bigr)$ where $q_i, k_j, v_j$0 is hyperbolic distance and $q_i, k_j, v_j$1 the root. This approach models hierarchy-aware similarity, with partial order and transitivity properties unattainable for pure dot products.

SE(2)-Invariance via Group Representations: In SE(2)-invariant attention (Pronovost et al., 24 Jul 2025), each token is augmented with a pose $q_i, k_j, v_j$2. The standard attention is replaced by: $q_i, k_j, v_j$3 where $q_i, k_j, v_j$4 encodes the relative pose as a block-diagonal or Fourier basis representation. Factorization into per-token transforms allows exact group-equivariant attention at linear memory cost.

Geometric Algebraic Attention: Geometric Algebra Attention (Spellings, 2021) replaces raw inner products by invariants of geometric products of point coordinates, ensuring permutation and rotation equivariance by construction. Attention scores are produced by passing these invariants through MLPs, providing built-in geometric priors and suitable for small 3D point clouds.

Geometric Transform Attention (GTA): In multi-view vision, GTA (Miyato et al., 2023) samples the group element for each token $q_i, k_j, v_j$5 in $q_i, k_j, v_j$6 and scores key-value interactions by aligning them via group representations: $q_i, k_j, v_j$7 where $q_i, k_j, v_j$8 and $q_i, k_j, v_j$9 is a homomorphism to $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$0. No extra learnable parameters are added, and equivariance is ensured.

4. Theoretical Advantages and Robustness

Geometric augmentations deliver explicit benefits:

• Reduced Representation Collapse: By stretching or warping the attention surface in anisotropic directions, mechanisms like Elliptical Attention decrease the risk of hidden states collapsing onto low-rank manifolds and preserve context diversity, as evidenced by more uniform token similarity over layers (Nielsen et al., 2024).
• Robustness to Noise and Contamination: The Mahalanobis formulation admits provable sensitivity bounds, with per-coordinate scaling factors $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$1 reducing the model's susceptibility to adversarial perturbations or contaminated samples (Nielsen et al., 2024). Hyperbolic and group-invariant approaches natively encode structural priors that can further discourage spurious context aggregation.
• Compact or Efficient Representation: Many non-Euclidean approaches (e.g., cone attention) capture hierarchical or geometric structure at lower embedding dimensions, yielding efficiency at inference and decreasing parameters while retaining or improving performance (Tseng et al., 2023).

5. Empirical Results Across Domains

Augmented dot-product attention mechanisms have been validated on a spectrum of large-scale, high-variance tasks:

Method	Domain/Task	Key Gains
Elliptical Attention	Language modeling (WikiText-103), vision (ImageNet, ADE20K), multi-modal	Clean PPL: 34.29→32.00; robust PPL: 74.56→52.59; ImageNet attacked top-1 +40–50%; mIoU +4.7% (Nielsen et al., 2024)
Cone Attention	NMT, vision transformer, graph learning	IWSLT De→En BLEU +1.0; DeiT-Ti Top-1 +1.09; more compact embedding required for matching accuracy (Tseng et al., 2023)
SE(2) Fourier Attention	Autonomous driving (agent simulation)	3–4% lower minADE on turning trajectories; linear memory scaling at large N (Pronovost et al., 24 Jul 2025)
Geometric Algebra Attn	Small point clouds physics/chemistry	Invariant/covariant prediction of forces, structures, or class; robust to spatial symmetries (Spellings, 2021)
GTA	Multi-view novel view synthesis	CLEVR-TR PSNR: SRT+GTA 39.63 (vs. 33.51 SRT baseline); MSN-Hard PSNR +1.45; RealEstate10k PSNR +1.2 (Miyato et al., 2023)

Additionally, these methods typically require only negligible increases in computational or memory overhead relative to standard $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$2 attention.

6. Implementation and Integration Considerations

In practice, geometric augmentations are realized as drop-in replacements for standard attention. For Elliptical Attention, the update is to replace $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$3 with $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$4, with $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$5 computed on-line from the value vectors of the preceding layer, and a fixed $$z_{ij} = \mathrm{softmax}_j\bigl(q_i \cdot k_j\bigr) = \frac{\exp(-\|q_i - k_j\|^2/2)}{\sum_{j'}\exp(-\|q_i - k_{j'}\|^2/2)}$$6 hyperparameter governing difference scaling (Nielsen et al., 2024). In SE(2)-Fourier attention, per-token transformations precede and follow efficient FlashAttention calls, ensuring scalability (Pronovost et al., 24 Jul 2025). Hyperbolic kernels and GTA require only local feature-wise multiplication with block matrices or group representations, with no added learnable parameters.

Careful ablation analyses demonstrate that the benefits of geometric augmentation diminish when random or uniform metrics are used, or when equivariant transformations are omitted from values as well as queries/keys. Performance further improves when scaling or centering schemes for metric matrices are optimized.

7. Perspectives and Extensions

The geometric projection lens for attention suggests multiple further directions (Sanger, 25 Jan 2026):

• Anisotropic or learned kernel metrics beyond diagonal Mahalanobis.
• Nonlinear (curved) projection surfaces, via local PCA or kernel-based submanifolds.
• Attention on known Riemannian manifolds, e.g., spherical or hyperbolic embedding spaces.
• Continuous-time or infinite impulse response versions.
• Head- or token-specific control of projection dimension or bandwidth.

These generalizations provide a systematic foundation for principled attention design tied directly to the domain's structure, improving generalization, robustness, and efficiency. The geometric reinterpretation fundamentally bridges deep learning, kernel methods, and structured signal processing, making geometric augmentation of dot-product attention a fertile ground for ongoing research and application.