Zeroth-Order Optimization Finds Flat Minima

TL;DR

It is shown that zeroth-order optimization with the standard two-point estimator favors solutions with small trace of Hessian, which is widely used in previous work to distinguish between sharp and flat minima.

Abstract

Zeroth-order methods are extensively used in machine learning applications where gradients are infeasible or expensive to compute, such as black-box attacks, reinforcement learning, and language model fine-tuning. Existing optimization theory focuses on convergence to an arbitrary stationary point, but less is known on the implicit regularization that provides a fine-grained characterization on which particular solutions are finally reached. We show that zeroth-order optimization with the standard two-point estimator favors solutions with small trace of Hessian, which is widely used in previous work to distinguish between sharp and flat minima. We further provide convergence rates of zeroth-order optimization to approximate flat minima for convex and sufficiently smooth functions, where flat minima are defined as the minimizers that achieve the smallest trace of Hessian among all optimal solutions. Experiments on binary classification tasks with convex losses and language model fine-tuning support our theoretical findings.

Figures (6)

• Figure 1: Loss and trace of Hessian on Example \ref{['exp:xy']} with $d=100$. (a) Comparisons among gradient descent (GD), zeroth-order (ZO) optimization (Algorithm \ref{['algo:zo']}) with $\lambda=0.1$, and ZO with $\lambda\to 0$ (directional derivatives). (b) Comparisons on different $\lambda$ used in ZO. GD and ZO with $\lambda\to 0$ fail to decrease the trace of Hessian. Larger $\lambda$ in ZO leads to larger errors in the loss but brings more regularization effect on minimizing the trace of Hessian.
• Figure 2: Training loss, trace of Hessian on the training data (Train Hess in the plot), and test accuracy on (a) SVMs and (b) Logistic regression. Zeroth-order (ZO) optimization is slower than gradient descent (GD) for minimizing the loss and maximizing the accuracy, but there is a clear trend of decreasing trace of Hessian for ZO.
• Figure 3: Training loss, trace of Hessian on the training data (Train Hess in the plot), and test accuracy on (a) SST-2 with $K=32$ and $K=256$, and (b) SST-5 and TREC with $K=32$. Both gradient descent (GD) and zeroth-order (ZO) optimization reduce the trace of Hessian during training. In most cases, GD achieves lower training loss, smaller trace of Hessian, and higher test accuracy.
• Figure 4: Loss and trace of Hessian (Hess in the plot) on Example \ref{['exp:xy']} using different random seeds. The observation aligns with Figure \ref{['fig:toy']} (random seed $13$). Zeroth-order optimization (ZO) decreases the trace of Hessian, and $\lambda$ controls the trade-offs between regularization on the trace of Hessian and the optimization error induced by additional bias terms.
• Figure 5: Training loss, trace of Hessian on the training data (Train Hess in the plot), and test accuracy on (a) SVMs and (b) Logistic regression using different random seeds. The observation is consistent with Figure \ref{['fig:cvx']} (seed $29$), where zeroth-order optimization (ZO) reduces the trace of Hessian.

Theorems & Definitions (26)

• Example 2.1
• Proposition 2.2
• Definition 3.1: Flat Minima
• Definition 3.2: Approximate Flat Minima
• Lemma 3.5
• Lemma 3.6
• Theorem 1
• Corollary 2: Iteration Complexity for Finding Flat Minima
• Remark 3.1
• Remark 3.2