arXiv subjects

Samuel S. Schoenholz

Publications and source records attributed to Samuel S. Schoenholz.

At least 19 recordsLinked to original sources

arXiv2024

Force field optimization by end-to-end differentiable atomistic simulation

The accuracy of atomistic simulations depends on the precision of force fields. Traditional numerical methods often struggle to optimize the empirical force field parameters for reproducing target properties. Recent approaches rely on training these force fields based on forces and energies from first-principle simulations. However, it is unclear whether these approaches will enable capturing complex material responses such as vibrational, or elastic properties. To this extent, we introduce a framework, employing inner loop simulations and outer loop optimization, that exploits automatic differentiation for both property prediction and force-field optimization by computing gradients of the simulation analytically. We demonstrate the approach by optimizing classical Stillinger-Weber and EDIP potentials for silicon systems to reproduce the elastic constants, vibrational density of states, and phonon dispersion. We also demonstrate how a machine-learned potential can be fine-tuned using automatic differentiation to reproduce any target property such as radial distribution functions. Interestingly, the resulting force field exhibits improved accuracy and generalizability to unseen temperatures than those fine-tuned on energies and forces. Finally, we demonstrate the extension of the approach to optimize the force fields towards multiple target properties. Altogether, differentiable simulations, through the analytical computation of their gradients, offer a powerful tool for both theoretical exploration and practical applications toward understanding physical systems and materials.

cond-mat.mtrl-sci↗

arXiv2023

Programmable patchy particles for materials design

Direct design of complex functional materials would revolutionize technologies ranging from printable organs to novel clean energy devices. However, even incremental steps towards designing functional materials have proven challenging. If the material is constructed from highly complex components, the design space of materials properties rapidly becomes too computationally expensive to search. On the other hand, very simple components such as uniform spherical particles are not powerful enough to capture rich functional behavior. Here, we introduce a differentiable materials design model with components that are simple enough to design yet powerful enough to capture complex materials properties: rigid bodies composed of spherical particles with directional interactions (patchy particles). We showcase the method with self-assembly designs ranging from open lattices to self-limiting clusters, all of which are notoriously challenging design goals to achieve using purely isotropic particles. By directly optimizing over the location and interaction of the patches on patchy particles using gradient descent, we dramatically reduce the computation time for finding the optimal building blocks.

cond-mat.soft↗

arXiv2022

What does a deep neural network confidently perceive? The effective dimension of high certainty class manifolds and their low confidence boundaries

Deep neural network classifiers partition input space into high confidence regions for each class. The geometry of these class manifolds (CMs) is widely studied and intimately related to model performance; for example, the margin depends on CM boundaries. We exploit the notions of Gaussian width and Gordon's escape theorem to tractably estimate the effective dimension of CMs and their boundaries through tomographic intersections with random affine subspaces of varying dimension. We show several connections between the dimension of CMs, generalization, and robustness. In particular we investigate how CM dimension depends on 1) the dataset, 2) architecture (including ResNet, WideResNet \& Vision Transformer), 3) initialization, 4) stage of training, 5) class, 6) network width, 7) ensemble size, 8) label randomization, 9) training set size, and 10) robustness to data corruption. Together a picture emerges that higher performing and more robust models have higher dimensional CMs. Moreover, we offer a new perspective on ensembling via intersections of CMs. Our code is at https://github.com/stanislavfort/slice-dice-optimize/

cs.LG↗

arXiv2022

Deep equilibrium networks are sensitive to initialization statistics

Deep equilibrium networks (DEQs) are a promising way to construct models which trade off memory for compute. However, theoretical understanding of these models is still lacking compared to traditional networks, in part because of the repeated application of a single set of weights. We show that DEQs are sensitive to the higher order statistics of the matrix families from which they are initialized. In particular, initializing with orthogonal or symmetric matrices allows for greater stability in training. This gives us a practical prescription for initializations which allow for training with a broader range of initial weight scales.

cs.LG↗

arXiv2022

Fast Finite Width Neural Tangent Kernel

