bibprimaryclass bibarchiveprefix

# Sampling-Based Methods for Multi-Block Optimization Problems over Transport
Polytopes††thanks: Funding: The work of the first author was supported by the
National Key R&D Program of China (2020YFA0711900, 2020YFA0711904). The work
of the second author was supported by the Outstanding Innovative Talents
Cultivation Funded Programs 2021 of Renmin University of China. The work of
the third author was supported in part by the National Natural Science
Foundation of China (12125108, 11971466, 12226008, 11991021, 11991020,
12021001, 12288201), Key Research Program of Frontier Sciences, Chinese
Academy of Sciences (ZDBS-LY-7022), and CAS-Croucher Funding Scheme for Joint
Laboratories “CAS AMSS-PolyU Joint Laboratory of Applied Mathematics:
Nonlinear Optimization Theory, Algorithms and Applications”. The work of the
fourth author was supported by Beijing Municipal Natural Science Foundation
(1232019), the National Natural Science Foundation of China (12101606), and
Renmin University of China research fund program for young scholars.

Yukuan Hu State Key Laboratory of Scientific and Engineering Computing,
Academy of Mathematics and Systems Science, Chinese Academy of Sciences, and
University of Chinese Academy of Sciences, Beijing, China<EMAIL_ADDRESS>liuxin@lsec.cc.ac.cn). Mengyu Li Institute of Statistics and Big Data, Renmin
University of China, Beijing, China (limengyu516@ruc.edu.cn). Xin
Liu22footnotemark: 2 Cheng Meng Center for Applied Statistics, Institute of
Statistics and Big Data, Renmin University of China, Beijing, China
(chengmeng@ruc.edu.cn).

###### Abstract

This paper focuses on multi-block optimization problems over transport
polytopes, which underlie various applications including strongly correlated
quantum physics and machine learning. Conventional block coordinate descent-
type methods for the general multi-block problems store and operate on the
matrix variables directly, resulting in formidable expenditure for large-scale
settings. On the other hand, optimal transport problems, as a special case,
have attracted extensive attention and numerical techniques that waive the use
of the full matrices have recently emerged. However, it remains nontrivial to
apply these techniques to the multi-block, possibly nonconvex problems with
theoretical guarantees. In this work, we leverage the benefits of both sides
and develop novel sampling-based block coordinate descent-type methods, which
are equipped with either entropy regularization or Kullback-Leibler
divergence. Each iteration of these methods solves subproblems restricted on
the sampled degrees of freedom. Consequently, they involve only sparse
matrices, which amounts to considerable complexity reductions. We explicitly
characterize the sampling-induced errors and establish convergence and
asymptotic properties for the methods equipped with the entropy
regularization. Numerical experiments on typical strongly correlated electron
systems corroborate their superior scalability over the methods utilizing full
matrices. The advantage also enables the first visualization of approximate
optimal transport maps between electron positions in three-dimensional
contexts.

##  1 Introduction

In this work, we consider multi-block optimization problems over transport
polytopes as follows:

$\begin{array}[]{rcl}\operatorname*{minimize}&(\min)&f(X_{1},\ldots,X_{N}),\\\
\text{subject
to}&(\mathrm{s.t.})&X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}),~{}i=1,\ldots,N,\end{array}$
(1.1)

where, for any $i\in\\{1,\ldots,N\\}$ ($1\leq N\in\mathbb{N}$),
$\mathbf{a}_{i}\in\mathbb{R}_{+}^{m_{i}}$,
$\mathbf{b}_{i}\in\mathbb{R}_{+}^{n_{i}}$,

$\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}):=\left\\{T\in\mathbb{R}_{+}^{m_{i}\times
n_{i}}\mid
T\mathbf{1}_{n_{i}}=\mathbf{a}_{i},~{}T^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i}\right\\}$

is called the transport polytope. The notation “$\mathbf{1}_{n}$” refers to
the all-ones vector in $\mathbb{R}^{n}$. The unknown matrix variables are
$X_{i}\in\mathbb{R}^{m_{i}\times n_{i}}$ ($i=1,\ldots,N$). The objective
function $f:\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times n_{i}}\to\mathbb{R}$ is
assumed to be block Lipschitz smooth over
$\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$ (for the
definition see section 4), possibly nonconvex. Problem (1.1) finds its
applications in several fields. For example, in quantum physics, it provides a
promising route for treating the elusive strongly correlated electron systems
(e.g., transition metal oxides [dagotto2005complexity]), describing the
electron-electron correlation explicitly [chen2014numerical, hu2023global]. It
can also act as a subproblem in finding the Wasserstein barycenter among
several discrete probability distributions [carlier2015numerical], which has
gained popularity so far in statistics [bigot2012consistent] and machine
learning [cuturi2014fast], as well as in label distribution learning
[zhao2018label], which reflects the relative importance of different labels in
supervised learning [geng2016label].

For solving general multi-block optimization problems, block coordinate
descent (bcd)-type methods rank among the top choices. These methods fully
exploit the separability of the feasible region, in that each subproblem
involves only one variable block and is much easier to solve than the original
problem. Representatives of the bcd-type methods are the block coordinate
descent (bcd) methods [bertsekas2016nonlinear, fercoq2016optimization,
wright2015coordinate], block conditional gradient (bcg) methods
[beck2015cyclic, braun2022conditional], proximal alternating linearized
minimization (palm) methods [bolte2014proximal, hu2023convergence], as well as
their stochastic versions [beck2015cyclic, chen2019extended,
driggs2021stochastic, hertrich2022inertial, lacoste2013block,
sun2020efficiency], where randomness is introduced to the gradient
calculations or update order. Nevertheless, for the specific problem (1.1),
all the existing bcd-type methods store and operate on the matrix variables
directly, requiring at least quadratically growing complexities per iteration.
This forms formidable memory and computation burdens when
$\\{m_{i}\\}_{i=1}^{N}$ and/or $\\{n_{i}\\}_{i=1}^{N}$ are of large magnitude.
Taking the aforementioned quantum physics application [hu2023global] for
instance, $m_{i}$ ($=n_{i}$) stands for the number of grid points and can be
of order $10^{4}$ or $10^{5}$ even for crude discretization.

When $N$ equals one ($m_{i}=m$, $n_{i}=n$) and $f$ is affine, problem (1.1)
reduces to the Kantorovich formulation of the classical optimal transport (OT)
problem [villani2003topics]. The exploration of this problem dates back to
Monge’s pioneering work in the 18th century [monge1781memoire], followed by
Kantorovich’s relaxation in the 20th century [kantorovich1942transfer]. Since
then, a plethora of numerical methods for solving OT problems have been
constantly emerging. Traditional ones solve differential equations
[benamou2002monge, brenier1997homogenized] or turn to linear programming
solvers [pele2009fast, rubner1997earth], resulting in unacceptable cubic
complexities. Nowadays, the widest usage may go to the entropy regularization-
based methods [cuturi2013sinkhorn, peyre2019computational], which allow for
the approximations of solutions in $\mathcal{O}(t_{\max}mn)$ scaling time with
the Sinkhorn algorithm [sinkhorn1967concerning], where $t_{\max}$ is the
number of iterations; see section 2.3 for more discussions. In recent years,
motivated by the need in large-scale contexts, there have been works dedicated
to alleviating the per-iteration quadratic costs by the conventional Sinkhorn
algorithm; for example, the low-rank approximation-based
[altschuler2019massively] and the entrywise sampling-based [li2023importance]
variants of the Sinkhorn algorithm. Remarkably, the latter variant essentially
deals with a restricted OT problem:

$\min_{X}~{}~{}\left\langle\hat{C},X\right\rangle,~{}~{}\mathrm{s.t.}~{}X\in\mathcal{U}(\mathbf{a},\mathbf{b}),~{}X_{\mathcal{I}^{c}}=0,$
(1.2)

where $\hat{C}\in\mathbb{R}^{m\times n}$ is an effective cost matrix (defined
later), $X\in\mathbb{R}^{m\times n}$, $\mathbf{a}\in\mathbb{R}^{m}$,
$\mathbf{b}\in\mathbb{R}^{n}$, $\mathcal{I}\subseteq\\{(j,k)\mid
j=1,\ldots,m,~{}k=1,\ldots,n\\}$ contains the indices sampled from the
beginning according to some probability distribution related to $\mathbf{a}$
and $\mathbf{b}$, and $\mathcal{I}^{c}$ denotes its complementary set. The
constraint “$X_{\mathcal{I}^{c}}=0$” enforces the entries in $X$ indexed by
$\mathcal{I}^{c}$ to be zero, which distinguishes the algorithm from the well-
known stochastic optimization methods. As a result, only
$\left|\mathcal{I}\right|$ entries in $X$ get involved in the calculations and
updates, leading to a nice scaling when $\left|\mathcal{I}\right|=o(mn)$.
However, it remains unclear whether the sampling technique can be adapted to
handle the multi-block, possibly nonconvex problem (1.1), while maintaining
favorable convergence properties. One possible way is to integrate the
sampling technique into the bcd-type methods and to solve restricted
subproblems like (1.2) in each iteration. Analyzing the accumulation of errors
induced by sampling will then become subtle.

###  1.1 Contributions and organization.

We develop in this paper sampling-based bcd-type methods for problem (1.1),
which are equipped with either entropy regularization or Kullback-Leibler
divergence. In particular, importance samplings are performed conforming to
the probability distributions associated with the previous iterates and each
iteration solves subproblems restricted over sampled supports. Consequently,
only $o(m_{i}n_{i})$ entries in $X_{i}$ ($i=1,\ldots,N$) take part in the
updates and derivatives calculations, which amounts to considerable
computational saving in large-scale contexts.

Following the theoretical results about randomized matrix sparsification, we
analyze the convergence and asymptotic properties for the methods equipped
with entropy regularization. We explicitly characterize the sampling-induced
errors and establish upper bounds for the average stationarity violations. The
average violation is further shown to vanish in the limit
$\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$ (with probability going to 1).
Notably, to the best of our knowledge, our work is the first attempt in
applying the matrix entrywise sampling technique to multi-block nonconvex
settings with theoretical guarantees.

We demonstrate the efficiency of the newly designed methods via numerical
simulations of typical strongly correlated electron systems. Their better
scalability enables the first visualization of the approximate OT maps between
electron positions in three-dimensional contexts.

The paper is organized as follows. We provide preliminaries in section 2 and
elaborate on the algorithmic developments in section 3. Section 4 contains our
theoretical results, whose proofs are deferred to the appendix. Numerical
experiments and results are described in section 5. Finally, we conclude in
section 6.

##  2 Preliminaries

This section offers some preliminaries, including notations, tools from OT,
and bibliographical notes on entrywise matrix sparsification.

###  2.1 Notations

This paper presents scalars, vectors, and matrices by regular-font, bold
lower-case, and upper-case letters, respectively. We denote the rounding down
operation by “$\lfloor\cdot\rfloor$”. The notation “$\mathbf{1}_{n}$” stands
for the all-ones vector in $\mathbb{R}^{n}$. The notations
“$\left\langle\cdot,\cdot\right\rangle$” and “$\|\cdot\|$” calculate,
respectively, the standard inner product and norm of vectors or matrices in
the ambient Euclidean space. We use “$\|\cdot\|_{2}$” particularly for the
2-norm of matrices. The notation “$\kappa(\cdot)$” refers to the spectral
condition number of a matrix. We use “$\mathrm{Diag}(\cdot)$” to form a
diagonal matrix with the input vector. We denote the entries or sub-blocks of
vectors by single subscripts (e.g., $\varrho_{k}$ or $\mathbf{a}_{i}$), the
sub-blocks of matrices by single subscripts (e.g., $X_{i}$), and the entries
of matrices by double subscripts (e.g., $x_{i,jk}$). A matrix with a set
subscript refers to the entries indexed by the set (e.g., $X_{\mathcal{I}}$).
Sometimes, we make abbreviations for the aggregation of the sub-blocks of a
matrix (e.g., $X_{\leq i}:=(X_{1},\ldots,X_{i})$,
$X_{>i}:=(X_{i+1},\ldots,X_{N})$). These abbreviations become null if the
index sets in the subscripts are empty. The entrywise product and division of
two vectors or matrices are denoted by “$\odot$” and “$\oslash$”,
respectively. Univariate functions, such as “$\exp(\cdot)$” and
“$\log(\cdot)$”, are extended to vectors and matrices as entrywise operations.

For a multivariate function $g$, $\nabla g$ is the gradient of $g$ at the
points where $g$ is differentiable. We add a subscript to indicate the block
with respect to which the derivative is taken (e.g., $\nabla_{i}g$).

We use $\mathbb{R}_{++}$ to denote the set of positive real numbers. Given a
set, its measure or cardinality is represented using “$\left|\cdot\right|$”.
We use “$\bigtimes$” or “$\times$” to refer to the Cartesian product of sets
or spaces (e.g., $\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times n_{i}}$ or
$\mathbb{R}^{m_{i}}\times\mathbb{R}^{n_{i}}$), use exponents to represent the
Cartesian product of identical sets or spaces (e.g., $(\mathbb{R}^{K\times
K})^{N_{e}-1}$). The complementary set is noted with a superscript “$c$”
(e.g., $\mathcal{I}^{c}$).

When describing algorithms, we use superscripts within brackets to refer to
the iteration numbers (e.g., $X_{i}^{(t)}$).

###  2.2 Negative entropy and Kullback-Leibler divergence

######  Definition 1 ([kullback1951information, shannon1948mathematical]).

For any $T=(t_{ij})\in\mathbb{R}_{+}^{m\times n}$, its negative entropy is
defined as $h(T):=\sum_{ij}t_{ij}(\log t_{ij}-1)$. Given any $T=(t_{ij})$,
$T^{\prime}=(t_{ij}^{\prime})\in\mathbb{R}_{+}^{m\times n}$, the Kullback-
Leibler (KL) divergence between $T$ and $T^{\prime}$ is defined as

$\mathrm{KL}(T;T^{\prime}):=\sum_{i,j}\left[t_{ij}(\log t_{ij}-\log
t_{ij}^{\prime})-(t_{ij}-t_{ij}^{\prime})\right].$ (2.1)

If $t_{ij}>0$ and $t_{ij}^{\prime}=0$ for some pair $(i,j)$, then
$\mathrm{KL}(T;T^{\prime})=+\infty$.

The negative entropy has been adopted in thermodynamics as a measure of
disorder in a system, or a measure of uncertainty in information theory. The
KL divergence can be treated as the Bregman distance [bregman1967relaxation]
associated with the negative entropy; it has been used as a measure of the
disparity between probability distributions [peyre2019computational].

###  2.3 Entropy regularized optimal transport

The Kantorovich formulation of the discrete OT problem is in general

$\min_{T}~{}\left\langle
W,T\right\rangle,~{}~{}\mathrm{s.t.}~{}T\in\mathcal{U}(\mathbf{p},\mathbf{q}),$
(2.2)

where $T=(t_{ij})\in\mathbb{R}^{m\times n}$ is the transport plan,
$W=(w_{ij})\in\mathbb{R}^{m\times n}$ is the cost matrix, and
$\mathbf{p}\in\mathbb{R}^{m}$, $\mathbf{q}\in\mathbb{R}^{n}$ are discrete
probability distributions. The solution to problem (2.2) is referred to as the
OT plan, which achieves minimal transportation efforts. Nowadays, OT has
attracted extensive attention from applications (e.g.,
[arjovsky2017wasserstein, ma2019optimal, meng2019large, xu2019learning]).

The computational complexity of directly solving problem (2.2) as a linear
programming usually grows cubically as $m$ and $n$ increase, which severely
hinders the wide applications of OT. To approximate the solution efficiently
within certain tolerance, the author of [cuturi2013sinkhorn] adds a (negative)
entropy regularization penalty, obtaining

$\min_{T}~{}\left\langle W,T\right\rangle+\lambda
h(T),~{}~{}\mathrm{s.t.}~{}T\mathbf{1}_{n}=\mathbf{p},~{}T^{\top}\mathbf{1}_{m}=\mathbf{q},$
(2.3)

where $\lambda>0$ is the regularization parameter. Note that the nonnegativity
requirement is unnecessary due to the definition of $h$.

On account of the entropy regularizer, problem (2.3) becomes strongly convex
and thus admits a unique optimal solution. Furthermore, as $\lambda\to 0$,
this optimal solution converges to an optimal solution of problem (2.2)
[peyre2019computational]. Computationally, the dual of problem (2.3) can be
solved by an alternating minimization scheme, known as the Sinkhorn algorithm
in the OT community [sinkhorn1967concerning]. The alternating scheme involves
only matrix-vector multiplications and entrywise divisions between vectors
(see section 3.1), particularly suited for GPU executions
[cuturi2013sinkhorn].

###  2.4 Entrywise matrix sparsification

Entrywise matrix sparsification is pioneered by Achlioptas and McSherry
[achlioptas2007fast], later developed in [braverman2021near, drineas2011note,
kundu2017recovering], where sampling-based algorithms are described to select
a small number of entries from an input matrix to construct a sparse sketch;
the sketch is close to the original one in the operator norm with a high
probability guarantee. The entries are sampled following some probability
distribution associated with the original matrix. Among others, importance
sampling, as a statistical technique, constructs the sampling probability
distribution in the spirit of variance reduction [elvira2021advances,
liu1996metropolized, liu2008monte, owen2013monte]. Owing to the substantial
computational complexity reduction thereby, entrywise matrix sparsification
has been adopted in various scenarios, e.g., for computing approximate
eigenvectors [achlioptas2007fast], solving OT problems [li2023importance], and
computing Gromov-Wasserstein distances [li2023efficient].

