Formation of Representations in Neural Networks

## Abstract

Abstract Understanding neural representations will help open the black box of neural networks and advance our scientific understanding of modern AI systems. However, how complex, structured, and transferable representations emerge in modern neural networks has remained a mystery. Building on previous results, we propose the Canonical Representation Hypothesis (CRH), which posits a set of six alignment relations to universally govern the formation of representations in most hidden layers of a neural network. Under the CRH, the latent representations (R), weights (W), and neuron gradients (G) become mutually aligned during training. This alignment implies that neural networks naturally learn compact representations, where neurons and weights are invariant to task-irrelevant transformations. We then show that the breaking of CRH leads to the emergence of reciprocal power-law relations between R, W, and G, which we refer to as the Polynomial Alignment Hypothesis (PAH). We present a minimal-assumption theory demonstrating that the balance between gradient noise and regularization is crucial for the emergence the canonical representation. The CRH and PAH lead to an exciting possibility of unifying major key deep learning phenomena, including neural collapse and the neural feature ansatz, in a single framework.

## 1 Introduction

The success of deep learning is often attributed to its ability to learn latent representations from data (Bengio et al.,, ). These latent representations, progressively formed as data passes through the network’s layers, are found to encode increasingly abstract features of the input:

$$ $x\to h^{1}\to h^{2}\to...\to h^{D}\to\hat{y},$ (1) $$

where $x$ is the input, $\hat{y}$ the output, $D$ the network depth, and $h^{i}$ the activation of the $i$-th layer. For neural networks to perform well, the transformations between layers must capture meaningful structures in the data. Understanding how these latent representations are formed and structured is a foundational problem in deep learning, with implications for both theoretical understanding and practical applications. Despite significant advances, how neural networks organize and transform these internal representations remains an open question. This gap in understanding hinders our ability to design more efficient, interpretable, and generalizable models.

In this work, we seek to bridge this gap by introducing the Canonical Representation Hypothesis (CRH). At its core, the CRH posits that neural networks, during training, inherently align their representations with the gradients and weights. Furthermore, the satisfaction and breaking of a subset of CRH equations are found to delineate the universal phases, which are empirically observable scaling relationships between the weights, activations, and gradients. The CRH reveals a striking aspect of representation learning: there may exist a set of universal equations that govern the formation of representations and universal phases which distinguish the layers in modern neural networks, independent of the task, architecture, or loss function.
Thus, the CRH provides an useful perspective on how neural networks evolve toward compact and interpretable solutions.

Our main contributions are the following:

- 1.
Proposal and justification of the CRH which states that within a layer, the neuron gradients, latent representation, and parameters are driven into mutual alignment after training, due to noise regularization balance (Section 3 and 4);
- 2.
Identification of mechanisms that break alignments, quantified via a Polynomial Alignment Hypothesis, which predicts power law scaling behaviors that characterize distinct phases of neural networks arising when the CRH is broken (Section 5).

The experiments are presented in Section 6. Section 7 discusses the implications of CRH to the formation of neural representations and the connection of the CRH to prior observations. Related works are discussed in Section 2. The proofs are left to the Appendix Section B.

## 2 Related Works

Empirical results show that the representations of well-trained neural networks share universal characteristics (Maheswaranathan et al.,, ; Huh et al.,, ).
For us, a closely related phenomenon is the neural collapse (NC) (Papyan et al.,, ), which studies how structured and noise-robust low-rank representations emerge in a classification model. The CRH can be seen as a generalization of NC because one can prove that when restricted to certain settings, the CRH is equivalent to the NC (Section 7). Another related phenomenon is the neural feature ansatz (NFA) (Radhakrishnan et al.,, ), which shows that the weight matrices of fully connected layers evolve during training. However, the NFA studies the weight evolution, not the representations.
Empirical power-laws are known to exist in large neural networks (Kaplan et al.,, ; Bahri et al.,, ), which relates the model performance to their sizes. The power laws discovered in our work are different because they are reciprocal relations that relate dual objects (R, G, W) to each other rather than to the performance. Other related works are discussed in the context where they become relevant. Also, due to space constraints, more related works regarding neural representations are discussed in Appendix A.

## 3 Canonical Representation Hypothesis

Let us consider an arbitrary hidden layer $h_{b}$ of any model after a linear transformation:

$$ $h_{b}=Wh_{a}(x).$ (2) $$

For convention, $h_{b}$ is called the “preactivation” of the next layer, and $h_{a}$ is the “postactivation” of the previous layer. The gradients of the representations are also of interest: $g_{a}=-\nabla_{h_{a}}\ell$ and $g_{b}=-\nabla_{h_{b}}\ell$, where $\ell$ is the sample-wise loss function. We note the generality of this setting, as $h_{a}$ can be an arbitrarily nonlinear function of $x$, and $f(x)=f(h_{b}(x))$ can be another arbitrary transformation. Also, letting $h_{a}=(h_{a}^{\prime}(x),1)$ accounts for when there is a trainable bias.(^1^11As an example, consider a two-layer network, $f(x)=W_{2}\sigma(W_{1}x)$, with fully connected layers. For the first layer, $h_{a}^{1}=x$, and $h_{b}^{1}=W_{1}x$. For the second layer, $h_{a}^{2}=\sigma(W_{1}x)$ and $h_{b}^{2}=W_{2}h_{a}^{2}$.)

Much recent literature has suggested that the quantities $h_{a},\ h_{b},\ W$ and their gradients are correlated with each other after or throughout training. The neural collapse phenomenon suggests that in a deep overparameterized classifier, $\mathbb{E}[h_{a}h_{a}^{\top}]\propto W^{\top}W$ for the penultimate fully connected layer (Papyan et al.,, ; ; Ji et al.,, ; Kothapalli,, ; Rangamani et al.,, ). In the study of kernel and feature learning, a primary mechanism of how the neural tangent kernel changes is that after a few steps of update, the representations become correlated with the weights (Everett et al.,, ) and so the quantity $Wh_{a}$ will be significantly away from zero, which also implies a strong relationship between $W^{\top}W$ and $\mathbb{E}[h_{a}h_{a}^{\top}]$. The recent work on neural feature ansatz shows that $W^{\top}W\propto\nabla_{h_{a}}f\nabla_{h_{a}}^{\top}f$ for fully connected networks (Radhakrishnan et al.,, ; Gan and Poggio,, ). The idea that the latent variables will become correlated with the weight updates is also a central notion in the feature learning literature (Yang and Hu,, ; Everett et al.,, ).

Taken together, these results suggest a simple and unifying set of equations that can describe a fully connected layer after (and perhaps during) training. Let $c\in\{a,b\}$ and define $H_{c}=\mathbb{E}[h_{c}h_{c}^{\top}]$, $G_{c}=\mathbb{E}[g_{c}g_{c}^{\top}]$, and $Z_{c}=M_{c}M_{c}^{\top}$, where $M_{a}=W^{\top}=M_{b}^{\top}$. One can imagine six alignment relations between all the quantities within the same layer:

$$ representation-gradient alignment (RGA): $\displaystyle H_{c}\propto G_{c},$ (3) representation-weight alignment (RWA): $\displaystyle H_{c}\propto Z_{c},$ (4) gradient-weight alignment (GWA): $\displaystyle G_{c}\propto Z_{c},$ (5) $$

where $\mathbb{E}$ denotes the averaging over the training set. Because $h_{b}$ comes after $h_{a}$ during computation, we will refer to the alignment between any of the $b$-subscript matrices as a forward relation and all $a$-subscript matrices as a backward relation.

For the formation of representations in neural networks, the most important relations in Eq. (3)-(5) are perhaps the forward and backward RGA, as they directly relate the representations $\mathbb{E}[hh^{\top}]$ to their gradients with respect to the loss function. Because there is no scientific reason for us to believe that the forward representation is more important than the backward representation or vice versa, one should study both directions carefully. The backward RWA in its general form has not been discussed in the literature but has been implicitly studied in the particular setting of neural collapse, which happens in the penultimate layer of an image classification task (Section 7), while the backward GWA seems to be unknown to the best of our knowledge. The backward GWA is not identical to the neural feature ansatz (NFA) but can be seen as an equivariant correction to the NFA (Section 7), and the forward GWA is also unknown to the community. Adding together the forward and backward versions, there are six alignment relations. That these relations simultaneously for any fully connected layer will be referred to as the canonical representation hypothesis (CRH).

Three scientifically fundamental questions are thus (1) does there exists a rigorous set of assumptions under which the CRH can be proved; (2) what mechanisms can cause the CRH to break; and (3) can predictions of the CRH and its breaking be empirically observed in realistic deep neural network settings? We devote the rest of the paper to answering these three questions in the given order, and then we collect insights from all these answers.

## 4 Noise-Regularization Balance Leads to Alignment

#### Notation

$\Delta A$ denotes the difference in the quantity $A(\theta)$ after one step of training algorihtm iteration at time step $t$: $\Delta A:=A(\theta_{t+1})-A(\theta_{t})$. $\eta$ denotes the learning rate and $\gamma$ denotes the weight decay. $\mathbb{E}$ denotes the empirical average over the training set. $\ell(x,\theta)$ denotes the per-sample loss function, where $x$ is the data point and $\theta$ is the parameters. Its empirical average is the empirical loss $L$: $L(\theta)=\mathbb{E}[\ell(x,\theta)]$.

In this section, we present a formal and rigorous framework under which the CRH can be proved. As will become clear in the next section, this proof also explains how and when the CRH may fail. The problem setting is the same as in Eq. (2). The training proceeds in an online learning setting in which the training proceeds with weight decay of strength $\gamma$. We make the following assumption.

###### Assumption 1 (Mean-field norms) .

The norms of $g$ and $h$ approximate their empirical averages: (A1) $\|h_{a}\|^{2}=\mathbb{E}[\|h_{a}\|^{2}]$, (A2) $\|g_{b}\|^{2}=\mathbb{E}[\|g_{b}\|^{2}]$.

This assumption holds, for example, for a high-dimensional Gaussian random vector, whose norm is of order $O(d)$ with a $\sqrt{d}$ standard deviation. A1 also holds automatically if the representations are normalized. Note that only a subset of all the assumptions is needed to prove each equation we derive below. For example, Eq. (6) below only requires A1 to prove. The minimal set of assumptions required to prove each equation in this section are stated in Section B.3. We discuss the main intuition for the proof in the main text, and present the formal theorem at the end of the section.

###### Assumption 1 (Mean-field norms) .

#### Forward alignment.

Consider the time evolution of $h_{b}h_{b}^{\top}$ during SGD training:

$$ $\displaystyle\Delta$ $\displaystyle(h_{b}(x)h_{b}^{\top}(x))=\eta(\|h_{a}\|^{2}g_{b}h_{b}^{\top}+\|h _{a}\|^{2}h_{b}g_{b}^{\top}-2\gamma h_{b}h_{b}^{\top})+\eta^{2}\|h_{a}\|^{4}g_ {b}g_{b}^{\top}+O(\eta^{2}\gamma+\|\Delta(h_{a}h_{a}^{\top})\|),$ $$

where $\eta$ is the learning rate. At the end of training, the learned representations should reach stationarity and so $\Delta\mathbb{E}[h_{b}h_{b}^{\top}]\approx 0$.(^2^22See Figure 21 for the evolution of $\Delta H$.) Taking the expectation of both sides, we obtain

$$ $0=\underbrace{z_{b}\mathbb{E}[g_{b}h_{b}^{\top}]+z_{b}\mathbb{E}[h_{b}g_{b}^{ \top}]}_{\text{learning}}-\underbrace{2\gamma\mathbb{E}[h_{b}h_{b}^{\top}]}_{ \text{regularization}}+\underbrace{\eta z_{b}^{2}\mathbb{E}[g_{b}g_{b}^{\top}] }_{\text{noise}},$ (6) $$

where $z_{b}=\mathbb{E}[\|h_{a}\|^{2}]$. The noise term is due to the discretization error of SGD and can be significant either when the step size is large, or the gradient is noisy. The mechanism behind this alignment is that while the gradient noise expands the representation, weight decay contracts it. When the dynamics reaches approximate stationarity, the dynamics due to learning plus these two effects must balance at the end of training.

Now, if additionally either (a) $\mathbb{E}[\Delta W]=0$ (namely, at a local minimum) or (b) $\mathbb{E}[\Delta(WW^{\top})]=0$ holds, the weight will also align with the cross terms between $g_{b}$ and $h_{b}$: $WW^{\top}\propto\mathbb{E}[g_{b}h_{b}^{\top}]+\mathbb{E}[h_{b}g_{b}^{\top}]$, which leads to the effect that the learning term above must also balance with the regularization term. Thus, eventually, the regularization effect will have to be balanced with the gradient noise. Lastly, if both (a) and (b) hold, the alignment between all three matrices emerges: $G_{b}\propto H_{b}\propto Z_{b}$.

