Source-linked AI summary
Studying the Effects of Task Diversity Using Small Mathematical Worlds
Emily Xia Meng, Xuanyu Yang, Junaid Hasan, Michael Ruofan Zeng, Henry Kvinge
TL;DR
Single-task synthetic studies leave task interactions underexplored, despite language models encountering many tasks. This paper builds controllable integer and permutation task ecosystems; broader task sets generally improve transformer probe scores, while transfer depends on task composition and fine-tuning learning rate.
Problem
Synthetic studies have often focused on one task, leaving task interactions unexplored despite language models encountering many tasks during pretraining.
Source: §1 · pp. 1–2Method
The paper builds procedurally generated, related-task ecosystems for integers and permutations and varies the number of training tasks.
Source: §1 · p. 2Results
Takeaways & Limitations
Broader task exposure can make information about unseen mathematical properties more accessible in learned representations.
Source: §5 · p. 5Varying task count also changes per-task exposure under a fixed training-step budget, so task diversity is not isolated from exposure.
Source: §3 · p. 3
Abstract
from arXiv · showhide
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.
1 Introduction
Single-task synthetic studies offer controlled ways to investigate neural-network training, but leave interactions among tasks unexplored. The paper introduces related-task ecosystems whose task counts can be varied to study task diversity.
1 Introduction
- Synthetic tasks support controlled training-data design, known algorithmic solutions, and verifiable outputs.Prior modular-arithmetic studies discovered grokking [1] and decoded a learned algorithm [2], [3].Source: §1 · p. 1
- Single-task analyses leave task interactions unexplored, limiting study of how task diversity affects learning.The authors motivate this gap by noting that language models encounter many tasks during pretraining.Source: §1 · pp. 1–2
- The paper introduces integer- and permutation-based task ecosystems with procedurally controlled task counts to study multi-task learning.Tasks within each world are distinct but related, and the experiments vary the number of training tasks.Source: §1 · p. 2
- The paper varies training-task counts and compares related-task effects and training exposure.It also asks what properties are represented, which tasks support transfer, and how performance changes on longer inputs.Source: §1 · p. 2
2 Related work
Related work examines shared representations, task grouping, and task diversity in controlled settings. This paper studies task-set variation in mathematical worlds, parameter fine-tuning transfer, and representations across models trained on related properties.
2 Related work
- Multitask learning studies benefits of shared representations, while task-grouping experiments show task choice affects joint-training outcomes.Raventós et al. [6] study pretraining-task count and in-context learning in controlled regression; this paper instead fine-tunes model parameters for transfer.Source: §2 · p. 2
- Procedurally generated mathematics benchmarks cover multiple tasks [7], while this work varies the training task set to study learning about a common object class.Source: §2 · p. 2
- The paper probes hidden states and uses CKA to compare representations, building on prior representation analyses [8] [9] [10].For permutations, it fixes one-line notation and uses exact property relations to study task composition and alignment; related work includes disjoint-task representation convergence [12], permutation notation, and multi-task, multi-encoding pretraining [14].Source: §2 · p. 2
3 Task Ecosystems from Small Mathematical Worlds
IntWorld and PermWorld collect distinct, related tasks that characterize integers and permutations. Their task sets and training-task counts are structured to support studies of transfer and task diversity, with a fixed-step budget that trades off against exposure to individual tasks.
3 Task Ecosystems from Small Mathematical Worlds
- The ecosystems treat integers and permutations as small mathematical worlds whose related tasks capture core facets of each object.The paper describes these worlds as objects defined in part by their associated tasks, using informal terminology from [15].Source: §3 · p. 2
- IntWorld includes arithmetic, sorting, and digit operations; sixteen tasks are used for pretraining and four, including sorting, are held out for transfer.Integer input scales can be modified by changing the number of digits, supporting length-generalization studies.Source: §3 · p. 2
- PermWorld uses 32 scalar permutation properties for lengths 2 ≤ n ≤ 30, split into disjoint pools of sixteen for training and held-out evaluation.Each replicate trains on prefixes of one pool and evaluates on the other.Source: §3 · p. 3
- Both ecosystems use nested training-task sets of 1, 2, 4, 8, and 16 tasks.IntWorld uses designated pretraining and held-out tasks; PermWorld swaps the training and held-out property pools across replicates.Source: §3 · p. 3
- Holding training steps fixed across task counts trades off total budget against exposure to each task.The authors chose equal step counts across T_k because they considered differing total training budgets the more concerning confound.Source: §3 · p. 3
4 Experiments
Experiments in IntWorld and PermWorld vary pretraining task count and assess held-out-property decoding, few-shot transfer, and representation similarity. Results depend on architecture, task relationships, and fine-tuning learning rate.
Linear probing
- Integer Transformer probe R² rises 0.200→0.371 (T1→T16); permutation mean R² rises 0.207→0.308 (T1→T8), then dips to 0.301 (T16).These probes predict properties excluded from pretraining; the MLP does not show the same trend.Source: §4 · p. 3, Figure 1 · p. 4
Few-shot transfer
- With 20 examples, T8 predecessor accuracy is 72.24 ± 10.51% versus 0% from random initialization; T8 transfer is 0.85% for least common multiple, 8.20% for modular addition, and 0% for sorting.Removing successor from T8 lowers predecessor accuracy to 16.62 ± 13.16%, though that control also changes mixture size and per-task exposure.Source: §4 · p. 4
- On PermWorld, twenty-shot exact accuracy rises from 16.64 ± 2.60% at T1 to 33.59 ± 3.61% at T16, while its gain over zero-shot grows from 3.48 ± 2.78 to 17.80 ± 5.76 percentage points; however, T16 does not outperform random initialization.Source: §4 · p. 4
Representation similarity
- For individual permutation properties, related pairs have mean CKA 0.1375 versus 0.0903 for other pairs (one-sided permutation test p = 0.015).Related pairs are prespecified; for example, descents and recoils are related, while descents and fixed points are not.Source: §4 · p. 5
- Across all eight permutation relations, the correct inverse or complement transformation yields higher CKA than both controls (two-sided exact sign tests p = 0.0078).This supports similarity being sensitive to the specific mathematical relationship between properties.Source: §4 · p. 5
- With four properties per model, CKA does not consistently rise as related pairs increase; unseen-task T16 zero-shot accuracy is 16.72%, below the 32.83% majority-answer baseline.Only one of 12 paired curves increases monotonically.Source: §4 · p. 5
5 Conclusion
Initial experiments in small mathematical worlds suggest that broader task exposure can make unseen properties more accessible in Transformer representations, while transfer depends on task relationships and fine-tuning settings. Task count also changes exposure, composition, and how well tasks are learned, so the results remain preliminary.
Conclusion
- Transformer probe scores generally improve as training task sets grow, while successful transfer depends on which tasks are learned together.Probe targets are excluded from pretraining, suggesting that broader task exposure can make information about unseen mathematical properties more accessible in learned representations.Source: §5 · p. 5
- Transfer benefits also vary with the fine-tuning learning rate.The matched-exposure control covers only T16 predecessor transfer.Source: §5 · p. 5
- Task count changes per-task exposure, task composition, and how well tasks are learned; future work will vary composition and exposure independently and examine shared computations among related tasks.Source: §5 · p. 5
A Model Architecture and Training Setup
The setup trains small models on nested task-count conditions in two mathematical domains, using procedural data and shared optimization settings. Integer inputs use base-100 tokenization, and the pretraining task collections cover arithmetic and other integer properties.
A Model Architecture and Training Setup
- Across both domains, models use nested task sets T1, T2, T4, T8, and T16, with 20,000 updates, effective batch size 64, and procedurally generated datasets.Transformers have four blocks, width 256, eight heads, and feed-forward width 1,024; IntWorld also uses a one-block MLP.Source: A Model Architecture and Training Setup · p. 7
B.1 Experimental setup
- The main experiment crosses five task-count settings, two architectures, and three seeds for 30 base models; nine additional controls test task removal and matched exposure.The matched-exposure T16 control doubles training to 40,000 updates so its per-task exposure matches T8.Source: §B.1 · p. 7
B.2 Integer representation and tokenization
- Integers use most-significant-first base-100 digits; values above 99 receive NUM_START and NUM_END markers, while values from 0 to 99 use one number token.Sequences include a size field, inputs, task token, answer, and sequence markers; binary inputs also delimit the second argument.Source: §B.2 · p. 7
B.3 Pretraining tasks
- Integer pretraining tasks include successor, arithmetic, modulo, comparisons, digit operations, divisibility, factorial, and list operations; factorial inputs are restricted to 0–99.The tasks are defined over non-negative integers, with positive moduli; digit-count answers are already represented in the size field.Source: §B.2 · p. 7, §B.3 · p. 8
B.4 Held-out transfer tasks
Predecessor transfer is strongest for T8, while results vary across the other held-out tasks and training mixtures. Evaluation fine-tunes models on fixed support sets, with the comparison also reflecting changes in task mixture and per-task exposure.
B.4 Held-out transfer tasks
- Models are fine-tuned for 500 steps on fixed support sets of 20, 50, or 100 examples, with checkpoint selection by validation loss.The held-out tasks include predecessor, least common multiple, modular addition, and ascending sort.Source: §B.4 · p. 8
- With twenty examples, T8 reaches 72.24 ± 10.51% predecessor accuracy versus zero from random initialization; removing successor reduces T8 accuracy to 16.62 ± 13.16%.Doubling T16's pretraining updates raises its accuracy from 21.02% to 36.12%; the task-removal comparison also changes mixture size and per-task exposure.Source: §B.4 · p. 8
- 98.01 ± 2.23%: with 100 examples, T8 reaches this predecessor accuracy, versus 82.91 ± 3.08% for T16 and 14.81 ± 3.17% from random initialization.Figure 3 presents predecessor transfer by support size; error bars show standard deviation across three seeds.Source: §B.4 · p. 8, Figure 3 · p. 9
- Transfer is weaker on the other holdouts: T8's twenty-shot accuracy is 0.85% for least common multiple, 8.20% for modular addition, and zero for sorting.Sorting remains below 1% even with one hundred examples.Source: §B.4 · p. 8
B.5 Generalization by task and input length
Integer accuracy is assessed across task mixtures and input lengths, but neither architecture improves monotonically on the longest inputs. Task-level results show that performance differs substantially by task, including a decline in addition accuracy as the mixture grows.
B.5 Generalization by task and input length
- At T16, mean accuracy on seven-to-ten-digit inputs is 9.77% for the Transformer and 2.78% for the MLP.Figure 4 reports integer accuracy by input length; task-count comparisons also change per-task exposure and training mixture.Source: §B.5 · p. 9, Figure 4 · p. 9
- Neither architecture improves monotonically on the longest inputs, despite evaluation spanning in-distribution, six-digit, and seven-to-ten-digit inputs.The task-count comparison changes both the training mixture and the number of presentations per task.Source: §B.5 · p. 9
- Transformer addition accuracy falls from 86.04 ± 8.11% at T8 to 34.57 ± 2.06% at T16, while addition, multiplication, successor, and subtraction score zero on seven-to-ten-digit inputs in every mixture that includes them.Digit sum, comparison, and successor have high in-distribution accuracy before T16, while multiplication and list sum perform substantially worse.Source: §B.5 · p. 9
B.6 Linear probing results by layer
The appendix reports layer-wise probing alongside the permutation task design, held-out evaluation, and experiment setup. For the Integer Transformer, probes are strongest at intermediate layers and generally improve with greater task diversity.
B.6 Linear probing results by layer
- Probe performance is generally strongest at intermediate Transformer layers, particularly Blocks 2 and 3, then decreases slightly toward the final normalized representation.Performance improvements with greater task diversity are also visible at intermediate layers.Source: §B.6 · p. 9, Figure 6 · p. 10
C.1 Pretraining tasks
- Table 1 defines 32 scalar permutation properties; the experiments split properties into pools, with each property held out from its model's training pool.Across three replicates, each property changes pools at least once, while related properties may occur across the two pools.Source: §C.1 · pp. 10–11, Table 1 · p. 11
C.2 Held-out transfer tasks
- Scalar-task transfer evaluates the sixteen properties in the opposite pool, keeping each replicate's held-out set fixed as the training prefix grows from one to sixteen properties.A separate twenty-task experiment holds out reduced-word translation, composition, and Lehmer-code translation.Source: §C.2 · p. 11
C.3 Experiment description
- Permutation experiments train four-layer Transformers on 1, 2, 4, 8, or 16 tasks per pool, using lengths 2–30 and 20,000 updates per model.Thirty models span five task counts and two pools; each is evaluated on the 16 properties in the other pool, with results averaged across three joint task-split/model-seed replicates.Source: §C.3 · p. 12
C.4 Zero Shot Evaluation
Across zero-shot evaluation, linear probes, and few-shot transfer, results vary with the number of training properties and evaluation setup. More training properties do not reliably predict unseen tasks, while probe scores and matched-rate adaptation comparisons show qualified gains.
C.4 Zero Shot Evaluation
- Zero-shot average accuracy rose from 12.25 ± 2.66% with 1 training property to 16.72 ± 2.50% with 16, but stayed below the 32.83% majority baseline.In a separate 20-task experiment, Transformer and MLP models both had 0% accuracy on three unseen tasks.Source: §C.4 · p. 12, Figure 7 · p. 12
C.5 Linear probes
- Final-layer average R2 on unseen properties was 0.2071, 0.2591, 0.2719, 0.3076, and 0.3013 for models trained on 1, 2, 4, 8, and 16 properties, respectively.All three seeds improved from 1 to 8 properties, while one run decreased at 16; the probes used frozen models and linear ridge regression on held-out properties.Source: §C.5 · pp. 12–13
- For held-out descents at T16, mean probe R2 rose from 0.164 after block 1 to 0.258 after block 4, then fell to 0.254 after final normalization.Appendix C.8 plots probe performance across all 32 targets and layers.Source: §C.5 · p. 13
C.6 Few-shot transfer
- At T16, the pretrained advantage was 11.71 ± 5.24 percentage points at 10−5 and 2.96 ± 1.67 at 3 × 10−4 learning rates.Under the original protocol, test accuracy rose from 16.64% at T1 to 33.59% at T16, while random initialization reached 34.28%.Source: §C.6 · p. 13, Figure 8 · p. 13
C.7 Representation similarity
CKA comparisons show stronger alignment for related properties when inputs receive the correct mathematical transformation, while increasing the number of related training pairs does not consistently raise alignment. The section also describes all-layer probe plots across permutation properties.
C.7 Representation similarity
- Linear CKA compares models’ internal representations; these comparisons use the final-normalized hidden state immediately before the task token.The feature matrices are column-centered, with rows representing paired input examples.Source: §C.7 · p. 13
- Related property pairs had average CKA of 0.1375, versus 0.0903 for other cross-property pairs (p = 0.015).The comparison covered eight related pairs and 112 other pairs of single-property models.Source: §C.7 · p. 14, Figure 9 · p. 14
- The correct inverse or complement transformation produced higher CKA than either control for all 8 property relations across three seeds (p = 0.0078).The transformation identities were checked exhaustively for permutations of lengths two through eight.Source: §C.7 · p. 14, Figure 10 · p. 15
- Mean CKA across 0, 1, 2, and 4 related training pairs was 0.2935, 0.2752, 0.2662, and 0.3648, with no consistent increase as related pairs increased.The mean paired difference between four and zero related pairs was 0.0712 ± 0.2121 (p = 0.7744); only seven of twelve differences were positive.Source: §C.7 · p. 15, Figure 11 · p. 15
C.8 All-layer probe plots
- The all-layer probe plots cover 32 permutation properties, with sixteen targets per page and curves for T1, T2, T4, T8, T16, and untrained models.Panels fix one property and use a shared R2 scale; Emb. marks embedding output and Norm. marks final normalization.Source: §C.8 · p. 15, Figure 12 · p. 16, Figure 13 · p. 17