Mirror variational transport: a particle-based algorithm for distributional optimization on constrained domains

Abstract

We consider the optimization problem of minimizing an objective functional, which admits a variational form and is defined over probability distributions on a constrained domain, which poses challenges to both theoretical analysis and algorithmic design. We propose Mirror Variational Transport (mirrorVT), which uses a set of samples, or particles, to represent the approximating distribution and deterministically updates the particles to optimize the functional. To deal with the constrained domain, in each iteration, mirrorVT maps the particles to an unconstrained dual domain, induced by a mirror map, and then approximately performs Wasserstein Gradient Descent on the manifold of distributions defined over the dual space to update each particle by a specified direction. At the end of each iteration, particles are mapped back to the original constrained domain. Through experiments on synthetic and real world data sets, we demonstrate the effectiveness of mirrorVT for the distributional optimization on the constrained domain. We also analyze its theoretical properties and characterize its convergence to the global minimum of the objective functional.

Appendix 1

1.1 Appendix 1.1: Proof of Theorem 2

Proof

We have assumed $F(p)=\text {KL}(p||p^{*})$ for $p,p^{*}\in \mathscr {P}_{2}(\mathscr {X})$, so $G(q)=\text {KL}(q||q^{*})$, for $q,q^{*}\in \mathscr {P}_{2}(\mathscr {Y})$. By the definition of the first variation of a functional, we have:

$$\begin{aligned} \frac{\textrm{d}}{\textrm{d}\epsilon } G(q + \epsilon \chi )\bigg |_{\epsilon =0}=\int _{\mathscr {Y}}\frac{\partial G}{\partial q}({\textbf {y}})\chi ({\textbf {y}})\textrm{d}{} {\textbf {y}} \text {, for all } \chi \in \mathscr {P}_{2}(\mathscr {Y}) \end{aligned}$$

We can compute the left-hand side as follows:

$$\begin{aligned} \frac{\textrm{d}}{\textrm{d}\epsilon } G(q + \epsilon \chi )\bigg |_{\epsilon =0}&=\frac{\textrm{d}}{\textrm{d}\epsilon } \text {KL}(q + \epsilon \chi || q^{*})\bigg |_{\epsilon =0}\\&= \frac{\textrm{d}}{\textrm{d}\epsilon }\int (q + \epsilon \chi )\log \left( \frac{q+\epsilon \chi }{q^{*}} \right) \textrm{d}{} {\textbf {y}}\bigg |_{\epsilon =0}\\&= \int \log \frac{q}{q^{*}}({\textbf {y}})\chi ({\textbf {y}})\textrm{d}{} {\textbf {y}} \end{aligned}$$

which indicates that $\partial G/\partial q= \log q - \log q^{*}$. For t-th iteration, the update direction $v_{t}$ is given by:

$$\begin{aligned} \begin{aligned} v_{t}({\textbf {x}})&= \nabla ^{2}\varphi ({\textbf {x}})^{-1}\nabla f^{*}_{t}({\textbf {x}})=\nabla g^{*}_{t}({\textbf {y}})\\&= \nabla \log q_{t}({\textbf {y}}) - \nabla \log q^{*}({\textbf {y}}) \end{aligned} \end{aligned}$$

(30)

for all ${\textbf {x}}\in \mathscr {X}, {\textbf {y}}=\nabla \varphi ({\textbf {x}})\in \mathscr {Y}$. By applying the integral operator $\mathscr {L}_{k, p_{t}}$ (see Definition 1) to $v_{t}$, we obtain:

