PaperlayerView exact revision on arXiv

Source-linked AI summary

Ensemble Kalman Inversion: A Derivative-Free Technique For Machine Learning Tasks

Nikola B. Kovachki, Andrew M. Stuart

arXiv:1808.03620v1cs.LGmath.OCstat.ML

TL;DR

The paper addresses limitations of stochastic gradient methods, especially their limited parallelization, by formulating machine-learning tasks as inverse problems solved with Ensemble Kalman Inversion. Experiments report competitive or superior performance across several supervised and semi-supervised settings, alongside important scope and implementation caveats.

  • Problem

    Derivative-free, parallelizable alternatives to stochastic gradient descent are needed because parallelizing SGD has had limited success and existing methods face challenges.

  • Method

    The paper applies Ensemble Kalman Inversion within a unified Bayesian inverse-problem framework, using ensemble dynamics and adaptive steps to solve machine-learning tasks.

  • Results

    Across experiments, EKI matches SGD on large feed-forward classification problems, outperforms SGD on recurrent classification and online regression, and is faster but less accurate than MCMC for the tested semi-supervised problem.

  • Takeaways & Limitations

    EKI shows wide applicability and competitiveness with state-of-the-art methods, while offering computational advantages for some large-scale semi-supervised problems.

  • Takeaways & Limitations

    The softmax output layer can only asymptotically match the data, although its combination with cross-entropy produced the best numerical results among tested thresholding/loss pairs.

Abstract

from arXiv · show

The standard probabilistic perspective on machine learning gives rise to empirical risk-minimization tasks that are frequently solved by stochastic gradient descent (SGD) and variants thereof. We present a formulation of these tasks as classical inverse or filtering problems and, furthermore, we propose an efficient, gradient-free algorithm for finding a solution to these problems using ensemble Kalman inversion (EKI). Applications of our approach include offline and online supervised learning with deep neural networks, as well as graph-based semi-supervised learning. The essence of the EKI procedure is an ensemble based approximate gradient descent in which derivatives are replaced by differences from within the ensemble. We suggest several modifications to the basic method, derived from empirically successful heuristics developed in the context of SGD. Numerical results demonstrate wide applicability and robustness of the proposed algorithm.

1. Introduction

The paper frames machine-learning tasks as inverse problems and investigates ensemble Kalman methods as derivative-free, parallelizable alternatives to gradient-based optimization. It evaluates these methods across supervised, semi-supervised, and online learning settings.

  • The Setting: SGD is widely used but faces challenges and has seen limited success when parallelized, motivating derivative-free parallelizable methods.
  • The Setting: Earlier derivative-free approaches were demonstrated mainly on shallow, dense neural networks with relatively few parameters.
  • The Setting: The paper uses ensemble Kalman methods to address derivative-free and parallelizable optimization for machine-learning tasks.
  • Our Contribution: Many machine-learning tasks are formulated within a unified Bayesian inverse-problem framework.
  • Our Contribution: The paper applies modified ensemble methods across supervised, semi-supervised, and online learning with deep dense, convolutional, and recurrent neural networks.
  • Our Contribution: The paper presents its work as a transfer of inverse-problem methodologies to machine learning, while leaving parallelization and further research directions open.

2. Problem Formulation

The paper formulates learning as empirical risk minimization from finite data, then recasts supervised, semi-supervised, and online tasks as parameterized inverse problems. Regularization and priors address ill-posedness, while the formulations accommodate labeled, unlabeled, and sequential data.

  • 2.1. Classical Framework: The classical formulation minimizes expected loss over mappings, but finite datasets provide only indirect access to the underlying data distribution.
  • 2.1. Classical Framework: Common losses include squared error for regression and cross-entropy for classification, with classification outputs restricted to probability vectors.
  • 2.1. Classical Framework: Parametric models replace the unrestricted mapping class with a computable family indexed by parameters, while regularization helps make minimization unambiguous.
  • 2.1.1. Supervised Learning: Supervised learning minimizes empirical loss over paired input-output data, with parameter regularization used to prevent overfitting or address ill-posedness.
  • 2.1.2. Semi-Supervised Learning: Semi-supervised learning assigns labels across all data while using only a small labeled subset and incorporating unlabeled data through regularization.
  • 2.1.2. Semi-Supervised Learning: The paper studies graph-based semi-supervised regularization, including a graph Laplacian construction with tunable parameters τ and α.
  • 2.1.3. Online Learning: Online learning processes data sequentially and updates parameters using current observations and previous estimates, reducing computational cost or matching sequential acquisition.
  • 2.2. Inverse Problems: The inverse-problem formulation introduces a forward map for concatenated observations, a noise or model-error distribution, and a prior measure on parameters.

