TRACE for Tracking the Emergence of Semantic Representations in Transformers

The lexicon is built by assigning tokens to POS and semantic role categories, then applying Zipfian scaling and semantic clustering. Each token is associated with a tuple [POS, SRL, Zipf_rank, ClusterID], ensuring interpretability and alignment across analysis stages.

To mimic natural lexical statistics, token frequencies follow a Zipfian distribution [55, 40]:

where $\alpha=1.05$, $\varepsilon_{i}\sim\mathcal{N}(0,0.05)$, and $V=9000$ is the vocabulary size. This maintains a realistic long-tail frequency spectrum.

Semantic clustering introduces collocational structure by forming token groups with variable intra- and inter-cluster association strengths. For tokens $tok_{i}$ and $tok_{j}$ belonging to clusters $c_{i}$ and $c_{j}$:

where $S_{base}$ and $S_{range}$ are tuned to create structured but noisy associations, emulating semantic co-occurrence patterns. Intra-cluster associations are drawn from $[0.4,0.7]$, while cross-cluster links are weaker $([0.05,0.2])$.

Each word receives a naturalised name (e.g., result1, noun5), allowing transparent reverse-mapping for analysis.

Frame-to-syntax mappings define how semantic frames are realised as surface forms while controlling contextual predictability through entropy calibration. All realisations follow valid English grammatical constructions, respecting part-of-speech ordering, agreement patterns, and canonical phrase structure. This grounding enables the study of syntactic and semantic abstraction within structurally coherent input sequences.

Each frame component is annotated with expected entropy based on its role in the frame:

• •

  Low entropy (0.5–1.5 bits): Grammatically determined positions (e.g., determiners required by nouns)

• •

  Medium entropy (1.5–3.0 bits): Semantically constrained positions with multiple valid fillers (e.g., theme roles that accept various object types)

• •

  High entropy (3.0–4.5 bits): Optional frame elements with high variability (e.g., adverbial modifier)

Frames are instantiated according to a target complexity distribution (55% simple, 35% medium, 10% complex), which guides the global entropy profile of the corpus. Example frame realisations include:

Simple TRANSFER:    [AGENT=NOUN] [ACTION=VERB] [THEME=NOUN]
                   "noun2 verb3 noun5"

Medium CREATION:    [AGENT=NOUN] [ACTION=VERB] [THEME=NOUN]
                   [PURPOSE=PREP+NOUN]
                   "noun1 verb2 noun4 prep3 noun7"

Complex MOTION:     [AGENT=NOUN] [REL] [ACTION=VERB] [SOURCE=NOUN]
                   [ACTION=VERB] [GOAL=NOUN]
                   "noun3 rel1 verb5 noun6 verb7 noun9"

To enforce statistical balance across complexity levels, the sentence generator incorporates an entropy-aware sampling mechanism. During generation, the system maintains a global entropy profile, defined as the frequency distribution of sentence positions assigned to low, medium, and high entropy tiers. This profile is updated in real time and compared to the desired target distribution.

If the observed distribution diverges from the target (e.g., too many low-entropy tokens have been sampled), the system increases the sampling weight for frames or token categories that contribute to underrepresented tiers. This feedback mechanism modulates the difficulty of the dataset without sacrificing grammaticality.

The final vocabulary consists of 9,000 tokens distributed across categories as shown in Table 1.

Each sentence is stored with a naturalised token sequence and associated structured annotations, including:

• •

  POS labels (e.g., NOUN, ADJ, VERB)

• •

  Semantic roles (e.g., AGENT, PATIENT, RESULT)

• •

  Entropy tier

• •

  Other Contextual complexity metadata

These annotations enable direct supervision for probing tasks. During model training, hidden states are extracted layer-wise and evaluated using linear probes trained on these annotations. This setup facilitates fine-grained analysis of how and when compositional representations emerge, and how they correlate with curvature, intrinsic dimensionality, and mutual information.