##  3 Algorithmic developments

In this part, we develop two classes of methods for the multi-block problem
over the transport polytopes (1.1), with the bcg and palm methods as starting
points. In particular, we add entropy regularizers to the subproblems in the
bcg method, while replacing the proximal term with the KL divergence in the
palm method. We further equip them with importance sampling-based entrywise
matrix sparsification.

###  3.1 Entropy regularized alternating linearized minimization

In each iteration, the bcg method obtains search directions via solving
subproblems as

$\min_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle,~{}\mathrm{s.t.}~{}X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}),$
(3.1)

where

$C_{i}^{(t)}:=\nabla_{i}f\big{(}X_{<i}^{(t+1)},X_{\geq
i}^{(t)}\big{)}\in\mathbb{R}^{m_{i}\times n_{i}}.$ (3.2)

Subproblem (3.1) can be identified as an OT problem with $C_{i}^{(t)}$ being
the cost matrix. As mentioned in section 2.3, solving subproblem (3.1) entails
cubic complexities using linear programming methods, which forms an
unacceptable computational burden. To this end, we add an entropy
regularization and instead resort to

$\min_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle+\lambda_{i}^{(t)}h(X_{i}),~{}\mathrm{s.t.}~{}X_{i}\mathbf{1}_{n_{i}}=\mathbf{a}_{i},~{}X_{i}^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i}$
(3.3)

in each iteration, where $\lambda_{i}^{(t)}>0$ is the regularization parameter
and $h$ is the negative entropy in Definition 1. Consequently, we obtain the
entropy regularized alternating linearized minimization (eralm) method; see
Algorithm 1.

Algorithm 1 The eralm method for solving problem (1.1).

0: $X_{i}^{(0)}\in\mathbb{R}^{m_{i}\times n_{i}}$,
$\mathbf{a}_{i}\in\mathbb{R}^{m_{i}}$, $\mathbf{b}_{i}\in\mathbb{R}^{n_{i}}$
($i=1,\ldots,N$), $t_{\max}\in\mathbb{N}$.

1: Set $t:=0$.

2: while certain conditions are not satisfied and $t<t_{\max}$ do

3: for $i=1,\ldots,N$ do

4: Select a regularization parameter $\lambda_{i}^{(t)}>0$ and a step size
$\alpha_{i}^{(t)}\in(0,1]$.

5: Compute $C_{i}^{(t)}\in\mathbb{R}^{m_{i}\times n_{i}}$ as in the formula
(3.2).

6: Solve subproblem (3.3) or (3.4) to obtain
$\tilde{X}_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times n_{i}}$.

7: Update
$X_{i}^{(t+1)}:=(1-\alpha_{i}^{(t)})X_{i}^{(t)}+\alpha_{i}^{(t)}\tilde{X}_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times
n_{i}}$.

8: end for

9: Set $t:=t+1$.

10: end while

10: Approximate solution
$(X_{1}^{(t)},\ldots,X_{N}^{(t)})\in\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times
n_{i}}$.

Observe that the number of variables in subproblem (3.3) scales quadratically,
while the number of equality constraints grows linearly. Therefore, it is more
advantageous to work from the dual perspective, especially when $m_{i}$ and/or
$n_{i}$ is of large magnitude. From [peyre2019computational], the dual of
subproblem (3.3) is

$\min_{\tilde{\mathbf{u}}_{i},\tilde{\mathbf{v}}_{i}}~{}q_{i}(\tilde{\mathbf{u}}_{i},\tilde{\mathbf{v}}_{i};\lambda_{i}^{(t)},\Psi_{i}^{(t)}):=\lambda_{i}^{(t)}\exp\left(\frac{\tilde{\mathbf{u}}_{i}}{\lambda_{i}^{(t)}}\right)^{\top}\Psi_{i}^{(t)}\exp\left(\frac{\tilde{\mathbf{v}}_{i}}{\lambda_{i}^{(t)}}\right)-\tilde{\mathbf{u}}_{i}^{\top}\mathbf{a}_{i}-\tilde{\mathbf{v}}_{i}^{\top}\mathbf{b}_{i},$
(3.4)

where $\tilde{\mathbf{u}}_{i}\in\mathbb{R}^{m_{i}}$,
$\tilde{\mathbf{v}}_{i}\in\mathbb{R}^{n_{i}}$ are the dual variables
associated with the equality constraints, and

$\Psi_{i}^{(t)}:=\exp\left(-C_{i}^{(t)}/\lambda_{i}^{(t)}\right)\in\mathbb{R}^{m_{i}\times
n_{i}}$ (3.5)

is called the kernel matrix. To tackle the block-structured problem (3.4), one
natural choice is the bcd method, which is also known as the Sinkhorn
algorithm [sinkhorn1967concerning] in the context of OT. Starting from a given
$\tilde{\mathbf{v}}_{i}^{(t,0)}\in\mathbb{R}^{n_{i}}$, the Sinkhorn algorithm
repeats the following two steps until fulfilling certain criteria:

$\displaystyle\tilde{\mathbf{u}}_{i}^{(t,s+1)}:$
$\displaystyle=\lambda_{i}^{(t)}\log\left(\mathbf{a}_{i}\oslash\left(\Psi_{i}^{(t)}\exp\big{(}\tilde{\mathbf{v}}_{i}^{(t,s)}/\lambda_{i}^{(t)}\big{)}\right)\right),$
$\displaystyle\tilde{\mathbf{v}}_{i}^{(t,s+1)}:$
$\displaystyle=\lambda_{i}^{(t)}\log\left(\mathbf{b}_{i}\oslash\left(\Psi_{i}^{(t)\top}\exp\big{(}\tilde{\mathbf{u}}_{i}^{(t,s+1)}/\lambda_{i}^{(t)}\big{)}\right)\right),$

where $s$ indicates the subiteration number. After letting
$\check{\mathbf{u}}_{i}^{(t,s)}:=\exp\big{(}\tilde{\mathbf{u}}_{i}^{(t,s)}/\lambda_{i}^{(t)}\big{)}\in\mathbb{R}^{m_{i}}$
and
$\check{\mathbf{v}}_{i}^{(t,s)}:=\exp\big{(}\tilde{\mathbf{v}}_{i}^{(t,s)}/\lambda_{i}^{(t)}\big{)}\in\mathbb{R}^{n_{i}}$,
the above schemes can be rewritten as

$\check{\mathbf{u}}_{i}^{(t,s+1)}:=\mathbf{a}_{i}\oslash\left(\Psi_{i}^{(t)}\check{\mathbf{v}}_{i}^{(t,s)}\right),~{}~{}\check{\mathbf{v}}_{i}^{(t,s+1)}:=\mathbf{b}_{i}\oslash\left(\Psi_{i}^{(t)\top}\check{\mathbf{u}}_{i}^{(t,s+1)}\right),$
(3.6)

involving matrix-vector multiplications and entrywise divisions between
vectors, thus favoring high parallel scalability [cuturi2013sinkhorn]. The
linear convergence rate of the Sinkhorn algorithm has been established
[luo1993convergence]. Further acceleration can be gained via warm starts
[xie2020fast].

###  3.2 Sampling-based variant of the eralm method

The eralm method still works on matrix variables and requires calculating
$C_{i}^{(t)}$ (3.5) explicitly. Below, we use a sparse matrix
$\hat{\Psi}_{i}^{(t)}\in\mathbb{R}^{m_{i}\times n_{i}}$ to approximate
$\Psi_{i}^{(t)}=(\psi_{i,jk}^{(t)})$, rendering most of the computational
costs dispensable.

The idea of the sparse approximation largely originates from the following
multiplicative expression for the unique optimal solution of subproblem (3.3):

$\tilde{X}_{i}^{(t+1,\star)}:=\mathrm{Diag}\left(\exp\left(\frac{\tilde{\mathbf{u}}_{i}^{(t,\star)}}{\lambda_{i}^{(t)}}\right)\right)\Psi_{i}^{(t)}\mathrm{Diag}\left(\exp\left(\frac{\tilde{\mathbf{v}}_{i}^{(t,\star)}}{\lambda_{i}^{(t)}}\right)\right)\in\mathbb{R}^{m_{i}\times
n_{i}},$ (3.7)

where
$\big{(}\tilde{\mathbf{u}}_{i}^{(t,\star)},\tilde{\mathbf{v}}_{i}^{(t,\star)}\big{)}\in\mathbb{R}^{m_{i}}\times\mathbb{R}^{n_{i}}$
is an optimal solution of the dual (3.4). The expression (3.7) indicates that
$\tilde{x}_{i,jk}^{(t+1,\star)}=0$ whenever $\psi_{i,jk}^{(t)}=0$. For
another, it has been theoretically established in [hosseini2022intrinsic] that
the solutions of subproblem (3.1) can be sparse. When $f$ is multi-affine,
there even exist sparse solutions with
$\mathcal{O}(\sum_{i=1}^{N}(m_{i}+n_{i}))$ nonzero entries to problem (1.1)
[hu2023exactness, hu2023global]. These two points together motivate us to
compute only a small portion of the entries in $\Psi_{i}^{(t)}$ to eliminate
most of the storage and computation overhead.

For this purpose, it suffices to estimate the sparsity pattern along
iterations. In this work, we make an attempt through entrywise sampling
[achlioptas2013near, braverman2021near, drineas2011note, kundu2017recovering,
li2023efficient]; that is, randomly pick a small portion from
$\\{(j,k):j=1,\ldots,m_{i},~{}k=1,\ldots,n_{i}\\}$ according to certain
sampling probability distribution. Specifically, the entrywise sampling is
implemented via Poisson sampling framework following the recent works
[braverman2021near, ai2021optimal, li2023importance], which independently
evaluate each index for inclusion in the sampled set. Compared to the sampling
with replacement, the Poisson sampling has been shown to provide higher
accuracy in some scenarios and is more practical for distributed systems
[wang2022sampling].

In light of the importance sampling (see section 2.4), the optimal sampling
probabilities should be
$p_{i,jk}^{(t,\star)}\propto\tilde{x}_{i,jk}^{(t+1,\star)}$ [li2023efficient].
However, $\tilde{X}_{i}^{(t+1,\star)}$ is completely unknown without the
knowledge of $\tilde{\mathbf{u}}_{i}^{(t,\star)}$ and
$\tilde{\mathbf{v}}_{i}^{(t,\star)}$ (see the formula (3.7)). An alternative
is to sample conforming to the values in the previous iteration, i.e.,
$p_{i,jk}^{(t)\prime}\propto x_{i,jk}^{(t)}$. This becomes reasonable when the
procedure gets close to the optimum. In addition, to recover entries $(j,k)$
from the optimal sparsity pattern that are missing in $X_{i}^{(t)}$, in the
sense that $x_{i,jk}^{(t)}=0$, we linearly interpolate between
$p_{i,jk}^{(t)\prime}$ and the sampling probability proposed by
[li2023efficient], i.e.,
$p_{i,jk}^{\prime\prime}\propto\sqrt{a_{i,j}b_{i,k}}$. Specifically,

$p_{i,jk}^{(t)}:=\gamma
p_{i,jk}^{(t)\prime}+(1-\gamma)p_{i,jk}^{\prime\prime}=\gamma\frac{x_{i,jk}^{(t)}}{\sum_{j^{\prime},k^{\prime}}x_{i,j^{\prime}k^{\prime}}^{(t)}}+(1-\gamma)\frac{\sqrt{a_{i,j}b_{i,k}}}{\sum_{j^{\prime},k^{\prime}}\sqrt{a_{i,j^{\prime}}b_{i,k^{\prime}}}}$
(3.8)

for $j=1,\ldots,m_{i}$, $k=1,\ldots,n_{i}$. Here, $\gamma\in[0,1]$ stands for
the interpolation factor. Such a shrinkage strategy is widely adopted in the
subsampling literature [ma2014statistical, yu2022optimal].

Given a sampling parameter $n_{s,i}\in\mathbb{N}$, the sparse approximation
$\hat{\Psi}_{i}^{(t)}=(\hat{\psi}_{i,jk}^{(t)})\in\mathbb{R}^{m_{i}\times
n_{i}}$ for the kernel matrix $\Psi_{i}^{(t)}$ is constructed in accordance
with the Poisson sampling principle, i.e.,

$\hat{\psi}_{i,jk}^{(t)}:=\left\\{\begin{array}[]{ll}\psi_{i,jk}^{(t)}/p_{i,jk}^{(t)\ast},&\text{with
probability}~{}p_{i,jk}^{(t)\ast}:=\min\left\\{1,n_{s,i}\cdot
p_{i,jk}^{(t)}\right\\},\\\ 0,&\text{otherwise},\end{array}\right.$ (3.9)

where the denominators ensure the unbiasedness of the random approximation
(see Appendix B). We denote the sampled set of indices as
$\mathcal{I}_{i}^{(t)}\subseteq\\{(j,k):j=1,\ldots,m_{i},~{}k=1,\ldots,n_{i}\\}$.
Note that
$\mathsf{E}(|\mathcal{I}_{i}^{(t)}|)=\sum_{j,k}p_{i,jk}^{(t)\ast}\leq
n_{s,i}\sum_{j,k}p_{i,jk}^{(t)}=n_{s,i}$, which indicates that $n_{s,i}$ is an
upper bound for the expected number of nonzero entries in
$\hat{\Psi}_{i}^{(t)}$. Although $|\mathcal{I}_{i}^{(t)}|$ fluctuates, it
still concentrates near its expectation with high probability [ai2021optimal],
ensuring that the computational cost remains manageable.

Replacing $\Psi_{i}^{(t)}$ in subproblem (3.4) with $\hat{\Psi}_{i}^{(t)}$, we
obtain

$\min_{\tilde{\mathbf{u}}_{i},\tilde{\mathbf{v}}_{i}}~{}q_{i}(\tilde{\mathbf{u}}_{i},\tilde{\mathbf{v}}_{i};\lambda_{i}^{(t)},\hat{\Psi}_{i}^{(t)}),$
(3.10)

which is the dual of1)1)1)The strong duality holds if subproblem (3.11) is
feasible with $\mathcal{I}_{i}^{(t)}$. The same will be true later, for
subproblem (3.20).

$\begin{array}[]{cl}\min\limits_{X_{i}}&\left\langle\hat{C}_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle+\lambda_{i}^{(t)}h(X_{i}),\\\
\mathrm{s.t.}&X_{i}\mathbf{1}_{n_{i}}=\mathbf{a}_{i},~{}X_{i}^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i},~{}(X_{i})_{(\mathcal{I}_{i}^{(t)})^{c}}=0.\end{array}$
(3.11)

Here, $\hat{C}_{i}^{(t)}=(\hat{c}_{i,jk}^{(t)})\in\mathbb{R}^{m_{i}\times
n_{i}}$ is the effective cost matrix, defined as

$\hat{c}_{i,jk}^{(t)}:=\left\\{\begin{array}[]{ll}c_{i,jk}^{(t)}+\lambda_{i}^{(t)}\log\left(p_{i,jk}^{(t)\ast}\right),&\text{if}~{}(j,k)\in\mathcal{I}_{i}^{(t)},\\\
c_{i,jk}^{(t)},&\text{otherwise}.\end{array}\right.$ (3.12)

Note that in $\hat{C}_{i}^{(t)}$ only the entries indexed by
$\mathcal{I}_{i}^{(t)}$ are required to form $\hat{\Psi}_{i}^{(t)}$. To solve
subproblem (3.10) and (3.11), it suffices to substitute $\Psi_{i}^{(t)}$ with
$\hat{\Psi}_{i}^{(t)}$ in the Sinkhorn algorithm (3.6).

We call the eralm method equipped with sampling the S-eralm method, which is
summarized in Algorithm 2. The maximum expected numbers of sampled indices are
controlled by $\\{n_{s,i}\\}_{i=1}^{N}\subseteq\mathbb{N}$.

Algorithm 2 The S-eralm method for solving problem (1.1).

0: $X_{i}^{(0)}\in\mathbb{R}^{m_{i}\times n_{i}}$,
$\mathbf{a}_{i}\in\mathbb{R}^{m_{i}}$, $\mathbf{b}_{i}\in\mathbb{R}^{n_{i}}$
($i=1,\ldots,N$), $\gamma\in[0,1]$,
$\\{n_{s,i}\\}_{i=1}^{N}\subseteq\mathbb{N}$, $t_{\max}\in\mathbb{N}$.

1: Set $t:=0$.

2: while certain conditions are not satisfied and $t<t_{\max}$ do

3: for $i=1,\ldots,N$ do

4: Select a regularization parameter $\lambda_{i}^{(t)}>0$ and a step size
$\alpha_{i}^{(t)}\in(0,1]$.

5: Randomly pick a subset
$\mathcal{I}_{i}^{(t)}\subseteq\\{(j,k):j=1,\ldots,m_{i},~{}k=1,\ldots,n_{i}\\}$
following the Poisson sampling with
$P_{i}^{(t)}=(p_{i,jk}^{(t)})\in\mathbb{R}_{+}^{m_{i}\times n_{i}}$ given by
(3.8) and $n_{s,i}$.

6: Construct
$\hat{\Psi}_{i}^{(t)}=(\hat{\psi}_{i,jk}^{(t)})\in\mathbb{R}^{m_{i}\times
n_{i}}$ as in the formula (3.9).

7: Solve subproblem (3.10) or (3.11) to obtain
$\tilde{X}_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times n_{i}}$.