3. Approximation Architectures

The paper represents neural networks as parameterized compositions of maps, with architecture choices specifying how dimensions, nonlinearities, convolutions, channels, and pooling transform inputs. These constructions cover dense, convolutional, and recurrent architectures used in the stated learning tasks.

  • Feed-forward networks: Neural networks are parameterized compositions whose layers can change input dimensions, followed by an affine map and problem-dependent output projection.The parameter vector concatenates layer parameters, while the final projection can be identity, softmax, or another response function.
  • Dense networks: Each dense layer applies an affine transformation followed by a point-wise activation, with weights and biases forming its parameters.Dense layers have no global sparsity assumptions, so their weight matrices can be fully populated.
  • Dense networks: Activation functions are layerwise design choices, including sigmoid, hyperbolic tangent, and ReLU; ReLU is numerically successful despite violating boundedness assumptions.The paper notes that ReLU is unbounded and therefore falls outside the classical universal approximation theorem assumptions cited there.
  • Convolutional networks: Convolutional layers compute multiple channel-wise convolutions, with channel counts, stencil sizes, input dimensions, and operator design determining the image-space dimension.The convolution operator is a modeling choice governing which input blocks are combined at each spatial location.
  • Convolutional networks: Pooling applies statistics such as maxima to convolution outputs, reduces dimension, and is commonly combined with convolutional and dense layers.Kernel and stride are design choices; global pooling produces one scalar per channel when the kernel spans the relevant spatial grid.

4. Algorithms

Section 4 introduces the loss function, gradient-based optimization, EKI, and subsequent modifications as the paper’s algorithmic framework.

  • Algorithms: The algorithms section proceeds from loss-function choices to gradient-based optimization, then defines EKI and presents modifications to the basic method.The section structure covers optimization algorithms and convergence heuristics before EKI variants.

4.1. Loss Function

The paper focuses on squared-error loss for regression and cross-entropy loss for classification, matching the latter with softmax outputs. Cross-entropy has useful probabilistic meaning but important boundary limitations.

  • Loss choices: Squared-error loss is the standard regression choice, while cross-entropy is the standard classification choice considered in the paper.Squared error is associated with an additive Gaussian noise model; cross-entropy is treated separately for classification.
  • Classification loss: Cross-entropy is defined on probability-vector outputs and is consistent with using softmax as the network’s projection map.The softmax image is the probability simplex, matching the domain on which the classification loss is well-defined.
  • Classification loss: The softmax output has the drawback that perfectly labeled data may only be approached asymptotically, because such labels lie on the closure rather than generally inside the probability simplex.The paper nevertheless reports softmax with cross-entropy as producing the best numerical results among the thresholding/loss pairs considered.
  • Classification loss: Cross-entropy measures the Shannon cross-entropy between the model’s categorical distribution and the data distribution.The paper notes that this loss cannot be represented purely as a function of the residual y − y′ and therefore lacks an associated additive-noise MAP model.

4.2. Gradient Based Optimization

