Fast Compute for ML Optimization

TL;DR

This paper introduces the Scale Mixture EM (SM-EM) algorithm for efficient optimization in machine learning, which removes the need for manual tuning and outperforms Adam in convergence speed and loss reduction on synthetic benchmarks.

Contribution

The paper proposes a novel EM-based optimization algorithm that automatically adapts learning rates and momentum, improving convergence and efficiency over traditional methods like Adam.

Findings

SM-EM with Nesterov acceleration achieves up to 13x lower loss than Adam.

Sharing statistics across penalty values reduces runtime by 10x.

EM guarantees nonincreasing objectives; acceleration improves convergence speed.

Abstract

We study optimization for losses that admit a variance-mean scale-mixture representation. Under this representation, each EM iteration is a weighted least squares update in which latent variables determine observation and parameter weights; these play roles analogous to Adam's second-moment scaling and AdamW's weight decay, but are derived from the model. The resulting Scale Mixture EM (SM-EM) algorithm removes user-specified learning-rate and momentum schedules. On synthetic ill-conditioned logistic regression benchmarks with $p \in \{20, \ldots, 500\}$, SM-EM with Nesterov acceleration attains up to $13\times$ lower final loss than Adam tuned by learning-rate grid search. For a 40-point regularization path, sharing sufficient statistics across penalty values yields a $10\times$ runtime reduction relative to the same tuned-Adam protocol. For the base (non-accelerated) algorithm, EM monotonicity guarantees nonincreasing objective values; adding Nesterov extrapolation trades this guarantee for faster empirical convergence.