TORA: Topological Representation Alignment for 3D Shape Assembly¶

Conference: ECCV 2026
arXiv: 2604.04050
Code: https://nahyuklee.github.io/tora (Project Page)
Area: 3D Vision
Keywords: 3D Shape Assembly, Representation Alignment, Knowledge Distillation, Flow Matching, Relational Topology

TL;DR¶

During the training of a flow-matching-based 3D assembly model, TORA additionally aligns its feature representation with the "who-is-similar-to-whom" relational structure of a frozen, pre-trained 3D encoder (matching Gram matrices via CKA loss). This distills geometric interaction priors into the network, accelerating convergence by up to \(6.9\times\), significantly enhancing cross-domain robustness, and incurring zero extra cost during inference.

Background & Motivation¶

Reassembling a set of part point clouds with unknown poses into a complete object is a fundamental geometric reasoning task in scenarios such as archaeological artifact restoration, computer graphics, and robotic peg-in-hole insertion. It spans a broad spectrum: semantic assembly (positioning functional parts like chair legs and backrests in reasonable locations), geometric assembly (reassembling fragments based on boundary complementarity), and cross-object assembly (matching parts from two different objects under cross-object constraints). These subtasks are dominated by different cues, but share a common bottleneck—how to robustly discover "pairing relationships" (i.e., which regions should contact and move together) under symmetry ambiguity and domain shifts. Recently, flow-matching-based methods have made significant progress in this direction, with Rectified Point Flow (RPF) emerging as the current SOTA. RPF learns a point-wise velocity field to transport noisy point clouds along straight lines to their assembled positions, then analytical Procrustes/SVD is used to solve for the rigid transformation of each part.

The issue is that RPF only supervises the entire flow using the "endpoint reconstruction loss"—the model must implicitly infer pairing relationships solely from the global signal of the final assembled geometry. However, critical cues often reside in sparse contact regions, which are frequently ambiguous under symmetry, and the loss function provides no explicit guidance on "which cross-part interactions should drive the motion." Lacking this explicit intermediate guidance, the model easily becomes brittle under distribution shifts. Meanwhile, RPF's own encoder is only pre-trained on a binary overlap prediction task, offering a rather narrow geometric signal. In contrast, academic research has produced a batch of pre-trained 3D point cloud encoders that have learned rich spatial features on large-scale shape data, which could be leveraged.