#### Backward Alignment.

One can similarly derive condition for $\mathbb{E}[\Delta(g_{a}g_{a}^{\top})]=0$: $z_{a}(\mathbb{E}[h_{a}g_{a}^{\top}]+\mathbb{E}[g_{a}h_{a}^{\top}])+\eta z_{a}^
{2}\mathbb{E}[h_{a}h_{a}^{\top}]=2\gamma\mathbb{E}[g_{a}g_{a}^{\top}]$, where we have defined $z_{a}=\mathbb{E}[\|g_{b}\|^{2}]$. The forward CRH can then be derived if $W^{\top}W$ and $W$ reaches stationarity.

The following theorem formalizes these results.

###### Theorem 1 .

Under Assumption 1, when $\mathbb{E}[\Delta(h_{a}h_{a}^{\top})]=0$, $\mathbb{E}[\Delta(g_{b}g_{b}^{\top})]=0$, $\mathbb{E}[\Delta(WW^{\top})]=0$, and $\mathbb{E}[\Delta(W^{\top}W)]=0$, there exist real-valued constants $c_{1},\ c_{2},\ c_{3},\ c_{4}>0$ such that

$$ $\displaystyle WW^{\top}+c_{1}\mathbb{E}[g_{b}g_{b}^{\top}]=c_{2}\mathbb{E}[h_{ b}h_{b}^{\top}],\quad W^{\top}W+c_{3}\mathbb{E}[h_{a}h_{a}^{\top}]=c_{4} \mathbb{E}[g_{a}g_{a}^{\top}].$ (7) $$

Additionally, if at a local minimum,

$$ $\displaystyle WW^{\top}\propto\mathbb{E}[g_{b}g_{b}^{\top}]\propto\mathbb{E}[h _{b}h_{b}^{\top}],\quad W^{\top}W\propto\mathbb{E}[h_{a}h_{a}^{\top}]\propto \mathbb{E}[g_{a}g_{a}^{\top}].$ (8) $$

###### Remark .

A strength of this derivation is that it is oblivious to the loss function, the model architecture, or the type of activation used (as long as the second moments exist). This may explain the wide applicability of the RGA observed in Section 6. The above alignment relations can be seen as a type of fluctuation-dissipation theorem in theoretical physics (Kubo,, ), which states that in a driven stochastic dynamics, the fluctuation of the force must balance with the rate of energy loss – a fundamental law first discovered by Einstein, (). Prior applications of the fluctuation-dissipation theorem to deep learning have focused on the covariance of the model parameters (Yaida,, ; Liu et al.,, ) and are not directly relevant to the representations.

###### Theorem 1 .

###### Remark .

## 5 CRH Breaking and Polynomial Alignment Hypothesis

While the CRH can be found to hold for many scenarios, it is highly unlikely that it always holds perfectly and for every layer (e.g., see Section 6). In this section, we study what happens if the CRH is broken; we then suggest two mechanisms which cause the CRH to break.
The following theorem shows that all six relations are intimately connected, even if only a subset of the CRH holds.
For a square matrix $A$, we use $A^{-n}:=(A^{+})^{n}$ to denote the $n-$th power of the pseudo inverse of $A$, and $A^{0}=AA^{+}$ is an orthogonal projection matrix to the column space of $A$.

###### Theorem 2 (CRH Master Theorem) .

Let $A$, $B$, $C$ be a permutation of $\mathbb{E}[hh^{\top}]$, $\mathbb{E}[gg^{\top}]$, and $Z$, and let $\tilde{D}:=PDP$ be a projected version of $D$ for a projection matrix $P$. Then,

- 1.
(Directional Redundancy) if any two forward (backward) alignments hold, all forward (backward) alignments hold;
- 2.
(Reciprocal Polynomial Alignments) if one of any forward alignments and one of any backward alignments hold, there exists scalars $\alpha_{c}$, $\beta_{c}$, and $\delta_{c}$ satisfying $-1\leq\alpha_{c},\beta_{c},\delta_{c}\leq 3$ such that
$\tilde{A}_{c}^{\alpha_{c}}\propto\tilde{B}_{c}^{\beta_{c}}\propto\tilde{C}_{c}
^{\delta_{c}},$
(9)
(as detailed in Table 1)
where $c\in\{a,b\}$ denotes the backward and forward relations respectively, and the corresponding projection $P_{c}\in\{Z_{c}^{0},\mathbb{E}[h_{c}h_{c}^{\top}]^{0},\mathbb{E}[g_{c}g_{c}^{
\top}]^{0}\}$, e.g. such that $\tilde{A}=P_{c}AP_{c}$.
- 3.
(Canonical Alignment I) If (any) one more relation holds in addition to part 2, then all six alignments hold in the $Z^{0}$ subspace; in addition, at a local minimum, all six alignments hold;
- 4.
(Canonical Alignment II) If all six alignments hold, $\mathbb{E}[hh^{\top}]\propto\mathbb{E}[gg^{\top}]\propto Z\propto P$, where $P$ is an orthogonal projection matrix.

**Table 1: The reciprocal polynomial relations of the CRH Master Theorem. When one forward relation and one backward relation hold simultaneously, all six matrices are polynomially aligned in a subspace (Theorem 2). Each scaling relationship can be regarded as a possible phase for the layer during actual training. The right panel shows how existing observations about neural networks fit into the phase diagram. A ✓ denotes that this phenomenon is compatible with the specified phase. NC refers to the neural collapse. NFA refers to the neural feature ansatz. CU (correlated update) refers to the (idealization of the) common observation that $h_{a}$ is correlated with $W$ a few steps after training (Everett et al.,, ). The llm column shows the compatibility of the scaling relation for transformer observed in Figure 3.**
| Phase | Back. Alignment | Forw. Alignment | Back. Power Law | Forw. Power Law | NC | NFA | CU | llm |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CRH | $H_{a}\propto Z_{a}\propto G_{a}$ | $H_{b}\propto Z_{b}\propto G_{b}$ | - | - | ✓ | ✓ | ✓ |  |
| \hdashlineback. CRH | $H_{a}\propto Z_{a}\propto G_{a}$ | - | - | $\tilde{H}_{b}^{0}\propto\tilde{Z}_{b}^{0}\propto\tilde{G}_{b}$ |  | ✓ | ✓ |  |
|  |  |  |  | ($H_{b}\propto Z_{b}^{2}$) |  |  |  |  |
| \hdashlineforw. CRH | - | $H_{b}\propto Z_{b}\propto G_{b}$ | $\tilde{H}_{a}\propto\tilde{Z}_{a}^{0}\propto\tilde{G}_{a}^{0}$ | - |  |  |  | ✓(3-6) |
|  |  |  | ($Z_{a}^{2}\propto G_{a}$) |  |  |  |  |  |
| 1 | $H_{a}\propto G_{a}$ | $H_{b}\propto G_{b}$ | $\tilde{H}_{a}^{0}\propto\tilde{Z}_{a}\propto\tilde{G}_{a}^{0}$ | $\tilde{H}_{b}^{0}\propto\tilde{Z}_{b}\propto\tilde{G}_{b}^{0}$ |  |  |  | ✓(3-6) |
| 2 | $H_{a}\propto Z_{a}$ | $H_{b}\propto Z_{b}$ | $\tilde{H}_{a}\propto\tilde{Z}_{a}\propto\tilde{G}_{a}^{0}$ | $\tilde{H}_{b}\propto\tilde{Z}_{b}\propto\tilde{G}_{b}^{0}$ |  |  | ✓ |  |
| 3 | $G_{a}\propto Z_{a}$ | $G_{b}\propto Z_{b}$ | $\tilde{H}_{a}^{0}\propto\tilde{Z}_{a}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}^{0}\propto\tilde{Z}_{b}\propto\tilde{G}_{b}$ |  | ✓ |  |  |
| 4 | $H_{a}\propto G_{a}$ | $H_{b}\propto Z_{b}$ | $\tilde{H}_{a}\propto\tilde{Z}_{a}^{0}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}\propto\tilde{Z}_{b}\propto\tilde{G}_{b}^{-1}$ |  |  |  |  |
| 5 | $H_{a}\propto Z_{a}$ | $H_{b}\propto G_{b}$ | $\tilde{H}_{a}^{3}\propto\tilde{Z}_{a}^{3}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}\propto\tilde{Z}^{2}\propto\tilde{G}_{b}$ |  |  | ✓ | ✓(3-6) |
| 6 | $H_{a}\propto G_{a}$ | $G_{b}\propto Z_{b}$ | $\tilde{H}_{a}\propto\tilde{Z}_{a}^{2}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}\propto\tilde{Z}_{b}^{3}\propto\tilde{G}_{b}^{3}$ |  |  |  | ✓(1) |
| 7 | $G_{a}\propto Z_{a}$ | $H_{b}\propto G_{b}$ | $\tilde{H}_{a}^{-1}\propto\tilde{Z}_{a}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}\propto\tilde{Z}_{b}^{0}\propto\tilde{G}_{b}$ |  | ✓ |  |  |
| 8 | $H_{a}\propto Z_{a}$ | $G_{b}\propto Z_{b}$ | $\tilde{H}_{a}^{2}\propto\tilde{Z}_{a}^{2}\propto\tilde{G}_{a}$ | $\tilde{H}_{b}\propto\tilde{Z}_{b}^{2}\propto\tilde{G}_{b}^{2}$ |  |  | ✓ | ✓(2) |
| 9 | $G_{a}\propto Z_{a}$ | $H_{b}\propto Z_{b}$ | $\tilde{H}_{a}\propto\tilde{Z}_{a}^{0}\propto\tilde{G}_{a}^{0}$ | $\tilde{H}_{b}^{0}\propto\tilde{Z}_{b}^{0}\propto\tilde{G}_{b}$ |  | ✓ |  |  |

The idea behind the proof of the theorem is that there is a degree of redundancy in the six matrices: every forward relation implies a backward relation and vice versa. As an example, if $Z_{a}\propto H_{a}$, then we also have $Z_{b}^{2}\propto H_{b}$, which can be obtained by multiplying $W$ on the left and $W^{\top}$ on the right.

The last part of the theorem clarifies what it means to perfectly satisfy the CRH: the latent representation is fully compact, where weight and representation are only nonvanishing in the subspaces where the gradient is nonvanishing. Moreover, the weight matrix does nothing but rotates the representation, implying that the information processing is invertible once the CRH is fully satisfied. This is consistent with the observation that once one layer has perfect alignment, all the layers after it also have almost perfect alignment, a sign that the representation cannot be further compressed (Section 6). Therefore, the CRH is consistent with the observation that last layers of neural networks are low-rank and invariant to irrelevant features.

Part (2) of the theorem is especially relevant when the the CRH is only partially satisfied. Depending on which subset of the hypotheses holds well, the learning dynamics and outcome may be classifiable into as many as $2^{6}=64$ phases. In different phases, the learning dynamics and the found solution will likely be completely different due to different scaling relations. For example, positive exponents between $Z_{a}$ and $H_{a}$ will imply that the layer is enhancing the principle components of $H_{a}$, while suppressing the lesser features; a negative exponent would imply the converse. Even if we remove the redundancy implied by the theorem, there are still at least $6+{6\choose 2}+1=22$ phases. One can also define additional phases according to the ordering of the matching of each relation when the matching is imperfect, which gives $6!=720$ phases, although these ordering phases may not have major influence on training. In many experiments we performed, at least one of the forward relations and one of the backward relations are observed to hold very well (e.g., Figure 8). This means that one is quite likely to observe the power-law relations predicted in Table 1. We also find it common for different layers to be in different phases, even within the same network.

Broadly interpreted, part (2) predicts a power law relation between the spectrum of all six matrices, which is also what we observe in almost all experiments (Section 6). What is quite surprising is that almost all positive exponents we observed are within the range $[1/3,3]$, which is exactly the range of exponents that the theorem predicts (e.g., see Section C.8). Formally, that the $H$, $Z$, and $G$ are polynomially related can be called the “Polynomial Alignment Hypothesis (PAH)” and is a natural extension of the CRH. That scaling relations can be used to characterize different phases of matter is an old idea in science. In physics, phases can be classified according to their scaling exponents, and having a different set of exponents implies that the underlying dynamics and mechanism are entirely different (Pelissetto and Vicari,, ). This connection corroborates our physics-inspired proof.

###### Theorem 2 (CRH Master Theorem) .

#### Breaking of CRH.

