Paper: Deep Network Partition Density Exhibits Double Descent (ICLR-2024)

Reviewer comment (Experiments):
The biggest weakness in the paper seems to be that a critical part of it is missing. The experimental setup in section 4.1 states “We control the size of the local neighborhood by controlling the diagonal lengths via the $r$ parameter mentioned in section 3.2”. But section 3.2 makes no mention of a parameter $r$, and it is difficult to make a full evaluation of the experiments in section 4 absent a definition of $r$.

Author response:
We thank the reviewer for raising this very important point. We have re-written section 4.1 and 3.1 adding the definition of the radius $r$ and its relationship with $conv(V)$. In the appendix, we have also added experiments showing the effect of varying $r$ and number of orthonormal vectors $P$.

To summarize, we start by sampling $P$ orthonormal vectors $B=${$ v_p : p=1...P$} and use them to denote vertices of a local neighborhood $V$ around any data point $x$ s.t. $V=${$x \pm rv_p : p = 1…P$} where $r$ is the distance from $x$ at which we choose vertices defining the neighborhood. Therefore, $x$ is the centroid of $V$ consisting of 2P vertices. We define the convex hull $conv(V)$ as the local neighborhood of $x$ in the input space. The dimensionality and volume of the neighborhood can be controlled via $P$ and $r$. Each neighborhood in the input space around any data point $x$ is a cross-polytope or the dual of a hypercube with dimensionality P and diagonal length $2r$. Each neighborhood can also be considered a unit $\ell_1$-norm ball scaled by radius $r$.

Reviewer comment (Novelty):
There is a line of prior work that seems closely related to the authors' own proposed work, which is the idea of polytope lenses by [Black et al.](https://arxiv.org/abs/2211.12312). Briefly, this work attempts to look at polysemy of neurons by using ideas from casting neural networks as compositions of splines, though they are chiefly interested in interpretation rather than complexity. This paper would be made stronger by contrasting the ideas in Black et al to this work's own, by showing the benefit of viewing DNs as splines has multiple uses.

Author response:
We thank the reviewer for bringing the paper to our attention. Indeed the analysis by Black et al. is very relevant to the results of our paper. While they are focusing on interpretability, the authors provide experiments showing that the density of polytope boundaries are related to semantic boundaries. In our paper, while we compute local complexity in terms of layerwise hyperplane intersections for a given neighborhood, the number of polytope edges formed via an arrangement of $N$ different $(d-1)$-hyperplanes is of the order $N^d$. Therefore, our proposed local complexity would be highly correlated with polytope boundary density. We will add a short discussion contrasting our method with Black et al. in our conclusion and connect it with potential future directions on interpretability.

Reviewer comment (Experiments):
The lack of a need for labels makes this measure ripe for helping to help explain the effects of self-supervised learning. The reviewer is surprised that the authors don't include experiments for why (e.g.) BYOL can work by measuring the local complexity of teacher & student networks as they train, and iteratively replace each other.

Author response:
We thank the reviewer for raising this excellent point. Indeed, evaluating BYOL especially the alignment between the teacher and student network, can be assessed in a sample-wise manner using our local complexity metric. In this paper however, we have mostly focused on the double descent phenomenon of local complexity that we observe in a wide variety of settings, while providing a method to do fast approximation of local complexity. The double descent phenomenon is important since it shows that based on the phase of training, neural networks tend to increase or decrease local complexity around training points, which opens up avenues for new theoretical studies on neural network training and generalization. We leave the SSL experiments for future work and will add a note on that in the conclusion.

Reviewer comment (Reproducibility):
It would be helpful for the authors to spend a few words about why the cost of computing the $P$ orthonormal basis vectors for each $x$ either isn't prohibitive, or why it might be in certain cases.

Author response:
In our experiments, we have seen no significant difference between sampling the orthonormal basis vectors for every data point $x$, vs using the same orthonormal frame being used for all the data points. Using the Pytorch implementation of sampling P orthonormal vectors, finding P=25 vectors takes approximately 9.49 ms, which needs to be done only once at the beginning of our experiments. We have added these details in the paper.

Reviewer comment (Theory):
In section 3.1, there is no characterization of the properties of $[\mathbf{v}_1, \dots, \mathbf{v}_P]$ other than they form an orthonormal set, and that $x$ is within $conv(\mathbf{V})$. This latter point is implied but not stated. Section 3.1 could use more care in specifying the relation between $\mathbf{x}$ and $\mathbf{V}$.