8: Update
$X_{i}^{(t+1)}:=(1-\alpha_{i}^{(t)})X_{i}^{(t)}+\alpha_{i}^{(t)}\tilde{X}_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times
n_{i}}$.

9: end for

10: Set $t:=t+1$.

11: end while

11: Approximate solution
$(X_{1}^{(t)},\ldots,X_{N}^{(t)})\in\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times
n_{i}}$.

###  3.3 KL divergence-based alternating linearized minimization

We now turn to methods involving a direct penalization of the distance between
$X_{i}$ and $X_{i}^{(t)}$. In each iteration, the palm method solves several
proximal linearized subproblems of the form

$\min_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle+\frac{\mu_{i}^{(t)}}{2}\|X_{i}-X_{i}^{(t)}\|^{2},~{}\mathrm{s.t.}~{}X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}),$
(3.13)

where $\mu_{i}^{(t)}>0$ refers to the proximal parameter and $C_{i}^{(t)}$ is
defined in equation (3.2). Solving subproblem (3.13) is equivalent to
projecting $X_{i}^{(t)}-C_{i}^{(t)}/\mu_{i}^{(t)}$ onto
$\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$. In this part, we replace the
proximal term $\|X_{i}-X_{i}^{(t)}\|^{2}/2$ by the KL divergence
$\mathrm{KL}(X_{i};X_{i}^{(t)})$, which is the Bregman distance induced by the
negative entropy (see section 2.2), and obtain

$\min_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle+\mu_{i}^{(t)}\mathrm{KL}(X_{i};X_{i}^{(t)}),~{}\mathrm{s.t.}~{}X_{i}\mathbf{1}_{n_{i}}=\mathbf{a}_{i},~{}X_{i}^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i}.$
(3.14)

This is motivated by the Bregman proximal point algorithms for convex
optimization [xie2020fast, yang2022bregman] and the Bregman distance-based
palm methods for nonconvex optimization with relatively smooth objectives
[ahookhosh2021multi, li2019provable]. We describe the KL divergence-based
alternating linearized minimization (klalm) method in Algorithm 3.

Algorithm 3 The klalm method for solving problem (1.1).

0: $X_{i}^{(0)}\in\mathbb{R}^{m_{i}\times n_{i}}$,
$\mathbf{a}_{i}\in\mathbb{R}^{m_{i}}$, $\mathbf{b}_{i}\in\mathbb{R}^{n_{i}}$
($i=1,\ldots,N$), $t_{\max}\in\mathbb{N}$.

1: Set $t:=0$.

2: while certain conditions are not satisfied and $t<t_{\max}$ do

3: for $i=1,\ldots,N$ do

4: Select a proximal parameter $\mu_{i}^{(t)}>0$.

5: Compute $C_{i}^{(t)}\in\mathbb{R}^{m_{i}\times n_{i}}$ as in the formula
(3.2).

6: Solve subproblem (3.14) or (3.15) to obtain
$X_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times n_{i}}$.

7: end for

8: Set $t:=t+1$.

9: end while

9: Approximate solution
$(X_{1}^{(t)},\ldots,X_{N}^{(t)})\in\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times
n_{i}}$.

As in the previous subsection, we write out the dual of subproblem (3.14):

$\min_{\mathbf{u}_{i},\mathbf{v}_{i}}~{}q_{i}(\mathbf{u}_{i},\mathbf{v}_{i};\mu_{i}^{(t)},\Phi_{i}^{(t)}),$
(3.15)

where $q_{i}$ is defined in problem (3.4),
$\mathbf{u}_{i}\in\mathbb{R}^{m_{i}}$, $\mathbf{v}_{i}\in\mathbb{R}^{n_{i}}$
are the dual variables associated with the equality constraints and the kernel
matrix

$\Phi_{i}^{(t)}:=\exp\left(-C_{i}^{(t)}/\mu_{i}^{(t)}\right)\odot
X_{i}^{(t)}\in\mathbb{R}^{m_{i}\times n_{i}}.$ (3.16)

The Sinkhorn algorithm (3.6) thus still applies, yet with $\Psi_{i}^{(t)}$
changed to $\Phi_{i}^{(t)}$.

###  3.4 Sampling-based variant of the klalm method

Since subproblem (3.14) in the klalm method admits the following
multiplicative expression for the unique optimal solution:

$X_{i}^{(t+1,\star)}:=\mathrm{Diag}\left(\exp\left(\frac{\mathbf{u}_{i}^{(t,\star)}}{\mu_{i}^{(t)}}\right)\right)\Phi_{i}^{(t)}\mathrm{Diag}\left(\exp\left(\frac{\mathbf{v}_{i}^{(t,\star)}}{\mu_{i}^{(t)}}\right)\right)\in\mathbb{R}^{m_{i}\times
n_{i}},$ (3.17)

where
$\big{(}\mathbf{u}_{i}^{(t,\star)},\mathbf{v}_{i}^{(t,\star)}\big{)}\in\mathbb{R}^{m_{i}}\times\mathbb{R}^{n_{i}}$
is an optimal solution of the dual (3.15), we could likewise employ sparse
approximation on $\Phi_{i}^{(t)}$ for better scalability. We define the sparse
approximation
$\hat{\Phi}_{i}^{(t)}=(\hat{\varphi}_{i,jk}^{(t)})\in\mathbb{R}^{m_{i}\times
n_{i}}$ for $\Phi_{i}^{(t)}=(\varphi_{i,jk}^{(t)})$ as

$\hat{\varphi}_{i,jk}^{(t)}:=\left\\{\begin{array}[]{ll}\varphi_{i,jk}^{(t)}/p_{i,jk}^{(t)\ast},&\text{with
probability}~{}p_{i,jk}^{(t)\ast},\\\ 0,&\text{otherwise},\end{array}\right.$
(3.18)

where $p_{i,jk}^{(t)\ast}$ is given as in equation (3.9). Replacing
$\Phi_{i}^{(t)}$ in subproblem (3.15) with $\hat{\Phi}_{i}^{(t)}$, we obtain

$\min_{\mathbf{u}_{i},\mathbf{v}_{i}}~{}q_{i}(\mathbf{u}_{i},\mathbf{v}_{i};\mu_{i}^{(t)},\hat{\Phi}_{i}^{(t)}),$
(3.19)

which is the dual of

$\begin{array}[]{cl}\min\limits_{X_{i}}&\left\langle\hat{C}_{i}^{(t)},X_{i}-X_{i}^{(t)}\right\rangle+\mu_{i}^{(t)}\mathrm{KL}(X_{i};X_{i}^{(t)}),\\\
\mathrm{s.t.}&X_{i}\mathbf{1}_{n_{i}}=\mathbf{a}_{i},~{}X_{i}^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i},~{}(X_{i})_{(\mathcal{I}_{i}^{(t)})^{c}}=0.\end{array}$
(3.20)

The Sinkhorn algorithm (3.6) then applies with $\Psi_{i}^{(t)}$ replaced by
$\hat{\Phi}_{i}^{(t)}$.

While everything seems to go smoothly, we shall point out some distinctions:
samplings over iterations will not improve the supports to be optimized and
can even result in infeasible subproblems. Recalling the definition (3.16) of
$\Phi_{i}^{(t)}$, one knows from relation (3.17) that $x_{i,jk}^{(t)}=0$
implies $x_{i,jk}^{(t+1,\star)}=0$. If the subproblems are exactly solved and
we perform samplings in two successive iterations (say, $t$ and $t+1$), the
$i$th support to be optimized in the $(t+1)$th iteration is
$\mathcal{I}_{i}^{(t-1)}\cap\mathcal{I}_{i}^{(t)}$, which is a subset of
$\mathcal{I}_{i}^{(t-1)}$. In implementation, we adopt the Sinkhorn algorithm
to inexactly solve the subproblems. Its iterative schemes (3.6) also imply the
inheritance of zero entries. In both cases, samplings over iterations do not
improve the supports to be optimized and can lead to infeasibility due to
repeated intersections. To this end, we choose to perform sampling only in
some critical iteration, say $\hat{t}\in\mathbb{N}$. For $t<\hat{t}$, the
kernel matrices $\Phi_{i}^{(t)}$ are fully computed. With properly chosen
$\hat{t}$, we can expect $X_{i}^{(\hat{t})}$ to capture well the sparsity
pattern. Then, for $t>\hat{t}$, the selected indices are fixed and no sampling
occurs.

We summarize the klalm method with sampling (the S-klalm method) in Algorithm
4.

Algorithm 4 The S-klalm method for solving problem (1.1).

0: $X_{i}^{(0)}\in\mathbb{R}^{m_{i}\times n_{i}}$,
$\mathbf{a}_{i}\in\mathbb{R}^{m_{i}}$, $\mathbf{b}_{i}\in\mathbb{R}^{n_{i}}$
($i=1,\ldots,N$), $\gamma\in[0,1]$,
$\\{n_{s,i}\\}_{i=1}^{N}\subseteq\mathbb{N}$, $\hat{t}$,
$t_{\max}\in\mathbb{N}$.

1: Set $t:=0$.

2: while certain conditions are not satisfied and $t<t_{\max}$ do

3: for $i=1,\ldots,N$ do

4: Select a proximal parameter $\mu_{i}^{(t)}>0$.

5: if $t=\hat{t}$ then

6: Randomly pick a subset
$\mathcal{I}_{i}^{(t)}\subseteq\\{(j,k):j=1,\ldots,m_{i},~{}k=1,\ldots,n_{i}\\}$
following the Poisson sampling with
$P_{i}^{(t)}=(p_{i,jk}^{(t)})\in\mathbb{R}_{+}^{m_{i}\times n_{i}}$ given by
(3.8) and $n_{s,i}$.

7: end if

8: if $t<\hat{t}$ then

9: Let $\hat{\Phi}_{i}^{(t)}:=\Phi_{i}^{(t)}$ defined in the formula (3.16).

10: else

11: Construct $\hat{\Phi}_{i}^{(t)}\in\mathbb{R}^{m_{i}\times n_{i}}$ as in
the formula (3.18) with $\mathcal{I}_{i}^{(\hat{t})}$ and $P_{i}^{(\hat{t})}$.

12: end if

13: Solve subproblem (3.19) or (3.20) to obtain
$X_{i}^{(t+1)}\in\mathbb{R}^{m_{i}\times n_{i}}$.

14: end for

15: Set $t:=t+1$.

16: end while

16: Approximate solution
$(X_{1}^{(t)},\ldots,X_{N}^{(t)})\in\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times
n_{i}}$.

###  3.5 Computational complexities

We compare the single-iteration computational complexities of the eralm,
S-eralm, klalm, and S-klalm methods. In particular, we focus on the cost of
computing the (sparse) kernel matrices, performing importance sampling, and
the subiterations within the Sinkhorn algorithm (3.6). For a summary, see
Table 1, where $s_{\max}\in\mathbb{N}$ is the maximum subiteration number, and
we assume $m_{i}\equiv m\in\mathbb{N}$, $n_{i}\equiv n\in\mathbb{N}$,
$n_{s,i}\equiv n_{s}\in\mathbb{N}$ ($i=1,\ldots,N$) for better readability.

Table 1: A comparison of computational complexities.

| Ingredients
---
in one iteration
eralm | klalm
Kernel matrices |  $Nmn$ entries |  $Nmn$ entries
Subiterations | $s_{\max}\times Nmn$ | $s_{\max}\times Nmn$
| Ingredients
---
in one iteration
S-eralm | S-klalm
Sampling | $\mathcal{O}(Nmn)$ |  $\mathcal{O}(Nmn)$ ($t=\hat{t}$)
Kernel matrices |  $Nn_{s}$ entries |  |  $Nmn$ entries ($t<\hat{t}$)
---
$Nn_{s}$ entries ($t\geq\hat{t}$)
Subiterations | $s_{\max}\times Nn_{s}$ |  |  $s_{\max}\times Nmn$ ($t<\hat{t}$)
---
$s_{\max}\times Nn_{s}$ ($t\geq\hat{t}$)

Armed with warm starts, the Sinkhorn algorithm usually terminates after few
subiterations. As such, the S-klalm method enjoys the lowest complexity per
iteration when $t>\hat{t}$. Given $n_{s}\sim(m+n)^{1+\tau}$ with
$\tau\in(0,1)$, this advantage becomes more evident as $\tau$ tends to $0$ or
$m+n$ goes to $+\infty$.

##  4 Convergence analysis

In this section, we present the convergence and asymptotic properties of both
the eralm and S-eralm methods. Since the KL divergence lacks local Lipschitz
smoothness, the analysis for the klalm and S-klalm methods can be highly
nontrivial and is left over as future work. Nevertheless, the theoretical
results obtained for the eralm and S-eralm methods already deserve a whistle.
To the best of our knowledge, our work is the first attempt in incorporating
randomized matrix sparsification into numerical methods for multi-block
nonconvex optimization, while establishing the convergence and asymptotic
properties.

All the results are obtained by assuming that the subproblems are exactly
solved. Furthermore, the subproblems of the S-eralm method are assumed to be
feasible in their primal forms in all iterations, so that the strong duality
holds. Although at present no analysis is present on the conditions under
which the latter assumption is fulfilled, we remark that sampling parameters
without careful selection (say, $n_{s,i}=\lfloor(m_{i}+n_{i})^{1.5}\rfloor$)
already meets the demand in numerical simulations (see section 5).

To characterize the stationarity violation for problem (1.1), we define the
residual functions: for any
$X:=(X_{1},\ldots,X_{N})\in\bigtimes_{i=1}^{N}\mathbb{R}^{m_{i}\times n_{i}}$
and $i=1,\ldots,N$,

$R_{i}(X):=\max_{T\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})}\left\langle\nabla_{i}f(X),X_{i}-T\right\rangle.$
(4.1)

Moreover, let $R:=\sum_{i=1}^{N}R_{i}$. It is not hard to verify that
$R(X)\geq 0$ for any
$X\in\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$ and that
$X\in\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$ is a
Karush-Kuhn-Tucker point of problem (1.1) if and only if $R(X)=0$ holds. For
the iterate $X^{(t)}$ generated by the eralm or S-eralm method, $R(X^{(t)})$
can thus characterize the stationarity violation at $X^{(t)}$.

We assume Lipschitz smoothness of $f$ over the feasible region, which holds
automatically, e.g., for the quantum physics application (see section 5.1).

######  Assumption 1.

The gradient of the function $f$ is Lipschitz continuous over
$\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$, i.e., there
exists an $L\geq 0$ such that, for $i=1,\ldots,N$,

$\|\nabla_{i}f(X)-\nabla_{i}f(X^{\prime})\|\leq
L\|X-X^{\prime}\|~{}~{}\text{for
all}~{}X,X^{\prime}\in\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}).$

We first give the convergence and asymptotic properties of the eralm method.
The proofs of the theorem and corollary below can be found in Appendix A.

######  Theorem 1.

Suppose that Assumption 1 holds. Let $\\{X^{(t)}\\}$ be the iterate sequence
generated by the eralm method when

$t_{\max}\geq\frac{f(X^{(0)})-\underline{f}}{2\bar{d}^{2}LN(2\sqrt{N}+1)},\quad\alpha_{1}^{(t)}=\cdots=\alpha_{N}^{(t)}\equiv\alpha:=\frac{1}{\bar{d}}\sqrt{\frac{f(X^{(0)})-\underline{f}}{2LN(2\sqrt{N}+1)t_{\max}}},$
(4.2)

and $\lambda_{1}^{(t)}=\cdots=\lambda_{N}^{(t)}\equiv\lambda$ for $0\leq t\leq
t_{\max}$ and subproblems (3.3) are exactly solved, where
$\underline{f}\in\mathbb{R}$ is less than or equal to the optimal value of
problem (1.1),

$d_{i}:=\min\left\\{\sqrt{m_{i}}\left\|\mathbf{a}_{i}\right\|_{\infty},\sqrt{n_{i}}\left\|\mathbf{b}_{i}\right\|_{\infty}\right\\}~{}(i=1,\ldots,N),\quad\bar{d}:=\max_{i=1}^{N}d_{i}.$

Then

$0\leq\frac{1}{t_{\max}}\sum_{t=0}^{t_{\max}-1}R(X^{(t)})\leq
2\bar{d}(2N+1)\sqrt{\frac{L(f(X^{(0)})-\underline{f})}{t_{\max}}}+N\lambda\bar{h},$
(4.3)

where $\bar{h}:=-\min_{i=1}^{N}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})$.

######  Remark 1.

In Theorem 1, as well as the subsequent Corollary 1, Theorem 2 and Corollary
2, the requirements $\alpha_{1}^{(t)}=\cdots=\alpha_{N}^{(t)}$ and
$\lambda_{1}^{(t)}=\cdots=\lambda_{N}^{(t)}$ are assumed for better
readability. The theorems and corollaries can be extended without too much
difficulty to the case where $\alpha_{i}^{(t)}$ and $\lambda_{i}^{(t)}$ vary
across $i\in\\{1,\ldots,N\\}$ and $0\leq t\leq t_{\max}$.

