GeoTransolver#

The GeoTransolver model extends Transolver with Geometry-Aware Latent Embeddings (GALE) attention. It combines physics-aware self-attention over learned state slices with cross-attention to geometry and global context, supporting both unstructured meshes and structured 2D or 3D grids.

GALE layers use either PhysicsAttentionBase (the default setting) or FLARE (with attention_type="GALE_FA") as the self-attention backend.

For more information on GeoTransolver, refer to the GeoTransolver paper.

class physicsnemo.models.geotransolver.geotransolver.GeoTransolver(*args, **kwargs)[source]#

Bases: Module

GeoTransolver: Geometry-Aware Physics Attention Transformer.

GeoTransolver is an adaptation of the Transolver architecture, replacing standard attention with GALE (Geometry-Aware Latent Embeddings) attention. GALE combines physics-aware self-attention on learned state slices with cross-attention to geometry and global context embeddings.

The model projects geometry and global features onto physical state spaces, which are then used as context in all transformer blocks. This design enables the model to incorporate geometric structure and global information throughout the forward pass.

Parameters:

  • functional_dim (int | tuple[int, ...]) – Dimension of the input values (local embeddings), not including global embeddings or geometry features. Input will be projected to n_hidden before processing. Can be a single int or tuple for multiple input types.

  • out_dim (int | tuple[int, ...]) – Dimension of the output of the model. Must have same length as functional_dim if both are tuples.

  • geometry_dim (int | None, optional) – Pointwise dimension of the geometry input features. If provided, geometry features will be projected onto physical states and used as context in all GALE layers. Default is None.

  • global_dim (int | None, optional) – Dimension of the global embedding features. If provided, global features will be projected onto physical states and used as context in all GALE layers. Default is None.

  • n_layers (int, optional) – Number of GALE layers in the model. Default is 4.

  • n_hidden (int, optional) – Hidden dimension of the transformer. Default is 256.

  • dropout (float, optional) – Dropout rate applied across the GALE layers. Default is 0.0.

  • n_head (int, optional) – Number of attention heads in each GALE layer. Must evenly divide n_hidden to yield an integer head dimension. Default is 8.

  • act (str, optional) – Activation function name. Default is "gelu".

  • mlp_ratio (int, optional) – Ratio of MLP hidden dimension to n_hidden. Default is 4.

  • slice_num (int, optional) – Number of learned physical state slices in the GALE layers, representing the number of learned states each layer should project inputs onto. Default is 32.

  • use_te (bool, optional) – Whether to use Transformer Engine backend when available. Default is False.

  • time_input (bool, optional) – Whether to include time embeddings. Default is False.

  • plus (bool, optional) – Whether to use Transolver++ features in the GALE layers. Default is False.

  • include_local_features (bool, optional) – Whether to include local features in the global context. Default is False.

  • radii (list[float], optional) – Radii for the local features. Default is [0.05, 0.25].

  • neighbors_in_radius (list[int], optional) – Neighbors in radius for the local features. Default is [8, 32].

  • n_hidden_local (int, optional) – Hidden dimension for the local features. Default is 32.

  • structured_shape (tuple[int, ...] | None, optional) – If set to (H, W) or (H, W, D), enables structured 2D/3D paths (Conv2d/Conv3d GALE; no ball-query local features). Inputs may be flattened \((B, N, C)\) with \(N = H W\) or \(H W D\), or spatial \((B, H, W, C)\) / \((B, H, W, D, C)\). Default is None.

  • attention_type ({"GALE", "GALE_FA"}, optional) – Attention implementation used inside each GALE block: "GALE" for the reference version, "GALE_FA" for the flash-attention one. Validated in GALEBlock, which raises on any other value. Default is "GALE".

  • state_mixing_mode (str, optional) – How to blend self-attention and cross-attention outputs in GALE layers. "weighted" uses a learnable sigmoid-gated weighted sum. "concat_project" concatenates the two along the head dimension and projects back with a linear layer. Default is "weighted".

