Studying the Effects of Task Diversity Using Small Mathematical Worlds

Emily Xia Meng; Xuanyu Yang; Junaid Hasan; Michael Ruofan Zeng; Henry Kvinge

cs.LGOct 9, 2026

Studying the Effects of Task Diversity Using Small Mathematical Worlds

Authors: Emily Xia Meng, Xuanyu Yang, Junaid Hasan, Michael Ruofan Zeng, Henry Kvinge

Organizations: University of Washington, Seattle · University of Washington Pacific Northwest National Laboratory

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.

Figures & tables

Figure 1: Linear probe performance generally improves with pretraining task diversity for Transformers

Figure 2: Few-shot predecessor transfer results

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Predecessor transfer by support size. Error bars show standard deviation across three seeds.

Figure 4: Integer accuracy by input length. Means and SDs across three seeds.

Figure 5: Integer Transformer accuracy by task (%). Cells average three seeds; $<1$ marks positive accuracy below $1\%$ , and dashes mark unavailable evaluations.

Figure 6: Linear probing by layer for Integer Transformer

Property	Definition
Local order
Descents	$\#\{i<n:\pi_{i}>\pi_{i+1}\}$ .
Recoils	Number of descents of $\pi^{-1}$ .
Peaks	$\#\{1<i<n:\pi_{i-1}<\pi_{i}>\pi_{i+1}\}$ .
Valleys	$\#\{1<i<n:\pi_{i-1}>\pi_{i}<\pi_{i+1}\}$ .
Double ascents	$\#\{1<i<n:\pi_{i-1}<\pi_{i}<\pi_{i+1}\}$ .

Appendix

Table 1: The 32 scalar permutation tasks.

Figure 7: Permutation results before and after retraining. The dashed line marks majority-answer accuracy. Error bars show SD across three random seeds.

Figure 8: Permutation adaptation at matched learning rates. Error bars show SD of paired accuracy differences.

Figure 9: Specialist CKA by task pair.

Property pair	Transformation
Descents / recoils	Inverse
Peaks / valleys	Value complement
Fixed / anti-fixed points	Value complement
Excedances / deficiencies	Inverse
Longest increasing / decreasing subsequence	Value complement
Longest increasing / decreasing run	Value complement

Appendix

Table 2: Relations used for the specialist comparison. Value complement is

c(\pi)_{i}=n+1-\pi_{i}

. The transformed right-hand property equals the left-hand property, except in the last row where it is larger by one.

Figure 10: Specialist CKA under the same input, the correct transformation, and a control transformation.

Figure 11: CKA with four training tasks per model. Thin lines show the twelve layout–seed combinations; the thick line shows their mean.

Figure 12: Permutation probes: order and positions.

Figure 13: Permutation probes: cycles and global structure.

Multi-Task Learning

Multi-task learning trains models to solve related tasks together by sharing representations, parameters, or supervision.

See all 389 papers

Open-world AI: a Tale of Three Requirements
Oct 14, 2026Jianyu Zhang, Léon BottouContinual Learning for LLMs
FAIR, Meta∗ · Flatiron Institute
ResearchMath-14K: Scaling Research-Level Mathematics via Agents
May 27, 2026Guijin Son, Seungyeop Yi, Minju Gwak +3Supervised Fine-Tuning LLM Mathematical Reasoning
Seoul National University · OneLineAI · Yonsei University
Beyond Compression: Diagnosing How Post-Training Changes Mathematical Reasoning
Sep 29, 2026Hongyang Li, Yiming Zhu, Xiao Li +3Language Model Distillation RL Post-Training
University of Luxembourg · Seafill Open-Source Community · Université Paris-Saclay

Oct 14, 2026J/K move · Enter open · S save

Open-world AI: a Tale of Three Requirements

Jianyu Zhang, Léon Bottou

FAIR, Meta∗ · Flatiron Institute

Continual Learning for LLMs

Open paper AskPDF arXiv

The closed world asks AI for specialized competence and lets it manufacture the iid data its theory assumes: a fixed criterion licenses collecting examples until the task is covered. The open world asks for versatility --- taking up tasks a machine was not built for, from few examples and with little prior from a human designer --- and offers data that is grouped, not shuffled. Its tasks are not known in advance, so a machine must do inference-time learning (1st requirement). Modeling tasks as sharing latent factors and decomposing the training objective, we derive two more: a rich representation, retaining more than iid generalization requires, and a predictive disentanglement that aligns inputs and outputs into matched low-dimensional groups. The analysis further predicts the need for a consistent inner learner, and the benefit of long contexts spanning related tasks. We test these claims on language models up to 8B, asking whether a model has a property (qualitative) rather than how much it wins (quantitative). Models that meet the requirements improve as data from related tasks accumulates; the rest degrade, and neither scale nor layer-type interleaving repairs this.

Abstract

Multi-Task Learning

Explore similar work

Open-world AI: a Tale of Three Requirements

ResearchMath-14K: Scaling Research-Level Mathematics via Agents

Beyond Compression: Diagnosing How Post-Training Changes Mathematical Reasoning