Transformers self-organize like newborn visual systems when trained in prenatal worlds

TL;DR

The paper addresses how brains acquire their characteristic visual structure and whether transformers can develop a similar organization when exposed to biologically plausible prenatal input. It trains a Vision Transformer with Contrastive Learning Through Time (ViT-CoT) on sequences generated by a retinal-wave simulator, using a space-time fitting objective over a $300$ ms window across three frames. The results show that, after training, the models spontaneously exhibit edge sensitivity in early layers, shape sensitivity in later layers, and progressively larger receptive fields, mirroring newborn visual systems; these effects vanish under temporally scrambled inputs or with insufficient data. This work supports a common fitting principle between brains and transformers, suggesting prenatal experience alone can shape large-scale visual organization in artificial systems and offering a framework to study brain-like development with transformer models.

Abstract

Do transformers learn like brains? A key challenge in addressing this question is that transformers and brains are trained on fundamentally different data. Brains are initially "trained" on prenatal sensory experiences (e.g., retinal waves), whereas transformers are typically trained on large datasets that are not biologically plausible. We reasoned that if transformers learn like brains, then they should develop the same structure as newborn brains when exposed to the same prenatal data. To test this prediction, we simulated prenatal visual input using a retinal wave generator. Then, using self-supervised temporal learning, we trained transformers to adapt to those retinal waves. During training, the transformers spontaneously developed the same structure as newborn visual systems: (1) early layers became sensitive to edges, (2) later layers became sensitive to shapes, and (3) the models developed larger receptive fields across layers. The organization of newborn visual systems emerges spontaneously when transformers adapt to a prenatal visual world. This developmental convergence suggests that brains and transformers learn in common ways and follow the same general fitting principles.

Figures (9)

• Figure 1: How do visual systems get their structure? Neuroscientists discovered that newborn visual systems are highly structured, including (1) edge sensitivity in early layers, (2) shape sensitivity in later layers, and (3) increasing receptive field sizes across layers. Instructional theories propose that genes instruct the organization of newborn visual systems. Selectional theories propose that genes produce a flexible space-time fitter (brain), which adapts to the continuous flow of prenatal sensory input. As brains adapt, they develop the structure of newborn visual systems.
• Figure 2: Testing Selectional Theories of visual development. If adaptation to space-time data distributions in prenatal worlds is sufficient to produce structured visual systems, then space-time fitting models should develop the same organization as newborn visual systems when trained in prenatal worlds. We tested this prediction by (a) selecting a space-time fitting transformer model, then (b) training the model on simulated prenatal experiences (retinal waves). (c) In their untrained state, the models were not organized like newborn visual systems: they lacked edge sensitivity in early layers, shape sensitivity in later layers, and increasing receptive field sizes across layers. (d) After fitting to retinal waves, the models spontaneously developed all three signatures of newborn visual systems. Space-time fitting transformers develop the same structure as newborn visual systems when exposed to the same prenatal data as newborns.
• Figure 3: Edge development. (a) We measured how models organize lines with different orientations, separated by 10° increments. (b) t-SNE visualizations show that untrained models do not have structured feature spaces that organize images as a function of line orientation, whereas trained models develop structured feature spaces. Early layers (L) have particularly structured edge representations. (c) As a baseline, we created a hardcoded orientation-tuned RDM where similarity between orientations decreased as a function of distance. (d) Correlations between untrained (red lines) and trained (green lines) RDMs for each layer of the model and the orientation-tuned RDM. After training, early layers of the models developed sensitivity to oriented edges.
• Figure 4: Shape development. (a) We tested whether models group novel objects by color versus shape, using two test sets (simple and realistic objects), each crossing four colors and four shapes. RDMs were used to evaluate whether trained fitting models developed color-based versus shape-based representational spaces. (b) Color scores were computed by averaging across RDM cells where objects matched in color (filled cells in Panel b). (c) Shape scores were computed by averaging across RDM cells where objects matched in shape (filled cells in Panel c). (d) Untrained models had color-based representational spaces, whereas trained models developed shape-based representational spaces, as demonstrated by the RDMs (left) and color/shape scores (right). (e) The correlations between untrained (red lines) and trained (test set 1: green lines; test set 2: blue lines) RDMs for each layer of the model and the shape-tuned RDM (panel c). After training, the models developed sensitivity to object shape.
• Figure 5: Development of larger receptive field (RF) sizes across layers. (a) We measured RF sizes across all layers of the models. In the trained models (colored lines), RF sizes increased across layers. Different colors indicate models trained on different numbers of retinal waves. In contrast, untrained models (black line) showed minimal increases in RF size across layers. (b) RF sizes were estimated using a gradient-based method that tracks sensitivity to input pixels. This method estimates RF sizes by identifying which input pixels influence the output of each layer. The gradient maps for untrained models (top) are similar across all layers, indicating a lack of hierarchical spatial organization and no effective increase in RF size. The gradient maps for trained models (bottom) become progressively broader across layers, indicating an increase in RF size and the emergence of hierarchical spatial organization.