The errors related to the entropy terms are inevitable because the objective
function in the subproblem of the eralm method is not a local approximation
for $f$. Nevertheless, the right-hand side of inequality (4.3) vanishes in the
limit $\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$ after choosing proper $t_{\max}$
and $\lambda$ and imposing Assumption 2 below. This is of particular
importance for large-scale applications.

######  Assumption 2.

  1. (i)

There exists an $\underline{f}\in\mathbb{R}$ such that the optimal value of
problem (1.1) is lower bounded by $\underline{f}$ for any
$\\{m_{i}\\}_{i=1}^{N}$, $\\{n_{i}\\}_{i=1}^{N}\subseteq\mathbb{N}$.

  2. (ii)

There exists a $q>0$ such that, for any $\\{m_{i}\\}_{i=1}^{N}$,
$\\{n_{i}\\}_{i=1}^{N}\subseteq\mathbb{N}$,
$\mathbf{a}_{i}^{\top}\mathbf{1}_{m_{i}}=\mathbf{b}_{i}^{\top}\mathbf{1}_{n_{i}}=1$,
$\max_{j}a_{i,j}\leq q\cdot\min_{j}a_{i,j}$, and $\max_{k}b_{i,k}\leq
q\cdot\min_{k}b_{i,k}$.

  3. (iii)

There exists a $\theta\geq 0$ such that, for any $\\{m_{i}\\}_{i=1}^{N}$,
$\\{n_{i}\\}_{i=1}^{N}\subseteq\mathbb{N}$, the block Lipschitz constant
$L=\mathcal{O}(\sum_{i=1}^{N}(m_{i}+n_{i})^{\theta})$.

  4. (iv)

There exists a $\xi\geq 0$ such that
$\max_{i=1}^{N}(m_{i}+n_{i})/\min_{i=1}^{N}(m_{i}+n_{i})\leq\xi$ for any
$\\{m_{i}\\}_{i=1}^{N}$, $\\{n_{i}\\}_{i=1}^{N}\subseteq\mathbb{N}$.

######  Remark 2.

Items (i) and (ii) in Assumption 2 are motivated by the application of
interest, where a continuous problem arises and one seeks to solve its
discretized version. For example, in the application of strongly correlated
quantum physics (see section 5.1 and also [chen2014numerical, hu2023global]),
the discretized problem has a natural objective lower bound $\underline{f}=0$
due to the nonnegativity of energy, $\mathbf{a}_{i}$ and $\mathbf{b}_{i}$ are
discretization of the so-called single-particle density, whose integral is a
prescribed constant. Incidentally, up to normalizing $f$, we assume the total
mass of marginals to be 1 in item (ii). For simplicity in theoretical
analysis, we adopt items (iii) and (iv), which can be further relaxed to some
extent.

######  Corollary 1.

Suppose that Assumptions 1 and 2 (i)-(iii) hold. Let $\\{X^{(t)}\\}$ be the
iterate sequence generated by the eralm method when

$\displaystyle t_{\max}$
$\displaystyle\geq\max\left\\{\Omega\left(\sum_{i=1}^{N}(m_{i}+n_{i})^{\eta}\right),\frac{f(X^{(0)})-\underline{f}}{2\bar{d}^{2}LN(2\sqrt{N}+1)}\right\\},\quad
f(X^{(0)})\leq M,$ $\displaystyle\alpha_{1}^{(t)}$
$\displaystyle=\cdots=\alpha_{N}^{(t)}\equiv\alpha,\quad\lambda_{1}^{(t)}=\cdots=\lambda_{N}^{(t)}\equiv\lambda=o\left(\frac{1}{\sum_{i=1}^{N}\log
m_{i}n_{i}}\right)$

for $0\leq t\leq t_{\max}$ and subproblems (3.3) are exactly solved, where
$\eta(>\theta)$ and $M$ are constants independent from $\\{m_{i}\\}_{i=1}^{N}$
and $\\{n_{i}\\}_{i=1}^{N}$ and $\alpha$ is defined in equation (4.2). Then
$\sum_{t=0}^{t_{\max}-1}R(X^{(t)})/t_{\max}\to 0$ as
$\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$.

With sampling, the analysis for the S-eralm method is much more involved, in
that the entrywise matrix sparsification results in additional errors. We
establish the convergence and asymptotic results with the help of the theory
of randomized matrix sparsification as well as the following Assumption 3;
details are provided in Appendix B.

######  Assumption 3.

  1. (i)

There exist constants $\nu\in(1/2,1]$, $c_{1}$, $c_{2}$, $\hat{c}_{2}>0$ such
that, for any $0\leq t\leq t_{\max}$ and $i=1,\ldots,N$,

$\|\Psi_{i}^{(t)}\|_{2}\geq\frac{(m_{i}+n_{i})^{\nu}}{c_{1}},\quad\kappa(\Psi_{i}^{(t)})\leq
c_{2},\quad\kappa(\hat{\Psi}_{i}^{(t)})\leq\hat{c}_{2}.$

  2. (ii)

The interpolation factor $\gamma$ is less than 1 and there exists an
$\varepsilon>0$ such that, for $i=1,\ldots,N$,

$\frac{1}{\max_{j,k,t}p_{i,jk}^{(t)}}\geq
n_{s,i}\geq\frac{8(m_{i}+n_{i})^{1-2\nu}\log^{4}(m_{i}+n_{i})}{(1-\gamma)w_{i}\log^{4}(1+\varepsilon)},$

where
$w_{i}:=\min_{j,k}\sqrt{a_{i,j}b_{i,k}}/\sum_{j^{\prime},k^{\prime}}\sqrt{a_{i,j^{\prime}}b_{i,k^{\prime}}}$.

######  Remark 3.

Assumption 3 is adopted for simplicity in theoretical analysis. In particular,
since $w_{i}\leq 1/m_{i}n_{i}$, item (ii) implies the following lower bound
for $n_{s,i}$:

$\frac{8(m_{i}+n_{i})^{1-2\nu}\log^{4}(m_{i}+n_{i})}{(1-\gamma)w_{i}\log^{4}(1+\varepsilon)}\geq\frac{8}{(1-\gamma)\log^{4}(1+\varepsilon)}\frac{m_{i}n_{i}}{(m_{i}+n_{i})^{2\nu-1}}\log^{4}(m_{i}+n_{i}),$

which is of lower order than $m_{i}n_{i}$ because $\nu\in(1/2,1]$. The
condition $n_{s,i}p_{i,jk}^{(t)}\leq 1$ is a convention prevalent in the
literature of Poisson sampling [wang2021comparative, wang2022sampling]. Note
that under Assumption 2 (ii), the entries in $\mathbf{a}_{i}$ or
$\mathbf{b}_{i}$ are of the same order. Therefore,
$w_{i}=\Theta(1/m_{i}n_{i})$ and $p_{i,jk}^{(t)}=\Theta(1/m_{i}n_{i})$ after
choosing a proper value for $\gamma$. Item (ii) then holds if we choose
$n_{s,i}=\Theta(m_{i}n_{i}\log^{4}(m_{i}+n_{i})/(m_{i}+n_{i})^{2\nu-1})$ and
the values of $\varepsilon$ and $\gamma$ are independent from
$\\{m_{i}\\}_{i=1}^{N}$ and $\\{n_{i}\\}_{i=1}^{N}$.

######  Theorem 2.

Suppose that Assumption 1 holds. Let $\\{X^{(t)}\\}$ be the iterate sequence
generated by the S-eralm method when

$t_{\max}\geq\frac{f(X^{(0)})-\underline{f}}{2\bar{d}^{2}LN(2\sqrt{N}+1)},$

$\alpha_{1}^{(t)}=\cdots=\alpha_{N}^{(t)}\equiv\alpha$, and
$\lambda_{1}^{(t)}=\cdots=\lambda_{N}^{(t)}\equiv\hat{\lambda}$ for $0\leq
t\leq t_{\max}$, subproblems (3.11) are feasible and exactly solved, and
Assumption 3 is fulfilled, where $\underline{f}\in\mathbb{R}$ is less than or
equal to the optimal value of problem (1.1) and $\alpha$ is defined in
equation (4.2). Then, for any
$m_{i}+n_{i}>\max\\{152,e^{\sqrt{c_{3}}}\\}~{}(i=1,\ldots,N)$, $\zeta>0$, and
$\iota>0$, with probability no less than

$\prod_{i=1}^{N}\left\\{\left[1-2\exp\left(-\frac{16\zeta^{2}}{\varepsilon^{4}}\log^{4}(m_{i}+n_{i})\right)\right]\left[1-\exp\left(-2\iota^{2}m_{i}n_{i}\right)\right]\right\\}^{t_{\max}},$

there holds

$\displaystyle 0$
$\displaystyle\leq\frac{1}{t_{\max}}\sum_{t=0}^{t_{\max}-1}R(X^{(t)})\leq
2\bar{d}(2N+1)\sqrt{\frac{L(f(X^{(0)})-\underline{f})}{t_{\max}}}+2N\hat{\lambda}\bar{h}$
(4.4)
$\displaystyle+\hat{\lambda}\bar{d}\sum_{i=1}^{N}\sqrt{n_{s,i}+\iota\cdot
m_{i}n_{i}}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}}+\hat{\lambda}\sum_{i=1}^{N}\frac{\hat{c}_{2}c_{3}}{\log^{2}(m_{i}+n_{i})-c_{3}},$

where $c_{3}:=c_{1}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)$.

######  Corollary 2.

Suppose that Assumptions 1 and 2 hold. Let $\\{X^{(t)}\\}$ be the iterate
sequence generated by the S-eralm method when

$\displaystyle t_{\max}$
$\displaystyle=\Theta\left(\sum_{i=1}^{N}(m_{i}+n_{i})^{\eta}\right)~{}\text{satisfying}~{}t_{\max}\geq\frac{f(X^{(0)})-\underline{f}}{2\bar{d}^{2}LN(2\sqrt{N}+1)},\quad
f(X^{(0)})\leq M,$ $\displaystyle\alpha_{1}^{(t)}$
$\displaystyle=\cdots=\alpha_{N}^{(t)}\equiv\alpha,~{}n_{s,i}=\Theta\left(\frac{m_{i}n_{i}}{(m_{i}+n_{i})^{2\nu-1}}\log^{4}(m_{i}+n_{i})\right),$
$\displaystyle\lambda_{1}^{(t)}$
$\displaystyle=\cdots=\lambda_{N}^{(t)}\equiv\hat{\lambda}=o\left(\frac{1}{\sum_{i=1}^{N}\sqrt{m_{i}n_{i}}\log(m_{i}+n_{i})}\right)$

for $0\leq t\leq t_{\max}$, subproblems (3.11) are feasible and exactly
solved, and Assumption 3 is fulfilled, where $c_{1}$, $c_{2}$, $\hat{c}_{2}$,
$\varepsilon$, $\eta(>\theta)$, $\gamma$, $\nu$, $\xi$, and $M$ are
independent from $\\{m_{i}\\}_{i=1}^{N}$ and $\\{n_{i}\\}_{i=1}^{N}$ and
$\alpha$ is defined in equation (4.2). Then
$\sum_{t=0}^{t_{\max}-1}R(X^{(t)})/t_{\max}\to 0$ as
$\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$ with probability going to $1$.

##  5 Numerical experiments

We demonstrate the efficiency of the proposed methods via numerical results on
model one/two/three-dimensional strongly correlated electron systems. We first
describe the related optimization problem of the form (1.1) mathematically and
provide experimental details, including systems to be simulated and default
algorithmic settings. Then numerical comparisons are conducted among the palm
method [bolte2014proximal, hu2023convergence] and the newly designed four
methods on the one-dimensional systems. We integrate those with favorable
performances into a cascadic multigrid optimization framework for the
simulations of two/three-dimensional systems. A first visualization of
approximate OT maps between electron positions in three-dimensional contexts
is provided. Finally, we test the scalability of methods with respect to the
problem size as well as the number of variable blocks.

###  5.1 Problem description

From the strong-interaction limit of density functional theory
[friesecke2023strong], the strongly correlated quantum systems in the strictly
correlated regime can be well understood by solving the multimarginal optimal
transport problems with Coulomb cost (MMOT) [alfonsi2021approximation,
alfonsi2022constrained, benamou2016numerical, buttazzo2012optimal,
chen2014numerical, cotar2013density, friesecke2022genetic, khoo2019convex,
khoo2020semidefinite, mendl2013kantorovich]:

$\begin{array}[]{cl}\min\limits_{\pi}&\displaystyle\int_{(\mathbb{R}^{d})^{N_{e}}}c_{\mathrm{ee}}(\mathbf{r}_{1},\ldots,\mathbf{r}_{N_{e}})\,\mathrm{d}\pi(\mathbf{r}_{1},\ldots,\mathbf{r}_{N}),\\\
\mathrm{s.t.}&\displaystyle\int_{(\mathbb{R}^{d})^{N_{e}-1}}\,\mathrm{d}\pi(\mathbf{r}_{1},\ldots,\mathbf{r}_{i-1},\cdot,\mathbf{r}_{i+1},\ldots,\mathbf{r}_{N_{e}})=\frac{\rho}{N_{e}},~{}~{}i=1,\ldots,N_{e}.\end{array}$
(5.1)

Here, $d\in\\{1,2,3\\}$ is the system dimension, $N_{e}\in\mathbb{N}$ is the
number of electrons, $\mathbf{r}_{i}\in\mathbb{R}^{d}$ refers to the position
of the $i$th electron ($i=1,\ldots,N_{e}$),
$\rho:\mathbb{R}^{d}\to\mathbb{R}_{+}$ is the single-particle density,
satisfying $\int\rho=N_{e}$,
$c_{\mathrm{ee}}(\mathbf{r}_{1},\ldots,\mathbf{r}_{N_{e}}):=\sum_{1\leq
i<j}1/\left\|\mathbf{r}_{i}-\mathbf{r}_{j}\right\|$ stands for the
$N$-particle Coulomb potential, and $\pi$ is a joint probability measure of
$N_{e}$ electron positions.

Since the dimension of the search space in the MMOT (5.1) scales exponentially
with the number of electrons, one could adopt a Monge-like ansatz
[chen2014numerical, hu2023global], which characterizes the electron-electron
couplings explicitly, to transform the MMOT into the following problem:

$\displaystyle\min\limits_{\\{\gamma_{i}\\}_{i=2}^{N_{e}}}$
$\displaystyle\displaystyle\sum_{i=2}^{N_{e}}\int_{(\mathbb{R}^{d})^{2}}\frac{\rho(\mathbf{r})\gamma_{i}(\mathbf{r},\mathbf{r}^{\prime})}{\left\|\mathbf{r}-\mathbf{r}^{\prime}\right\|}\,\mathrm{d}\mathbf{r}\,\mathrm{d}\mathbf{r}^{\prime}+\sum_{2\leq
i<j}\int_{(\mathbb{R}^{d})^{3}}\frac{\rho(\mathbf{r})\gamma_{i}(\mathbf{r},\mathbf{r}^{\prime})\gamma_{j}(\mathbf{r},\mathbf{r}^{\prime\prime})}{\left\|\mathbf{r}^{\prime}-\mathbf{r}^{\prime\prime}\right\|}\,\mathrm{d}\mathbf{r}\,\mathrm{d}\mathbf{r}^{\prime}\,\mathrm{d}\mathbf{r}^{\prime\prime},$
(5.2) $\displaystyle\quad\mathrm{s.t.}$
$\displaystyle\displaystyle\int_{\mathbb{R}^{d}}\gamma_{i}(\cdot,\mathbf{r}_{i})\,\mathrm{d}\mathbf{r}_{i}=1,~{}~{}\int_{\mathbb{R}^{d}}\gamma_{i}(\mathbf{r}_{1},\cdot)\rho(\mathbf{r}_{1})\,\mathrm{d}\mathbf{r}_{1}=\rho,~{}~{}\gamma_{i}\geq
0,~{}~{}i=2,\ldots,N_{e}.$

Here, $\gamma_{i}:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}_{+}$ encodes
the coupling between the first and $i$th electrons, $\rho\gamma_{i}$ can be
understood as the joint probability density between their positions
($i=2,\ldots,N_{e}$).

To numerically solve problem (5.2), we confine the integral domain to some
bounded $\Omega\subseteq\mathbb{R}^{d}$ and adopt finite elements-like
discretization. In particular, we first define a mesh
$\mathcal{T}=\\{e_{k}\\}_{k=1}^{K}$ to divide $\Omega$ into $K$ non-
overlapping elements, i.e., $\cup_{k=1}^{K}e_{k}=\Omega$ and $e_{k}\cap
e_{k^{\prime}}=\emptyset$ whenever $k\neq k^{\prime}$. Then we use a finite
summation of Dirac measures to approximate $\rho$ as
$\rho\approx\sum_{k=1}^{K}\varrho_{k}\delta_{\mathbf{d}_{k}}$, where
$\varrho_{k}:=\int_{e_{k}}\rho$, $\mathbf{d}_{k}\in\mathbb{R}^{d}$ is the
barycenter of the element $e_{k}$ ($k=1,\ldots,K$). Let
$\bm{\varrho}:=[\varrho_{1},\ldots,\varrho_{K}]\in\mathbb{R}_{++}^{K}$ and
$\Lambda:=\mathrm{Diag}(\bm{\varrho})\in\mathbb{R}^{K\times K}$. The two-
particle Coulomb potential and couplings are respectively discretized into
$K\times K$ matrices $C=(c_{kl})\in\mathbb{R}^{K\times K}$ and
$X_{i}=(x_{i,kl})_{kl}\in\mathbb{R}^{K\times K}$ ($i=2,\ldots,N_{e}$), where,
for $i=2,\ldots,N_{e}$ and $k,l=1,\ldots,K$,

