# A Challenge in Reweighting Data
with Bilevel Optimization

Anastasia Ivanova1 work done while at Apple Pierre Ablin2

( 1 Univ. Grenoble Alpes
2 Apple )

###### Abstract

In many scenarios, one uses a large training set to train a model with the
goal of performing well on a smaller testing set with a different
distribution. Learning a weight for each data point of the training set is an
appealing solution, as it ideally allows one to automatically learn the
importance of each training point for generalization on the testing set. This
task is usually formalized as a bilevel optimization problem. Classical
bilevel solvers are based on a warm-start strategy where both the parameters
of the models and the data weights are learned at the same time. We show that
this joint dynamic may lead to sub-optimal solutions, for which the final data
weights are very sparse. This finding illustrates the difficulty of data
reweighting and offers a clue as to why this method is rarely used in
practice.

##  1 Introduction

In many practical learning scenarios, there is a discrepancy between the
training and testing distribution. For instance, when training large language
models, we may have access to a training set that contains many low-quality
data points from different sources and want to train a model on this dataset
to perform well on a testing set that contains a few high-quality points [7,
3, 18]. An appealing way to solve this problem is _data reweighting_ [21, 23,
26], where one attributes one weight to each data point in the training set.
The weight of a training sample should reflect how much this sample resembles
the testing set and helps the model perform well on it. Figure 1 illustrates
the general principle.

Learning the optimal weights can be cast as a bilevel optimization problem
[9], where the optimal weights are such that training the model with these
weights leads to the smallest test loss possible. The weights are usually
constrained to sum to one, leading to an optimization problem on the simplex,
which is usually solved with mirror descent [19]. Despite its promise of
automatically learning the importance of data points, data reweighting is
still seldom used in practice. In this paper, we try to provide a possible
explanation for this lack of adoption, showing, in short, that the underlying
optimization problem is hard.

In a large-scale setting where fitting the model once is expensive, it is
prohibitively costly to iteratively update the weights with a model fit at
each iteration [20]. Hence, practitioners often resort to _warm-started_
bilevel optimization, where the parameters of the model and the weights evolve
simultaneously [17, 14, 6].

Figure 1: Data reweighting principle: test and train distributions are
different. We aim to estimate a weight for each training sample that reflects
its contribution to the model’s performance on the test set. For instance, if
the test set only contains dog images, we want a large weight on dog images in
the training set, a small weight on other animals which might help the model,
and zero weight on irrelevant training images. These weights should be learned
automatically during training.

This paper aims to provide a detailed analysis of the corresponding joint
dynamics. Our main result indicates that warm starting with mirror descent
leads to _sparse weights_ : after training, only a few training points have a
non-zero weight, which is detrimental to the generalization power of the
corresponding parameters. Our results show a weakness of warm-started bilevel
approaches for this problem and are a first step toward explaining the
hardness of data-reweighting.

Contributions and paper organization In Sec. 2, we introduce data reweighting
as a bilevel problem and explain how warm-started bilevel aims at solving the
problem. In Sec. 3, we study the corresponding dynamics. We focus on two
settings. When the parameters are updated at a much _greater_ pace than the
weights, we formalize the intuition that this recovers the standard, non-warm-
started, bilevel approach, which leads to satisfying solutions but takes a
long time to converge. In the opposite setting, where the parameters are
updated at a much _slower_ pace than the weights, where we show that this
leads to extremely sparse weights, which in turn hinders the generalization of
the model since the model is effectively trained with few samples. Finally,
Sec. 4 gives numerical results illustrating the theory presented in the paper.

Notation: The vector containing all $1$’s is
$\mathbb{1}_{k}\in\mathbb{R}^{k}$. The simplex is
$\Delta_{n}=\\{w\in\mathbb{R}^{n}_{+}|\sum_{i=1}^{n}w_{i}=1\\}$. The
multiplication of two vectors $u,v\in\mathbb{R}^{n}$ is $u\odot
v\in\mathbb{R}^{n}$ of entries $(u_{i}v_{i})$. The set $\\{1,n\\}$ is the set
of integers $1$ to $n$. The support $\mathrm{Supp}(w)$ of a vector
$w\in\mathbb{R}^{n}$ is the set of indices $i$ in $\\{1,n\\}$ such that
$w_{i}\neq 0$. Given a set of indices $S=(s_{1},\dots,s_{l})\subset\\{1,n\\}$
of size $s$, the restriction to $S$ of a vector $w\in\mathbb{R}^{n}$ is the
vector $w\mid_{S}\in\mathbb{R}^{l}$ of entries $w_{s_{j}}$ for
$j\in\\{1,l\\}$. The gradient (resp. Hessian) of a loss $\ell(\theta;x)$ is
$\nabla_{\theta}\ell(\theta;x)$ (resp.
$\nabla^{2}_{\theta\theta}\ell(\theta;x)$). The $\mathrm{range}$ of a matrix
is the span of its column.

##  2 Data reweighting as a bilevel problem

We consider a train dataset $X=[x_{1},\dots,x_{n}]$, and a testing dataset
$X^{\prime}=[x^{\prime}_{1},\dots,x^{\prime}_{m}]$. Our goal is to train a
machine learning model on the train set that has good performance on the
testing set. Letting $\theta\in\mathbb{R}^{p}$ the parameters of the model, we
let $\ell(\theta;x)$ the loss corresponding to the train set, and
$\ell^{\prime}(\theta;x^{\prime})$ the loss of the test set. The classical
_empirical risk minimization_ cost function is
$G(\theta)=\frac{1}{n}\sum_{i=1}^{n}\ell(\theta;x_{i})$ for the train set and
$F(\theta):=\frac{1}{m}\sum_{j=1}^{m}\ell^{\prime}(\theta;x_{j}^{\prime})$ for
the test set, where each sample has the same weight.

The goal of data reweighting is to give a different weight from $1/n$ to each
training data in order to get a solution $\theta$ that leads to a good
performance on the test set, i.e., leads to a small $F(\theta)$. To this end,
we introduce

$G(\theta,w):=\sum_{i=1}^{n}w_{i}\ell(\theta;x_{i})$ (1)

defined for $w$ belonging to the _simplex_
$\Delta_{n}=\\{w\in\mathbb{R}_{+}^{n}|\enspace\sum_{i=1}^{n}w_{i}=1\\}$.
Ideally, we would want $w_{i}$ to be large when the training sample $x_{i}$
helps the model’s performance on the testing set, and conversely, $w_{i}$
should be small when $x_{i}$ does not help the model on the testing set. We
make the following blanket assumption which makes the bilevel problem well-
defined:

######  Assumption 1 (Strong convexity).

The loss functions $\ell,\ell^{\prime}$ are differentiable. Additionally, the
function $\theta\mapsto\ell(\theta;x)$ is $\mu-$ strongly convex with $\mu>0$
for all $x$, i.e. for all $x,\theta$, we have
$\lambda_{\min}(\nabla^{2}_{\theta\theta}\ell(\theta;x))\geq\mu$.

For instance, such assumption is verified if $\ell(\theta;x)$ is of the form
$\mathrm{fit}(\theta;x)+\mu\|\theta\|^{2}/2$, where the $\mathrm{fit}$
function is a data-fit term that is convex (like a least-squares or logistic
loss) and $\mu\|\theta\|^{2}/2$ is a regularizer.

Changing the weights $w$ modifies the cost function $G(\theta,w)$, hence its
minimizer is now a function of $w$, denoted
$\theta^{*}(w)=\arg\min_{\theta}G(\theta,w)$. Note that the strong-convexity
assumption 1 implies that the function $G$ itself is $\mu-$strongly convex,
guaranteeing the existence and uniqueness of $\theta^{*}(w)$. Data reweighting
is formalized as the _bilevel problem_

$\min_{w\in\Delta_{n}}h(w):=F(\theta^{*}(w)),$ (2)

where $\theta^{*}$ depends implicitly on $w$ through $G$.

###  2.1 Importance sampling as the hope behind data reweighting?

One can take a distributional point of view on the problem to get better
insights into the sought-after solutions. We let
$\mu=\sum_{i=1}^{n}\delta_{x_{i}}$ the empirical training distribution, and
$\nu=\sum_{j=1}^{m}\delta_{x^{\prime}_{j}}$ the empirical test distribution.
Instead of having one weight per data sample, we now have a weighting function
$\omega:\mathbb{R}^{d}\to\mathbb{R}_{+}$ such that $\int\omega(x)d\mu(x)=1$.
The bilevel problem becomes