Addressing this gap, the authors reformulate training into a "teacher-student distillation" framework—inspired by REPA in 2D generation, but adapted for 3D point flow assembly. First, they find that the most naive approach (token-wise cosine matching, forcing each point token to match the teacher's feature content) is already a strong alignment strategy, injecting the geometric descriptors learned by the teacher. However, this independent token-wise matching treats each token as an isolated individual and does not explicitly constrain the relational topology between points—yet the success of assembly heavily depends on the interaction structure between parts. Core Idea: Instead of aligning feature vectors token-by-token, explicitly align the pairwise relational structure (topology) of "who is similar to whom" in the teacher's representation—using CKA loss to match the centered Gram matrices of the student and teacher. This distills the relational topology into the flow backbone, while only aligning at late transformer layers and selecting a teacher with strong geometric/contact-awareness.

Method¶

Overall Architecture¶

TORA is built upon the RPF architecture, appending a "topological representation alignment branch" alongside its flow matcher. This branch exists only during training and is completely discarded during inference, resulting in zero overhead at test time.

Specifically, the input consists of \(K\) part point clouds with unknown poses \(\{\mathbf{P}_k\}\). A frozen overlap-aware encoder first extracts point-wise conditional features \(\mathbf{c}\). A DiT-style transformer \(V_\theta\) takes the noisy point positions \(\mathbf{X}(t)\) and conditions \(\mathbf{c}\) as input, generating intermediate representations \(\mathbf{h}^{(l)}\) layer by layer, and ultimately predicting the point-wise velocity field supervised by a conditional flow matching loss \(\mathcal{L}_{\text{CFM}}\). The alignment branch selects the intermediate features of a specific layer \(l^*\), projects them into the teacher's feature space via a lightweight projection head, and aligns them with a frozen 3D teacher encoder (target representations extracted from clean, noise-free point clouds). The total loss is \(\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{CFM}}+\lambda\,\mathcal{L}_{\text{align}}\). The authors compare three alignment targets ordered along the axis of "how explicitly the relational structure is transferred": token-wise cosine distance, contrastive NT-Xent, and the default CKA relational alignment in TORA. Around this branch, there are two key engineering choices—using the right teacher and aligning at the right layer—which determine whether the alignment is effective.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["K Unposed<br/>Part Point Clouds"] --> B["Flow Matcher (DiT)<br/>Layer-wise Intermediate Features h^(l)"]
    B -->|"Extract Layer l*"| C["Projection Head φ<br/>Map to Teacher Space"]
    D["Clean Point Clouds"] --> E["Frozen 3D Teacher<br/>Uni3D-L"]
    C --> F["CKA Relational Alignment<br/>Match Centered Gram Matrices"]
    E --> F
    G["What Makes a Good Teacher?<br/>Geometric/Contact Probing"] -.Select.-> E
    H["Where to Align?<br/>Late Layers Spatial Structure Emergence"] -.Select.-> C
    B --> I["Predict Point-wise Velocity Field<br/>Procrustes for Pose Solving"]
    F -.Training-only Regularization.-> B

Key Designs¶

1. CKA Relational Topological Alignment: Aligning "who is similar to whom" instead of the features themselves

RPF only utilizes an endpoint loss, requiring the model to implicitly guess pairing relationships. Although naive token-wise cosine matching can inject the teacher's geometric descriptors, it aligns each point token as an independent individual and does not constrain the relational structure between points—which is crucial for assembly. TORA's main objective is changed to aligning the pairwise similarity structure: first, compute an \(n\times n\) Gram matrix (recording pairwise dot products of sampled tokens) from the student features \(\hat{\mathbf{h}}\) and teacher features \(\mathbf{y}\) respectively. After centering, the normalized Frobenius inner product (i.e., CKA) is used to measure how similar the two structures are, and the negative of this value is taken as the loss:

\[\mathcal{L}_{\text{CKA}}(\hat{\mathbf{h}},\mathbf{y})=1-\frac{\langle\tilde{\mathbf{G}}_S,\tilde{\mathbf{G}}_T\rangle_F}{\|\tilde{\mathbf{G}}_S\|_F\,\|\tilde{\mathbf{G}}_T\|_F}\]

The fundamental difference from token-wise matching is: CKA explicitly preserves the "who-is-similar-to-whom" relationship between all token pairs, and is naturally invariant to isotropic scaling and orthogonal transformations of the feature space. In other words, it does not force the student to replicate the teacher's feature vectors exactly, but only requires the relational topology to match. This perfectly matches the requirements of the assembly task: what matters is not the absolute feature of a point, but its relative relationship with potential mating points. Experiments confirm this: on cross-object (TwoByTwo) and semantic assembly (PartNet) settings where interaction structures are strong and domain shifts exist, CKA's gain is most pronounced. Combining CKA with token-wise targets actually degrades performance (see Ablation), showing that the relational signal itself is sufficient, and token-wise targets introduce conflicting gradients that dilute alignment.

2. Random Subsampling of Gram Matrix: Making \(O(N^2)\) relational alignment computationally feasible

Aligning pairwise structures naturally requires computing an \(N\times N\) Gram matrix. Since the total number of points \(N\) in a multi-part object is large, computing the full Gram matrix introduces prohibitive computational overhead. TORA's solution is to randomly and uniformly sample \(n\) token indices (\(n\ll N\)) in each batch, and only construct an \(n\times n\) Gram matrix on these \(n\) points for CKA. The intuition behind this is that even sampling a moderate number of random tokens is sufficient to provide a faithful estimation of the overall Gram structure. Ablation studies confirm that the convergence curves for \(n\in\{256, 512, 1024, 2048\}\) almost overlap, demonstrating high robustness to the sample size. The authors choose \(n=1{,}024\) as a conservative default. Because the main overhead comes from teacher feature lookup and the projection head's forward pass—rather than the Gram calculation itself—and with offline caching of teacher features, CKA only adds about \(3\text{ ms}\) (\(+2.8\%\)) per step and \(0.18\text{ GB}\) (\(+5.5\%\)) of peak VRAM, which is virtually free.

3. Teacher Selection: Look at geometric/contact probes, not classification accuracy

The alignment performance highly depends on the choice of the teacher, but it is not clear a priori which teacher attributes are beneficial for assembly. The authors systematically probe six pre-trained 3D encoders, measuring four lightweight metrics on frozen features: object classification accuracy (proxy for global semantics), mating surface segmentation F1 (proxy for contact-aware/interaction geometry), local-to-distant similarity (LDS, for geometric sensitivity), and part contour score (part-level geometric prior). They then correlate each probing score with the downstream assembly Part Accuracy obtained by aligning with that teacher. The conclusion is highly counter-intuitive: semantic classification accuracy has almost zero predictive power for assembly transfer (Pearson \(r\approx-0.04\)), while mating surface segmentation F1 (\(r=+0.94\), still robust after controlling for parameters) and part contour are strongly correlated, and LDS is moderately correlated. In other words, the criterion for a good teacher is "how much interaction geometry (potential contact zones, cross-part shared geometric context) it encodes" rather than "whether it can classify objects." Based on this, TORA defaults to Uni3D as the teacher—which consistently ranks high in geometric/contact probes and yields the highest downstream assembly accuracy. Conversely, scene-level encoders (Sonata, Concerto) bring almost no improvement due to the domain gap between scene-level pre-training and the object-level geometry that assembly relies on.

4. Late-Layer Alignment: Spatial structure emerges in deep layers

Besides choosing the right teacher, alignment must occur at the right layer. The authors analyze how spatial structures evolve across layers in the unaligned RPF flow backbone using four metrics: boundary contrast, LDS, part contour, and pose discriminability. They find that all four increase monotonically as the layer depth increases: deeper layers exhibit sharper feature transitions at part boundaries, more coherent local geometry, and better part-level clusterability. Crucially, pose discriminability (sensitivity to rigid motion, which is essential for 6-DoF pose solving) emerges primarily in late layers as global context is integrated. Since the model forms the global structures and interactions that determine alignment in the later layers, the teacher prior should be injected there. Ablations verify that deeper alignment layers yield better Part Accuracy, with \(l^*=5\) being optimal, significantly outperforming early-layer alignment and the unaligned baseline. This explains why the same teacher signal is only effective when applied to the correct layer.

Loss & Training¶

The total loss is \(\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{CFM}}+\lambda\,\mathcal{L}_{\text{align}}\), with \(\mathcal{L}_{\text{align}}=\mathcal{L}_{\text{CKA}}\) by default, alignment weight \(\lambda=0.5\), and CKA subsampling size \(n=1{,}024\). The teacher is Uni3D-L (303M), and alignment is applied to the \(l^*=5\)-th layer. The projection head is a 3-layer MLP (SiLU activation) with input/hidden dimension equal to the student's intermediate dimension of 1536 and output dimension equal to the teacher's dimension, discarded after training. The compared NT-Xent variant uses a temperature of \(\tau=0.07\). Training is conducted on 8 GH200 GPUs with a total batch size of 256 for 2000 epochs, using AdamW with a learning rate of \(5\times10^{-4}\) (halved every 200 epochs after 1000 epochs). Teacher features can be pre-computed and cached offline, further reducing the training overhead to near zero.

