FLARE Attention#

FLARE (Fast Low-rank Attention Routing Engine) is a low-rank self-attention mechanism that aggregates token features into learned global query slots before routing information back to the tokens. It provides an alternative to PhysicsAttentionBase and can use either PyTorch scaled dot-product attention or Transformer Engine by setting use_te=True.

For details of the method, see the FLARE paper.

class physicsnemo.nn.module.flare_attention.FLARE(
dim,
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
n_global_queries: int = 64,
use_te: bool = False,
)[source]#

Bases: Module

FLARE: Fast Low-rank Attention Routing Engine attention layer. Adopted: - FLARE attention: Fast Low-rank Attention Routing Engine

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, default=False) – Whether to use Transformer Engine backend when available.

Forward:

x (torch.Tensor[Batch, N_points, N_Channels] ([B, N, C]))

Outputs:

torch.Tensor[Batch, N_points, N_Channels] ([B, N, C])

Examples

>>> import torch
>>> flare = FLARE(dim=256, heads=8, dim_head=32)
>>> x = torch.randn(2, 100, 256)
>>> outputs = flare(x)
>>> outputs.shape
torch.Size([2, 100, 256])
On this page