A major remaining question is whether we can find mechanisms such that the CRH breaks. The theory in the previous section implies one primary mechanism that the CRH breaks. For all six alignments to hold, it needs to be the case that both $\mathbb{E}[\Delta W]$, $\mathbb{E}[\Delta Z]$ are zero, but these two conditions may not be easily compatible with each other, as they together imply that $\Delta W$ has zero variance. While this may be possible for some subspaces (e.g., see Ziyin et al., ()), it does not hold for every subspace. In fact, it is easy to show that unless the minibatch size is large enough the SGD updates will never reach zero variance that is $\Delta W\neq 0$ even for $t\to\infty$ (see Lemma 8 in ; see also ).
Thus, the competition between reaching a training loss of zero and the need to reach a stationary fluctuation is a primary cause of the CRH breakage. Because the rank tends to decrease for later layers in networks performing classification, we expect that CRH holds better for later layers than for earlier layers.

This problem is especially troublesome in the layers where $\mathbb{E}[\Delta W]$ has a high rank, which holds true for the earlier layers of the network but not the later layers (). This is consistent with the observation that the CRH holds much better in the latter layers than in the beginning layers. This analysis is consistent with the fact that a stronger alignment is strongly correlated with a more compact representation (Figure 4).

#### Finite-Time Breaking of CRH.

To leading order in $\gamma$, the proof of the CRH implies that

$$ $\displaystyle G_{b}+O(\gamma)\propto\gamma^{2}H_{b}\propto\gamma^{2}WW^{\top}, \quad H_{a}+O(\gamma)\propto\gamma^{2}G_{a}\propto\gamma^{2}W^{\top}W.$ (10) $$

This means that when $\gamma$ is small, the forward alignment between $H_{b}$ and $WW^{\top}$ are strong (because the prefactor is independent of $\gamma$), while the other two are weak – because the huge disparity between the two matrices, it might take gradient descent longer than practical to reach such a solution. Similarly, the backward alignment between $G_{a}$ and $W^{\top}W$ is strong for a small $\gamma$. This prediction will be verified in Section 7 when we discuss the neural feature ansatz.

Figure: Figure 1: Six alignment relations in the penultimate layer and output layer of a ResNet18 trained on CIFAR-10 (res1). Left: forward CRH. Right: backward CRH. We see that all six relations hold significantly across two fully connected layers. Also, we show that the matrix ${\rm cov}(g,h)$ is well aligned with $WW^{\top}$ in the appendix Section C.7, which is a strong piece of evidence supporting the key theoretical step that the cross terms will be aligned with the weights (and $G,\ H$).
Refer to caption: extracted/5900109/plot/resnet_f-CRH_0.png

Figure: Figure 2: Penultimate layer of the conjugate matrices ($H,G,Z$) after training (fc2). This is an example of CRH being well satisfied, where all three matrices are well aligned after training.
Refer to caption: extracted/5900109/plot/Z_example.png

## 6 Experiments

In this section, we present experimental evidence that supports predictions resulting from the CRH. We also perform experiments to test mechanisms that break the CRH. When the CRH is broken, we put special emphasis on verifying the RGA, as it directly links to the formation of representations and may arguably be the most important relation of the three subsets.

#### Metric for alignment.

We would like to measure how similar and well-aligned the six matrices in the CRH are. Let $A$ and $B$ be two square matrices, each with $d^{2}$ real-valued elements. We use the Pearson correlation to measure the alignment between two matrices (Herdin et al.,, ):

$$ $\alpha(A,B):=\frac{1}{K}\left(\frac{1}{d^{2}}\sum_{ij}A_{ij}B_{ij}-\frac{1}{d^ {4}}\sum_{ij}A_{ij}\sum_{ij}B_{ij}\right),$ (11) $$

where $K$ is a normalization factor that ensures $\alpha\in[-1,1]$. $\alpha(A,B)$ will be referred to as the alignment between the two matrices $A$ and $B$. Note that $\alpha=\pm 1$ if and only if $A=c_{0}B$ for some constant $c_{0}$. Therefore, $\alpha$ can be seen as a quantitative metric for the alignment and if $|\alpha|=1$, the alignment is perfect. Our initial pilot experiments suggest that the alignment effects are the strongest when the $h$ and $g$ are normalized and if the mean of each is subtracted.(^3^33These may be due to the fact that the gradient and activation have rare outliers that tend to disturb the balance. See Section C.2 for the RGA on transformers without subtracting the mean.) Thus, we always normalize $h$ and $g$ and subtract the mean to be consistent in the experiments. This is equivalent to measuring ${\rm cov}(\hat{h},\hat{h})$ and ${\rm cov}(\hat{g},\hat{g})$, where ${\rm cov}$ denotes covariance and $\hat{a}=a/\|a\|$. As a notational shorthand, we use $\alpha_{ab,cd}$ to denote $\alpha({\rm cov}(a,b),{\rm cov}(c,d))$ for the rest of the paper.

#### Settings.

We experiment with the following settings and name each setting with a unique identifier. fc1: Fully connected neural networks trained on a synthetic dataset that we generated using a two-layer teacher network. This experiment is used for a controlled study of the effect of different hyperparameters. fc2: the same as fc1, except that the output dimension is extended to 100 and the input distribution interpolates between an isotropic and nonisotropic distribution.
res1: ResNet-18 ($11$M parameters) for the image classification; res2: ResNet-18 self-supervised learning tasks with the CIFAR-10/100 datasets. llm: a six-layer eight-head transformer ($100$M parameters) trained on the OpenWebText (OWT) dataset. The details of training methods are described in Section C.

#### CRH.

We start with the supervised learning setting with ResNet-18 trained on CIFAR-10. We measure the covariances matrices with data points from the test set. Figure (3) shows that very good alignment $\alpha>0.7$ is achieved quite early in the training, and continues to improve during the later stage of training. This case might remind some readers of neural collapse (NC) – because NC is also most significant in the penultimate layer of large image classifiers. As we will show in the next section, in the interpolation regime of classification tasks, the NC is equivalent to the CRH.

Figure: Figure 3: The power-law alignment between the eigenvalues $\lambda_{h}$ and $\lambda_{g}$ of $H_{b}$ and $G_{b}$ in a six-hidden layer transformer (llm). Left to Right: first to the penultimate layers. The grey dashed lines show the power-law relations $\lambda_{h}\propto\lambda_{g}^{\alpha}$ for $\alpha=1,\ 2,3$ respectively. We see that the first layer has an exponent of $3$, the second has an exponent of $2$, and all the layers after it are observed to have an exponent of $1$. Different colors show different heads within the same layer. The range of the power exponents is in almost perfect agreement with the predicted range in Table 1. Referring to the table, this implies that these layers are in phases 5, 8, and 6, respectively. The setting is the same as the LLM experiment. Also, see Section C.8 for fully connected nets.
Refer to caption: extracted/5900109/plot/gpt/power_law_0.png

Figure: Figure 4: The rank deficiency and the backward $\alpha_{gg,hh}$ in fully connected nets (fc2). The rank of representation is strongly negatively correlated with $\alpha$. Here, every color is a different weight decay (from $10^{-6}$ to $10^{-4}$), and every point is a different layer in the net. The setting is the same as the fully connected net experiment.
Refer to caption: extracted/5900109/plot/rank-alignment.png

Another example of CRH is provided in Figure 8 below, in case of a fully connected network in a regression task. The result shows that when the weight decay is not too small, the CRH is quite close to being perfectly satisfied for intermediate layers. Examples of the representation and the dual matrices are presented in Figure 2 for a layer that (almost) satisfies the CRH.

#### Breaking of CRH.

A clear evidence for the correctness of the theory is that at a small weight decay, the strongest alignments are the forward-RWA and the backward-RGA, which will be presented in the experiment in Figure 8 after we discuss the relationship of the CRH to the NFA.

Two indirect evidences are that (1) the compactness of the representation is found to be negatively correlated with the alignment level (Figure 4), and (2) the observed positive exponents between the eigenvalues of the matrices are almost aways within the range $[1/3,3]$, which is the predicted range by Theorem 2 (Figure 3).

#### RGA.

A relation of particular interest to the formation of representations is the RGA, which predicts that the represenations are aligned with the gradients across them. Now, we show that the RGA holds well across a broad range of tasks in common training settings. When the CRH holds, the RGA holds as well, and so the experiments in the CRH section already shows that the RGA holds for the last layers of ResNet.

Large Language Models (llm). We measure the covariances of the output of each attention head in every layer.
See Figure 5 for the evolution of the alignment and Figure 14 for examples of the representation. Three baselines of comparisons are (1) the alignment between the covariances of a rank-$40$ (roughly equivalent to the actual ranks of ${\rm cov}(h,h)$ and ${\rm cov}(g,g)$) random projection of two $200$ dimensional isotropic Gaussian, which stays close to $0.14$, (2) the alignment between the feature covariance of different heads, which starts high but drops to a value significantly lower than $\alpha_{gg,hh}$, (3) the alignment between the initial feature and current feature for the same head, which starts from $1$ and also drops quickly. The RGA is found to hold stronger than the other baselines.

Figure: Figure 5: Alignment between ${\rm cov}(h,h)$ and ${\rm cov}(g,g)$ in a six-layer transformer trained on the OWT dataset (llm). From left to right: layer 1, 2, 4, 5. Also, see layer 3 in Figure 15. The shaded region shows the variation (min and max) across eight different heads in the same layer. The RGA is significnatly stronger than the alignment between initial and final representation, and the alignment between different heads.
Refer to caption: extracted/5900109/plot/gpt/first_alignment_0.png

Self-Supervised Learning (res2). Self-supervised learning focuses on learning a good and versatile representation without knowing the labels for the problem. See Figure 6 for the time evolution of $\alpha$. We see that the RGA holds well for the fully connected layers, but not so strongly for the conv layers.
Experiments do show increased $\alpha$ values if we decrease the batch size and increase $\gamma$. Also, see Figure 13 for examples of the representations in the convolutional layers. We see very good qualitative alignment between the two matrices.

Figure: Figure 6: The alignment for different layers during 1000 training epochs (res2). Layers 1-3 are convolutional layers, and layers 4-5 are fully connected ones.
Refer to caption: extracted/5900109/plot/simclr_total.png

Fully Connected Nets (fc1). We also perform a systematic exploration of how hyperparameters of training and model architectures affect the formation of representations. Results presented in Section C.6 show the following phenomena: (1) Gradual Alignment: deeper layers tend to have better alignment; if a layer has an almost perfect alignment ($\alpha\approx 1$), then any layer after it will also have almost perfect alignment (note the similarity of this phenomenon with neural collapse); this is consistent with the tunnel effect discovered in (Masarczyk et al.,, ); (b) Critical Depth and Alignment Separation: batch size $B$ affects the alignment significantly; earlier layers have worse alignment as $B$ increases; later layers have better alignment as $B$ increases; wider networks have similar level of alignment as thinner ones for SGD; for Adam, it is more subtle: early layers have worse alignment for a large width, later layers have better alignment for a large width; (c) weight decay $\gamma$ affects the alignment significantly and systematically; larger $\gamma$ always leads to better alignment.

## 7 Insights

This section first studies the implication of the CRH (mainly the RGA) for the representation of neural networks (7.1-2). We then clarify the relation of the CRH to NC and NFA (7.3-4).

### 7.1 Plasticity-Expressivity Coupling

Figure: Figure 7: The response diversity (the diagonal terms of the covariance) and correlation is coupled to the plasticity of neurons (fc1).
Refer to caption: extracted/5900109/plot/coupling.png

A direct interpretation of the RGA is that the plasticity of neurons in neural networks is strongly coupled to their expressivity after training:

$$ $\underbrace{\mathbb{E}[h_{b}h_{b}^{\top}]}_{\text{expressivity}}\propto \underbrace{\mathbb{E}[g_{b}g_{b}^{\top}]}_{\text{plasticity}}.$ (12) $$

The first term directly measures the variability of the neuron response to different inputs, thus measuring how expressive or capable this neuron is. The second term measures the degree of plasticity because the gradient for the weight matrix prior to $h_{b}$ is proportional to $g_{b}h_{a}^{\top}$, and if $g_{b}$ is zero with zero variance, the input weights to this neuron will never be updated. It is, therefore, a measure of plasticity. This coupling is quite unexpected because the expressivity should be independent of plasticity at initialization. See Figure 7 for a scattering plot of how the elements of $\mathbb{E}[hh^{\top}]$ are strongly aligned with that of $\mathbb{E}[gg^{\top}]$. While this experiment shows a strong alignment between $H$ and $G$, it also shows a limitation of our theory. If the activation norm $h$ is really independent of its direction $h/\|h\|$, one would only see one branch in the figure, whereas there are three branches. This can be another reason for CRH breaking.

### 7.2 Feature-Importance Alignment and Invariant Learning

