Studying the Effects of Task Diversity Using Small Mathematical Worlds

AI summary

Overview

Research area: Machine learning interpretability and training dynamics, specifically multi-task / task-diversity effects studied in small procedurally generated mathematical settings.

Technical level: Intermediate. The paper uses standard tools (transformers, MLPs, linear ridge probes, centered kernel alignment) but the concepts are accessible to readers with basic deep learning familiarity.

Scope: The paper introduces two synthetic multi-task benchmarks, IntWorld (integer tasks) and PermWorld (permutation tasks), and reports initial experiments on how the number of training tasks affects probe decodability, few-shot transfer, and representation similarity.

What This Paper Is About

Most mechanistic studies of neural networks use a single synthetic task, such as modular addition, so they cannot say anything about how multiple tasks interact during training. Since real language models are pretrained on enormous mixtures of tasks with a next-token prediction objective, the authors build two controlled "small mathematical worlds" — one of integer tasks, one of permutation tasks — where the number of training tasks can be varied exactly because the data is procedurally generated. The goal is to start separating the effect of task variety from other confounding factors in training.

Key Contributions

1. Two task ecosystems. IntWorld, built from integer tasks (arithmetic, digit operations, sorting, divisibility, and so on), with sixteen pretraining tasks and four related held-out transfer tasks; and PermWorld, built from 32 scalar permutation properties evaluated on permutations of length 2 ≤ n ≤ 30, where each replicate splits the 32 properties into two disjoint pools of sixteen.

2. Controlled task-count experiments. Nested task sets T₁, T₂, T₄, T₈, and T₁₆ let the authors vary training-task diversity while holding architecture and training steps fixed, using small MLPs and four-block transformers.

3. Representation and transfer measurements. Linear ridge probes on frozen hidden states measure how linearly decodable held-out mathematical properties are, and centered kernel alignment (CKA) compares internal representations across models trained on mathematically related versus unrelated properties.

4. Evidence that task diversity enriches internal representations. Transformer probe scores generally rise with training task count, while behavioral transfer proves more conditional — dependent on which related tasks were learned together and on the fine-tuning learning rate.

Main Findings

• Linear probing improves with task count for transformers. For integer tasks, the transformer final-layer probe R² rises steadily from 0.200 at T₁ to 0.371 at T₁₆. For permutations, probe mean R² rises from 0.207 at T₁ to 0.308 at T₈ and then decreases slightly to 0.301 at T₁₆. The permutation probe baseline from random initialization is R² = 0.215, and the reported final-layer averages across the five task counts are 0.2071, 0.2591, 0.2719, 0.3076, and 0.3013.

• The MLP does not show the same trend. For integers, the MLP lacks the monotonic probe improvement seen in the transformer; the authors state it is unclear why, suggesting the transformer's token-based structure may make cross-task-representation building easier.

• Transfer depends on task relevance. Fine-tuning on twenty examples of predecessor (a task excluded from pretraining), the T₈ transformer reaches 72.24 ± 10.51% exact accuracy, versus 0% from random initialization and 0.27 ± 0.27% after T₄ pretraining. Successor is in T₈ but not T₄; removing successor from T₈ drops accuracy to 16.62 ± 13.16%. T₁₆ performs worse than T₈, and doubling its updates raises accuracy from 21.02% to 36.12%, still below T₈.

• Other integer holdouts transfer poorly. T₈ twenty-shot accuracy is 0.85% for least common multiple, 8.20% for modular addition, and zero for sorting. With one hundred examples, T₈ reaches 98.01 ± 2.23% on predecessor versus 82.91 ± 3.08% for T₁₆ and 14.81 ± 3.17% from random initialization. Sorting stays below 1% even with one hundred examples.

• Permutation few-shot transfer improves with task count but does not beat random initialization. Twenty-shot exact accuracy rises from 16.64 ± 2.60% at T₁ to 33.59 ± 3.61% at T₁₆, and the gain over zero-shot grows from 3.48 ± 2.78 to 17.80 ± 5.76 percentage points. However, T₁₆ did not outperform random initialization (random controls reached 34.28%), and because only task count varied rather than the presence of a specific related property, the authors interpret this as improved adaptability rather than a causal relevance effect.

• Permutation transfer is sensitive to learning rate. With matched learning rates at T₁₆, the pretrained model beat the randomly initialized model by 11.71 ± 5.24 percentage points at 10⁻⁵, but by only 2.96 ± 1.67 percentage points at 3 × 10⁻⁴. The original protocol used different learning rates for pretrained (10⁻⁵) and random (3 × 10⁻⁴) models, which the authors identify as a confound.

• Zero-shot performance stays below a majority baseline. Permutation zero-shot accuracy rises from 12.25 ± 2.66% with one training property to 16.72 ± 2.50% with sixteen, but remains below the 32.83% majority-answer baseline. In a separate twenty-task permutation experiment, both transformer and MLP models score 0% on three unseen tasks: reduced-word translation, composition, and Lehmer-code translation.

