Wu Lin

ICML Conference 2024 Conference Paper

Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective

• Wu Lin
• Felix Dangel
• Runa Eschenhagen
• Juhan Bae
• Richard E. Turner
• Alireza Makhzani

Adaptive gradient optimizers like Adam(W) are the default training algorithms for many deep learning architectures, such as transformers. Their diagonal preconditioner is based on the gradient outer product which is incorporated into the parameter update via a square root. While these methods are often motivated as approximate second-order methods, the square root represents a fundamental difference. In this work, we investigate how the behavior of adaptive methods changes when we remove the root, i. e. strengthen their second-order motivation. Surprisingly, we find that such square-root-free adaptive methods close the generalization gap to SGD on convolutional architectures, while maintaining their root-based counterpart’s performance on transformers. The second-order perspective also has practical benefits for the development of non-diagonal adaptive methods through the concept of preconditioner invariance. In contrast to root-based methods like Shampoo, the root-free counterparts do not require numerically unstable matrix root decompositions and inversions, thus work well in half precision. Our findings provide new insights into the development of adaptive methods and raise important questions regarding the currently overlooked role of adaptivity for their success.

ICML Conference 2024 Conference Paper

Structured Inverse-Free Natural Gradient Descent: Memory-Efficient & Numerically-Stable KFAC

• Wu Lin
• Felix Dangel
• Runa Eschenhagen
• Kirill Neklyudov
• Agustinus Kristiadi
• Richard E. Turner
• Alireza Makhzani

Second-order methods such as KFAC can be useful for neural net training. However, they are often memory-inefficient since their preconditioning Kronecker factors are dense, and numerically unstable in low precision as they require matrix inversion or decomposition. These limitations render such methods unpopular for modern mixed-precision training. We address them by (i) formulating an inverse-free KFAC update and (ii) imposing structures in the Kronecker factors, resulting in structured inverse-free natural gradient descent (SINGD). On modern neural networks, we show that SINGD is memory-efficient and numerically robust, in contrast to KFAC, and often outperforms AdamW even in half precision. Our work closes a gap between first- and second-order methods in modern low-precision training.

NeurIPS Conference 2024 Conference Paper

Training Data Attribution via Approximate Unrolling

• Juhan Bae
• Wu Lin
• Jonathan Lorraine
• Roger Grosse

Many training data attribution (TDA) methods aim to estimate how a model's behavior would change if one or more data points were removed from the training set. Methods based on implicit differentiation, such as influence functions, can be made computationally efficient, but fail to account for underspecification, the implicit bias of the optimization algorithm, or multi-stage training pipelines. By contrast, methods based on unrolling address these issues but face scalability challenges. In this work, we connect the implicit-differentiation-based and unrolling-based approaches and combine their benefits by introducing Source, an approximate unrolling-based TDA method that is computed using an influence-function-like formula. While being computationally efficient compared to unrolling-based approaches, Source is suitable in cases where implicit-differentiation-based approaches struggle, such as in non-converged models and multi-stage training pipelines. Empirically, Source outperforms existing TDA techniques in counterfactual prediction, especially in settings where implicit-differentiation-based approaches fall short.

EAAI Journal 2023 Journal Article

A Kriging model-based evolutionary algorithm with support vector machine for dynamic multimodal optimization

• Xunfeng Wu
• Qiuzhen Lin
• Wu Lin
• Yulong Ye
• Qingling Zhu
• Victor C.M. Leung

ICML Conference 2023 Conference Paper

Simplifying Momentum-based Positive-definite Submanifold Optimization with Applications to Deep Learning

• Wu Lin
• Valentin Duruisseaux
• Melvin Leok
• Frank Nielsen
• Mohammad Emtiyaz Khan
• Mark Schmidt 0001

Riemannian submanifold optimization with momentum is computationally challenging because, to ensure that the iterates remain on the submanifold, we often need to solve difficult differential equations. Here, we simplify such difficulties for a class of structured symmetric positive-definite matrices with the affine-invariant metric. We do so by proposing a generalized version of the Riemannian normal coordinates that dynamically orthonormalizes the metric and locally converts the problem into an unconstrained problem in the Euclidean space. We use our approach to simplify existing approaches for structured covariances and develop matrix-inverse-free $2^\text{nd}$-order optimizers for deep learning in low precision settings.

ICML Conference 2021 Conference Paper

Tractable structured natural-gradient descent using local parameterizations