Forward:

  • local_embedding (torch.Tensor | tuple[torch.Tensor, …]) – Local embedding: unstructured \((B, N, C)\); structured 2D \((B, H, W, C)\) or flattened \((B, H W, C)\); structured 3D \((B, H, W, D, C)\) or flattened. Can be a tuple for multiple input types.

  • local_positions (torch.Tensor | tuple[torch.Tensor, …] | None, optional) – Local positions for each input, each of shape \((B, N, 3)\). Required if include_local_features=True. Default is None.

  • global_embedding (torch.Tensor | None, optional) – Global embedding of the input data of shape \((B, N_g, C_g)\) where \(N_g\) is number of global tokens and \(C_g\) is global_dim. If None, global context is not used. Default is None.

  • geometry (torch.Tensor | None, optional) – Geometry features of the input data of shape \((B, N, C_{geo})\) where \(C_{geo}\) is geometry_dim. If None, geometry context is not used. Default is None.

  • time (torch.Tensor | None, optional) – Time embedding (currently not implemented). Default is None.

Outputs:

torch.Tensor | tuple[torch.Tensor, …] – When return_embedding_states=False (default): output tensor(s) of shape \((B, N, C_{out})\). Returns a single tensor if input was a single tensor, or a tuple of tensors if input was a tuple (multi-stream). For structured grids, output matches the input layout—flattened \((B, N, C_{out})\) or spatial \((B, H, W, C_{out})\) / \((B, H, W, D, C_{out})\) when inputs were 4D/5D.

When return_embedding_states=True, returns a 2-tuple (output, embedding_states) where output follows the same rules above, and embedding_states is of shape \((B, H, S, D_c)\) (geometry/global context), or None if no context sources were provided.

Raises:

  • ValueError – If n_hidden is not evenly divisible by n_head.

  • ValueError – If functional_dim and out_dim have different lengths when both are tuples.

  • NotImplementedError – If time is provided (not yet implemented).

Notes

Unstructured mesh uses linear GALE projection; structured structured_shape uses the same Conv2d/Conv3d slice projection as Transolver. Ball-query local features are disabled when structured_shape is set.

For more details on Transolver, see:

See also

GALE

The attention mechanism used in GeoTransolver.

GALEBlock

Transformer block using GALE attention.

ContextProjector

Projects context features onto physical states.

Examples

Basic usage with local embeddings only:

>>> import torch
>>> from physicsnemo.models.geotransolver import GeoTransolver
>>> model = GeoTransolver(
...     functional_dim=64,
...     out_dim=3,
...     n_hidden=256,
...     n_layers=4,
...     use_te=False,
... )
>>> local_emb = torch.randn(2, 1000, 64)  # (batch, nodes, features)
>>> output = model(local_emb)
>>> output.shape
torch.Size([2, 1000, 3])

Usage with geometry, global context, and embedding states:

>>> model = GeoTransolver(
...     functional_dim=64,
...     out_dim=3,
...     geometry_dim=3,
...     global_dim=16,
...     n_hidden=256,
...     n_layers=4,
...     use_te=False,
... )
>>> local_emb = torch.randn(2, 1000, 64)
>>> geometry = torch.randn(2, 1000, 3)  # (batch, nodes, spatial_dim)
>>> global_emb = torch.randn(2, 1, 16)  # (batch, 1, global_features)
>>> output = model(local_emb, global_embedding=global_emb, geometry=geometry)
>>> output.shape
torch.Size([2, 1000, 3])

To also retrieve the geometry/global context embeddings:

>>> output, emb_states = model(
...     local_emb,
...     global_embedding=global_emb,
...     geometry=geometry,
...     return_embedding_states=True,
... )
>>> emb_states.shape[0] == 2  # batch dimension preserved
True

Structured 2D grid:

>>> model = GeoTransolver(
...     functional_dim=3,
...     out_dim=1,
...     structured_shape=(8, 8),
...     n_hidden=64,
...     n_head=4,
...     n_layers=2,
...     use_te=False,
... )
>>> y = model(torch.randn(2, 8, 8, 3))
>>> y.shape
torch.Size([2, 8, 8, 1])

Building blocks#

class physicsnemo.models.geotransolver.context_projector.ContextProjector(
dim: int,
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
slice_num: int = 64,
use_te: bool = False,
plus: bool = False,
concrete_dropout: bool = False,
)[source]#

Bases: _SliceToContextMixin, Module

Projects context features onto physical state space.