• Representations of mathematically related properties align more closely. Comparing models trained on individual properties (48 single-property models from 16 properties and three seeds), mean CKA is 0.1375 for the eight prespecified related property pairs versus 0.0903 for the other 112 pairs, with a one-sided task-label permutation test giving p = 0.015.

• Alignment is sensitive to the specific mathematical relationship. When inputs are transformed by the correct inverse or complement, CKA exceeds both control conditions for all eight relations across three seeds, with relation-level two-sided sign tests giving p = 0.0078 (= 2/2⁸).

• Relatedness does not accumulate monotonically. Training each model on exactly four properties with r ∈ {0, 1, 2, 4} directly related pairs between the two sets gives average CKA of 0.2935, 0.2752, 0.2662, and 0.3648 respectively; only one of the 12 paired curves increases monotonically.

• Length generalization does not improve monotonically. Neither architecture improves monotonically on the longest inputs. At T₁₆, mean accuracy on seven-to-ten-digit inputs is 9.77% for the transformer and 2.78% for the MLP, and addition, multiplication, successor, and subtraction score zero on that range in every mixture containing them.

• Task averages hide per-task differences. Digit sum, comparison, and successor have high in-distribution accuracy before T₁₆, while multiplication and list sum perform substantially worse. Transformer addition accuracy falls from 86.04 ± 8.11% at T₈ to 34.57 ± 2.06% at T₁₆.

• Probe information peaks at intermediate layers. Probe performance is generally strongest at intermediate layers, particularly Blocks 2 and 3, and decreases slightly toward the final normalized representation. As one example, mean probe R² for held-out descents at T₁₆ rises from 0.164 after block 1 to 0.258 after block 4, then falls to 0.254 after final normalization.

Methodology in Plain English

The authors generate all training data programmatically, which lets them control exactly how many distinct tasks a model sees. For integers, they use sixteen pretraining tasks (successor, addition, multiplication, modulo, comparison, digit sum, GCD, list sum, subtraction, integer division, digit count, digit reversal, digit occurrence count, even/odd classification, divisibility classification, and factorial) and hold out four related tasks (predecessor, least common multiple, modular addition, ascending sort). Integers are written in base 100.

For permutations, they use 32 scalar properties — things like descents, peaks, valleys, fixed points, cycle counts, longest increasing subsequence length, and maximum displacement — computed on permutations of length 2 to 30 written in one-line notation. Each of three replicates splits these 32 properties into two disjoint pools of sixteen, so a model trained on one pool is evaluated on the other.

Models are trained on nested task sets T₁ through T₁₆ for 20,000 update steps with effective batch size 64. The architectures are small: a four-block transformer with width 256, eight heads, and feed-forward width 1,024, and a one-block MLP with width 256 and hidden width 1,024. Three seeds are used throughout (integer seeds 17, 42, 314159; permutation seeds 17, 42, 101 paired with different property partitions).

To measure what the model has learned, the authors freeze the base model and train linear ridge probes to predict properties that were never seen during pretraining, reporting R². They also fine-tune on a few examples of held-out tasks to test transfer, and use linear CKA on 4,096 matched validation examples to compare internal representations across models.

Why This Matters

This work gives interpretability researchers a controllable sandbox for asking questions that single-task synthetic settings cannot answer — namely, what happens to a model's internal representations and transfer behavior as the number of training tasks grows. It also sharpens a methodological point: small differences in training setup (which related task happens to be included, how many updates were run, which learning rate fine-tuning uses) can change transfer conclusions substantially, so claims about task diversity need careful controls.

Potential real-world applications (note: the paper itself does not report deployed applications; these are plausible domains the setup could inform):

• Designing pretraining data mixtures for language models by reasoning about which task families reinforce each other.
• Improving few-shot adaptation in low-data specialist settings, where knowing which pretraining tasks support transfer matters.
• Diagnosing whether a model's representations support properties it was never explicitly trained on.
• Building educational or tutoring models for mathematics, where coverage of related skills is a design question.

Industry relevance: Task-mixture design is a central open problem for anyone pretraining or fine-tuning models. The findings that transfer can hinge on the presence of a single closely related task, and that measured benefits can disappear under a different fine-tuning learning rate, are directly relevant to practitioners evaluating whether multi-task pretraining actually helps a downstream capability.

Future Directions

• Vary task composition and training exposure independently, since the current design changes both at once; the matched-exposure control only covers T₁₆ predecessor transfer.
• Test whether the learning-rate sensitivity observed in PermWorld few-shot transfer also occurs in IntWorld.
• Investigate how related tasks share computations, beyond measuring representational alignment through CKA.
• Extend the two task collections beyond the current lists, which the authors explicitly flag as a plan for future work.
• Explain why the MLP does not show the same probe improvement with task diversity as the transformer.

Target Audience

Researchers and graduate students working on mechanistic interpretability, training dynamics, and multi-task or continual learning, especially those who already use synthetic mathematical tasks like modular arithmetic as testbeds. Practitioners involved in pretraining data mixture design or few-shot fine-tuning will also find the transfer and learning-rate controls instructive. Readers looking for large-scale language model results should note that this paper deliberately works at small scale with small MLPs and four-block transformers.

Authors’ 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.