• Wu Lin
• Frank Nielsen
• Mohammad Emtiyaz Khan
• Mark Schmidt 0001

Natural-gradient descent (NGD) on structured parameter spaces (e. g. , low-rank covariances) is computationally challenging due to difficult Fisher-matrix computations. We address this issue by using \emph{local-parameter coordinates} to obtain a flexible and efficient NGD method that works well for a wide-variety of structured parameterizations. We show four applications where our method (1) generalizes the exponential natural evolutionary strategy, (2) recovers existing Newton-like algorithms, (3) yields new structured second-order algorithms, and (4) gives new algorithms to learn covariances of Gaussian and Wishart-based distributions. We show results on a range of problems from deep learning, variational inference, and evolution strategies. Our work opens a new direction for scalable structured geometric methods.

ICML Conference 2020 Conference Paper

Handling the Positive-Definite Constraint in the Bayesian Learning Rule

• Wu Lin
• Mark Schmidt 0001
• Mohammad Emtiyaz Khan

The Bayesian learning rule is a natural-gradient variational inference method, which not only contains many existing learning algorithms as special cases but also enables the design of new algorithms. Unfortunately, when variational parameters lie in an open constraint set, the rule may not satisfy the constraint and requires line-searches which could slow down the algorithm. In this work, we address this issue for positive-definite constraints by proposing an improved rule that naturally handles the constraints. Our modification is obtained by using Riemannian gradient methods, and is valid when the approximation attains a block-coordinate natural parameterization (e. g. , Gaussian distributions and their mixtures). Our method outperforms existing methods without any significant increase in computation. Our work makes it easier to apply the rule in the presence of positive-definite constraints in parameter spaces.

ICML Conference 2019 Conference Paper

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

• Wu Lin
• Mohammad Emtiyaz Khan
• Mark Schmidt 0001

Natural-gradient methods enable fast and simple algorithms for variational inference, but due to computational difficulties, their use is mostly limited to minimal exponential-family (EF) approximations. In this paper, we extend their application to estimate structured approximations such as mixtures of EF distributions. Such approximations can fit complex, multimodal posterior distributions and are generally more accurate than unimodal EF approximations. By using a minimal conditional-EF representation of such approximations, we derive simple natural-gradient updates. Our empirical results demonstrate a faster convergence of our natural-gradient method compared to black-box gradient-based methods. Our work expands the scope of natural gradients for Bayesian inference and makes them more widely applicable than before.

ICML Conference 2018 Conference Paper

Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam

• Mohammad Emtiyaz Khan
• Didrik Nielsen
• Voot Tangkaratt
• Wu Lin
• Yarin Gal
• Akash Srivastava

Uncertainty computation in deep learning is essential to design robust and reliable systems. Variational inference (VI) is a promising approach for such computation, but requires more effort to implement and execute compared to maximum-likelihood methods. In this paper, we propose new natural-gradient algorithms to reduce such efforts for Gaussian mean-field VI. Our algorithms can be implemented within the Adam optimizer by perturbing the network weights during gradient evaluations, and uncertainty estimates can be cheaply obtained by using the vector that adapts the learning rate. This requires lower memory, computation, and implementation effort than existing VI methods, while obtaining uncertainty estimates of comparable quality. Our empirical results confirm this and further suggest that the weight-perturbation in our algorithm could be useful for exploration in reinforcement learning and stochastic optimization.

UAI Conference 2016 Conference Paper

Faster Stochastic Variational Inference using Proximal-Gradient Methods with General Divergence Functions

• Mohammad Emtiyaz Khan
• Reza Babanezhad 0001
• Wu Lin
• Mark Schmidt 0001
• Masashi Sugiyama

Several recent works have explored stochastic gradient methods for variational inference that exploit the geometry of the variational-parameter space. However, the theoretical properties of these methods are not well-understood and these methods typically only apply to conditionallyconjugate models. We present a new stochastic method for variational inference which exploits the geometry of the variational-parameter space and also yields simple closed-form updates even for non-conjugate models. We also give a convergence-rate analysis of our method and many other previous methods which exploit the geometry of the space. Our analysis generalizes existing convergence results for stochastic mirror-descent on non-convex objectives by using a more general class of divergence functions. Beyond giving a theoretical justification for a variety of recent methods, our experiments show that new algorithms derived in this framework lead to state of the art results on a variety of problems. Further, due to its generality, we expect that our theoretical analysis could also apply to other applications. 1 Reza Babanezhad University of British Columbia Vancouver, Canada babanezhad@gmail. com