Self-Attention Amortized Distributional Projection Optimization for Sliced Wasserstein Point-Cloud Reconstruction

We denote a point-cloud of $m$ points $x_{1},\ldots,x_{m}\in\mathbb{R}^{d}$ ($d\geq 1$) as $X=(x_{1},\ldots,x_{m})\in\mathbb{R}^{dm}$ which is a vector of a concatenation of all points in the point-cloud. We denote the set of all possible point-clouds as $\mathcal{X}\subset\mathbb{R}^{dm}$.

Permutation invariant metric space. Given a permutation one-to-one mapping function $\sigma:[m]\to[m]$, we have $\sigma(X)\in\mathcal{X}$ for all $X\in\mathcal{X}$. Moreover, we need a metric $\mathcal{D}:\mathcal{X}\times\mathcal{X}\to\mathbb{R}^{+}$ such that $\mathcal{D}(X,\sigma(X))=0$ for all $X\in\mathcal{X}$ where $\sigma(X)=(x_{\sigma(1)},\ldots,x_{\sigma(m)})$. Here, $\mathcal{D}$ is a metric, namely, it needs to satisfy the non-negativity, symmetry, triangle inequality, and identity property. The pair $(\mathcal{X},\mathcal{D})$ forms a point-cloud metric space.

Learning representation via reconstruction. The raw representation of point-clouds is hard to work with in applications due to the complicated metric space. Therefore, a famous approach is to map point-clouds to points in a different space e.g., Euclidean, which is easier to apply machine learning algorithms. In more detail, we want to estimate a function $f_{\phi}:\mathcal{X}\to\mathcal{Z}$ ($\phi\in\Phi$) where $\mathcal{Z}$ is a set that belongs to another metric space. Then, we can apply machine learning algorithms on $\mathcal{Z}$ instead of $\mathcal{X}$. The most well-known and effective way to estimate the function $f_{\phi}$ is through reconstruction loss. Namely, we estimate $f_{\phi}$ jointly with a function $g_{\gamma}:\mathcal{Z}\to\mathcal{X}$ ($\gamma\in\Gamma$) given a point-cloud dataset $p(X)$ (distribution over $\mathcal{X}$) by minimizing the objective:

The loss $\mathbb{E}_{X\sim p(X)}\mathcal{D}(X,g_{\gamma}(f_{\phi}(X)))$ is known as the reconstruction loss. If the reconstruction loss is 0, we have $g_{\gamma}=f_{\phi}^{-1}$ p-almost surely. Therefore, we can move from $\mathcal{X}$ to $\mathcal{Z}$ and move back from $\mathcal{Z}$ to $\mathcal{X}$ without losing information through the functions $f_{\phi}$ (referred as the encoder) and $g_{\gamma}$ (referred as the decoder). We show an illustration of the framework (Achlioptas et al., 2018) in Figure 1. After learning how to do the reconstruction well, other point-cloud tasks can be done using the autoencoder (the pair $(f_{\phi},g_{\gamma})$) e.g., shape interpolation, shape editing, shape analogy, shape completion, point-cloud classification, and point-cloud generation (Achlioptas et al., 2018).

We now review some famous choices of the metric $\mathcal{D}$ which are Chamfer distance (Barrow et al., 1977), Wasserstein distance (Villani, 2008), sliced Wasserstein (SW) distance (Bonneel et al., 2015), and max sliced Wasserstein (Max-SW) (Deshpande et al., 2019) distance.

Chamfer distance. For any two point-clouds $X$ and $Y$, the Chamfer distance is defined as follows: $\text{CD}(X,Y)=$

where $|X|$ denotes the number of points in $X$.

Wasserstein distance. Given two probability measures $\mu\in\mathcal{P}_{p}(\mathbb{R}^{d})$ and $\nu\in\mathcal{P}_{p}(\mathbb{R}^{d})$, the Wasserstein distance between $\mu$ and $\nu$ is defined as follows:

where $\Pi(\mu,\nu)$ is set of all couplings whose marginals are $\mu$ and $\nu$ respectively. Since the Wasserstein distance is originally defined on probability measures space, we need to convert a point-cloud $X=(x_{1},\ldots,x_{m})\in\mathcal{X}$ to the corresponding empirical probability measure $P_{X}=\frac{1}{m}\sum_{i=1}^{m}\delta_{x_{i}}\in\mathcal{P}(\mathbb{R}^{d})$. Therefore, we can use $\mathcal{D}(X,Y)=\text{W}_{p}(P_{X},P_{Y})$ for $X,Y\in\mathcal{X}$.

Sliced Wasserstein distance. As discussed, the Wasserstein distance is expensive to compute with the time complexity $\mathcal{O}(m^{3}\log m)$ and the memory complexity $\mathcal{O}(m^{2})$. Therefore, an alternative choice is sliced Wasserstein (SW) distance between two probability measures $\mu\in\mathcal{P}_{p}(\mathbb{R}^{d})$ and $\nu\in\mathcal{P}_{p}(\mathbb{R}^{d})$ is:

The benefit of SW is that $\text{W}_{p}(\theta\sharp\mu,\theta\sharp\nu)$ has a closed-form solution which is

with $F^{-1}$ denotes the inverse CDF function. The expectation is often approximated by Monte Carlo sampling, namely, it is replaced by the average from $\theta_{1},\ldots,\theta_{L}$ that are drawn i.i.d from $\mathcal{U}(\mathbb{S}^{d-1})$. The computational complexity and memory complexity of SW becomes $\mathcal{O}(Lm\log_{2}m)$ and $\mathcal{O}(Lm)$.

Max sliced Wasserstein distance. It is well-known that SW has a lot of less discriminative projections due to the uniform sampling. Therefore, max sliced Wasserstein distance is proposed to use the most discriminative projecting direction. Max sliced Wasserstein (Max-SW) distance (Deshpande et al., 2019) between $\mu\in\mathcal{P}_{p}(\mathbb{R}^{d})$ and $\nu\in\mathcal{P}_{p}(\mathbb{R}^{d})$ is introduced as follows:

Max-SW is often computed by a projected sub-gradient ascent algorithm. When the projected sub-gradient ascent algorithm has $T\geq 1$ iterations, the computation complexity of Max-SW is $\mathcal{O}(Tm\log_{2}m)$ and the memory complexity is $\mathcal{O}(m)$. Both SW and Max-SW are applied successfully in point-cloud reconstruction (Nguyen et al., 2021c).

Amortized Optimization. We start with the definition of amortized optimization.

Amortized Projection Optimization. We now revisit the point-cloud reconstruction objective with $\mathcal{D}(X,Y)=\text{Max-SW}_{p}(P_{X},P_{Y})$:

where the expectation is with respect to $X\sim p(X)$. For each point-cloud $X\in\mathcal{X}$, we need to compute a Max-SW distance with an iterative optimization procedure. Therefore, it is computationally expensive.

Authors in  (Nguyen & Ho, 2022a) propose to use amortized optimization (Shu, 2017; Amos, 2022) to speed up the problem. Instead of solving all optimization problems independently, an amortized model is trained to predict optimal solutions to all problems. In greater detail, given a parametric function $a_{\psi}:\mathcal{X}\times\mathcal{X}\to\mathbb{S}^{d-1}$ ($\psi\in\Psi$), the amortized objective is:

where the expectation is with respect to $X\sim p(X)$, and $\theta_{\psi,\gamma,\phi}=a_{\psi}(X,g_{\gamma}(f_{\phi}(X)))$. The above optimization is solved by an alternative stochastic (projected)-gradient descent-ascent algorithm. Therefore, it is faster to compute in each update iteration of $\phi$ and $\gamma$. It is worth noting that the previous work (Nguyen & Ho, 2022a) considers the generative model application which is unstable and hard to understand. Here, we adapt the framework to the point-cloud reconstruction application which is easier to explore the behavior of amortized optimization. We refer the reader to Algorithms 2-3 in Appendix A.3 for algorithms on training an autoencoder with Max-SW and amortized projection optimization.