$\displaystyle\begin{split}\min_{\omega}\int\ell^{\prime}(\theta^{*}(\omega);x)d\nu(x)\text{
subject to}\\\
\theta^{*}(\omega)\in\arg\min_{\theta}\int\ell(\theta;x)\omega(x)d\mu(x)\end{split}$
(3)

In practice, the distributions $\mu$ and $\nu$ are sums of Diracs; we can
wonder what happens instead if they are continuous.

######  Proposition 1.

If $\nu$ is absolutely continuous w.r.t. $\mu$, and $\ell=\ell^{\prime}$, a
global solution to the bilevel problem (3) is
$\omega^{*}(x)=\frac{d\nu}{d\mu}(x)$.

This ratio $d\nu/d\mu$ is precisely the one that _importance sampling_
techniques [24] try to estimate: in this specific case, bilevel optimization
recovers importance sampling. We now turn to the resolution of the bilevel
problem.

###  2.2 Solving the bilevel problem

The bilevel optimization problem corresponding to data-reweighting is the
_single-level_ optimization of the non-convex function $h$ over the simplex
$\Delta_{n}$ (2), for which _mirror descent_ [19, 2] is an algorithm of
choice. Starting from an initial guess $w^{0}\in\Delta_{n}$, mirror descent
iterates $\tilde{w}^{k+1}=w^{k}\odot\exp(-\eta\nabla h(w^{k}))$ and
$w^{k+1}=\frac{\tilde{w}^{k+1}}{\|\tilde{w}^{k+1}\|_{1}}$, where $\odot$ is
the element-wise multiplication.

This method involves the gradient of the value function $h$, which is obtained
with the chain rule [22, 8]

$\nabla h(w)=\left[\frac{\partial\theta^{*}}{\partial w}\right]^{T}\nabla
F(\theta^{*}(w)).$ (4)

The implicit function theorem then gives, thanks to the invertibility of
$\nabla^{2}_{\theta\theta}$ guaranteed by Assumption 1:

$\frac{\partial\theta^{*}}{\partial
w}=-\left[\nabla_{\theta\theta}^{2}G(\theta^{*}(w),w)\right]^{-1}\nabla_{\theta
w}^{2}G(\theta^{*}(w),w).$ (5)

In the data reweighting problem, $G$ has a special structure, which gives
$\nabla_{\theta\theta}^{2}G(\theta,w)=\sum_{i=1}^{n}w_{i}\nabla^{2}\ell(\theta;x_{i})$
and $\nabla_{\theta
w}^{2}G(\theta,w)=[\nabla\ell(\theta,x_{1}),\dots,\nabla\ell(\theta,x_{n})]\in\mathbb{R}^{p\times
n}$. We finally obtain $\nabla h(w)=\Psi(\theta^{*}(w),w)\in\mathbb{R}^{n}$
where the $i-$th coordinate of $\Psi$ is given by

$\Psi(\theta,w)_{i}=-\left\langle\nabla\ell(\theta;x_{i}),\left[\nabla_{\theta}^{2}G(\theta,w)\right]^{-1}\nabla
F(\theta)\right\rangle.$ (6)

This hyper-gradient has an intuitive structure: letting $\langle
u,v\rangle_{\theta}$ the scalar product defined over $\mathbb{R}^{p}$ by
$\langle u,v\rangle_{\theta}=\langle
u,\left[\nabla_{\theta\theta}^{2}G(\theta,w)\right]^{-1}v\rangle$, the hyper-
gradient corresponding to sample $i$ is simply the (opposite) of the alignment
measured with this new scalar product between the gradient of the sample $i$,
and the gradient of the outer function $F$. Therefore, this gradient increases
weights for which $\nabla\ell(\theta;x_{i})$ aligns with the outer gradient.

Algorithm 1 Exact Bilevel Algorithm

Input : Initial point $w^{0}\in\Delta$, step-size $\eta$, number of iterations
$N$.

for $k=1$ to $N$ do

Compute $\theta^{*}(w^{k})$

Compute $\nabla h(w^{k})=\Psi(\theta^{*}(w^{k}),w^{k})$

Update $w^{k+1}=\dfrac{w^{k}\odot\exp(-\eta\nabla
h(w^{k}))}{\|w^{k}\odot\exp(-\eta\nabla h(w^{k}))\|_{1}}$

end for

Return : $w^{N}$.

The mirror descent algorithm to solve the bilevel problem (2) is described in
Algorithm 1.

Unfortunately, this algorithm is purely theoretical as it is not
implementable; indeed, computing $\theta^{*}(w^{k})$ is equivalent to finding
the minimum of $G$, which is generally impossible. We can instead try to
approximate it by replacing $\theta^{*}(w^{k})$ by the output of many
iterations of an optimization algorithm, but then the cost of the algorithm
becomes prohibitive: each iteration requires the approximate resolution of an
optimization problem.

###  2.3 The practical solution: warm-started bilevel optimization

A sound idea for a scalable algorithm is to use an iterative method instead to
minimize $G$ and have $\theta$ and $w$ evolve simultaneously. In this case, we
have two sets of variables $\theta^{k},w^{k}$. The parameters $\theta^{k}$ are
updated using standard algorithms like (stochastic) gradient descent, and then
to update $w^{k}$, we approximate the hyper-gradient by $\nabla
h(w^{k})\simeq\Psi(\theta^{k},w^{k})$ in Eq. (6). For simplicity, we use
gradient descent with step-size $\rho$ to update $\theta^{k}$i, and each
iteration consists of one update of $\theta^{k}$ followed by one update of
$w^{k}$. The full procedure is described in Algorithm 2.

Algorithm 2 Warm Started Bilevel Algorithm

Input : Initial points $\theta^{0}\in\mathbb{R}^{p}$ and $w^{0}\in\Delta_{n}$,
step-sizes $\eta$ and $\rho$, number of iterations $N$.

for $k=1$ to $N$ do

Compute $\Psi(\theta^{k},w^{k})$

Update $\theta^{k+1}=\theta^{k}-\rho\nabla G(\theta^{k},w^{k})$

Update
$w^{k+1}=\dfrac{w^{k}\odot\exp(-\eta\Psi(\theta^{k},w^{k}))}{\|w^{k}\odot\exp(-\eta\Psi(\theta^{k},w^{k}))\|_{1}}$

end for

Return : $w^{N}$.

As a side note, this algorithm can still be expensive to implement because i)
computing $\Psi$ requires solving a linear system and ii) in a large-scale
setting when $n$ and $m$ are large, computing the inner and outer gradients
and the Hessian scales linearly with the number of samples. Several works try
to fix these issues by proposing bilevel algorithms that do not have to form a
Hessian or invert a system and that are stochastic, i.e., only use one sample
at each iteration to progress [10, 27, 11, 16]. For our theoretical analysis,
we do not consider such modifications and focus on the bare-bones case of
Algorithm 2, which slightly departs from practice, but is already insightful.

Figure 2: Data reweighting for a linear regression problem. The first two
figures show the problem setup. The train set consists of points from two
clusters. Points color corresponds to the value of target
$y=\langle\hat{\theta}_{i},d\rangle+\mathrm{noise}$, where each cluster has
its own linear model $\hat{\theta}_{1}$ and $\hat{\theta}_{2}$. The test set
consists of points only from the first cluster with
$\hat{\theta}=\hat{\theta}_{1}$. The last four figures show four reweighting
solutions. The size of the points in the picture corresponds to the weight
value $w_{i}$ for the corresponding sample. We give a small non-zero radius to
zero weights for visualization purposes. Uniform weights: same weights for all
data points regardless of the cluster. Since the full dataset is not separable
with a linear model, the linear model finds a compromise that leads to a poor
solution. Optimal weights: equal weights for points from the first cluster,
zero weights for the other points. The wrong cluster is fully discarded,
leading to optimal estimation of the parameters. Exact bilevel: output of
Algorithm 1, gets to a solution close to the optimal weights, and the cluster
is correctly identified with almost uniform weights. Warm started bilevel:
output of Algorithm 2, gets to a very sparse solution which does not
generalize well.

###  2.4 A toy experiment

