Implicit Bias and Convergence of Matrix Stochastic Mirror Descent

This paper proves that in the overparameterized regime, Stochastic Mirror Descent with matrix parameters converges exponentially to a global interpolator while exhibiting an implicit bias toward the solution that minimizes the Bregman divergence induced by the mirror function from the initialization.

The Gist

Imagine you are trying to solve a giant, messy puzzle. You have a picture (the data), but many pieces are missing. Your goal is to fill in the blanks to recreate the original image.

In the world of machine learning, this is called Matrix Completion. Usually, we assume the picture has a simple structure (like a low-resolution sketch rather than a chaotic scribble), which helps us guess the missing pieces.

This paper introduces a new, smarter way to solve this puzzle using a method called Matrix Stochastic Mirror Descent (SMD). Here is the breakdown in everyday language:

1. The Problem: Too Many Answers, Which One is Right?

When you have a puzzle with many missing pieces, there are often thousands of ways to fill them in that technically fit the pieces you do have.

• The Old Way: Standard algorithms (like Gradient Descent) are like a hiker walking down a hill. They just look for the lowest point nearby. If there are many paths to the bottom, they might get stuck in a random one, or they might pick a solution that looks "messy" or overly complicated.
• The New Insight: The authors realized that the shape of the hill you walk down matters just as much as the destination. By changing the "shape" of the terrain, you can force the algorithm to find a specific, cleaner solution.

2. The Solution: The "Mirror" Map

The authors use a concept called a Mirror Map.

• The Analogy: Imagine you are navigating a city.
  • Standard GPS (Gradient Descent): Tells you to go "North 5 blocks, East 3 blocks." It treats the city as a perfect grid.
  • Mirror Descent: Tells you to go "Walk until you feel the wind change," or "Follow the river." It uses a different set of rules (a "mirror") to navigate.
• Why it helps: In this paper, the "Mirror" is designed to prefer simple, low-rank solutions. Think of "low-rank" as a solution that is smooth and organized, rather than chaotic. The algorithm is biased (in a good way) to find the "cleanest" possible picture that fits the data.

3. The Magic: "Implicit Bias"

This is the coolest part. The authors prove that you don't need to tell the algorithm, "Hey, find the simplest picture!" explicitly.

• The Metaphor: Imagine you are teaching a dog to fetch. If you train the dog in a specific way (using a specific toy and a specific command), the dog will naturally learn to fetch only that toy, even if you never explicitly said "Don't fetch the ball."
• The Result: The algorithm naturally "biases" itself toward the simplest, most elegant solution (the one that minimizes a specific mathematical distance called Bregman Divergence) just by virtue of how it takes its steps. It finds the "Goldilocks" solution—not too complex, not too simple, but just right.

4. Speed and Stability

The paper proves two major things:

1. It Converges: The algorithm is guaranteed to eventually find the perfect solution that fits all the known data points.
2. It's Fast: It doesn't just wander aimlessly; it zooms toward the solution exponentially fast (like a rocket, not a snail).

5. The Real-World Test: Filling in the Blanks

To prove this works, the authors tested it on a classic problem: Matrix Completion.

• The Setup: They took a 100x100 grid of numbers (like a spreadsheet) and hid 90% of the numbers.
• The Competition: They pitted their new "Mirror Descent" method against two old-school methods (SVT and Soft-Impute) that are the industry standard for this task.
• The Outcome: The new method was better. It reconstructed the missing numbers more accurately, especially when very few numbers were visible (the hardest scenarios). It was like having a detective who could guess the missing parts of a crime scene photo with much higher accuracy than the usual suspects.

Summary

Think of this paper as upgrading the navigation system for AI.

• Old System: "Just go downhill until you stop." (Good, but might get lost in messy solutions).
• New System: "Go downhill, but follow a special map that naturally guides you to the cleanest, most organized solution."