Another direct interpretation of $\mathbb{E}[gg^{\top}]$ is the saliency of a latent feature. The quantity $\mathbb{E}[gg^{\top}]$ measures how the loss function changes as a neuron activation is changed by a small degree and is often used on the input layer to measure the importance of a feature (Selvaraju et al.,, ; Adadi and Berrada,, ). This interpretation justifies many heuristics for understanding neural networks: take the principal components of a hidden layer and study its largest eigenvalues. If there is no direct link between the magnitude of the eigenvalues in these latent representations and their importance in affecting the learning (measured by the training loss), these analyses will be unreasonable.

Let $\hat{n}$ be any unit vector that encodes a feature irrelevant to the task. The following theorem shows that the invariance of the representation to such features is equivalent to the invariance of the model.

###### Proposition 1 .

Let $f(h(x))$ be a model whose hidden states $h$ obey RGA. Let $\hat{n}$ be a vector and $\epsilon$ a scalar. The following statements are equivalent: (1) $\ell(f(h+\epsilon\hat{n}))=\ell(f(h))+O(\epsilon^{2})$; (2) $\mathbb{E}[hh^{\top}]\hat{n}=0$; (if the CRH also holds) $W\hat{n}=0$, and $G\hat{n}=0$.

###### Corollary 1 .

When CRH holds for a layer with activations $h_{a}$ and $h_{b}$, $f(x)$ is invariant to infinitesimal transforms of $x$ if and only if $h_{b}(x)$ is invariant to infinitesimal transforms of $x$.

In some sense, Proposition 1 can be seen as a generalization of NC (next section) because it applies to both regression and classification tasks. Note that there are two ways for the condition $\ell(f(h+\epsilon\hat{n}))=\ell(f(h))$ to be satisfied: (1) $f(h+\epsilon\hat{n})=f(h)$, which means that when the model output itself invariant to such a variation, the latent representation in this direction will vanish; (2) $\ell(f(h)+\epsilon\hat{n}^{T}\nabla f(h))=\ell(f(h))$, which means that any variation that does not change the loss function value will vanish. This fact also means that the model will learn a compact representation: any irrelevant latent space will have zero variation. This is consistent with the often observed matching between the rank of the representation and the inherent dimension of the problem (Papyan,, ; Ziyin,, ). A major phenomenon of well-trained CNNs is that the latent representations learn to become essentially invariant to task-irrelevant features (Zeiler and Fergus,, ; Selvaraju et al.,, ), an observation that matches the high-level features of the human visual cortex (Booth and Rolls,, ). The CRH thus suggests a reason for these observed invariances.

This result also means that the invariances of the models achieved through the CRH are robust: the objective function is invariant to small noises in the irrelevant activations. While artificial neural networks rarely contain such noises in the activations, it is plausibly the case that biological neurons do suffer from environmental noises, and our result suggests a mechanism to achieve robustness against irrelevant noises or perturbation.

###### Proposition 1 .

###### Corollary 1 .

### 7.3 Reduction to Neural Collapse in Perfect Classification

We now prove that NC is equivalent to the CRH in an interpolating classifier. We focus on the first two properties here because NC3-4 are essentially consequences of NC1-2 (Rangamani and Banburski-Fahey,, ). NC1 states that the inner-class variations of the penultimate representation vanishes. NC2 states that the average representation $\mu_{c}$ of the class $c$ is orthogonal to each other: $\mu_{c}^{\top}\mu_{c^{\prime}}=\delta_{cc^{\prime}}$. The ground truth model must be invariant to inner class variations in a classification task. Let $c$ denote the index of a class, it should be the case that the ground truth model $f^{\prime}$ satisfies: $f^{\prime}_{c}(x_{c})=\zeta\mathbf{1}_{c}$, where $\mathbf{1}_{c}$ is a one hot vector on the $c$-th dimension and $\zeta>0$ is an arbitrary scalar. Proposition 1 thus suggests that any model that recovers the ground truth must have such invariance in the output and in the latent representation $h$, which is exactly NC1.

Now, we show that when the CRH holds, a perfectly trained model must have neural collapse. Let $x_{c}$ denote any data point belonging to the $c$-class among a set of $C$ labels.

###### Theorem 3 .

Consider a classification task and the penultimate layer postactivation $h_{a}$. If the model is quasi-interpolating: $f(x_{c})=h_{b}=Wh_{a}(x_{c})=\zeta\mathbf{1}_{c}$, and the loss covariance is proportional to identity: $\mathbb{E}[\nabla_{f}\ell\nabla_{f}^{\top}\ell]\propto I$, then $h$ satisfies all four properties of neural collapse (NC1-NC4) if and only if $h$ satisfies the CRH.

The assumption $\mathbb{E}[\nabla_{f}\ell\nabla_{f}^{\top}\ell]\propto I$ is empirically found to hold very well in standard classification tasks at the end of training. See Section C.3 for an experiment with Resnet on CIFAR-10, where the phenomenon of neural collapse is primarily observed.

###### Theorem 3 .

### 7.4 Equivariance of CRH and Neural Feature Ansatz

A closely related hypothesis is the NFA, which is related to the feature learning in fully connected layers (Radhakrishnan et al.,, ; Beaglehole et al.,, ). Using our notation, the NFA states $W^{\top}W\propto\mathbb{E}\left[\nabla_{h_{a}}f\nabla_{h_{a}}f^{\top}\right]$.
To see its connection to the CRH, note that it can be written as

$$ $W^{\top}W\propto W^{\top}\mathbb{E}[g_{b}g_{b}^{\top}]W=\mathbb{E}[g_{a}g_{a}^ {\top}]=\mathbb{E}\left[(\nabla_{h_{b}}f)B(\nabla_{h_{b}}f^{\top})\right],$ (13) $$

where $B=\nabla_{f}\ell(\nabla_{f}\ell)^{\top}$.
Therefore, the NFA is identical to the GWA (Eq. (5)) when $B\propto I$.

Figure: Figure 8: After the training of a six-layer eight-head fully connected network, all six relations (Eq. (3)-(5)) hold strongly at a large weight decay (fc2). For a small weight decay, at least one forward and backward relation holds strongly. The shaded region shows the variation across five hidden layers, and the solid lines show the median of these alignments. At a small $\gamma$, the best alignment is between $H_{a}$ and $W^{\top}W$ for the forward relation, and $G_{a}$ and $WW^{\top}$ for the backward, in agreement with the theoretical prediction.
Refer to caption: extracted/5900109/plot/forward_alignment_wd2.png

This is an instance of a consistency problem of the NFA: the NFA is not invariant to trivial redefinitions of the loss function and model. Let $f(x)$ be a model trained on loss function $\ell(f)$. Let us assume that NFA applies to this model. Now, let us construct a trivially redefined model and loss: $f^{\prime}(x):=Zf(x)$ where $Z$ is an invertible matrix, and $\ell^{\prime}(f):=\ell(Z^{-1}f)$. Now, for NFA to hold for $f$, we must have $W^{\top}W\propto\mathbb{E}[\nabla f\nabla^{\top}f]$, but for NFA to hold for $f^{\prime}$, we need $W^{\top}W\propto\mathbb{E}[\nabla fZZ^{\top}\nabla^{\top}f]$. These cannot hold simultaneously.

In contrast, the GWA is invariant to such redefinitions:

$$ $\nabla\ell(f)\nabla\ell(f)=\nabla\ell^{\prime}(f^{\prime})\nabla\ell^{\prime}( f^{\prime}).$ (14) $$

Therefore, the CRH is more likely to be a fundamental law of learning as it is invariant to a subjective choice of basis. For this reason, one may also refer to the GWA as an “equivariant NFA” (eNFA). Furthermore, if we treat NFA as a special case of the forward alignment, the backward alignment relations imply a novel variant of the NFA, which can be referred to as the forward NFA: $WW^{\top}\propto\mathbb{E}\left[\nabla_{h_{b}}g\nabla_{h_{b}}g^{\top}\right]$. In this picture, the original NFA should thus be called the backward NFA. See Figure 8, which validates both forward and backward eNFA.

## 8 Conclusion

In this work, we propose the Canonical Representation Hypothesis (CRH), a new perspective for studying the formation of representations in neural networks. The CRH suggests that representations align with the weights and gradients after training. It is a generalization of the neural collapse phenomenon for any fully connected layer in a neural network. In this view, representations are formed based on the degree and modes of deviation from the CRH. This deviation leads to the Polynomial Alignment Hypothesis (PAH), which posits that when the CRH is broken, distinct phases emerge in which the representations, gradients, and weights become polynomial functions of each other. A key future direction is to understand the conditions that lead to each phase and how these phases affect the behavior and performance of models. The CRH may also have algorithmic implications. If representations align with the gradients (as in RGA), it might be possible to manually inject noise into neuron gradients to engineer specific structures in the model’s representations. However, the CRH and PAH have several limitations. They apply only to fully connected layers, and a future step is extending them to convolutional layers. Additionally, our research has primarily focused on characterizing the final stage of representation formation. A more comprehensive theory of representation dynamics could lead to better training algorithms.

## Acknowledgements

ILC acknowledges support in part from the Institute for Artificial Intelligence and Fundamental Interactions (IAIFI) through NSF Grant No. PHY-2019786. This work was also supported by the Center for Brains, Minds and Machines (CBMM), funded by NSF STC award CCF - 1231216.

## Appendix A More Related Works

The neural representation is a central object of study in both AI and neuroscience (Esser et al.,, ; Wang et al.,, ). The widespread use of techniques like t-SNE (Van der Maaten and Hinton,, ) for analyzing the penultimate layer representation in deep learning highlights the importance of understanding representation space. However, the theory of how the representations are formed in neural networks is scarce. Liu et al., () proposes an effective theory to study representation learning in modular addition. A recent work studies representation formation from the angle of symmetries (Ziyin,, ), showing that permutation symmetries in the latent layers lead to neuron merging and low rankness in the representation. Ziyin et al., () showed that after training, the latent representation of a deep linear net is aligned with a linear transformation of the prediction residual. A recent work showed that the evolution of the representation during the initial stage of training has universal properties shared by different types of models (van Rossem and Saxe,, ). Empirically, Huh et al., () showed that the neural representations learned by different networks are essentially similar. Also, an emergent field in neuroscience studies how the representations learned by neural networks closely resemble those of biological neurons (Rajalingham et al.,, ). These results suggest the existence of some fundamentally shared mechanisms of the formation of representations. Also, neural collapse has been found to improve the transfer learning of neural networks Galanti et al., (), and an interesting future step may be understanding how CRH might improve transfer learning.

## Appendix B Theory and Proofs

### B.1 Gradient Moment is Dominated by Gradient Covariance

We show that when there is a bias term in the layer, the expected neuron gradient must be negligible close to a local minimum:

$$ $h_{b}=W^{\prime}h_{a}^{\prime}=Wh_{a}+\beta,$ (15) $$

At the local minimum, we have

$$ $0=\mathbb{E}[\Delta b]=-\eta(\mathbb{E}[\nabla_{h_{b}}\ell]+\gamma b)=\eta( \mathbb{E}[g_{b}]-\gamma b),$ (16) $$

which implies that

$$ $\gamma b=\mathbb{E}[g_{b}].$ (17) $$

This implies that

$$ $\mathbb{E}[g_{b}]\mathbb{E}[g_{b}^{\top}]=O(\gamma^{2}),$ (18) $$

which is negligible. This agrees with the experimental observation in Section C.2.

### B.2 Proof of Theorem 2

We need a few lemmas.

###### Lemma 1 .

Let $A$, $B$ be symmetric matrices such that $B=ABA$, then,

$$ ${\rm ker}(A)\subseteq{\rm ker}(B)=I-B^{0},$ (19) $$

and $B^{0}AB^{0}$ is a projection matrix.

###### Proof.

Let $n$ be in the null space of $A$. We have

$$ $Bn=ABAn=0.$ (20) $$

Thus, $n$ is also in the null space of $B$. This means that the kernel of $A$ is in the kernel of $B$. Now, Let $P=B^{0}$ and $\tilde{D}:=PDP$, we have that

$$ $B=\tilde{A}B\tilde{A},$ (21) $$

where the rank of $A$ is the same as the rank of $B$. This implies that there is an orthonormal matrix $O$ such that

$$ $B^{\prime}=OBO^{\top},\ A^{\prime}=O\tilde{A}O^{\top}$ (22) $$

are full rank, $A^{\prime}$ is diagonal, and

$$ $B^{\prime}=A^{\prime}B^{\prime}A^{\prime}.$ (23) $$

This implies that

$$ $B^{\prime}_{ij}=a_{i}a_{j}B^{\prime}_{ij},$ (24) $$