The paper presents SGD and related gradient methods as the prevailing optimization approach, while emphasizing initialization, normalization, stochastic mini-batching, and momentum as practical design issues. These methods rely on gradients and face challenges in parallelization and training setup.

  • Gradient descent: Gradient descent updates parameters in the negative gradient direction of a regularized loss and is expected to converge to a local minimum for almost all initial data.The update uses a step size h_k and produces bounded trajectories under the stated conditions.
  • Stochastic gradient descent: Practical SGD discretizes the dynamics with a learning rate and estimates the full gradient from randomly selected mini-batches.A new subset is drawn each iteration, and one complete pass through the dataset is called an epoch.
  • Stochastic gradient descent: Mini-batch noise has been suggested to guide descent toward parameter regions that generalize better than standard descent methods.The paper reports typical training durations of 10 to 500 epochs, depending on the model and data.
  • Momentum methods: Momentum-augmented methods are used to accelerate convergence, but the machine-learning discretization called Nesterov momentum differs from the corresponding continuous-time dynamic.The paper characterizes the practitioner algorithm as a discretization of a rescaled gradient flow.
  • Gradient-based variants: Adaptive methods such as Adam, RMSProp, and Adagrad have emerged from gradient-descent variants, with no consensus on which performs best.The paper also describes online gradient descent as the standard gradient-based approach for online learning.
  • Training challenges: Initialization and normalization are identified as major challenges because they respectively affect starting parameters and the output distributions of network layers.Batch normalization addresses changing output distributions during training but relies on data batches and therefore cannot be used in the online setting.
  • Training challenges: Batch normalization introduces scale and shift parameters after layer outputs and estimates normalization statistics during training.The parameters γ and β are learned through SGD, with separate normalization performed per channel for multi-channel outputs.

4.3. Ensemble Kalman Inversion

EKI reformulates machine-learning training as an inverse-problem optimization method that approximates gradient descent with ensemble differences rather than explicit derivatives. The section develops modifications for general losses, momentum, adaptive discretization, and practical training heuristics.

  • Core EKI method: EKI evolves an ensemble of parameter estimates through data interaction, using empirical cross-covariances and ensemble differences as approximate derivatives.The method is analyzed as gradient descent projected onto the ensemble-generated subspace, without explicitly computing model gradients.
  • Core EKI method: For nonlinear forward maps, linearization shows that EKI behaves like gradient descent with a projection onto the ensemble subspace.This interpretation relies on particles being sufficiently close for local linearization.
  • Loss functions: A modified EKI dynamic replaces the usual squared-error objective with a loss-gradient formulation, enabling losses such as cross-entropy while remaining derivative-free in the model parameters.The loss is differentiated with respect to the auxiliary output variable rather than the model parameter u; direct regularization is left for future work.
  • Optimization modifications: Momentum EKI adds Nesterov-style momentum and can accelerate convergence, but it doubles memory cost because additional state must be retained.The practical scheme fixes λ ∈ (0, 1) and uses a discretization motivated by standard machine-learning optimization.
  • Discrete scheme: Adaptive time-stepping sets hk inversely proportional to ∥D(uk)∥F + ϵ, allowing larger stable steps when particle interactions or data alignment are small.The Frobenius-norm scaling is intended to maximize h0 without losing numerical stability and to avoid floating-point issues.
  • Practical heuristics: Practical variants borrow mini-batching, ensemble initialization from a prior, and mean-particle prediction from stochastic-gradient methods.The mean particle is selected because its performance is reported to be comparable to that of individual particles.

5. Numerical Experiments

The experiments evaluate EKI across supervised, semi-supervised, and online learning tasks using image, voting-record, and time-series data. EKI matches or exceeds SGD in several settings, while trading computational speed for parallelizable forward evaluations.

  • Supervised learning: EKI matches SGD on feed-forward classification with up to half a million parameters and ensembles two orders of magnitude smaller.It also shows less test-accuracy oscillation and requires less hyper-parameter tuning, but is significantly slower because each iteration requires many forward passes.
  • Supervised learning: EKI significantly outperforms SGD on supervised classification with recurrent neural networks.The authors suggest that EKI may avoid steep barriers on recurrent-network loss surfaces through its noisy Jacobian estimates.
  • Semi-supervised learning: EKI is worse than MCMC but better than the naive solution on the considered semi-supervised learning problem.Even with many ensemble members, EKI is much faster and computationally cheaper than MCMC.
  • Online learning: EKI converges significantly faster and to a better solution than SGD for online recurrent-network regression with O(1) ensemble members.The experiments use simple univariate time-series, so domain-specific applications remain future work.
  • Experimental setup: The benchmark covers three classification datasets and two regression datasets, evaluating generalization with held-out test portions.Classification uses test accuracy, while regression uses average squared ℓ2 prediction error.
  • Experimental setup: The datasets comprise MNIST and SVHN images, U.S. House voting records, Melbourne temperatures, and Zürich sunspots.Voting records support semi-supervised learning, while the two time-series datasets support one-step-ahead regression.