This is a big deal for fields like recommender systems (Netflix guessing what you'll like next), medical imaging (reconstructing blurry MRI scans), and data science, where we often have to guess missing information based on a few clues. The authors show that by changing how we calculate the steps, we get better, faster, and cleaner results automatically.

Technical Summary

1. Problem Statement

The paper addresses the optimization of matrix parameters ($W \in \mathbb{R}^{d \times k}$) in overparameterized regimes (where the number of parameters $d \times k$ exceeds the number of training samples $p$). This setting is relevant to:

• Matrix Completion: Recovering a low-rank matrix from a subset of observed entries.
• Multi-class Linear Classification: Learning a weight matrix that perfectly interpolates one-hot encoded labels for multiple classes.

The core challenge is understanding the implicit bias of optimization algorithms in these settings. While standard Stochastic Gradient Descent (SGD) is known to find the minimum $\ell_2$-norm solution, the authors investigate how Stochastic Mirror Descent (SMD) behaves when the parameters are matrices and the geometry is defined by a general mirror function $\psi(\cdot)$.

2. Methodology

A. Matrix Stochastic Mirror Descent (Matrix SMD)

The authors propose an extension of SMD to matrix parameters. Given a strongly convex potential function (mirror map) $\psi: \mathbb{R}^{d \times k} \to \mathbb{R}$, the update rule at iteration $t$ is:
$$\nabla\psi(W_t) = \nabla\psi(W_{t-1}) - \eta \nabla_W L_t(W_{t-1})$$
where $L_t$ is the loss computed on a random batch.

• Generalization: When $\psi(W) = \frac{1}{2}\|W\|_F^2$, this recovers standard SGD.
• Flexibility: By choosing different $\psi$, the algorithm induces different geometries, leading to different implicit biases.

B. Theoretical Framework

The authors establish convergence and implicit bias under the following conditions:

1. Loss Function: The loss $L(W)$ is a sum of convex functions $\ell_i$ applied to linear constraints $A(W)_i = b_i$.
2. Mirror Function: $\psi$ is differentiable and $\nu$-strongly convex.
3. Overparameterization: The linear operator $A$ has full rank ($\sigma_p(A) > 0$), ensuring a solution space exists.
4. Step Size: The learning rate $\eta$ is small enough to maintain convexity of the composite function.

C. Key Theoretical Results

Theorem 1 (Convergence and Implicit Bias):

• Implicit Bias: The algorithm converges to the unique solution $W^*$ that minimizes the Bregman divergence $D_\psi(W, W_0)$ from the initialization $W_0$, subject to interpolating the data ($A(W)=b$).
  $$W^* = \arg\min_{W: A(W)=b} D_\psi(W, W_0)$$
• Convergence Rate: Under additional assumptions (specifically regarding the boundedness of the Hessian of $\psi$ on the feasible set), the algorithm converges exponentially to $W^*$:
  $$\mathbb{E}\|W^* - W_t\|_F^2 \leq 2 \left( 1 - \frac{\eta \mu \sigma_p(A)^2}{2pL} \right)^t D_\psi(W^*, W_0)$$
  where $L$ is a Lipschitz-like constant related to the Bregman divergence.

Special Case: Schatten-p Norms
The authors analyze $\psi(W) = \|W\|_{Schatten, p}^p$ (sum of singular values raised to power $p$).

• For $p \geq 2$, exponential convergence is guaranteed.
• For $1 < p < 2$, convergence is guaranteed, but the exponential rate requires further relaxation of assumptions (specifically regarding singular matrices in the feasible set).

3. Key Contributions

1. Extension to Matrix Parameters: The paper generalizes existing implicit bias theory (previously limited to vector parameters) to matrix-valued parameters and vector-valued outputs.
2. Exponential Convergence Proof: It provides rigorous proofs for the exponential convergence rate of Matrix SMD in the overparameterized regime, relaxing the common requirement for $L$-smoothness of the loss function.
3. Characterization of Inductive Bias: It demonstrates that Matrix SMD converges to the solution minimizing the Bregman divergence induced by the mirror map. This allows practitioners to "program" the inductive bias (e.g., low-rankness) directly into the optimization geometry via $\psi$.
4. Practical Application to Matrix Completion: The authors apply the theory to the matrix completion problem using $\psi(W) = \|W\|_{Schatten, p}^p$ with $p \approx 1$. This approximates the nuclear norm (sum of singular values) but via a smooth mirror map rather than a proximal operator.

4. Experimental Results

The authors evaluated their method on a low-rank matrix completion task ($100 \times 100$ matrix, rank 5) with varying sampling probabilities (0.1 to 0.9).

• Comparison: They compared Schatten-p SMD ($p=1.05$) against two standard baselines:
  1. Singular Value Thresholding (SVT): Uses a proximal operator with a soft-thresholding step.
  2. Soft-Impute: An iterative imputation method using soft-thresholding.
• Findings:
  • Schatten-p SMD consistently outperformed both SVT and Soft-Impute across all sampling rates.
  • The performance gap was most significant in low-sampling regimes (e.g., 10% observed entries), where the problem is most ill-posed.
  • The SMD approach achieved lower relative Frobenius norm error ($\|W - M\|_F / \|M\|_F$) despite running for the same number of iterations as the baselines.

5. Significance and Conclusion

• Theoretical Insight: The paper bridges the gap between optimization dynamics and the structural properties of the learned model in high-dimensional matrix problems. It proves that the choice of mirror map $\psi$ dictates the inductive bias, effectively selecting the "simplest" solution (in terms of Bregman divergence) that fits the data.
• Practical Impact: The results suggest that Matrix SMD is a superior alternative to traditional proximal methods (like SVT) for matrix completion. By using a smooth mirror map with $p \approx 1$, the algorithm naturally induces low-rank structure without needing explicit non-smooth regularization terms or complex projection steps.
• Future Directions: The authors note that proving exponential convergence rates for the $1 < p < 2$ regime (crucial for approximating the nuclear norm) requires relaxing the assumption that the feasible set does not contain singular matrices, which remains an open problem.

In summary, this work establishes a robust theoretical foundation for Matrix SMD, proving its exponential convergence and demonstrating its practical superiority in recovering low-rank matrices compared to standard thresholding techniques.