Studying the Effects of Task Diversity Using Small Mathematical Worlds

Abstract

Basic mathematical tasks like modular addition have been a valuable tool for illuminating aspects of neural network behavior that are challenging to see in noisy, unverifiable, large-scale, real-world settings. These studies have typically focused on a single task. However, modern language models are exposed to a vast number of tasks despite being trained with a next-token prediction objective. In this work, we introduce two mathematical task ecosystems, IntWorld and PermWorld, built from distinct but related tasks on integers and permutations. Both are designed to simulate task diversity within a cohesive mathematical 'world' where tasks are distinct but related. Because the datasets are procedurally generated, we can precisely vary the number of tasks, allowing us to begin to disentangle the role of task variety from other factors in training. After describing the datasets, we provide some initial results obtained from these datasets which support the idea that increasing task diversity increases the richness of a model's internal representations.