AB-GATr: Anchored-Branched Geometric Algebra Transformer

Symmetry in the Wild: The Role of Equivariance in Neural Fluid Surrogates

Authors: Patryk Rygiel, Julian Suk, Kak Khee Yeung, Christoph Brune and Jelmer M. Wolterink

Preprint

Abstract

Neural surrogates enable orders-of-magnitude acceleration of computational fluid dynamics (CFD) simulations, but their use in real-world applications requires addressing scalability to large, high-resolution surface and volume meshes and accounting for limited training data through the use of inductive biases. Group-equivariant architectures are a principled way to introduce such bias, yet they can be detrimental when the learning problem itself breaks symmetry, for example, due to strong distributional alignment in the dataset. We find that on strongly aligned datasets (e.g., automotive aerodynamics), i.e., those that break symmetry, enforcing equivariance can degrade in-distribution performance. In contrast, across irregular benchmarks (e.g., hemodynamics) with diverse geometries and varying alignment, equivariance is consistently beneficial. Moreover, we find that across all benchmarks, the explicit equivariance reliably outperforms implicit symmetry learning through data augmentation. We conclude that equivariance is not universally beneficial across domains, yet it brings tangible advantages in problems lacking strong data regularities

AB-GATr

Designing neural surrogates for real-world CFD involves simultaneously addressing two competing challenges: scalability to large, high-resolution surface and volume meshes, and a physically meaningful inductive bias. AB-GATr aims to bridge these paradigms by proposing a scalable, equivariant architecture. AB-GATr is an $E(3)$-equivariant neural surrogate built upon Anchored-Branched Universal Physics Transformer (AB-UPT), the current state-of-the-art on automotive aerodynamics benchmarks (Alkin et al. (2025)). AB-GATr retains the key architectural components that enable AB-UPT to scale to large CFD meshes, namely anchor-attention and multi-branching, while incorporating equivariance and geometric inductive bias through projective geometric algebra (PGA) as introduced by Brehmer et al. (2023).

Getting started

# Create virtual environment
conda create --name abgatr python=3.11
conda activate abgatr

# Install via pip
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
pip install git+https://github.com/PatRyg99/abgatr.git --find-links https://data.pyg.org/whl/torch-2.5.0+cu124.html

Using AB-GATr

1) Construct HeteroData object to hold volume and surface data

Default embedding_cls constructs following PGA embedding:

• x (N, F) -> scalars
• x_vecs (N, K, 3) -> planes
• pos (N, 3) -> points

This behaviour can be overwritten by providing custom embedding_cls to ABGATr.

import torch
from abgatr.data import HeteroData, GetRadiusPooling
from abgatr.models import ABGATr, ABUPT

S, V = 1000, 10_000

num_surface_scalars = 2
num_surface_channels = 3
num_volume_scalars = 4
num_volume_channels = 5

device = torch.device("cuda")

data = HeteroData({
    "surface": {
        "x": torch.rand(S, num_surface_scalars, device=device),
        "x_vecs": torch.rand(S, num_surface_channels, 3, device=device),
        "pos": torch.rand(S, 3, device=device),
        "batch": torch.zeros(S, device=device).long(),
    },
    "volume": {
        "x": torch.rand(V, num_volume_scalars, device=device),
        "x_vecs": torch.rand(V, num_volume_channels, 3, device=device),
        "pos": torch.rand(V, 3, device=device),
        "batch": torch.zeros(V, device=device).long(),
    },
})

2) Assigning query and anchor tokens

# Assign every 10th point as anchor
data["surface"].anchors_index = torch.randperm(S, device=device)[:S // 10]
data["volume"].anchors_index = torch.randperm(V, device=device)[:V // 10]

# Assign all points as queries
data["surface"].query_index = torch.arange(S, device=device)
data["volume"].query_index = torch.arange(V, device=device)

3) Precomputing radius pooling

Required when using message_passing pooling akin to AB-UPT, alternatively pooling_mode can be changed to cross_attention akin to LaB-GATr - then the precomputing of radius pooling is not required. The scale0_sampling_index is used for denoting pooling tokens - we set it to the same tokens as anchors.

data["surface"].scale0_sampling_index = data["surface"].anchors_index
pooler = GetRadiusPooling("surface", "scale0_sampling_index", r=0.1)
data = pooler(data)

3) Constructing AB-GATr and running forward pass

