Equivariant Flow Matching for Point Cloud Assembly

We thank the reviewer for the time and effort. We address the concerns below.

1. while the paper provides solid technical contributions, these contributions are hidden under a very cluttered and heavy notation. Furthermore, no intuitive figures are provided which could help the reader understand key concepts more easily.

To improve the readability, we have now added a complete walk-through of a toy example in the appendix (Please see the "more intuition" part in our reply to reviewer i6up). We have also provided brief readable comments for Prop.4.5 and 4.6 (See our reply to reviewer ia2b quesion 2 for details) We hope those improvements can make the reading easier.

2. evaluation metrics which are used in the experiments are not explained and with that evaluation is not fully clear.

The evaluation metric is the pair-wise rotation/translation error. For each pair of piece $(i, j)$, we compute the error of $i$ w.r.t. $j$. We have now added the following clarification in Sec. 6.1: This metric is the pair-wise rotation/translation error: it measures the averaged error of $g_i$ w.r.t $g_j$ for all $(i, j)$ pairs of pieces.

3. the rendering of point clouds are hard to interpret as lack of shadows make it hard to perceive their 3D geometry

We apologize for the lack of clarity. We used Open3D for visualization, however its rendering ability is limited.

To provide better visualizations, we have now added more visualizations of Eda on 3DZ in Fig 10 in the appendix, where the camera is set to look at the room from above, so the main structures of the rooms are clearer. (Note that we are not allowed to upload figures during this rebuttal.)

5. proofs in the appendix should be referenced from the main paper, e.g. the proof of Theorem 4.2 should be referenced after Theorem 4.2.

Thanks for the suggestion. We have now added a sentence: All proofs can be found in Appx.C in the introduction section.

Q1. l. 345: the qualitative comparison between methods seems impossible since you chose different objects to assemble for the different methods in Fig 6 and 8 in the appendix

To avoid cherry picking, we used random samples in each class instead of fixed samples. Nevertheless, there are some overlaps of sample selection, for example, the ring shape has appeared in Fig. 6(a), 6(b), 7(a) and 8, where a direct comparison is possible.

Q2. l.373-375: the concluding sentence about kernel level optimization is unlcear and should be rewritten to improve clarity.

We apologize for the confusion. The "kernel" means CUDA kernel of the Nvidia GPU [8]. We have now made this clear: ...CUDA kernel level...

Q3. is there a specific reason why the scaling law was not studied? if not it should be added.

The focus of this work is a new theory and its algorithm. Since we have experimentally shown that the algorithm is effective on existing datasets, we think it is appropriate to leave the study of the scaling law to future work. We have now modified the last sentence to make this clearer: Finally, the scaling law [19] of Eda is an interesting research direction left for future work...

Q4. l. 698 citation is empty

Thanks for pointing this out. We have now fixed that citation ([28]).

Q5. fig 5. right: why is there such a large increase in training error in the one of the last training iterations? (even higher than the initial error)

This is a spike in the training loss. Such spikes are commonly seen in the loss curve of transformer-type models. See [1] for an in-depth discussion.

[1] Takase, Sho, et al. "Spike no more: Stabilizing the pre-training of large language models." arXiv preprint arXiv:2312.16903 (2023).

Thanks for your reply!

