Skip to content

serialized_attention

Serialized attention variants used by Point Transformer V3 and descendants.

Classes:

RelativePositionalEncoding 

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.

Parameters:

  • patch_size (int) –

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

  • num_heads (int) –

    Number of attention heads.

SerializedAttention 

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 

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 

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.