where $a_{i}$ is the $i$-th diagonal term of $A^{\prime}$. Because $B^{\prime}$ is full-rank. We have that for all $i$,

$$ $a_{i}=1.$ (25) $$

This implies that $\tilde{A}=O^{\top}A^{\prime}O$ is an orthogonal projection matrix. Because it also has the same kernel as $B$, we have

$$ $\tilde{A}=B^{0}.$ (26) $$

∎

The following two lemmas are straightforward to prove by multiplying $A^{-1}$ from the left and right.(^4^44Recall that $A^{-1}$ is the pseudo inverse, and $A_{0}=AA^{-1}$ is an orthogonal projector.)

###### Lemma 2 .

Let $A$, $B$ be symmetric matrices such that $A^{2}=ABA$, then,

$$ $A^{0}=A^{0}BA^{0}.$ (27) $$

###### Lemma 3 .

Let $A$, $B$ be symmetric matrices such that $A=ABA$, then,

$$ $A^{-1}=A^{0}BA^{0}.$ (28) $$

###### Lemma 4 .

If ${\rm ker}(A)={\rm ker}(B)$, then ${\rm ker}(CAC^{\top})={\rm ker}(CBC^{\top})$, for any symmetric $A$ and $B$ and arbitrary matrix $C$.

###### Proof.

For $n$ to be in the null space of $CAC^{\top}$, it must satisfy one of the following two conditions:

$$ $C^{\top}n=0,$ (29) $$

which implies that $n\in{\rm ker}(CBC^{\top})$. Or

$$ $\sqrt{A}C^{\top}n=0,$ (30) $$

which implies that $C^{\top}n\in{\rm ker}(\sqrt{A})={\rm ker}(A)={\rm ker}(B)$, which again implies that

$$ $n\in{\rm ker}(CBC^{\top}).$ (31) $$

This finishes the proof.
∎

Now, we are ready to prove Theorem 2.

###### Proof.

(Part 1) Fix the direction to be forward or backward. Let $A,B,C$ be a permutation of three moment matrices. Then, that two alignments hold implies that

$$ $A\propto B,$ (32) $$

and $B\propto C$. By transitivity, $A\propto B\propto C$. This proves part 1.

(Part 2) There are many combinations, which can be divided into a few types. We thus prove a prototype for each type of relation, and the rest can be derived in a similar but tedious manner. The key mechanism is that every forward relation implies a backward relation and vice versa. For example, if we know $\mathbb{E}[h_{b}h_{b}^{\top}]\propto\mathbb{E}[g_{b}g_{b}^{\top}]$, we also have

$$ $W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}\propto\mathbb{E}[g_{b}g_{b}^{\top}],$ (33) $$

which implies that

$$ $W^{\top}W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}W\propto W^{\top}\mathbb{E}[g_{b }g_{b}^{\top}]W=\mathbb{E}[g_{a}g_{a}^{\top}].$ (34) $$

Now, let $A$ be either $W^{\top}W$ or $WW^{\top}$, $B,C$ be some permutation of the gradient and neuron covariance.

Type 0: this is the simplest and most straightforward type of relations. First, consider relation 5 in the table. We have, by assumption,

$$ $H_{a}\propto Z_{a},$ (35) $$

$$ $WH_{a}W^{\top}=H_{b}\propto G_{b}.$ (36) $$

Together, we have

$$ $WH_{a}W^{\top}\propto WZ_{a}W^{\top}=Z_{b}^{2}.$ (37) $$

So,

$$ $Z_{a}^{3}\propto W^{\top}G_{b}W=G_{a}.$ (38) $$

This means that

$$ $H_{a}^{3}\propto Z_{a}^{3}\propto G_{a}.$ (39) $$

Also, this implies that

$$ $H_{b}\propto Z_{b}^{2}\propto G_{b}.$ (40) $$

A similar proof derives relation 6 is thus not shown.

Type 1: $ABA\propto ACA\propto B\propto C$. Using the above lemmas and defining $\tilde{D}:=PDP$, where $P=B^{0}=C^{0}$, we obtain

$$ $\tilde{A}=B^{0}=C^{0}.$ (41) $$

Type 2: $ABA\propto A^{2}\propto ACA$. This type of equation simply solves to

$$ $\tilde{B}\propto A^{0}\propto\tilde{C},$ (42) $$

where the tilde is defined with respect to $A^{0}$.

Type 3: $B\propto A\propto ACA$. By the above lemmas, we have

$$ $A^{-1}\propto\tilde{C}.$ (43) $$

We thus have

$$ $B\propto A\propto\tilde{C}^{-1}.$ (44) $$

Type 4: $ABA\propto A^{2}\propto C^{2}$. This implies that

$$ $\tilde{B}\propto A^{0}\propto C^{0},$ (45) $$

where tilde is defined with respect to $P=A^{0}$.

Type 5: This type is a little strange, and so we explicitly solve them. There are two cases:

$$ $\displaystyle\mathbb{E}[g_{a}g_{a}^{\top}]\propto Z_{a}\propto Z_{a}^{2}$ (46) $\displaystyle\mathbb{E}[g_{b}g_{b}^{\top}]\propto Z_{b}\propto Z_{b}^{2}$ (47) $$

This directly implies that $Z_{a}$ and $Z_{b}$ are projection matrices. For $\mathbb{E}[hh^{\top}]$, we have that by definition

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]\propto W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}.$ (48) $$

This means that the kernel of $\mathbb{E}[h_{b}h_{b}^{\top}]$ must contain the kernel of $WW^{\top}=Z_{b}$. Therefore, letting $P=\mathbb{E}[h_{b}h_{b}^{\top}]^{0}$

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]^{0}\propto PZ_{b}P\propto P\mathbb{E}[g_{b}g_{b} ^{\top}]P.$ (49) $$

Likewise, we have that

$$ $W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}=\mathbb{E}[h_{b}h_{b}^{\top}],$ (50) $$

and so

$$ ${\rm ker}(W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top})={\rm ker}(\mathbb{E}[h_{b}h_ {b}^{\top}])={\rm kerZ_{b}}.$ (51) $$

This implies that

$$ ${\rm ker}(Z_{a}\mathbb{E}[h_{a}h_{a}^{\top}]Z_{a})={\rm kerZ_{a}}.$ (52) $$

This means that we have proved

$$ $\tilde{\mathbb{E}[h_{a}h_{a}^{\top}]}^{0}=Z^{0}_{a}=\mathbb{E}[g_{a}g_{a}^{ \top}].$ (53) $$

A similar derivation applies to the case when

$$ $\displaystyle\mathbb{E}[h_{a}h_{a}^{\top}]\propto Z_{a}\propto Z_{a}^{2}$ (54) $\displaystyle\mathbb{E}[h_{b}h_{b}^{\top}]\propto Z_{b}\propto Z_{b}^{2}.$ (55) $$

(Part 3) By the pigeonhole principle, two of either the forward or backward relations must be satisfied. This means that by part 2 of the theorem, there are two cases: (a) three forward relations are satisfied, (b) three backward relations are satisfied. Since the argument is symmetric in forward and backward directions, we focus on case (a).

Case (a). We have

$$ $\mathbb{E}[h_{a}h_{a}^{\top}]\propto Z_{a}\propto\mathbb{E}[g_{a}g_{a}^{\top}] =W^{\top}\mathbb{E}[g_{b}g_{b}^{\top}]W.$ (56) $$

There are now three subcases depending on which of the backward relations are satisfied.

Case (a1): $Z_{b}\propto\mathbb{E}[g_{b}g_{b}^{\top}]$, which implies that

$$ $Z_{a}\propto\mathbb{E}[g_{a}g_{a}^{\top}]=W^{\top}\mathbb{E}[g_{b}g_{b}^{\top} ]W\propto W^{\top}Z_{b}W=Z_{a}^{2},$ (57) $$

and so all the three forward matrices are (scalar multiples of) projections.

But $W^{\top}W$ and $WW^{\top}$ share eigenvalues and so $Z_{b}\propto\mathbb{E}[g_{b}g_{b}^{\top}]$ are also projection matrices. Now,

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]=W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}\propto WZ _{a}W^{\top}=Z_{b}^{2},$ (58) $$

which is also a projection matrix. This, in turn, implies that all three backward relations hold.

Case (a2): $\mathbb{E}[h_{b}h_{b}^{\top}]\propto\mathbb{E}[g_{b}g_{b}^{\top}]$. This implies that

$$ $Z_{b}^{2}=WZ_{a}W^{\top}\propto W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}=\mathbb {E}[h_{b}h_{b}^{\top}]\propto\mathbb{E}[g_{b}g_{b}^{\top}].$ (59) $$

In turn, this implies that

$$ $Z_{a}\propto\mathbb{E}[g_{a}g_{a}^{\top}]=W^{\top}\mathbb{E}[g_{b}g_{b}^{\top} ]W\propto W^{\top}Z_{b}^{2}W=Z_{a}^{3}.$ (60) $$

This implies that $Z_{a}$ only contains zero and one as eigenvalues, which implies that $Z_{a}=Z_{a}^{2}=Z_{a}^{3}$ is a projection. Similarly, $Z_{b}$ is a projection. Together, this implies that all six relations hold.

Case (a3): $Z_{b}\propto\mathbb{E}[h_{b}h_{b}^{\top}]$. This case is subtly different. Here,

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]=W\mathbb{E}[h_{a}h_{a}^{\top}]W^{\top}=WZ_{a}W^{ \top}=Z_{b}^{2}=Z_{b},$ (61) $$

and so all three forward matrices are projections.

For the backward relations,

$$ $Z_{b}\propto\mathbb{E}[h_{b}h_{b}^{\top}]=W\mathbb{E}[h_{a}h_{a}^{\top}]W^{ \top}\propto Z_{b}^{2},$ (62) $$

which is also a projection matrix.
Now, we have essentially no relation for $\mathbb{E}[g_{b}g_{b}^{\top}]$ as it cannot be directly derived from any other quantity. But noting that

$$ $W^{\top}\mathbb{E}[g_{b}g_{b}^{\top}]^{\top}=\mathbb{E}[g_{a}g_{a}^{\top}] \propto Z_{a},$ (63) $$

one obtains that

$$ $P\mathbb{E}[g_{b}g_{b}^{\top}]P=P,$ (64) $$

where $P=Z_{b}$. Therefore, when restricted to the subspace of $W$, $\mathbb{E}[g_{b}g_{b}^{\top}]$ is also a projection matrix.

For case (b), it is the same, except for the case when $Z_{a}\propto\mathbb{E}[g_{a}g_{a}^{\top}]$. When this is the case, the alignment happens with

$$ $P\mathbb{E}[h_{a}h_{a}^{\top}]P=P,$ (65) $$

for $P=Z_{a}^{0}$.

At local minimum. At any local minimum or first-order stationary point, Lemma 5 applies, which implies that

$$ ${\rm ker}(Z_{c})\subseteq{\rm ker}(\mathbb{E}[h_{c}h_{c}^{\top}]),$ (66) $$

$$ ${\rm ker}(Z_{c})\subseteq{\rm ker}(\mathbb{E}[g_{c}g_{c}^{\top}]),$ (67) $$

for $c\in\{a,b\}$. This is because

$$ $Z_{c}\propto\mathbb{E}[g_{c}h_{c}^{\top}],$ (68) $$

and for any $n\in{\rm ker}(\mathbb{E}[h_{c}h_{c}^{\top}])$, it holds with probability one that

$$ $h_{c}^{\top}n=0,$ (69) $$

and so

$$ $Z_{c}n=\mathbb{E}[h_{c}g_{c}^{\top}]n=0.$ (70) $$

The same argument applies to $g_{c}$.

Now, it suffices to prove the two subtle cases. For case (a3), we have that

$$ $P\mathbb{E}[g_{b}g_{b}^{\top}]P=P.$ (71) $$

However, the local minimum condition implies that

$$ ${\rm ker}\mathbb{E}[g_{b}g_{b}^{\top}]\subseteq{\rm ker}(P).$ (72) $$

This can only happen if $P\mathbb{E}[g_{b}g_{b}^{\top}]P=\mathbb{E}[g_{b}g_{b}^{\top}]=P$. Therefore, all six relations hold. The same applies to the other subtle case.

(Part 4) We have

$$ $\mathbb{E}[h_{a}h_{a}^{\top}]\propto Z_{a}\propto\mathbb{E}[g_{a}g_{a}^{\top}] =W^{\top}\mathbb{E}[g_{b}g_{b}^{\top}]W,$ (73) $$

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]=W\mathbb{E}[h_{a}h_{a}^{\top}]W\propto Z_{b} \propto\mathbb{E}[g_{b}g_{b}^{\top}].$ (74) $$

Plugging the forward relation into the backward relation, we obtain that