$c_{kl}:=\left\\{\begin{array}[]{ll}\left\|\mathbf{d}_{k}-\mathbf{d}_{l}\right\|^{-1},&\text{if}~{}k\neq
l,\\\ 0,&\text{otherwise},\end{array}\right.\quad
x_{i,kl}:=\frac{1}{\left|e_{k}\right|}\int_{e_{k}}\int_{e_{l}}\gamma_{i}(\mathbf{r},\mathbf{r}^{\prime})\,\mathrm{d}\mathbf{r}^{\prime}\,\mathrm{d}\mathbf{r}.$

Note that the diagonal entries in $C$ are set to 0 to avoid numerical
instability. We impose the following extra constraints on
$\\{X_{i}\\}_{i=2}^{N_{e}}$ to maintain the model equivalence:

$\mathrm{Tr}(X_{i})=0,~{}i=2,\ldots,N_{e};~{}\left\langle
X_{i},X_{j}\right\rangle=0,~{}i,j=2,\ldots,N_{e}:i\neq j.$ (5.3)

Intuitively, the constraints (5.3) exclude the cases where two electrons
collide. After penalizing the constraints (5.3) in $\ell_{1}$ form and
transforming $X_{i}$ to $Y_{i}:=\Lambda X_{i}\in\mathbb{R}^{K\times K}$
($i=2,\ldots,N_{e}$), we obtain a multi-block optimization problem over the
transport polytopes of the form (1.1) (with $N=N_{e}-1$):

$\begin{array}[]{cl}\min\limits_{\\{Y_{i}\\}_{i=2}^{N_{e}}}&\displaystyle\sum_{i=2}^{N_{e}}\left\langle
Y_{i},C+\beta\Lambda^{-1}\right\rangle+\sum_{2\leq i<j}\left\langle
Y_{i},\Lambda^{-1}Y_{j}C+\beta\Lambda^{-2}Y_{j}\right\rangle,\\\
\mathrm{s.t.}&Y_{i}\in\mathcal{U}(\bm{\varrho},\bm{\varrho})\subseteq\mathbb{R}^{K\times
K},~{}~{}i=2,\ldots,N_{e}.\end{array}$ (5.4)

The matrix variable $Y_{i}$ can be understood as the transport plan between
the positions of the first and $i$th electrons ($i=2,\ldots,N_{e}$). It has
been shown in [hu2023exactness] that there exists a $\hat{\beta}\geq 0$ such
that the optimal solutions of (5.4) satisfy the constraints (5.3) whenever
$\beta\geq\hat{\beta}$.

###  5.2 Systems under simulations

We consider eight one/two/three-dimensional (1D/2D/3D) systems. Table 2
contains their single-particle densities, domains of interest, and numbers of
electrons.

Table 2: $1$D/$2$D/$3$D systems used for simulations. The second column lists
the unnormalized single-particle densities $\rho$, the third gives the domains
$\Omega$, and the last indicates the numbers of electrons $N_{e}$ in systems.

System No. | $\rho$ | $\Omega$ | $N_{e}$
---|---|---|---
$1$D systems
1 | $\cos(\pi r)+1$ | $[-1,1]$ | 3
2 | $2\rho_{6}(r;-0.5)+1.5\rho_{4}(r;0.5)$ | $[-1.5,1.5]$ | 3
3 | $\rho_{1/\sqrt{\pi}}(r)$ | $[-2,2]$ | 7
4 |  | $\rho_{4}(r;-2)+\rho_{4}(r;-1.5)+\rho_{4}(r;-1)+\rho_{4}(r;-0.5)$
---
$+\rho_{4}(r;2/3)+\rho_{4}(r;4/3)+\rho_{4}(r;2)$
$[-3,3]$ | 7
$2$D systems
5 |  | $\rho_{3}(\mathbf{r};[0,0.96]^{\top})+\rho_{3}(\mathbf{r};[1.032,-0.84]^{\top})+\rho_{3}(\mathbf{r};[-1.032,-0.84]^{\top})$
---
$[-3,3]^{2}$ | 3
6 |  | $2\rho_{3}(\mathbf{r};[0,1.2]^{\top})+\rho_{3}(\mathbf{r};[1.29,-1.05]^{\top})+\rho_{3}(\mathbf{r};[-1.29,-1.05]^{\top})$
---
$[-3,3]^{2}$ | 4
$3$D systems
7 | $\rho_{3}(\mathbf{r};[-1,-1,-1]^{\top})+\rho_{3}(\mathbf{r};[1,1,-1]^{\top})+\rho_{3}(\mathbf{r};[-1,1,1]^{\top})$ | $[-2,2]^{3}$ | 3
8 | $3\rho_{4}(\mathbf{r};[-1,0,0]^{\top})+\rho_{4}(\mathbf{r};[1,0,0]^{\top})$ | $[-2,2]\times[-1,1]^{2}$ | 4

The component function $\rho_{\alpha}(\cdot;\mathbf{c})$ ($\alpha>0$,
$\mathbf{c}\in\mathbb{R}^{d}$) is defined as

$\rho_{\alpha}(\mathbf{r};\mathbf{c}):=\exp\left(-\alpha\left\|\mathbf{r}-\mathbf{c}\right\|^{2}\right),~{}~{}\forall~{}\mathbf{r}\in\mathbb{R}^{d}.$

We illustrate the single-particle densities in Figure 1.

Figure 1: The illustrations of the single-particle densities listed in Table
2. For the 3D systems, we only show the regions where the values of the
single-particle densities are larger than 0.01.

######  Remark 4.

Nearly all the single-particle densities in Table 2 comprise of Gaussian
functions. These settings are reasonable because electrons tend to concentrate
around the nuclei, which are represented by the potential wells. For example,
the single-particle density of system 8 comprises of two Gaussian functions
with weights 1 and 3, respectively. It can describe a dissociating lithium
hydride [filatov2015spin], which has two nuclei with charge numbers 1 and 3,
respectively. We shall note that the applicability of the proposed methods is
independent from the constructions of the single-particle densities.

###  5.3 Default settings

We employ either equimass or equisize discretization (to be fixed later) for
problem (5.4). We set $\beta=1$ for any $K$ following [hu2023global], which is
reasonable because the nonzero entries in both $C$ and $\Lambda^{-1}$ are of
order $K$. For both the eralm and S-eralm methods, we adopt decreasing step
sizes $\alpha_{i}^{(t)}=1/(t+1)^{0.75}$, $i=2,\ldots,N_{e}$, to pursue
convergence. In the S-eralm and S-klalm methods, we let $n_{s,i}=\lfloor
K^{1.5}\rfloor$ ($i=2,\ldots,N_{e}$). All the proposed methods set the
regularization or proximal parameters adaptively as

$\lambda_{i}^{(t)}=\sigma\|\tilde{\mathbf{v}}_{i}^{(t)}\|_{\infty}/(20\log(K)),~{}~{}\mu_{i}^{(t)}=\sigma\|\mathbf{v}_{i}^{(t)}\|_{\infty}/(20\log(K)),~{}~{}i=2,\ldots,N_{e},$
(5.5)

where $\sigma=1$. The selections of $\gamma$ and $\hat{t}$ will be presented
in section 5.4. We shall point out that better performances of the proposed
methods can be expected with more careful parameter tuning. All the
subproblems are solved using the Sinkhorn algorithm (3.6), with warm starts
for acceleration. Incidentally, to circumvent possible underflow and achieve
better convergence rates [dvurechensky2018computational] when using the
Sinkhorn algorithm, we discard the entries in $\bm{\varrho}$ that are smaller
than 0.1% of $\|\bm{\varrho}\|_{\infty}$; this is also practically reasonable,
in that the regions of low probabilities are far less important. With the
abuse of notations, we still denote the truncated vector by $\bm{\varrho}$ so
that other symbols remain unchanged. Regarding the stopping criteria, we
terminate the Sinkhorn algorithm whenever the feasibility violation
$\|Y_{i}^{(t+1)}\mathbf{1}_{K}-\bm{\varrho}\|_{\infty}$ is less than $10^{-6}$
or the subiteration number arrives at $s_{\max}=20$ ($i=2,\ldots,N_{e}$). We
stop the outer loop if

$\Delta^{(t)}:=\frac{1}{N_{e}-1}\sum_{i=2}^{N_{e}}\left\|\Lambda^{-1}\big{(}Y_{i}^{(t)}-Y_{i}^{(t-1)}\big{)}\right\|$

falls below a prescribed $tol>0$ or the iteration number reaches a prescribed
$t_{\max}\in\mathbb{N}$. The specific values of $tol$ and $t_{\max}$ are
detailed in the subsequent subsections. All the experiments presented are run
in a platform with Intel(R) Xeon(R) Gold 6242R CPU @ 3.10GHz and 510GB RAM
running Matlab R2019b under Ubuntu 20.04.

For quantities of interest, we monitor the converged objective value (obj) and
approximate the so-called strictly-correlated-electrons (SCE) potential
[chen2014numerical, di2020optimal, hu2023global] with the output dual
variables. The SCE potential is the functional derivative of the optimal value
of the MMOT with respect to $\rho$ and is important for electronic structure
calculations [chen2014numerical]. Taking the eralm method for example, we
approximate the SCE potential by
$\mathbf{v}:=\tilde{\mathbf{v}}-\min_{k=1}^{K}\\{\tilde{v}_{k}\\}\cdot\mathbf{1}_{K}\in\mathbb{R}^{K}$,
where

$\tilde{\mathbf{v}}:=\frac{1}{N_{e}-1}\sum_{i=2}^{N_{e}}\tilde{\mathbf{v}}_{i}\in\mathbb{R}^{K},$

and $\\{\tilde{\mathbf{v}}_{i}\\}_{i=2}^{N_{e}}$ are the dual solutions
yielded by the Sinkhorn algorithm. In addition, for the cases where explicit
constructions of the optimal solutions to problem (5.4) and the SCE potentials
of the MMOT are available (e.g., in one-dimensional settings
[colombo2015multimarginal, hu2023exactness]), we also evaluate the qualities
of the converged solutions via the relative errors of the objective values
(err_obj) and SCE potentials (err_sce). They are defined respectively as

$\mathrm{err\\_obj}:=\left|\frac{\text{obj}-\text{obj}^{\star}}{\text{obj}^{\star}}\right|,~{}\mathrm{err\\_sce}:=\frac{\|\mathbf{v}-\mathbf{v}^{\star}\|_{\infty}}{\|\mathbf{v}^{\star}\|_{\infty}}.$

Here, $\text{obj}^{\star}\in\mathbb{R}$ denotes the optimal objective value of
problem (5.4) and $\mathbf{v}^{\star}\in\mathbb{R}^{K}$ refers to the vector
made up by the values of the SCE potential at barycenters. For the efficiency
comparison, we record the CPU time in seconds (T).

###  5.4 Algorithm comparisons

We conduct comparisons among the palm method [bolte2014proximal,
hu2023convergence] and the proposed four methods on the 1D systems in Table 2.
In particular, we first test the S-eralm methods with different sampling
probabilities and the S-klalm methods with different choices of $\hat{t}$. As
a byproduct, we select default values of $\gamma$ and $\hat{t}$. Secondly, we
compare the proposed four methods. Those with favorable performances are
tested against the palm method in the third part and used for the simulations
of 2D and 3D systems in section 5.6.

####  5.4.1 Comparisons among the S-eralm methods with different sampling
probabilities

We consider randomly generated sampling probabilities and importance sampling-
based probabilities (3.8) ($\gamma\in\\{0.1,0.3,0.5,0.7,0.9,0.99,0.999\\}$) on
system 1 with $K=90$ (equimass discretization). For each setting of sampling
probability, 10 random trials are generated by the built-in function “rand” in
Matlab. The stopping parameters are $tol=5\times 10^{-3}$ and
$t_{\max}=+\infty$. We record the achieved err_obj, err_sce, and required T
averaged over 10 trials for each setting in Table 3.

Table 3: The achieved err_obj, err_sce, and required T averaged over 10 trials given by the S-eralm methods with different sampling probabilities on system 1 with $K=90$ (equimass discretization). Sampling Prob. | err_obj | err_sce | T
---|---|---|---
Random | 0.4184 | 0.86 | 128.84
$\gamma=0.1$ | 0.3098 | 0.69 | 120.16
$\gamma=0.3$ | 0.1729 | 0.54 | 93.30
$\gamma=0.5$ | 0.1044 | 0.42 | 65.10
$\gamma=0.7$ | 0.0693 | 0.39 | 48.13
$\gamma=0.9$ | 0.0597 | 0.37 | 35.05
$\gamma=0.99$ | 0.0525 | 0.36 | 21.76
$\gamma=0.999$ | 0.1118 | 0.34 | 5.66

Though lacking theoretical justifications, the sampling probabilities
incorporating information of previous iterates are found to yield lower errors
within less CPU time than randomly generated ones. Increasing the value of
$\gamma$ contributes to less CPU time for fulfilling the stopping criterion,
yet worsening the accuracy once surpassing some threshold. We select
$\gamma=0.99$ in the S-eralm and S-klalm methods for a compromise between
accuracy and efficiency in the subsequent experiments.

####  5.4.2 Comparisons among the S-klalm methods with different $\hat{t}$

We conduct numerical comparisons on system 1 in Table 2 with equimass
discretization and $K\in\\{90,180,360,720\\}$. We call the S-klalm method with
$\hat{t}\in\\{0,5,10\\}$. For each pair of $(K,\hat{t})$, 10 random trials are
performed. The stopping parameters are
$tol=10^{-3}\times\sqrt{2}^{\log_{2}(K/90)}$2)2)2)We increase $tol$ since the
size of $\mathcal{U}(\bm{\varrho},\bm{\varrho})$ decreases as
$\mathcal{O}(1/\sqrt{K})$; see Lemma 2. and $t_{\max}=+\infty$. We depict the
achieved err_obj, err_sce, and required T averaged over 10 trials for each
pair of $(K,\hat{t})$ in Figure 2.

Figure 2: The achieved err_obj, err_sce, and required T averaged over 10
trials for each pair of $(K,\hat{t})$ given by the S-klalm method with
different values of $\hat{t}$ on system 1 (equimass discretization). The blue,
green, and purple dashed lines with right-pointing triangle markers are the
results of the S-klalm method with $\hat{t}=0,5,10$, respectively. Left:
err_obj. Middle: err_sce. Right: T.

From Figure 2, we observe that the S-klalm method with a positive $\hat{t}$
does yield higher-quality solutions within comparable CPU time. But the
advantage becomes less obvious as $K$ grows large. Since large-scale problems
are common in the applications of interest and a positive $\hat{t}$ indicates
the computations and storage of full matrices during the first iterations, we
choose $\hat{t}=0$ in the S-klalm method for the ensuing numerical
experiments.

####  5.4.3 Comparisons among the S-eralm and S-klalm methods

We first test the proposed four methods with $\sigma$ taking its value in
$\\{2^{0},2^{2},2^{4},2^{6},2^{8}\\}$ on system 1 with $K=90$ and equimass
discretization. For each value of $\sigma$, 10 random trials are performed.
The stopping parameters are $tol=5\times 10^{-3}/\sqrt{\sigma}$ and
$t_{\max}=+\infty$. We record the achieved err_obj, err_sce and required T
averaged over 10 random trials for each value of $\sigma$ in Figure 3.

Figure 3: The achieved err_obj, err_sce, and required T averaged over 10
trials for each value of $\sigma$ given by the eralm, klalm, S-eralm, and
S-klalm methods on system 1 with $K=90$ (equimass discretization). The red
solid and dashed lines with triangle markers represent the results of the
eralm and S-eralm methods, respectively. The blue solid and dashed lines with
right-pointing triangle markers represent the results of the klalm and S-klalm
methods, respectively. Left: err_obj. Middle: err_sce. Right: T.

From Figure 3, we observe that (i) the objective errors of the eralm and
S-eralm methods rise quickly as $\sigma$ increases, which conforms to the
theoretical results in section 4; (ii) the klalm and S-klalm methods yield
high-quality solutions regardless of the choices of $\sigma$, demonstrating
their robustness to the choices of proximal parameters. The robustness is
practically desirable because a tiny $\sigma$, as is needed by the (S-)eralm
methods to achieve high accuracy, can result in numerical underflow.

In the above settings, we also notice that the S-eralm method arrives at
better solutions than those given by the S-klalm method when a small $\sigma$
is used. But we shall point out that the advantage no longer persists as $K$
increases. We test the S-eralm and S-klalm methods with $\sigma=1$ on system 1
with $K\in\\{90,180,360,720\\}$ (equimass discretization). For each value of
$K$, 10 random trials are performed. The stopping parameters are
$tol=10^{-3}\times\sqrt{2}^{\log_{2}(K/90)}$ and $t_{\max}=+\infty$. We depict
the achieved err_obj, err_sce, and required T averaged over 10 random trials
for each value of $K$ in Figure 4.

