Out-of-Support Generalisation via Weight-Space Sequence Modelling

TL;DR

This paper introduces WeightCaster, a novel approach that models neural network weights as sequences to improve out-of-support generalisation, providing more reliable and interpretable predictions in safety-critical applications.

Contribution

The paper proposes a new sequence modelling framework in weight space for OoS generalisation, achieving competitive results without explicit inductive biases.

Findings

Outperforms state-of-the-art on synthetic and real-world datasets

Produces plausible and uncertainty-aware predictions

Maintains high computational efficiency

Abstract

As breakthroughs in deep learning transform key industries, models are increasingly required to extrapolate on datapoints found outside the range of the training set, a challenge we coin as out-of-support (OoS) generalisation. However, neural networks frequently exhibit catastrophic failure on OoS samples, yielding unrealistic but overconfident predictions. We address this challenge by reformulating the OoS generalisation problem as a sequence modelling task in the weight space, wherein the training set is partitioned into concentric shells corresponding to discrete sequential steps. Our WeightCaster framework yields plausible, interpretable, and uncertainty-aware predictions without necessitating explicit inductive biases, all the while maintaining high computational efficiency. Emprical validation on a synthetic cosine dataset and real-world air quality sensor readings demonstrates performance competitive or superior to the state-of-the-art. By enhancing reliability beyond in-distribution scenarios, these results hold significant implications for the wider adoption of artificial intelligence in safety-critical applications.