Key Experimental Results¶

Main Results¶

Evaluated on three complementary scenarios under a unified protocol: Breaking Bad (geometric assembly), PartNet-Assembly (semantic assembly), and TwoByTwo (cross-object assembly), compared with RPF, etc. Part Accuracy (PA (%), percentage of parts with Chamfer distance below a threshold of 0.01) is the main metric, and RE/TE represent rotation/translation errors.

Dataset	Metric	RPF (Baseline)	Ours-Cos-dist	Ours-CKA
Breaking Bad [2, 20 parts]	PA(%)↑	93.2	95.7	95.7
Breaking Bad [2, 20 parts]	RE(°)↓	16.0	9.0	8.6
Breaking Bad [21, 33 parts]	PA(%)↑	62.1	72.4	71.7
PartNet-Assembly	PA(%)↑	59.8	67.8	69.1
TwoByTwo	PA(%)↑	65.4	68.9	71.5
TwoByTwo	TE(cm)↓	11.9	9.5	7.6

The multi-part (21–33 parts) scenario widens the performance gap the most: correspondence-based methods like Jigsaw/CMNet/PMTR deteriorate rapidly due to the combinatorial complexity explosion, while GARF performs acceptably with few parts but drops sharply with more parts. RPF manages to maintain performance, whereas TORA further blows up its scalability. On cross-object and semantic assembly, CKA yields the largest gain, reflecting that "explicitly preserving relational structures is most valuable when structured interaction and domain shift coexist."

Zero-shot transfer (trained on Breaking Bad, directly evaluated on three unseen datasets without fine-tuning):

Dataset	Metric	GARF	RPF	Ours-CKA
BBad-Artifact	PA(%)↑	91.4	88.3	94.4
FRACTURA	PA(%)↑	44.2	68.1	76.0
Fantastic Breaks	RE(°)↓	8.2	6.3	3.5

Ablation Study¶

Configuration	TwoByTwo PA↑	Description
\(\mathcal{L}_{\text{CKA}}\) (Default)	71.5	Relational topology alignment alone is optimal
\(\mathcal{L}_{\text{CKA}}+\mathcal{L}_{\text{cos-dist}}\)	70.0	Adding token-wise target degrades performance
\(\mathcal{L}_{\text{CKA}}+\mathcal{L}_{\text{NT-Xent}}\)	68.5	Same as above, due to gradient clash
All three stacked	67.7	Worst
Alignment layer \(l^\): early layers \(\to l^=5\)	Deeper is better	Spatial structures emerge in late layers; deep alignment is optimal
Teacher: Uni3D vs Find3D/OpenShape/PatchAlign3D	Uni3D is optimal	Teachers stronger in geometric/contact probes yield the best downstream results

