OmniLatent: Unified Latent for 3D Part Modeling

AI summary

Overview

Research area: Computer Vision / 3D generative modeling — specifically part-aware 3D asset generation and 3D part decomposition.

Technical level: Advanced. The paper assumes familiarity with variational autoencoders, diffusion/flow-matching transformers (MM-DiT, rectified flow), voxel-query set representations, rotary position embeddings, and 3D surface extraction via Marching Cubes.

Scope: OmniLatent is a unified latent framework that jointly encodes and generates 3D geometry, deformation, appearance, and binary surface part IDs from a single image, with adapter extensions for point-prompted segmentation/refinement and amodal part completion.

What This Paper Is About

Most image-to-3D generators produce a single monolithic surface: the object looks right, but its semantic components are never explicitly separated. The usual fix is to generate a mesh first and then run a separate segmentation model on it, which makes part quality hostage to mesh artifacts and often fails near thin structures, sharp boundaries, or texture-defined part transitions. OmniLatent instead asks whether a shared latent space can encode geometry, appearance, and part identity together, so that part decomposition is produced by the same generative pass that builds the object — and can then be edited interactively or completed where parts are occluded.

Key Contributions

1. A unified latent representation for four surface attributes. Geometry, deformation, appearance, and 6-bit binary part IDs are compressed by a shared encoder into one token space and decoded by attribute-specific heads, so part IDs are predicted on the very surface that defines shape and texture.

2. Single-pass part-aware generation from one image. A 1.9B-parameter MM-DiT trained with rectified flow generates geometry, deformation, texture, and surface part IDs together, optionally steered by a 2D segmentation mask for granularity control — and works without the mask when none is given.

3. A point-prompt adapter for interactive part control. Arbitrary 3D point prompts, each tagged with a 6-bit binary part ID, are mapped by a lightweight MLP into latent tokens appended to the denoising sequence — enabling full segmentation from sparse prompts, refinement of an existing segmentation, and subdivision of a coarse part, without a positive/negative click formulation.

4. An amodal part completion model. Given full geometry plus partial surface masks, the adapted model completes all parts into a coherent full-object decomposition using dynamic per-part token allocation, RoPE-anchored versus no-RoPE tokens, a cross-attended geometry latent, and alternating part-local / global attention.

Main Findings

• Surface part segmentation outperforms post-hoc baselines. On PartObjaverse-Tiny-wt (over 3,000 fine-grained part annotations across 200 shapes, mIoU), OmniLatent reaches 74.18 with mask guidance, versus 58.10 for P³-SAM, 52.61 for PartField, 50.02 for SAMPart3D, and 20.76 for Find3D. Per-category OmniLatent scores range from 58.94 (Buildings) to 82.68 (Animals), with Human-Shape 79.36, Daily-Used 76.98, Transportations 65.04, Plants 78.73, Food 80.72, and Electronics 70.98.

• The advantage holds on a shared surface. When every method is evaluated on OmniLatent's own generated meshes, OmniLatent stays at 74.18 while the strongest baseline P³-SAM drops to 56.01 (SAMPart3D 48.73, PartField 51.93, Find3D 19.21).

• Mask guidance helps but is optional. On the ground-truth watertight surface, OmniLatent scores 62.42 mIoU without mask guidance and 76.53 with it — a single base model supports both modes because mask conditions are randomly dropped during training.

• Progressive training transfers data-rich priors to data-scarce parts. The model is trained in stages: 900K geometry/deformation meshes, then 300K textured assets, then 20K part-annotated objects, with missing fields zeroed so the latent space stays compatible across stages.

• Binary part codes were chosen for stability. The authors argue 6-bit binary IDs (up to 2⁶ = 64 parts) are more robust under denoising than RGB-coded masks, where small continuous color deviations can split one semantic part into inconsistent IDs.

• Point-prompted and amodal completion results are not reported in the provided text. The paper's experiment section is truncated at the description of the refinement/subdivision evaluation setup; quantitative numbers for those two capabilities, and for the supporting geometry/appearance quality claims, do not appear in the supplied content.

Methodology in Plain English

Represent the object's surface as a bundle of aligned attributes. The team samples points on an object's surface and attaches several fields to each point: geometry (position, normal, and a sharp-edge flag), color, a 6-bit part code, and later a small learned offset that sharpens edges. Because all fields live on the same surface, part labels are inherently aligned with shape and texture.

Compress everything into one latent with a single encoder. A variational autoencoder with a shared encoder compresses this bundle into an ordered set of latent tokens. Decoding is cascaded: first the signed distance field is reconstructed and the surface extracted with Marching Cubes, then the resulting vertices query deformation, color, and part ID from the same tokens. Part labels are converted to per-face masks, so decoded labels never overlap on internal interfaces.

