Computation of Rate-Distortion-Perception Functions With Wasserstein Barycenter

Consider a discrete memoryless source $X\in\mathcal{X}$ and a reconstruction $\hat{X}\in\hat{\mathcal{X}}$, where $\mathcal{X}=\{x_{1},\cdots,x_{M}\},\hat{\mathcal{X}}=\{\hat{x}_{1},\cdots,\hat{x}_{N}\}$ are finite alphabets. Suppose $p_{X}$ and $p_{\hat{X}}$ are defined on the probability space $(\mathcal{X},\mathcal{F},\mathbb{P}).$ We consider the single-letter distortion measure $\Delta:\mathcal{X}\times\hat{\mathcal{X}}\mapsto[0,\infty)$ and perception measure between distributions $d:\mathbb{P}\times\mathbb{P}\mapsto[0,\infty).$

Note that the information RDP function (1) degenerates to RD functions when the constraint (1c) is removed.

Since the alphabets $\mathcal{X}$ and $\hat{\mathcal{X}}$ are finite, we denote

and $w_{ij}=W(\hat{x}_{j}\mid x_{i})$ for all $1\leq i\leq M,1\leq j\leq N$. Here $W:\mathcal{X}\rightarrow\hat{\mathcal{X}}$ is the channel transition mapping. Thus the discrete form of problem (1) can be written as

One commonly used measure of perceptual quality $d(\bm{p},\bm{r})$ is the Wasserstein metric,

where $c_{ij}$ denotes the cost matrix between $x_{i}$ and $\hat{x}_{j}$. In some cases, the total variation (TV) distance $\delta(\bm{p},\bm{r})=\frac{1}{2}\|\bm{p}-\bm{r}\|_{1}$ and the Kulback-Leibler (KL) divergence $\text{KL}(\bm{p}\|\bm{r})=\sum_{i=1}^{M}p_{i}\left[\log p_{i}-\log r_{i}\right]$ are also considered as a measure of perceptual quality [10]. We will focus on the Wasserstein metric, as generalizations to the TV and KL metrics follow straightforwardly as will be discussed in Section III.

Obviously, solving the RDP functions is more difficult than solving the RD function due to the introduction of the new perception constraint (2d). Numerically, we need to frequently update the $d(\bm{p},\bm{r})$ to ensure that the perception constraint is satisfied. For metrics such as the Wasserstein metric (3), the computational cost for this is high. Moreover, the original feasible domain of RD functions is changed due to the introduction of (2d). The feasible solution of RD functions is inside the simplex structure since (2b) and (2c) are both linear constraints. This is the reason why the BA algorithm and the AS algorithm can solve the RD functions at low costs [16]. However, the perception constraint (2d) breaks the desirable simplex structure, which brings computational challenges. In this section, we develop a new model and an algorithm for solving RDP functions efficiently and accurately. Our starting point is to convert the perception constraint (2d) into a linear constraint in a higher dimensional space. Consequently, all the constraints, (2b)-(2d), also preserve the simplex structure in the higher-dimensional space. This would bring great convenience to our algorithm design. The main idea is summarized into the following theorem:

Theorem 1 establishes an equivalence between the RDP problem (2) and the optimization (4). Also, we observe that the model (4) has the Wasserstein Barycenter structure. The optimization variable $r_{j}$ can be regarded as the Barycenter, and $\text{diag}(\bm{p})\cdot\bm{w}$ and $\bm{\Pi}$ are two couplings between each input $\bm{p}$ and Barycenter $\bm{r}$. Thus, we have obtained the Wasserstein Barycenter model for RDP functions. In general, the numerical methods for the Wasserstein Barycenter problem are computationally intensive [18, 20]. However, as $\bm{\Pi}$ is not explicitly included in the objective function of (4), it is possible to design a fast and efficient algorithm for the above WBM-RDP.

As discussed above, the WBM-RDP (4) is not strictly convex on $\bm{\Pi}$, the most direct way [21] is to introduce an extra entropy regularized term in the objective optimization function (4a), i.e.,

This leads to the entropy regularized WBM-RDP:

Here, $\varepsilon>0$ is a newly introduced regularization parameter. Thus, we get the entropy regularized WBM-RDP with strict convexity. Moreover, during the alternative optimization iterations, closed-form expressions of the dual variables can always be obtained, which accelerates the algorithm while improving the accuracy. Not only that, the following theorem guarantees that the solution to entropy regularized WBM-RDP (5) converges to WBM-RDP (4) as $\varepsilon\rightarrow 0$.

We construct the Lagrangian function of the regularized WBM-RDP (5) by introducing dual variables $\bm{\alpha},\bm{\theta}\in\mathbb{R}^{M},\bm{\beta},\bm{\tau}\in\mathbb{R}^{N},\lambda,\gamma\in\mathbb{R}^{+}$ and $\eta\in\mathbb{R}$:

Here we note that the Lagrangian function (7) is convex with respect to each variable. Furthermore, we can improve the algorithm by designing the directions of the alternating iterations according to how the variables appear in (7). Next, we sketch the main ingredients of our algorithm. The pseudo-code is presented in Algorithm 1.

