Skip to content

Utonia

Utonia pretrained encoder and linear-probing segmentation model.

First page of Utonia: Toward One Encoder for All Point Clouds

2603.03283 · March 2026

Classes:

UtoniaSegmentation 

UtoniaSegmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "z",
        "z-trans",
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (3, 3, 3, 12, 3),
    encoder_channels: Sequence[int] = (
        54,
        108,
        216,
        432,
        576,
    ),
    encoder_num_heads: Sequence[int] = (3, 6, 12, 24, 32),
    encoder_patch_size: Sequence[int] = (
        1024,
        1024,
        1024,
        1024,
        1024,
    ),
    norm: Union[str, Callable] = "layer_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    pooling: str = "grid",
    stem_type: str = "linear",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
)

Bases: SegmentationModel

Utonia linear-probing segmentation model.

Linear-probe variant from Utonia: Toward One Encoder for All Point Clouds (ICML 2026). Architecturally similar to Sonata / Concerto's linear-probe head, with one key change: every attention layer adds a 3D rotary position embedding (Point3DRoPE) on top of (q, k), indexed by the real-valued metric position rather than the integer voxel grid. The position is mean-pooled at every encoder stage so each level operates at its natural scale.

Note

The default (and registered) configuration enables flash attention (use_flash_attn=True), which requires the flash-attn package and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run the forward on CPU.

Methods:

  • configure_encoder –

    Build the PointTransformerV3Encoder backbone with rotary position embeddings.

  • forward –

    Forward pass.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

num_features property 

num_features: int

Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

configure_encoder 

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone with rotary position embeddings.

forward 

forward(
    x: Tensor, pos: Tensor, pos_grid: Tensor, batch: Tensor
) -> Tensor

Forward pass.

Parameters:

  • x (Tensor) –

    Per-point features of shape \((N, C)\).

  • pos (Tensor) –

    Real-valued metric positions of shape \((N, 3)\) used by 3D RoPE.

  • pos_grid (Tensor) –

    Integer voxel-grid coordinates of shape \((N, 3)\) used by the encoder for serialization and sparse convolutions.

  • batch (Tensor) –

    Per-point batch index of shape \((N,)\).