We describe a toy experiment illustrating the practical issues of Algorithm 2.
We consider a 2-d regression setting, where the train data
$[x_{1},\dots,x_{n}]$ comes from a mixture of two Gaussian, and each cluster
has a different regression parameter. Formally, given two centroids
$\mu_{1},\mu_{2}\in\mathbb{R}^{2}$ and parameters
$\hat{\theta}_{1},\hat{\theta}_{2}\in\mathbb{R}^{2}$, each training sample
$x_{i}=[d_{i},y_{i}]$ is generated by

$\displaystyle z_{i}\sim\mathrm{Bern}(1,2),\enspace
d_{i}\sim\mathcal{N}(\mu_{z_{i}},I),y_{i}\sim\mathcal{N}(\langle
d_{i},\theta^{*}_{z_{i}}\rangle,\sigma^{2}),$

where $\mathrm{Bern}(1,2)$ is the Bernoulli law with probability $1/2$ over
$\\{1,2\\}$. Here, $z_{i}$ represents the random cluster associated to
$x_{i}$. Meanwhile, the test data $x_{j}^{\prime}$ is only drawn from the
first cluster with parameter $\hat{\theta}=\hat{\theta}_{1}$. The
corresponding points $d$ are plotted in the two leftmost figures in Figure 2
in a color reflecting the target’s value $y$. The model is linear with a
least-squares loss, meaning that the train and test loss are

$\ell(x,\theta)=\ell^{\prime}(x,\theta)=\frac{1}{2}\|\langle
d,\theta\rangle-y\|^{2}\text{ with }x=[d,y].$ (7)

The four rightmost figures in Figure 2 display four possible reweighting
solutions. Taking uniform weights $w_{i}=1/n$, the train loss is minimized by
finding a compromise between the two cluster’s parameters, and the model does
not recover $\hat{\theta}_{1}$ nor $\hat{\theta}_{2}$. This leads to a high
training and testing loss, and $\|\hat{\theta}-\theta^{*}\|$ is large.
Intuitively, in the data reweighting framework, we want the weights to focus
only on the first cluster and to discard points from the second cluster, which
is irrelevant to the test problem. Additionally, in a noiseless setting
($\sigma=0$), setting $w_{i}$ proportional to $\delta_{z_{i}=1}$ leads to a
global minimum in the bilevel problem since $\theta^{*}(w)=\hat{\theta}$ and
hence $h(w)=F(\hat{\theta})=0$. In this ideal scenario, the model is fit only
on the correct cluster, leading to a small error
$\|\hat{\theta}-\theta^{*}\|$, due only to the label noise.

We then apply Algorithm 1 to this problem. Here, since the inner problem is
quadratic, we have a closed-form
$\theta^{*}(w)=\left(\sum_{i=1}^{n}w_{i}d_{i}d_{i}^{T}\right)^{-1}\sum_{i=1}^{n}w_{i}y_{i}d_{i}$,
making Algorithm 1 implementable. It gives a reasonable solution where most of
the weights in the second cluster are close to $0$, and weights in the first
cluster are nearly uniform. Thus, the error $\|\hat{\theta}-\theta^{*}\|$ is
small.

We finish with the warm-started Algorithm 2. We observe that the algorithm
outputs sparse weights: only a few weights are non-zero. This, in turn, leads
to a higher estimation error $\|\hat{\theta}-\theta^{*}\|$. As an important
note, this effect is worsened when taking a higher learning rate $\eta$ for
the update of $w$, and is mitigated by taking small learning rates $\rho$ that
lead to slow convergence to better solutions. This illustrates a problem with
the warm-started method: it has a trajectory distinct from the exact bilevel
method and recovers sub-optimal solutions. In the following section, we
develop a mathematical framework that explains this phenomenon.

##  3 A Dynamical System View on the Warm-Started Bilevel Algorithm

The study of iterative methods like Algorithm 2 is notoriously hard; we focus
on the dynamical system obtained by letting the step sizes $\eta$, $\rho$ go
to $0$ at the same speed. Algorithm 2 can be seen as the discretization of the
Ordinary Differential Equation (ODE)

$\begin{cases}\dot{\theta}&=-\alpha\nabla G(\theta,w)\\\ \dot{w}&=-\beta
P(w)\Psi(\theta,w)\end{cases}\enspace,$ (8)

where $\alpha>0$ controls the speed of convergence of $\theta$, $\beta>0$
controls the speed of $w$, and $P(w)=\mathrm{diag}(w)-ww^{T}$ is a
preconditioning matrix, which is the inverse Hessian metric $\mathrm{diag}(w)$
associated to the entropy applied to the projector on the tangent space
$I_{n}-\mathbb{1}_{n}w^{T}$, recovering a Riemannian gradient flow [12]. This
ODE can be recovered from Algorithm 2 in the following sense:

######  Proposition 2.

Let $(\theta(t),w(t))$ the solution of the ODE (8), and $(\theta^{k},w^{k})$
the iterates of Algorithm 2. Assume that the steps are such that
$(\rho,\eta)=\tau(\alpha,\beta)$ with $\tau>0$. Then, for any $T>0$, we have
$\lim_{\tau\to 0}\sup_{t\in[0,T]}\|(\theta(t),w(t))-(\theta^{\lfloor
t/\tau\rfloor},w^{\lfloor t/\tau\rfloor})\|=0$.

When $\alpha\gg\beta$, the dynamics in $\theta$ are much faster, and we
recover the classical bilevel approach:

######  Theorem 1.

Let $w^{*}(t)$ the solution of the ODE $\dot{w}=-P(w)\nabla h(w)$, and
$\theta^{\alpha,\beta},w^{\alpha,\beta}$ the solution of (8). Then for all
$\alpha$, for all time horizon $T$, we have

$\lim_{\beta\to 0}\sup_{t\in[0,T]}\|w^{\alpha,\beta}(t/\beta)-w^{*}(t)\|=0$
$\lim_{\beta\to
0}\|\theta^{\alpha,\beta}(t)-\theta^{*}(w(t))\|\leq\|\theta_{0}-\theta^{*}(w_{0})\|e^{-\mu\alpha
t}$

The proof of this result uses the classical tools in bilevel literature [10].
This result highlights that, as expected, if $\alpha\gg\beta$, the variable
$\theta$ is tracking $\theta^{*}(w)$, which in turn means that $w$ follows the
direction of the true gradient of $h$: we recover the Exact Bilevel dynamics,
that corresponds to the gradient flow of $h$. The rest of this section is
devoted to understanding what happens in the other regime where
$\beta\gg\alpha$, when the dynamics in $w$ are much faster than that in
$\theta$.

###  3.1 Mirror descent flows on the simplex

Before understanding the joint dynamics, we consider the dynamics of the warm-
started bilevel ODE (8) when only $w$ evolves, i.e., when $\alpha=0$ and
$\theta(t)=\theta^{0}$. For a vector field
$\phi:\mathbb{R}^{n}\to\mathbb{R}^{n}$, we consider mirror descent updates
given by $\tilde{w}^{k+1}=w^{k}\odot\exp(-\eta\phi(w^{k}))$ and
$w^{k+1}=\frac{\tilde{w}^{k+1}}{\|\tilde{w}^{k+1}\|_{1}}$. Letting the step
size $\eta$ to $0$, we recover the _mirror descent flow_ [13], that is the ODE

$\dot{w}=-\Phi(w)\text{ with }\Phi(w)=P(w)\phi(w).$ (9)

When the vector field $\phi$ is the gradient of a function
$f:\mathbb{R}^{n}\to\mathbb{R}$, we recover mirror descent to minimize $f$ on
the simplex, with standard guarantees. In our warm-started formulation, the
field $\phi(w)=\Psi(\theta_{0},w)$ does not correspond to a gradient since its
Jacobian may not be symmetric. We analyze the stationary points of this ODE
and their stability. We recall that the support of a vector $w$,
$\mathrm{Supp}(w)$, is the set of indices in $\\{1,n\\}$ such that $w_{i}\neq
0$, and let $l$ its cardinal.

######  Proposition 3.

The stationary points of the mirror descent flow (9) are the $w$ such that
$\phi(w)\mid_{\mathrm{Supp}(w)}$ is proportional to $\mathbb{1}_{l}$.