$$\begin{aligned} \begin{aligned} \mathscr {L}_{k, p_{t}} v_{t}({\textbf {x}})&= \int _{\mathscr {X}}k({\textbf {x}}, {\textbf {x}}^\prime )v_{t}({\textbf {x}}^\prime )p_{t}({\textbf {x}}^\prime ) \textrm{d}{} {\textbf {x}}^\prime \\&= \int _{\mathscr {Y}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla \log q_{t}({\textbf {y}}^\prime )q_{t}({\textbf {y}}^\prime )\textrm{d}{} {\textbf {y}}^\prime -\int _{\mathscr {Y}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla \log q^{*}({\textbf {y}}^\prime )q_{t}({\textbf {y}}^\prime )\textrm{d} {\textbf {y}}^\prime \\&= \int _{\mathscr {Y}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla q_{t}({\textbf {y}}^\prime )\textrm{d}{} {\textbf {y}}^\prime -\int _{\mathscr {Y}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla \log q^{*}({\textbf {y}}^\prime )q_{t}({\textbf {y}}^\prime )\textrm{d}{} {\textbf {y}}^\prime \\&= -\int _{\mathscr {Y}}\nabla k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime ) q_{t}({\textbf {y}}^\prime ) \textrm{d}{} {\textbf {y}}^\prime -\int _{\mathscr {Y}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla \log q^{*}({\textbf {y}}^\prime )q_{t}({\textbf {y}}^\prime ) \textrm{d}{} {\textbf {y}}^\prime \\&= -\,\mathbb {E}_{{\textbf {y}}^\prime \sim q_{t}}k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime )\nabla \log q^{*}({\textbf {y}}^\prime )+ \nabla k_{\varphi }({\textbf {y}}, {\textbf {y}}^\prime ) \end{aligned} \end{aligned}$$

(31)

The first equality is obtained by the definition of the integral operator (see Definition 1), the second equality is obtained by using (30) and the forth equality is obtained by applying the integration by parts to the first term. The proof is completed. $\square$

1.2 Appendix 1.2: Proof of Theorem 3

Proof

We analyze the performance of one step of mirrorVT. Under the Assumption 1.1 ($L_{2}$-smoothness of G), for any $t\ge 0$, we have:

$$\begin{aligned} \begin{aligned} G(q_{t+1})&\le G(q_{t}) + \langle \texttt {grad}G(q_{t}), \texttt {Exp}_{q_{t}}^{-1}(q_{t+1})\rangle _{q_{t}} +1/2 L_{2}\cdot \mathscr {W}^{2}_{2}(q_{t+1},q_{t})\\&= G(q_{t}) - \eta _{t} \langle \texttt {grad}G(q_{t}), \texttt {grad}G(q_{t})+\tilde{\delta }_{t} \rangle _{q_{t}}\\&\quad +1/2 L_{2}\eta _{t}^{2}\langle \texttt {grad}G(q_{t})+\tilde{\delta }_{t} , \texttt {grad}G(q_{t})+\tilde{\delta }_{t} \rangle _{q_{t}}\\&= G(q_{t}) - \eta _{t} \langle \texttt {grad}G(q_{t}), \texttt {grad}G(q_{t}) \rangle _{q_{t}} - \eta _{t} \langle \texttt {grad}G(q_{t}), \tilde{\delta }_{t}\rangle _{q_{t}}\\&\quad + 1/2 L_{2}\eta _{t}^{2}\langle \texttt {grad}G(q_{t})+\tilde{\delta }_{t} , \texttt {grad}G(q_{t})+\tilde{\delta }_{t} \rangle _{q_{t}} \end{aligned} \end{aligned}$$

(32)

where $\eta _{t}\in \left( 0, \alpha /h\right]$ (see (11)) and $\tilde{\delta }_{t}=-\texttt {div}(q_{t}(\nabla \tilde{g_{t}^{*}}-\nabla g^{*}_{t}))$ is the difference between the true 2-Wasserstein gradient at $q_{t}$ given by $\texttt {grad}G(q_{t})=-\texttt {div}(q_{t}\nabla g^{*}_{t})$ and its estimate given by $-\texttt {div}(q_{t}\nabla \tilde{g_{t}^{*}})$. The corresponding expected gradient error for G is defined as:

$$\begin{aligned} \tilde{\epsilon }_{t}=\mathbb {E}\langle \tilde{\delta }_{t}, \tilde{\delta }_{t} \rangle _{q_{t}}=\mathbb {E}\int \Vert \nabla ^{2}\varphi ({\textbf {x}})^{-1}\left( \nabla \tilde{f_{t}^{*}}({\textbf {x}})-\nabla f^{*}_{t}({\textbf {x}})\right) \Vert ^{2}_{2}p_{t}({\textbf {x}})\textrm{d} {\textbf {x}} \end{aligned}$$

(33)

Also since $0 \prec \alpha {\textbf {I}}\preceq \nabla ^{2}\varphi ({\textbf {x}})$ for all ${\textbf {x}}\in \mathscr {X}$, we have

$$\begin{aligned} \tilde{\epsilon }_{t} \le \frac{1}{\alpha ^{2}}\epsilon _{t} \end{aligned}$$

(34)

By applying the basic inequality: $\langle \texttt {grad}G(q_{t}), \tilde{\delta }_{t}\rangle \le \frac{1}{2}\langle \texttt {grad}G(q_{t}), \texttt {grad}G(q_{t})\rangle + \frac{1}{2} \langle \tilde{\delta }_{t}, \tilde{\delta }_{t}\rangle$ and combining with (34), we have:

$$\begin{aligned} \begin{aligned} G(q_{t+1})&\le G(q_{t}) - 1/2 \cdot \eta _{t} (1 - 2\eta _{t} L_{2})\cdot \langle \texttt {grad}G(q_{t}), \texttt {grad}G(q_{t})\rangle _{q_{t}} + \frac{\eta _{t} (1 + 2\eta _{t} L_{2})}{2\alpha ^{2}} \epsilon _{t}\\ \end{aligned} \end{aligned}$$

(35)

By the definition of the inner product on the tangent space and the assumption of $\mu$-strong convexity of F, we obtain the following inequality:

$$\begin{aligned} \begin{aligned} \langle \texttt {grad}G(q_{t}), \texttt {grad}G(q_{t})\rangle _{q_{t}}&=\int _{\mathscr {Y}} \Vert \nabla g^{*}_{t}({\textbf {y}})\Vert _{2}^{2}q_{t}({\textbf {y}})\textrm{d}{} {\textbf {y}}\\&=\int _{\mathscr {X}} \Vert \nabla ^{2}\varphi ({\textbf {x}}) \nabla f^{*}_{t}({\textbf {x}})\Vert _{2}^{2}p_{t}({\textbf {x}})\textrm{d} {\textbf {x}}\\&\ge \frac{1}{\alpha ^{2}}\int _{\mathscr {X}} \Vert \nabla f^{*}_{t}({\textbf {x}})\Vert _{2}^{2}p_{t}({\textbf {x}})\textrm{d}{} {\textbf {x}}=\frac{1}{\alpha ^{2}} \langle \texttt {grad}F(p_{t}), \texttt {grad}F(p_{t})\rangle _{p_{t}}\\&\ge \frac{\mu }{\alpha ^{2}} \left( F(p_{t})-F(p^{*})\right) \end{aligned} \end{aligned}$$

(36)

where the first inequality is obtained by $\nabla ^{2}\varphi ({\textbf {x}})\preceq \beta {\textbf {I}}$ for all ${\textbf {x}}\in \mathscr {X}$ and the second inequality is obtained by Assumption 1.3 (see 26). Thus combining (35) and use the identity: $F(p_{t})=G(q_{t})$, we have:

$$\begin{aligned} F(p_{t+1}) - F(p_{t}) \le \left[ 1 - \frac{\mu \eta _{t}}{2\beta ^{2}} (1 - 2\eta _{t}L_{2}))\right] \left( F(p_{t})-F^{*}\right) + \frac{\eta _{t} (1 + 2\eta _{t} L_{2})}{2\alpha ^{2}} \epsilon _{t} \end{aligned}$$

(37)

By setting $\eta _{t} =\eta \le \min \left\{ \frac{1}{2L_{2}},\frac{1}{\mu \beta ^{2}}\right\}$, we have:

$$\begin{aligned} 1 - \frac{\mu \eta _{t}}{2\beta ^{2}} \left( 1 - 2\eta _{t}L_{2})\right) \le 1-\frac{\mu \eta _{t}}{2\beta ^{2}}, 0 \le 1-\frac{\mu \eta }{2\beta ^{2}} \text { and } \frac{\eta (1 + 2\eta L_{2})}{2\alpha ^{2}} \le \frac{\eta }{\alpha ^{2}} \end{aligned}$$

(38)

In the sequel, we define $\rho = 1-\frac{\mu \eta }{2\beta ^{2}}\in \left[ 0,1 \right]$, we have:

$$\begin{aligned} F(p_{t+1}) - F(p_{t}) \le \rho \left( F(p_{t})-F^{*}\right) + \frac{\eta _{t}}{\alpha ^{2}} \epsilon _{t} \end{aligned}$$

(39)

By forming a telescoping sequence and combining the upper bound of $\epsilon _{t}$ given in Liu et al. (2021), we have:

$$\begin{aligned} F(p_{t+1}) - \inf _{p\in \mathscr {P}_{2}(\mathscr {X})} F(p) \le \rho ^{t}\left( F(p_{1})-\inf _{p\in \mathscr {P}_{2}(\mathscr {X})} F(p)\right) + \frac{1-\rho ^{t}}{1-\rho }\frac{\eta }{\alpha ^{2}}\cdot \texttt {Error} \end{aligned}$$

(40)

Finally, by taking the expectation over the initial particle set, we complete the proof. $\square$

1.3 Appendix 1.3: Details of the mirror maps

In this section, we describe more details of the mirror maps used in our simulated experiments.

1.3.1 Appendix 1.3.1: Mirror map on the unit ball

For the mirror map defined in (28), we can easily shown that:

$$\begin{aligned} \frac{\partial \varphi }{\partial {\textbf {x}}_{i}}=\frac{{\textbf {x}}_{i}}{1-\Vert {\textbf {x}} \Vert _{2}}, \frac{\partial \varphi ^{*}}{\partial {\textbf {y}}_{i}}=\frac{{\textbf {y}}_{i}}{1+\Vert {\textbf {y}} \Vert _{2}}, \frac{\partial ^{2}\varphi }{\partial {\textbf {x}}_{i}\partial {\textbf {x}}_{j}}=\frac{\delta _{ij}}{1-\Vert {\textbf {x}} \Vert _{2}} + \frac{{\textbf {x}}_{i}{} {\textbf {x}}_{j}}{\Vert {\textbf {x}} \Vert _{2} \left( 1 - \Vert {\textbf {x}} \Vert _{2} \right) ^{2}} \end{aligned}$$

Hence the Hessian matrix can be written as: $\nabla ^{2}\varphi ({\textbf {x}})=\frac{1}{1-\Vert {\textbf {x}} \Vert _{2}}{} {\textbf {I}}+\frac{1}{\Vert {\textbf {x}} \Vert _{2} \left( 1- \Vert {\textbf {x}} \Vert _{2}\right) ^{2}}{} {\textbf {x}} {\textbf {x}}^\top$, where ${\textbf {I}}$ is the identity matrix. In order to obtain the inversion of Hessian matrix, we apply the celebrated Woodbury matrix identity and show that

$$\begin{aligned} \left( \nabla ^{2}\varphi ({\textbf {x}})\right) ^{-1}=\left( 1 - \Vert {\textbf {x}} \Vert _{2}\right) \left( {\textbf {I}}-\frac{1}{\Vert {\textbf {x}} \Vert _{2}}{} {\textbf {x}} {\textbf {x}}^\top \right) \end{aligned}$$

1.3.2 Appendix 1.3.2: Mirror map on the simplex

For the entropic mirror map (see Beck & Teboulle 2003), we can consider ${\textbf {x}}=\left[ {\textbf {x}}_{1},\ldots ,{\textbf {x}}_{d-1}\right] \in \mathbb {R}^{d-1}$ by discarding the last entry ${\textbf {x}}_{d}=1 - \sum _{i=1}^{d-1}{} {\textbf {x}}_{i}$ and easily show that:

$$\begin{aligned} \frac{\partial \varphi }{\partial {\textbf {x}}_{i}}=\log {\textbf {x}}_{i} - \log {\textbf {x}}_{d}, \frac{\partial \varphi ^{*}}{\partial {\textbf {y}}_{i}}=\frac{\exp {\left( {\textbf {y}}_{i}\right) }}{\sum _{j=1}^{d-1}\exp {\left( {\textbf {y}}_{j}\right) }}, \frac{\partial ^{2}\varphi }{\partial {\textbf {x}}_{i}\partial {\textbf {x}}_{j}}=\frac{\delta _{ij}}{{\textbf {x}}_{i}} + \frac{1}{{\textbf {x}}_{d}}, \forall i \in \left[ d-1\right] \end{aligned}$$

Hence the Hessian matrix can be written as: $\nabla ^{2}\varphi ({\textbf {x}})=\texttt {diag}(1/{\textbf {x}}_{1},1/{\textbf {x}}_{2},\ldots ,1/{\textbf {x}}_{d-1})+1/{\textbf {x}}_{d} {\textbf {1}}{} {\textbf {1}}^\top$. By applying the Sherman-Morrison formula, we obtain the inverse Hessian matrix of the following form:

$$\begin{aligned} \left( \nabla ^{2}\varphi ({\textbf {x}})\right) ^{-1}=\texttt {diag}({\textbf {x}}) - {\textbf {x}} {\textbf {x}}^\top \end{aligned}$$

1.4 Appendix 1.4: Details of network architectures

We present the neural network architectures of DirVAE on the image data set MNIST in Table 4 and on the text data set in Table 5.

Table 4 Network architecture of DirVAE for MNIST

Table 5 Network Architecture of DirVAE for IMDB