We implemented decoder-only transformer architectures based on the original design of Vaswani et al. [50]. Each model consists of $L$ layers, each with hidden size $d_{\text{model}}$, $H$ attention heads (where $d_{\text{head}}=d_{\text{model}}/H$), and a feedforward dimension $d_{\text{ffn}}$. We focus on decoder-only models given their increasing prevalence in production LLMs and recent arguments that causal architectures provide a cleaner demonstration for emergence [47].

Our architectural choices are guided by recent Transformer scaling laws, notably those articulated by Hoffmann et al. [22]. Namely, we evaluate three configurations that adhere to the Chinchilla scaling law (see below), with approximate parameter counts and architectural specifications as presented in Table 2. All models are trained with a maximum sequence length of $T=16$ tokens and a dropout rate of $0.1$.

We use a simple whitespace-based tokeniser, where each token corresponds to a space-separated word or symbol. This choice allows us to maintain interpretability and simplify downstream representational analyses, while aligning with our controlled, low-scale experimental setup.

All models are trained using next-token prediction with the Adam optimiser [28] with learning rate $1e-3$, and 1000 warm-up steps. Training was conducted for significantly extended epochs beyond the computed minimum requirements of the discussed scaling laws (Section B.1), consistent with the observations of Nanda et al. [36] on the "grokking" phenomenon, where generalisation emerges abruptly after an initial memorisation phase. We consider this an extension of the training strategies outlined in Csordás et al. [12], which emphasise the importance of avoiding early stopping to fully exploit the learning capacity of neural models.

Specifically, we trained:

• •

  Small model: 500 epochs ($\sim$82× the minimum of 6 epochs)

• •

  Medium model: 500 epochs ($\sim$26× the minimum of 19 epochs)

• •

  Large model: 500 epochs ($\sim$12× the minimum of 42 epochs)

We record dense checkpoints throughout training (every 500 training steps), extracting hidden states and gradients from all layers to compute our diagnosis metrics.

We applied minimal regularisation and continuously monitored loss, accuracy, gradient dynamics, representational similarity, and mutual information (MI) throughout training to detect potential phase transitions. This methodological rigour enables us to assess the relationship between model structure and task complexity under controlled and interpretable conditions, while ensuring that all models are trained in accordance with modern scaling principles. To ensure correctness and reproducibility, all experiments were repeated at least 5 times with different random seeds, and we reported the averaged results.

All models demonstrate strong performance on the downstream generation task after their respective phase transitions. Table 3 summarises the evaluation metrics across model scales.

To investigate the interpretability and internal structure of our models, we implemented a probing framework that trains lightweight classifiers on the frozen hidden representations extracted from each layer of our trained models. Specifically, for every layer in the tested models, we trained and evaluated both a part-of-speech (POS) probe and a semantic role labelling (SRL) probe.

Each probe is implemented as a feedforward neural network comprising three linear layers interleaved with ReLU activations and dropout. The network is trained using binary cross-entropy loss, with sigmoid activations at the output layer to support multi-label classification. Given an input representation $\mathbf{h}\in\mathbb{R}^{d}$, the probe computes:

where $W_{1},W_{2},W_{3}$ are trainable weight matrices and $\sigma$ denotes the element-wise sigmoid function.

Probes are trained independently for each layer in the model, allowing us to analyse the emergence and distribution of linguistic and functional features across depth. Representations from each layer are frozen and taken from an already trained model following our training setting explained in Section B, and the underlying model weights are not updated during probing.

Due to the synthetic nature of our dataset, both POS and semantic labels are deterministically derived from token names. For example, a token such as noun3 is assigned the POS label NOUN and may additionally be annotated with semantic roles such as AGENT or ENTITY, depending on the symbolic structure of the task. This approach eliminates annotation ambiguity and ensures consistent supervision across examples.

Given the supervised design of our dataset, we employ linear probes due to their demonstrated effectiveness in similar contexts. While alternative methods such as Sparse Autoencoders (SAEs) have shown promise in unsupervised settings Cunningham et al. [13], linear probes remain a robust and interpretable choice for supervised feature probing Kantamneni et al. [25], Smith et al. [48].