All vectors of the form $w=(0,\dots,1,\dots,0)$ fulfill this condition, but
there might be other solutions with a larger support. To study the stability
of these points, we turn to the Jacobian of $\Phi$. Letting for short
$\phi=\phi(w)\in\mathbb{R}^{n}$ and $J=D\phi(w)\in\mathbb{R}^{n\times n}$, we
find

$D\Phi(w)=\mathrm{diag}(\phi)+\mathrm{diag}(w)J-\langle w,\phi\rangle
I_{n}-w(\phi^{T}+w^{T}J)$

Because of the simplex constraint, we only care about this Jacobian in
directions that are in the tangent space
$T_{\Delta}=\\{\delta\in\mathbb{R}^{n}|\enspace\sum_{i=1}^{n}\delta_{i}=0\\}$.
Since $\Phi$ makes the flow stay in the simplex, we have that
$D\Phi(w)[\delta]\in T_{\Delta}$ for any $\delta\in T_{\Delta}$. Without loss
of generality, we assume that $w_{1},\dots,w_{l}\neq 0$ and
$w_{l+1}=\dots=w_{n}=0$. The Jacobian greatly simplifies for coordinates $i$
such that $w_{i}=0$; indeed in this case for any vector $\delta$ in
$\mathbb{R}^{n}$ we have $D\Phi(w)[\delta]_{i}=(\phi_{i}-\langle
w,\phi\rangle)\delta_{i}$. On the other hand, for a coordinate $i$ in the
support, under the optimality conditions, taking a displacement of the form
$\delta=(\tilde{\delta}_{l},0,\dots,0)$, we find that
$D\Phi(w)[\delta]_{i}=w_{i}([\tilde{J}\tilde{\delta}]_{i}-\sum_{j=1}^{l}(\phi_{j}+[\tilde{J}^{T}w]_{j})\tilde{\delta}_{j})$,
where $\tilde{J}$ is the upper-left $l\times l$ block of $J$. In other words,
$D\Phi(w)$ has the following structure when $w$ is a stationary point of the
flow:

$D\Phi(w)=\begin{bmatrix}P(\tilde{w})\tilde{J}&(*)\\\
0&\mathrm{diag}(\phi_{i}-\langle w,\phi\rangle)\end{bmatrix}$

We readily obtain the stability condition, which is that this matrix has all
positive eigenvalues:

######  Proposition 4.

A stationary point $w$ of (9) is stable if and only if for all
$i\notin\mathrm{Supp}(w)$ we have $\phi(w)_{i}>\langle w,\phi(w)\rangle$ and
the matrix $P(\tilde{w})\tilde{J}$, as a linear operator $T_{\Delta}\to
T_{\Delta}$, has eigenvalues with positive real parts.

We can finally quantify the local speed of convergence towards these
stationary points:

######  Proposition 5.

Let $w^{*}$ be a stable stationary point of (9), and let $\delta$ be in the
tangent cone at $w^{*}$. Then, letting $w(t)$ the trajectory of the ODE (9)
starting from $w^{*}+\delta$, we have
$w(t)=w^{*}+\exp(-D\Phi(w^{*})t)\delta+o(\delta)$.

This result is classical in ODE theory [5]. We now turn to the case of
interest for bilevel optimization, where $\phi$ corresponds to the hyper-
gradient field with frozen parameters $\theta$, i.e.,
$\phi(w)=\Psi(\theta_{0},w)$.

###  3.2 The Bilevel Flow with Frozen Parameters

Figure 3: Weights dynamics from the mirror-flow ODE (9) with the field $\phi$
in Eq. (10), with $n=5$ training samples and $p=3$ parameters. Matrix $\Gamma$
is randomly sampled with i.i.d. Gaussian entries, same for $\nabla
F(\theta_{0})$, and the individual Hessians are taken as
$\nabla^{2}_{\theta\theta}\ell(\theta;x_{i})=u_{i}u_{i}^{T}+0.1I_{p}$ where
the $u_{i}$ are drawn i.i.d. from random Gaussians. The weights follow a non-
trivial dynamic that eventually converges to a solution with at most $p$ non-
zero weights, as predicted by Prop. 8. Most random initializations lead to
only one non-zero weight; we display a rarer dynamic here.

With fixed parameters $\theta_{0}$, the field of interest
$\phi(w)=\Psi(\theta_{0},w)$ is given by the simple equation

$\phi(w)=\Gamma g(w),$ (10)

where
$\Gamma=[\nabla\ell(\theta_{0};x_{1}),\dots,\nabla\ell(\theta_{0};x_{n})]\in\mathbb{R}^{n\times
p}$ is a matrix containing all the inner gradients, and
$g(w)=\left(\sum_{i=1}^{n}w_{i}\nabla^{2}_{\theta}\ell(\theta_{0};x_{i})\right)^{-1}\nabla
F(\theta_{0})\in\mathbb{R}^{p}$ is a the outer gradient $\nabla F(\theta_{0})$
transformed by the metric induced by the Hessian of the inner function $G$.
Figure 3 illustrates the weights’ dynamics with few samples on a low-
dimensional problem and its “sparsifying” effect.

The field $\phi$ only depends mildly on $w$, through the impact of $w$ on the
Hessian of the inner problem. We start our analysis with the simple case where
all Hessians $\nabla_{\theta}^{2}\ell(\theta_{0},x_{i})$ are the same, in
which case $g$ is constant:

######  Proposition 6.

If $g(w)$ is constant in Eq. 10, then the field $\phi(w)$ is constant, equal
to $\phi\in\mathbb{R}^{n}$. The mirror flow (9) converges to $w^{*}$ such that
$w^{*}_{i}=0$ if $i\notin\arg\max\phi$, and $w^{*}_{i}$ is proportional to
$w^{*}_{0}$ otherwise.

In particular, if $\phi$ only has a unique maximal coefficient, the mirror
flow converges to a $w^{*}$ with only _one_ non-zero weight – which is as
sparse as it gets. We now turn to the general case with non-constant Hessians.

The vector field $\phi$ has a “low rank” structure, as it is parameterized by
$g(w)$ of dimension $p$, which generally is much lower than $n$. It is,
therefore, natural that it is hard to satisfy the stability conditions of
Prop. 3 with non-sparse weights: we have $n$ conditions to verify with a
family of vectors $g(w)$ that is of dimension $p$. To formalize this
intuition, we define the following set:

######  Definition 1.

The set $\mathcal{I}_{l}^{p}$ is the set of $l\times p$ matrices $Z$ such that
$\mathbb{1}_{l}\in\mathrm{range}(Z)$ or $\mathbb{0}_{l}\in\mathrm{range}(Z)$.

In other words, this is the set of matrices $Z$ for which the equation
$Zx=\mathbb{1}_{l}$ or the equation $Zx=\mathbb{0}_{l}$ has a non-zero
solution $x$. This set either contains most matrices if $l\leq p$ or very few
if $p<l$:

######  Proposition 7.

Assume that $Z\in\mathbb{R}^{l\times p}$ has i.i.d. entries drawn from a
continuous distribution. If $l\leq p$, then
$\mathbb{P}(Z\in\mathcal{I}_{l}^{p})=1$, while if $l>p$ then
$\mathbb{P}(Z\in\mathcal{I}_{l}^{p})=0$.

This set is linked to the stationary points of the flow (9):

######  Proposition 8.

If $w$ is a stationary point of (9) with the hypergradient (10), letting $l$
the size of the support of $w$, we have
$\Gamma|_{\mathrm{Supp}(w)}\in\mathcal{I}_{l}^{p}$.

This proposition immediately shows that, in general, we cannot find a
stationary point of (9) with support larger than $p$, the number of features:
the mirror descent dynamics applied with the hypergradient, if it converges,
will converge to a sparse solution with at most $p$ non-zero coefficients.
Note that we could not show that the corresponding flow always converges; our
results only indicate that _if_ the flow converges, then it must be towards a
sparse solution. Through numerical simulations, we have identified
trajectories of the ODE that oscillate and do not seem to converge.

###  3.3 Two Variables Dynamics

We conclude our analysis by going back to the two variables ODE (8) and use
the previous analysis to show that the sparsity issue also impacts the warm-
started bilevel problem in the case where $\beta\gg\alpha$. To do so, we first
assume that the mirror flow ODE with frozen parameters $\theta$ in Eq (9)
converges for all $\theta$.

######  Assumption 2.

