Convergent World Representations and Divergent Tasks

TL;DR

This work tackles how neural representations encode an underlying world and adapt when the world changes. It introduces a World–Data–Model framework that decouples world structure (city coordinates) from data generation (seven geometric tasks) and model training (autoregressive transformers), enabling precise probing of representation geometry and adaptation through Atlantis as new entities. Key findings show that single-task autoregressive training yields emergent, separable world representations, while multi-task learning progressively aligns representations across disjoint tasks; however, certain divergent tasks can catastrophically impede the integration of new entities during fine-tuning, revealing a gap between forward representations and backward adaptation. The study provides a controlled setting to examine representation formation, supports partial evidence for the Multitask Scaling Hypothesis, and highlights the need to account for task-divergence when planning continual-learning and world-model updates.

Abstract

While neural representations are central to modern deep learning, the conditions governing their geometry and their roles in downstream adaptability remain poorly understood. We develop a framework clearly separating the underlying world, the data generation process and the resulting model representations to study these questions in a controlled setup. 5,075 city coordinates define the world and 7 geometric tasks generate the training data for autoregressive training. We find that different tasks give rise to qualitatively and quantitatively distinct world representation geometries. However, multi-task training drives convergence of world representations: models trained on non-overlapping tasks develop aligned geometric representations, providing controlled evidence for the Multitask Scaling Hypothesis of the Platonic Representation Hypothesis. To study adaptation, we pretrain models on all tasks, then test whether new entities (cities) can be consistently integrated into the representation space via fine-tuning. Surprisingly, we find that despite multi-task pretraining, some tasks, which we call divergent, actively harm the representational integration of new entities and harm generalization. Our results show that training on multiple relational tasks reliably produces convergent world representations, but lurking divergent tasks can catastrophically harm new entity integration via fine-tuning.

Figures (20)

• Figure 1: Overview of the World-Data-Model framework.Top: The world consists of 5,075 real city coordinates; we test adaptation by adding 100 synthetic Atlantis cities (App. \ref{['app:world']}). Middle: Seven geometric tasks generate training data from city coordinates (App. \ref{['app:data']}). Bottom: Training dynamics of one model, showing loss curves, linear probing accuracy for coordinate reconstruction and visualizations of internal representations (PCA and linear probe projections) at different training stages. See App. Fig. \ref{['fig:app_training']} for all training curves.
• Figure 2: World representation geometry depends on the data generation process. (a) Different tasks create distinct geometries: distance (thread-like), angle (2D manifold), compass (fragmented), inside (diffuse). Row 1: PCA. Row 2: Linear probe projections. Row 3: Rotated views showing hidden structure. See App. Fig. \ref{['fig:app_reprs']} for more seeds. (b) CKA matrix at layer 5, estimated across 3 seeds. Crossing (Cr) fails to train alone. See App. Fig. \ref{['fig:app_cka_pt1']} for SEM and layers 3, 4, 6. 3D visualizations: https://osf.io/jb8an/?view_only=da001f31c0534dc0b6476141f30db90d .
• Figure 3: Multi-task pretraining drives representational convergence. (a,b) Two-task training creates more regular structures than single-task models. (c) CKA matrix (7$\times$7) for two-task models shows higher alignment (see App. Fig. \ref{['fig:app_cka_pt2']} for SEM). (d) Average CKA increases with task count (1$\rightarrow$2$\rightarrow$3), saturating at $\sim$0.85 for layers 4-6 while layer 3 continues improving (see App. Fig. \ref{['fig:app_cka_3seed']} for SEM). Crossing, which failed to learn in single-task training, is excluded; including it would only strengthen the convergence finding.
• Figure 4: 7-task model. (a) PCA projection of layer 5 representations naturally reveals world map structure. (b) Training curves showing successful learning of all 7 tasks, including crossing which failed in single-task training.
• Figure 5: Fine-tuning generalization and its correlation with representational similarity. (a) Generalization matrix (averaged over 4 seeds; see App. Fig. \ref{['fig:app_ft_vs_ni_4seed']} for individual seeds): each row is a model that integrated Atlantis via one task; columns show normalized improvement on Atlantis queries for each task (see App. \ref{['app:eval']} for metric details). (b) For each task pair (X, Y), we plot the single-task CKA between X and Y against the normalized improvement on task Y after fine-tuning on task X (see App. Fig. \ref{['fig:app_cka_vs_ni_annotated']} for annotated version).