$$ $WZ_{a}W^{\top}\propto Z_{b}^{2}\propto Z_{b},$ (75) $$

which implies that $Z_{b}$ is proportional to a projection matrix. This implies that $\mathbb{E}[h_{b}h_{b}^{\top}]$ and $\mathbb{E}[g_{b}g_{b}^{\top}]$ are also projection matrices.

Likewise, one can plug the backward relation into the forward, which implies that

$$ $Z_{a}^{2}\propto Z_{a},$ (76) $$

and is thus proportional to a projection matrix. This completes the proof.
∎

###### Lemma 1 .

###### Proof.

###### Lemma 2 .

###### Lemma 3 .

###### Lemma 4 .

###### Proof.

###### Proof.

### B.3 Proof of Theorem 1

The proof will be divided into several theorems, each of which proves a claimed relation in Section 4. The following theorem proves Eq. 6.

###### Theorem 4 .

Assume assumption 1. If $\mathbb{E}[\Delta(h_{a}h_{a}^{\top})]=0$ and $\mathbb{E}[\Delta(h_{b}h_{b}^{\top})]=0$, then,

$$ $0=z_{b}\mathbb{E}[g_{b}h_{b}^{\top}]+z_{b}\mathbb{E}[h_{b}g_{b}^{\top}]+\eta z _{b}^{2}\mathbb{E}[g_{b}g_{b}^{\top}]-2\gamma\mathbb{E}[\Delta(h_{b}h_{b}^{ \top})]+O(\eta^{2}\gamma),$ (77) $$

where $z_{b}=\mathbb{E}\|h_{a}\|^{2}$.

###### Proof.

Let $B_{a}=h_{a}h_{a}^{\top}$, and $B_{b}=h_{b}h_{b}^{\top}$. Let $H_{a}=\mathbb{E}[B_{b}]$ and $H_{b}=\mathbb{E}[B_{a}]$. We consider an online learning setting. Let $x$ denote the input at time step $t$. Consider the time evolution of the feature covariance during the SGD training:

$$ $\displaystyle\Delta$ $\displaystyle(h_{b}(x)h_{b}^{\top}(x))=\Delta(Wh_{a}h_{a}^{\top}W^{\top})$ (78) $\displaystyle=\Delta WB_{a}W^{\top}+{W}B_{a}\Delta{W}^{\top}+\Delta WB_{a} \Delta W^{\top}+O(\Delta B_{a})$ (79) $\displaystyle=-\eta(\nabla_{h_{b}}\ell h_{a}^{\top}+\gamma W)B_{a}W^{\top}- \eta WB_{a}(h_{a}\nabla_{h_{b}}^{\top}\ell+\gamma W^{\top})+\eta^{2}(\nabla_{h _{b}}\ell h_{a}^{\top}+\gamma W)B_{a}(h_{a}\nabla_{h_{b}}^{\top}\ell+\gamma W^ {\top})$ (80) $\displaystyle=-\eta(-\|h_{a}\|^{2}g_{b}h_{b}^{\top}+\gamma B_{b})-\eta(-\|h_{a }\|^{2}h_{b}g_{b}^{\top}+\gamma B_{b})+\eta^{2}\|h_{a}\|^{4}g_{b}g_{b}^{\top}+ O(\eta^{2}\gamma)$ (81) $\displaystyle=\eta(\|h_{a}\|^{2}g_{b}h_{b}^{\top}+\|h_{a}\|^{2}h_{b}g_{b}^{ \top}-2\gamma B_{b})+\eta^{2}\|h_{a}\|^{4}g_{b}g_{b}^{\top}.$ (82) $$

Taking expectation of both sides, we obtain that

$$ $0=z_{b}\mathbb{E}[g_{b}h_{b}^{\top}]+z_{b}\mathbb{E}[h_{b}g_{b}^{\top}]+\eta z ^{2}\mathbb{E}[g_{b}g_{b}^{\top}]-2\gamma H_{b},$ (83) $$

where $z_{b}=\mathbb{E}\|h_{a}\|^{2}$.
∎

Similarly, one can show that the pre-activations follow a similar relationship.

###### Theorem 5 .

Assume assumption 1. If $\mathbb{E}[\Delta(g_{a}g_{a}^{\top})]=0$ and $\mathbb{E}[\Delta(g_{b}g_{b}^{\top})]=0$, then,

$$ $z_{a}(\mathbb{E}[h_{a}g_{a}^{\top}]+\mathbb{E}[g_{a}h_{a}^{\top}])+\eta z_{a}^ {2}\mathbb{E}[h_{a}h_{a}^{\top}]=2\gamma\mathbb{E}[g_{a}g_{a}^{\top}],$ (84) $$

where we have defined $z_{a}=\mathbb{E}[\|g_{b}\|^{2}]$

###### Proof.

We have

$$ $\displaystyle\Delta(g_{a}g_{a}^{\top})$ $\displaystyle=\Delta(W^{\top}g_{b}g_{b}^{\top}W)$ (85) $\displaystyle=\Delta W^{\top}g_{b}g_{b}^{\top}W+W^{\top}g_{b}g_{b}^{\top} \Delta W+\Delta W^{\top}g_{b}g_{b}^{\top}\Delta W+O(\|\Delta g_{b}g_{b}^{\top}\|)$ (86) $\displaystyle=\eta(z_{a}h_{a}g_{a}^{\top}-\gamma g_{a}g_{a}^{\top})+\eta(z_{a} h_{a}g_{a}^{\top}-\gamma g_{a}g_{a}^{\top})^{\top}+\eta^{2}z_{a}^{2}h_{a}h_{a} ^{\top}+O(\eta^{2}\gamma),$ (87) $$

where we have used the relations $g_{a}=W^{\top}g_{b}$ and $\Delta W=\eta(g_{b}h_{a}^{\top}-\gamma W)$. We have also defined $z_{a}=\|g_{b}\|^{2}=\mathbb{E}[\|g_{b}\|^{2}]$.

Taking expectation, this leads to

$$ $z_{a}(\mathbb{E}[h_{a}g_{a}^{\top}]+\mathbb{E}[g_{a}h_{a}^{\top}])+\eta z_{a}^ {2}\mathbb{E}[h_{a}h_{a}^{\top}]=2\gamma\mathbb{E}[g_{a}g_{a}^{\top}].$ (88) $$

∎

The following lemma proves the balance condition at a local minimum.

###### Lemma 5 .

(First-order stationary point condition / local minimum balance.) At any stationary point of the loss function, we have

$$ $\mathbb{E}[g_{b}h_{b}^{\top}]=\mathbb{E}[g_{b}h_{b}^{\top}]^{\top}=\gamma WW^{ \top},$ (89) $$

$$ $\mathbb{E}[g_{a}h_{a}^{\top}]=\mathbb{E}[g_{a}h_{a}^{\top}]^{\top}=\gamma W^{ \top}W.$ (90) $$

###### Proof.

Close to a stationary point, we have

$$ $0=\mathbb{E}[\Delta W]=\eta(\mathbb{E}[g_{b}h_{a}^{\top}]-\gamma W).$ (91) $$

Multiplying $W^{\top}$ from the left and the right, we obtain, respectively:

$$ $0=\mathbb{E}[g_{b}h_{a}^{\top}]W^{\top}-\gamma WW^{\top}=\mathbb{E}[g_{b}h_{b} ^{\top}]-\gamma WW^{\top},$ (92) $$

$$ $0=W^{\top}\mathbb{E}[g_{b}h_{a}^{\top}]-\gamma W^{\top}W=\mathbb{E}[g_{a}h_{a} ^{\top}]-\gamma W^{\top}W,$ (93) $$

where we have used the definition $h_{b}=Wh_{a}$ and the chain rule $g_{a}=W^{\top}g_{b}$.

This condition implies that

$$ $\mathbb{E}[g_{b}h_{b}^{\top}]=\mathbb{E}[g_{b}h_{b}^{\top}]^{\top}=\gamma WW^{ \top},$ (94) $$

$$ $\mathbb{E}[g_{a}h_{a}^{\top}]=\mathbb{E}[g_{a}h_{a}^{\top}]^{\top}=\gamma W^{ \top}W.$ (95) $$

∎

Now, we derive the stationarity condition for $\Delta WW^{\top}$ and $\Delta W^{\top}W$.

###### Lemma 6 .

(Stationary alignment of parameter-outer product.) If at a local minimum and assuming Assumption 1, then,

- 1.
if $\mathbb{E}[\Delta(WW^{\top})]=0$,
$z_{b}\mathbb{E}[g_{b}g_{b}^{\top}]=\gamma^{2}WW^{\top}.$
(96)
- 2.
if $\mathbb{E}[\Delta(W^{\top}W)]=0$,
$z_{a}\mathbb{E}[h_{a}h_{a}^{\top}]=\gamma^{2}W^{\top}W.$
(97)

###### Proof.

We have

$$ $\displaystyle 0=\mathbb{E}[\Delta(WW^{\top})]$ $\displaystyle=\mathbb{E}[\Delta W]W^{\top}+W\mathbb{E}[\Delta W^{\top}]+ \mathbb{E}[\Delta W\Delta W^{\top}]$ (98) $\displaystyle=0+\mathbb{E}[\Delta W\Delta W^{\top}]$ (99) $\displaystyle=\mathbb{E}[(g_{b}h_{a}^{\top}-\gamma W)(g_{b}h_{a}^{\top}-\gamma W )^{\top}]$ (100) $\displaystyle=z_{b}\mathbb{E}[g_{b}g_{b}^{\top}]-\gamma\mathbb{E}[g_{b}h_{b}^{ \top}]-\gamma\mathbb{E}[g_{b}h_{b}^{\top}]^{\top}+\gamma^{2}WW^{\top}$ (101) $\displaystyle=z_{b}\mathbb{E}[g_{b}g_{b}^{\top}]-\gamma^{2}WW^{\top}$ (102) $$

where we have used Eq. (89) in the last line.

Likewise, when $\mathbb{E}[\Delta W^{\top}W]=0$,

$$ $\displaystyle 0=\mathbb{E}[\Delta(W^{\top}W)]$ $\displaystyle=\mathbb{E}[(g_{b}h_{a}^{\top}-\gamma W)^{\top}(g_{b}h_{a}^{\top} -\gamma W)]$ (103) $\displaystyle=z_{a}\mathbb{E}[h_{a}h_{a}^{\top}]-\gamma\mathbb{E}[h_{a}g_{a}^{ \top}]-\gamma\mathbb{E}[h_{a}g_{a}^{\top}]^{\top}+\gamma^{2}W^{\top}W$ (104) $\displaystyle=z_{a}\mathbb{E}[h_{a}h_{a}^{\top}]-\gamma^{2}W^{\top}W,$ (105) $$

where we have defined $z_{a}=\mathbb{E}[\|g_{b}\|^{2}]$.

∎

Now, we are ready to prove Theorem 1.

###### Proof.

The above two lemmas imply that

$$ $\frac{2z_{b}}{\gamma}\mathbb{E}[g_{b}g_{b}^{\top}]=\mathbb{E}[g_{b}h_{b}^{\top }]+\mathbb{E}[g_{b}h_{b}^{\top}]^{\top}.$ (106) $$

This relation can be substituted into Eq. (6) to obtain

$$ $z_{b}^{2}(\eta\gamma+2)\mathbb{E}[g_{b}g_{b}^{\top}]=2\gamma^{2}H_{b}=z_{b}( \eta\gamma+2)\gamma^{2}WW^{\top}.$ (107) $$

Likewise,

$$ $\frac{2z_{a}}{\gamma}\mathbb{E}[h_{a}h_{a}^{\top}]=\mathbb{E}[h_{a}g_{a}^{\top }]+\mathbb{E}[h_{a}g_{a}^{\top}]^{\top}.$ (108) $$

This relation can be plugged into Eq. (84) to obtain that

$$ $z_{a}^{2}(\eta\gamma+2)\mathbb{E}[h_{a}h_{a}^{\top}]=2\gamma^{2}\mathbb{E}[g_{ a}g_{a}^{\top}]={z_{a}(\eta\gamma+2)}\gamma^{2}W^{\top}W.$ (109) $$

This completes the proof.
∎

###### Theorem 4 .

###### Proof.

###### Theorem 5 .

###### Proof.

###### Lemma 5 .

###### Proof.

###### Lemma 6 .

###### Proof.

###### Proof.

### B.4 Proof of Proposition 1

###### Proof.

We first prove that part (1) implies part (2). By assumption, we have that $\mathbb{E}[hh^{\top}]\propto\mathbb{E}[gg^{\top}]$, and so(^5^55Note that the proof still works if we replace the second moments with the covariances.)

