Universal Neural Functionals

TL;DR

This paper introduces universal neural functionals (UNFs), algorithms that automatically create permutation equivariant models for any neural network weight space, improving learned optimizers for various architectures.

Contribution

It presents a novel algorithm to construct permutation equivariant models applicable to any neural network architecture's weight space.

Findings

UNFs improve optimization of small image classifiers.

UNFs enhance language model training.

Open-source library available for constructing UNFs.

Abstract

A challenging problem in many modern machine learning tasks is to process weight-space features, i.e., to transform or extract information from the weights and gradients of a neural network. Recent works have developed promising weight-space models that are equivariant to the permutation symmetries of simple feedforward networks. However, they are not applicable to general architectures, since the permutation symmetries of a weight space can be complicated by recurrence or residual connections. This work proposes an algorithm that automatically constructs permutation equivariant models, which we refer to as universal neural functionals (UNFs), for any weight space. Among other applications, we demonstrate how UNFs can be substituted into existing learned optimizer designs, and find promising improvements over prior methods when optimizing small image classifiers and language models. Our results suggest that learned optimizers can benefit from considering the (symmetry) structure of the weight space they optimize. We open-source our library for constructing UNFs at https://github.com/AllanYangZhou/universal_neural_functional.