From Alexnet to Transformers: Measuring the Non-linearity of Deep Neural Networks with Affine Optimal Transport

We thank the reviewer for taking the time to review our work and for their comments.

I would like to see how the distribution in Fig. 3(C) changes when other metrics such as R^2 are used instead of the proposed $\rho_{aff}$.

We thank the reviewer for this suggestion. We've added histograms for $R^2$ score similar to those presented in Figure 3 for the affinity score in the Fig.1 of the PDF included in the global rebuttal. We can see vividly that $R^2$ score is less informative, with several models having modes around negative values. This is in line with our Remark 3.6 explaining why such a metric may be inappropriate for the considered task.

one would expect the nonlinearity score to behave symmetrically at $x=0$ for activation functions like ReLU

This is an interesting question, as it is a counterintuitive behavior. We can study this behaviour analytically by restricting to $X \sim \mathcal{U}[b,a]$, $b < 0$ and $a > 0$, $f: x \mapsto \mathrm{ReLU}(x)$, $Y = f(X)$ and verifying whether $\rho_{aff}(X, Y)$ is symmetrical in $a$ and $b$, that is, if it's invariant with respect to the assignment $(a,b)\mapsto (-b, -a)$.

A straightforward computation yields

• $\mu(X) = \frac{a+b}{2}$, $\Sigma(X) = \frac{(a-b)^2}{12}$,
• $\mu(Y) =\frac{a^2}{2(a-b)}$ and $\Sigma(Y) = \frac{a^3(a-4b)}{12(a-b)^2}$.

Plugging these in the affine transport formula yields

• $A_{aff}=\frac{\sqrt{a^3(a-4b)}}{(a-b)^2}$
• $b_{aff} = \frac{a}{2(a-b)}\left(a - \sqrt{a(a-4b)}\left(\frac{a+b}{a-b}\right)\right)$.

Finally, we can compute

$W_2^2(T_{aff}(X), Y) = \frac{a^3}{6(a-b)^2}\left((a-4b) + \sqrt{a(a-4b)}\left(\frac{-a+3b}{a-b}\right)\right)$.

Here we can see that this term is not strictly symmetric, but empirically we can assess that this term is not too far from being symmetric (see leftmost part of Figure 3 in the attached PDF).

A greater source of asymmetry in the nonlinearity score

$\rho_{aff}(X,Y) = 1 - \frac{W_2(T_{aff}(X), Y)}{\sqrt{2Tr(\Sigma(Y))}}$

is the trace term, as can be seen from the expression for $\Sigma(Y)$ and the middle plot of Figure 3. We illustrate this in the attached PDF in the global rebuttal.

I could not fully understand the definition of the non-linearity signature in Def. 3.1. It is defined based on $F_i \cap \mathcal{A}$, but $F_i$ is a layer in N, so it seems to be empty.

We consider a layer $i$ as a function $F_i$ acting as a block used in a repeated way in a given architecture. In that sense, $F_i \cap \mathcal{A}$ would just be a set of activation functions in such block without all other eventual operations (such as for instance, convolutions).

Affinity score of a neuron in the same model depending on the input class.

We thank the reviewer for this suggestion. We provide a visualization (Fig.4 of the attached PDF in the general reply) of the variance of the affinity scores when passing batches consisting of samples belonging to a single class of CIFAR10 dataset (so we have 10 batches, each batch contains samples from 1 class). For the sake of clarity, we use models with less than 40 layers. One can observe that early layers seem to be the most sensitive to the class distribution for all models. For ViT, some MLP block of intermediate layers also exhibit a high variability potentially hinting at their mono- and poly-semanticity.

We thank the reviewer for taking the time to review our work and for their comments.

I do not agree with the following statement Without non-linear activation functions, most of DNNs, no matter how deep, reduce to a linear function unable to learn complex patterns.

We believe that there is a misunderstanding regarding this claim. First, we believe that the presence of activation functions in all models available on Torchvision suggests that they are necessary to learn complex patterns. We are not aware of any DNN architecture that is widely acclaimed and that removes the activation functions altogether. Yet, we agree that for some models, such as transformers, even removing the activation function from attention and MLP blocks, doesn't make them linear due to the quadratic term in the attention equation. We propose to reformulate this phrase by saying: "For instance, without non-linear activation functions, popular multi-layer perceptrons, no matter how deep, reduce to a linear function unable to learn complex patterns."

I also disagree with the following Activation functions were also early identified [29, 30, 31, 32] as a key to making even a shallow network capable of approximating any function, however complex it may be, to arbitrary precision.

We do not claim that "nonlinearity is the key ingredient to function approximation in general" as the reviewer says. Here, we refer to universal approximation theorems [29, 30, 31, 32] of feedforward networks that require at least one hidden layer and a non-linearity to be proved. If there are other universal approximation theorems that the reviewer has in mind and that can be proved for models without an activation function, we will gladly discuss them in our work and adjust the phrasing accordingly.

Many formal results such as Theorem 3.3 are well known

We respectfully disagree with this comment as they seem to confuse the common results for OT between Gaussian distributions, duly cited in our work, and Theorem 3.3. Our statement holds for arbitrary (non-gaussian) random variables $X$ and $Y$ that satisfy $Y=T(X)$ for a symmetric positive-definite linear transformation $T$. The classical result about normal distributions is then a direct corollary: if $X$ and $Y$ are normally distributed, then such a $T$ can always be found, and the classical result follows.