All probes were trained for 30 epochs using the Adam optimiser [28] with a learning rate of $1\times 10^{-4}$, a hidden dimension of 256, and a dropout rate of 0.5. Input dimensionality matched the model’s hidden state size.

We report the average POS probe models’ performance on Tables 4, 5, and 6, for the small, medium, and large models, receptively. We also report the average semantic probe models’ performance on Tables 7, 8, and 9, for the small, medium, and large models, receptively.

We present extended probing results across all model scales (Small, Medium, and Large) for both part-of-speech (POS) and semantic role categories. Figures 6, 7, and 8 show model confidence scores for each layer across training steps.

To assess whether abstraction emerges during training rather than being encoded by architecture alone, we also trained probes on frozen hidden states from randomly initialised models. These models exhibited near-zero performance across all linguistic categories, confirming the absence of structured representations at initialisation.

Figures 9, 10, and 11 show the performance of semantic and POS probes on randomly initialised models.

All the experiments were on an NVIDIA RTX A6000. We used PyTorch (v2.6.0), scikit-learn (v1.6.1), scipy (v1.12.0), seaborn (v0.13.2), Python (v3.9.21), matplotlib (v3.9.4), numpy (v1.26.4), and matplotlib (v3.9.4). Training the small models required around 21 minutes without tracking, and 116 with tracking of all metrics. For medium models, it takes around 28 minutes to train and 130 minutes to train with full tracking. For large models, it takes around 34 minutes to train and 139 minutes to train with full tracking. Training probes require approximately 9 minutes, 19 minutes, and 28 minutes for small, medium and large models, respectively.

MINE approximates the Donsker-Varadhan lower bound on mutual information using a neural critic function $T_{\theta}:\mathcal{X}\times\mathcal{Z}\rightarrow\mathbb{R}$, parameterised by a multilayer perceptron (MLP). Given joint samples $(x,z)\sim P_{XZ}$ and marginal samples formed by pairing $x\sim P_{X}$ with independently sampled $z^{\prime}\sim P_{Z}$, the MI is estimated via:

Our implementation uses a 3-layer MLP with hidden dimensions $[128,128,1]$ and ReLU activations. MINE is used to estimating $I(X;Z_{\ell})$ — the MI between the input embeddings and layer $\ell$ — it is also used to compute $I(Z_{\ell};Z_{\ell+1})$, capturing how information is transmitted between adjacent layers. This allows us to analyse information bottlenecks, compression phases, and abstraction dynamics across the depth of the model.

For each chosen training step $t$ of the model and its layers, we train a separate MINE estimator to convergence. The training protocol is as follows:

• •

  Batch size: 128 examples

• •

  Optimiser: Adam optimiser [28] with learning rate 0.001

• •

  Training steps: 200 iterations

• •

  Positive samples: Joint pairs $(x,z)$ from the same forward pass

• •

  Negative samples: $(x,z^{\prime})$ where $z^{\prime}$ is obtained by shuffling across the batch

Figure 12 shows how mutual information evolves during training across model scales. In small models (top-left), MI between the embedding and subsequent layers drops rapidly and stabilises early, suggesting early-stage compression and limited representational differentiation.

In medium models (top-right), we observe a pronounced dip in MI that aligns temporally with the phase transition identified in our diagnostics. This suggests that representational compression may act as a precursor or trigger for abstraction. Following this dip, MI values remain volatile across layers, reflecting a noisier or less stable reconfiguration of internal representations.

Large models (bottom) follow a broadly similar trend, though with key differences: the MI dip appears in most transitions, but is less evident or absent in others (Layer 1 $\rightarrow$ Layer 2).

While the timing of these dips often aligns with the abstraction phase transition, the metric remains highly volatile. This instability suggests that mutual information, although partially correlated with representational restructuring, may be too noisy and inconsistent to serve as a reliable standalone indicator of abstraction onset.