Fast Adaptation with Linearized Neural Networks

Abstract

The inductive biases of trained neural networks are difficult to understand and, consequently, to adapt to new settings. We study the inductive biases of linearizations of neural networks, which we show to be surprisingly good summaries of the full network functions. Inspired by this finding, we propose a technique for embedding these inductive biases into Gaussian processes through a kernel designed from the Jacobian of the network. In this setting, domain adaptation takes the form of interpretable posterior inference, with accompanying uncertainty estimation. This inference is analytic and free of local optima issues found in standard techniques such as fine-tuning neural network weights to a new task. We develop significant computational speed-ups based on matrix multiplies, including a novel implementation for scalable Fisher vector products. Our experiments on both image classification and regression demonstrate the promise and convenience of this framework for transfer learning, compared to neural network fine-tuning. Code is available at https://github.com/amzn/xfer/tree/master/finite_ntk.

arXiv commentAISTATS 2021

References in corpus (15)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms
• How transferable are features in deep neural networks?
• On Calibration of Modern Neural Networks
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• On Exact Computation with an Infinitely Wide Neural Net
• Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation
• Enhanced Convolutional Neural Tangent Kernels
• Continual Deep Learning by Functional Regularisation of Memorable Past
• Metrics for Probabilistic Geometries
• Adaptive Deep Kernel Learning
• Approximate Inference Turns Deep Networks into Gaussian Processes
• Gradients as Features for Deep Representation Learning
• Improving predictions of Bayesian neural nets via local linearization
• Adversarial Fisher Vectors for Unsupervised Representation Learning