1. Thanks for your suggestion. We have now removed the (*) symbol in the subscript in Prop 4.5 and the paragraph after it, as the expression would be computationally (not mathematically) equivalent. However, we feel it is necessary to keep other subscripts, like $X$, (#), and $i$, as they represent the conditions, the pushforward, and the index of point cloud respectively. Removing them will cause some confusions.
2. We have now added a reference to the proof after each proposition.
3. Thanks for your suggestion. We will select more results, and place them side-by-side for better comparison.

We thank the reviewer for the time and effort. We address the concerns below.

1. Insufficient ablation study: The paper does not sufficiently validate the individual contributions of its components. This significantly weakens the experimental section. For example, the distribution of noise; the choice of equivariant and permutation-invariant network f; optimization in SE(3) space; rotation correction in Eq 6.

Thanks for your suggestion. As for ablation study, we did an ablation on the Rounge-Kutta degree (1/4) and step (from $5$ to $50$) in Fig. 5 in the appendix. We have now added an ablation study of the rotation correction. We compare Eda with and without rotation correction on 3DL and 3DM in the following table. We observe that the Eda without rotation correction performs worse than Eda, while it still performs better than all baselines in Tab.2. We have now added this table to the appendix.

As for other components, we do not feel necessary to discuss the use of them because we never claim any contribution on these components, and they are alreay sufficiently simple or are standard.

• The noise $(U_{SO(3)} \otimes \mathcal{N})^N$ is the simplest noise we can think of: the rotation component is a uniform distribution, which does not contain any preference; its translation component is a Gaussian noise, which is standard for Euclidean space, and they are independently coupled. We do not know any simpler alternative to this noise.
• The network is designed to be nothing more than an "imitation of the standard transformer structure" as stated in line 227. We combined Croco [42], Equiformer[23] and DiT (diffusion transformer) [29], which are already standard.
• The optimiztion is in $se(3)$ (a Lie algebra), which is a linear space. The training loss is the standard mean square error loss (Eqn 2) in flow matching.

2. The definition of the symbol is unclear (L121-125) and some background knowledge should be introduced in a more detailed way before moving on. Proposition 4.5 and 4.6 seems to be very important, maybe a brief comment using natural language will help the readers to understand the paper.

Due to the limited space, we can not present more preliminaries on related vector fields, so we cited an introductional text book [37], where the reader can get a quick idea of the concepts. We have now made this clearer: More details can be found in Sec.14.6 in [37].

Intuitively, the "relatedness" of two vector fields $v_{X}$ and $v_{gX}$ means that $v_{X}$ is a "transformed" version of $v_{gX}$. Corollary 4.4 first suggests that to keep the $SO(3)^N$-equivariance of solution, $v_{X}$ must be a transformed version of $v_{gX}$. Then, Proposition 4.5 suggests that to satisfy such requirement of $v_X$ and $v_{gX}$, equation $v_X(g)=v_{gX}(e)g$ must hold. This is useful because it provides a way to represent $v_{X}$ as a neural network, i.e., in Eqn (4), when we parameterize $v_{gX}(e)$ on the right hand side as $f(gX)$, then $v_X(g)=f(gX)g$. Futhermore, Proposition 4.6 translates the other two requirements of $v_{X}$ (permutation and $SO(3)$-relatedness) to the requirements of $f$.

The comments of these results were already in line 163-165 and line 170-172. Following your suggestion, we have now rewritten them for readability. Line 163: Prop. 4.5 provides a way to represent $v_X$ by a neural network. Specifically,.... Line 170: To guarantee $\sigma$-relatedness and $\mathcal{L}_r$-invariance of $v_X$, the following requirements of $f$ are needed.. We have further made proposition 4.6 clearer by replacing the relatedness notation \\# by natural language: in (1) ...., then $v_X$ is $\sigma$-related to $v_{\sigma X}$. and in (2) ..., then $v_X$ is $\mathcal{L}_r$-invariant.

Please also refer to the added "more intuition" section in our reply to reviewer i6up, where we provide a walk-through of a toy example. A concrete formulation of Prop.4.5 and 4.6 is also provided there.

3. Also I think the bottom-right corner of se3 matrix in Eq 5 should be 1.

No. It is $0$. a $se(3)$ matrix (a Lie algebra) is not a $SE(3)$ (a Lie group) matrix. $se(3)$ can be seen as a gradient of $SE(3)$. The bottom-right corner of $SE(3)$ is a constant $1$, so after taking a gradient, that element is $0$ in $se(3)$.

4. Experiments in zero-overlap setting: It seems that the rotation error is very large (78.32) and it is said that the distribution of error is centered at 0, 90, 180. The meaning of this experiment is unclear.

The meaning of the experiments is to show that Eda can "learn global geometric properties beyond correspondences" even when the true assembly is ambiguous due to the lack of overlap as noted in line 328. This is "a key advantage of Eda" compared to all baseline methods in Tab.2, which all rely on correspondence, i.e., they cannot learn such geometry properties in this experiment.

The statistics of rotation error was reported to show that the model learned plausible solutions instead of doing random guessing despite a large mean error, i.e. if the model is doing random guessing, then the rotation error distribution would be close to a uniform distribution on interval [0, 180].

5. Although I agree that "Eda learns to keeps the orthogonality or parallelism of walls, floors and ceilings of the indoor scenes", it seems that all it learns is this global geometric constraints instead of part assembly. Moreover, the visualization in Figure 2 is hard to understand.

Being able to keep the orthogonality/parallelism of planes already shows that Eda can learn the rotation component of the assembly. In addition, as shown in the visualization in Fig.2(a), Eda can also learn the translation faithfully, i.e., it keeps a plausible distance between walls of the assembled room, while keeping the ceilings (floors) on the same plane.

As for the visualiztion, we apologize for the lack of clarity. We have now added more visualizations of Eda on 3DZ in Fig 10 in the appendix for better understanding, where the camera is set to look at the room from above. (Note that we are not allowed to upload figures during this rebuttal.)

Thanks for your reply.

We believe that the experiments are sufficient to show the usefulness of the theory becaseu our model is built solely on the theory, and it is shown to be SOTA. We do not think more ablation studies are neede because we have never claimed the performance gain of using any individual part of the theory, like the networks, or the mean square error training.

In particular, the advantages of the SO(3)-equivariant networks are proved to be the ensurance of the SO(3)-invariance of the solution, but we have never claimed (or proved) the performance gain of using those networks, so we do not see why an ablation on the networks is needed: replacing the equivariant network by a non-equivariant one may lead to better, worse, or comparable performance, and none of these results conflicts with our theory.

If the reviewer wants to see the contribution of equivariances on performance, a recent diffusion-based model AMR in Tab. 2 could be compared, which is permutation-equivariant but not SO(3)-equivariant.

We thank the reviewer for the positive feedback. We now address the concerns below.

1. The idea of parameterizing vector fields using equivariant networks has been explored in other domains, e.g., 3D module generation [1,2]. Despite the differences in applications, the network architectures and training strategy (Optimal Transport Training) seem to be highly related. It would be beneficial to cite and discuss these works.

Thanks for your suggestion. Those two papers are indeed related to our work.

Compared to [1,2], a key difference of our work in theory is that it compares different "equivariant" or "related" vector fields, e.g, $v_X$ and $v_{gX}$, instead of one "invariant" vector field (the word "invariant" in our work means "equivariant" in [1,2]. Aslo see our reply 3 for this difference.). This is noted in line 65. On the other hand, [1, 2] developed an OT training strategy for flow matching on invariant distributions which is beyond the discussion of our work. Nevertheless, we use the rotation correction (6) which is similar to the first equation in (16) in [1], and we believe that the use of the second equation in (16), i.e., OT, is useful for future applications.

We have now added the following sentence in line 66 in the related work section: Furthermore, when the distribution is invariant, the optimal-transport-based integral path [53, 54] was explored for efficient sampling. and we have cited them in line 196: A similar rotation correction in the Euclidean space was studied in [53, 54].

2. ...In particular, the variance of the evaluation metrics should be reported, and multiple solutions could be visualized given the same set of inputs, especially for cases without significant overlaps.

Thanks for the suggestion. The variance was already reported in Tab. 5 and Tab. 6 in the appendix.

We have now added a figure showing 3 different runs of Eda on 6 samples of 3DZ in Fig. 10 in the appendix. (Note that we are not allowed to upload figures during this rebuttal.)

3. In the third condition of Corollary 4.4, it seems that equivariance of the vector field to SO(3) is already sufficient to ensure the invariance of the generation distribution (as suggested in [1]). Considering that invariance is a stricter requirement, I wonder if it is necessary?

Yes. It is necessary. This is because the "equivariance" of vector field in [1,2] is called "invariance" in our work, so Corollary 4.4 in our work and Theorem 4.1 in [2] use the same concept. (We save the word "equivariance" for comparing two different vector fields.) We noted such difference in line 624.

4. The number of pieces is predetermined in the current framework. It would be valuable to discuss the possibility of extending the framework to handle a variable number of pieces.

No. The number of pieces can vary. This is because the $SO(3)^N$-equivariance is ensured by construction: the network computes the velocity $(w, t)$ for each of the $N$ pieces, where the network (transformer structure) does not assume a fixed length $N$. For example, in our experiments, the BB dataset contains $2 \sim 8$ pieces, and the kitti dataset contains $2 \sim M$ pieces, where $M$ is $3$ or $4$.

Dear reviewer xTSh, please let us know whether we have addressed your questions and concerns.

We thank the reviewer for the time and effort. We first present the added section to provide more intuitions, and then respond to each of the specific concern.

More intuition

We have now added the following paragraphs to the appendix to provide some intuitive explanations of the main results without the tools for describing transformations of probabilities (pushforward) and vector fields (relatedness). We have now added a sentence in the introduction: A walk-through of the theory using a toy example is provided in Appx. B.

Sec B. A toy example

Consider the following two-piece deterministic example. Assume that a solution for the input point clouds $(X, Y)$ is $(r_1, r_2)$, meaning $r_1X$ and $r_2Y$ are assembled, where $r$ is a rotation matrix.

The equivariances in Def 3.1 are natural properties of the solution: when $(X, Y)$ are transformed, the solution will change accordingly:

1. $SO(3)^2$-equivariance: a solution for $(r_3X, r_4Y)$ is $(r_1r_3^{-1}, r_2r_4^{-1})$. Because $(r_1r_3^{-1}r_3X, r_2 r_4^{-1}r_4Y)=(r_1X, r_2Y)$ are assembled by assumption.
2. Permutation-equivariance: a solution for $(Y, X)$ is $(r_2, r_1)$. Because $(r_2Y, r_1X)$ are assembled by assumption.
3. SO(3)-invariance: another solution for $(X, Y)$ is $(rr_1, rr_2)$. Because $(rr_1X, rr_2Y)$ are just the assembled point clouds $(r_1X, r_2Y)$ rotated by $r$.

Corollary 4.4 incorporates these equivariances into flow matching. Denote $v_{(X, Y)}$ the vector field learned for $(X, Y)$. For $SO(3)^2$-equivariance, the goal is to ensure $v_{(r_3X, r_4Y)}$ flows to $(r_1r_3^{-1}, r_2r_4^{-1})$ when $v_{(X, Y)}$ flows to $(r_1, r_2)$. This corollary shows that the goal can be achieved if $v_{(r_3X, r_4Y)}$ is a proper "transformation" of $v_{(X, Y)}$ (related), and the noise is invariant.

Proposition 4.5 provides a way to construct $v_{(X, Y)}$ satisfying the $SO(3)^2$-equivariance requirement of Corollary 4.4: $v_{(X, Y)}(r_7, r_8) =f(r_7X, r_8Y)(r_7 \oplus r_8)$, here $f(X, Y) = (w_1, t_1) \oplus(w_2, t_2)$ is a neural network mapping $(X, Y)$ to their respective rotation/translation velocity component $w$ and $t$, and $\oplus$ is the concatenation. (Vector $w$ and $t$ are combined into a matrix as in Eqn. 5.)

Proposition 4.6 suggests that, to ensure the other two requirements (permutation and SO(3)-relatedness) of $v_X$, $f$ needs to satisfy

Proposition 4.7 suggests that some data augmentations are not needed. Specifically, for data $(X, Y)$ we learn a vector field $v_{(X, Y)}$. We can randomly augment the data $(r_3X, r_4Y)$ and learn $v_{(r_3X, r_4Y)}$. However, this proposition suggests that this is not necessary when the path and $v_{(X, Y)}$ are "$SO(3)^2$-equivariance" (relatedness) and the noise is invariant. Similar results hold for the other two types of augmentations.

We now address each of the concerns below.

1. The paper is a bit difficult to read and could benefit from modifying the text (esp Sec 4.2). 3. The intuition of the proposed method is not clear, and part of the method seems to lack clarity. Q2. If the distribution is equivariant, how are the parts in the assembly aligned? Q3. Proposition 4.7 does not seem clear. Could you elaborate on this?

For more intuition, we have now added a complete walk-through of a toy example in the appendix (see above). In particular, the meaning of 3 equivariances and the meaning of Prop. 4.7 is included. For clarification, we have now provided brief and readable comments for Prop.4.5 and 4.6 (A brief comment of Prop. 4.7 are in line 208.) Please see our reply to reviewer ia2b quesion 2 for details. We hope the above modifications make the reading easier.

2. The model evaluation, both wr.t to datasets as well as ablation, is limited.

We have included all standard assembly datasets in our experiments. i.e., 3DM, 3DL and kitti. See a recent paper AMR[6] for example. We additionally considered BB and a novel dataset 3DZ. If the reviewer thinks some datasets are missing, please let us know.

We did an ablation study on the Rounge-Kutta degree (1/4) and step (from $5$ to $50$) in Fig. 5 in the appendix. We have added an ablation study for rotation correction. Please see our reply to reviewer ia2b question 1 for details.

Q1. Why is there a need to compute the rotation error (equation 6) if the path is SO(3) equivariant?

As we stated in line 197, the purpose is to remove the "redundant rotation component". For an oversimplified example, for noise $g_0=(r_1, r_2)$ and a data point $\tilde{g_1}=(rr_1, rr_2)$, we can get a clean $g_1=(r_1, r_2)$ where $r$ is removed by solving (6). This process is compatible with the path selection.

Q4. In the experimental setup, the center-of-mass frame is considered. Now, suppose all the objects are aligned, i.e., there is a central axis across which the object is present (which is more common for ShapeNet, etc.). How does the Eda use this to be computationally efficient?

Eda is designed to be $SO(3)$-invariant, meaning that it has no preference of canonical poses (like those in ShapeNet). To use those canonical poses, one needs to remove the $SO(3)$-invariance, which would remove Proposition 4.6 (2), meaning equivariant networks are not needed. This might lead to more efficient computation.

Q5. Could you provide intuition for the necessity of using training data of similar length to the test data?

Apology for the confusion of this sentence. We mean the test data lengths should be included in training data. We correct the sentence in line 367 as ...necessity of using training data whose lengths subsume that of the test data.

Such requirement of the training data is common for transformer models. See Table 2 in [1] for an example.

[1] Sun, Yutao, et al. "A length-extrapolatable transformer." arXiv preprint arXiv:2212.10554 (2022).