Uncovering Symmetry Transfer in Large Language Models via Layer-Peeled Optimization

TL;DR

This paper investigates how training large language models induces geometric symmetries in their weights and embeddings, revealing that target distribution symmetries are transferred to model parameters through a mathematically tractable analysis.

Contribution

The paper introduces a layer-peeled optimization framework to analytically demonstrate symmetry transfer in LLMs, linking target symmetries to model weight geometries.

Findings

Symmetries in target distributions lead to circulant logit matrices in optimal models.

Optimal output projections form equiangular tight frames under permutation symmetries.

Empirical evidence shows open-source LLMs naturally exhibit predicted symmetries.

Abstract

Large language models (LLMs) are pretrained by minimizing the cross-entropy loss for next-token prediction. In this paper, we study whether this optimization strategy can induce geometric structure in the learned model weights and context embeddings. We approach this problem by analyzing a constrained layer-peeled optimization program, which serves as a mathematically tractable surrogate for LLMs by treating the output projection matrix and last-layer context embeddings as optimization variables. Our analysis of this nonconvex optimization program demonstrates that symmetries in the target next-token distributions are transferred to the global minimizers of the layer-peeled model in a precise group-theoretic sense. Specifically, we prove that when the target tokens exhibit a cyclic-shift symmetry (such as the seven days of the week or the twelve months of the year), the optimal logit matrix is exactly circulant, and the Gram matrices of both the output projections and the context embeddings form circulant geometries as well. Next, for exchangeable target distributions invariant under the symmetric group and, more generally, under two-transitive group actions, we show that the global optimal output projection matrix forms a simplex equiangular tight frame, while the optimal logit matrix and context embeddings inherit the permutation symmetries present in the input data. A key technical step is to reduce the constrained nonconvex factorized problem to an explicit logit-level convex characterization for cyclic symmetry and to a symmetry-based lower bound for permutation symmetry, together with a sharp characterization of the optimal factorization. Finally, we empirically demonstrate that open-source LLMs naturally exhibit symmetries consistent with our theoretical predictions, despite being trained without any explicit regularization promoting such geometric structure.