Generate the latent from an image with a diffusion transformer. A 1.9B-parameter MM-DiT trained with rectified flow learns to produce these latents conditioned on an RGB image, and optionally on a 2D segmentation mask that hints at the desired part granularity. A condition-aware rotary position embedding keeps image conditions, masks, and 3D latent tokens distinguishable. To stop the model from over-trusting 2D masks — which may be incomplete or view-inconsistent — masks are augmented during training with random dropping, merging, splitting, component decomposition, internal noise, and erosion.

Train in a curriculum so scarce part labels ride on abundant geometry data. Geometry and deformation first, texture second, part annotations last. This is the paper's central bet: parts are learned inside a latent space already shaped by strong geometric and visual priors.

Add a prompt adapter for interaction. Each 3D click is a coordinate plus a binary part ID; a small MLP turns these into extra tokens that join the denoising sequence. Because part IDs themselves define the groups, one prompt set can specify several parts at once, letting the model infer a full segmentation, fix a coarse one, or split a region into finer parts. This adapter is fine-tuned from the base generator rather than trained from scratch.

Adapt for amodal completion with a token-budget scheme. Parts differ wildly in visible area, so tokens are allocated by visible voxel count, clamped to 0.5–1.5 times the average per-part budget. Each part gets RoPE-anchored tokens pinned to observed surface evidence plus no-RoPE tokens free to represent unseen regions; the loss on completion tokens is upweighted by 5, and the no-RoPE ratio is sampled from 20%–40% in training and fixed at 30% at inference. Geometry enters via cross-attention, and attention alternates between within-part and full-object passes.

Practical figures. VAE training takes roughly 3 days and DiT training roughly 7 days; the DiT uses a sequence length of 12,248 tokens; inference produces a full part-aware asset in about 30 seconds at 8192 latent tokens.

Why This Matters

For research: The paper reframes part decomposition from a post-hoc labeling problem into a property of the generative representation itself. It offers a concrete recipe — shared encoder, cascaded surface decoding, curriculum training, binary part codes — for letting data-rich modalities (geometry, texture) bootstrap a data-scarce one (part annotations), and it bundles automatic, interactive, and amodal completion behaviors into one pretrained backbone rather than three separate models.

Real-world applications:

• Game and film production: assets arrive already split into addressable parts, so materials, rigs, and destructible pieces can be assigned per component instead of hand-cut.
• Digital twins and robotics simulation: functionally separable components (doors, wheels, lids, handles) are needed for physical interaction and articulated simulation.
• Interactive design tools: click-a-few-points segmentation and part subdivision fit naturally into authoring UIs where artists refine an automatically generated decomposition.
• Physical fabrication and 3D printing: printable assemblies require parts split along sensible seams; binary part IDs up to 64 per object are sufficient for the "5 to 30 parts" the authors report as typical for their annotated set.

Industry relevance: The work comes from a university–industry collaboration (University of Hong Kong, Chinese University of Hong Kong, Nanjing University, Tencent Technology) and targets editable asset pipelines — the point where generative 3D currently loses most of its value, since a beautiful but monolithic mesh is hard to integrate into production tooling.

Future Directions

• Part identity is object-local, not semantic. Training assigns each object's parts random distinct 6-bit codes, so "wheel" on one object and "wheel" on another share no correspondence. A cross-object semantic taxonomy would let prompts or edits transfer meaningfully between assets.

• Granularity control still leans on external cues. The mask-free mode gives autonomous segmentation, but controlling granularity otherwise requires a 2D mask or point prompts. Learning to specify part granularity directly in latent space — without any 2D supervision — is an open step.

• Quantitative validation of the extensions is incomplete in the supplied text. Point-prompted segmentation, refinement, subdivision, and amodal completion are described qualitatively and evaluated in the paper, but the numbers are not in the provided content; whether these adapters reach the same accuracy bar as the base segmentation remains to be reported.

• Part fields live on the outer surface. The Part-ID output contains no overlapping internal interfaces, and complete internal or occluded part geometry is recovered only by the separate completion model — leaving room for a single model that outputs full volumetric part structure directly.

Target Audience

Researchers and graduate students working on 3D generative models, 3D segmentation, or multi-modal latent representations; ML engineers building asset-generation pipelines for games, AR/VR, or simulation; and practitioners who need editable, part-structured 3D output rather than a single textured mesh. Readers should be comfortable with diffusion/flow-matching transformers and implicit surface representations — the paper's exposition of the multi-attribute VAE and the prompt/completion adapters is compact and assumes that background.

Authors’ abstract

3D part decomposition is a central capability for editable asset creation, yet existing workflows typically segment a completed mesh as a post-processing step, making the result sensitive to geometric artifacts, texture ambiguity, and limited part-level annotations. We present OmniLatent, a unified latent modeling framework for 3D part decomposition. Instead of predicting part labels in isolation, OmniLatent learns a shared latent space that jointly encodes geometry,…