DeepOSets: Non-Autoregressive In-Context Learning with Permutation-Invariance Inductive Bias

TL;DR

DeepOSets introduces a non-autoregressive neural architecture with permutation-invariance that can learn in-context without parameter updates, demonstrating strong performance in regression tasks with fewer parameters than transformers.

Contribution

This paper presents DeepOSets, a novel permutation-invariant architecture combining DeepSets and DeepONets, with a proven universal approximation property for regression operators.

Findings

DeepOSets effectively learn various regression tasks.

DeepOSets outperform transformer-based models in parameter efficiency.

Replacing DeepSets with Set Transformer improves high-dimensional accuracy.

Abstract

In-context learning (ICL) is the remarkable ability displayed by some machine learning models to learn from examples provided in a user prompt without any model parameter updates. ICL was first observed in the domain of large language models, and it has been widely assumed that it is a product of the attention mechanism in autoregressive transformers. In this paper, using stylized regression learning tasks, we demonstrate that ICL can emerge in a non-autoregressive neural architecture with a hard-coded permutation-invariance inductive bias. This novel architecture, called DeepOSets, combines the set learning properties of the DeepSets architecture with the operator learning capabilities of Deep Operator Networks (DeepONets). We provide a representation theorem for permutation-invariant regression learning operators and prove that DeepOSets are universal approximators of this class of operators. We performed comprehensive numerical experiments to evaluate the capabilities of DeepOSets in learning linear, polynomial, and shallow neural network regression, under varying noise levels, dimensionalities, and sample sizes. In the high-dimensional regime, accuracy was enhanced by replacing the DeepSets layer with a Set Transformer. Our results show that DeepOSets deliver accurate and fast results with an order of magnitude fewer parameters than a comparable transformer-based alternative.