Paper: Task structure and nonlinearity jointly determine learned representational geometry (ICLR-2024)

Reviewer comment (Experiments):
Most of the experiment focused on toy scale examples of binary classification with 2-layer network, even experiments with CNN has only 2 conv-layer networks. I feel it’s within the scope of this paper to show empirical evidence that the observations may generalize to larger scale CNN and larger dataset. (e.g. resnet and ImageNet) Will it be feasible to swap the activation function and show some similar effects?

Author response:
– Regarding the concerns about our focus on simple networks and tasks, see our “response to all reviewers” and Section 8 in the revised manuscript for a description additional results with deep convolutional networks on realistic tasks

Reviewer comment (Writing):
Notation in Eq. (in Sec. 3.1.1), or the method description in Sec. 3.1.1 is a bit confusing. Is it using multi-output setup so $W_O$ is also a vector?

Author response:
– We agree the w_o vector notation was confusing.  In principle, w_o and the labels y can be multi-dimensional vectors (as in our Fig. 4 experiments).  However, our experiments in Section 4 focus on the single-output case, so we have changed the notation to defining w_o as a single scalar for each hidden-layer neuron to avoid confusion.

Reviewer comment (Writing):
There seems to be super interesting intuition going on in Figure 1C, but the text in Sec. 3.1.1 didn’t seem to walk through the logic, leaving the reader to parse the schematics themselves. —— though after working through the math of an example it starts making sense.

Author response:
– Regarding the explanation of Figure 1C, we agree that the explanation in our initial submission was inadequate, and so we have added a paragraph to accompany the schematics (in red text in section 4.1.1 in the revised paper).

Reviewer comment (Writing):
Is there a typo of formula in Sec 3.1.1 “*the covariance between input and output: xxx*" should it be $y_i$?

Author response:
– Thank you for pointing out the typo in section 3.1.1 (now section 4.1.1 in the revised manuscript), we have fixed it to say y_i.

Reviewer comment (Experiments):
The results generated by the current study showing that Tanh nonlinearity helps generate target-aligned representation are limited to simple networks and simple input-output mapping.

Author response:
– Regarding the concerns about our focus on simple networks and tasks, see our “response to all reviewers” and Section 8 in the revised manuscript for a description additional results with deep convolutional networks on realistic tasks

Reviewer comment (Theory):
The representation geometry in these simple networks is probably not sufficient for many real-world problems that require capturing intricate patterns.

Author response:
– Regarding the concerns about our focus on simple networks and tasks, see our “response to all reviewers” and Section 8 in the revised manuscript for a description additional results with deep convolutional networks on realistic tasks

Reviewer comment (Theory):
Why does the noise level have non-monotonic effects on the ReLU network, consistently observed in all geometric matrices and for all tested separability of trained dichotomy? The author suggested smoothing gradients — what is the evidence supporting such a conclusion?

Author response:
– Regarding the effects of noise on learned representations in ReLU networks (its non-monotonic effect on target alignment and other representational metrics): we have conducted an additional analysis investigating this phenomenon in Appendix E in the revised manuscript.  In short, increased training noise makes it more likely for units that are initialized with high intra-class selectivity to “escape” regions of weight space with high intra-class / low inter-class selectivity, increasing target alignment.  We believe this phenomenon explains why intermediate levels of noise increase target alignment relative to zero noise.  However, we do not yet have a full mechanistic explanation for the non-monotonicity (i.e. why very high noise levels reduce target alignment); we suspect it arises due to noisy weight update dynamics outside the 2D subspace we are focusing on in our analysis.  We leave a more thorough investigation of this phenomenon to future work (or potentially the camera-ready version if the reviewers think it is a high-priority item).

Reviewer comment (Experiments):
Are the results robust to the input data range, e.g., input data ranges between 0 to 1 versus input data ranges between -1 to 1?

Author response:
– Regarding the input data range, we do not believe this factor has a substantial impact on our results, given that our weight initialization is random with expected mean zero (and thus the alignment of hidden unit input weights with the input data points at initialization is random, regardless of the input data range). Rerunning our Figure 1 experiments with inputs constrained to be nonnegative reproduces all the qualitative effects observed.

Reviewer comment (Novelty):
When training data becomes more complicated, as suggested by the noise input analysis and the XOR task, the difference between Tanh and ReLU vanishes. This suggests that Tanh helps generate a target-aligned representation when the input data is readily separated for clustering, or that in a multi-layer network, using Tanh at the final layer seems to be beneficial. This is not a novel conclusion — what are the new insights learned from the current study? Provide a discussion on how the results from the current study would have a general impact.