Figure 4: The achieved err_obj, err_sce, and required T averaged over 10
trials for each value of $K$ given by the S-eralm and S-klalm methods with
$\sigma=1$ on system 1 (equimass discretization). The red dashed lines with
triangle markers represent the results of the S-eralm method. The blue dashed
lines with right-pointing triangle markers represent the results of the
S-klalm method. Left: err_obj. Middle: err_sce. Right: T.

From Figure 4, we conclude that the advantage of the S-eralm method over the
S-klalm method in terms of accuracy disappears as the problem size increases,
along with fast growing computational time. This is in line with Table 1,
i.e., the matrix entrywise sampling consumes quadratic complexity in each
iteration of the S-eralm method.

Finally, we compare the klalm and S-klalm methods. Aiming at a problem
restricted on the sampled support, it is impractical for the S-klalm method to
outperform the klalm method in accuracy. However, if the computational budget
is limited, the S-klalm method can yield relatively high-quality solutions
within much less CPU time, particularly when the problem size goes large. We
demonstrate this point by numerical comparisons on system 1. We employ
equimass discretization with $K\in\\{90,180,360,720,1440,2880\\}$. We call the
klalm and S-klalm methods for each value of $K$ with 10 random trials. The
stopping parameters are $tol=10^{-3}\times\sqrt{2}^{\log_{2}(K/90)}$ and
$t_{\max}=+\infty$. We depict the convergence curves of err_obj along with the
CPU time averaged over 10 trials for each value of $K$ in Figure 5.

Figure 5: The convergence curves of err_obj along with the CPU time averaged
over 10 trials for each value of $K$ given by the klalm and S-klalm methods on
system 1 (equimass discretization). The blue solid and dashed lines stand for
the results of the klalm and S-klalm methods, respectively. The notation
$T_{\mathrm{inter}}$ refers to the CPU time where the curves of two methods
intersect for the last time.

The notation $T_{\mathrm{inter}}$ refers to the CPU time where the curves of
two methods intersect for the last time.

From Figure 5, we observe that the klalm method attains worse accuracy than
the S-klalm method until the CPU time touches $T_{\mathrm{inter}}$. Moreover,
as we increase the value of $K$, $T_{\mathrm{inter}}$ grows at a cubic rate;
see Figure 6 for an illustration of cubic polynomial fitting.

Figure 6: Cubic polynomial fitted relation between $T_{\mathrm{inter}}$ and
$K$ on system 1 (equimass discretization). The fitted polynomial is
$9.09\times 10^{-7}K^{3}-1.86\times 10^{-3}K^{2}+1.00K-106.38$. The blue
pentagons are the obtained $T_{\mathrm{inter}}$ in experiments and the blue
dashdotted line is the fitted relation.

These results imply that (i) if the problem size is small (e.g., of order
$10^{2}$), the klalm method achieves high accuracy within acceptable CPU time;
(ii) if the problem size is relatively large (e.g., of order $10^{3}$ or
higher), as is usually the case in the applications of interest, and the
computational budget is limited, the S-klalm method is more preferable from
the practical perspective.

Based on the above numerical findings, we choose the klalm and S-klalm methods
for the comparisons with the palm method as well as the simulations of the 2D
and 3D systems in section 5.6.

####  5.4.4 Comparisons among the palm, klalm, and S-klalm methods

We perform numerical comparisons on the 1D systems in Table 2 with equimass
discretization. The implementation of the palm method follows [hu2023global],
except that we tune the proximal parameters adaptively as in equation (5.5)
for fair comparisons. For systems 1 and 2, we consider
$K\in\\{90,180,360,720\\}$. For each value of $K$, 10 random trials are
performed. The stopping parameters are
$tol=10^{-3}\times\sqrt{2}^{\log_{2}(K/90)}$ and $t_{\max}=+\infty$. For
systems 3 and 4, we consider $K\in\\{140,280,560,1120\\}$. For each value of
$K$, 10 random trials are performed. The stopping parameters are
$tol=10^{-3}\times\sqrt{2}^{\log_{2}(K/140)}$ and $t_{\max}=+\infty$. The
achieved err_obj, err_sce, and required T averaged over 10 random trials given
by the three methods are gathered in Figures 7 and 8.

(a) System 1

(b) System 2

Figure 7: The achieved err_obj, err_sce, and required T averaged over 10
trials for each value of $K$ given by the palm, klalm, and S-klalm methods on
the three-electron 1D systems (equimass discretization). The orange solid
lines with square markers are the results of the palm method. The blue solid
and dashed lines with right-pointing triangle markers are the results of the
klalm and S-klalm methods, respectively. From left to right: err_obj, err_sce,
and T. (a) System 1. (b) System 2.

(a) System 3

(b) System 4

Figure 8: The achieved err_obj, err_sce, and required T averaged over 10
trials for each value of $K$ given by the palm, klalm, and S-klalm methods on
the seven-electron 1D systems (equimass discretization). The orange solid
lines with square markers are the results of the palm method. The blue solid
and dashed lines with right-pointing triangle markers are the results of the
klalm and S-klalm methods, respectively. From left to right: err_obj, err_sce,
and T. (a) System 3. (b) System 4.

Under the same settings of proximal and stopping parameters, the palm method
suffers from premature convergence, while the klalm method yields much smaller
objective errors across the 1D systems. The multiplicative expression (3.17)
further enables the use of randomized matrix sparsification, leading to the
S-klalm method with better scalability.

###  5.5 A cascadic multigrid optimization framework

Before the simulations of 2D and 3D systems, we introduce a cascadic multigrid
(cmg) optimization framework for problem (5.4). The cmg optimization framework
has demonstrated its power in accelerating the numerical solution of large-
scale scientific problems, especially those with governing partial
differential equations; see, e.g., [borzi2008multigrid, chen2017efficient,
xia2018cascadic, liu2021multilevel]. In general, the framework begins with a
coarse mesh, with which a small-scale problem is associated. An accurate
solver is then called to solve the small-scale problem and yields a solution.
Later, it basically repeats the following three steps: (i) refine the previous
coarse mesh to a finer one; (ii) prolongate the previous solution and
construct an initial point over the current mesh; (iii) start a local solver
from the initial point and obtain a solution over the current mesh. With
carefully designed prolongation operator and local solver, the framework can
fully utilize the solution information in the coarser meshes and considerably
accelerate the numerical solution of the large-scale problems over the finer
meshes.

In our context, we call accurate solvers, such as the klalm method, to tackle
the problems over the initial coarse meshes. After uniform or adaptive mesh
refinements, we prolongate the previous solution (say,
$\big{(}Y_{2}^{(\ell-1,\star)},\ldots,Y_{N_{e}}^{(\ell-1,\star)}\big{)}$) to
the current mesh (say, $\\{e_{k}^{(\ell)}\\}_{k=1}^{K^{(\ell)}}$) as follows:

$y_{i,kl}^{(\ell,0)}:=\frac{1}{K_{k^{\prime}l^{\prime}}^{(\ell)}}y_{i,k^{\prime}l^{\prime}}^{(\ell-1,\star)},\quad
k,l=1,\ldots,K^{(\ell)},~{}~{}i=2,\ldots,N_{e},$ (5.6)

where $k^{\prime}$, $l^{\prime}\in\\{1,\ldots,K^{(\ell-1)}\\}$ are such that
$e_{k}^{(\ell)}\subseteq e_{k^{\prime}}^{(\ell-1)}$, $e_{l}^{(\ell)}\subseteq
e_{l^{\prime}}^{(\ell-1)}$, and

$K_{k^{\prime}l^{\prime}}^{(\ell)}:=\left|\left\\{(k,l):e_{k}^{(\ell)}\subseteq
e_{k^{\prime}}^{(\ell-1)},~{}e_{l}^{(\ell)}\subseteq
e_{l^{\prime}}^{(\ell-1)},~{}k,l\in\\{1,\ldots,K^{(\ell)}\\}\right\\}\right|.$

The design of the prolongation operator follows that $\rho\gamma_{i}$ can be
understood as a joint probability density between the positions of the first
and $i$th electrons ($i=2,\ldots,N_{e}$). Finally, cheap local solvers, such
as the S-klalm method, start from
$\big{(}Y_{2}^{(\ell,0)},\ldots,Y_{N_{e}}^{(\ell,0)}\big{)}$ and solve the
problem over the current mesh. Consequently, apart from the previously
mentioned warm starts, the prolongation operator also enables “warm
samplings”; please refer to the importance sampling probability (3.8). We
summarize our cmg optimization framework in Framework 5.

Framework 5 The cmg optimization framework for problem (5.4).

0: Discretization and refinement oracles, accurate and cheap local solvers,
initial number of finite elements $K^{(0)}\in\mathbb{N}$.

1: Set $\ell:=0$.

2: while certain conditions are not satisfied do

3: if $\ell=0$ then

4: Discretization: discretize the MMOT into problem (5.4) with $K^{(0)}$
finite elements $\\{e_{k}^{(0)}\\}_{k=1}^{K^{(0)}}\subseteq\mathbb{R}^{d}$.

5: Construct a random initial point
$(Y_{2}^{(0,0)},\ldots,Y_{N_{e}}^{(0,0)})\in(\mathbb{R}^{K^{(\ell)}\times
K^{(\ell)}})^{N_{e}-1}$.

6: Local solution: start the accurate local solver from
$(Y_{2}^{(0,0)},\ldots,Y_{N_{e}}^{(0,0)})$ for problem (5.4) at level 0 and
obtain $(Y_{2}^{(0,\star)},\ldots,Y_{N_{e}}^{(0,\star)})$.

7: else

8: Grid refinement: refine $\\{e_{k}^{(\ell-1)}\\}_{k=1}^{K^{(\ell-1)}}$ to
$\\{e_{k}^{(\ell)}\\}_{k=1}^{K^{(\ell)}}\subseteq\mathbb{R}^{d}$ with
$K^{(\ell)}\in\mathbb{N}$.

9: Prolongation: prolongate
$(Y_{2}^{(\ell-1,\star)},\ldots,Y_{N_{e}}^{(\ell-1,\star)})$ to the current
mesh through equation (5.6) and obtain
$(Y_{2}^{(\ell,0)},\ldots,Y_{N_{e}}^{(\ell,0)})\in(\mathbb{R}^{K^{(\ell)}\times
K^{(\ell)}})^{N_{e}-1}$.

10: Local solution: start the cheap local solver from
$(Y_{2}^{(\ell,0)},\ldots,Y_{N_{e}}^{(\ell,0)})$ for problem (5.4) at level
$\ell$ and obtain $(Y_{2}^{(\ell,\star)},\ldots,Y_{N_{e}}^{(\ell,\star)})$.

11: end if

12: Set $\ell:=\ell+1$.

13: end while

13: Approximate solution
$(Y_{2}^{(\ell-1,\star)},\ldots,Y_{N_{e}}^{(\ell-1,\star)})$.

To demonstrate the utility of Framework 5, we include numerical experiments on
the 1D systems. We call the cmg optimization method, which takes respectively
the klalm and S-klalm methods as the accurate and cheap local solvers, the
S-klalm-cmg method. We compare the numerical performances of the klalm,
S-klalm, and S-klalm-cmg methods. For systems 1 and 2, the klalm and S-klalm
methods directly solve the problems with $K=720$ and equimass discretization
starting from 10 random trials. The stopping parameters are
$tol=2\sqrt{2}\times 10^{-3}$ and $t_{\max}=+\infty$. The S-klalm-cmg method
starts 10 random trials from $K^{(0)}=90$ with equimass discretization and
reaches the desired the mesh after equimass refinements for three times. The
stopping parameters at level $\ell\in\\{0,1,2,3\\}$ are
$tol=10^{-3}\times\sqrt{2}^{\log_{2}(K^{(\ell)}/K^{(0)})}$ and
$t_{\max}=+\infty$. The settings for the simulations of systems 3 and 4 are
analogous, except that the problem size handled by the klalm and S-klalm
methods is $K=1120$ and the initial problem size handled by the S-klalm-cmg
method is $K^{(0)}=140$. We report the achieved err_obj, err_sce, and required
T averaged over 10 random trials given by the three methods on the 1D systems
in Table 4.

Table 4: The achieved err_obj, err_sce, and required T averaged over 10 trials
given by the klalm, S-klalm, and S-klalm-cmg methods on the 1D systems
(equimass discretization).
Algorithms | System 1 ($K=720$) | System 2 ($K=720$)
---|---|---
err_obj | err_sce | T | err_obj | err_sce | T
klalm | 0.0006 | 0.29 | 7014.57 | 0.0016 | 0.31 | 14190.02
S-klalm | 0.0159 | 0.34 | 180.51 | 0.0183 | 0.34 | 184.96
S-klalm-cmg | 0.0032 | 0.34 | 96.32 | 0.0044 | 0.38 | 109.39
Algorithms | System 3 ($K=1120$) | System 4 ($K=1120$)
err_obj | err_sce | T | err_obj | err_sce | T
klalm | 0.0031 | 0.27 | 288688.75 | 0.0035 | 0.30 | 323734.12
S-klalm | 0.0231 | 0.28 | 2034.49 | 0.0238 | 0.33 | 1910.36
S-klalm-cmg | 0.0074 | 0.32 | 1204.31 | 0.0079 | 0.36 | 1252.50

The S-klalm-cmg method is found to yield relatively high-quality solutions
within the least CPU time.

###  5.6 Simulations on 2D and 3D systems

We use the S-klalm-cmg method for simulating the $2$D and $3$D systems in
Table 2. The initial step employs equisize discretization and the latter ones
refine the meshes uniformly. For the $2$D systems, the initial mesh size is
$K^{(0)}=900$ before truncation. After grid refinements for three times, we
get $K^{(3)}=57600$ before truncation. For the $3$D system 7, we set
$K^{(0)}=1728$ before truncation and perform grid refinements twice to arrive
at $K^{(2)}=110592$ before truncation. For the $3$D system 8, we set
$K^{(0)}=1000$ before truncation and perform grid refinements twice to arrive
at $K^{(2)}=64000$ before truncation. The stopping parameters are
$t_{\max}=10^{4}$ and

