Source-linked AI summary
Fast large-scale optimization by unifying stochastic gradient and quasi-Newton methods
Jascha Sohl-Dickstein, Ben Poole, Surya Ganguli
TL;DR
Stochastic optimization can evaluate one subfunction per step, but combining those gradients to estimate the full objective’s inverse Hessian is not straightforward. SFO maintains separate quadratic approximations in a shared adaptive low-dimensional subspace, and outperformed the compared methods on seven objective functions.
Problem
With one subfunction evaluated per step, combining gradients to estimate the full objective’s inverse Hessian is not straightforward, and SGD typically requires hyperparameter tuning.
Source: §1 · p. 2Method
SFO maintains a separate quadratic approximation for each subfunction and stores them in a shared adaptive low-dimensional subspace.
Source: §1 · pp. 2–3Results
Takeaways & Limitations
The optimizer and test suite are released as open-source Python and MATLAB packages.
Source: §1 · p. 3Quadratic functions can be a poor match for the geometry of an objective function.
Source: §6 · p. 8
Abstract
from arXiv · showhide
We present an algorithm for minimizing a sum of functions that combines the computational efficiency of stochastic gradient descent (SGD) with the second order curvature information leveraged by quasi-Newton methods. We unify these disparate approaches by maintaining an independent Hessian approximation for each contributing function in the sum. We maintain computational tractability and limit memory requirements even for high dimensional optimization problems by storing and manipulating these quadratic approximations in a shared, time evolving, low dimensional subspace. Each update step requires only a single contributing function or minibatch evaluation (as in SGD), and each step is scaled using an approximate inverse Hessian and little to no adjustment of hyperparameters is required (as is typical for quasi-Newton methods). This algorithm contrasts with earlier stochastic second order techniques that treat the Hessian of each contributing function as a noisy approximation to the full Hessian, rather than as a target for direct estimation. We experimentally demonstrate improved convergence on seven diverse optimization problems. The algorithm is released as open source Python and MATLAB packages.
1. Introduction
SFO combines stochastic updates with quasi-Newton curvature by maintaining a separate quadratic approximation for each subfunction and updating only one subfunction per step. It uses a shared adaptive low-dimensional subspace to manage computation and memory, and the paper reports faster optimization on seven problems.
- SFO estimates each subfunction’s Hessian directly, maintaining separate quadratic approximations so one subfunction can be evaluated per update without large minibatches.This differs from earlier stochastic second-order methods, which treat a subset’s Hessian as a noisy approximation to the full Hessian and may require many subfunctions per update or yield highly regularized approximations.Source: §1 · pp. 2–3
- A shared, adaptive low-dimensional subspace determined by recent gradients and iterates keeps SFO’s computation and memory tractable for high-dimensional problems.The subspace is used to store and manipulate the subfunctions’ approximations.Source: §1 · pp. 2–3
- SFO’s second-order approximations are intended to reduce hyperparameter adjustment, while its single-subfunction updates avoid requiring large minibatches to estimate the full Hessian.The introduction contrasts this with SGD’s often painstaking hyperparameter tuning and previous stochastic Hessian methods’ reliance on regularization and large minibatches.Source: §1 · pp. 2–3
- The paper reports faster optimization on seven disparate problems and releases SFO with its test suite as open-source Python and MATLAB packages.Source: §1 · p. 3
2. Algorithm
The algorithm minimizes an objective built from stored quadratic approximations, updating one subfunction at a time. It keeps the optimization tractable by representing its history and updates in a shared, adaptive low-dimensional subspace and maintaining per-subfunction Hessian approximations.
2.1. Approximating Functions
- The objective is approximated by stored quadratic functions, one for each contributing function, with one such approximation updated per learning step.Together, the approximations form Gt, intended to approximate F.Source: §2.1 · p. 3
2.2. Update Steps
- At each step, the method minimizes the current approximating objective, then updates one selected subfunction using its value and gradient and a BFGS Hessian estimate.The position is chosen by a Newton step; the other subfunctions remain unchanged.Source: §2.2 · p. 3
2.3. A Shared, Adaptive, Low-Dimensional Representation
- To control large-problem costs, the method stores history and computes updates in a shared subspace containing each subfunction’s most recent gradient and position.Its minimum dimension is at least 2N, ensuring inclusion of the full-batch steepest-gradient direction and recent Newton update directions; experiments use Kmin = 2N and Kmax = 3N.Source: §2.3 · pp. 3–4
2.3.1. EXPANDING THE SUBSPACE WITH A NEW OBSERVATION
- Each iteration adds the component of the updated subfunction’s latest gradient that lies outside the existing subspace, while the new position is already represented there.Existing vectors are carried into the expanded subspace by appending a zero coordinate.Source: §2.3.1 · p. 4
2.3.2. RESTRICTING THE SIZE OF THE SUBSPACE
- When the subspace exceeds Kmax, it is collapsed to the latest gradient and position measurements for each subfunction, using QR decomposition and projection.Projection can discard components outside the new subspace; the QR decomposition and projection are computed within the prior subspace.Source: §2.3.2 · p. 4
2.4.1. HISTORY MATRICES
- An independent online Hessian approximation is maintained for each subfunction using BFGS on that subfunction’s gradient and position history.The history stores successive gradient and position changes and is truncated to L columns; the paper uses L = 10.Source: §2.4 · p. 4, §2.4.1 · p. 4
2.4.2. BFGS UPDATES
- BFGS processes each history pair to satisfy the secant equation while minimizing the weighted Frobenius difference from the prior estimate, and uses the final estimate as the subfunction Hessian.The update is computed successively from the oldest to the most recent history column.Source: §2.4.2 · p. 4
3. Implementation Details
This section reviews additional algorithm design choices, with fuller treatments in Appendix C.
3.1. BFGS Initialization
- For a subfunction’s first evaluation, SFO initializes its Hessian to the identity scaled by the median eigenvalue of other active subfunctions’ average Hessian; later it uses B0 = βI.For later evaluations, β is the minimum eigenvalue found by solving the squared secant equation for the full history.Source: §3.1 · p. 5
3.2. Enforcing Positive Definiteness
- SFO enforces positive-definite Hessian approximations by eigendecomposition, replacing any too-small eigenvalues with the median positive eigenvalue.The passage says this is computationally cheap because the approximations use a shared low-dimensional subspace.Source: §3.2 · p. 5
3.3. Choosing a Target Subfunction
- The target subfunction is the one farthest from the current location under the current Hessian approximation as metric, outperforming random selection.The passage points to Supplemental Figure C.1 for the comparison.Source: §3.3 · p. 5
3.4. Growing the Number of Active Subfunctions
- SFO begins with two active subfunctions and expands the set when gradient standard error across subfunctions exceeds gradient length; performance differs from all-active initialization only in the first several passes.Figure 2(c) reports near-maximal performance on two convex problems with only a small number of minibatches.Source: §3.4 · p. 5, Figure 2 · p. 5
3.5. Detecting Bad Updates
- SFO detects bad proposed updates by comparing the measured subfunction value at the current location with its quadratic approximation.The passage notes that quasi-Newton methods commonly detect and reject bad steps, for instance through line search.Source: §3.5 · pp. 5–6
4. Properties
SFO limits computational and memory costs by maintaining its approximations in a shared low-dimensional subspace. The section also describes minibatch sizing to balance evaluation and projection costs and discusses convergence under conditions related to prior work.
4.1. Computational Overhead and Storage Cost
- SFO’s dominant per-iteration costs are O(MN) for projection and O(Q) for evaluating one subfunction, while storing the active subspace costs O(MN) memory.The active subspace has dimension N within the M-dimensional parameter space.Source: §4.1 · p. 6
4.1.1. IDEAL MINIBATCH SIZE
- With minibatches of size S, N=D/S and subfunction evaluation costs O(MS); choose S to balance evaluation and projection costs.The proportionality constant should be small enough that most time is spent evaluating the subfunction. The ideal S may be larger for sparse objectives, and performance remains very good even for small N.Source: §4.1.1 · p. 6
4.2. Convergence
- (Mairal, 2013) establishes monotone decrease and asymptotic stationarity under majorization and smoothness conditions, and linear convergence when the fi are strongly convex.The passage suggests a near-identical proof for simplified SFO with random subfunction update order and regularized Ht_i to ensure majorization.Source: §4.2 · p. 6
5. Experimental Results
Across seven objective functions, SFO outperformed all competing techniques. The experiments compared optimizers across six Figure 3 tasks and a separate CURVES neural-network task.
5. Experimental Results
- SFO outperformed all other compared techniques across seven objective functions, with comparisons shown in Figures 3 and 4.Source: §5 · p. 6, Figure 3 · p. 7
5.1. Optimizers
- Figure 3 compares SFO with SAG, SGD, AdaGrad, LBFGS, minibatch LBFGS, and Hessian-free; unlike most, SFO requires no hyperparameter tuning.SAG, SGD, and AdaGrad settings were grid-searched; training data was divided into 100 minibatches for the other experiments, while Hessian-free used four chunks.Source: §5.1 · pp. 6–7, Figure 3 · p. 7
5.2. Objective Functions
- The seven objectives covered logistic regression, autoencoding, ICA, Ising/Hopfield, multilayer perceptron, CIFAR-10 convolution, and a 12-layer CURVES network.The logistic regression and Ising/Hopfield objectives were convex and plotted relative to the smallest objective value achieved by any optimizer.Source: Figure 3 · p. 7, §5.2 · p. 8
6. Future Directions
SFO converges in fewer effective data passes than Hessian-free optimization in the reported CURVES experiment, while the section outlines possible extensions to its subspace, parallelism, approximations, initialization, and analysis.
6. Future Directions
- An online subspace built from leading eigenvectors of gradient and update-step covariances could reduce memory and computational overhead without substantially hurting performance.The proposed reduction is from the current O(N)-dimensional subspace and is intended to remain useful for large numbers of subfunctions.Source: §6 · p. 8
- Asynchronous updates, including evaluations at stale positions, could enable massive-scale parallelization, while rolling subfunction replacement could adapt SFO to online or infinite data.The subfunction updates are described as naively parallelizable.Source: §6 · p. 8
- Because quadratic functions can poorly match objective geometry, exploring non-quadratic approximating subfunctions could improve performance.The dynamically updated subspace and independent approximations fitted to the true subfunctions do not depend on the approximations being quadratic.Source: §6 · p. 8
- On the twelve-layer CURVES network, SFO converges in approximately one tenth the effective data passes of Hessian-free optimization, in an experiment matching Martens (2010).The Hessian-free convergence trace was generated using source code from that paper.Source: Figure 4 · p. 8
- Initializing each subfunction’s approximate Hessian from the average of the others could exploit their similarity, with individual subfunctions overwriting the initialization where they diverge.This is proposed as an alternative to the diagonal initialization described in Section 3.1.Source: §6 · p. 8
- Further directions include analyzing SFO’s non-asymptotic convergence and incorporating the natural gradient through a change of variables that equates it with the ordinary gradient.The proposed natural-gradient technique is described as straightforward to incorporate into SFO.Source: §6 · p. 8
7. Conclusion
The method combines LBFGS-style quasi-Newton optimization with stochastic gradient descent by maintaining an independent BFGS quadratic approximation for each objective subfunction. It alternates descent on the full-objective approximation with a single-subfunction update, using a shared low-dimensional subspace to keep computation and memory tractable.
- The method combines LBFGS-style quasi-Newton optimization and stochastic gradient descent by maintaining an independent BFGS quadratic approximation for each objective subfunction.Source: §7 · p. 8
- Each step alternates between descending the full-objective quadratic approximation and evaluating one subfunction to update its approximation.Source: §7 · p. 9
- A shared low-dimensional subspace defined by gradient-evaluation history makes the procedure tractable in memory and computational time.Source: §7 · p. 9
quasi-Newton methods - Supplemental Material … B.2. Autoencoder
The supplemental material describes the computational costs of SFO’s components and specifies the objective functions used in the experiments. It details the optimizer’s per-pass operations and the setups for logistic regression and a contractive autoencoder.
A. Computational Complexity
- The computational-complexity discussion covers the costs of SFO’s individual components.Source: A. Computational Complexity · p. 10
A.1. Function and Gradient Computation
- Computing each subfunction’s value and gradient costs O(Q), giving O(QN) for a full pass through N subfunctions.Source: §A.1 · p. 10
A.2. Subspace Projection
- Each iteration projects parameters and gradients between the full M-dimensional space and the N-dimensional subspace, and appends the gradient-projection residual, at O(MN) per operation.Source: §A.2 · p. 10
A.3. Subspace Collapse
- To keep the subspace dimension K_t at O(N), SFO collapses it every O(N) steps using QR decomposition and projects the old subspace into the new one.Source: §A.3 · p. 10
A.4. Minimize Gt (x)
- SFO explicitly inverts the matrix minimizing G_t(x) N times per pass; a Woodbury update is possible, but improving this step would have little overall effect.Source: §A.4 · p. 10
A.5. BFGS
- BFGS runs in an L-dimensional subspace with rank-two updates; recomputing the full chain each iteration preserves history-dependent initialization, while incremental updates sacrifice it.Source: §A.5 · p. 10
B. Objective Functions; B.1. Logistic Regression
- The supplement details experimental objectives, including logistic regression matched to Roux et al. (2012), except that the total objective is divided by minibatch size but not minibatch count.Source: B. Objective Functions · p. 10, §B.1 · p. 10
B.2. Autoencoder
- The MNIST contractive autoencoder used 784 visible units, 256 hidden units, sigmoid encoder and decoder nonlinearities, and a regularization penalty of 1.Source: §B.2 · p. 11
B.3. Independent Components Analysis … B.6. Deep Convolutional Network
The sections describe four optimization experiments: ICA and Ising/Hopfield models on MNIST, a multilayer perceptron on MNIST, and a deep convolutional network on CIFAR-10. They specify each model’s objective, architecture, preprocessing, or implementation, without reporting comparative results.
B.3. Independent Components Analysis
- The ICA experiment estimated receptive fields and a Student’s t shape parameter by minimizing negative log likelihood on MNIST digits.The model used a Student’s t-distribution prior.Source: §B.3 · p. 11
- ICA inputs were PCA-whitened, components below 10^-4 times maximum variance were discarded, and the objective was divided by minibatch example count.Source: §B.3 · p. 11
B.4. Ising Model / Hopfield Network via MPF
- The Ising model was trained via MPF, a method described as achieving optimal Hopfield storage capacity (Hillar et al., 2012; Sohl-Dickstein et al., 2011b;a).The experiment used MNIST digits and code from Hillar et al. (2012).Source: §B.4 · p. 11
- For each Ising-model minibatch, the MPF objective was divided by the number of examples and an L2 penalty with coefficient 0.01 was added.Source: §B.4 · p. 11
B.5. Multilayer Perceptron
- The MNIST multilayer perceptron used 784 inputs, two 1200-unit hidden layers, and 10 outputs, with rectified linear or sigmoid units and softmax regression.Its architecture was similar to Hinton et al. (2012); Theano implemented the model and computed gradients.Source: §B.5 · p. 11
B.6. Deep Convolutional Network
- The CIFAR-10 convolutional network used max pooling, rectified linear units, convolutional layers with 48 and 128 units, and a 240-unit fully connected layer.The architecture was loosely based on Goodfellow & Warde-Farley (2013a); Pylearn2 (Goodfellow & Warde-Farley, 2013b) and Theano implemented it.Source: §B.6 · p. 11
C. Implementation Details
The appendix expands on the paper’s implementation details and uses convergence traces to motivate its design choices; even when those choices are changed, convergence remains more rapid than for the competing techniques in Figure 3.
C. Implementation Details
- Figure C.1 shows how changing the design choices affects convergence traces.The figure provides empirical motivation for several design choices.Source: C. Implementation Details · p. 11
- Even with the design choices changed, convergence remains more rapid than for the competing techniques in Figure 3.The passage does not specify the competing techniques or quantify the difference.Source: C. Implementation Details · p. 11
- This section expands on implementation details outlined in Section 3 of the main text.Source: C. Implementation Details · p. 11
C.1. BFGS Initialization … C.4. Growing the Number of Active Subfunctions
The method initializes and stabilizes subfunction Hessian approximations using curvature estimates, selects updates based on distance from the current location, and grows the active set as learning progresses. Its choices are designed to keep curvature estimates useful while making updates and computations tractable.
C.1. BFGS Initialization
- With only one gradient evaluation, a subfunction’s Hessian approximation is initialized to the identity times the median eigenvalue of other active subfunctions’ average Hessians; the first subfunction instead uses 10^6 times identity.Source: §C.1 · p. 11
- BFGS initializes unexplored directions with a scaled identity, setting β to the smallest positive eigenvalue of a minimum-Frobenius-norm matrix consistent with the squared secant equations, computed in a low-dimensional subspace.The initial scaled identity is overwritten for explored directions; its role is to estimate curvature in unexplored directions.Source: §C.1 · p. 11
C.2. Enforcing Positive Definiteness
- SFO enforces positive definiteness by replacing eigenvalues below γλ_max with the median positive eigenvalue; experiments use γ = 10^-8, and the shared low-dimensional representation makes this computation tractable.The median provides a typical-curvature estimate when an eigenvalue is negative or near zero.Source: §C.2 · p. 12
C.3. Choosing a Target Subfunction
- SFO updates the subfunction last evaluated farthest from the current location, measuring distance with either its Hessian approximation or the full-objective Hessian, each chosen with probability 1/2.Using the full-objective metric helps avoid a poor single-subfunction Hessian estimate preventing that subfunction from being evaluated.Source: §C.3 · p. 12
- Unlike cyclic selection and random selection, SFO targets the least accurate approximating functions at the current location; Supplemental Figure C.1 compares update orderings.Source: §C.3 · p. 12
C.4. Growing the Number of Active Subfunctions
- Because subfunction gradients are often nearly identical early in learning, SFO starts with a small active set and adds one subfunction when the average gradient shrinks to within α of its standard error, measured using the inverse approximate Hessian.Source: §C.4 · p. 12
- Experiments use α = 1 and an initial active-set size of two; the set also grows after a bad update or a full pass through the active batch without a batch-size increase.Supplemental Figure C.1 compares this approach with initially activating all subfunctions.Source: §C.4 · p. 12
C.5. Detecting Bad Updates
For some ill-conditioned problems, SFO detects bad updates using a single subfunction rather than a full-objective line search. It rejects detected updates and temporarily shortens the step length.
C.5. Detecting Bad Updates
- A proposed update is labeled bad when the subfunction value has increased since its previous evaluation and exceeds its approximation by more than the corresponding reduction in the summed approximation.This criterion is used because each update evaluates only one subfunction, making a full-objective line search impossible.Source: §C.5 · p. 12
- After a bad proposal, SFO resets the parameters to their previous value and updates the BFGS history matrices with the gradient change in the failed direction.The update step length is also temporarily shortened, then decays back toward 1 with a time constant of one data pass.Source: §C.5 · p. 12
- SFO shortens the step because an inaccurate Hessian approximation for one subfunction often signals inaccuracies for the remaining subfunctions, so failed steps tend to co-occur.Bad-update detection was found necessary for some ill-conditioned problems, including ICA with a Student’s t-prior.Source: §C.5 · p. 12