Author response:
– Regarding the point about differences between tanh / ReLU networks vanishing for more difficult tasks: we thank the reviewer for bringing up this important point.  We agree that our results in Figure 3 suggest that differences between learned representations are less prominent when the target geometry is insufficiently separable in input space relative to the input noise level.  This is also observed in Fig. 4B – the differences between activation functions vanish when input-output kernel alignment values approach 0.0, which corresponds to tasks where the output is not linearly separable in the input space.   We think it is an empirical question whether realistic tasks are in this regime or not, and we believe our (now more thorough) experiments with convolutional networks provide evidence that (at least some) realistic tasks fall in the regime where the differences between activation functions are relevant.  We believe that this is because for deep networks, the relevant notion of task difficulty becomes more complicated, as tasks in which the outputs are nonlinearly separable in input space may nevertheless be “easily separable” given the expressive capacity of a deeper network.  See our “response to all reviewers” for a discussion of new experiments we performed with multi-layer networks on synthetic tasks that explore a notion of task difficulty for deep networks.  Notably, we find tanh networks more target-aligned representations than ReLU, and this difference is in fact more pronounced for more difficult tasks.

Reviewer comment (Experiments):
While the analysis is interesting and it it wasn't clear to me how much it can impact can the paper have, its current form: (i) it is not clear how much the analysis extends to real tasks: the experiments of Cifar are somewhat limited (just the alignment metric is reported and it is not clear if this behavior holds for deeper networks: i.e. some ablations on the network depth should be incorporated in the experiment in my opinion) and the assumptions done on the synthetic tasks are unlikely to hold on larger networks (ii) there is no theory or additional experimental evidence that support why tanh and ReLU behave differently (see question section).

Author response:
– Regarding the concerns about the strength of our empirical evidence outside the synthetic task setting, see our “response to all reviewers” and Section 8 in the revised manuscript for additional results with deep convolutional networks on realistic tasks

Reviewer comment (Novelty):
The paper should report a better contextualization in the literature and comparison with similar works (it misses a related work section). I reported some works that should be discussed:

- Hayou, S., Doucet, A.; Rousseau, J.. On the Impact of the Activation function on Deep Neural Networks Training. Proceedings of the 36th International Conference on Machine Learning

- Ramachandran, Prajit, Barret Zoph, and Quoc V. Le. "Searching for activation functions." arXiv preprint arXiv:1710.05941 (2017).

- Ding, Bin, Huimin Qian, and Jun Zhou. "Activation functions and their characteristics in deep neural networks." 2018 Chinese control and decision conference (CCDC). IEEE, 2018.

- Papyan, Vardan, X. Y. Han, and David L. Donoho. "Prevalence of neural collapse during the terminal phase of deep learning training." Proceedings of the National Academy of Sciences 117, no. 40 (2020): 24652-24663.

- Zhu, Zhihui, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. "A geometric analysis of neural collapse with unconstrained features." Advances in Neural Information Processing Systems 34 (2021): 29820-29834.

Author response:
– We agree that our initial submission did not adequately address related work, especially the literature on neural collapse (which is highly related to our target alignment metric).  We have added a more comprehensive related work section, incorporating the references suggested by the reviewer and some others, to our revised manuscript.

Reviewer comment (Theory):
Can the authors elaborate on the intuition of why Tanh and ReLU behave in this way? And it would be possible to derive theoretical results on this?

Author response:
– Regarding intuition and theoretical explanation for the behavior of Tanh vs. ReLU: we have added a paragraph in section 4.1.1 (see red text) that provides further intuition, and additional explanation of how our approximation theoretical derivation of the input weight gradients in equation (1) gives rise to the vector fields plotted in Fig. 1D.

Reviewer comment (Theory):
The target alignment phenomenon of Tanh relates to neural collapse [a] phenomenon: i.e. when training is kept under zero error the representation in the last layer tend to collapse in equidistance clusters aligned to the targets. However, to the best of my knowledge this phenomenon should be agnostic of the activation function. Can the authors elaborate on this perspective?  
[a] Papyan, Vardan, X. Y. Han, and David L. Donoho. "Prevalence of neural collapse during the terminal phase of deep learning training." Proceedings of the National Academy of Sciences 117, no. 40 (2020): 24652-24663.  
[b] Zhu, Zhihui, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. "A geometric analysis of neural collapse with unconstrained features." Advances in Neural Information Processing Systems 34 (2021): 29820-29834.

Author response:
– Regarding the question about neural collapse, it is our impression that prior work has not directly investigated the role of activation function and input data structure on neural collapse.  Prior theoretical work on the subject (such as Zhu et al’s “unconstrained features model”) uses a network parameterization that ignores the constraints imposed by the activation function / network architecture and the input data geometry, and shows that representations tend to collapse toward the label geometry under this parameterization.  Our work shows that, while this may be a good first approximation (as target alignment does tend to increase / input alignment tends to decrease over training across all networks), the specifics of network architecture and task structure play an important role in determining the extent to which neural collapse occurs. Our suite of measures of representational geometry (including CCGP, PS, and input alignment) also provides a more detailed window into the learned representational geometry.

Reviewer comment (Theory):
What is the mathematical formulation of the nonlinearity analyzed in section 7?

Author response:
– We have included the specific form of the nonlinearities analyzed in section 7 (now section 9 in the revised manuscript), along with increased explanation of these experiments and the reasoning behind them.