We also note that "Computational Optimal Transport" by Peyré and Cuturi (2019), which arguably represents a modern reference for applied OT doesn't contain any results having the generality of Theorem 3.3. Yet, it does refer to the classical result for OT between Gaussians in Remark 2.31 as well as its generalization to elliptical distributions in Remark 2.32.

Fig 2. is also misleading since the "nonlinearity" of any activation function depends on the range of the inputs.

Kindly note that we illustrate exactly what the reviewer claims in Fig. 2 (A) for ReLU and in Appendix D Fig.6 for 8 more activation functions. In Appendix D, we explicitly say, "We observe that each activation function can be characterized by 1) the lowest values of its non-linearity obtained for some subdomain of the considered interval and 2) the width of the interval in which it maintains its non-linearity." (lines 514-516). We propose to clarify this in the caption of Fig. 2 (A) by saying "Non-linearity of activation functions depends on the range of input values (red), illustrated here on ReLU".

the statement No other metric extracted from the activation functions of the considered networks exhibits a strong consistent correlation with the non-linearity signature. is again an overstatement

As explained in our work, we are not aware of any other quantifiable way to measure the non-linearity of an activation function. Yet, we considered reasonable alternatives to it: such as $R^2$ score, CKA commonly used to measure the similarity across the layers of NNs, and sparsity of activations recently used to better understand the semantic properties of neurons in MLPs and transformers. Our results show that the non-linearity signatures capture something different from these metrics, just as we claim it. We will gladly discuss this with the reviewer.

the statement We proposed the first sound approach to measure non-linearity of activation functions in neural networks is also incorrect

None of the works mentioned by the reviewers provides a measure of non-linearity for activation functions. If we missed the introduction of such a metric upon reading the above-mentioned works, we would like to kindly ask the reviewer to point out the exact equations where it is defined in each of these works.

the study does not provide any actionable insights or understanding on the "why"

Our experiments put forward several novel contributions factually explaining "why" different neural architectures differ from the lens of their activation functions' non-linearity. We show that DNNs of the last decade have several trends for which we do not find any mention in prior work: 1) Till 2020, state-of-the-art CNNs included more (on average) linear activation functions, 2) ViTs behave differently from CNNs with a much more non-linear behavior of the activation function in their MLP blocs, 3) skip connections help the model learn the (linear) identity function, as suggested by He et al. (2016).

As highlighted in the Discussion section, the affinity score does provide actionable insights: it allows comparing neural architectures and identifying those that have a potential to disrupt the ML field, just as the transformers did.

the study leverages common optimal transport results on DNN's internal representations, under some strong assumption about the distribution of those representations.

We stand by our claim expressed before and hope that our reply regarding Theorem 3.3 which holds for arbitrary random variables will clear out the misunderstanding that the reviewer insists on.

We thank the reviewer for their positive comments about our work.

The authors should discuss more activation functions

We thank the reviewer for this suggestion. We kindly note that the individual behavior of 9 different activation functions (Sigmoid, ReLU, GeLU, ReLU6, LeakyReLU, tanh, hardtanh, SiLU and Hardswish) is shown in Fig. 6 in Appendix D. Additionally, we also consider the original activation functions used for each neural architecture studied further in our work (GeLU for ViTs, Swin, ConvNext, and SiLU in EfficientNets).

There is currently some research on the nonlinearity of deep neural networks that should be compared and discussed.

There was a recent paper by Teney et al. ("Neural Redshift: Random Networks are not Random Functions", CVPR, June 2024) that discusses the biases of the different activation functions for randomly initialized DNNs. Although it doesn't introduce a non-linearity measure for activation functions, it does illustrate the inductive bias of varying activation functions and their sensitivity to a change of weight magnitudes in Glorot initialization for randomly initialized MLPs. The authors of the latter work conclude that different activation functions visually carry different inductive biases (low complexity for ReLU, higher complexity for Tanh) that are more or less sensitive to weight magnitudes. Our work allows quantifying their findings using the affinity score (panel (a) of Fig.2 in the attached PDF). We also extend their work by showing that these biases can be changed by shifting the domain of the weight initialization (toward negative values) following the intuition provided in our work (panel (b) of Fig.2). This also affects the affinity score that becomes lower confirming our findings from the main manuscript.

We are open to discussing any other relevant work we may have missed further.

It would be beneficial to showcase examples from domains beyond computer vision.

We thank the reviewer for this suggestion. We considered the computer vision task, as the torchvision repository allows for a reproducible comparison of state-of-the-art models trained on Imagenet dataset. Additionally, the available models are very varied and span a decade of research. We agree that extending our analysis to other tasks mentioned by the reviewer is also very important, and we believe that our work can spur such contributions in the future. We will mention this in the Discussion.

How would the model’s nonlinearity affect different downstream tasks? Does the impact vary across datasets?

We show in Appendix H and Fig. 15 the deviation of non-linearity signatures for the same model between ImageNet and different datasets (CIFAR10/CIFAR100/Random data). We can see that deviations between ImageNet and CIFAR10/CIFAR100 are much lower than for random data, implying that these datasets are processed similarly to Imagenet. Measuring the distance between the non-linearity signature of the same DNN for different datasets thus could be a promising way to measure semantic similarity between datasets with potential applications in meta- and transfer learning where this notion is of primary importance.