The Neural Tangent Kernel (NTK), defined as $Θ_θ^f(x_1, x_2) = \left[\partial f(θ, x_1)\big/\partial θ\right] \left[\partial f(θ, x_2)\big/\partial θ\right]^T$ where $\left[\partial f(θ, \cdot)\big/\partial θ\right]$ is a neural network (NN) Jacobian, has emerged as a central object of study in deep learning. In the infinite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite width NTK is notoriously expensive to compute, which severely limits its practical utility. We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. Leveraging the structure of neural networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency. Our algorithms can be applied in a black box fashion to any differentiable function, including those implementing neural networks. We open-source our implementations within the Neural Tangents package (arXiv:1912.02803) at https://github.com/google/neural-tangents.

cs.LG↗

arXiv2022

Gradients are Not All You Need

Differentiable programming techniques are widely used in the community and are responsible for the machine learning renaissance of the past several decades. While these methods are powerful, they have limits. In this short report, we discuss a common chaos based failure mode which appears in a variety of differentiable circumstances, ranging from recurrent neural networks and numerical physics simulation to training learned optimizers. We trace this failure to the spectrum of the Jacobian of the system under study, and provide criteria for when a practitioner might expect this failure to spoil their differentiation based optimization algorithms.

cs.LG↗

arXiv2021

dPV: An End-to-End Differentiable Solar-Cell Simulator

We introduce dPV, an end-to-end differentiable photovoltaic (PV) cell simulator based on the drift-diffusion model and Beer-Lambert law for optical absorption. dPV is programmed in Python using JAX, an automatic differentiation (AD) library for scientific computing. Using AD coupled with the implicit function theorem, dPV computes the power conversion efficiency (PCE) of an input PV design as well as the derivative of the PCE with respect to any input parameters, all within comparable time of solving the forward problem. We show an example of perovskite solar-cell optimization and multi-parameter discovery, and compare results with random search and finite differences. The simulator can be integrated with optimization algorithms and neural networks, opening up possibilities for data-efficient optimization and parameter discovery.

physics.comp-ph↗

arXiv2021

Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping

Using an extended and formalized version of the Q/C map analysis of Poole et al. (2016), along with Neural Tangent Kernel theory, we identify the main pathologies present in deep networks that prevent them from training fast and generalizing to unseen data, and show how these can be avoided by carefully controlling the "shape" of the network's initialization-time kernel function. We then develop a method called Deep Kernel Shaping (DKS), which accomplishes this using a combination of precise parameter initialization, activation function transformations, and small architectural tweaks, all of which preserve the model class. In our experiments we show that DKS enables SGD training of residual networks without normalization layers on Imagenet and CIFAR-10 classification tasks at speeds comparable to standard ResNetV2 and Wide-ResNet models, with only a small decrease in generalization performance. And when using K-FAC as the optimizer, we achieve similar results for networks without skip connections. Our results apply for a large variety of activation functions, including those which traditionally perform very badly, such as the logistic sigmoid. In addition to DKS, we contribute a detailed analysis of skip connections, normalization layers, special activation functions like RELU and SELU, and various initialization schemes, explaining their effectiveness as alternative (and ultimately incomplete) ways of "shaping" the network's initialization-time kernel.

cs.LG↗

arXiv2021

Whitening and second order optimization both make information in the dataset unusable during training, and can reduce or prevent generalization

Machine learning is predicated on the concept of generalization: a model achieving low error on a sufficiently large training set should also perform well on novel samples from the same distribution. We show that both data whitening and second order optimization can harm or entirely prevent generalization. In general, model training harnesses information contained in the sample-sample second moment matrix of a dataset. For a general class of models, namely models with a fully connected first layer, we prove that the information contained in this matrix is the only information which can be used to generalize. Models trained using whitened data, or with certain second order optimization schemes, have less access to this information, resulting in reduced or nonexistent generalization ability. We experimentally verify these predictions for several architectures, and further demonstrate that generalization continues to be harmed even when theoretical requirements are relaxed. However, we also show experimentally that regularized second order optimization can provide a practical tradeoff, where training is accelerated but less information is lost, and generalization can in some circumstances even improve.

cs.LG↗

arXiv2021

Tilting the playing field: Dynamical loss functions for machine learning