Key Findings¶

The gains of CKA stem from the "relational form" rather than better utilizing stronger teachers: Across both strong teachers (Uni3D-G/L) and weak teachers (Find3D/PatchAlign3D), CKA consistently outperforms Cos-dist (\(+0.7\%\) to \(+2.6\%\) PA) and does not vanish as the teacher weakens—demonstrating the intrinsic effectiveness of relational topology.
Significant convergence acceleration: On Breaking Bad, TORA reaches the baseline's peak accuracy approximately \(6.9\times\) faster, and \(3.3\times\) faster on PartNet (Cos-dist: \(2.2\times\), NT-Xent: \(1.8\times\)), while also achieving higher final accuracy.
NT-Xent only fails in the cross-object setting: NT-Xent fluctuates around the baseline across five single-object baselines, but drops by \(5.0\%\) PA solely on TwoByTwo. This is because its negative samples include tokens from a different object, which the teacher (trained on isolated objects) already pushes apart; forcing them further apart suppresses cross-object assembly cues.

Highlights & Insights¶

"Aligning relational topology rather than feature vectors" is the core "aha" moment: Matching Gram matrices via CKA permits any representation that shares the same "who-is-similar-to-whom" structure, regardless of feature space scaling/rotation. This tolerance for invariants perfectly aligns with the nature of assembly tasks—which prioritize relative relationships over absolute feature content. This offers a highly transferable recipe for other distillation tasks where "structure matters more than appearance."
Transforming "what makes a good teacher" into a testable probing problem: Instead of arbitrarily choosing a teacher, the authors use mating surface segmentation F1 and part contour as geometric probes to correlate with downstream accuracy, yielding a clear criterion of "geometry/contact \(\succ\) semantic classification" (backed by Fisher confidence intervals and partial correlation). This selection methodology is highly valuable.
Zero inference overhead + near-zero training overhead: The alignment branch is only attached during training, and the projection head is discarded afterwards. With offline cached teacher features, CKA only adds \(\sim 3\text{ ms}\) per step. Exchanging a "free" training-time regularization for a \(6.9\times\) training speedup and higher generalization is highly cost-effective.

Limitations & Future Work¶

The authors acknowledge that under highly symmetric or repetitive parts (such as a stack of nearly identical bars with minimal mating surfaces in PartNet), part-to-part signals become under-determined, which can lead to globally incoherent predictions. Since the alignment target does not directly constrain surface contact, occasional "geometric hallucinations" like boundary gaps may occur, which look jarring visually but may not be fully reflected by error metrics.
The geometric expressiveness of the teacher sets a ceiling for the framework, and the sensitivity to teacher choice remains an open question.
Future Work: Extending point-wise attention to larger assemblies or denser point clouds, incorporating visual and functional semantics into alignment, and reducing inference complexity to support real-time robotic assembly registration.

vs RPF (Rectified Point Flow): RPF relies solely on the endpoint reconstruction loss to implicitly learn mating relationships, and its encoder is only pre-trained on a binary overlap task. TORA adds a training-only alignment branch to distill relational geometric priors from a pre-trained 3D teacher, leaving the inference pipeline unchanged (zero overhead), yielding faster convergence and stronger cross-domain robustness.
vs REPA / REPA-E (Representation Alignment for 2D Generation): REPA-like methods align diffusion transformer features with frozen vision encoders to accelerate 2D image generation, typically using token-wise matching. TORA migrates this paradigm to 3D point flow assembly, revealing that the guide signals are dominated by geometric compatibility and sparse coordinate-aligned interactions. Thus, it requires a relational topology (CKA) target and geometric/contact-centric teachers, rather than directly copying token-wise alignment.
vs Correspondence-based Assembly (Jigsaw / PMTR / CMNet): These methods explicitly establish correspondences and solve poses via SVD, which suffers from combinatorial explosion when establishing reliable correspondences across multiple parts. TORA adopts a generative route to bypass explicit correspondences and maintains scalability in multi-component and cross-object scenarios using relational distillation.

Rating¶

Novelty: ⭐⭐⭐⭐☆ Introducing "relational topological distillation (CKA)" to 3D point flow assembly and systematically answering "what makes a good teacher/layer" offers a fresh perspective, though the overall alignment paradigm inherits from REPA.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Six baselines + zero-shot transfer + multi-dimensional ablations (on convergence/overhead/teacher selection/depth), supplemented with multi-seed statistical significance checks.
Writing Quality: ⭐⭐⭐⭐☆ Clear logic, robust tables, and highly insightful probing analysis. The notation is dense but readable.
Value: ⭐⭐⭐⭐☆ Zero inference overhead and highly plug-and-play, offering direct practical value for archaeological restoration and robotic assembly.