This context projector is conceptually similar to half of a GALE attention layer. It projects context values (geometry or global embeddings) onto a learned physical state space, but unlike a full attention layer, it never projects back to the original space. The projected features are used as context in all GALE blocks of the GeoTransolver model.

Parameters:

  • dim (int) – Input dimension of the context features.

  • heads (int, optional) – Number of projection heads. Default is 8.

  • dim_head (int, optional) – Dimension of each projection head. Default is 64.

  • dropout (float, optional) – Dropout rate. Default is 0.0.

  • slice_num (int, optional) – Number of learned physical state slices. Default is 64.

  • use_te (bool, optional) – Whether to use Transformer Engine backend when available. Default is False.

  • plus (bool, optional) – Whether to use Transolver++ features. Default is False.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, N, C)\) where \(B\) is batch size, \(N\) is number of tokens, and \(C\) is number of channels.

Outputs:

torch.Tensor – Slice tokens of shape \((B, H, S, D)\) where \(H\) is number of heads, \(S\) is number of slices, and \(D\) is head dimension.

Notes

The global features are reused in all blocks of the model, so the learned projections must capture globally useful features rather than layer-specific ones.

See also

GALE

Full GALE attention layer that uses these projected context features.

GeoTransolver

Main model that uses ContextProjector for geometry and global embeddings.

Examples

>>> import torch
>>> projector = ContextProjector(dim=64, heads=8, dim_head=32, slice_num=32, use_te=False)
>>> x = torch.randn(2, 100, 64)  # (batch, tokens, features)
>>> slice_tokens = projector(x)
>>> slice_tokens.shape
torch.Size([2, 8, 32, 32])
project_input_onto_slices(
x: Float[Tensor, 'batch tokens channels'],
) → Float[Tensor, 'batch tokens heads dim'] | tuple[Float[Tensor, 'batch tokens heads dim'], Float[Tensor, 'batch tokens heads dim']][source]#

Project the input onto the slice space.

Parameters:

x (torch.Tensor) – Input tensor of shape \((B, N, C)\) where \(B\) is batch size, \(N\) is number of tokens, and \(C\) is number of channels.

Returns:

If plus=True, returns single tensor of shape \((B, N, H, D)\) where \(H\) is number of heads and \(D\) is head dimension. If plus=False, returns tuple of two tensors both of shape \((B, N, H, D)\), representing the query and key projections respectively.

Return type:

torch.Tensor or tuple[torch.Tensor, torch.Tensor]

class physicsnemo.models.geotransolver.context_projector.StructuredContextProjector(
dim: int,
spatial_shape: tuple[int, ...],
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
slice_num: int = 64,
kernel: int = 3,
use_te: bool = False,
plus: bool = False,
concrete_dropout: bool = False,
)[source]#

Bases: _SliceToContextMixin, Module

Context projector with Conv2d/Conv3d geometry encoding on structured grids.

Same output interface as ContextProjector—slice tokens \((B, H, S, D)\)—but projects per-cell geometry via spatial convolutions aligned with structured GALE attention.

class physicsnemo.models.geotransolver.context_projector.GeometricFeatureProcessor(
radius: float,
neighbors_in_radius: int,
feature_dim: int,
hidden_dim: int,
)[source]#

Bases: Module

Processes geometric features at a single spatial scale using BQWarp.

This is a simple, reusable component that handles neighbor querying and feature processing for one radius scale. It encapsulates the BQWarp + MLP pattern used throughout the model.

Parameters:

  • radius (float) – Query radius for neighbor search.

  • neighbors_in_radius (int) – Maximum number of neighbors within the radius.

  • feature_dim (int) – Dimension of the input features to query.

  • hidden_dim (int) – Output dimension after MLP processing.

Forward:

  • query_points (torch.Tensor) – Query coordinates of shape \((B, N, 3)\) where \(B\) is batch size and \(N\) is number of query points.

  • key_features (torch.Tensor) – Features to query from of shape \((B, N, C)\) where \(C\) is feature_dim.

Outputs:

torch.Tensor – Processed features of shape \((B, N, D)\) where \(D\) is hidden_dim.

See also

MultiScaleFeatureExtractor

Uses multiple GeometricFeatureProcessor instances.