We show that learning can be improved by using loss functions that evolve cyclically during training to emphasize one class at a time. In underparameterized networks, such dynamical loss functions can lead to successful training for networks that fail to find a deep minima of the standard cross-entropy loss. In overparameterized networks, dynamical loss functions can lead to better generalization. Improvement arises from the interplay of the changing loss landscape with the dynamics of the system as it evolves to minimize the loss. In particular, as the loss function oscillates, instabilities develop in the form of bifurcation cascades, which we study using the Hessian and Neural Tangent Kernel. Valleys in the landscape widen and deepen, and then narrow and rise as the loss landscape changes during a cycle. As the landscape narrows, the learning rate becomes too large and the network becomes unstable and bounces around the valley. This process ultimately pushes the system into deeper and wider regions of the loss landscape and is characterized by decreasing eigenvalues of the Hessian. This results in better regularized models with improved generalization performance.

cs.LG↗

arXiv2020

JAX, M.D.: A Framework for Differentiable Physics

We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional code. Since the simulations themselves are differentiable functions, entire trajectories can be differentiated to perform meta-optimization. These features are built on primitive operations, such as spatial partitioning, that allow simulations to scale to hundreds-of-thousands of particles on a single GPU. These primitives are flexible enough that they can be used to scale up workloads outside of molecular dynamics. We present several examples that highlight the features of JAX MD including: integration of graph neural networks into traditional simulations, meta-optimization through minimization of particle packings, and a multi-agent flocking simulation. JAX MD is available at www.github.com/google/jax-md.

physics.comp-ph↗

arXiv2020

Self-assembling kinetics: Accessing a new design space via differentiable statistical-physics models

The inverse problem of designing component interactions to target emergent structure is fundamental to numerous applications in biotechnology, materials science, and statistical physics. Equally important is the inverse problem of designing emergent kinetics, but this has received considerably less attention. Using recent advances in automatic differentiation, we show how kinetic pathways can be precisely designed by directly differentiating through statistical-physics models, namely free energy calculations and molecular dynamics simulations. We consider two systems that are crucial to our understanding of structural self-assembly: bulk crystallization and small nanoclusters. In each case we are able to assemble precise dynamical features. Using gradient information, we manipulate interactions among constituent particles to tune the rate at which these systems yield specific structures of interest. Moreover, we use this approach to learn non-trivial features about the high-dimensional design space, allowing us to accurately predict when multiple kinetic features can be simultaneously and independently controlled. These results provide a concrete and generalizable foundation for studying non-structural self-assembly, including kinetic properties as well as other complex emergent properties, in a vast array of systems.

physics.comp-ph↗

arXiv2020

Finite Versus Infinite Neural Networks: an Empirical Study

We perform a careful, thorough, and large scale empirical study of the correspondence between wide neural networks and kernel methods. By doing so, we resolve a variety of open questions related to the study of infinitely wide neural networks. Our experimental results include: kernel methods outperform fully-connected finite-width networks, but underperform convolutional finite width networks; neural network Gaussian process (NNGP) kernels frequently outperform neural tangent (NT) kernels; centered and ensembled finite networks have reduced posterior variance and behave more similarly to infinite networks; weight decay and the use of a large learning rate break the correspondence between finite and infinite networks; the NTK parameterization outperforms the standard parameterization for finite width networks; diagonal regularization of kernels acts similarly to early stopping; floating point precision limits kernel performance beyond a critical dataset size; regularized ZCA whitening improves accuracy; finite network performance depends non-monotonically on width in ways not captured by double descent phenomena; equivariance of CNNs is only beneficial for narrow networks far from the kernel regime. Our experiments additionally motivate an improved layer-wise scaling for weight decay which improves generalization in finite-width networks. Finally, we develop improved best practices for using NNGP and NT kernels for prediction, including a novel ensembling technique. Using these best practices we achieve state-of-the-art results on CIFAR-10 classification for kernels corresponding to each architecture class we consider.

cs.LG↗

arXiv2020

Unifying framework for strong and fragile liquids via machine learning: a study of liquid silica

