SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ilyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=YPpSngE-ZU.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Avishek Joey Bose, Marcus Brubaker, and Ivan Kobyzev. Equivariant finite normalizing flows, 2022.
Andrew Campbell, Wenlong Chen, Vincent Stimper, Jose Miguel Hernandez-Lobato, and Yichuan Zhang. A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization. In Proceedings of the 38th International Conference on Machine Learning, pages 1238-1248. PMLR, 2021.
Taco S. Cohen and Max Welling. Steerable cnns. International Conference on Learning Representations, 2017.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: non-linear independent components estimation. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings, 2015. URL http://arxiv.org/abs/1410.8516.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview.net/forum?id=HkpbnH91x.
John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. Neural spline flows. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7509-7520, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/7ac71d433f282034e088473244df8c02-Abstract.html.
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems 33, 2020.
Vincent Stimper, Bernhard Scholkopf, and José Miguel Hernández-Lobato. Resampling Base Distributions of Normalizing Flows. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pages 4915-4936, 2022.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian FATRAS, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023. URL https://openreview.net/forum?id=HgDwiZrpVq.
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=8pvnfTAbu1f.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/bit/41d80bfc327ef980528426fc810a6d7a-AAbstract.html.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openreview.net/forum?id=JprMOp-qOCo.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation. In International Conference on Learning Representations, 2022.
Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi S. Jaakkola. SE(3) diffusion model with application to protein backbone generation. 202:40001-40039, 2023.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generation flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2310.02679, 2023.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=ucb2ynRu7Y.
A.1 Tensors and representation
We assume a group  $(G,\cdot)$ , an algebraic structure that consists of a set  $G$  with a group operator  $\cdot$  that follows the closure, inverse and associativity axioms.
A representation is an invertible linear transformation  $\rho(g): V \to V$ , indexed by group elements  $g \in G$ , and defines how group elements  $g$  acts on elements of the vector space  $V$ . It satisfies the group structure  $\rho(gh) = \rho(g)\rho(h)$ .
A tensor  $T$  is a geometrical object that is acted on by group elements  $g \in G$  in a particular way, characterized by its representation  $\rho \colon g \cdot T = \rho(g)T$ . Typical examples include scalar with  $\rho_{\mathrm{triv}}(g) = 1$  and vectors living in  $V$  with  $\rho_{\mathrm{Id}}(g) = g$ . Higher order tensors which live in  $V \times \dots \times V$  can be constructed as  $g \cdot T = (\rho_{\mathrm{Id}}(g) \otimes \dots \otimes \rho_{\mathrm{Id}}(g)) \operatorname{vec}(T)$ . A collection of tensors  $(T_1, \ldots, T_n)$  with representations  $(\rho_1, \ldots, \rho_n)$  can be stacked with fiber representation as the direct sum  $\rho = \rho_1 \oplus \dots \oplus \rho_n$ .
A.2 Translational invariant densities
Proposition A.1 (Disintegration of measures for translation invariance [Pollard, 2002, Yim et al., 2023]). Under mild assumptions, for every SE(3)-invariant measure  $\mu$  on  $\mathbb{R}^{3\times n}$ , there exists  $\eta$  an SO(3)-invariant probability measure on  $\{\mathbb{R}^{3\times n}|\sum_{i = 1}^{n}x_{i} = 0\}$  and  $\bar{\mu}$  proportional to the Lebesgue measure on  $\mathbb{R}^3$  s.t.  $\mathrm{d}\mu ([x_1,\ldots ,x_n]) = \mathrm{d}\bar{\mu} (\frac{1}{n}\sum_{i = 1}^{n}x_i)\times \mathrm{d}\eta ([x_1 - \frac{1}{n}\sum_{i = 1}^{n}x_i,\ldots ,x_n - \frac{1}{n}\sum_{i = 1}^{n}x_i])$ .
The converse is also true, which motivates constructing an SO(3)-invariant probability measure on the subspace of  $\mathbb{R}^{3\times n}$  with zero center of mass.
B.1 Transforming variables on a subspace
This section shows that a flow transform  $\mathcal{F}:\mathcal{X}\to \mathcal{X}$  applied to observations  $\tilde{\pmb{x}}$  embedded in the zero-CoM subspace  $\tilde{\mathcal{X}}$  does not have a well-defined Jacobian. We will show this by constructing such a transform, which leaves zero-CoM input unchanged, but for which the log-determinant is arbitrary.
Without loss of generality, we will assume in this subsection and the next that our observed and augmented variables are one dimensional, i.e.,  $\mathcal{X} = \mathcal{A} = \mathbb{R}^n$ . Consider the orthonormal basis  $V \in \mathbb{R}^{n \times n}$  constructed by concatenating orthonormal basis vectors  $\pmb{v}_i \in \mathbb{R}^n$  as  $[\pmb{v}_1, \dots, \pmb{v}_n]^T$ . We choose  $\pmb{v}_1$  to be equal to  $[1, 1, \dots, 1]^T / \sqrt{n}$  such that inner products with  $\pmb{v}_1$  retrieve the CoM of a system of particles up to multiplication by a scalar:  $\pmb{v}_1^T \pmb{x} = \sqrt{n} \bar{\pmb{x}}$  and  $\pmb{v}_1^T \tilde{\pmb{x}} = 0$ . Let  $\mathcal{F}(\pmb{x}) = V^T S V \pmb{x}$  be a mapping which rotates onto the basis defined by  $V$ , multiplies by the diagonal matrix  $S = \mathrm{diag}([s, 1, \dots, 1])$  and then rotates back to the original space with  $V^T$ . This transformation only acts on the CoM. Applying this transformation to a zero-CoM variable  $\tilde{\pmb{x}}$  we have
$$
\mathcal {F} (\tilde {\boldsymbol {x}}) = V ^ {T} S V \tilde {\boldsymbol {x}} = V ^ {T} \left[ s \boldsymbol {v} _ {1} ^ {T} \tilde {\boldsymbol {x}}, \boldsymbol {v} _ {2} ^ {T} \tilde {\boldsymbol {x}}, \dots , \boldsymbol {v} _ {n} ^ {T} \tilde {\boldsymbol {x}} \right] = V ^ {T} \left[ 0, \boldsymbol {v} _ {2} ^ {T} \tilde {\boldsymbol {x}}, \dots , \boldsymbol {v} _ {n} ^ {T} \tilde {\boldsymbol {x}} \right] = \tilde {\boldsymbol {x}},
$$
leaving the zero CoM input unchanged. However, the log-determinant is
$$
\log \left| \frac {\partial \mathcal {F} (\boldsymbol {x})}{\partial \boldsymbol {x}} \right| = \log | V ^ {T} S V | = \log | V ^ {T} | + \log | S | + \log | V | = \log s,
$$
since for the unitary matrices  $|V| = |V^T| = 1$ . Because  $s$  is arbitrary, so is the density returned by any flow applying such a layer to zero-CoM variables. In the maximum likelihood setting, the aforementioned flow would lead to a degenerate solution of  $s \to 0$  such that the log determinant would approach negative infinity and the log likelihood would approach infinity.
While our flow does use transforms which operate on non-CoM variables, we lift our CoM observations to the non-zero-CoM space before applying said transformations. We do this using the ShiftCoM operation, described in Sec. 3 and below.
B.2 Jacobian of ShiftCoM
We now show our ShiftCoM:  $\tilde{\mathcal{X}}\times \mathcal{A}\mapsto \mathcal{X}\times \tilde{\mathcal{A}}$  transformation to have unit Jacobian log determinant. We first, re-state the function definition
$$
\operatorname {S h i f t C o M} (\tilde {\boldsymbol {x}}, \boldsymbol {a}) \triangleq (\tilde {\boldsymbol {x}} - \bar {\boldsymbol {a}}, \boldsymbol {a} - \bar {\boldsymbol {a}}) \quad \text {w i t h} \quad \bar {\boldsymbol {a}} \triangleq \frac {1}{n} \sum_ {i = 1} ^ {n} [ \boldsymbol {a} ] ^ {i}.
$$
We now re-write it using the orthonormal basis  $V \in \mathbb{R}^{n \times n}$  defined in the previous subsection. For this, we use  $I_2 \otimes V \in \mathbb{R}^{2n \times 2n}$  to refer to a  $2 \times 2$  block diagonal matrix with  $V$  in both blocks, and we also stack the inputs and outputs of ShiftCoM, yielding ShiftCoM $([\tilde{\boldsymbol{x}}, \boldsymbol{a}]) = (I_2 \otimes V)^T P(I_2 \otimes V)[\tilde{\boldsymbol{x}}, \boldsymbol{a}]$ , where  $P$  is a permutation matrix that exchanges the first and  $n + 1$ th elements of the vector it acts upon. It also flips this element's sign. To see this note
$$
\begin{array}{l} \operatorname {S h i f t C o M} ([ \tilde {\boldsymbol {x}}, \boldsymbol {a} ]) = (I _ {2} \otimes V) ^ {T} P (I _ {2} \otimes V) [ \tilde {\boldsymbol {x}}, \boldsymbol {a} ] \\ = \left(I _ {2} \otimes V\right) ^ {T} P \left[ 0, v _ {2} ^ {T} \tilde {\boldsymbol {x}}, \dots , v _ {n} ^ {T} \tilde {\boldsymbol {x}}, v _ {1} ^ {T} \boldsymbol {a}, \dots , v _ {n} ^ {T} \boldsymbol {a} \right] \\ = \left(I _ {2} \otimes V\right) ^ {T} \left[ - \boldsymbol {v} _ {1} ^ {T} \boldsymbol {a}, \boldsymbol {v} _ {2} ^ {T} \tilde {\boldsymbol {x}}, \dots , \boldsymbol {v} _ {n} ^ {T} \tilde {\boldsymbol {x}}, 0, \boldsymbol {v} _ {2} ^ {T} \boldsymbol {a}, \dots , \boldsymbol {v} _ {n} ^ {T} \boldsymbol {a} \right] \\ = [ \bar {\boldsymbol {x}} - \bar {\boldsymbol {a}}, \boldsymbol {a} - \bar {\boldsymbol {a}} ]. \\ \end{array}
$$
For the determinant we use that  $|I_2| = 1$  and  $|V| = 1$ , and thus the determinant of their Kronecker product will be equal to one  $|(I_2 \otimes V)| = 1$ . Combining this with the fact that permutation matrices and matrices that flip signs have unit determinant, we arrive at
$$
\begin{array}{l} \log \left| \frac {\partial \operatorname {S h i f t C o M} ([ \tilde {\boldsymbol {x}} , \boldsymbol {a} ])}{\partial [ \tilde {\boldsymbol {x}} , \boldsymbol {a} ]} \right| = \log \left| (I _ {2} \otimes V) ^ {T} P (I _ {2} \otimes V) \right| \\ = \log | (I _ {2} \otimes V) ^ {T} | + \log | P | + \log | I _ {2} \otimes V | = 1. \\ \end{array}
$$
It is worth noting that ShiftCoM preserves the inherent dimensionality of its input; both its inputs and outputs consist of one zero-CoM vector and one non-zero-CoM vector. Because ShiftCoM does not apply any transformations to the CoM, it does not suffer from the issue of ill-definednes discussed in the previous subsection.
B.3 Projection variants and their respective transformations
In this section we provide a detailed description of the core transformation of our equivariant coupling flow layer, see Fig. 1. It consists of a projection into an invariant space via an equivariant reference, followed by a flow transformation in this space and a back-projection.
We denote the inputs with superscript  $\ell$  and use  $\ell + 1$  for outputs to define our transformation block  $\mathcal{F}:(\pmb{x}^{\ell};\pmb{a}^{\ell}) \mapsto (\pmb{x}^{\ell + 1},\pmb{a}^{\ell + 1})$ . Without loss of generality, we assume that we are transforming  $\pmb{x}^{\ell}$ . Then,  $\mathcal{F}$  takes the following form.
$$
\begin{array}{l} \boldsymbol {x} ^ {\ell + 1} = \gamma_ {\boldsymbol {r}} ^ {- 1} \cdot \tau_ {\theta} \left(\gamma_ {\boldsymbol {r}} \cdot \boldsymbol {x} ^ {\ell}\right), \quad \text {w i t h} (\boldsymbol {r}, \theta) = h (\boldsymbol {a} ^ {\ell}), \\ \boldsymbol {a} ^ {\ell + 1} = \boldsymbol {a} ^ {\ell}. \\ \end{array}
$$
Here,  $h$  is a (graph) neural network that returns a set of  $u$  equivariant reference vectors  $\boldsymbol{r} = [r^1, \ldots, r^u]$ , and invariant parameters  $\theta$ . The equivariant vectors  $\boldsymbol{r}$  parametrize  $\gamma_r$ , a projection operator onto a rotation invariant feature space. This means that any rotations applied to the inputs will be cancelled by  $\gamma_r$ , i.e.  $\gamma_{g \cdot r} = \gamma_r^{-1} \cdot g^{-1}$  or equivalently  $(\gamma_{g \cdot r})^{-1} = g \cdot \gamma_r^{-1}$  for all  $g \in \mathrm{SO}(3)$ .
In the next paragraphs, we discuss three projection operations and their respective flow transformations.
Cartesian projection Here, we predict a 3D Cartesian coordinate system as illustrated in Fig. 4a. The coordinate system has an origin  $o = r^1$  and an orthonormal basis  $Q$ , both being equivariant.  $Q$  is constructed from two equivariant vectors  $v_{1} = r^{2} - r^{1}$  and  $v_{2} = r^{3} - r^{1}$  in the following way. To obtain the first basis vector, we simply normalize  $v_{1}$ , i.e.  $b_{1} = \frac{v_{1}}{\|v_{1}\|}$ . We get the second basis vector via a Gram-Schmidt [Petersen, 2012] iteration, i.e. we compute the component of  $v_{2}$  orthogonal to  $v_{1}$ , which is  $v_{2,\perp} = v_{2} - b_{1}^{\top}v_{2}$ , and then normalize it to obtain  $b_{2} = \frac{v_{2,\perp}}{\|v_{2,\perp}\|}$ . Note that this is ill-posed if  $v_{1}$  and  $v_{2}$  are collinear. To avoid this case, we introduce an auxiliary loss,
$$
\text {a u x - l o s s} = - \log \left(\epsilon + \arccos  \left(\operatorname {a b s} \left(v _ {1} ^ {T} v _ {2} / \| v _ {1} \| / \| v _ {2} \|\right)\right)\right), \tag {11}
$$
(a) Cartesian projection.
Figure 4: Illustration of the Cartesian and spherical projection.
(b) Spherical projection.
which uses a logarithmic barrier function on the angle between the vectors to prevent them from being collinear.  $\epsilon$  is a small constant that prevents the aux-loss from being equal to infinity. Finally, the third basis vector is  $b_{3} = b_{1}\times b_{2}$ , and we can assemble  $Q = [b_{1},b_{2},b_{3}]$ .
The projection into this coordinate system and the respective inverse operation are given by
$$
\gamma_ {\boldsymbol {r}} (\boldsymbol {x}) = Q ^ {- 1} (\boldsymbol {x} - o)), \tag {12}
$$
$$
\gamma_ {r} ^ {- 1} \left(\boldsymbol {x} _ {\mathrm {p}}\right) = Q \boldsymbol {x} _ {\mathrm {p}} + o, \tag {13}
$$
where  $x_{\mathrm{p}}$  is the transformed projected variable.
In the projected space, we can apply any standard flow transform, as long as it acts on each particle individually. We choose  $\tau_{\theta}$  to an affine transformation as in the Real NVP architecture [Dinh et al., 2017].
Note that since  $b_{3}$  is constructed via a cross product, it is not parity equivariant, which is why the entire transformation block is not parity equivariant either. However, since  $\mathrm{sign}(b_3^\top r^2)$  is a pseudo-scalar, we can replace  $b_{3}$  with  $b_{3}^{\prime} = \mathrm{sign}(b_{3}^{\top}r^{2})b_{3}$ , which is a vector, to retain parity equivariance.
The determinant of the Jacobian of  $\mathcal{F}$  when using the Cartesian projection is just the determinant of the Jacobian of  $\tau_{\theta}$ , which we show below.
$$
\begin{array}{l} \left| \frac {\partial \mathcal {F} (\boldsymbol {x} ; \boldsymbol {a})}{\partial (\boldsymbol {x} , \boldsymbol {a})} \right| = \left| \frac {\partial \mathcal {F} (\boldsymbol {x} ; \boldsymbol {a}) | _ {\boldsymbol {x}}}{\partial \boldsymbol {x}} \right| \cdot \left| \frac {\partial \mathcal {F} (\boldsymbol {x} ; \boldsymbol {a}) | _ {\boldsymbol {a}}}{\partial \boldsymbol {a}} \right| \\ = \left| \left[ \frac {\partial \gamma_ {\boldsymbol {r}} ^ {- 1} (\boldsymbol {z})}{\partial \boldsymbol {z}} \right] _ {\boldsymbol {z} = \tau_ {\theta} (\gamma_ {\boldsymbol {r}} (\boldsymbol {x}))} \left[ \frac {\partial \tau_ {\theta} (\boldsymbol {z})}{\partial \boldsymbol {z}} \right] _ {\boldsymbol {z} = \gamma_ {\boldsymbol {r}} (\boldsymbol {x})} \frac {\partial \gamma_ {\boldsymbol {r}} (\boldsymbol {x})}{\partial \boldsymbol {x}} \right| \cdot 1 \\ = \left| Q \left[ \frac {\partial \tau_ {\theta} (\boldsymbol {z})}{\partial \boldsymbol {z}} \right] _ {\boldsymbol {z} = \gamma_ {r} (\boldsymbol {x})} Q ^ {- 1} \right| \\ = | Q | \left| \left[ \frac {\partial \tau_ {\theta} (\boldsymbol {z})}{\partial \boldsymbol {z}} \right] _ {\boldsymbol {z} = \gamma_ {r} (\boldsymbol {x})} \right| | Q ^ {- 1} | \\ = \left| \left[ \frac {\partial \tau_ {\theta} (\boldsymbol {z})}{\partial \boldsymbol {z}} \right] _ {\boldsymbol {z} = \gamma_ {\boldsymbol {r}} (\boldsymbol {x})} \right|. \\ \end{array}
$$
Spherical projection For the spherical projection, which is illustrated in Fig. 4b, we compute an origin  $o$  and three Cartesian basis vectors  $b_{1}, b_{2}$ , and  $b_{3}$  as in the Cartesian projection. Instead of
computing the Cartesian coordinates of each  $x^i$ , we compute the spherical coordinates  $(\rho, \beta, \varphi)$ . Hence, for the  $i$ -th atom the projection operation and its inverse is given by
$$
\left(\gamma_ {\boldsymbol {r}} (\boldsymbol {x})\right) ^ {i} = \left(\| x ^ {i} \|, \operatorname {a r c c o s} \left(\frac {x _ {\parallel} ^ {i}}{\| x ^ {i} \|}\right), \operatorname {a t a n 2} \left(b _ {3} ^ {\top} x _ {\perp} ^ {i}, b _ {2} ^ {\top} x _ {\perp} ^ {i}\right)\right) ^ {\top}, \tag {14}
$$
$$
\left(\gamma_ {\boldsymbol {r}} ^ {- 1} \left(\boldsymbol {x} _ {\mathrm {p}}\right)\right) ^ {i} = \rho^ {i} \left(\cos \left(\beta^ {i}\right) b _ {1} + \sin \left(\beta^ {i}\right) \left(\cos \left(\varphi^ {i}\right) b _ {2} + \sin \left(\varphi^ {i}\right) b _ {3}\right)\right), \tag {15}
$$
Here, we used  $x_{\parallel}^{i} = b_{1}^{\top}x^{i}, x_{\perp}^{i} = x^{i} - x_{\parallel}^{i}$ , and  $x_{\mathfrak{p}}^{i} = (\rho^{i},\beta^{i},\varphi^{i})$  for brevity.
The spherical coordinates are transformed with monotonic rational-quadratic splines [Durkan et al., 2019]. The radial coordinate  $\rho$  is transformed through a spline with the interval starting at 0 to ensure that the transformed coordinate is also non-negative to ensure invertibility. The two angles are transformed via circular splines [Rezende et al., 2020]. Similarly to the Cartesian projection, we use the auxiliary loss (Eq. (11)) to prevent the vectors output from the GNN from being co-linear.
When using the spherical projection, we compute the log-determinant of the Jacobian of  $\mathcal{F}$  by using the chain rule in the same way as we do for standard normalizing flows. The Jacobian determinant of the projection operation is given by the well-known term
$$
\left| \frac {\partial (\gamma_ {r} (\boldsymbol {x})) ^ {i}}{\partial \boldsymbol {x}} \right| = (\rho^ {i}) ^ {2} \sin \beta^ {i}. \tag {16}
$$
For the inverse, we can just raise this term to the power of  $-1$ .
Vector projection This is a special case of the spherical projection, where we only transform the radial coordinate but leave the angles the same. Since in practice we only need to compute  $\rho$ , we call this case vector projection. Note that just predicting the origin  $o$  as a reference with the equivariant neural network is sufficient here. Unlike the Cartesian and spherical projections, the vector projection does not rely on the cross product, and hence is parity equivariant by default.
B.4 Multiple augmented variables
Consider augmented variables consisting of multiple sets of observation-sized arrays, i.e.  $\mathcal{A} = \mathcal{X}^k$  with  $k > 1$ . For the base distribution, and also the target, we deal with multiple augmented variables by letting them follow conditional normal distributions centered at  $\pmb{x}$ . We thus have  $\pi(\pmb{a}|\pmb{x}) = \prod_{i=1}^{k} \mathcal{N}(\pmb{a}_i; \pmb{x}, \sigma^2 I)$  and  $q_0(\pmb{x}, \pmb{a}) = \tilde{\mathcal{N}}(\pmb{x}; 0, I) \prod_{i=1}^{k} \mathcal{N}(\pmb{a}_i; \pmb{x}, \sigma^2 I)$ .
For our flow block, we group our variables into two sets  $(\tilde{\boldsymbol{x}},\boldsymbol{a}_1,\dots,\boldsymbol{a}_{(k-1)/2})$  and  $(\boldsymbol{a}_{(k-1)/2+1},\dots,\boldsymbol{a}_k)$ , and update one set conditional on the other, and vice-versa, while always ensuring the arrays being transformed have non-zero CoM.
Initially,  $\pmb{a}$  is non-zero-CoM with each augmented variable  $a_{i}$  having a CoM located at a different point in space. Before the first coupling transform  $\mathcal{F}_M^{(1)}$ , the ShiftCoM function applies the mapping
$$
\operatorname {S h i f t C o M} (\tilde {\boldsymbol {x}}, \boldsymbol {a}) \triangleq (\tilde {\boldsymbol {x}} - \bar {\boldsymbol {a}} _ {1 + \frac {k - 1}{2}}, \boldsymbol {a} - \bar {\boldsymbol {a}} _ {1 + \frac {k - 1}{2}}) \quad \text {w i t h} \quad \bar {\boldsymbol {a}} _ {1 + \frac {k - 1}{2}} \triangleq \frac {1}{n} \sum_ {i = 1} ^ {n} [ \boldsymbol {a} _ {1 + \frac {k - 1}{2}} ] ^ {i}. \quad (1 7)
$$
taking  $\pmb{a}_{(k-1)/2+1}$  to the zero-COM hyper-plane, while  $\pmb{x}$  and every other augmented variable remain non-zero-CoM. Before the second coupling transform  $\mathcal{F}_M^{(2)}$ , we apply the ShiftCoM maps in reverse, taking  $\pmb{x}$  back to the zero-COM hyper-plane.
The coupling transforms  $\mathcal{F}_M^{(1)}$  and  $\mathcal{F}_M^{(2)}$  work as previously, except now they transform  $(k + 1) / 2$  sets of variables (holding the other set constant), where for each of the  $k$  variables, each of the  $n$  nodes get their own local reference frame and inner flow transform  $\tau_{\theta}$ .
B.5 Invariant base and target distributions
We now show the  $\mathrm{SE}(3)\times S_n$  invariance of our extended base distribution choice, which trivial extends to our choice of augmented target distribution. Consider the choice  $\pi (\pmb {a}|\pmb {x}) = \mathcal{N}(\pmb {a};\pmb {x},\eta^2 I)$  for  $\eta^2$  a scalar hyperparameter. With this choice, the constraint of interest is satisfied since the Isotropic Gaussian is rationally invariant and centered at  $\pmb{x}$ . To see this, let  $t\in \mathrm{SE}(3)$  be decomposed
into two components  $t = (R, u)$  where  $R \in \mathrm{SO}(3)$  is a  $3 \times 3$  rotation matrix and  $u \in \mathbb{R}^3$  represents a translation,