BQWarp

The ball query operation used internally.

Examples

>>> import torch
>>> processor = GeometricFeatureProcessor(
...     radius=0.1, neighbors_in_radius=16, feature_dim=3, hidden_dim=64
... )
>>> query_points = torch.randn(2, 100, 3)  # (batch, points, xyz)
>>> key_features = torch.randn(2, 100, 3)  # (batch, points, features)
>>> output = processor(query_points, key_features)
>>> output.shape
torch.Size([2, 100, 64])
class physicsnemo.models.geotransolver.context_projector.MultiScaleFeatureExtractor(
geometry_dim: int,
radii: list[float],
neighbors_in_radius: list[int],
hidden_dim: int,
n_head: int,
dim_head: int,
dropout: float = 0.0,
slice_num: int = 64,
use_te: bool = False,
plus: bool = False,
concrete_dropout: bool = False,
)[source]#

Bases: Module

Multi-scale geometric feature extraction with minimal complexity.

Manages multiple GeometricFeatureProcessor instances for different radii. Provides both tokenized context and concatenated local features.

Parameters:

  • geometry_dim (int) – Dimension of geometry features.

  • radii (list[float]) – Radii for multi-scale processing.

  • neighbors_in_radius (list[int]) – Neighbors per radius (must have same length as radii).

  • hidden_dim (int) – Hidden dimension for processing.

  • n_head (int) – Number of attention heads.

  • dim_head (int) – Dimension per head.

  • dropout (float, optional) – Dropout rate. Default is 0.0.

  • slice_num (int, optional) – Number of slices for context tokenization. Default is 64.

  • use_te (bool, optional) – Whether to use Transformer Engine. Default is False.

  • plus (bool, optional) – Whether to use Transolver++ features. Default is False.

Forward:

  • This class does not implement a standard ``forward`` method. Instead, use

  • - :meth:`extract_context_features` (Get tokenized features for GALE context.)

  • - :meth:`extract_local_features` (Get concatenated features for local pathway.)

See also

GeometricFeatureProcessor

Single-scale processor used by this class.

ContextProjector

Tokenizer used for context features.

GlobalContextBuilder

High-level builder that uses this class.

Examples

>>> import torch
>>> extractor = MultiScaleFeatureExtractor(
...     geometry_dim=3,
...     radii=[0.05, 0.25],
...     neighbors_in_radius=[8, 32],
...     hidden_dim=32,
...     n_head=8,
...     dim_head=32,
...     use_te=False,
... )
>>> spatial_coords = torch.randn(2, 100, 3)
>>> geometry = torch.randn(2, 100, 3)
>>> context_feats = extractor.extract_context_features(spatial_coords, geometry)
>>> len(context_feats)  # One per scale
2
>>> local_feats = extractor.extract_local_features(spatial_coords, geometry)
>>> local_feats.shape  # Concatenated across scales
torch.Size([2, 100, 64])
extract_context_features(
spatial_coords: Float[Tensor, 'batch points spatial_dim'],
geometry: Float[Tensor, 'batch points geometry_dim'],
) → list[Float[Tensor, 'batch heads slices dim']][source]#

Extract and tokenize features for context.

Parameters:

  • spatial_coords (torch.Tensor) – Spatial coordinates of shape \((B, N, 3)\).

  • geometry (torch.Tensor) – Geometry features of shape \((B, N, C_{geo})\).

Returns:

List of tokenized context features, one per scale, each of shape \((B, H, S, D)\).

Return type:

list[torch.Tensor]

extract_local_features(
spatial_coords: Float[Tensor, 'batch points spatial_dim'],
geometry: Float[Tensor, 'batch points geometry_dim'],
) → Float[Tensor, 'batch points total_hidden'][source]#

Extract and concatenate features for local pathway.

Parameters:

  • spatial_coords (torch.Tensor) – Spatial coordinates of shape \((B, N, 3)\).

  • geometry (torch.Tensor) – Geometry features of shape \((B, N, C_{geo})\).

Returns:

Concatenated local features of shape \((B, N, D_{total})\) where \(D_{total}\) is hidden_dim * num_scales.

Return type:

torch.Tensor