Amortized models. Authors in (Nguyen & Ho, 2022a) propose three types of amortized models that are based on the literature on linear models (Christensen, 2002). In particular, the linear amortized model is defined as:

Similarly, the generalized linear amortized model and the non-linear amortized model are defined by injecting non-linearity into the linear model. We review the definitions of the generalized linear amortized model and non-linear amortized model in Definitions 4-5 in Appendix A.1.

Sub-optimality. Despite being faster, amortized optimization often cannot recover the global optimum of optimization problems. Namely, we denote

and $\psi^{\star}=$

Then, it is well-known that the amortization gap $\mathbb{E}_{X\sim p(X)}[c(\theta^{\star}(X),a_{\psi^{\star}}(X,g_{\gamma}(f_{\phi}(X))))]>0$ for a metric $c:\mathbb{S}^{d-1}\times\mathbb{S}^{d-1}\to\mathbb{R}^{+}$. A great amortized model is one that can minimize the amortization gap. However, in the amortized projection optimization setting, we cannot obtain $\theta^{\star}(X)$ since the projected gradient ascent algorithm can only yield the local optimum. Therefore, a careful investigation of the amortization gap is challenging.

The current amortized projection optimization framework is for predicting the “max” projecting direction in Max-SW. However, the projected one-dimensional Wasserstein is only a metric on space of probability measure at the global optimum of Max-SW. Therefore, the local optimum from the projected sub-gradient ascent algorithm (Nietert et al., 2022) and the prediction from the amortized model only yield pseudo-metricity for the projected Wasserstein.