The ODE $\dot{w}=-\Psi(\theta,w)$ starting from $w_{0}$, with fixed $\theta$,
is such that $w(t)$ goes to a limit as $t$ goes to infinity. We call
$\Omega(\theta,w_{0})$ this limit.

Note that, following Prop. 8, the limit $\Omega$ is in general sparse with a
support smaller than $p$. We now give a result similar to Thm. 1 but when the
dynamics in $w$ gets much faster than that in $\theta$:

######  Theorem 2.

Let $\theta^{*}(t)$ the solution of the ODE $\dot{\theta}=-\nabla
G(\theta,\Omega(\theta,w_{0}))$, and $\theta^{\alpha,\beta},w^{\alpha,\beta}$
the solution of (8),where $w$ starts from $w_{0}$. Under technical assumptions
described in Appendix, for all $\beta$, for all time horizon $T$, we have

$\lim_{\alpha\to
0}\sup_{t\in[0,T]}\|\theta^{\alpha,\beta}(t/\alpha)-\theta^{*}(t)\|=0.$

As a consequence, the parameters $\theta$ track the gradient flow of $G$
obtained with the sparse weights $\Omega(\theta,w_{0})$. This leads to sub-
optimal parameters, which are estimated using only a few training samples, and
explains the behavior observed in Figure 2. Note that our theory only works in
the regime where $\alpha\gg\beta$ (Thm. 1) or $\alpha\ll\beta$ (Thm. 2), the
behavior of the warm-started bilevel in practice is therefore interpolating
between these two regimes. However, in practice, we observe that the warm-
started bilevel methods are attracted to sparse solutions, hinting at the fact
that Thm. 2 might better describe reality than Thm. 1.

##  4 Experiments

All experiments are run using the Jax framework [4] on CPUs.

###  4.1 The role of mirror descent

We place ourselves in the same toy setup with a mixture of two Gaussians that
is described in Sec. 2.4. We take $n=500$ and $m=100$ in dimension $2$. This
experiment aims to understand if mirror descent is critical to observing the
behavior described in the paper. We use another method to enforce the simplex
constraint, namely, we introduce the following inner function:
$G(\theta,\lambda)=\sum_{i=1}^{n}\frac{\sigma(\lambda_{i})}{\sum_{j}\sigma(\lambda_{j})}\ell(\theta;x_{i})$,
where $\sigma$ is the sigmoid function. Thus, we take the vector $\lambda$ as
outer parameters, which defines the weights as
$w_{i}=\frac{\sigma(\lambda_{i})}{\sum_{j}\sigma(\lambda_{j})}$ and we use
gradient descent on the $\lambda$ instead of mirror descent on $w$ to solve
the bilevel problem. We use
$F(\theta)=\frac{1}{m}\sum_{j=1}^{m}\ell^{\prime}(\theta;x_{j}^{\prime})$ as
outer function. We use the same algorithm as Algorithm 1 but without the
mirror step since the $\lambda$’s are unconstrained.

Figure 4: Convergence of Algorithm 2 on a toy mixture problem as described in
Sec. 2.4. We replace mirror descent by reparameterizing the weights as the
softmax of new parameters. Top: Weights at different iterations. Bottom:
training dynamics.

Figure 4 displays the results. In blue, we display the error between the
parameters and the target $\theta^{*}$. In orange, we display the error
between the parameters found by minimizing the inner function with the current
weights. Finally, the red curve tracks the entropy of the weights defined as
$\mathrm{entropy}=-\sum_{i=1}^{n}w_{i}\log(w_{i})$. We use entropy as a proxy
for sparsity: entropy is maximized when the weights are uniform and minimized
when all weights but one are $0$. Entropy decreases during training, as
suggested by our theory. Between iterations 1000 and 5000, the weight
distribution is not yet sparse and correctly identifies the good cluster;
hence minimizing the train loss with those weights leads to good results: the
orange curve is low. However, because of the warm started dynamics, the
parameters take some time to catch up, eventually converging only when the
weights are already sparse, leading to a large error. Here, replacing mirror
descent with reparameterization leads to the same behavior described in the
paper.

###  4.2 Hyper data-cleaning

Figure 5: Applying SOBA, a scalable warm-start bilevel method, for
datacleaning on the MNIST dataset, with a fixed inner learning rate and
different outer learning rates. Top: training curves. Bottom: Weights after
training, sorted in descending order, zooming in on the first $1500$ weights
out of $8000$.

We conduct a hyper data-cleaning experiment in the spirit of [9]. It is a
classification task on the MNIST dataset [15], where the training and testing
set consists of pairs of images of size $28\times 28$ and labels in
$\\{0,9\\}$. The gist of this experiment is that the training set is
_corrupted_ : with a probability of corruption $p_{c}=0.9$, each training
sample’s label is replaced by a random, _different_ label; corrupted samples
always have an incorrect label. The testing set is uncorrupted. We take $n=8K$
training samples; hence we only have $800$ clean training samples hidden in
the training set.

We use a linear model and a cross-entropy loss, with a small $\ell_{2}$
regularization of $10^{-2}$ on the training loss. We are, therefore, in the
strongly convex setting described in this paper: for any weight $w$, the inner
problem has one and exactly one set of optimal parameters. However, in this
case, Algorithm 1 is not implementable since we do not have a closed-form
solution to the regularized multinomial logistic regression.

The linear system resolution in Algorithm 2 is also impractical; we,
therefore, use the scalable algorithm SOBA [6], which has an additional
variable that tracks the solution to the linear system updated using Hessian-
vector products. We first run the algorithm with a fixed small inner learning
rate $\rho=10^{-3}$ and several outer learning rates $\eta$. Figure 5 displays
the results. The validation loss is computed on samples that are not part of
the test set. When the outer learning rate $\eta$ is too high (red curves), as
predicted by our theory, the weights go to very sparse solutions, leading to
sub-optimal performance. When the outer learning rate $\eta$ is too small, the
system converges to a good solution, but slowly (the bluest curve has not yet
converged). Overall, the range of learning rates where the algorithm converges
quickly to a good solution is very narrow.

Then, we run take different choices of inner learning rate $\rho$ and outer
learning rate $\eta$. We cover a large range of learning rate ratios
$r=\eta/\rho$ while ensuring the algorithm’s convergence. Hence, for a fixed
target ratio $r$, we pick $\eta$ and $\rho$ so that neither go over a
prescribed value $\eta_{\max}=1$ and $\rho_{\max}=10^{-2}$, by taking
$\eta=\min(r\rho_{\max},\eta_{\max})$ and $\rho=\eta/r$. The algorithm always
converges in this setup. We perform $12K$ iterations of the algorithm and then
compute two metrics: validation accuracy and entropy of the final weights. As
baselines for entropy, we compute the entropy of the _uniform distribution_ ,
equal to $\log(n)$, and of the _perfect distribution_ , which would put a
weight of $0$ on all corrupted data points, and uniform weight elsewhere,
equal to $(1-p_{c})\log(n)$.

Figure 6: Output of SOBA on the datacleaning task for a wide range of learning
rate ratios. Entropy is a proxy for weight’s sparsity.

Figure 6 displays the results. We see that when the outer learning rate $\eta$
is too small, nothing happens because the weights have not yet converged. The
corresponding accuracy is around $10\%$: training the model on corrupted data
completely fails. Then, when the ratio gets larger, reweighting starts
working: for ratios between $10^{-1}$ to $10^{2}$, the weights learned are not
sparse and correctly identify the correct data. This leads to an accuracy
close to the accuracy of the model trained on the perfect distribution, i.e.,
on the $800$ clean train samples. However, the time it takes to convergence is
roughly linear with the ratio: taking a ratio of $10^{-1}$ leads to training
about $1000$ times slower than with $10^{2}$. Finally, when the ratio is much
higher, we arrive at the region predicted by our theory (Thm. 2): weights
become extremely sparse, as highlighted by the entropy curve, and accuracy
decreases. We even reach a point where the entropy goes to $0$, i.e., only one
weight is non-zero. The corresponding accuracy is not catastrophic thanks to
warm-starting; it is much better than that of a model trained on that single
sample: the model has had time to learn something in the period where the
weights were non-zero. This observation is reminiscent of [25], which also
mention that the weights trajectory impacts the final parameters.

## Discussion & Conclusion