• scalars = x
• channels = pos|x_vecs

The default embedding_cls provides dislodging into a vector and scalar per each output channel, this behaviour can be overwritten by providing custom embedding_cls to AB-GATr.

abgatr = ABGATr(
    num_surface_input_channels=num_surface_channels + 1,
    num_surface_input_scalars=num_surface_scalars,
    num_volume_input_channels=num_volume_channels + 1,
    num_volume_input_scalars=num_volume_scalars,
    num_surface_output_channels=1,
    num_volume_output_channels=1,
    pooling_mode="message_passing",
).to(device=device)

(
    VA_vecs, VA_scalars, VQ_vecs, VQ_scalars, # volume 
    SA_vecs, SA_scalars, SQ_vecs, SQ_scalars # surface
) = abgatr(data)

5*) Running baselines

AB-UPT

Original AB-UPT implementation is wrapped so it can be run with the same data structure. For AB-UPT both x and x_vecs will be flatten to one vector together with pos if use_pos=True.

abupt = ABUPT(
    volume_input_dim=3 * (num_volume_channels + 1) + num_volume_scalars,
    surface_input_dim=3 * (num_surface_channels + 1) + num_surface_scalars,
    output_dim_surface=4, # vector + scalar
    output_dim_volume=4, # vector + scalar
    radius=0.1, # ABUPT runs radius pooling inside
    use_pos=True, # append pos to input features
).to(device=device)

(
    VA_vecs, VA_scalars, VQ_vecs, VQ_scalars, # volume 
    SA_vecs, SA_scalars, SQ_vecs, SQ_scalars # surface
) = abupt(data)

LaB-GATr

Original LaB-GATr implementation is enhanced with sine-cosine and RoPE positional encodings - this change yields better performance and serves as the stronger baseline to AB-UPT and AB-GATr that use the same postional encoding scheme.

• As AB-GATr, LaB-GATr utilizes embedding_cls for PGA embedding and dislodging which can be overwritten for custom behaviour,
• As AB-GATr, LaB-GATr has two pooling modes message_passing and cross_attention - be default it utilizes cross_attention which does not require initial radius pooling,
• Contrary to AB-GATr, LaB-GATr requires farthest-point-sampling (FPS) preprocessing instead of radii pooling - this step is reffered to as 'tokenisation' and can be viewed similarly to anchor selection step in AB-GATr,
• LaB-GATr cannot process volume and surface data at once, hence, two seperate models need to be trained.

Volume:

from abgatr.data import GetFPS

fps_preprocessor = GetFPS("volume", k=V // 10, simplex="tetrahedron")
data = fps_preprocessor(data)

labgatr = LaBGATr(
    num_input_channels=num_volume_channels + 1,
    num_input_scalars=num_volume_scalars,
    num_output_channels=1,
    pooling_mode="cross_attention",
).to(device=device)

V_vecs, V_scalars = labgatr(data["volume"])

Surface:

fps_preprocessor = GetFPS("surface", k=S // 10, simplex="tetrahedron")
data = fps_preprocessor(data)

labgatr = LaBGATr(
    num_input_channels=num_surface_channels + 1,
    num_input_scalars=num_surface_scalars,
    num_output_channels=1,
    pooling_mode="cross_attention",
).to(device=device)

S_vecs, S_scalars = labgatr(data["surface"])

Reproducing experiments

Coming soon...

Citation

If you find our code useful, please cite:

@misc{rygiel2026symmetrywildroleequivariance,
      title={Symmetry in the Wild: The Role of Equivariance in Neural Fluid Surrogates}, 
      author={Patryk Rygiel and Julian Suk and Kak Khee Yeung and Christoph Brune and Jelmer M. Wolterink},
      year={2026},
      eprint={2605.18816},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.18816}, 
}