Author response:
We apologize for the lack of clarity, we are rewriting Sec 3.1 and clarifying the notations as discussed above.

Reviewer comment (Theory):
How stable is local complexity about a point $\mathbf{x}$ as a consequence of the choice of $\mathbf{V}$? And choice of $\mathbf{P}$? These seem like important questions to answer.

Author response:
In appendix, we have added experiments exploring how the parameters of neighborhood $conv(V)$, namely radius $r$ and dimensionality $P$ affects LC computation. In summary, we provide the following observations:

1. As we increase P, after a certain point, there is marginal increase in LC with increasing P for the same input space point $x$. In Fig.16 left we show that for a 784 dimensional input space and a randomly initialized MLP with width 100 and depth 18, $P=42$ has marginal difference in LC compared to $P=382$.

2. As we increase radius $r$, the volume of the neighborhood $conv(V)$ increases. While higher neighborhood volume allows computing LC at coarser scales (Fig.16 right), it can also introduce distortion in the embedded neighborhood reducing approximation quality. In this regard, we compute the eccentricity and diameter of the graph formed via $V^\ell$ which is the layer $\ell-1$ embeddings of vertices $V$.  We use $1000$ training points from CIFAR10 to define the neighborhoods for a trained CNN. We see that for larger $r$ both of the deformation metrics exponentially increase, whereas for $r \leq 0.014$ the deformation is lower and more stable. This shows that for lower $r$ our LC approximation for deeper CNN networks is stable since the neighborhood does not get deformed significantly.

Reviewer comment (Presentation):
The authors intentionally leave out extension of this to deeper networks, as a clear continuation of this line of work. The conclusion would be strengthened by spending more of the word budget to help lead the reader into how the authors plan to do so.

Author response:
In appendix, we have added an experiment computing the LC while training a Resnet18 on Imagenet upto 69% top-1 accuracy. Computing LC for 1000 imagenet samples require ~28s on a Quadro RTX 8000 GPU. We use a hull dimensionality of $P=25$ and radius $r=0.001$. We have also previously presented, local complexity dynamics for Resnet18 trained on CIFAR10, and shown that without Batch Normalization, it exhibits similar dynamics as we have seen in CNNs and MLPs.

Reviewer comment (Theory):
Do the authors have any conjecture about how the partition migrating dynamics relate to the phenomenon of neural collapse?

Author response:
We thank the reviewer for this excellent point. In Fig.1, we can see that as training progresses, the neurons (denoted by black lines) start orienting themselves with each other and grouping together. We call this phenomenon partition migration. This can be further seen in the animation we have provided as supplementary material, which presents an exact visualization of the dynamics presented in Fig.1 as training progresses. The grouping of neurons is equivalent to the neurons losing linear independence, which is directly related to neural collapse. There are two additional observations in our result that to the best of our knowledge have not been explored previously in Neural Collapse literature. The first, neurons from different layers of the network may ‘collapse’ together, meaning that neural collapse is not a layerwise phenomenon. Secondly, as can be seen in the supplementary animation, the grouped neurons can jump between different orientations while staying grouped together, indicating that a network may keep changing its features locally while in a collapsed form. We will add further discussions on this point taking into account space limitations.

Reviewer comment (Writing):
I think that the writing of the paper can be improved, as in many parts it was quite hard to understand the actual information. Also, in some figures, some extra information is necessary, for example, in Fig. 6 what is the difference in each panel?

Author response:
We apologize for the lack of clarity especially Sec. 3.1 and Sec. 4.1 which are integral to understanding how the proposed method works. 