In this work, we have illustrated a challenge for data reweighting with
bilevel optimization: bilevel methods must use warm-starting to be practical,
but warm-starting induces sparse data weights, leading to sub-optimal
solutions. To remedy the situation, a small outer learning rate should
therefore be used, which might, in turn, lead to slow convergence. Classical
bilevel optimization theory [1] demonstrates the convergence of warm-started
bilevel optimization to the solutions of the true bilevel problem. This may
seem paradoxical at first and in contradiction with our results. Two
explanations lift the paradox: i)[1] require the ratio $\frac{\alpha}{\beta}$
to be smaller than some intractable constant of the problem, hence not
explaining the dynamics of the system in the setting where $\alpha$ is not
much smaller than $\beta$, which is the gist of this paper, ii) convergence
results in bilevel optimization are always obtained as non-convex results,
only proving that the gradient of $h$ goes to $0$. In fact, for the data
reweighting problem, several stationary points of $h$ are sparse (see Prop.
8). Hence, our results on the sparsity of the resulting solution can be seen
as implicit bias results, where the ODE converges to different solutions on
the manifold of stationary points. Finally, our results are orthogonal to
works considering the implicit bias of bilevel optimization [1, 25] These
results are based on over-parameterization, implying that the inner problem is
not strongly convex and has multiple minimizers, while our results do not
require such a structure.

## References

  * [1] M. Arbel and J. Mairal.  Amortized implicit differentiation for stochastic bilevel optimization.  arXiv preprint arXiv:2111.14580, 2021.
  * [2] A. Beck and M. Teboulle.  Mirror descent and nonlinear projected subgradient methods for convex optimization.  Operations Research Letters, 31(3):167–175, 2003.
  * [3] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al.  On the opportunities and risks of foundation models.  arXiv preprint arXiv:2108.07258, 2021.
  * [4] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * [5] C. Chicone.  Ordinary differential equations with applications, volume 34.  Springer Science & Business Media, 2006.
  * [6] M. Dagréou, P. Ablin, S. Vaiter, and T. Moreau.  A framework for bilevel optimization that enables stochastic and global variance reduction algorithms.  Advances in Neural Information Processing Systems, 35:26698–26710, 2022.
  * [7] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  * [8] J. Domke.  Generic methods for optimization-based modeling.  In Artificial Intelligence and Statistics, pages 318–326. PMLR, 2012.
  * [9] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil.  Forward and reverse gradient-based hyperparameter optimization.  In International Conference on Machine Learning, pages 1165–1173. PMLR, 2017.
  * [10] S. Ghadimi and M. Wang.  Approximation methods for bilevel programming.  arXiv preprint arXiv:1802.02246, 2018.
  * [11] R. Grazzi, M. Pontil, and S. Salzo.  Convergence properties of stochastic hypergradients.  In International Conference on Artificial Intelligence and Statistics, pages 3826–3834. PMLR, 2021.
  * [12] S. Gunasekar, J. Lee, D. Soudry, and N. Srebro.  Characterizing implicit bias in terms of optimization geometry.  In International Conference on Machine Learning, pages 1832–1841. PMLR, 2018.
  * [13] S. Gunasekar, B. Woodworth, and N. Srebro.  Mirrorless mirror descent: A natural derivation of mirror descent.  In International Conference on Artificial Intelligence and Statistics, pages 2305–2313. PMLR, 2021.
  * [14] K. Ji, J. Yang, and Y. Liang.  Bilevel optimization: Convergence analysis and enhanced design.  In International conference on machine learning, pages 4882–4892. PMLR, 2021.
  * [15] Y. LeCun, C. Cortes, C. Burges, et al.  Mnist handwritten digit database, 2010.
  * [16] J. Li, B. Gu, and H. Huang.  A fully single loop algorithm for bilevel optimization without hessian inverse.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7426–7434, 2022.
  * [17] J. Luketina, M. Berglund, K. Greff, and T. Raiko.  Scalable gradient-based tuning of continuous regularization hyperparameters.  In International conference on machine learning, pages 2952–2960. PMLR, 2016.
  * [18] D. Mahajan, R. Girshick, V. Ramanathan, K. He, M. Paluri, Y. Li, A. Bharambe, and L. Van Der Maaten.  Exploring the limits of weakly supervised pretraining.  In Proceedings of the European conference on computer vision (ECCV), pages 181–196, 2018.
  * [19] A. S. Nemirovskij and D. B. Yudin.  Problem complexity and method efficiency in optimization.  1983\.
  * [20] F. Pedregosa.  Hyperparameter optimization with approximate gradient.  In International conference on machine learning, pages 737–746. PMLR, 2016.
  * [21] M. Ren, W. Zeng, B. Yang, and R. Urtasun.  Learning to reweight examples for robust deep learning.  In International conference on machine learning, pages 4334–4343. PMLR, 2018.
  * [22] K. G. Samuel and M. F. Tappen.  Learning optimized map estimates in continuously-valued mrf models.  In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 477–484. IEEE, 2009.
  * [23] J. Shu, Q. Xie, L. Yi, Q. Zhao, S. Zhou, Z. Xu, and D. Meng.  Meta-weight-net: Learning an explicit mapping for sample weighting.  In NeurIPS, 2019.
  * [24] S. T. Tokdar and R. E. Kass.  Importance sampling: a review.  Wiley Interdisciplinary Reviews: Computational Statistics, 2(1):54–60, 2010.
  * [25] P. Vicol, J. P. Lorraine, F. Pedregosa, D. Duvenaud, and R. B. Grosse.  On implicit bias in overparameterized bilevel optimization.  In International Conference on Machine Learning, pages 22234–22259. PMLR, 2022.
  * [26] X. Wang, H. Pham, P. Michel, A. Anastasopoulos, J. Carbonell, and G. Neubig.  Optimizing data usage via differentiable rewards.  In International Conference on Machine Learning, pages 9983–9995. PMLR, 2020.
  * [27] J. Yang, K. Ji, and Y. Liang.  Provably faster algorithms for bilevel optimization.  Advances in Neural Information Processing Systems, 34:13670–13682, 2021.

##  Appendix A Proofs

###  A.1 Proof of Prop. 1

Let $\tilde{\theta}$ the global minimizer of the test loss
$\theta\mapsto\int\ell(\theta;x)d\nu(x)$. By definition, for all $\omega$, we
have
$\int\ell(\theta^{*}(\omega);x)d\nu(x)\geq\int\ell(\tilde{\theta};x)d\nu(x)$.
Furthermore, for $\omega^{*}=d\nu/d\mu$, we have that the inner loss becomes
$\int\ell(\theta;x)\omega^{*}(x)d\mu(x)=\int\ell(\theta;x)d\nu(x)$, the outer
loss. Minimizing it leads to $\theta^{*}(\omega^{*})=\tilde{\theta}$, which
is, therefore, the global minimizer of the bilevel problem.

###  A.2 Proof of Prop. 2

This is a standard discretization argument from ODE theory. The gist of the
proof is the following result. For a fixed $\alpha,\beta>0$, define

$\mathcal{F}(\theta,w,\tau)=\begin{bmatrix}-\alpha\nabla G(\theta,w)\\\
\frac{1}{\tau}\left(\frac{w\odot\exp(-\tau\beta\Psi(\theta,w))}{\|w\odot\exp(-\tau\beta\Psi(\theta,w))\|_{1}}-w\right)\end{bmatrix},$

so that the iterations of Algorithm 2 can be compactly rewritten as
$(\theta^{k+1},w^{k+1})=(\theta^{k},w^{k})+\tau\mathcal{F}(\theta^{k},w^{k},\tau)$.

We find that

$\lim_{\tau\to 0}\mathcal{F}(\theta,w,\tau)=\begin{bmatrix}-\alpha\nabla
G(\theta,w)\\\ -\beta P(w)\Psi(\theta,w)\end{bmatrix},$

hence Algorithm 2 is indeed a discretization of the ODE (8). The result
follows from classical ODE discretization theory (e.g. [5]).

###  A.3 Proof of Prop. 3

