Beyond Co-occurence: A Study of Early-stage Semantic Geometry in Next-Token Prediction

Neural Collapse predicts that balanced one-hot classification pushes model representations to be equally far from each other; a symmetric configuration that ignores any semantic similarity in the inputs. This creates a puzzle: next-token prediction language models are trained predominantly (as context length increases) with one-hot labels, yet they clearly learn that "red" and "blue" are more similar than "red" and "circle." How does gradient descent find such semantic structure when co-occurrence statistics collapse to one-hot sparsity, eliminating any shared next-tokens among different contexts? To investigate this tension we identify a controlled setting where inputs have latent semantic factors but are mapped to distinct one-hot labels. We find that semantic geometry emerges early in training: representations cluster by shared attributes despite receiving no explicit supervision to do so. This structure is transient: with sufficient capacity and time, the model eventually reaches the predicted symmetric state where all representations are equally separated. We study this phase transition through Gram matrix analysis and propose a preliminary modification to the commonly used unconstrained features model to capture the emerged semantic geometry.