1. 1)

   Update $\bm{w}$ and associated dual variables $\bm{\alpha,\beta},\lambda$ while fixing $\bm{r},\bm{\Pi}$ as constant parameters. First, we can update $\alpha_{i}$ and $\beta_{j}$:

   $$\psi_{j}\leftarrow 1\Big{/}\sum_{i=1}^{M}K_{ij}\phi_{i}p_{i},\quad\phi_{i}\leftarrow 1\Big{/}\sum_{j=1}^{N}K_{ij}\psi_{j}r_{j},$$

   (8)

   where $\phi_{i}=\exp(-\alpha_{i}/p_{i}-1/2),\psi_{j}=\exp(-\beta_{j}-1/2)$ and $K_{ij}=\exp(-\lambda d_{ij})$. Second, we update $\lambda$ by a root-searching operation on the following monotonic single-variable function on $\mathbb{R}^{+}$:

   $$F(\lambda)\triangleq\sum_{i=1}^{M}\sum_{j=1}^{N}d_{ij}p_{i}\phi_{i}e^{-\lambda d_{ij}}\psi_{j}r_{j}-D=0.$$

   (9)

2. 2)

   Update $\bm{\Pi}$ and associated dual variables $\bm{\theta,\tau},\gamma$ while fixing $\bm{r},\bm{w}$ as constant parameters. We alternatively update $\theta_{i},\tau_{j}$:

   $$\varphi_{j}\leftarrow r_{j}\Big{/}\sum_{i=1}^{M}M_{ij}\xi_{i},\quad\xi_{i}\leftarrow p_{i}\Big{/}\sum_{j=1}^{N}M_{ij}\varphi_{j},$$

   (10)

   where $\xi_{i}=\exp(-\theta_{i}/\varepsilon-1/2),\varphi_{j}=\exp(-\tau_{j}/\varepsilon-1/2)$ and $M_{ij}=\exp(-\gamma c_{ij}/\varepsilon)$. Again, we can update $\gamma$ by the root-searching operation on the following monotonic single-variable function on $\mathbb{R}^{+}$:

   $$G(\gamma)\triangleq\sum_{i=1}^{M}\sum_{j=1}^{N}c_{ij}\xi_{i}e^{-\gamma c_{ij}/\varepsilon}\varphi_{j}-P=0.$$

   (11)

3. 3)

   Update $\bm{r}$ and associated dual variables $\eta$ while fixing $\bm{w},\bm{\Pi}$ as constant parameters. We can update $\eta$ by finding the root of the following single-variable function on its largest monotone interval $(\max_{j}(\beta_{j}+\tau_{j}),+\infty)$:

   $$H(\eta)\triangleq\sum_{j=1}^{N}\left[\left(\sum_{i=1}^{M}w_{ij}p_{i}\right)\Big{/}\left(\eta-\beta_{j}-\tau_{j}\right)\right]-1=0,$$

   (12)

   and we finally get

   $$r_{j}=\left(\sum_{i=1}^{M}w_{ij}p_{i}\right)\Big{/}(\eta-\beta_{j}-\tau_{j}).$$

   (13)

We stress that the three steps above do not contain inner iterations, because $\bm{\Pi}$ is not explicitly included in the objective optimization function of the WBM-RDP as discussed above. Therefore, compared to other existing algorithms for solving the Wasserstein Barycenter, our proposed algorithm gain much efficiency and simplicity. On the other hand, our proposed algorithm adopts a similar alternating iteration technique to the Alternating Sinkhorn algorithm for RD functions. However, a vital operation we conduct is altering the projection direction of different joint distribution variables, i.e., $\bm{w}$ and $\bm{\Pi}$, which is essentially different from the AS algorithm for RD functions. Thus we name it the Improved Alternating Sinkhorn Algorithm.

We verify the convergence of the improved AS algorithm in this subsection. Here we consider the residual errors of the Karish-Kuhn-Tucker (KKT) condition of the optimization problem (5) to be the indicator of convergence. We define $L_{1}$ residual errors $r_{\psi}$ as $r_{\psi}=\sum_{j=1}^{N}\left|\psi_{j}\sum_{i=1}^{M}K_{ij}\phi_{i}p_{i}-1\right|,$ and $r_{\phi}$,$r_{\lambda}$,$r_{\eta}$,$r_{\varphi}$,$r_{\xi}$,$r_{\gamma}$ can be defined similarly. We define the overall residual error $r$ to be the root mean square of the above residual errors.

In Fig. 3, we respectively output the convergent trajectories of $r$ of the binary source with TV distance and Gaussian source with Wasserstein-2 metric against iteration numbers. For the binary source, we set $P=0.06$. For the Gaussian source, we set $P=2$. The results show different convergence behaviors with different distortion parameters, and all of these cases will converge below $1\times 10^{-10}$ at last.

We have theoretically demonstrated that when $\varepsilon\rightarrow 0$ the entropy regularized problem converges to the original problem. However, according to the general results of entropy regularized OT problem, problems with smaller $\varepsilon$ have higher precision but require more computation [20]. Moreover, according to the theory on Sinkhorn [20], an excessive $\lambda$ might trigger numerical stability problems. Thus we wish to investigate the impact of $\varepsilon$ on the WBM-RDP.

Results are shown in Table I, where the error represents the $L_{1}$ difference between the algorithm results and the explicit results. Here for the binary source with TV distance, we set $P=0.06$. And we set $P=2$ for the Gaussian source with Wasserstein-2 metric. When $\varepsilon$ decreases, the computational time rises while the results are more accurate regardless of the source. Furthermore, for $\varepsilon=1e-4$ in both sources, the small $\varepsilon$ causes numerical instability. Therefore, from a practical perspective, $\varepsilon=0.01$ seems to be an ideal choice for RDP functions, as it ensures a certain level of accuracy and does not consume too much time.