5.3. Implementation Details

The implementation standardizes the compared methods around common initialization and mini-batching, while distinguishing EKI variants by their discretization and randomization choices.

  • Common setup: All methods use the same model-based prior initialization and mini-batching.This provides a common implementation setup across the compared EKI techniques.
  • Method variants: The forward Euler discretization is called EKI, while the momentum discretization is called MEKI.The paper assigns separate names to methods using different techniques so they remain distinguishable.

5.4. Supervised Learning

Across dense, convolutional, and recurrent supervised-learning tasks, EKI variants are often competitive with SGD, with randomization becoming important for larger networks and EKI(R) showing smooth or reliable training behavior.

  • Dense Neural Networks: On dense MNIST networks, EKI methods are comparable to SGD for relatively small models, while larger networks require randomization, especially with smaller ensembles.Momentum SGD generally performs best, but EKI(R) trails closely.
  • Dense Neural Networks: The mean particle achieves performance comparable to individual particles on DNN-(1,2), motivating its use as the final parameter estimate.The authors note that a weighted ensemble might achieve better results and potentially help avoid adversarial examples.
  • Dense Neural Networks: With J = 6000 on DNN 2, EKI becomes comparable to SGD, while EKI(R) performs almost identically with 2,000 and 6,000 ensemble members.The experiments identify EKI(R) as the preferred method for subsequent tasks.
  • Convolutional Neural Networks: On CNN-MNIST, EKI(R) avoids the early descent-direction difficulty observed for SGD and reaches almost perfect classification within 30 epochs.The network uses ReLU without batch normalization, while EKI(R) produces a smooth test-accuracy curve.
  • Recurrent Neural Networks: For the MNIST recurrent network, EKI(R) performs significantly better than SGD and appears more reliable overall.The ensemble expands from J = 1000 to J = 4000 during training.

5.5. Semi-supervised Learning

For graph-based semi-supervised classification on Voting Records, vanilla EKI performs only slightly better than the Fiedler-vector baseline because its particles remain near the prior-defined subspace.

  • Voting Records: Vanilla EKI uses J = 1000 prior-drawn ensemble members and performs only slightly better than the Fiedler-vector classifier.The particles quickly collapse toward a solution close to the Fiedler vector.
  • Voting Records: The prior covariance restricts EKI to the subspace orthogonal to the first eigenvector of the graph Laplacian, keeping solutions close to the Fiedler-vector direction.The authors attribute the behavior to the i.i.d. prior ensemble and note that MCMC can explore outside this subspace.

5.6. Online Learning

In online recurrent-network learning on two univariate time series, EKI converges faster than OGD, reaches a slightly better solution, and captures small-scale oscillations more effectively.

  • Online Recurrent Learning: EKI converges more quickly and to a slightly better solution than OGD on both the Melbourne temperatures and Zürich sunspots tasks.Both tasks provide one data-label pair at a time to a single-layer recurrent network.
  • Online Recurrent Learning: The EKI-learned model captures small-scale oscillations better in both online time-series problems.The authors describe these results as promising for harder recurrent-network problems.

6. Conclusion and Future Directions

The paper frames many machine-learning tasks as Bayesian inverse problems and finds modified EKI methods broadly applicable and competitive, while identifying several directions for extending theory, scale, and task coverage.

  • Conclusions: The authors report wide applicability and competitiveness of modified Ensemble Kalman Inversion methods across machine-learning tasks.Their framework treats many machine-learning problems as Bayesian inverse problems.
  • Future Directions: Future work includes theoretical analysis of momentum and general-loss EKI methods and their possible use for physical inverse problems.
  • Future Directions: GPU parallelization is proposed for applying EKI methods to large-scale machine-learning tasks.
  • Future Directions: The authors propose applying EKI to more difficult recurrent-network and reinforcement-learning problems.
  • Future Directions: Using the entire particle ensemble is proposed as a way to improve accuracy and possibly combat adversarial examples.
Loading 1808.03620v1…