Positional Embeddings

Positional Embedding¤

Appends grid coordinate features (e.g., $(x, y)$ positions) to the input tensors channel dimension.

`neojax.nn.positional_embedding.GridEmbeddingNd` ¤

Regular grid embedding for n-dim signals.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of channels of the input signal. The embedding appends one coordinate channel per spatial dimension.	required
`ndim`	`int`	Number of spatial dimensions.	required
`grid_boundaries`	`tuple[tuple[int, int], ...] \| None`	Boundaries of regular grid per dimension ((low, high), ...) or None. Default is None which is a ((0, 1), ...) bounded grid.	`None`

Internal Attributes

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

in_channels (int): Number of channels of the input signal.
grid_boundaries (tuple[tuple[int, int], ...]): Boundaries of regular grid.

`n_extra_channels` `property` ¤

Number of channels appended to the input signal.

`call(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...']` ¤

Generates n-dim regular grid and appends it to input signal.

Parameters:

Name	Type	Description	Default
`x`	`Inexact[Array, 'in_c ...']`	Input signal, shaped (channels, d1, ..., dN).	required

Returns:

Type	Description
`Inexact[Array, 'out_c ...']`	Concatenated signal and grid array (along channel dim), shaped (in_channels + ndim, d1, ..., dN).

Raises:

Type	Description
`ValueError`	If the number of input channels does not match `in_channels`.

`neojax.nn.positional_embedding.SinusoidalEmbedding` ¤

Sinusoidal embedding of point coordinates.

Lifts low-dimensional coordinates to a richer spectral representation by encoding each coordinate as sines and cosines at multiple frequencies. This helps coordinate-based networks, like the kernels of GNOBlock, to resolve fine-scale variations.

For coordinate $p$ and frequency level $k \in \{0, \dots, L-1\}$ the embedding contains the pair $(\sin(\omega_k p), \cos(\omega_k p))$ with

"transformer": $\omega_k = M^{-2k/L}$ with $M$ = max_positions,
"nerf": $\omega_k = 2^k \pi$.

The pairs are laid out per coordinate and then per frequency, so the embedding of d coordinates has 2 * L * d channels.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of coordinates per point, i.e. the spatial dimension.	required
`num_frequencies`	`int`	Number of frequency levels $L$.	required
`embedding_type`	`Literal['transformer', 'nerf']`	Type of embedding, `"transformer"` or `"nerf"`. Default is `"transformer"`.	`'transformer'`
`max_positions`	`int`	Maximum number of positions $M$ of the `"transformer"` embedding. Not used for `"nerf"`. Default is 10000.	`10000`

Raises:

Type	Description
`ValueError`	If `in_channels` or `num_frequencies` are not positive integers.
`ValueError`	If `embedding_type` is invalid.
`ValueError`	If `max_positions` is not a positive number.

Internal Attributes

in_channels (int): Number of coordinates per point.
num_frequencies (int): Number of frequency levels.
embedding_type (Literal["transformer", "nerf"]): Type of embedding.
max_positions (int): Maximum number of positions of the "transformer" embedding.

Aliasing

Make sure the highest frequency resolves the discretization. E.g. for "nerf", $2^{L-1} < N/2$ for $N$ sample points.

Cite

Attention Is All You Need

@inproceedings{vaswani2017attention,
    title={Attention is all you need},
    author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and
        Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and
        Kaiser, {\L}ukasz and Polosukhin, Illia},
    booktitle={Advances in Neural Information Processing Systems},
    year={2017}
}

NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis

@inproceedings{mildenhall2020nerf,
    title={NeRF: Representing Scenes as Neural Radiance Fields
        for View Synthesis},
    author={Mildenhall, Ben and Srinivasan, Pratul P and Tancik, Matthew
        and Barron, Jonathan T and Ramamoorthi, Ravi and Ng, Ren},
    booktitle={European Conference on Computer Vision},
    year={2020}
}

`out_channels` `property` ¤

Number of channels of the embedding.

`call(x: Float[Array, '... in_c']) -> Float[Array, '... out_c']` ¤

Embeds the coordinates.

Note that the coordinates are along the last axis, which is the convention for point coordinates (as opposed to channel-first signals).

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, '... in_c']`	Point coordinates shaped (..., in_channels), e.g. (n_points, in_channels).	required

Returns:

Type	Description
`Float[Array, '... out_c']`	The embedded coordinates shaped (..., 2 * num_frequencies * in_channels).

Raises:

Type	Description
`ValueError`	If the size of the last axis does not match `in_channels`.

`neojax.nn.positional_embedding.ReferenceGridEmbeddingNd` ¤

Reference-grid distance embedding for n-dim coordinates.

Embeds every point $x \in \mathbb{R}^d$ by its Euclidean distances to the nodes $\{r_j\}_{j=1}^{\text{ref}^d}$ of a small, fixed, regular reference grid (a.k.a. "unified position" embedding in Transolver), i.e. $$ e_j(x) = \lVert x - r_j \rVert_2, \quad j = 1, \dots, \text{ref}^d. $$ The embedding is well-defined for structured grids and unstructured meshes alike.

Parameters:

Name	Type	Description	Default
`ndim`	`int`	Number of spatial dimensions (coordinate dimension).	required
`ref`	`int`	Number of reference grid nodes per dimension. The embedding has `ref ** ndim` features. Default is 8.	`8`
`grid_boundaries`	`tuple[tuple[int \| float, int \| float], ...] \| None`	Boundaries of the reference grid per dimension ((low, high), ...) or None. Default is None which is a ((0, 1), ...) bounded grid.	`None`

Raises:

Type	Description
`ValueError`	If `ndim` or `ref` is not positive.
`ValueError`	If the number of grid boundaries doesn't match `ndim`.

Internal Attributes

ndim (int): Number of spatial dimensions.
ref (int): Number of reference grid nodes per dimension.
grid_boundaries (tuple[tuple[int | float, int | float], ...]): Boundaries of the reference grid.

`n_features` `property` ¤

Number of embedding features per point (ref ** ndim).

`regular_grid(resolutions: tuple[int, ...]) -> Float[Array, 'd ...']` ¤

Creates a n-dim regular grid of coordinates on the grid boundaries.

Parameters:

Name	Type	Description	Default
`resolutions`	`tuple[int, ...]`	Dimensions (d1, ..., dN) without coordinate axis.	required

Returns:

Type	Description
`Float[Array, 'd ...']`	Regular grid of coordinates (d, d1, ..., dN).

Raises:

Type	Description
`ValueError`	If the number of resolutions doesn't match `ndim`.

`call(coords: Float[Array, 'd ...']) -> Float[Array, 'g ...']` ¤

Embeds coordinates by their distances to the reference grid nodes.

Parameters:

Name	Type	Description	Default
`coords`	`Float[Array, 'd ...']`	Coordinates shaped (d, d1, ..., dN) or (d, N).	required

Returns:

Type	Description
`Float[Array, 'g ...']`	Distance features shaped (ref d, d1, ..., dN) or (ref d, N).

Raises:

Type	Description
`ValueError`	If the size of the leading axis doesn't match `ndim`.

Positional Embeddings