Neural collapse in the orthoplex regime

1. Problem Statement

The paper investigates the geometric phenomenon known as Neural Collapse (NC) in deep neural networks trained for classification.

• Context: In the "terminal phase of training" (where loss approaches zero), feature vectors of training data collapse to class-specific prototypes, and the classifier weights align with these prototypes.
• Standard Regime: Previous research established that when the number of classes $n$ is small relative to the feature dimension $d$ ($n \le d+1$), these prototypes form the vertices of a regular simplex.
• The Gap: Many modern applications (e.g., large language models) operate in a regime where $n \gg d$. While it is known that for $n \to \infty$ with fixed $d$, features distribute uniformly on a sphere, the specific geometric structures for the intermediate regime where $d + 2 \le n \le 2d$ (the orthoplex regime) were not fully characterized.
• Objective: The authors aim to characterize the emergent geometric figures (specifically softmax codes) in the orthoplex regime, determine their relationship to spherical codes (solutions to the Tammes problem), and analyze how temperature $\tau$ in the cross-entropy loss affects which specific configurations are optimal.

2. Methodology

The authors employ a combination of convex geometry, combinatorial topology, and optimization theory under an Unconstrained Features Model (UFM).

• Model: They treat feature vectors $H$ and weight vectors $W$ as free variables on the unit sphere $S^{d-1}$, decoupling them from the specific neural network architecture.
• Hardmax Limit: They analyze the limit as temperature $\tau \to 0$, where minimizing cross-entropy loss converges to the Hardmax problem. This involves minimizing the maximum inner product between distinct classes.
• Key Definitions:
  • Softmax Code: A configuration $X$ that maximizes $\delta(X) = \min_j \text{dist}(x_j, \text{conv}(X \setminus \{x_j\}))$.
  • Spherical Code: A configuration that minimizes the maximum pairwise inner product $\alpha(X)$.
  • Rattlers: Points in a configuration that are not "tight" (i.e., they can move without violating distance constraints).
• Analytical Tools:
  • Radon's Theorem: Used to prove that any set of $n \ge d+2$ points in $\mathbb{R}^d$ can be partitioned into two sets whose convex hulls intersect.
  • Convexity Analysis: Used to study the behavior of the cross-entropy loss function with respect to the partition of dimensions.
  • Decomposition: Breaking down the configuration space into mutually orthogonal subspaces (batches) to simplify the optimization landscape.

3. Key Contributions and Results

A. Characterization of Softmax Codes (Section 2)

The paper establishes a fundamental equivalence in the orthoplex regime ($d+2 \le n \le 2d$):

• Theorem 2: In this regime, Softmax Codes are precisely Spherical Codes.
• Result: A configuration $X$ is a softmax code if and only if it minimizes the maximum pairwise inner product $\alpha(X)$. Specifically, $\delta(X) = 1$ if and only if $\alpha(X) = 0$.
• Implication: This confirms that the geometric structures emerging from neural collapse in this regime are the solutions to the Tammes problem (maximizing minimum pairwise distance on a sphere).

B. Self-Duality (Section 3)

The authors prove that these codes possess a "self-dual" property:

• Theorem 5: For every hardmax code $(W, H)$ in the orthoplex regime, the feature vectors and weight vectors are identical ($H = W$).
• Mechanism: This is proven by showing that optimal codes in this regime have no "rattlers" (no points that can move freely). The absence of rattlers ensures that the optimal feature configuration is identical to the optimal weight configuration.

C. Temperature-Dependent Optimality (Section 4)

The paper investigates the non-asymptotic case where $\tau > 0$. It reveals that not all softmax codes are equally optimal; the preference depends on the temperature:

• Low-Entropy vs. High-Entropy Codes:
  • Low-Entropy: Consists of a regular simplex and an orthoplex (or similar structures) that are orthogonal.
  • High-Entropy: Consists of mutually orthogonal regular simplices of roughly equal size.
• Theorem 12: There exist temperature thresholds $\tau^-(n)$ and $\tau^+(n)$ such that:
  • For low temperatures ($\tau < \tau^-$), the minimizers of the cross-entropy loss are low-entropy codes.
  • For high temperatures ($\tau > \tau^+$), the minimizers are high-entropy codes.
• Mathematical Insight: The authors reduce the optimization to a discrete problem involving a function $f_{n,\tau}(x)$. They prove this function is strictly concave for small $\tau$ (favoring extreme, uneven partitions) and strictly convex for large $\tau$ (favoring balanced partitions).

4. Significance

• Theoretical Unification: The paper bridges the gap between the well-understood simplex regime ($n \le d+1$) and the high-dimensional limit ($n \to \infty$), providing a complete characterization for the critical intermediate range $d+2 \le n \le 2d$.
• Geometric Insight: It confirms that the "regular simplex" intuition of neural collapse generalizes to orthoplex-based structures (combinations of simplices and orthoplexes) when the number of classes exceeds the dimensionality.
• Practical Implications for LLMs: Since language models often operate with $n \gg d$, understanding that the optimal geometry shifts based on temperature $\tau$ offers new perspectives on how training dynamics and hyperparameters influence the final representation space.
• Methodological Advancement: The use of Radon's theorem and the "rattler-free" analysis provides a robust framework for analyzing high-dimensional geometric configurations in machine learning, potentially applicable to other regimes ($n > 2d$).

Conclusion

The paper demonstrates that in the orthoplex regime, neural collapse results in configurations that are both optimal spherical codes and self-dual. Furthermore, it reveals a phase transition in the optimal geometry driven by the temperature parameter: low temperatures favor "clustered" (low-entropy) structures, while high temperatures favor "balanced" (high-entropy) structures. This work significantly deepens the theoretical understanding of the geometric limits of neural network representations.