$tol=\left\\{\begin{array}[]{ll}5\times 10^{-3},&\ell=0,\\\
10^{-2}\times(\sqrt{2^{d}})^{\log_{2}(K/K^{(0)})},&\ell>0.\end{array}\right.$

Note that for these systems, no explicit constructions of the optimal
solutions are available. Hence, we monitor instead the evolution of objective
values and the approximate SCE potentials. Moreover, following [hu2023global],
we approximate the OT maps between electron positions in the so-called Monge
ansatz [seidl1999strictly, seidl1999strong, seidl2000simulation], which are of
particular physical interest, as

$\mathscr{T}_{i}(\mathbf{d}_{j}):=\sum_{1\leq k\leq
K}y_{i,jk}\mathbf{d}_{k}/\varrho_{j},~{}j=1,\ldots,K,~{}i=2,\ldots,N_{e},$

where $\mathscr{T}_{i}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ represents the
transport map between the positions of the first and $i$th electrons. We
collect the results in Table 5 as well as Figures 9 and 10, where
$K_{\mathrm{trunc}}\in\mathbb{N}$ refers to the dimension of the truncated
$\bm{\varrho}$.

Table 5: The objective values calculated along the iteration of the S-klalm-cmg method when solving the $2$D and $3$D systems. The notation $K_{\mathrm{trunc}}\in\mathbb{N}$ is the dimension of the truncated $\bm{\varrho}$. Step | System 5 | System 6
---|---|---
$K$ | $K_{\mathrm{trunc}}$ | obj | $K$ | $K_{\mathrm{trunc}}$ | obj
0 | 900 | 424 | 1.1339 | 900 | 408 | 3.0690
1 | 3600 | 1622 | 1.1337 | 3600 | 1534 | 3.0690
2 | 14400 | 6410 | 1.1335 | 14400 | 6068 | 3.0677
3 | 57600 | 25562 | 1.1334 | 57600 | 24176 | 3.0667
Step | System 7 | System 8
$K$ | $K_{\mathrm{trunc}}$ | obj | $K$ | $K_{\mathrm{trunc}}$ | obj
0 | 1728 | 780 | 1.0202 | 1000 | 720 | 4.6193
1 | 13824 | 5628 | 1.0209 | 8000 | 5272 | 4.6716
2 | 110592 | 42936 | 1.0209 | 64000 | 40764 | 4.6833

(a) System 5

(b) System 6

(c) System 7

(d) System 8

Figure 9: The evolution of the approximate SCE potentials for the $2$D and
$3$D systems given by the S-klalm-cmg method. (a) System $5$. (b) System
$6$.(c) System $7$. (d) System $8$.

(a) System 5 (b) System 6

(c) System 7 (d) System 8

Figure 10: The approximate OT maps between electron positions for the $2$D and
$3$D systems given by the S-klalm-cmg method. The pre-image $\omega$ in red
stands for the positions of the first electron, while the areas in other
colors are the associated positions of the other electrons, respectively,
determined by $\\{\mathscr{T}_{i}\\}_{i=2}^{N_{e}}$. (a) System $5$. (b)
System $6$. (c) System $7$. (d) System $8$.

We illustrate the approximate OT maps $\\{\mathscr{T}_{i}\\}_{i=2}^{N_{e}}$
through their images of the barycenters of the finite elements within some
given subregion $\omega\subseteq\Omega$.

The results in Table 5 and Figure 9 showcase the convergence of the cmg
optimization framework, also manifesting the effectiveness of the prolongation
operator therein. The approximate OT maps in Figure 10 are in line with
physical intuitions. In particular, for the three-Gaussian $2$D system 5,
Figure 10 (a) implies that if the first electron is around one Gaussian
center, the other two electrons will go near the other two Gaussian centers,
respectively. For the four-Gaussian $2$D system 6, Figure 10 (b) shows that if
one electron revolves around but stays away from the top Gaussian center, the
other three electrons will stay near the three Gaussian centers, respectively,
with one of which surrounded by the first one but keeping away from each
other. For the three-Gaussian $3$D system 7, the results depicted in Figure 10
(c) are analogous to those in Figure 10 (a). For the two-Gaussian $3$D system
8, Figure 10 (d) indicates that if one electron lies around the Gaussian
center $[1,0,0]^{\top}$, the other three will be located around the other
center, with their positions trisecting a sphere. In Figure 10, the positions
of the electrons determined by the approximate OT maps
$\\{\mathscr{T}_{i}\\}_{i=2}^{N_{e}}$ conform to the repulsive law and are
improved along our optimization process. Notably in Figure 10, we provide the
first visualization of the approximate OT maps between electron positions in
$3$D contexts.

######  Remark 5.

There have been extensive works dedicated to the numerical solutions of the
MMOT (5.1). The authors of [mendl2013kantorovich] consider the Kantorovich
dual of the MMOT, whose number of inequality constraints increases
exponentially with $N_{e}$. The authors of [benamou2016numerical] investigate
the entropy regularized MMOT, where curse of dimensionality still resides. The
authors of [khoo2019convex] derive a convex semidefinite programming
relaxation for the discretized $N$-representability form of the MMOT, where
the problem size is independent from the value of $N_{e}$ but the gap induced
by the convex relaxation remains elusive. The authors of
[alfonsi2022constrained] exploit a moment-constrained relaxation of the MMOT,
which admits sparse optimal solutions but entails careful selections of
parameters and test functions to achieve satisfactory approximations. In both
[chen2014numerical] and our work, the reformulation of the MMOT under the
Monge-like ansatz is adopted, where the problem size increases linearly with
respect to $N_{e}$. However, at the moment, the ansatz is provably true only
for special (e.g., two-electron/one-dimensional/special radially symmetric)
cases [friesecke2023strong].

The considerable distinctions in models make a fair numerical comparison
difficult and go beyond our scope. Nevertheless, we shall note that the model
under the Monge-like ansatz explicitly characterizes the electron-electron
couplings. Therefore, in contrast with others, it enables the evaluation of
solution qualities through approximate OT maps between electron positions; see
Figure 10. Additionally, in comparison with [chen2014numerical] where the
authors only consider two-electron systems (i.e., H2 molecule) with
discretization size $K\sim 3000$, we solve larger-scale problems ($K\sim
10^{5}$) and simulate systems with more electrons ($N_{e}=3\sim 7$).

###  5.7 Scalability tests

Finally, we conduct scalability tests for the klalm and S-klalm methods with
respect to $K$ and $N_{e}$. The system under simulation shares the same
normalized single-particle density $\rho$ with system 1 in Table 2, yet with
varying $N_{e}$. We employ equimass discretization.

For the tests with respect to $K$, we fix $N_{e}=3$ and consider
$K\in\\{90,180,360,720\\}$. For each value of $K$, the two methods are called
with 10 random trials. The proximal parameter is set to $\mu_{i}^{(t)}\equiv
0.05$. The stopping parameters are $tol=10^{-3}$ and $t_{\max}=+\infty$. The
achieved err_obj, err_sce, and required T averaged over 10 trials are gathered
in Figure 11 (a).

(a) Scalability with respect to $K$

(b) Scalability with respect to $N_{e}$

Figure 11: The achieved err_obj, err_sce, and required T averaged over 10
trials for each value of $K$ and $N_{e}$ given by the klalm and S-klalm
methods on the cosine-type system (equimass discretization). The blue solid
and dashed lines represent the results of the klalm and S-klalm methods,
respectively. The pink solid and dashed lines denote the fitted relations
between T and $K$ (or $N_{e}$) when using the klalm and S-klalm methods,
respectively. Left: err_obj. Middle: err_sce. Right: T. (a) Scalability with
respect to $K$. For the klalm method, $\text{T}\sim K^{2.6}$. For the S-klalm
method, $\text{T}\sim K^{2.2}$. (b) Scalability with respect to $N_{e}$. For
the klalm method, $\text{T}\sim N_{e}^{1.3}$. For the S-klalm method,
$\text{T}\sim N_{e}^{1.1}$.

In light of Table 1, we fit the linear relation between $\log(\text{T})$ and
$\log(K)$ with linear least squares and obtain the following:

$\displaystyle\text{{{klalm}}:}\quad\log(\text{T})$ $\displaystyle~{}\approx
2.59\log(K)-7.22;$ $\displaystyle\text{{{S-klalm}}:}\quad\log(\text{T})$
$\displaystyle~{}\approx 2.21\log(K)-8.31.$

For the tests with respect to $N_{e}$, we fix $K=144$ and vary $N_{e}$ in
$\\{3,6,12,24,48\\}$. For each value of $N_{e}$, the two methods are called
with 10 random trials. We fix the proximal parameter $\mu_{i}^{(t)}\equiv
20/\log(K)$ because $\|\mathbf{v}_{i}^{(t)}\|_{\infty}$ grows with $N_{e}$.
The stopping parameters are $tol=5\times 10^{-3}$ and $t_{\max}=+\infty$. The
achieved err_obj, err_sce, and required T averaged over 10 trials are
collected in Figure 11 (b). Likewise, we obtain the following relations:

$\displaystyle\text{{{klalm}}:}\quad\log(\text{T})$ $\displaystyle~{}\approx
1.31\log(N_{e})+0.32;$ $\displaystyle\text{{{S-klalm}}:}\quad\log(\text{T})$
$\displaystyle~{}\approx 1.06\log(N_{e})+0.49.$

The obtained scalings of the klalm and S-klalm methods with respect to $N_{e}$
are in rough accordance with the computational complexities summarized in
Table 1, where $\tau=0.5$. The marked deviation in the scalings of the klalm
and S-klalm methods with respect to $K$ can be ascribed to the implementation
in software as well as the difference in the numbers of iterations for each
value of $K$.

##  6 Conclusions

Leveraging the tools from OT and the importance sampling technique, we
introduce novel bcd-type methods, the (S-)eralm and (S-)klalm methods, for the
multi-block optimization problems over the transport polytopes. These methods
enjoy highly scalable schemes for the subproblems and save considerable
expenditure for the calculations and storage in large-scale contexts. With the
theory of randomized matrix sparsification, we establish for the eralm and
S-eralm methods that the average stationarity violations tend to 0 as the
problem size increases to $+\infty$ (with probability going to $1$). To the
best of our knowledge, our work is the first attempt in applying the matrix
entrywise sampling technique to multi-block nonconvex settings with
theoretical guarantees. In the simulations of strongly correlated electron
systems, though lacking convergence analysis, the (S-)klalm methods exhibit
desirable robustness to the choices of proximal parameters; compared with the
klalm method, the S-klalm method trades accuracy for efficiency. Both the
klalm and S-klalm methods are further brought together under a cascadic
multigrid optimization framework, pursuing a decent balance for large-scale
simulations. The numerical results conform to both theoretical predictions and
physical intuitions. The better scalability of the sampling-based methods
allows the first visualization of the approximate OT maps for the 3D systems.

For future work, we are eager to provide some insights into the theoretical
advantages of importance sampling, as evidenced by the numerical results in
section 5.4.1. It is also of special interest in optimization to investigate
the convergence properties for the proposed KL divergence-based methods by
exploring routes that do not rest on the local Lipschitz smoothness of the
proximal terms. Computationally, further acceleration can be gained via other
techniques, such as support identification [lee2023accelerating]. In terms of
quantum physics application, it is necessary to study the solution landscape
for problem (5.4). Our numerical findings suggest that the objective errors at
stationary points decrease as the discretization becomes finer.

##  Appendix A Convergence of the eralm method

This part includes the proofs of the results concerning the eralm method. We
first adapt a lemma from [kerdoncuff2021sampled], characterizing the objective
error induced by entropy regularization on an arbitrary OT problem.

######  Lemma 1.

Let $W\in\mathbb{R}^{m\times n}$, $\mathbf{a}\in\mathbb{R}^{m}$,
$\mathbf{b}\in\mathbb{R}^{n}$, $\lambda>0$, and
$\mathcal{I}\subseteq\\{(j,k):j=1,\ldots,m,~{}k=1,\ldots,n\\}$. Suppose that
$T^{\prime}$, $T^{\prime\prime}\in\mathbb{R}^{m\times n}$ are the optimal
solutions of

$\min_{T}~{}\left\langle
W,T\right\rangle,~{}\mathrm{s.t.}~{}T\in\mathcal{U}(\mathbf{a},\mathbf{b}),~{}T_{\mathcal{I}^{c}}=0$

and

$\min_{T}~{}\left\langle W,T\right\rangle+\lambda
h(T),~{}\mathrm{s.t.}~{}T\in\mathcal{U}(\mathbf{a},\mathbf{b}),~{}T_{\mathcal{I}^{c}}=0,$

respectively. Then $0\leq\left\langle
W,T^{\prime\prime}-T^{\prime}\right\rangle\leq-\lambda
h(\mathbf{a}\mathbf{b}^{\top})$.

######  Proof.

See the proof of [kerdoncuff2021sampled, Lemma 1], which leverages the fact
that $0\geq h(T)\geq h(\mathbf{a}\mathbf{b}^{\top})$ for all
$T\in\mathcal{U}(\mathbf{a},\mathbf{b})$. ∎

Next, we provide an upper bound for the size of
$\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$ ($i\in\\{1,\ldots,N\\}$).

######  Lemma 2.

For $i\in\\{1,\ldots,N\\}$, there holds

$\|T-T^{\prime}\|\leq 2d_{i},~{}~{}\text{for
all}~{}T,T^{\prime}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}).$

######  Proof.

It suffices to bound the norm of an arbitrary member in
$\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$. Note that, for any
$T\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$,

$\displaystyle\|T\|$
$\displaystyle=\sqrt{\sum_{j=1}^{m_{i}}\sum_{t=1}^{n_{i}}t_{jk}^{2}}\leq\sqrt{\sum_{j=1}^{m_{i}}\left(\sum_{t=1}^{n_{i}}t_{jk}\right)^{2}}=\|T\mathbf{1}_{n_{i}}\|=\left\|\mathbf{a}_{i}\right\|$
$\displaystyle=\sqrt{\sum_{j=1}^{m_{i}}a_{i,j}^{2}}\leq\sqrt{m_{i}\|\mathbf{a}_{i}\|_{\infty}^{2}}=\sqrt{m_{i}}\|\mathbf{a}_{i}\|_{\infty}.$

Similarly, we have $\|T\|\leq\sqrt{n_{i}}\|\mathbf{b}_{i}\|_{\infty}$. Then by
triangle inequality and the definition of $d_{i}$ in Theorem 1, we complete
the proof. ∎

Based upon Assumption 1 and Lemma 2, the residual function $R_{i}$ in equation
(4.1) enjoys a Lipschitz-like property, which will be useful for the
convergence proof.

######  Lemma 3.

Let $X^{\prime}$,
$X^{\prime\prime}\in\bigtimes_{i=1}^{N}\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})$.
Suppose that Assumption 1 holds and $X_{i}^{\prime}=X_{i}^{\prime\prime}$ for
some $i\in\\{1,\ldots,N\\}$. Then
$\left|R_{i}(X^{\prime})-R_{i}(X^{\prime\prime})\right|\leq
2d_{i}L\|X^{\prime}-X^{\prime\prime}\|$.

######  Proof.

Let $\bar{X}_{i}^{\prime}$ and $\bar{X}_{i}^{\prime\prime}$ be the optimal
solutions of

$\min_{X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})}~{}\left\langle\nabla_{i}f(X^{\prime}),X_{i}\right\rangle\quad\text{and}\quad\min_{X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i})}~{}\left\langle\nabla_{i}f(X^{\prime\prime}),X_{i}\right\rangle,$

respectively. Simple algebraic calculations yield

$\displaystyle
R_{i}(X^{\prime})=\left\langle\nabla_{i}f(X^{\prime}),X_{i}^{\prime}-\bar{X}_{i}^{\prime}\right\rangle~{}(\text{equation
\eqref{eqn:residual function}})$
$\displaystyle=\left\langle\nabla_{i}f(X^{\prime\prime}),X_{i}^{\prime}-\bar{X}_{i}^{\prime}\right\rangle+\left\langle\nabla_{i}f(X^{\prime})-\nabla_{i}f(X^{\prime\prime}),X_{i}^{\prime}-\bar{X}_{i}^{\prime}\right\rangle$
$\displaystyle\leq\left\langle\nabla_{i}f(X^{\prime\prime}),X_{i}^{\prime}-\bar{X}_{i}^{\prime}\right\rangle+2d_{i}L\|X^{\prime}-X^{\prime\prime}\|~{}(\text{Assumption
\ref{assume:Lip grad} and Lemma \ref{lemma:diameter}})$
$\displaystyle=\left\langle\nabla_{i}f(X^{\prime\prime}),X_{i}^{\prime\prime}-\bar{X}_{i}^{\prime}\right\rangle+2d_{i}L\|X^{\prime}-X^{\prime\prime}\|~{}(\text{since}~{}X_{i}^{\prime}=X_{i}^{\prime\prime})$
$\displaystyle\leq\left\langle\nabla_{i}f(X^{\prime\prime}),X_{i}^{\prime\prime}-\bar{X}_{i}^{\prime\prime}\right\rangle+2d_{i}L\|X^{\prime}-X^{\prime\prime}\|~{}(\text{the
definition of }\bar{X}_{i}^{\prime\prime})$
$\displaystyle=R_{i}(X^{\prime\prime})+2d_{i}L\|X^{\prime}-X^{\prime\prime}\|.~{}(\text{equation}~{}\eqref{eqn:residual
function})$

Similarly, one can show $R_{i}(X^{\prime\prime})\leq
R_{i}(X^{\prime})+2d_{i}L\|X^{\prime}-X^{\prime\prime}\|$. These two together
give the desired result. ∎

With the above tools in place, we are ready to establish an upper bound for
the average residual over the iterate sequence generated by the eralm method.

###### Proof of Theorem 1.

For any $i\in\\{1,\ldots,N\\}$ and $t\geq 0$,

$\displaystyle f(X_{\leq i}^{(t+1)},X_{>i}^{(t)})\leq f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})+\left\langle
C_{i}^{(t)},X_{i}^{(t+1)}-X_{i}^{(t)}\right\rangle+\frac{L}{2}\|X_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}$
$\displaystyle=$ $\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})+\alpha\left\langle
C_{i}^{(t)},\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\right\rangle+\frac{\alpha^{2}L}{2}\|\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}$
$\displaystyle\leq$ $\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})+\alpha\left\langle
C_{i}^{(t)},\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\right\rangle+2(d_{i}\alpha)^{2}L$
$\displaystyle\leq$ $\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})-\alpha\lambda h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})-\alpha
R_{i}(X_{<i}^{(t+1)},X_{\geq i}^{(t)})+2(d_{i}\alpha)^{2}L,$

where the first inequality uses Assumption 1, the second one uses Lemma 2, the
last one relies on Lemma 1 and the definition (4.1) of the residual function
$R_{i}$. The above relation further gives

$\alpha R_{i}(X_{<i}^{(t+1)},X_{\geq i}^{(t)})\leq f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})-f(X_{\leq i}^{(t+1)},X_{>i}^{(t)})-\alpha\lambda
h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})+2(d_{i}\alpha)^{2}L.$ (A.1)

Note that

$\displaystyle\left|R_{i}(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})-{R_{i}(X^{(t)})}\right|^{2}$ $\displaystyle\leq$
$\displaystyle~{}4(d_{i}L)^{2}\|X^{(t+1)}-X^{(t)}\|^{2}=4(d_{i}L)^{2}\sum_{i=1}^{N}\|X_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}$
$\displaystyle=$
$\displaystyle~{}4(d_{i}L\alpha)^{2}\sum_{i=1}^{N}\|\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}\leq
16\bar{d}^{4}L^{2}\alpha^{2}N,$

where the first inequality follows from Lemma 3 and the last one is due to
Lemma 2. Combining the above inequality and relation (A.1) yields

$\displaystyle~{}\alpha
R_{i}(X^{(t)})=\alpha\left[R_{i}(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})+\left({R_{i}(X^{(t)})-R_{i}(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})}\right)\right]$ $\displaystyle\leq$
$\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq i}^{(t)})-f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})-\alpha\lambda
h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})+4\bar{d}^{2}\alpha^{2}L\sqrt{N}+2(d_{i}\alpha)^{2}L$
$\displaystyle\leq$ $\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})-f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})+\alpha\lambda\bar{h}+2\bar{d}^{2}\alpha^{2}L(2\sqrt{N}+1).$

Summing the above inequality over $i$ from $1$ to $N$ and dividing both sides
by $\alpha$, one obtains