$$ $\displaystyle E[hh^{\top}]\hat{n}$ $\displaystyle\propto\mathbb{E}[gg^{\top}]\hat{n}$ (110) $\displaystyle=\mathbb{E}[\nabla_{h}\ell\nabla_{h}^{\top}\ell]\hat{n}$ (111) $\displaystyle=0,$ (112) $$

where the last line follows from the fact that

$$ $0=\ell(f(h+\epsilon\hat{n}))-\ell(f(h))=\epsilon\hat{n}^{\top}\nabla_{h}\ell(h).$ (113) $$

One can derive the other two relations simply using the definition of CRH: $\mathbb{E}[hh^{\top}]\propto\mathbb{E}[gg^{\top}]\propto Z$.

For the backward direction, we have that

$$ $E[hh^{\top}]\hat{n}=\mathbb{E}[gg^{\top}]\hat{n}=0.$ (114) $$

This implies that $n^{\top}g=0$ with probability $1$. Now,

$$ $\ell(f(h+\epsilon\hat{n}))-\ell(f(h))=\epsilon\hat{n}^{\top}g+O(\epsilon).$ (115) $$

The proof is complete.

∎

###### Proof.

### B.5 Proof of Corollary 1

###### Proof.

We first prove the forward direction: $f(x)$ is invariant implies $h(x)$ is invariant. Assume that $f(x)$ is invariant to a change in $x$: $x\to x+\epsilon\hat{n}$, where $\hat{n}$ is a vector and $\epsilon$ a scalar. This implies that

$$ $f(x+\epsilon\hat{n})=f(x)+O(\epsilon^{2}).$ (116) $$

Now, if $h_{a}(x+\epsilon\hat{n})=h_{a}(x)+O(\epsilon^{2})$, we are done. Thus, it suffices to consider the case where

$$ $h_{a}(x+\epsilon\hat{n})=h_{a}^{\prime}(x)+\epsilon v+O(\epsilon^{2}),$ (117) $$

for some vector $v\neq 0$. But because of Eq. (116), we have

$$ $\ell(h_{a}+\epsilon v+O(\epsilon^{2}))=\ell(h_{a})+O(\epsilon^{2}).$ (118) $$

By the CRH and Proposition 1, this implies that

$$ $\mathbb{E}[h_{a}h_{a}^{\top}]v\propto\mathbb{E}[W^{\top}W]v=0.$ (119) $$

This implies that

$$ $Wv=0,$ (120) $$

and so

$$ $h_{b}=W(h_{a}+v)=0.$ (121) $$

The reverse direction is trivially true. This completes the proof.
∎

###### Proof.

### B.6 Proof of Theorem 3

As NC4 is a trivial consequence of NC1-3, we focus on NC1-3 here. For notational simplicity, we state neural collapse in the case when there is no bias in the last layer. The first three properties are defined as, at the end of training

- 1.
NC1: $h(x_{c})=\mu_{c}$, where $x_{c}$ is any data point in class $c$;
- 2.
NC2: $\mu^{\top}_{c}\mu_{c^{\prime}}=\delta_{cc^{\prime}}$;
- 3.
NC3: $W^{\top}W=\sum_{c}^{C}\mu_{c}\mu_{c}^{\top}$.

###### Proof.

We first prove that NC1-4 implies the CRH. When NC1 holds,

$$ $h_{a}(x_{c})=\mu_{c}.$ (122) $$

This means that

$$ $H_{a}\propto\sum_{c}\mu_{c}\mu_{c}^{\top}.$ (123) $$

By NC2, we have

$$ $\mu_{c}^{\top}\mu_{c^{\prime}}=\delta_{cc^{\prime}}.$ (124) $$

This means that $H_{a}$ is proportional to an orthogonal projection.

By NC3, we have that

$$ $W^{\top}W\propto\sum_{c}\mu_{c}\mu_{c}^{\top}.$ (125) $$

Additionally,

$$ $G_{a}=W^{\top}G_{b}W=W^{\top}W,$ (126) $$

where we have used the assumption that $G_{b}=\mathbb{E}[\nabla_{f}\ell\nabla_{f}^{\top}\ell]=0$.

Together, this implies the backward CRH

$$ $H_{a}\propto G_{a}\propto Z_{a}.$ (127) $$

For the forward CRH, because $WW^{\top}\in\mathbb{R}^{C\times C}$ has rank $C$ and they have equal eigenvalues, it must be proportional to identity:

$$ $Z_{b}\propto I.$ (128) $$

By the interpolation assumption, we also have

$$ $\mathbb{E}[h_{b}h_{b}^{\top}]\propto\sum_{c}^{C}\mathbf{1}_{c}\mathbf{1}_{c}^{ \top}\propto I.$ (129) $$

Therefore, we have proved the forward CRH. This proves one direction of the theorem.

Now, we prove that the CRH implies NC1-NC3.
We first prove NC1. By the (backward) alignment hypothesis for the last layer, we have

$$ $H_{a}\propto G_{a}=W^{\top}G_{b}W=W^{\top}=W^{\top}W,$ (130) $$

where $W$ is the weight matrix for the last layer. This means that there exists an orthonormal matrix $U$ such that

$$ $W\propto U\sqrt{H_{a}},$ (131) $$

and that the rank of $H_{a}$ must be $C$. By the interpolation hypothesis, it must be the case that

$$ $Wh_{c}=\mathbf{1}_{c},$ (132) $$

which implies that for a fixed $c$. This implies that

$$ $h_{c}=z_{c}+v,$ (133) $$

where $z_{c}$ is a constant vector and $Wv=0$. However, by Proposition 1, $v$ must vanish, and so $h_{c}=z_{c}=\mu_{c}$. This proves NC1.

This means that

$$ $Z_{a}\propto H_{a}$ (134) $$

is essentially a projection matrix. So,

$$ $Wh_{a}(x_{c})=UH_{a}h_{a}(x_{c})=Uh_{a}(x_{c})=\mathbf{1}_{c}.$ (135) $$

This implies that in turn, $h_{a}(x_{c})=\mu_{c}=U_{c:}$. This proves NC3. Due to the orthogonality of $U$, we have that

$$ $\mu_{c}^{\top}\mu_{c^{\prime}}=U_{c:}^{\top}U_{c^{\prime}:}=\delta_{cc^{\prime }}.$ (136) $$

This proves NC2. The proof is complete.
∎

###### Proof.

## Appendix C Experiments

### C.1 Experimental Details

#### fc1:

Fully connected neural networks trained on a synthetic dataset that we generated using a two-layer teacher network. This experiment is used for a controlled study of the effect of different hyperparameters. To control the variables, we consider a synthetic task where the input $x\in\mathbb{R}^{100}$ is sampled from an isotropic Gaussian distribution, and the label generated by a nonlinear function of the form: $y(x)=\sum_{i=1}^{100}u^{*}_{i}\sin((W_{i}^{*})^{\top}x+b_{i}^{*})\in\mathbb{R}$, where $u^{*}$, $w_{i}^{*}$, and $b_{i}^{*}$ are fixed variables drawn from a Gaussian distribution. In this form, the target function can be seen as a two-layer network with sin activation. Our model is a fully connected neural network trained with SGD in an online fashion, and the representations are computed with unseen data points. Unless specified to be the independent variable, the controlled variables of the experiment are: depth of the network ($D=4$), the width of the network ($d=100$), weight decay strength ($\gamma=2\times 10^{-5}$), minibatch size ($B=100$).

#### fc2:

Here, we choose a high-dimensional setting where the teacher net is given by $y(x)=\sum_{i=1}^{100}u^{*}_{i}\sin((W_{i}^{*})^{\top}x+b_{i}^{*})\in\mathbb{R}$, where $u_{i}^{*}\in\mathbb{R}^{100}$, $w_{i}^{*}\in\mathbb{R}^{100}$, and $b_{i}^{*}\in\mathbb{R}$ are fixed variables drawn from a Gaussian distribution. In this form, the target function can be seen as a two-layer network with sin activation. The distribution of $x$ is controlled by an independent variable $\phi_{x}$ such that $x^{\prime}\sim\mathcal{N}(0,I_{100})$, and $x=(1-\phi_{x})Z+\phi_{x}I_{100}$, where $Z$ is a zero-one matrix generated by a Bernoulli distribution with probability $0.8$. When $\phi_{x}$ is small, the input features are thus highly correlated to each other, and the covariance matrix deviates far from $I$. The training proceeds with SGD with a learning rate of $0.1$ with momentum $0.9$ and $\gamma=10^{-4}$ for $10^{5}$ steps when the loss function value has stopped decreasing. The training proceeds with a batch size of $100$. All the expectation and covariance matrices are estimated using $3000$ independently sampled data points. The trained model is a $5$-hidden layer fully connected network with the ReLU activation.

#### res1:

ResNet-18 ($11$M parameters) for CIFAR-10; we apply the standard data augmentation techniques and train with SGD with a learning rate $0.01$, momentum $0.9$, cosine annealing for 200 epochs, and batch size $128$. The model has four convolutional blocks followed by two fully connected layers with ReLU activations. The model has $11$M parameters and achieves $94\%$ test accuracy after training, in agreement with the standard off-the-shelf ResNet-18 for the dataset.

#### res2:

ResNet-18 for self-supervised learning tasks with the CIFAR-10/100 datasets. The model is the same as res1, except that the last fc layer output becomes $128$-dimensional, which is known as the projection dimension in SSL. The training follows the default procedure in the original paper (Chen et al.,, ), proceeding with a batch size of $512$ and $\gamma=5\times 10^{-5}$ for $1000$ epochs.

#### llm:

a six-layer transformer ($100$M parameters) trained on the OpenWebText (OWT) dataset (Gokaslan and Cohen,, ); the number of parameters of this model matches the smallest version of GPT2. The model has six layers with eight heads per layer, having $100$M trainable parameters in total. For each experiment, we train with Adam with a weight decay strength of $1\times 10^{-4}$ for $10^{5}$ iterations, when the training loss stops changes significantly. Since every representation has three dimensions: data $N$, token $T$, and feature $F$, we treat each token as if they are a separate sample in computing the covariances. Namely, we contract the representation tensor along the data and token dimension, resulting in a $F\times F$ covariance matrix.

### C.2 Second Moment Alignments

This section shows the results for the alignment of the matrices $\mathbb{E}[hh^{\top}]$ and $\mathbb{E}[gg^{\top}]$. See Figure 9. The results are qualitatively similar to the result for the alignment between ${\rm cov}(h,h)$ and ${\rm cov}(g,g)$, but with a larger variation. The reason for the similarity is that it is often the case that the covariance term dominates the second moments at the end of training. See Figure 10.

Figure: Figure 9: Alignment of the matrices $\mathbb{E}[hh^{\top}]$ and $\mathbb{E}[gg^{\top}]$. The experimental setting is the same as the LLM experiment.
Refer to caption: extracted/5900109/plot/nocov_first_alignment_0.png

Figure: Figure 10: The ratios of the traces of the second moment matrices for transformer: $r_{g}={\rm Tr}[\mathbb{E}[g]\mathbb{E}[g^{\top}]]/{\rm Tr}[\mathbb{E}[gg^{\top }]]$ and $r_{g}={\rm Tr}[\mathbb{E}[h]\mathbb{E}[h^{\top}]]/{\rm Tr}[\mathbb{E}[hh^{\top }]]$. We see that $r_{g}$ essentially converges to zero, which means that $\mathbb{E}[gg^{\top}]={\rm cov}(g,g)$ at the end of training. $r_{h}$ is generally non-zero but is essentially negligible. The experiment setting is the same as the LLM experiments.
Refer to caption: extracted/5900109/plot/gpt/norm_ratio.png

### C.3 Representations of Resnet18

See Figure 12 for the representations of the last convolutional layer of Resnet18 before and after training on CIFAR-10. See Figure 12 for the representations of the output layer. Interestingly, for the classification task, both ${\rm cov}(g,g)$ and ${\rm cov}(h,h)$ become proportional to the identity.

Figure: Figure 11: Representation of the output layer of Resnet18. Essentially, these are the covariances of the output at the end of training. First Row: Initialization. Second Row: End of training.
Refer to caption: extracted/5900109/plot/resnet_rep_conv_init_gg_4.png

Figure: Figure 13: Examples of representation for Resnet-18 after a self-supervised contrastive training. Left: second convolution block representation, Right: penultimate convolution block representation.
Refer to caption: extracted/5900109/plot/simclr_rep_conv_end_gg_0.png

### C.4 Representations of Self-Supervised Learning

See Figure 13 for the representations of the last and the penultimate convolutional layers. They have significant alignments, but the agreement is perfect. For fully connected layers, the alignment is much better (see the main text; examples not shown).

### C.5 Large Language Model

