Search papers, labs, and topics across Lattice.
This paper analyzes the convergence of data mixing, a bilevel optimization problem for reweighting training domains, when using a finite number of inner loop steps ($T$). It demonstrates that a single inner step ($T=1$) can fail even in simple cases. The paper proves that, under a fixed parameter update budget $N$ and strong convexity, the optimal number of inner steps scales as $Θ(\log N)$ for full gradients and $Θ({(N \log N)}^{1/2})$ for stochastic gradients, suggesting fewer weight updates over a longer horizon can be beneficial.
Surprisingly, using only a single inner loop update in data mixing can lead to failure, and the optimal number of inner loop steps scales logarithmically with the parameter update budget.
Data mixing--the strategic reweighting of training domains--is a critical component in training robust machine learning models. This problem is naturally formulated as a bilevel optimization task, where the outer loop optimizes domain weights to minimize validation loss, and the inner loop optimizes model parameters to minimize the weighted training loss. Classical bilevel optimization relies on hypergradients, which theoretically require the inner optimization to reach convergence. However, due to computational constraints, state-of-the-art methods use a finite, often small, number of inner update steps before updating the weights. The theoretical implications of this approximation are not well understood. In this work, we rigorously analyze the convergence behavior of data mixing with a finite number of inner steps $T$. We prove that the "greedy" practical approach of using $T=1$ can fail even in a simple quadratic example. Under a fixed parameter update budget $N$ and assuming the per-domain losses are strongly convex, we show that the optimal $T$ scales as $Θ(\log N)$ (resp., $Θ({(N \log N)}^{1/2})$) for the data mixing problem with access to full (resp., stochastic) gradients. We complement our theoretical results with proof-of-concept experiments.