class physicsnemo.models.geotransolver.context_projector.GlobalContextBuilder(
functional_dims: tuple[int, ...],
geometry_dim: int | None = None,
global_dim: int | None = None,
radii: list[float] | None = None,
neighbors_in_radius: list[int] | None = None,
n_hidden_local: int = 32,
n_hidden: int = 256,
n_head: int = 8,
dropout: float = 0.0,
slice_num: int = 32,
use_te: bool = False,
plus: bool = False,
include_local_features: bool = False,
structured_shape: tuple[int, ...] | None = None,
concrete_dropout: bool = False,
)[source]#

Bases: Module

Orchestrates all context construction with a clean, simple interface.

Manages geometry tokenization, global embedding tokenization, and optional multi-scale local features. This is the main entry point for building context in the GeoTransolver model.

Parameters:

  • functional_dims (tuple[int, ...]) – Dimensions of each functional input type.

  • geometry_dim (int | None, optional) – Geometry feature dimension. If None, geometry context is disabled. Default is None.

  • global_dim (int | None, optional) – Global embedding dimension. If None, global context is disabled. Default is None.

  • radii (list[float], optional) – Radii for local features. Default is [0.05, 0.25].

  • neighbors_in_radius (list[int], optional) – Neighbors per radius. Default is [8, 32].

  • n_hidden_local (int, optional) – Hidden dim for local features. Default is 32.

  • n_hidden (int, optional) – Model hidden dimension. Default is 256.

  • n_head (int, optional) – Number of attention heads. Default is 8.

  • dropout (float, optional) – Dropout rate. Default is 0.0.

  • slice_num (int, optional) – Number of slices for tokenization. Default is 32.

  • use_te (bool, optional) – Whether to use Transformer Engine. Default is False.

  • plus (bool, optional) – Whether to use Transolver++ features. Default is False.

  • include_local_features (bool, optional) – Enable local feature extraction. Default is False.

  • structured_shape (tuple[int, ...] | None, optional) – If set, disables ball-query extractors and uses StructuredContextProjector for geometry when geometry_dim is set. Default is None.

Forward:

  • This class does not implement a standard ``forward`` method. Instead, use

  • :meth:`build_context` to construct context, local features, and the

  • detached geometry context.

See also

ContextProjector

Used for tokenizing geometry and global embeddings.

MultiScaleFeatureExtractor

Used for multi-scale local features.

GeoTransolver

Main model that uses this builder.

Examples

>>> import torch
>>> builder = GlobalContextBuilder(
...     functional_dims=(64,),
...     geometry_dim=3,
...     global_dim=16,
...     n_hidden=256,
...     n_head=8,
...     use_te=False,
... )
>>> local_embeddings = (torch.randn(2, 100, 64),)
>>> geometry = torch.randn(2, 100, 3)
>>> global_embedding = torch.randn(2, 1, 16)
>>> context, local_feats, geo_ctx = builder.build_context(
...     local_embeddings, None, geometry, global_embedding
... )
>>> context.shape
torch.Size([2, 8, 32, 64])
build_context(
local_embeddings: tuple[Float[Tensor, 'batch tokens features'], ...],
local_positions: tuple[Float[Tensor, 'batch tokens spatial_dim'], ...] | None,
geometry: Float[Tensor, 'batch tokens geometry_dim'] | None = None,
global_embedding: Float[Tensor, 'batch global_tokens global_dim'] | None = None,
) → tuple[Float[Tensor, 'batch heads slices context_dim'] | None, list[Float[Tensor, 'batch tokens local_features']] | None, Float[Tensor, 'batch heads slices dim_head'] | None][source]#

Build all context and local features.

Parameters:

  • local_embeddings (tuple[torch.Tensor, ...]) – Input embeddings, each of shape \((B, N, C_i)\) where \(B\) is batch size, \(N\) is number of tokens, and \(C_i\) is the feature dimension for input type \(i\).

  • local_positions (tuple[torch.Tensor, ...] | None) – Local positions, each of shape \((B, N, 3)\). These are used to query neighbors for local features. Required if include_local_features=True.

  • geometry (torch.Tensor | None, optional) – Geometry features of shape \((B, N, C_{geo})\). Default is None.

  • global_embedding (torch.Tensor | None, optional) – Global embedding of shape \((B, N_g, C_g)\). Default is None.