See Figure 14 and 15.

Figure: Figure 14: Examples of the representation learned by the transformer in the third layer. Left-Mid: examples. Right: The spectra of the two matrices are exactly the same for the leading eigenvalues. The difference is mainly in the smaller eigenvalues, and this difference gets smaller as the training proceeds.
Refer to caption: extracted/5900109/plot/gpt/gpt_rep_end_gg_1_30.png

Figure: Figure 15: The alignment of feature and gradient covariance ($\alpha_{gg,hh}$) remains high during most of the training (llm). The shaded region shows the variation across 8 different heads in the same layer.
Refer to caption: extracted/5900109/plot/gpt/first_alignment_intro.png

### C.6 Fully Connected Nets

See Figure 16 and 17. We see that the alignment effect is significant for both SGD and Adam. Also, see Figure 18 for the effect of having different depths.

Figure: Figure 16: The effect of the width, weight decay $\gamma$, and batch size on the alignment for a fully connected network. The training proceeds for $10^{5}$ iterations, when the training stops decreasing significantly. The training proceeds with SGD.
Refer to caption: extracted/5900109/plot/width_fc_sgd.png

Figure: Figure 17: Same as the previous figure, except that the training proceeds with Adam.
Refer to caption: extracted/5900109/plot/width_fc.png

Figure: Figure 18: Alignment of of different layers of a fully connected ReLU network at different layers (layer 0 0 is the input layer). Left: a 6-layer network. Right: a 12-layer network.
Refer to caption: extracted/5900109/plot/depth_fc_shallow_t_-1.png

### C.7 CRH in ResNet-18

See Figure 1.

Figure: Figure 19: Evolution of $\alpha$ during training for different layers of Resnet-18 on CIFAR-10. For reference, the training accuracy (grey) and testing accuracy (black) are shown in the dashed line. Left to Right: (1) penultimate convolution block representation, (2) last convolution block representation, (3) penultimate fully connected layer, (4) output layer.
Refer to caption: extracted/5900109/plot/resnet_alpha_1.png

### C.8 CRH and PAH in Fully Connected Nets

See Figure 20. The task is the same as other fully connected net experiments. The model is a 4-hidden-layer tanh net with the same width.

Figure: Figure 20: The alignment scalings in fully connected nets. The dashed lines show power laws with exponents $1/3,\ 1/2,\ 1,\ 2,\ 3$, respectively.
Refer to caption: extracted/5900109/plot/fc_scaling_0.png

### C.9 Stationarity of the covariance matrix

See Figure 21 for the evolution of $\Delta{\rm cov}(h,h)$ to zero for three convolutional layers (0-2) and the two fully connected layers (3-4) in a Resnet-18 during training.

Figure: Figure 21: The change in the representation covariance converges to (near) zero at the end of training. The model is Resnet-18 trained on CIFAR-10 with standard SGD.
Refer to caption: extracted/5900109/plot/convergence_resnet.png

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: Please refer to the weakness section.
soundness: 3
strengths: The CRH provides a new framework to interpret basic statistical aspects of how neural networks develop representations. The system of CRH equations generalize and organize into a coherent view previous, partial attempts to reveal alignment phenomena between representations, gradients and weights.  By relating CRH to neural collapse, the Authors offer a broader view that positions this phenomenon as a specific case of a more universal behavior in neural networks. 

The framework is very general, focusing on alignment rather than on specific network designs or loss functions. This  increases the potential impact of CRH on future studies, as it provides a general model for understanding basilar statistical aspects of representation formation.

The Authors present a detailed theoretical derivation supporting the CRH. The theoretical derivations adds credibility to their claims, encouraging further studies, in particular related to dynamics.
summary: First of all, I apologize for the late review with the Authors, the AC, Senior AC and the PC.

This paper introduce the Canonical Representation Hypothesis (CRH), positing that in neural networks, as an effect of training, the representations, weights, and gradients (with respect to the representations) within a layer align with each other. 

This alignment is seen as an essential part of how neural networks form structured, invariant representations. 
The Authors introduce the CRH and support it with a theory that shows how this alignment emerge as a consequence of stationarity towards the end of training, through a balance of gradient noise and regularization.   

Additionally, they propose the Polynomial Alignment Hypothesis (PAH), which describes patterns in cases where CRH is broken, leading to specific scaling relationships in neural networks, of a different nature of scaling laws involving network performances.

This framework offers a new perspective on understanding important aspects of neural representations, unifying phenomena such as neural collapse and feature learning.
weaknesses: While this work aims at universality, the Authors are the first to recognize that their focus is on specific types of networks and tasks, commenting that further empirical tests across more diverse datasets and architectures are necessary to confirm the CRH and PAH in practice.

The weakness I found is that certain assumptions - which are critical for the theoretical foundation of CRH - may not hold in general, like the self-averaging conditions (Hypothesis A.1, A.2). I think the work woud benefit for an empirical assessment on how severe is the departure from self-averaging in a concrete, relevant case.

Furthermore, I found particularly difficult to follow Section B (THEORY AND PROOFS) of the Supplementary Materials. For example, I could not figure out, Lemma 1 in Section B, mainly because some of the "actors" in the scene (like the A tilde in eq. 21) were not presented first. But it may be also a lack of understanding from my side.

However, adding a few more explanations in this section would improve readibility.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: - Can you explain the limitation of the theory exposed by Figure 7? Why do you believe that you see the three branches as opposed to one?
soundness: 4
strengths: - The paper is well written, and carefully presented to ensure correctness and clarity.
- The framework appears general, encompassing significant past work. 
- The theorems are general, avoiding reference to any specific loss function or activation functions, helping to explain their ubiquity. 
- The experimental results appear to strongly support the developed theory and appear both rigorous and extensive.
- The insights section is very much welcomed, and provides significant value to readers coming from outside the primary domain. 
- The limitations are clearly stated.
summary: The authors propose an overarching theory for neural network learning which encompasses the alignment of representations, gradients, and weights throughout training. This framework uniquely combines a diverse set past work under a unified 'canonical representation hypothesis' (CRH) umbrella, and prove theoretically that the CRH does hold. Succinctly, they show that gradient nose 'expands the representation' and weight decay 'contracts it', and they must reach stationarity by balancing these terms. They then use this framework to make predictions about neural network training when each of the six elements of the CRH is broken, provided a set of 64 potential phases for neural network training.  

In conclusion,the paper is well written, organized, and situated in the literature. While I am not an expert in this domain, I therefore defer to other more informed reviewers to make statements about the novelty and accuracy of the authors' claims. Otherwise, the results appear impressive and impactful for the theoretical neural network community, and thus carry significant value.
weaknesses: - The authors state that the CRH holds better for later layers than earlier layers.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - How might different initialization strategies impact the different observed alignment phenomena?
- Could the authors elaborate on how the phases correspond to observable behavioural or performance changes in the model? 
- Beyond the two mechanisms identified in the paper, could other factors contribute to the disruption of alignment? 
- How might the different alignments discussed in this paper relate to concepts from the rich and lazy learning literature? My understanding is that the neural feature ansatz describes and quantifies the rich regime by analyzing feature evolution. Could these alignments and phases potentially provide a way to further describe and quantify distinct learning regimes ( Specifically further define the different types of rich regimes)?

I hope this review will be helpful in the further development of the paper. I encourage the author to continue this research and incorporate the feedback provided. With improvements in structure, clarification of key concepts, and development of the practical significance of these hypotheses, this paper has the potential to earn a higher score.
soundness: 3
strengths: First of all, congratulations on your ICLR submission! I’d like to commend the author for this great piece of research. Below is my review of the work:

**Originality**  
This paper introduces a unifying framework for studying neural network representations, with a notable emphasis on backward representation—a relatively unexplored area. 

**Quality**  
The paper is of generally high quality:
- It acknowledges its own limitations effectively.
- Experimental results are detailed and well-aligned with theoretical insights.
- Findings show that the theory holds across different types of networks.
- Two Mechanisms that disrupt alignment are identified. Quantification is provided via the Polynomial Alignment Hypothesis.

**Clarity**  
The paper is clearly written in good English and is well-structured. Figures are clear, informative, and effectively support the content.

**Significance**  
This paper offers a potentially unifying framework for studying neural network representations and identifying key alignment metrics that inform network alignment. It contributes to an active research area and presents predictions relevant to widely used models like ResNet and LLMs, enhancing its impact in the field. Furthermore, this research could have potential interdisciplinary applications such as neurosciences.
summary: The paper introduces six alignment relations to quantify and analyze model representations, collectively termed the Canonical Representation Hypothesis (CRH). This hypothesis proposes that neural networks naturally develop compact representations, with neurons and weights remaining unchanged by transformations that do not affect the task. The study theoretically examines conditions under which CRH holds and demonstrates that when CRH is disrupted, reciprocal power-law relationships arise among representations, weights, and gradients—a phenomenon referred to as the Polynomial Alignment Hypothesis (PAH). Together, CRH and PAH provide a potential unified theory that encompasses several significant phenomena in representation learning, such as Neural Collapse and the Neural Feature Ansatz.
weaknesses: **Originality**
Further clarifying how this work uniquely contributes to and advances the field would enhance its perceived originality.

**Quality**
No clear weakness in the quality.

**Clarity**
-  The term "hypothesis" may be somewhat misleading, as the proposed concept functions more as a measurement framework rather than a traditional hypothesis. Reframing it as such would emphasize its role in quantifying neural network representations without introducing potentially confusing terminology.
- The organization of the paper could benefit from refinement. Early references to later sections (such as mentioning Figure 9 before the relevant context is provided) disrupt the flow, and frequent shifts in reference to the later text and figures make the argument more challenging to follow.
- Providing more detailed explanations for terms such as "neural collapse" and "alignment" would enhance the paper's accessibility.
- Figure 4 lacks a colour bar for the weight decay parameter, which complicates the interpretation of results.

**Significance**
To improve the paper’s significance, it would be helpful to expand on the implications of the different alignment phases. For instance, as mentioned in the text, positive exponents between Za and Ha suggest that the layer enhances the principal components of Ha while diminishing lesser features, whereas a negative exponent indicates the opposite. Greater detail on these implications would improve interpretability, especially regarding how these phases contribute to understanding network behaviour. Additionally, exploring factors that lead to alignment breakdown would deepen the paper's contribution to interpretability in this area.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: 23: emergence + of

44: Furthermore, ... - this sentence makes no sense without more context

49: How does the CRH have anything to do with 'interpretable' solutions?

Related work, you might want to include: https://arxiv.org/abs/2007.00810

68: weight matrices evolve during training? This seems rather tautological, what is the point?

117: missing verb?

149: In practice, it is often not the case that stationarity is reached. Comment?

159-160: Can you explain?

217: Neural collapse / invariance to irrelevant features seems to be at odds with these works (https://proceedings.mlr.press/v139/zimmermann21a.html; https://arxiv.org/abs/2410.21869). Comment?

232: Why is this 'surprising' if it is exactly what the theory predicts?
soundness: 3
strengths: The paper makes an effort to be thorough in examining all possible combinations of the forward and backward alignment relations. Moreover, they include a significant number of theoretical results which distinguishes this paper from most work in machine learning.
summary: The paper proposes and studies the 'canonical representation hypothesis'. They study the evolution of weights, representations and gradients in neural networks and establish that these quantities become aligned. The paper includes theoretical derivations and experimental measurements that study how the CRH is fulfilled or broken.
weaknesses: Why should we be interested in the alignment of weights, gradients and representations? How does this contribute to the bigger picture of understanding why deep neural network work. If the authors provided a clear motivation there would be much more incentive to dig through the notation and derivations.

Meta-review: **Summary** This paper proposes the Canonical Representation Hypothesis (CRH), a hypothesis expressed in terms of six relations that after training the representations, weights and gradients align in each layer. The alignment is proposed as a fundamental part of understanding how neural networks form compact invariant representations.  The paper also proposes the Polynomial Alignment Hypothesis (PAH) which describes cases where CRH is broken. 

**Strengths** Reviewers were excited that the paper offered a new framework for understanding representation learning.  In particular, CRH seems to unify and generalize several previous efforts in the literature, including explaining neural collapse and  neural feature ansatz. The framework is quite general, applying broadly to different architectures and datasets.  Reviewers also praised the theoretical derivation supporting CRH and noted the experimental results supported the theory.  

**Weaknesses** Few weaknesses were noted.  Initially there were questions about the paper clarity and significance, but these were addressed during the author-reviewer discussion.  It was also noted the theoretical results rely on assumptions which may not hold in practice. 

**Conclusion** Reviewers were unanimously in favor of acceptance and generally praised the paper as a significant contribution to the theory of representation learning in neural networks.

Decision: Accept (Spotlight)