Physics Attention

Physics Attention¤

`neojax.nn.physics_attention.PhysicsAttention` ¤

Physics-Attention layer of the Transolver.

Physics-Attention adaptively groups the \(N\) points of the discretized domain into \(G\) learnable slices of points with similar physical states, runs multi-head self-attention among the \(G\) slice tokens only and scatters the result back onto the points. Its cost is thus linear in \(N\). Per head, it computes

Slice: soft slice weights \(w_{g,n} = \operatorname{softmax}_g \big(\mathrm{Linear}(x_n) / \tau\big)\) and the slice tokens \(z_g = \sum_n w_{g,n} f_n / \sum_n w_{g,n}\), where \(x_n, f_n\) are two linear (or convolutional) projections of the input and \(\tau\) is a learnable temperature.
Attention: \(z' = \operatorname{softmax}(QK^\top / \sqrt{d_h}) V\) among the slice tokens, with \(Q, K, V\) projections that are shared across heads.
Deslice: \(f'_n = \sum_g w_{g,n} z'_g\).

The heads are concatenated and linearly projected back to dim channels.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`dim`	`int`	Number of input and output channels.	required
`n_heads`	`int`	Number of attention heads. Default is 8.	`8`
`dim_head`	`int \| None`	Channels per head. Default is `None`, which uses `dim // n_heads`.	`None`
`slice_num`	`int`	Number of slices (slice tokens) `G`. Default is 32.	`32`
`dropout`	`float`	Dropout probability applied to the attention probabilities and the output. Default is 0.0.	`0.0`
`ndim`	`int \| None`	Number of spatial dimensions of a structured grid input. If an integer, the input projections are `ndim`-dimensional convolutions of size `kernel_size`, so the input must be shaped `(dim, d1, ..., d_ndim)`. If `None`, they are pointwise linear maps and the input can have arbitrary trailing axes (e.g. `(dim, N)` for unstructured meshes). Default is `None`.	`None`
`kernel_size`	`int`	Kernel size of the convolutional input projections. Only used if `ndim` is not `None`. Must be odd. Default is 3.	`3`
`temperature_range`	`tuple[float, float] \| None`	Interval the learnable slice temperature is clipped to in the forward pass for stability. `None` disables clipping. Default is `(0.1, 5.0)`.	`(0.1, 5.0)`

Raises:

Type	Description
`ValueError`	If `dim`, `n_heads`, `dim_head`, `slice_num` are not positive or `dim` is not divisible by `n_heads` and `dim_head` is `None`.
`ValueError`	If `dropout` is not in the half-open interval [0, 1).
`ValueError`	If `ndim` is not positive or `kernel_size` is not a positive odd number.
`ValueError`	If `temperature_range` is not a valid positive interval.

Internal Attributes

These fields store the internal layers state (and weights).

in_project_x (PointwiseMLP | eqx.nn.Conv): Projection of the input to slice logits features.
in_project_fx (PointwiseMLP | eqx.nn.Conv): Projection of the input to the features to be sliced.
in_project_slice (PointwiseMLP): Maps per head features to the slice_num slice logits.
to_q (eqx.nn.Linear): Query projection shared across heads.
to_k (eqx.nn.Linear): Key projection shared across heads.
to_v (eqx.nn.Linear): Value projection shared across heads.
to_out (PointwiseMLP): Output projection from n_heads * dim_head to dim channels.
dropout (eqx.nn.Dropout): Dropout layer.
temperature (Float[Array, "1 h 1"]): Learnable slice temperature per head.
dim (int): Number of input and output channels.
n_heads (int): Number of attention heads.
dim_head (int): Channels per head.
slice_num (int): Number of slices.
ndim (int | None): Number of spatial dimensions of grid inputs or None.
temperature_range (tuple[float, float] | None): Clipping interval of the temperature.

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import PhysicsAttention

key = jr.key(0)

# unstructured mesh with 500 nodes
attn = PhysicsAttention(key, dim=64, n_heads=4, slice_num=16)
out = attn(jnp.ones((64, 500)))

# structured 2D grid
attn = PhysicsAttention(key, dim=64, n_heads=4, slice_num=16, ndim=2)
out = attn(jnp.ones((64, 32, 32)))

`call(x: Inexact[Array, 'c ...'], *, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'c ...']` ¤

Forward pass of the Physics-Attention layer.

Parameters:

Name	Type	Description	Default
`x`	`Inexact[Array, 'c ...']`	Input array shaped (dim, N) or (dim, d1, ..., dN). If `ndim` is not `None`, it must be shaped (dim, d1, ..., d_ndim).	required
`key`	`PRNGKeyArray \| None`	PRNG key used for dropout masks.	`None`
`inference`	`bool`	If `True`, dropout is disabled.	`False`

Returns:

Type	Description
`Inexact[Array, 'c ...']`	Output array with the same shape as `x`.

Raises:

Type	Description
`ValueError`	If the number of channels of `x` doesn't match `dim`.
`ValueError`	If `ndim` is not `None` and `x` has not `ndim` spatial axes.
`ValueError`	If dropout is active but no `key` is passed.

Physics Attention