Returns:

  • context: Concatenated context tensor of shape \((B, H, S, D_c)\) where \(D_c\) is the total context dimension, or None if no context sources are provided.

  • local_features: List of local feature tensors, one per input type, each of shape \((B, N, D_l)\), or None if local features are disabled.

  • geometry_context_detached: Detached geometry-tokenizer output of shape \((B, H, S, D)\), intended for downstream observers such as the embedded OOD guard. None when geometry tokenization is disabled or no geometry was provided.

Return type:

tuple[torch.Tensor | None, list[torch.Tensor] | None, torch.Tensor | None]

Raises:

ValueError – If local_positions is None but local features are enabled.

get_context_dim() → int[source]#

Return total context dimension.

Returns:

Total dimension of the concatenated context features.

Return type:

int

FLARE Attention Backend#

For large meshes, setting attention_type="GALE_FA" swaps the physics-attention slice mechanism for the FLARE (Fast Low-rank Attention Routing Engine) backend. GALE_FA keeps GeoTransolver’s geometry- and context-aware cross-attention while using FLARE for the self-attention pass over learned physical-state slices, reducing attention cost at scale. Refer also the FLARE model documentation.

class physicsnemo.nn.module.gale.GALE_FA(
dim,
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
n_global_queries: int = 64,
use_te: bool = False,
context_dim: int = 0,
concrete_dropout: bool = False,
state_mixing_mode: str = 'weighted',
)[source]#

Bases: Module

GALE_FA: Geometry-Aware Latent Embeddings with FLARE self-Attention attention layer.

Adopted:

GALE_FA is an alternative to the GALE attention mechanism of the GeoTransolver. It supports cross-attention with a context vector, built from geometry and global embeddings. GALE_FA combines FLARE self-attention on learned physical state slices with cross-attention to geometry-aware context, using a learnable mixing weight to blend the two.

Parameters:

  • dim (int) – Input dimension of the features.

  • heads (int, optional) – Number of attention heads. Default is 8.

  • dim_head (int, optional) – Dimension of each attention head. Default is 64.

  • dropout (float, optional) – Dropout rate. Default is 0.0.

  • n_global_queries (int, optional) – Number of learned global queries. Default is 64.

  • use_te (bool, optional) – Whether to use Transformer Engine backend when available. Default is False.

  • context_dim (int, optional) – Dimension of the context vector for cross-attention. Default is 0.

  • concrete_dropout (bool, optional) – Whether to use learned concrete dropout instead of standard dropout. Default is False.

  • state_mixing_mode (str, optional) – How to blend self-attention and cross-attention outputs. "weighted" uses a learnable sigmoid-gated weighted sum. "concat_project" concatenates the two along the head dimension and projects back with a linear layer. Default is "weighted".

Forward:

  • x (tuple[torch.Tensor, …]) – Tuple of input tensors, each of shape \((B, N, C)\) where \(B\) is batch size, \(N\) is number of tokens, and \(C\) is number of channels.

  • context (tuple[torch.Tensor, …] | None, optional) – Context tensor for cross-attention of shape \((B, H, S_c, D_c)\) where \(H\) is number of heads, \(S_c\) is number of context slices, and \(D_c\) is context dimension. If None, only self-attention is applied. Default is None.

Outputs:

list[torch.Tensor] – List of output tensors, each of shape \((B, N, C)\), same shape as inputs.

Notes

The mixing between self-attention and cross-attention is controlled by a learnable parameter state_mixing which is passed through a sigmoid function to ensure the mixing weight stays in \([0, 1]\).

See also

GALE

Original GeoTransolver GALE attention class.

GALEBlock

Transformer block that calls GALE or GALE_FA attention.

Examples

>>> import torch
>>> gale_fa = GALE_FA(dim=256, heads=8, dim_head=32, context_dim=32)
>>> x = (torch.randn(2, 100, 256),)  # Single input tensor in tuple
>>> context = torch.randn(2, 8, 64, 32)  # Context for cross-attention
>>> outputs = gale_fa(x, context)
>>> len(outputs)
1
>>> outputs[0].shape
torch.Size([2, 100, 256])