Geom-Seg VecSet: Latent 3D Segmentation

• Geom-Seg VecSet is a unified latent representation that encodes both geometry and part-level segmentation for fine-grained 3D object generation.
• It employs a transformer-based encoder with cross-attention and specialized decoders to predict geometry, segmentation, and latent anchor positions in one framework.
• The approach integrates dual-space latent diffusion to balance global object structure with local part details, eliminating the need for external segmentation models.

Geom-Seg VecSet is a unified latent representation for 3D point-clouds and part-level segmentation, proposed for controllable, decomposable 3D object generation in the context of latent diffusion frameworks. The method encodes both geometry and segmentation structure into a compact set of latent vectors, enabling fine-grained, promptable segmentation and shape generation without requiring external segmenters or separated part-supervision. Geom-Seg VecSet serves as the core representational and interface component of the UniPart framework for part-level 3D generation (He et al., 10 Dec 2025).

1. Formal Definition and Mathematical Formulation

For an object mesh $\mathcal{O}$ with $N$ parts, Geom-Seg VecSet begins by uniformly sampling $C$ points from $\mathcal{O}$, assembling the set

$$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$

where $x_k \in \mathbb{R}^3$ is position, $n_k \in \mathbb{R}^3$ is the normal, and $s_k$ is a one-hot encoding for part label ($1,\ldots,N$). A transformer-based encoder $\mathcal{E}$ maps $N$0 to $N$1, where $N$2 and $N$3 is the number of latent tokens. The joint set $N$4 constitutes a Geom-Seg VecSet. Each $N$5 reflects both geometric properties and local part membership.

Mathematically, the framework supports three decoding tasks:

• Geometry decoding for query point $N$6:

$N$7

where $N$8 is the ground-truth implicit field.

• Segmentation decoding (promptable) for token index $N$9:

$C$0

assigning part labels to latent tokens.

• Latent anchor position decoding (auxiliary):

$C$1

estimating anchor positions for each $C$2.

2. Encoder/Decoder Architecture

Encoder ($C$3)

• Input: $C$4.
• Architecture:
  • Initial cross-attention layer with $C$5 learnable queries ($C$6) attending over the point set.
  • $C$7 stacked blocks of multi-head self-attention (with $C$8 heads, $C$9 channel width) and MLPs, each followed by layer normalization and residual connections.
• Output: Geom-Seg VecSet $\mathcal{O}$0.

Decoders

• Geometry Decoder ($\mathcal{O}$1): 8-layer MLP (hidden width 512, ReLU) with cross-attention mechanisms between latents and spatial query points.
• Segmentation Decoder ($\mathcal{O}$2): Based on promptable SAM2 segmentation, employs a transformer head taking (token, index embedding), outputs logits over part labels.
• Position Decoder ($\mathcal{O}$3): 3-layer MLP, trained to regress anchor point via $\mathcal{O}$4 loss.

3. Training Objectives and Procedures

Training employs a VAE framework, with the total objective: $\mathcal{O}$5

• $\mathcal{O}$6: Squared error over predicted and ground-truth implicit fields.
• $\mathcal{O}$7: Cross-entropy segmentation loss, following SAM2.
• $\mathcal{O}$8: KL divergence term for latent regularization; $\mathcal{O}$9 controls balance.

Pretraining typically occurs on geometry-only data; fine-tuning incorporates segmentation supervision. The design encourages the emergence of part-level structure during joint geometry encoding without external part annotations for each training step.

4. Integration into Two-Stage Latent Diffusion Pipeline

Geom-Seg VecSet enables latent-level control in the UniPart two-stage diffusion process:

Stage 1: Whole-object diffusion and latent segmentation

• Uses a DiT (Diffusion Transformer) backbone with rectified flow matching.
• Latent trajectory: $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$0, $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$1, $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$2.
• Training minimizes flow-matching loss conditioned on input image.
• After diffusion, segmentation tokens are extracted via frozen $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$3 and $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$4 for part assignment, then post-processed for mask generation.

Stage 2: Part-level diffusion with dual-space conditioning

• Each part $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$5 obtains dual latents in global coordinate space (GCS) and normalized canonical space (NCS):

$$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$6

• DiT $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$7 predicts part flows, conditioned on $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$8.
• Each transformer block fuses local (per-space) and global (cross-space) attention among part latents.

A dual-space approach enforces both global composition and localized detail preservation at the part level.

5. Sampling and Mesh Decoding Workflow

Given an input image $$P = \left\{ (x_k, n_k, s_k) \in \mathbb{R}^7 \right\}_{k=1}^C$$9, the generative pipeline proceeds as:

1. Apply whole-object latent diffusion to obtain $x_k \in \mathbb{R}^3$0.
2. Perform latent segmentation, resulting in $x_k \in \mathbb{R}^3$1 sets $x_k \in \mathbb{R}^3$2.
3. For each part $x_k \in \mathbb{R}^3$3:
   • Initialize noise $x_k \in \mathbb{R}^3$4.
   • Denoise via part-level DiT to yield $x_k \in \mathbb{R}^3$5, $x_k \in \mathbb{R}^3$6.
   • Decode implicit field mesh: $x_k \in \mathbb{R}^3$7 and $x_k \in \mathbb{R}^3$8.
   • Compute rigid transform $x_k \in \mathbb{R}^3$9 in GCS, then apply to NCS mesh.
   • Compose all part meshes into the final 3D object.

This approach avoids repeated lossy marching cubes re-encoding, increases consistency of part-geometry mapping, and provides direct control over part-level generation at the latent stage.

6. Properties, Benefits, and Limitations

Benefits

• Emergent segmentation: Part-awareness is learned “for free” through geometry-centric objectives.
• Fine-grained granularity: Latent segmentation enables precise control and adapts to varying object complexity.
• No external segmenter: The model does not require separate part-annotated data or models at inference.
• Dual-space diffusion: Maintains coherence and fidelity for both holistic object shape and localized part attributes.
• Efficient training: Domain-aligned latent conditioning enhances diffusion efficiency and entangles part–whole relations.

Limitations and Open Challenges

• Currently limited to single-object synthesis; full-scene composition is not addressed.
• Relies on point-set latent format; extension to voxel or mesh-graph VecSets (e.g., SLAT) would generalize application.
• The segmentation head is fixed-size ($n_k \in \mathbb{R}^3$0 tokens); scalability to extremely complex objects may require dynamic token allocation.
• Segmentation accuracy for thin or fine parts is not explicitly regularized; some topological degradation may occur.
• Absence of adversarial losses or explicit topological priors; part consistency relies solely on geometric and segmentation objectives.

7. Comparative Context and Applications

Geom-Seg VecSet marks a shift from prior approaches—either implicit, non-controllable part segmentation, or reliance on external semantic masks—toward a unified, latent-based geometry-segmentation paradigm with promptable control. Applications span part-level generative design, interactive 3D editing, robotic manipulation where part decomposition is crucial, and any downstream task benefiting from decomposable, high-fidelity 3D synthesis with semantic control (He et al., 10 Dec 2025).

A plausible implication is that the VecSet format, through its compositionality and promptability, offers a generalizable scaffold for integrating discrete part structure into continuous generative models, providing a pathway toward scene-level compositional synthesis and broader generative 3D understanding.