Attention & Transformer

Three hyperbolic attention variants from the Hypformer paper (Yang et al. 2025, Section 4.3). All operate on hyperboloid points and support independent curvatures for input (c_in), attention computation (c_attn), and output (c_out), plus causal (autoregressive) masking via causal=True. For complexity/use-case trade-offs see the NN Layers guide.

The supporting utilities (spatial_to_hyperboloid, lorentz_midpoint, focus_transform) are documented under Primitives. The Hypformer transformation component HTCLinear (used for MLP sublayers) is on the Linear page; HRC normalization is on the Normalization page.

Attention modules

hyperbolix.nn_layers.HyperbolicLinearAttention

HyperbolicLinearAttention(
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    power: float = 2.0,
    init_bound: float = 0.02,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
    rngs: Rngs,
)

Bases: _HyperbolicAttentionBase

Hyperbolic linear attention with focus function (Eq 14-19).

The paper's main contribution: O(N) attention using the kernel trick in the spatial domain of the hyperboloid. Focus function φ sharpens query and key.

Parameters:

Name	Type	Description	Default
in_features	int	Ambient input dimension (d_in + 1).	required
out_features	int	Spatial output dimension per head.	required
num_heads	int	Number of attention heads (default: 1).	1
power	float	Focus function sharpening exponent (default: 2.0).	2.0
init_bound	float	Uniform init bound for weights (default: 0.02).	0.02
eps	float	Numerical stability floor (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by the manifold's dtype.	float32
rngs	Rngs	Random number generators.	required

Source code in hyperbolix/nn_layers/hyperboloid_attention.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    power: float = 2.0,
    init_bound: float = 0.02,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
    rngs: nnx.Rngs,
):
    super().__init__(
        in_features,
        out_features,
        num_heads=num_heads,
        init_bound=init_bound,
        eps=eps,
        param_dtype=param_dtype,
        rngs=rngs,
    )
    self.power = power
    self.temperature = nnx.Param(jnp.array(1.0, dtype=param_dtype))
    # Spatial residual projection ψ: D → D (shared across heads)
    self.residual_proj = nnx.Linear(out_features, out_features, param_dtype=param_dtype, rngs=rngs)

hyperbolix.nn_layers.HyperbolicSoftmaxAttention

HyperbolicSoftmaxAttention(
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    init_bound: float = 0.02,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
    rngs: Rngs,
)

Bases: _HyperbolicAttentionBase

Hyperbolic softmax attention in the spatial domain.

Standard scaled dot-product attention applied to spatial components of query, key, value, followed by the same HRC pipeline (residual + time calibration) as the linear variant.

Parameters:

Name	Type	Description	Default
in_features	int	Ambient input dimension (d_in + 1).	required
out_features	int	Spatial output dimension per head.	required
num_heads	int	Number of attention heads (default: 1).	1
init_bound	float	Uniform init bound for weights (default: 0.02).	0.02
eps	float	Numerical stability floor (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by the manifold's dtype.	float32
rngs	Rngs	Random number generators.	required

Source code in hyperbolix/nn_layers/hyperboloid_attention.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    init_bound: float = 0.02,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
    rngs: nnx.Rngs,
):
    super().__init__(
        in_features,
        out_features,
        num_heads=num_heads,
        init_bound=init_bound,
        eps=eps,
        param_dtype=param_dtype,
        rngs=rngs,
    )
    # Spatial residual projection ψ: D → D (shared across heads)
    self.residual_proj = nnx.Linear(out_features, out_features, param_dtype=param_dtype, rngs=rngs)

hyperbolix.nn_layers.HyperbolicFullAttention

HyperbolicFullAttention(
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    init_bound: float = 0.02,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
    rngs: Rngs,
)

Bases: _HyperbolicAttentionBase

Full Lorentzian attention with midpoint aggregation.

Uses the Lorentzian inner product for similarity and weighted Lorentzian midpoint for aggregation — operating on full hyperboloid points throughout.

Parameters:

Name	Type	Description	Default
in_features	int	Ambient input dimension (d_in + 1).	required
out_features	int	Spatial output dimension per head.	required
num_heads	int	Number of attention heads (default: 1).	1
init_bound	float	Uniform init bound for weights (default: 0.02).	0.02
eps	float	Numerical stability floor (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by the manifold's dtype.	float32
rngs	Rngs	Random number generators.	required

Source code in hyperbolix/nn_layers/hyperboloid_attention.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    num_heads: int = 1,
    init_bound: float = 0.02,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
    rngs: nnx.Rngs,
):
    super().__init__(
        in_features,
        out_features,
        num_heads=num_heads,
        init_bound=init_bound,
        eps=eps,
        param_dtype=param_dtype,
        rngs=rngs,
    )
    self.scale = nnx.Param(jnp.array(1.0, dtype=param_dtype))
    self.attn_bias = nnx.Param(jnp.array(0.0, dtype=param_dtype))

Causal (autoregressive) masking

All three variants support causal=True: position n attends only to positions m ≤ n, required for autoregressive tasks. Masking is JIT-compatible.

How causal masking is implemented

• HyperbolicSoftmaxAttention / HyperbolicFullAttention: lower-triangular -inf mask on the score matrix before softmax — O(N²) in both modes.
• HyperbolicLinearAttention: a cumulative-sum recurrence (jax.lax.scan, Katharopoulos et al. 2020): S_i = Σ_{j≤i} φ(K_j) V_jᵀ, O(1) per step → O(N) total, well-suited to long autoregressive sequences.

Example

import jax, jax.numpy as jnp
from flax import nnx
from hyperbolix.manifolds import Hyperboloid
from hyperbolix.nn_layers import HyperbolicLinearAttention

hyperboloid = Hyperboloid()
B, N, A_in, D_out = 4, 8, 9, 8  # 8-dim spatial + 1 time
spatial = jax.random.normal(jax.random.PRNGKey(0), (B, N, A_in - 1)) * 0.1
time = jnp.sqrt(jnp.sum(spatial**2, axis=-1, keepdims=True) + 1.0)
x = jnp.concatenate([time, spatial], axis=-1)  # (B, N, A_in) on the hyperboloid

attn = HyperbolicLinearAttention(in_features=A_in, out_features=D_out, num_heads=2, power=2.0, rngs=nnx.Rngs(0))
y = attn(x, c_in=1.0, c_attn=1.0, c_out=1.0)          # bidirectional
y_causal = attn(x, c_in=1.0, c_attn=1.0, c_out=1.0, causal=True)
print(y.shape)  # (4, 8, 9) — D_out spatial + 1 time

See the NN Layers guide (Pattern 3) for a full hyperbolic transformer block.