TORA: Topological Representation Alignment for 3D Shape Assembly

Abstract: Flow-matching methods for 3D shape assembly learn point-wise velocity fields that transport parts toward assembled configurations, yet they receive no explicit guidance about which cross-part interactions should drive the motion. We introduce TORA, a topology-first representation alignment framework that distills relational structure from a frozen pretrained 3D encoder into the flow-matching backbone during training. We first realize this via simple instantiation, token-wise cosine matching, which injects the learned geometric descriptors from the teacher representation. We then extend to employ a Centered Kernel Alignment (CKA) loss to match the similarity structure between student and teacher representations for enhanced topological alignment. Through systematic probing of diverse 3D encoders, we show that geometry- and contact-centric teacher properties, not semantic classification ability, govern alignment effectiveness, and that alignment is most beneficial at later transformer layers where spatial structure naturally emerges. TORA introduces zero inference overhead while yielding two consistent benefits: faster convergence (up to 6.9$\times$) and improved accuracy in-distribution, along with greater robustness under domain shift. Experiments on five benchmarks spanning geometric, semantic, and inter-object assembly demonstrate state-of-the-art performance, with particularly pronounced gains in zero-shot transfer to unseen real-world and synthetic datasets. Project page: https://nahyuklee.github.io/tora.

• The paper demonstrates that explicit teacher-student topological alignment via CKA improves assembly accuracy and reduces rotation/translation errors in flow-based models.
• It leverages geometry-centric teacher embeddings from ViT-based Uni3D models to guide student transformers, achieving up to 6.9× faster convergence.
• Extensive experiments validate TORA’s effectiveness across diverse datasets, showing robust zero-shot generalization and scalability in complex 3D assemblies.

Topological Representation Alignment for 3D Shape Assembly

Introduction

3D shape assembly from unposed part point clouds remains a central task in geometric reasoning with broad implications in computational archaeology, graphics, and robotics. The paper "TORA: Topological Representation Alignment for 3D Shape Assembly" (2604.04050) introduces TORA, a framework that leverages teacher-student topological representation alignment to systematically integrate pretrained geometric priors into flow-matching-based 3D shape assembly models. The proposed method emphasizes relational knowledge distillation using Centered Kernel Alignment (CKA), demonstrating improved performance and convergence, particularly in scenarios requiring robust part mating under severe ambiguity and domain distribution shifts.

Flow-Matching and the Challenge of Implicit Mating Relations

Recent 3D assembly approaches such as Rectified Point Flow (RPF) predict point-wise velocity fields that transport unposed fragments towards their target rigid transformations. While effective, these models rely on endpoint reconstruction losses and provide no explicit intermediate supervision linking which regions amongst parts should physically interact (i.e., the sparse and often ambiguous contact sets critical for assembly fidelity).

TORA introduces explicit relational inductive bias into flow-based assembly models via teacher-student representation distillation.

Figure 1: Multi-part assembly results across regimes, showing superior and consistent improvements via teacher-student distillation with geometric priors.

The TORA Framework: Topology-First Alignment

TORA extends the flow-matching backbone with a parallel alignment branch, injecting topological priors from a frozen pretrained 3D encoder ("teacher") into the student at selected transformer layers.

Figure 2: Overview of the TORA framework. TORA distills relational geometric topology from a frozen teacher into a flow-matching student via CKA alignment.

Key technical elements:

• Alignment Objectives: Alongside classical cosine and contrastive (NT-Xent) matchings, TORA advocates Gram matrix-based Centered Kernel Alignment (CKA) loss, directly matching the pairwise similarity structures ("who-is-similar-to-whom") between student and teacher representations. This objective is invariant to feature scaling and sensitive to global relational topology.
• Layer Selection: Analysis shows that spatially coherent, part-aware structure emerges primarily in the later transformer layers of the student, motivating late-stage alignment.
• Teacher Quality: Teacher representations with strong mating- or geometry-centric features, not semantic/global recognition capacity, drive downstream assembly performance.

  Figure 3: Schematic of alignment objectives—CKA preserves relational topology, in contrast to independent token-wise matching.