The fragility of a glassforming liquid characterizes how rapidly its relaxation dynamics slow down with cooling. The viscosity of strong liquids follows an Arrhenius law with a temperature-independent barrier height to rearrangements responsible for relaxation, whereas fragile liquids experience a much faster increase in their dynamics, suggesting a barrier height that increases with decreasing temperature. Strong glassformers are typically network glasses, while fragile glassformers are typically molecular or hard-sphere-like. As a result of these differences at the microscopic level, strong and fragile glassformers are usually treated separately from a theoretical point of view. Silica is the archetypal strong glassformer at low temperatures, but also exhibits a mysterious strong-to-fragile crossover at higher temperatures. Here we show that softness, a structure-based machine learned parameter that has previously been applied to fragile glassformers provides a useful description of model liquid silica in the strong and fragile regimes, and through the strong-to-fragile crossover. Just as for fragile glassformers, the relationship between softness and dynamics is invariant and Arrhenius in all regimes, but the average softness changes with temperature. The strong-to-fragile crossover in silica is not due to a sudden, qualitative change in structure, but can be explained by a simple Arrhenius form with a continuously and linearly changing local structure. Our results unify the study of liquid silica under a single simple conceptual picture.

cond-mat.soft↗

arXiv2020

Disentangling Trainability and Generalization in Deep Neural Networks

A longstanding goal in the theory of deep learning is to characterize the conditions under which a given neural network architecture will be trainable, and if so, how well it might generalize to unseen data. In this work, we provide such a characterization in the limit of very wide and very deep networks, for which the analysis simplifies considerably. For wide networks, the trajectory under gradient descent is governed by the Neural Tangent Kernel (NTK), and for deep networks the NTK itself maintains only weak data dependence. By analyzing the spectrum of the NTK, we formulate necessary conditions for trainability and generalization across a range of architectures, including Fully Connected Networks (FCNs) and Convolutional Neural Networks (CNNs). We identify large regions of hyperparameter space for which networks can memorize the training set but completely fail to generalize. We find that CNNs without global average pooling behave almost identically to FCNs, but that CNNs with pooling have markedly different and often better generalization performance. These theoretical results are corroborated experimentally on CIFAR10 for a variety of network architectures and we include a colab notebook that reproduces the essential results of the paper.

cs.LG↗

arXiv2020

On the infinite width limit of neural networks with a standard parameterization

There are currently two parameterizations used to derive fixed kernels corresponding to infinite width neural networks, the NTK (Neural Tangent Kernel) parameterization and the naive standard parameterization. However, the extrapolation of both of these parameterizations to infinite width is problematic. The standard parameterization leads to a divergent neural tangent kernel while the NTK parameterization fails to capture crucial aspects of finite width networks such as: the dependence of training dynamics on relative layer widths, the relative training dynamics of weights and biases, and overall learning rate scale. Here we propose an improved extrapolation of the standard parameterization that preserves all of these properties as width is taken to infinity and yields a well-defined neural tangent kernel. We show experimentally that the resulting kernels typically achieve similar accuracy to those resulting from an NTK parameterization, but with better correspondence to the parameterization of typical finite width networks. Additionally, with careful tuning of width parameters, the improved standard parameterization kernels can outperform those stemming from an NTK parameterization. We release code implementing this improved standard parameterization as part of the Neural Tangents library at https://github.com/google/neural-tangents.

cs.LG↗

arXiv2019

Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent

A longstanding goal in deep learning research has been to precisely characterize training and generalization. However, the often complex loss landscapes of neural networks have made a theory of learning dynamics elusive. In this work, we show that for wide neural networks the learning dynamics simplify considerably and that, in the infinite width limit, they are governed by a linear model obtained from the first-order Taylor expansion of the network around its initial parameters. Furthermore, mirroring the correspondence between wide Bayesian neural networks and Gaussian processes, gradient-based training of wide neural networks with a squared loss produces test set predictions drawn from a Gaussian process with a particular compositional kernel. While these theoretical results are only exact in the infinite width limit, we nevertheless find excellent empirical agreement between the predictions of the original network and those of the linearized version even for finite practically-sized networks. This agreement is robust across different architectures, optimization methods, and loss functions.

stat.ML↗

arXiv2019

Neural Tangents: Fast and Easy Infinite Neural Networks in Python

Neural Tangents is a library designed to enable research into infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finite-width as usual or in their infinite-width limit. Infinite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks in either function space or weight space. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. Neural Tangents is available at www.github.com/google/neural-tangents. We also provide an accompanying interactive Colab notebook.

stat.ML↗