To summarize, in this paper we provide a method to approximate the local complexity (LC) of DNNs and show that for a wide range of training settings, the local complexity exhibits an epoch-wise double descent phenomenon. To compute local complexity around any given point $x$ in the input space, we perform the following steps.
1. First we start by sampling $P$ orthonormal vectors $B=${$ v_p : p=1...P$}.
2.  We use $B$ to denote vertices of a local neighborhood $V$ around any data point $x$ s.t. $V=${$x \pm rv_p : p = 1…P$} where $r$ is a radius parameter. $x$ is therefore the centroid of $V$ consisting of 2P vertices each $r$ distance away from $x$. We define the convex hull $conv(V)$ as the local neighborhood of $x$ in the input space. The dimensionality and volume of the neighborhood can be controlled via $P$ and $r$. Each neighborhood can be considered a non-axis aligned unit $\ell_1$-norm ball scaled by radius $r$.
3. Given an input space neighborhood $conv(V)$, we embed $V$ to the input space of layers $\ell$ of the network which precede a non-linearity. We denote the embedded vertices in the input space of layer $\ell$ as $V^\ell$.
4. For a given embedded neighborhood $conv(V^\ell)$, we check if there is a change in the pre-activation sign of any given neuron between the vertices. If there is, the neuron (hyperplane) intersects $conv(V^\ell)$. We repeat this for all neurons in all layers and count the total number of intersections as the local complexity.

> Some extra information is necessary, for example, in Fig. 6 what is the difference in each panel

We apologize for the lack of clarity in Fig 6. From left to right each panel presents the local complexity for train, test and random points. For training LC, the LC during the ascent phase changes from low to high to low again while the number of parameters of the network is increase. We do not see such a phenomenon for LC around the test or random points.

Reviewer comment (Theory):
Even if the proposed measure seems to capture the local complexity of the classifier, apart from some connections to related works, I think that a thorough analysis is missing.

Author response:
In appendix, we have added empirical results analyzing how the parameters of neighborhood $conv(V)$, namely radius $r$ and dimensionality $P$ affects LC computation. In summary, we provide the following observations:

1. As we increase P, after a certain point, there is marginal increase in LC with increasing P for the same input space point $x$. In Fig.16 left we show that for a 784 dimensional input space and a randomly initialized MLP with width 100 and depth 18, $P=42$ has marginal difference in LC compared to $P=382$.

2. As we increase radius $r$, the volume of the neighborhood $conv(V)$ increases. While higher neighborhood volume allows computing LC at coarser scales (Fig.16 right), it can also introduce distortion in the embedded neighborhood reducing approximation quality. In this regard, we compute the eccentricity and diameter of the graph formed via $V^\ell$ which is the layer $\ell-1$ embeddings of vertices $V$.  We use $1000$ training points from CIFAR10 to define the neighborhoods for a trained CNN. We see that for larger $r$ both of the deformation metrics exponentially increase, whereas for $r \leq 0.014$ the deformation is lower and more stable. This shows that for lower $r$ our LC approximation for deeper CNN networks is stable since the neighborhood does not get deformed significantly.

Reviewer comment (Evaluation):
I think that the proposed measure is an interesting and simple approach to approximate the number of convex regions (in some sense the complexity) of the classifier around the training data. However, I do not understand how it should be used and how it helps to analyze the actual behavior of the deep classifier.

Author response:
While previous work has clearly motivated the need for methods to compute the local complexity of DNNs, exact computation of complexity measures, e.g., the number of linear regions, has combinatorial complexity and is intractable for large deep networks with arbitrary dimensional input spaces. The following are ways in which our results help understand neural networks:
1. Our method can be used to perform fast approximation of the DNN local complexity and can be scaled up to Imagenet experiments as well. In appendix we have presented the LC computed during Resnet18 training on Imagenet upto a top 1 accuracy of 69%. We can also use LC to compare between trained models, compare between sub-groups within the training/test set to quantify bias. 
2. Using our method we have presented a novel observation that based on the phase of training, the local complexity around training points might be increasing or decreasing with training epochs. To the best of our knowledge this is the first observation of this kind. This raises new questions regarding our current theoretical understanding of DNN training dynamics, opening new avenues for exploration.
3. Our results indicate that grokking could be a result of 'partition migration' - a phenomenon that we have presented where during the ending stages of training, layerwise neurons (hyperplanes) start orienting themselves with each other and grouping together. 
4. In supplementary materials, we present an animation showing the partition migration phenomenon during training, as was shown in Fig.1. Note that while the visualizations are produced via earlier work on partition visualization, our proposed method allows monitoring the partitions at scale during training, which can be used to probe the input space to find where to look. Partition migration can be used to explain the neural collapse phenomenon in DNN training. Because partitions as well as neurons (black lines in Fig.1) start grouping together, loosing linear independence and causing neural collapse.