The stationary points of the ODE (9) are the $w\in\Delta_{n}$ such that
$\Phi(w)=0$. We see that we need to solve the equation $P(w)u=0$ for
$u\in\mathbb{R}^{n}$. The $i$-th coordinate of $P(w)u$ is $w_{i}(u_{i}-\langle
w,u\rangle)$. Hence, this cancels for all $i$ if and only if $w_{i}=0$ or
$u_{i}=\langle w,u\rangle$. This means that for the $i\in\mathrm{Supp}(w)$,
$u_{i}$ is constant, i.e., $u\mid_{\mathrm{Supp}(w)}$ is proportional to
$\mathbb{1}_{l}$. This gives the adverstized result.

###  A.4 Proof of Prop. 4

A stationary point is stable if and only if the operator

$D\Phi(w)=\begin{bmatrix}P(\tilde{w})\tilde{J}&(*)\\\
0&\mathrm{diag}(\phi_{i}-\langle w,\phi\rangle)\end{bmatrix}$

has positive eigenvalues. This matrix is block diagonal, hence, the condition
is that both $P(\tilde{w})\tilde{J}$ and $\mathrm{diag}(\phi_{i}-\langle
w,\phi\rangle)$ have positive eigenvalues. The condition on
$\mathrm{diag}(\phi_{i}-\langle w,\phi\rangle)$ can simply be rewritten as
$\phi_{i}>\langle w,\phi\rangle$, finishing the proof.

###  A.5 Proof of Prop. 5

The proof is found in [5], corollary 4.23.

###  A.6 Proof of Prop. 6

In this case, the flow becomes the ODE

$\dot{w}=-w\odot\phi+w\langle w,\phi\rangle$

The solution to this ODE is more obvious by looking at the corresponding
discrete mirror descent. In the discrete mirror descente case, it corresponds
to the iterates
$w^{k+1}=\frac{w^{k}\exp(-\tau\phi)}{\|w^{k}\exp(-\tau\phi)\|_{1}}$, hence a
simple recursion shows that $w^{k}=\frac{w^{0}\exp(-\tau
k\phi)}{\|w^{0}\exp(-\tau k\phi)\|_{1}}$. We then infer the solution in
continuous time, it is simply

$w(t)=\frac{w^{0}\odot\exp(-t\phi)}{\|w^{0}\odot\exp(-t\phi)\|_{1}}$

We see that is has the adverstized behavior: as $t$ goes to infinity, the i-th
coefficient of $\frac{w^{0}\odot\exp(-t\phi)}{\|w^{0}\odot\exp(-t\phi)\|_{1}}$
goes to $0$ if $i$ is not in the argmax of $\phi$, and is proportionnal to
$w_{0}$ otherwise. We also note that this result also holds in the discrete
mirror descent case, here the fact that we simplify things to an ODE plays not
role in the behavior.

###  A.7 Proof of Prop. 7

In the case where $l\leq p$, then with probability one, $Z$ is surjective,
i.e. the equation $Zx=y$ has a solution $x$ for all $y$. Therefore, $Z$ is in
$\mathcal{I}^{p}_{l}$ with probability one. One the contrary, if $l>p$, then
$\mathrm{range}(Z)$ is a random subspace of dimension $p$, hence the
probability that $\mathbb{1}_{l}$ or $\mathbb{0}_{l}$ are in this space is $0$
– note that for $\mathbb{0}_{l}$ we take the subspace range of $Z$ without the
$0$ term.

###  A.8 Proof of Prop. 8

This is simply Prop. 3 specialized to the field of the form $\phi(w)=\Gamma
g(w)$.

Indeed, Prop. 3 shows that a stationary point must satisfy that $\Gamma
g(w)\mid_{\mathrm{Supp}(w)}$ is proportional to $\mathbb{1}_{l}$, hence
$\Gamma\mid_{\mathrm{Supp}(w)}g(w)=\gamma\mathbb{1}_{l}$ for some $\gamma$.
Therefore, $\Gamma\mid_{\mathrm{Supp}(w)}$ is in $\mathcal{I}_{l}^{p}$:
depending on whether $\gamma=0$ or not, we must have a nonzero solution to
$\Gamma\mid_{\mathrm{Supp}(w)}x=\mathbb{0}_{l}$ or
$\Gamma\mid_{\mathrm{Supp}(w)}x=\mathbb{1}_{l}$.

###  A.9 Proof of Thm. 1

We will use many times the fact that the simplex $\Delta_{n}$ is a compact
set, hence any continuous function over $\Delta_{n}$ is bounded. We will use
the constant $d_{\theta}=\sup_{w\in\Delta_{n}}\|\theta^{*}(w)\|$, which is
finite thanks to the compacity of $\Delta_{n}$.

We first show a very useful lemma: the trajectory of $\theta$ remains bounded.

######  Lemma 1.

We let $\kappa=\frac{L_{G}}{\mu}$ the conditioning of $G$. If
$\|\theta(0)\|\leq
R:=\frac{\sqrt{2}(2+\kappa)+\sqrt{2(2+\kappa)^{2}+8\kappa}}{4}d_{\theta}$ then
for all $t,$ we have $\|\theta(t)\|\leq R$.

###### Proof.

We let $r(t)=\frac{1}{2}\|\theta(t)\|^{2}$. We find

$\displaystyle\frac{dr}{dt}$ $\displaystyle=\langle\dot{\theta},\theta\rangle$
(11) $\displaystyle=-\alpha\langle\nabla G(\theta,w),\theta\rangle$ (12)
$\displaystyle=-\alpha\left(\langle\nabla
G(\theta,w),\theta-\theta^{*}(w)\rangle+\langle\nabla
G(\theta,w),\theta^{*}(w)\rangle\right)$ (13)
$\displaystyle\leq\alpha\left(-\mu\|\theta-\theta^{*}(w)\|^{2}+\|G(\theta,w)\|\|\theta^{*}(w)\|\right)$
(14)
$\displaystyle\leq\alpha(-\mu(\|\theta\|^{2}-2\|\theta\|d_{\theta})+L_{G}d_{\theta}(d_{\theta}+\|\theta\|))$
(15)

where we have used the $\mu-$strong convexity and $L_{G}$-smoothness of $G$,
the inequality
$\|\theta-\theta^{*}\|^{2}\geq\|\theta\|^{2}-2\|\theta\|\|\theta^{*}\|$ and
the triangular inequality. Hence, we get the differential inequality

$\frac{dr}{dt}\leq-\alpha(2\mu
r^{2}-\sqrt{2}d_{\theta}(2\mu+L_{G})r-L_{G}d_{\theta}^{2}):=z(r)$

The advertised radius $R$ is exactly the zero of the right-hand side $z(r)$.
By Gronwall’s lemma, we have that $r(t)\leq u(t)$ where $u(t)$ is the solution
to the ODE $\dot{u}=z(u)$ with $u(0)=r(0)$. Since $R$ is a zero of $z$, the
trajectory of $u$ stays in $[0,R]$ following Picard-Lindelof theorem, hence
that of $r$ as well, proving the result.

∎

Note that the radius here does not depend on either $\alpha$ and $\beta$: the
dynamics in $\theta$ and $w$ both remain in a compact set that does not depend
on the steps. We will actively use this result throughout the rest of the
analysis because it allows us to bound any quantity that depends on $\theta$
and $w$. This lemma also immediately implies that the ODE has a solution for
all time horizons if $\|\theta(0)\|$ is small enough regardless of $\alpha$
and $\beta$.

To prove the main result, we first control the distance from $\theta(t)$ to
$\theta^{*}(w(t))$. We let
$\phi_{1}(t)=\frac{1}{2}\|\theta(t)-\theta^{*}(w(t))\|^{2}$. We find

$\displaystyle\frac{d\phi_{1}}{dt}(t)$ $\displaystyle=-\alpha\langle\nabla
G(\theta,w),\theta-\theta^{*}\rangle-\beta\langle
P(w)\Psi(\theta,w)\frac{d\theta^{*}(w)}{dw},\theta-\theta^{*}\rangle$ (16)
$\displaystyle\leq-2\alpha\mu\phi_{1}(t)+\beta\sqrt{2}K\|\Psi(\theta,w)\|\phi_{1}^{\frac{1}{2}},$
(17)

where $K>0$ is a constant. The first inequality comes from the strong
convexity of $G$, and the second one is Cauchy-Schwarz combined with the
boundedness of $P(w)$ and $\frac{d\theta^{*}(w)}{dw}$. We now use Lemma 1 to
get a crude bound on $\|\Psi\|$ :$\|\Psi\|\leq a_{1}$. Hence, we obtain the
differential inequation:

$\frac{d\phi_{1}}{dt}(t)\leq-2\alpha\mu\phi_{1}+\sqrt{2}\beta
Ka_{1}\phi_{1}^{\frac{1}{2}}.$

This equation is integrated in closed form using the following lemma:

######  Lemma 2.

Let $a,b>0$. The solution to the differential equation
$x^{\prime}=-ax+b\sqrt{x}$ starting from $x(0)\in[0,(b/a)^{2}]$ is
$x(t)=\left(\frac{b}{a}(1-\exp(-\frac{at}{2}))+\sqrt{x(0)}\exp(-\frac{at}{2})\right)^{2}$.

Hence, we get

$\phi_{1}(t)\leq\left(\frac{\beta Ka_{1}}{\sqrt{2}\alpha\mu}(1-\exp(-\alpha\mu
t))+\sqrt{\phi_{1}(0)}\exp(-\alpha\mu t)\right)^{2}\leq\left(\frac{\beta
Ka_{1}}{\sqrt{2}\alpha\mu}+\sqrt{\phi_{1}(0)}\exp(-\alpha\mu t)\right)^{2}$
(18)

demonstrating the last inequality in Thm. 1, as the first part of the previous
inequality goes to $0$ when $\beta$ goes to $0$.

Next, we consider
$\phi_{2}(t)=\frac{1}{2}\|w(\frac{t}{\beta})-w^{*}(t)\|^{2}$, where $w^{*}(t)$
is the solution to the ODE $\dot{w}=-P(w)\nabla h(w)$. We find, where for
short $\theta=\theta(\frac{t}{\beta})$, and $w=w(\frac{t}{\beta})$:

$\displaystyle\frac{d\phi_{2}}{dt}$ $\displaystyle=-\langle
P(w)\Psi(\theta,w)-P(w^{*})\Psi(\theta^{*}(w^{*}),w^{*}),w-w^{*}\rangle$ (19)
$\displaystyle\leq\|P(w)\Psi(\theta,w)-P(w^{*})\Psi(\theta^{*}(w^{*}),w^{*})\|\|w-w^{*}\|$
(20) $\displaystyle\leq
L_{P}\left(\|w-w^{*}\|+\|\theta-\theta^{*}(w^{*})\|\right)\|w-w^{*}\|,$ (21)

where $L_{P}$ is the Lipschitz constant of the map $(w,\theta)\mapsto
P(w)\Psi(\theta,w)$, and we have used the Cauchy-Schwarz inequality to control
the scalar product.

The term $\|\theta-\theta^{*}(w^{*})\|$ is controled with the triangular
inequality by doing

$\displaystyle\|\theta-\theta^{*}(w^{*})\|$
$\displaystyle\leq\|\theta-\theta^{*}(w)\|+\|\theta^{*}(w^{*})-\theta^{*}(w^{*})\|$
(22)

The first term is $\phi_{1}(\frac{t}{\beta})$ which is controlled by Eq. 18,
while the second term is controled by Lipschitzness of $\theta^{*}$ [10],
yielding

$\|\theta-\theta^{*}(w^{*})\|\leq\frac{\beta
Ka_{1}}{\sqrt{2}\alpha\mu}+\sqrt{\phi_{1}(0)}\exp(-\frac{\alpha}{\beta}\mu
t)+L_{\theta}\|w-w^{*}\|$

Overall, we get

$\frac{d\phi_{2}}{dt}\leq
L_{P}(3+L_{\theta})\phi_{2}+L_{P}(\phi_{1}(0)\exp(-2\frac{\alpha}{\beta}\mu
t)+(\frac{\beta Ka_{1}}{\sqrt{2}\alpha\mu})^{2})$

where we have used the inequality $\|w-w^{*}\|(\frac{\beta
Ka_{1}}{\sqrt{2}\alpha\mu}+\sqrt{\phi_{1}(0)}\exp(-\frac{\alpha}{\beta}\mu
t))\leq\phi_{2}\phi_{1}(0)\exp(-2\frac{\alpha}{\beta}\mu t)+(\frac{\beta
Ka_{1}}{\sqrt{2}\alpha\mu})^{2}$. Intuitively, this shows that $\phi_{2}$ is
small because as $\beta$ goes to $0$ we get to the inequality
$\frac{d\phi_{2}}{dt}\leq L_{P}(3+L_{\theta})\phi_{2}$, which implies that
$\phi_{2}=0$ since $\phi_{2}(0)=0$. We use Gronwall’s lemma and the fact that
$\phi_{2}(0)=0$ to prove this intuition. This gives

$\phi_{2}(t)\leq\frac{L_{P}\phi_{1}(0)}{2\frac{\alpha}{\beta}\mu+L_{P}(3+L_{\theta})}(\exp(L_{P}(3+L_{\theta})t)-\exp(-2\frac{\alpha}{\beta}\mu
t))+\frac{(\beta
Ka_{1})^{2}}{2\alpha^{2}\mu^{2}(L_{P}(3+L_{\theta}))}(\exp(L_{P}(3+L_{\theta})t)-1)$

which indeed goes to $0$ as $\beta$ goes to $0$.

###  A.10 Proof of Thm. 2

We need to be able to control the speed at which the mirror flow converges to
the solution $\Omega(\theta,w^{0}),$ to do so we posit a Lyapunov inequality:

######  Assumption 3.

There exists $\nu>0$ such that for all $w,\theta$, it holds

$\langle
P(w)\Psi(\theta,w),w-\Omega(\theta,w^{0})\rangle\geq\nu\|w-\Omega(\theta,w^{0})\|^{2}$

This assumption implies, in particular, that the flow
$\dot{w}=-P(w)\Psi(\theta,w)$ with fixed $\theta$ goes to
$\Omega(\theta,w^{0})$ at an exponential speed. We also assume that $\Omega$
is $L_{\Omega}$\- Lipschitz.

The analysis is then extremely similar to the previous one: we let
$\phi_{1}(t)=\frac{1}{2}\|w(t)-\Omega(\theta,w^{0})\|^{2}$ and find

$\displaystyle\frac{d\phi_{1}}{dt}$ $\displaystyle=-\beta\langle
P(w)\Psi(\theta,w),w-\Omega(\theta,w^{0})\rangle+\alpha\langle\frac{d\Omega}{d\theta}\nabla
G(\theta,w),w-\Omega(\theta,w^{0})\rangle$ (23)
$\displaystyle\leq-2\beta\nu\phi_{1}(t)+\alpha K\sqrt{\phi_{1}(t)}$ (24)

where $K$ is an upper bound on $\|\frac{d\Omega}{d\theta}\nabla
G(\theta,w)\|$. Therefore, Lemma 2 gives

$\phi_{1}(t)\leq\left(\frac{\alpha
K}{2\beta\nu}+\sqrt{\phi_{1}(0)}\exp(-\beta\nu t)\right)^{2}$

Next, we let
$\phi_{2}(t)=\frac{1}{2}\|\theta(\frac{t}{\alpha})-\theta^{*}(t)\|^{2}$ with
$\theta^{*}(t)$ the solution to the ODE $\dot{\theta}=-\nabla
G(\theta,\Omega(\theta,w^{0}))$. We get

$\displaystyle\frac{d\phi_{2}}{dt}$ $\displaystyle=-\langle\nabla
G(\theta,w)-\nabla
G(\theta^{*},\Omega(\theta^{*},w^{0})),\theta-\theta^{*}\rangle$ (25)
$\displaystyle\leq\|\nabla G(\theta,w)-\nabla
G(\theta^{*},\Omega(\theta^{*},w^{0}))\|\theta-\theta^{*}\|$ (26)
$\displaystyle\leq
L_{G}\left(\|\theta-\theta^{*}\|+\|w-\Omega(\theta^{*},w^{0})\|\right)\|\theta-\theta^{*}\|,$
(27)

The term $\|w-\Omega(\theta^{*},w^{0})\|$ is controlled by the triangular
inequality:

$\|w-\Omega(\theta^{*},w^{0})\|\leq\|w-\Omega(\theta,w^{0})\|+\|\Omega(\theta,w^{0})-\Omega(\theta^{*},w^{0})\|\leq\phi_{1}(\frac{t}{\alpha})+L_{\Omega}\|\theta-\theta^{*}\|$

and the rest of the proof follows in the same way as that of Thm. 1.