The proof for Proposition 1 is given in Appendix B.1. This result implies that the if reconstruction loss $\mathbb{E}_{X\sim p(X)}[\text{PW}_{p}(P_{X},P_{g_{\gamma}(f_{\phi}(X))};\hat{\theta}(X))=0$, it does not imply $X=g_{\gamma}(f_{\phi}(X))$ for p-almost surely $X\in\mathcal{X}$. Therefore, a local maximum for $\max_{\theta\in\mathbb{S}^{d-1}}$ in Max-SW reconstruction (Equation 8) and the global maximum for $\max_{\psi\in\Psi}$ in amortized Max-SW reconstruction (Equation 9 with a misspecified amortized model) cannot guarantee perfect reconstruction even when their objectives obtain $0$ values.

Amortized Distributional Projection Optimization. To overcome the issue, we propose to replace Max-SW in Equation 8 with the von Mises Fisher distributional sliced Wasserstein (v-DSW) distance (Nguyen et al., 2021b):

where $\text{vMF}(\epsilon,\kappa)$ is the von Mises Fisher distribution with the mean location parameter $\epsilon\in\mathbb{S}^{d-1}$ and the concentration parameter $\kappa>0$, and $\text{v-DSW}_{p}(\mu,\nu;\kappa)=\max_{\epsilon\in\mathbb{S}^{d-1}}\Big{(}\mathbb{E}_{\theta\sim\text{vMF}(\epsilon,\kappa)}\text{W}_{p}^{p}(\theta\sharp\mu,\theta\sharp\nu)\Big{)}^{\frac{1}{p}}$ is von Mises Fisher distributional sliced Wasserstein distance. The optimization can be solved by a stochastic projected gradient ascent algorithm with the vMF reparameterization trick. In particular, $\theta_{1},\ldots,\theta_{L}$ ($L\geq 1$) is sampled i.i.d from $\text{vMF}(\epsilon,\kappa)$ via the reparameterized acceptance-rejection sampling (Davidson et al., 2018a) to approximate $\nabla_{\epsilon}\mathbb{E}_{\text{vMF}(\epsilon,\kappa)}[\text{W}_{p}^{p}(\theta\sharp\mu,\theta\sharp\nu)]$ via Monte Carlo integration. We refer the reader to Section A.2 for more detail about the vMF distribution, its sampling algorithm, its reparameterization trick, and the stochastic gradient estimators. We present a visualization of the difference between the new amortized distributional projection optimization framework and the conventional amortized projection optimization framework in Figure 2. The corresponding amortized objective is:

where $\epsilon_{\psi,\gamma,\phi}=a_{\psi}(X,g_{\gamma}(f_{\phi}(X)))$. The optimization is solved by an alternative stochastic (projected)-gradient descent-ascent algorithm with the vMF reparameterization.

The proof of Theorem 1 is given in Appendix B.2. The proof is based on proving the metricity of the non-optimal von Mises Fisher distributional sliced Wasserstein distance (v-DSW) with the smoothness condition of the vMF distribution. It is worth noting that the proof of metricity of von Mises Fisher distributional sliced Wasserstein distance is new since the original work (Nguyen et al., 2021b) only shows the pseudo-metricity with the global optimality condition. Theorem 1 indicates that a perfect reconstruction can be obtained with a local optimum for $\max_{\epsilon\in\mathbb{S}^{d-1}}$ in v-DSW reconstruction (Equation 3.1) and a local optimum for $\max_{\psi\in\Psi}$ in amortized v-DSW reconstruction (Equation 3.1).

Comparison with SW and Max-SW: When $\kappa\to 0$, the vMF distribution converges weakly to the uniform distribution over the unit hypersphere. Hence, we can get back the conventional sliced Wasserstein reconstruction in both Equation 3.1 and Equation 3.1. When $\kappa\to\infty$, vMF distribution converges weakly to the Dirac delta at the location parameter. Therefore, we obtain Max-SW reconstruction and amortized Max-SW reconstruction in Equation 3.1 and Equation 3.1, respectively. However, when $0<\kappa<\infty$, v-DSW reconstruction and amortized v-DSW reconstruction can find a region of discriminative projecting directions while preserving the metricity for perfect reconstruction.

We now discuss the parameterization of the amortized model for amortized optimization.

Permutation Invariance and Symmetry. Let $X$ and $Y$ be two point-clouds, the optimal slicing distribution $\text{vMF}(\epsilon^{\star},\kappa)$ of v-DSW between $P_{X}$ and $P_{Y}$ can be obtained by running Algorithm 4 in Appendix A.3. Clearly, $\text{vMF}(\epsilon^{\star},\kappa)$ is invariant to the permutation of the supports since $P_{\sigma(X)}=P_{X}$ and $P_{\sigma(Y)}=P_{Y}$ for a permutation function $\sigma$. Moreover, the optimal slicing distribution $\text{vMF}(\epsilon^{\star},\kappa)$ is also unchanged when we exchange $P_{X}$ and $P_{Y}$ since v-DSW is symmetric. However, the current amortized models (see Definition 2, Definitions 4-5 in Appendix A.1) are not permutation invariant and symmetric, namely, $a_{\psi}(X,Y)\neq a_{\psi}(X,\sigma(Y))$ and $a_{\psi}(X,Y)\neq a_{\psi}(Y,X)$ . Therefore, the current amortized models could be strongly misspecified. We show a visualization of an amortized model that is not symmetric and permutation invariant in Figure 3. To address the issue, we propose amortized models that are symmetric and permutation invariant based on the self-attention mechanism.

Self-Attention Mechanism. Attention is well-known for its effectiveness in learning long-range dependencies when data are sequences such as text (Devlin et al., 2019; Liu et al., 2019; Brown et al., 2020) or speech (Li et al., 2019; Wang et al., 2020b). This mechanism was then successfully generalized to other data types including image (Carion et al., 2020; Dosovitskiy et al., 2020), video (Sun et al., 2019), graph (Dwivedi & Bresson, 2021), point-cloud (Zhao et al., 2021; Guo et al., 2021), to name a few. We now revisit the attention mechanism (Vaswani et al., 2017). Given $Q,K\in{\mathbb{R}}^{m\times d_{k}},V\in{\mathbb{R}}^{m\times d_{v}}$, the scaled dot-product attention operator is defined as:

where $\text{softmax}_{\rm{row}}$ denotes the row-wise softmax function. In the self-attention mechanism, the query matrix Q, the key matrix K, and the value matrix V are usually computed by projecting the input sequence X into different subspaces. Thus, the self-attention mechanism is given as follows. Given $X\in{\mathbb{R}}^{m\times d}$, the self-attention operator is:

where $W_{q},W_{k}\in{\mathbb{R}}^{d\times d_{k}},W_{v}\in{\mathbb{R}}^{d\times d_{v}}$ and $\zeta=(W_{q},W_{k},W_{v})$. The self-attention operator is infamous for its quadratic memory and computational costs. In particular, given an input sequence of length $m$, both the time and space complexity are ${\mathcal{O}}(m^{2})$. Since we focus on the sliced Wasserstein setting where the computational complexity should be at most $\mathcal{O}(m\log m)$, the conventional self-attention is not appropriate. Several works (Li et al., 2020; Katharopoulos et al., 2020; Wang et al., 2020a; Shen et al., 2021) have been proposed to reduce the overall complexity from ${\mathcal{O}}(m^{2})$ to ${\mathcal{O}}(m)$. In this paper, we utilize two linear complexity variants of attention which are efficient attention (Shen et al., 2021) and linear attention (Wang et al., 2020a). Given $X\in{\mathbb{R}}^{m\times d}$, the efficient self-attention is defined as:

where $W_{q},W_{k}\in{\mathbb{R}}^{d\times d_{k}},W_{v}\in{\mathbb{R}}^{d\times d_{v}}$, $\zeta=(W_{q},W_{k},W_{v})$, and $\text{softmax}_{\rm{col}}$ denotes applying the softmax function column-wise. The linear self-attention is:

where $W_{q},W_{k2}\in{\mathbb{R}}^{d\times d_{k}},W_{v2}\in{\mathbb{R}}^{d\times d_{v}}$, $W_{k1},W_{v1}\in{\mathbb{R}}^{k\times n}$, and $\zeta=(W_{q},W_{k1},W_{k2},W_{v1},W_{v2})$. The projected dimension $k$ is chosen such that $m\gg k$ to reduce the memory and space consumption significantly.

Self-Attention Amortized Models: Based on the self-attention mechanism, we introduce the self-attention amortized model which is permutation invariant and symmetric. Formally, the self-attention amortized model is defined as:

By replacing the conventional self-attention with the linear self-attention and the efficient self-attention, we obtain the linear self-attention amortized model and the efficient self-attention amortized model.

The proof of Proposition 2 is given in Appendix B.3. The symmetry follows directly from the definition of the self-attention amortized models. The permutation invariance is proved by showing that the self-attention operators combined with average pooling are permutation invariant.

Comparison with Set Transformer. The authors in (Lee et al., 2019b) also proposed a method to guarantee the permutation invariant of sets. There are two main differences between our works and theirs. Firstly, Set Transformer introduced a new attention mechanism and a new Transformer architecture while we only present an approach to apply any attention mechanism to preserve the permutation invariance property of amortized models. Secondly, Set Transformer maintains the permutation invariance property by using a learnable multi-head attention as the aggregation scheme. We instead still rely on average pooling, a conventional permutation invariant aggregation scheme, to accumulate features learned by self-attention operations. Nevertheless, our works are orthogonal to Set Transformer, in other words, it is possible to apply techniques in Set Transformer to our attention-based amortized models. We leave this investigation for future work.