FLARE#

The FLARE model adapts Transolver by replacing its physics-attention blocks with FLARE attention. FLARE uses learned global queries to aggregate and redistribute token information through a low-rank attention mechanism, and supports structured and unstructured data.

For details of the attention mechanism, see the FLARE paper.

class physicsnemo.models.flare.flare.FLARE(*args, **kwargs)[source]#

Bases: Transolver

Transolver with FLARE attention.

Inherits from the core Transolver and replaces all physics attention blocks with FLARE (Fast Low-rank Attention Routing Engine) blocks.

Parameters:

functional_dim (int) – Dimension of input values, not including embeddings.
out_dim (int) – Dimension of model output.
embedding_dim (int | None, optional) – Dimension of input embeddings. Required if unified_pos=False.
n_layers (int, optional) – Number of transformer blocks. Default is 4.
n_hidden (int, optional) – Hidden dimension. Default is 256.
dropout (float, optional) – Dropout rate. Default is 0.0.
n_head (int, optional) – Number of attention heads. Default is 8.
act (str, optional) – Activation function name. Default is "gelu".
mlp_ratio (int, optional) – MLP hidden ratio. Default is 4.
slice_num (int, optional) – Number of global queries for FLARE attention. Default is 32.
unified_pos (bool, optional) – Whether to use unified positional embeddings. Default is False.
ref (int, optional) – Reference grid size for unified position. Default is 8.
structured_shape (None | tuple[int, ...], optional) – Shape of structured data. None for unstructured. Default is None.
time_input (bool, optional) – Whether to include time embeddings. Default is False.
use_te (bool, optional, default=False) – Whether to use Transformer Engine layers and attention.

Forward:

Same as :class:`~physicsnemo.models.transolver.Transolver`.

Outputs:

Same as Transolver.