$R(X^{(t)})\leq\frac{f(X^{(t)})-f(X^{(t+1)})}{\alpha}+N\lambda\bar{h}+2\bar{d}^{2}LN(2\sqrt{N}+1)\alpha.$

Summing the above relation over $t$ from 0 to $t_{\max}-1$ and dividing the
both sides by $t_{\max}$, we complete the proof after noting
$f(X^{(t_{\max})})\geq\underline{f}$, the definition of $\alpha$ in equation
(4.2), and $2N(2\sqrt{N}+1)<(2N+1)^{2}$. Incidentally, the lower bound for
$t_{\max}$ in equation (4.2) ensures $\alpha\leq 1$. ∎

###### Proof of Corollary 1.

By Assumption 2 (i)-(iii), it is not hard to derive that

$\displaystyle\bar{d}$
$\displaystyle=\max_{i=1}^{N}\min\\{\sqrt{m_{i}}\left\|\mathbf{a}_{i}\right\|_{\infty},\sqrt{n_{i}}\left\|\mathbf{b}_{i}\right\|_{\infty}\\}=\Theta\left(\max_{i=1}^{N}\min\left\\{\frac{1}{\sqrt{m_{i}}},\frac{1}{\sqrt{n_{i}}}\right\\}\right),$
$\displaystyle\bar{h}$
$\displaystyle=\max_{i=1}^{N}\sum_{j,k}a_{i,j}b_{i,k}(1-\log
a_{i,j}b_{i,k})=\Theta\left(\sum_{i=1}^{N}\log m_{i}n_{i}\right).$

Based on the above bounds, the first term on the right-hand side of inequality
(4.3) goes to $0$ if $t_{\max}=\Omega(\sum_{i=1}^{N}(m_{i}+n_{i})^{\eta})$
with $\eta>\theta$ and $M$ independent from $\\{m_{i}\\}_{i=1}^{N}$ and
$\\{n_{i}\\}_{i=1}^{N}$, and the second one goes to $0$ if
$\lambda=o(1/\sum_{i=1}^{N}\log m_{i}n_{i})$. The proof is complete. ∎

##  Appendix B Convergence of the S-eralm method

We define the following auxiliary sequences:

$\displaystyle\bar{X}_{i}^{(t+1)}$
$\displaystyle\in\operatorname*{arg\,min}_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}\right\rangle,~{}\mathrm{s.t.}~{}X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}),~{}i=1,\ldots,N,~{}t\geq
0,$ $\displaystyle\breve{X}_{i}^{(t+1)}$
$\displaystyle=\operatorname*{arg\,min}_{X_{i}}~{}\left\langle
C_{i}^{(t)},X_{i}\right\rangle+\hat{\lambda}h(X_{i}),~{}\mathrm{s.t.}~{}X_{i}\in\mathcal{U}(\mathbf{a}_{i},\mathbf{b}_{i}),~{}i=1,\ldots,N,~{}t\geq
0.$

###### Proof of Theorem 2.

For any $i\in\\{1,\ldots,N\\}$ and $t\geq 0$,

$\displaystyle~{}\left(f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})-f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})\right)/\alpha-2d_{i}^{2}L\alpha$ (B.1) $\displaystyle\leq$
$\displaystyle~{}\left(f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})-f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})\right)/\alpha-L\alpha\|\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}/2$
$\displaystyle=$ $\displaystyle~{}\left(f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})-f(X_{<i}^{(t+1)},X_{\geq
i}^{(t)})-L\|X_{i}^{(t+1)}-X_{i}^{(t)}\|^{2}/2\right)/\alpha$
$\displaystyle\leq$ $\displaystyle~{}\left\langle
C_{i}^{(t)},X_{i}^{(t+1)}-X_{i}^{(t)}\right\rangle/\alpha=\left\langle
C_{i}^{(t)},\tilde{X}_{i}^{(t+1)}-X_{i}^{(t)}\right\rangle$ $\displaystyle=$
$\displaystyle~{}\left\langle\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle+\left\langle
C_{i}^{(t)}-\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle-\left\langle
C_{i}^{(t)},X_{i}^{(t)}\right\rangle$ $\displaystyle=$
$\displaystyle~{}\underbrace{\left\langle\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle-\left\langle
C_{i}^{(t)},\bar{X}_{i}^{(t+1)}\right\rangle}_{I_{1}}+\underbrace{\left\langle
C_{i}^{(t)}-\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle}_{I_{2}}-R_{i}(X_{<i}^{(t+1)},X_{\geq
i}^{(t)}),$

where the first inequality follows from Lemma 2, the second one uses
Assumption 1, and the last equality uses the definition (4.1) of the residual
function $R_{i}$.

Next, we seek to bound $I_{1}$ and $I_{2}$ in relation (B.1). For the latter,

$\displaystyle\left\langle
C_{i}^{(t)}-\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle\leq
d_{i}\|\hat{C}_{i}^{(t)}-C_{i}^{(t)}\|=d_{i}\hat{\lambda}\sqrt{\sum_{j,k:(j,k)\in\mathcal{I}_{i}^{(t)}}\log^{2}\big{(}n_{s,i}\cdot
p_{i,jk}^{(t)}\big{)}}$ $\displaystyle\quad\leq
d_{i}\hat{\lambda}\sqrt{|\mathcal{I}_{i}^{(t)}|\log^{2}\frac{1}{(1-\gamma)w_{i}n_{s,i}}}=d_{i}\hat{\lambda}\sqrt{|\mathcal{I}_{i}^{(t)}|}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}},$

where the first inequality comes from the proof of Lemma 2 and the second one
uses the formula (3.8) and Assumption 3 (ii), the first equality follows from
equation (3.12) and Assumption 3 (ii). Following Hoeffding’s inequality, for
any $\iota>0$,

$\mathsf{P}\left(|\mathcal{I}_{i}^{(t)}|\geq n_{s,i}+\iota\cdot
m_{i}n_{i}\right)\leq\exp\left(-2\iota^{2}m_{i}n_{i}\right).$

Therefore, with probability no less than $1-\exp(-2\iota^{2}m_{i}n_{i})$, we
have

$\left\langle
C_{i}^{(t)}-\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle\leq
d_{i}\hat{\lambda}\sqrt{n_{s,i}+\iota\cdot
m_{i}n_{i}}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}}.$ (B.2)

Regarding the former term $I_{1}$,

$\displaystyle\left\langle\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle-\left\langle
C_{i}^{(t)},\bar{X}_{i}^{(t+1)}\right\rangle$ (B.3) $\displaystyle\leq$
$\displaystyle~{}\left\langle\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle-\left\langle
C_{i}^{(t)},\breve{X}_{i}^{(t+1)}\right\rangle-\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})$
$\displaystyle=$
$\displaystyle~{}\left[\left\langle\hat{C}_{i}^{(t)},\tilde{X}_{i}^{(t+1)}\right\rangle+\hat{\lambda}h(\tilde{X}_{i}^{(t+1)})\right]-\left[\left\langle
C_{i}^{(t)},\breve{X}_{i}^{(t+1)}\right\rangle+\hat{\lambda}h(\breve{X}_{i}^{(t+1)})\right]$
$\displaystyle~{}+\hat{\lambda}\left[h(\breve{X}_{i}^{(t+1)})-h(\tilde{X}_{i}^{(t+1)})\right]-\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})$
$\displaystyle\leq$
$\displaystyle~{}q_{i}(\tilde{\mathbf{u}}_{i}^{(t,\star)},\tilde{\mathbf{v}}_{i}^{(t,\star)};\hat{\lambda},\hat{\Psi}_{i}^{(t)})-q_{i}(\tilde{\mathbf{u}}_{i}^{(t,\star)},\tilde{\mathbf{v}}_{i}^{(t,\star)};\hat{\lambda},\Psi_{i}^{(t)})-2\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})$
$\displaystyle\leq$
$\displaystyle~{}\hat{c}_{2}\hat{\lambda}\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}\left|1-\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}\right|^{-1}-2\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top}),$

where the first inequality uses Lemma 1 (with
$\mathcal{I}=\mathcal{I}_{i}^{(t)}$), the second one leverages
$h(\breve{X}_{i}^{(t+1)})\leq 0$, $h(\tilde{X}_{i}^{(t+1)})\geq
h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})$, and the definition (3.4) together
with the strong duality and the optimality of
$(\tilde{\mathbf{u}}_{i}^{(t,\star)},\tilde{\mathbf{v}}_{i}^{(t,\star)})$ in
problem (3.10), the last one follows from [li2023efficient]. To bound the
first term on the right-hand side of the inequality (B.3), we resort to
[achlioptas2007fast, Theorem 3.1]. Specifically, for any
$j\in\\{1,\ldots,m_{i}\\}$, $k\in\\{1,\ldots,n_{i}\\}$, it holds by equation
(3.8) and Assumption 3 that

$\mathsf{E}\left(\frac{\hat{\psi}_{i,jk}^{(t)}}{\|\Psi_{i}^{(t)}\|_{2}}\right)=n_{s,i}\cdot
p_{i,jk}^{(t)}\cdot\frac{\psi_{i,jk}^{(t)}}{n_{s,i}\cdot
p_{i,jk}^{(t)}}\cdot\frac{1}{\|\Psi_{i}^{(t)}\|_{2}}=\frac{\psi_{i,jk}^{(t)}}{\|\Psi_{i}^{(t)}\|_{2}},$
$\displaystyle\mathsf{Var}\left(\frac{\hat{\psi}_{i,jk}^{(t)}}{\|\Psi_{i}^{(t)}\|_{2}}\right)$
$\displaystyle<\frac{(\psi_{i,jk}^{(t)})^{2}}{n_{s,i}\cdot
p_{i,jk}^{(t)}\|\Psi_{i}^{(t)}\|_{2}^{2}}\leq\frac{1}{(1-\gamma)w_{i}n_{s,i}\|\Psi_{i}^{(t)}\|_{2}^{2}}$
$\displaystyle\leq\frac{c_{1}^{2}(m_{i}+n_{i})^{-2\nu}}{(1-\gamma)w_{i}n_{s,i}}\leq\frac{c_{1}^{2}\log^{4}(1+\varepsilon)}{8(m_{i}+n_{i})\log^{4}(m_{i}+n_{i})}.$

In addition, $\hat{\psi}_{i,jk}^{(t)}/\|\Psi_{i}^{(t)}\|_{2}$ lies in an
interval with length no larger than

$\displaystyle\frac{\psi_{i,jk}^{(t)}}{n_{s,i}\cdot
p_{i,jk}^{(t)}\|\Psi_{i}^{(t)}\|_{2}}$
$\displaystyle\leq\frac{1}{(1-\gamma)w_{i}n_{s,i}\|\Psi_{i}^{(t)}\|_{2}}\leq\frac{c_{1}(m_{i}+n_{i})^{-\nu}}{(1-\gamma)w_{i}n_{s,i}}$
$\displaystyle\leq\frac{c_{1}\log^{4}(1+\varepsilon)}{8(m_{i}+n_{i})^{1-\nu}\log^{4}(m_{i}+n_{i})}.$

Then, by Theorem 3.1 in [achlioptas2007fast],

$\mathsf{P}\left(\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}\geq
c_{1}(1+\varepsilon+\zeta)\frac{\log^{2}(1+\varepsilon)}{\log^{2}(m_{i}+n_{i})}\right)<2\exp\left(-\frac{16\zeta^{2}}{\varepsilon^{4}}\log^{4}(m_{i}+n_{i})\right)$

holds for any $\zeta>0$ and $m_{i}+n_{i}\geq 152$. That is to say, with
probability no less than
$1-2\exp\big{(}-16\zeta^{2}\log^{4}(m_{i}+n_{i})/\varepsilon^{4}\big{)}$,

$\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}<c_{1}(1+\varepsilon+\zeta)\frac{\log^{2}(1+\varepsilon)}{\log^{2}(m_{i}+n_{i})},$

and

$\displaystyle\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}\left|1-\frac{\|\hat{\Psi}_{i}^{(t)}-\Psi_{i}^{(t)}\|_{2}}{\|\Psi_{i}^{(t)}\|_{2}}\right|^{-1}$
(B.4) $\displaystyle<$
$\displaystyle~{}\frac{c_{1}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)}{\log^{2}(m_{i}+n_{i})-c_{1}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)}.$

Combining relations (B.1) to (B.4), we have, with probability no less than

$\left[1-2\exp\left(-\frac{16\zeta^{2}}{\varepsilon^{4}}\log^{4}(m_{i}+n_{i})\right)\right]\left[1-\exp\left(-2\iota^{2}m_{i}n_{i}\right)\right],$

that

$\displaystyle f(X_{\leq i}^{(t+1)},X_{>i}^{(t)})\leq$
$\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq i}^{(t)})-\alpha
R_{i}(X_{<i}^{(t+1)},X_{\geq i}^{(t)})+2(d_{i}\alpha)^{2}L$
$\displaystyle~{}-2\alpha\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})+d_{i}\alpha\hat{\lambda}\sqrt{n_{s,i}+\iota\cdot
m_{i}n_{i}}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}}$
$\displaystyle~{}+\alpha\hat{\lambda}\frac{c_{1}\hat{c}_{2}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)}{\log^{2}(m_{i}+n_{i})-c_{1}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)},$

which implies

$\displaystyle\alpha R_{i}(X_{<i}^{(t+1)},X_{\geq i}^{(t)})\leq$
$\displaystyle~{}f(X_{<i}^{(t+1)},X_{\geq i}^{(t)})-f(X_{\leq
i}^{(t+1)},X_{>i}^{(t)})+2(d_{i}\alpha)^{2}L$
$\displaystyle~{}-2\alpha\hat{\lambda}h(\mathbf{a}_{i}\mathbf{b}_{i}^{\top})+d_{i}\alpha\hat{\lambda}\sqrt{n_{s,i}+\iota\cdot
m_{i}n_{i}}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}}$
$\displaystyle~{}+\alpha\hat{\lambda}\frac{c_{1}\hat{c}_{2}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)}{\log^{2}(m_{i}+n_{i})-c_{1}(1+\varepsilon+\zeta)\log^{2}(1+\varepsilon)}.$

Following similar arguments for proving Theorem 1 and noticing the definition
of $c_{3}$, one obtains the desired result. ∎

###### Proof of Corollary 2.

Assumption 2 implies bounds for $\bar{d}$, $\bar{h}$ (see the proof of
Corollary 1 in Appendix A) and also $w_{i}=\Theta(1/m_{i}n_{i})$
($i=1,\ldots,N$).

The first term on the right-hand side of inequality (4.4) tends to $0$ if
$t_{\max}=\Theta(\sum_{i=1}^{N}(m_{i}+n_{i})^{\eta})$ with $\eta>\theta$ and
$M$ independent from $\\{m_{i}\\}_{i=1}^{N}$ and $\\{n_{i}\\}_{i=1}^{N}$. For
a fixed $\iota>0$, since

$\displaystyle~{}\sum_{i=1}^{N}\sqrt{n_{s,i}+\iota\cdot
m_{i}n_{i}}\log\frac{1}{(1-\gamma)w_{i}n_{s,i}}$ $\displaystyle=$
$\displaystyle~{}\mathcal{O}\left(\sum_{i=1}^{N}\left[\frac{\sqrt{m_{i}n_{i}}}{(m_{i}+n_{i})^{\nu-1/2}}\log^{2}(m_{i}+n_{i})+\sqrt{m_{i}n_{i}}\right]\log\frac{(m_{i}+n_{i})^{2\nu-1}}{\log^{4}(m_{i}+n_{i})}\right),$

with the choices of $\\{n_{s,i}\\}_{i=1}^{N}$, the second and third terms tend
to $0$ if

$\hat{\lambda}=o\left(\frac{1}{\sum_{i=1}^{N}\sqrt{m_{i}n_{i}}\log(m_{i}+n_{i})}\right)$

and $\varepsilon$, $\nu$, $\gamma$ are independent from
$\\{m_{i}\\}_{i=1}^{N}$ and $\\{n_{i}\\}_{i=1}^{N}$. Incidentally, the choices
of $\\{n_{s,i}\\}_{i=1}^{N}$ do not conflict with Assumption 3 (ii) by virtue
of Remark 3. Since $c_{1}$, $c_{2}$, and $\hat{c}_{2}$ are also independent
from $\\{m_{i}\\}_{i=1}^{N}$ and $\\{n_{i}\\}_{i=1}^{N}$, the last term
vanishes as $\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$. Finally, the probability
is not less than

$\displaystyle~{}\prod_{i=1}^{N}\left\\{\left[1-2\exp\left(-\frac{16\zeta^{2}}{\varepsilon^{4}}\log^{4}(m_{i}+n_{i})\right)\right]\left[1-\exp\left(-2\iota^{2}m_{i}n_{i}\right)\right]\right\\}^{t_{\max}}$
$\displaystyle\geq$
$\displaystyle~{}\prod_{i=1}^{N}\left\\{\left[1-\frac{2}{(m_{i}+n_{i})^{16\zeta^{2}/\varepsilon^{4}}}\right]\left[1-\exp\left(-2\iota^{2}m_{i}n_{i}\right)\right]\right\\}^{t_{\max}},$

which, by Assumption 2 (iv), goes to $1$ as
$\sum_{i=1}^{N}(m_{i}+n_{i})\to+\infty$ after choosing $\zeta>0$ such that
$\zeta>\sqrt{\eta}\varepsilon^{2}/4$. The proof is completed. ∎

## References