Teacher Model Selection and Topology-Dependent Generalization

Extensive probing using segmentation F1, classification, local-vs-distant-similarity, and part-silhouette metrics demonstrates that geometry- and contact-focused teacher embeddings (not semantic accuracy) robustly correlate with high assembly accuracy.

Figure 4: Correlation analysis demonstrating that geometry/contact-centric teacher properties predict downstream assembly accuracy; global semantics do not.

Among leading representations, ViT-based Uni3D encoders (especially Uni3D-L) provide the strongest improvements in zero-shot and in-distribution settings.

Emergence of Spatial Structure and Layer-wise Analysis

Evaluation of intermediate student representations without alignment reveals that part-aware topology, boundary contrast, and pose discrimination increase monotonically in later transformer layers. Alignment is most effective in these later stages, further resolving assembly-relevant spatial structure.

Figure 5: Spatial metrics measured across flow model layers, all increasing with depth, indicating progressive resolution of assembly-relevant structure.

Experimental Results

Strong Numerical Results

• On Breaking Bad (2–20 parts), TORA with CKA or token-wise cosine alignment increases Part Accuracy (PA) from 93.2% (RPF) to 95.7%, and reduces rotation/translation errors.
• In high-part (21–33) settings, PA improves from 62.1% (RPF) to 72.4% (CKA).
• On PartNet-Assembly (semantic), $\text{PA}$ is raised from 59.8% (RPF) to 69.1% (CKA).
• In TwoByTwo (inter-object), CKA raises PA from 65.4% to 71.5%.
• Under domain shift (unseen datasets like FRACTURA, Fantastic Breaks), TORA with CKA consistently achieves the lowest rotation and translation errors.

Zero Overhead Inference

TORA introduces no test-time computational overhead, as the teacher is used only at train time and features can be precomputed offline. Training overhead is minimal (nearly negligible with caching).

Qualitative and Robustness Analysis

TORA consistently produces assemblies with correct mating (visual coherence), even under distribution shifts. The method is robust to part/intra-object symmetry and ambiguous contact scenarios.

Figure 6: Qualitative example—TORA outputs structurally coherent assemblies matching ground-truth, outperforming RPF in ambiguous configurations.

Convergence and Optimization

Topology alignment via CKA produces substantial convergence speedup (up to $6.9\times$ relative to RPF), with both faster and higher-quality optima across all tested regimes.

Figure 7: Convergence curves showing accelerated optimization and improved final accuracy with TORA (CKA) across multiple datasets.

Implications and Future Directions

The study demonstrates that topology-first alignment—explicitly matching pairwise relations stemming from geometric pretraining—facilitates robust, fast, and accurate 3D shape assembly. The key findings include:

• Relational over Semantic Distillation: Geometry/contact-centric signals, not conventional semantic classification, are essential for effective knowledge transfer in part assembly.
• Zero-Shot Robustness: Generalization to out-of-distribution, real-world, and synthetic, fracture-heavy datasets is significantly improved via topological alignment.
• Scalability and Efficiency: The framework scales efficiently to dense point clouds and is practical for large-scale or robotics applications.

Potential directions include further scaling of attention mechanisms to handle even larger assemblies, hybridization with semantic/visual cues for joint reasoning, and integration into real-time robotic manipulation pipelines.

Conclusion

TORA establishes the efficacy of topological representation alignment—via CKA-based relational distillation from geometry-aware teachers—for robust and efficient 3D shape assembly. The framework systematically improves both in-distribution and zero-shot performance, amplifies convergence, and provides a clear prescription for teacher model selection grounded in geometric, not semantic, properties. TORA has direct relevance for assembly tasks in robotics, graphics, and scientific analysis demanding compositional spatial reasoning under ambiguity and domain shift.