Function-space Parameterization of Neural Networks for Sequential Learning

Code accompanying ICLR 2024 submission Function-space Parameterization of Neural Networks for Sequential Learning. This repository contains code for reproducing the experiments in the ICLR 2024 paper. Please see this repo for a clean and minimal implementation of Sparse Function-space Representation of Neural Networks (SFR). We recommend using the clean and minimal repo.

Install

Install using virtual environment

Make a virtual environment: sh python -m venv .venv Activate it with: sh source .venv/bin/activate Install the dependencies with: sh python -m pip install --upgrade pip pip install laplace-torch==0.1a2 pip install -e ".[experiments]" We install laplace-torch separately due to version conflicts with backpacpk-for-pytorch. Note that laplace-torch is only used for running the baselines.

Install using pip

Alternatively, manually install the dependencies with: sh pip install laplace-torch==0.1a2 pip install -r requirements.txt

Reproducing experiments

See experiments for details on how to reproduce the results in the paper. This includes code for generating the tables and figures.

Useage

See the notebooks/README.md for how to use our code for both regression and classification.

Example

Here's a short example: ```python import src import torch

torch.setdefaultdtype(torch.float64)

def func(x, noise=True): return torch.sin(x * 5) / x + torch.cos(x * 10)

Toy data set

Xtrain = torch.rand((100, 1)) * 2 Ytrain = func(Xtrain, noise=True) data = (Xtrain, Y_train)

Training config

width = 64 numepochs = 1000 batchsize = 16 learningrate = 1e-3 delta = 0.00005 # prior precision dataloader = torch.utils.data.DataLoader( torch.utils.data.TensorDataset(*data), batchsize=batchsize )

Create a neural network

network = torch.nn.Sequential( torch.nn.Linear(1, width), torch.nn.Tanh(), torch.nn.Linear(width, width), torch.nn.Tanh(), torch.nn.Linear(width, 1), )

Instantiate SFR (handles NN training/prediction as they're coupled via the prior/likelihood)

sfr = src.sfr.SFR( network=network, prior=src.priors.Gaussian(params=network.parameters, delta=delta), likelihood=src.likelihoods.Gaussian(sigmanoise=2), outputdim=1, numinducing=32, dualbatch_size=None, # this reduces the memory required for computing dual parameters jitter=1e-4, )

sfr.train() optimizer = torch.optim.Adam([{"params": sfr.parameters()}], lr=learningrate) for epochidx in range(numepochs): for batchidx, batch in enumerate(dataloader): x, y = batch loss = sfr.loss(x, y) optimizer.zerograd() loss.backward() optimizer.step()

sfr.set_data(data) # This builds the dual parameters

Make predictions in function space

Xtest = torch.linspace(-0.7, 3.5, 300, dtype=torch.float64).reshape(-1, 1) fmean, fvar = sfr.predictf(X_test)

Make predictions in output space

ymean, yvar = sfr.predict(X_test) ```

Citation

Please consider citing our ICLR 2024 paper. bibtex @inproceedings{scannellFunction2024, title = {Function-space Prameterization of Neural Networks for Sequential Learning}, booktitle = {Proceedings of The Twelth International Conference on Learning Representations (ICLR 2024)}, author = {Aidan Scannell and Riccardo Mereu and Paul Chang and Ella Tami and Joni Pajarinen and Arno Solin}, year = {2024}, month = {5}, }