Skip to content

serialized_attention

Serialized attention variants used by Point Transformer V3 and descendants.

Classes:

RelativePositionalEncoding [source] 

RelativePositionalEncoding(patch_size: int, num_heads: int)

Bases: Module

Table-based relative position bias added to the attention logits of a serialized patch.

Each axis of the clamped relative grid coordinates indexes its own slice of a shared table, and the three per-head biases are summed. Clamping leaves \(\text{rpe\_num}^3\) distinct offsets, so the three rows are summed once per combination into a lookup table and the bias is a single gather. Summing them per point pair instead would hold \((\text{pairs}, 3)\) offsets and \((\text{pairs}, 3, H)\) gathered rows, which on a large cloud is the peak allocation of the whole model.

Parameters:

  • patch_size (int) –

    Number of points attending to each other, setting the clamping boundary.

  • num_heads (int) –

    Number of attention heads.

Shape
  • Input: \((P, K, 3)\) integer grid coordinates, \(K\) points per patch.
  • Output: \((P, H, K, K)\) bias, aligned with the attention logits.
Example 
rpe = RelativePositionalEncoding(patch_size=8, num_heads=4)
bias = rpe(torch.randint(0, 16, (2, 8, 3)))
bias.shape  # (2, 4, 8, 8)

SerializedAttention [source] 

SerializedAttention(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Vanilla serialized attention from Point Transformer V3.

No positional information is added inside attention itself: relative structure comes from the conditional position embedding (CPE) applied around each block.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

SerializedAttentionRPE [source] 

SerializedAttentionRPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Serialized attention with the relative position bias from PT-V3.

Adds a learned per-head bias indexed by the integer voxel-grid offset between query and key inside each patch. Flash Attention does not support arbitrary attention biases, so this variant always uses the manual softmax path.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

SerializedAttentionRoPE [source] 

SerializedAttentionRoPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
    rope_base: float = 10.0,
)

Bases: Module

Serialized attention with 3D rotary position embedding from Utonia.

Rotates \(Q\), \(K\) via Point3DRoPE using the real-valued metric position of each token. Flash Attention is supported and uses bfloat16 (matching upstream Utonia's reference implementation).

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.