SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Laurence I. Midgley*
lim24@cam.ac.uk
Vincent Stimper*
vs488@cam.ac.uk
Javier Antorán*
ja666@cam.ac.uk
ebm32@cam.ac.uk
Max Planck Institute
for Intelligent Systems
bs@tue.mpg.de
Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.
Modeling the distribution of a molecule's configurations at equilibrium, known as the Boltzmann distribution, is a promising application of deep generative models [Noé et al., 2019]. While the unnormalized density of the Boltzmann distribution can be obtained via physical modeling, sampling from it typically requires molecular dynamics (MD) simulations, which are expensive and produce correlated samples. A promising alternative is to rely on surrogate deep generative models, known as Boltzmann generators. We can draw independent samples from these and debias any expectations
Code: https://github.com/lollcat/se3-augmented-coupling-flows
estimated with the samples via importance weighting. The Boltzmann distribution typically admits rotation and translation symmetries, known as the Special Euclidean group SE(3), as well as permutation symmetry. These constraints are important to incorporate into the model as they improve training efficiency and generalization [Cohen and Welling, 2017, Batzner et al., 2022, Kohler et al., 2020]. Other key desiderata of Boltzmann generators are that they allow for fast sampling and density evaluation. These are necessary for energy-based training, that is, training using the Boltzmann distribution's unnormalized density [Noé et al., 2019, Stimper et al., 2022, Midgley et al., 2023]. Training by energy is critical, as it prevents the model quality from being constrained by the quality and quantity of MD samples.
Existing coupling flows which approximate the Boltzmann distribution of molecules are at least partially defined over internal coordinates, i.e. bond distances, bond angles, and dihedral angles [Wu et al., 2020, Campbell et al., 2021, Kohler et al., 2023a, Midgley et al., 2023], which are SE(3) invariant. However, the definition of these depends on the molecular graph and they are non-unique for most graphs. Furthermore, models parameterized in internal coordinates struggle to capture interactions among nodes far apart in the graph and cannot capture the permutation invariance of some atoms. Thus, they are not suitable for particle systems such as LJ13 [Kohler et al., 2020]. SE(3) equivariance constraints have been applied to continuous normalizing flows (CNFs) operating on Cartesian coordinates [Kohler et al., 2020, Satorras et al., 2021a], and to the closely related diffusion models [Xu et al., 2022, Hoogeboom et al., 2022, Yim et al., 2023]. These models are built upon SE(3) equivariant graph neural networks (GNNs) [Satorras et al., 2021b, Geiger and Smidt, 2022, Batzner et al., 2022]. These architectures can be applied to any molecular graph [Jing et al., 2022], enabling a single generative model to generalize across many molecules. Alas, sampling and evaluating the density of CNFs and diffusion models typically requires thousands of neural network evaluations [Xiao et al., 2022], preventing them from being trained by energy. As such, presently no Boltzmann generator exists that (i) acts on Euclidean coordinates of atoms (ii) enforces SE(3) equivariance, and (iii) allows for fast sampling.
To address this gap, we propose a flexible SE(3) equivariant coupling flow that operates on the Cartesian coordinates of atoms, allowing for fast sampling and density evaluation. Our contributions are:
- We extend coupling layers to be SE(3) equivariant by augmenting their input space with auxiliary variables [Huang et al., 2020] which can be acted upon on by SE(3). We update the atom positions conditioned on the auxiliary variables by first projecting the atoms into an SE(3)-invariant space and then applying a standard normalizing flow transform before projecting its output back onto the equivariant space.
- We demonstrate that, when trained by maximum likelihood, our flow matches the performance of existing SE(3) CNFs and diffusion models as well as coupling flows operating on internal coordinates on molecular generation tasks. Our flow is more than 10 times faster to sample from than SE(3) CNFs and diffusion models. Concurrently with Klein et al. [2023b], we are the first to learn the full Boltzmann distribution of alanine dipeptide solely in Cartesian coordinates.
- We demonstrate our flow in the energy-based training setting on DW4 and LJ13, where parameters are learned using only the molecular energy function. Energy-based training of CNFs or diffusion models is intractable due to slow sampling and density evaluation. Flows operating on internal coordinates are not able to capture the permutation invariance of these problems. Hence, our flow is the only existing permutation and SE(3) equivariant method that can tractably be applied there.
2 Background: coupling flows and invariant models
2.1 Normalizing flows and coupling transforms
A (discrete-time) normalizing flow is a flexible parametric family of densities on  $\mathcal{X}$  as the pushforward of a base density  $q_{0}$  along an invertible automorphism  $f_{\theta}:\mathcal{X}\to \mathcal{X}$  with parameters  $\theta \in \Theta$  [Papamakarios et al., 2021]. The density is given by the change of variable formula:
$$
q _ {\theta} (x) = q _ {0} \left(f ^ {- 1} (x)\right) \left| \det \frac {\partial f _ {\theta} ^ {- 1} (x)}{\partial x} \right|. \tag {1}
$$
We can efficiently sample from the flow by sampling from  $q_{0}$  and mapping these samples through  $f_{\theta}$  in a single forward pass. A popular way to construct  $f_{\theta}$  is to use coupling transforms. The  $D$
dimensional input  $\pmb{x} \in \mathcal{X}$  is split into two sets, transforming the first set conditional on the second, while leaving the second set unchanged:
$$
\boldsymbol {y} _ {1: d} = \mathcal {T} (\boldsymbol {x} _ {1: d}; \boldsymbol {x} _ {d + 1: D}),
$$
$$
\boldsymbol {y} _ {d + 1: D} = \boldsymbol {x} _ {d + 1: D}. \tag {2}
$$
They induce a lower triangular Jacobian, such that its determinant becomes  $\left| \partial \mathcal{T}(\boldsymbol{x}_{1:d}; \boldsymbol{x}_{d+1:D}) / \partial \boldsymbol{x}_{1:d} \right|$ . Further, choosing  $\mathcal{T}$  to have an easy to compute determinant, such as an elementwise transformation [Dinh et al., 2015, 2017, Durkan et al., 2019], allows for fast density evaluation and sampling at low computational cost.
2.2 Equivalence and invariance for coupling flow models of molecular conformations
Throughout, we deal with observations of an  $n$ -body system represented by a matrix  $\boldsymbol{x} = [x^{1},\ldots ,x^{n}]\in \mathcal{X} = \mathbb{R}^{3\times n}$ , where the rows index Cartesian coordinates and the columns index individual particles. We seek to construct flows on  $\mathcal{X}$  endowed with the symmetries present in molecular energy functions. These are invariant to rotations and translations of  $\boldsymbol{x}$  (SE(3)), and to permutation of atoms of the same type  $(S_{n})$ . We will formalize them in this section.
Symmetry groups The special Euclidean group  $\mathrm{SE}(3)$  is the set of orientation preserving rigid transformations in Euclidean space. Its elements  $t\in \mathrm{SE}(3)$  can 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; for a coordinate  $v\in \mathbb{R}^3$ ,  $t\cdot v = Rv + u$  denotes the action of  $t$  on  $v$ . The symmetric group  $S_{n}$  defined over a set of  $n$  atoms consists of all  $n!$  permutations that can be performed with said atoms. Its elements  $\sigma \in S_n$  act on an  $n$ -body system as  $\sigma \cdot \pmb {x} = [x^{\sigma (1)},\dots ,x^{\sigma (n)}]$ .
Equivariant maps A map  $f: \mathcal{X} \to \mathcal{Y}$  is said to be  $G$ -equivariant if it commutes with the group action  $\cdot$ , i.e. if for any  $x \in \mathcal{X}$  and  $g \in G$  we have  $f(g \cdot x) = g \cdot f(x)$ . Invariance is a special case where for any  $x \in \mathcal{X}$  and  $g \in G$  we have  $f(g \cdot x) = f(x)$ . There has been a plethora of recent work on constructing graph neural network functions equivariant to the action of  $G = \mathrm{SE}(3) \times S_n$  [e.g. Thomas et al., 2018, Satorras et al., 2021b, Geiger and Smidt, 2022], which we will leverage to construct our equivariant coupling flow model.
Invariant density A density  $p: \mathcal{X} \to \mathbb{R}_+$  is  $G$ -invariant if for any  $x \in \mathcal{X}$  and  $g \in G$  we have  $p(g \cdot x) = p(x)$ . Combining an invariant base density  $q_0$  with an equivariant invertible transform  $f$ , as in (1), yields an invariant flow density [Papamakarios et al., 2021, Kohler et al., 2020]. This gives a practical way to design invariant densities models.
Challenges in constructing  $\mathrm{SO}(3) \times S_{n}$  invariant flow models Unfortunately, no coupling transform can be simultaneously equivariant to both permutation of the particles and their rotations; coupling splits must be performed either across particles or spatial dimensions which would break either permutation or rotational symmetry [Köhler et al., 2020, Bose et al., 2022].
Furthermore, there does not exist a translation invariant probability measure, as any such measure would be proportional to the Lebesgue measure and therefore not have unit volume. This precludes us from defining an invariant base distribution directly on  $\mathcal{X}$ . Fortunately, Proposition A.1 and its converse allow us to disintegrate the probability measure into a translational measure proportional to the Lebesgue measure and an SO(3)-invariant probability measure on the subspace of  $\mathbb{R}^{3\times n}$  with zero center of mass. We can drop the former and only model the latter.
3 Method:  $\operatorname{SE}(3) \times S_n$  equivariant augmented coupling flow model
This section describes our main contribution, an  $\mathrm{SE}(3)\times S_n$  equivariant coupling flow. We first lay the groundwork for achieving translation invariance by defining our flow density on a lower-dimensional "zero Center of Mass (CoM)" space. To preserve permutation and rotation equivariance, we leverage the augmented flow framework of Huang et al. [2020]. Specifically, we use sets of augmented variables as a pivot for coupling transforms. Sec. 3.1 introduces a novel class of coupling transforms that achieve the aforementioned permutation and rotation equivariance by operating on atoms projected using a set of equivariant bases. Sec. 3.2 describes our choice of invariant base distribution and, finally, in Sec. 3.3, we discuss several schemes to train the augmented flow from either samples or energy functions and how to perform efficient density evaluation.
Figure 1: Illustration of the equivariant coupling layer of our augmented normalizing flow, where our variable with zero center of mass (CoM)  $\tilde{\pmb{x}}$  is transformed with the augmented variable  $\pmb{a}$ .
Translation invariance is obtained by modelling the data on the quotient space  $\mathbb{R}^{3\times n} / \mathbb{R}^3\triangleq \tilde{\mathcal{X}}\subseteq \mathcal{X}$  where all  $n$  -body systems that only differ by a translation are "glued" together, i.e. where  $\pmb {x}\sim \pmb{x}'$  if  $\pmb {x} = \pmb{x}^{\prime} + p$  with  $p\in \mathbb{R}^3$  . Constructing a parametric probabilistic model over  $\mathcal{X}$  , automatically endowing it with translation invariance. In practice, we still work with Cartesian coordinates, but center the data so as to zero its CoM:  $\tilde{\pmb{x}}\triangleq \pmb {x} - \bar{\pmb{x}}$  with  $\bar{\pmb{x}}\triangleq \frac{1}{n}\sum_{i = 1}^{n}[\pmb {x}]^{i}$  . Thus,  $\tilde{\pmb{x}}$  lies on  $\tilde{\mathcal{X}}$  , an  $(n - 1)\times 3$  dimensional hyperplane embedded in  $\mathcal{X}$
Augmented variable pivoted coupling To allow for simultaneously permutation and rotation equivariant coupling transforms, we introduce augmented variables  $\mathbf{a} \in \mathcal{A}$ . Our coupling layers update the particle positions  $\mathbf{x}$  conditioned on  $\mathbf{a}$  and vice-versa. The augmented variables need to "behave" similarly to  $\mathbf{x}$ , in that they can also be acted upon by elements of  $\mathrm{SE}(3) \times S_n$ . We achieve this by choosing  $\mathbf{a}$  to be a set of  $k$  of observation-sized arrays  $\mathcal{A} = \mathcal{X}^k$ , which we will discuss further in App. B.4. Importantly, we do not restrict  $\mathbf{a}$  to be zero-CoM.
Invariant flow density on the extended space We parameterize a density  $q$  over the extended space  $\tilde{\mathcal{X}} \times \mathcal{A}$  w.r.t. the product measure  $\lambda_{\tilde{\mathcal{X}}} \otimes \lambda_{\mathcal{A}}$ , where  $\lambda_{\tilde{\mathcal{X}}} \in \mathcal{P}(\tilde{\mathcal{X}})$  and  $\lambda_{\mathcal{A}} \in \mathcal{P}(\mathcal{A})$  respectively denote the Lebesgue measure on  $\tilde{\mathcal{X}}$  and  $\mathcal{A}$ . We use  $q(\boldsymbol{x}, \boldsymbol{a})$  as a shorthand for the density of the corresponding zero-CoM projection  $q(\tilde{\boldsymbol{x}}, \boldsymbol{a})$ . The density  $q$  is constructed to be invariant to the action of  $G \triangleq \mathrm{SO}(3) \times S_n$  when simultaneously applied to the observed and augmented variables, that is,  $q(\boldsymbol{x}, \boldsymbol{a}) = q(g \cdot \boldsymbol{x}, g \cdot \boldsymbol{a})$  for any  $g \in G$ . We aim to construct a  $G$ -equivariant flow  $f: \tilde{\mathcal{X}} \times \mathcal{A} \to \tilde{\mathcal{X}} \times \mathcal{A}$  on this extended space, and combine it with  $q_0: \tilde{\mathcal{X}} \times \mathcal{A} \to \mathbb{R}^+$ , a  $G$ -invariant base density function, to yield the invariant flow density
$$
q (\boldsymbol {x}, \boldsymbol {a}) = q _ {0} \left(f ^ {- 1} (\boldsymbol {x}, \boldsymbol {a})\right) \left| \det  \frac {\partial f ^ {- 1} (\boldsymbol {x} , \boldsymbol {a})}{\partial (\boldsymbol {x} , \boldsymbol {a})} \right|. \tag {3}
$$
Proposition 3.1 (Invariant marginal). Assume  $q: \mathcal{X} \times \mathcal{A} \to \mathbb{R}_+$  is a  $G$ -invariant density over the probability space  $(\mathcal{X} \times \mathcal{A}, \lambda_{\mathcal{X}} \otimes \lambda_{\mathcal{A}})$ , then  $q_x \triangleq \int_{\mathcal{A}} q(\cdot, \mathbf{a}) \lambda_{\mathcal{A}}(\mathrm{d}\mathbf{a}): \mathcal{X} \to \mathbb{R}_+$  is a  $G$ -invariant density w.r.t. to the measure  $\lambda_{\mathcal{X}}$ .
Proof. For any  $g \in G$ ,  $\pmb{x} \in \mathcal{X}$  and  $\pmb{a} \in \mathcal{A}$
$$
q _ {\boldsymbol {x}} (g \cdot \boldsymbol {x}) = \int_ {\mathcal {A}} q (g \cdot \boldsymbol {x}, \boldsymbol {a}) \lambda_ {\mathcal {A}} (\mathrm {d} \boldsymbol {a}) = \int_ {g ^ {- 1} \mathcal {A}} q (g \cdot \boldsymbol {x}, g \cdot \boldsymbol {a}) \lambda_ {\mathcal {A}} (\mathrm {d} g \cdot \boldsymbol {a}) = \int_ {\mathcal {A}} q (\boldsymbol {x}, \boldsymbol {a}) \lambda_ {\mathcal {A}} (\mathrm {d} \boldsymbol {a}) = q _ {\boldsymbol {x}} (\boldsymbol {x}),
$$
where we used the  $G$ -invariance of  $q$  and of the measure  $\lambda_{\mathcal{A}}$ , as well as  $g^{-1}\mathcal{A} = \mathcal{A}$ .
3.1 SE(3) and permutation equivariant coupling transform
We now derive our  $\mathrm{SE}(3)\times S_n$  equivariant map  $f:\tilde{\mathcal{X}}\times \mathcal{A}\to \tilde{\mathcal{X}}\times \mathcal{A}$  defined on the extended space. We introduce two modules, a shift-CoM transform, swapping the center of mass between the observed and augmented variables, and an equivariant core transformation, which updates  $\tilde{\pmb{x}}$  conditional on  $\pmb{a}$  and vice versa. Composing them yields our equivariant coupling layer illustrated in Fig. 1.
$\mathbf{SE}(3) \times S_{n}$  equivariant coupling We dissociate the equivariance constraint from the flexible parameterized flow transformation by (1) projecting the atoms' Cartesian coordinates into a learned,
local (per-atom) invariant space, (2) applying the flexible flow to the invariant representation of the atoms' positions, and (3) then projecting back into the original Cartesian space. Specifically, we construct a core coupling transformation that composes (a) an invariant map  $\gamma : \Xi \times \mathcal{X} \to \mathcal{Y}$ , where  $\mathcal{Y}$  is isomorphic to  $\mathcal{X}$ , and  $r \in \Xi$  parametrizes the map. We denote the parametrized map as  $\gamma_r$ . It is followed by (b) a standard flexible normalizing flow transform, e.g. a neural spline,  $\tau : \mathcal{Y} \to \mathcal{Y}$ , and (c) the inverse map  $\gamma^{-1}$ . Denoting the inputs with superscript  $\ell$  and with  $\ell + 1$  for outputs, our core transformation  $\mathcal{F} : (\boldsymbol{x}^\ell; \boldsymbol{a}^\ell) \mapsto (\boldsymbol{x}^{\ell + 1}, \boldsymbol{a}^{\ell + 1})$  is given as
$$
\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} \quad (\boldsymbol {r}, \theta) = h (\boldsymbol {a} ^ {\ell}),
$$
$$
\boldsymbol {a} ^ {\ell + 1} = \boldsymbol {a} ^ {\ell}. \tag {4}
$$
Here,  $h$  is a (graph) neural network that returns a set of equivariant reference vectors  $\mathbf{r}$ , which parametrize the map  $\gamma_{\mathbf{r}}$ , and invariant parameters  $\theta$ .  $\mathcal{V}$  is a rotation invariant space. This means that any rotations applied to the inputs will be cancelled by  $\gamma_{\mathbf{r}}$ , i.e.  $\gamma_{g\cdot r} = \gamma_r\cdot g^{-1}$  or equivalently  $(\gamma_{g\cdot r})^{-1} = g\cdot \gamma_r^{-1}$  for all  $g\in \mathrm{SO}(3)$ . We use the inverse projection  $\gamma_r^{-1}$  to map the invariant features back to equivariant features. The function  $\tau_{\theta}$  is a standard invertible inner-transformation such as an affine or spline based transform, that we apply to the invariant features [Papamakarios et al., 2021].
We now show that the above described transform is rotation and permutation equivariant.
Proposition 3.2 (Equivariant augmented coupling flow). If  $h: \mathcal{A} \to \mathcal{X}^n \times \Theta^n$  is SO(3)-equivariant for its first output, SO(3)-invariant for its second, and  $S_n$  equivariant for both, and  $(\gamma_{g\cdot r})^{-1} = g\cdot \gamma_r^{-1}$  for all  $g \in \mathrm{SO}(3)$ , the transform  $\mathcal{F}$  given by (4) is SO(3)  $\times S_n$  equivariant.
Proof. For SO(3): We first notice that  $h(g \cdot \mathbf{a}) = (g \cdot \mathbf{r}, \theta)$ , and then since  $(\gamma_{g \cdot \mathbf{r}})^{-1} = g \cdot \gamma_{\mathbf{r}}^{-1}$  we have  $\mathcal{F}(g \cdot \mathbf{x}, g \cdot \mathbf{a}) = (\gamma_{g \cdot \mathbf{r}})^{-1} \cdot \tau_{\theta}(\gamma_{g \cdot \mathbf{r}} \cdot g \cdot \mathbf{x}^{\ell}) = g \cdot \gamma_{\mathbf{r}}^{-1} \cdot \tau_{\theta}(\gamma_{\mathbf{r}} \cdot g^{-1} \cdot g \cdot \mathbf{x}^{\ell}) = g \cdot \mathcal{F}(\mathbf{x}, \mathbf{a})$ .
For  $S_{n}$ : We first note that  $h(\sigma \cdot \pmb{a}) = (\sigma \cdot \pmb{r}, \sigma \cdot \theta)$ . Then, using that  $\gamma_{\pmb{r}}$  and  $\tau$  act on  $\pmb{x}$  atom-wise, we have  $\mathcal{F}(\sigma \cdot \pmb{x}, \sigma \cdot \pmb{a}) = \gamma_{\sigma \cdot \pmb{r}}^{-1} \cdot \tau_{\sigma \cdot \theta} (\gamma_{\sigma \cdot \pmb{r}} \cdot (\sigma \cdot \pmb{x})) = (\sigma \cdot \gamma_{\pmb{r}}^{-1}) \cdot (\sigma \cdot \tau_{\theta}) ((\sigma \cdot \gamma_{\pmb{r}}) \cdot (\sigma \cdot \pmb{x})) = \sigma \cdot \mathcal{F}(\pmb{x}, \pmb{a})$ .
For the Jacobian of the coupling described above to be well-defined, the variable being transformed must be non-zero CoM (see App. B.1 for a derivation). Thus, although our observations live on  $\tilde{\mathcal{X}}$ , for now, assume that the inputs to the transform are not zero-CoM and we will deal with this assumption in the following paragraphs. This choice also allows us to use standard equivariant GNNs for  $h$  [Satorras et al., 2021b, Geiger and Smidt, 2022] which leverage per-node features defined in the ambient space, such as atom type and molecular graph connectivity.
Choices of projection  $\gamma$  The equivariant vectors  $r$  parameterize a local (per-atom) SO(3) equivariant reference frame used in the projection  $\gamma_{r}$ . We introduce three different projection strategies. (i) The first strategy, is for  $r$  to parameterize frame composed of an origin and orthonormal rotation matrix into which we project each atom's positions. We then take  $\tau$  to be a dimension-wise transformation for each of the projected atoms' coordinates. We dub this method CARTESIAN-PROJ. (ii) Alternatively, we let  $r$  parameterize a origin, zenith direction and azimuth direction for spherical coordinates, as in Liu et al. [2022]. We then apply elementwise transforms to each atom's radius, polar angle and azimuthal angle. We call this SPHERICAL-PROJ. (iii) Lastly, consider a variant of SPHERICAL-PROJ where just the radius is transformed and the polar and azimuth angles are held constant. Here,  $r$  parameterizes a single reference point, a per-atom origin. We refer to this last variant as VECTOR-PROJ.
Architectural details For the transformations applied in the invariant projected space we consider affine mappings [Dinh et al., 2017] and monotonic rational-quadratic splines [Durkan et al., 2019]. Additionally, to limit computational cost, we have our GNNs  $h$  output  $M$  sets of reference vectors  $r$  and invariant parameters  $\theta$ . These parametrize  $M$  core coupling transformations with a single GNN forward pass. For the CARTESIAN-PROJ and SPHERICAL-PROJ flow we include a loss term that discourages certain reference vectors from being collinear, which improves the projection's stability. We provide further details for this and the various projection types in App. B.3.
Center of mass shift The shift-CoM transform allows us to apply the aforementioned  $\mathrm{SE}(3) \times S_{n}$  equivariant coupling in the ambient space rather than zero-COM subspace. In particular, before transforming our observed vector  $\tilde{\pmb{x}} \in \tilde{\mathcal{X}}$  with  $\mathcal{F}$ , we lift it onto  $\mathcal{X}$ . We achieve this by swapping the center of mass between  $\tilde{\pmb{x}}$  and  $\pmb{a}$ . For now, assume  $\mathcal{A} = \mathcal{X}$ , i.e.  $k = 1$ , with App. B.4 providing details for  $k > 1$ . Letting  $\tilde{\mathcal{A}} \subseteq \mathcal{A}$  be the subspace where all augmented variables that differ by
a translation occupy the same point, and  $\tilde{\pmb{a}}\in \tilde{\mathcal{A}}$  be defined analogously to  $\tilde{\pmb{x}}$ , we apply the map ShiftCoM:  $\tilde{\mathcal{X}}\times \mathcal{A}\mapsto \mathcal{X}\times \tilde{\mathcal{A}}$  which acts on both of its arguments by subtracting from each of them the latter's CoM, that is,