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

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:40:25

Abstract: In the last decade, we have witnessed the introduction of several novel deep neural network (DNN) architectures exhibiting ever-increasing performance across diverse tasks. Explaining the upward trend of their performance, however, remains difficult as different DNN architectures of comparable depth and width -- common factors associated with their expressive power -- may exhibit a drastically different performance even when trained on the same dataset. In this paper, we introduce the concept of the non-linearity signature of DNN, the first theoretically sound solution for approximately measuring the non-linearity of deep neural networks. Built upon a score derived from closed-form optimal transport mappings, this signature provides a better understanding of the inner workings of a wide range of DNN architectures and learning paradigms, with a particular emphasis on the computer vision task. We provide extensive experimental results that highlight the practical usefulness of the proposed non-linearity signature and its potential for long-reaching implications.

Main paper: 
 introduction: 
 Deep neural networks (DNNs) are undoubtedly the most powerful AI models currently available [1, 2, 3, 4, 5] . Their performance on many tasks, including natural language processing (NLP) [6] and computer vision [7] , is already on par or exceeds that of a human being. One of the reasons explaining such progress is of course the increasing computational resources [8, 9] . Another one is the endeavour for finding ever more efficient neural architectures pursued by researchers over the last decade. As of today, the transformer architecture [10] has firmly imposed itself as a number one choice for most, if not all, of the recent breakthroughs [11, 12, 13] in the machine learning and artificial intelligence fields.Limitations But why transformers are more capable than other architectures? Answering this question requires finding a meaningful measure to compare the different famous models over time gauging the trend of their intrinsic capacity. For such a comparison to be informative, it is particularly appropriate to consider the computer vision field that produced many of the landmark neural architectures improving upon each other over the years. Indeed, the decade-long revival of deep learning started with Alexnet's [14] architecture, the winner of the ImageNet Large Scale Visual Recognition Challenge [15] in 2012. By achieving a significant improvement over the traditional approaches, Alexnet was the first truly deep neural network to be trained on a dataset of such scale, suggesting that deeper models were likely to bring even more gains. In the following years, researchers proposed novel ways to train deeper models with hundreds of layers [16, 17, 18, 19] pushing the performance frontier even further. The AI research landscape then reached a turning Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.point with the proposal of transformers [10] , starting their unprecedented dominance first in NLP and then in computer vision [20] . Surprisingly, transformers are not particularly deep, and the size of their landmark vision architecture is comparable to that of Alexnet, and this despite a significant performance gap between the two. Ultimately, this gap should be explained by the differences in the expressive power [21] of the two models: a term used to denote the ability of a DNN to approximate functions of a certain complexity. Unfortunately, the existing theoretical results related to this either associate higher expressive power with depth [22, 23, 24] or width [25, 26, 27, 28] falling short in comparing different families of architectures. This, in turn, limits our ability to understand what underpins the achieved progress and what challenges and limitations still exist in the field, guiding future research efforts.Contributions We argue that quantifying the non-linearity of a DNN may be what we were missing so far to understand the evolution of the deep learning models at a more fine-grained level. To verify this hypothesis in practice, we put forward the following contributions:1. We propose a first theoretically sound measure, called the affinity score, that estimates the non-linearity of a given (activation) function using optimal transport (OT) theory. We use the proposed affinity score to introduce the concept of the non-linearity signature of DNNs defined as a set of affinity scores of all its activation functions.2. We compare non-linearity signatures of a wide range of popular DNNs used in computer vision: from Alexnet to vision transformers (ViT) and their more recent variations. Through this, we clearly illustrate the disruptive patterns in the evolution of the deep learning field.3. We demonstrate that non-linearity signature can be predictive of DNNs performance and used to meaningfully identify the family of approaches to which a given DNN belongs. We further show that the non-linearity signature is unique as it doesn't correlate strongly with other potential candidates used for this task.The rest of the paper is organized as follows. We start by presenting the relevant background knowledge on OT in Section 2. Then, we introduce the affinity score together with its different theoretical properties in Section 3. Section 4 presents experimental evaluations on a wide range of popular convolutional neural networks. Finally, we conclude in Section 5. 
background: 
 Optimal Transport Let (X, d) be a metric space equipped with a lower semi-continuous cost function c : X × X → R ≥0 , e.g the Euclidean distance c(x, y) = ∥x -y∥. Then, the Kantorovich formulation of the OT problem between two probability measures µ, ν ∈ P(X) is given byEQUATIONwhere ADM(µ, ν) is the set of joint probabilities with marginals µ and ν, and E ν [f ] denotes the expected value of f under ν. The optimal γ minimizing equation 1 is called the OT plan. Denote by L(X) the law of a random variable X. Then, the OT problem extends to random variables X, Y and we write OT c (X, Y ) meaning OT c (L(X), L(Y )).Assuming that either of the considered measures is absolutely continuous, then the Kantorovich problem is equivalent to the Monge problem OT c (µ, ν) = minEQUATIONwhere the unique minimizing T is called the OT map, and T # µ denotes the push-forward measure, which is equivalent to the law of T (X), where X ∼ µ.Wasserstein distance Let X be a random variable over R d satisfying E[∥X -x 0 ∥ 2 ] < ∞ for somex 0 ∈ R d , and thus for any x ∈ R d . We denote this class of random variables by P 2 (R d ). Then, the 2-Wasserstein distance W 2 between X, Y ∈ P 2 (R d ) is defined asEQUATION)3 Non-linearity signature of deep neural networks Among all non-linear operations introduced into DNNs in the last several decades, activation functions remain the only structural piece that they all inevitably share. Without non-linear activation functions, most of DNNs, no matter how deep, reduce to a linear function unable to learn complex patterns.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 thus build our study on the following intuition: if activation functions play in important role in making DNNs non-linear, then measuring their degree of non-linearity can provide us with an approximation of the DNN's non-linearity itself. To implement this intuition in practice, however, we first need to find a way to measure the non-linearity of an activation function. Surprisingly, there is no widely accepted measure for this, neither in the field of mathematics nor in the field of computer science. To fill this gap, we will use the OT theory to develop a so-called affinity score below.T aff as in (4), W 2 (T aff X, Y ) ≤ 2 Tr [Σ(Y )].To have a more informative non-linearity measure, we now need to normalize the non-negative Wasserstein distance W 2 (T aff X, Y ) to an interpretable interval of [0, 1]. The bound given in Proposition 3.5 lets us define the following affinity scoreEQUATIONThe proposed affinity score quantifies how far a given activation σ is from an affine transformation.It is equal to 1 for any input for which the activation function is linear, and 0 when it is maximally non-linear, i.e., when T aff X and σ(X) are independent random variables.Remark 3.6. One may wonder whether a simpler alternative to the affinity score can be to use, instead of T aff , a mapping T W (x) = W x defined as a solution of a linear regression problemmin W ||Y -W X|| 2 F .Then, one can use the coefficient of determination (R 2 score) to measure how well T W fits the observed data. This approach, however, has two drawbacks. First, following the famous Gauss-Markov theorem, T W is an optimal linear (linear in Y ) estimator. On the contrary, T aff is a globally optimal non-linear mapping aligning X and Y . Second, R 2 compares the fit of T W with that of a mapping outputting µ(Y ) for any value of X. This is contrary to ρ aff that compares how well T aff fits the data wrt to the worst possible cost incurred by T aff as quantified in Proposition 3.5.This gives us a bounded score, i.e. ρ aff ∈ [0, 1], whereas R 2 is not lower bounded, i.e.R 2 ∈ [-∞, 1].We confirm experimentally in Section 4 that the two coefficients do not correlate consistently across the studied DNNs suggesting that R 2 is a poor proxy to ρ aff . Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations in Appendix B.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
affinity score: 
 Identifiability We consider the pre-activation signal X of an activation function within a neural network, and the post-activation signal σ(X) denoted by Y as input and output random variables.Our first step to build the affinity score then is to ensure that we can identify when σ is linear with respect to (wrt) X (for instance, when an otherwise non-linear activation is locally linear at the support of X). To show that such an identifiability condition can be satisfied with OT, we first recall the following classic result from the literature characterizing the OT maps.Theorem 3.1 ([33]). Let X ∈ P 2 (R d ), T (x) = ∇ϕ(x) for a convex function ϕ with T (X) ∈ P 2 (R d ).Then, T is the unique optimal OT map between µ and T # µ.Using this theorem about the uniqueness of OT maps expressed as gradients of convex functions, we can prove the following result (all proofs can be found in the Appendix C):Corollary 3.2. Without loss of generality, let X, Y ∈ P 2 (R d ) be centered, and let Y = σ(X) = T X,where T is a positive definite linear transformation. Then, T is the OT map from X to Y .Whenever the activation function σ is linear, the solution to the OT problem T exactly reproduces it.Characterization We now seek to understand whether T can be characterized more explicitly. For this, we prove the following theorem stating that T can be computed in closed-form using the normal approximations of X and Y .Theorem 3.3. Let X, Y ∈ P 2 (R d ) be centered and Y = T X for a positive definite matrix T . LetN X ∼ N (µ(X), Σ(X)) and N Y ∼ N (µ(Y ), Σ(Y ))be their normal approximations where µ and Σ denote mean and covariance, respectively. Then, W 2 (N X , N Y ) = W 2 (X, Y ) and T = T aff , where T aff is the OT map between N X and N Y and can be calculated in closed-formEQUATIONUpper bound When the activation σ is non-linear wrt X, the affine OT mapping T aff (X) will deviate from the true activation outputs Y . One important step toward quantifying this deviation is given by the famous Gelbrich bound, formalized by means of the following theorem:Theorem 3.4 (Gelbrich bound [34]). Let X, Y ∈ P 2 (R d ) and let N X , N Y be their normal approxi- mations. Then, W 2 (N X , N Y ) ≤ W 2 (X, Y ).This upper bound provides a first intuition of why OT can be a great tool for measuring non-linearity:the cost of the affine map solving the OT problem on the left-hand side increases when the map becomes non-linear. We now upper bound the difference between W 2 (N X , N Y ) and W 2 (X, Y ), two quantities that coincide only when σ is linear.Proposition 3.5. Let X, Y ∈ P 2 (R d ) and N X , N Y be their normal approximations. Then,1. |W 2 (N X , N Y ) -W 2 (X, Y )| ≤ 2 Tr (Σ(X)Σ(Y )) 1 2 √ Tr[Σ(X)]+Tr[Σ(Y )]. The affinity score is calculated based on 3 main steps. First, given an input (grey) and an output (red) of an activation function (left), we estimate the best affine OT fit T aff (X) (green) transporting the input to the output (middle-left). Second, we measure the mismatch between the two by summing the transportation costs (middle-right) to obtain the Wasserstein distance W 2 (T aff X, Y ). Finally, this distance is normalized with the magnitudes of variance (arrows in the rightmost plot) of the output data based on its covariance matrix.We now turn our attention to the definition of a non-linearity signature of deep neural networks. We define a neural network N as a composition of layers F i where each layer F i is a function taking as input a tensor X i ∈ R hi×wi×ci (for instance, an image of size 224 × 224 × 3 for i = 1) and outputting a tensor Y i ∈ R hi+1×wi+1×ci+1 used as an input of the following layer F i+1 . This definesN = F L ⊙ ... ⊙ F i ... ⊙ F 1 = k=1,...,L F k where ⊙ stands for a composition.We now present the definition of a non-linearity signature of a network N. Below, we abuse the compositional structure of F i and see it as an ordered sequence of functions.Definition 3.1. Let N = k=1,...,L F k be a neural network. Define by A a finite set of common activation functions such that A := {σ|σ : R h×w×c → R h×w×c }. Let r be a pooling operation such that r : R h×w×c → R c . Then, the non-linearity signature of N given an input X is defined as follows:ρ aff (N; X) = {ρ aff (r(X i ), σ(r(X i ))), ∀σ ∈ F i ∩ A, i = {1, . . . , L}}.Non-linearity signature, illustrated in Figure 1 , associates to each network N a vector of affinity scores calculated over the inputs and outputs of all activation functions encountered across its layers.What makes an activation function non-linear? We now want to understand the mechanism behind achieving a lower or higher non-linearity with a given (activation) function. This will explain what the different values of the affinity scores stand for when defining the non-linearity signature of a DNN. In Figure 2 (A), we show how the ReLU function [35] , defined element-wise as ReLU(x) = max(0, x), achieves its varying degree of non-linearity. Interestingly, this degree depends only on the range of the input values. Second, in Figure 2 (B) we also show how the shape of activation functions impacts their non-linearity for a fixed input: surprisingly, piece-wise linear ReLU function is more non-linear than Sigmoid(x) = 1/(e -x + 1) [36] or Tanh(x) = (e -x -e x )/(e -x + e x ).Similar observations also apply to compare polynomials of varying degrees (Figure 2 (C)). We refer the reader to Appendix D for more visualizations of the affinity score of popular activation functions.Layer-wise similarity analysis of DNNs A line of work that can be distantly related to our main proposal is that of quantifying the similarity of the hidden layers of the DNNs as proposed [37] and [38] (see [39] for a complete survey of the subsequent works). [37] extracts activation patterns of the hidden layers in the DNNs and use CCA on the singular vectors extracted from them to measure how similar the two layers are. Their analysis brings many interesting insights regarding the learning dynamics of the different convnets, although they do not discuss the non-linearity propagation in the convnets, nor do they propose a way to measure it. [38] proposed to use a normalized Frobenius inner product between kernel matrices calculated on the extracted activations of the hidden layers and argued that such a similarity measure is more meaningful than that proposed by [37] .Impact of activation functions [40] provides the most comprehensive survey on the activation functions used in DNNs. Their work briefly discusses the non-linearity of the different activation functions suggesting that piecewise linear activation functions with more linear components are more non-linear (e.g., ReLU vs. ReLU6). [41] show theoretically that smooth versions of ReLU allow for more efficient information propagation in DNNs with a positive impact on their performance.Our work provides a first extensive comparison of all popular activation functions; we also show that smooth version of ReLU exhibit wider regions of high non-linearity (see Appendix D).Non-linearity measure The only work similar to ours in spirit is the paper by [42] proposing the non-linearity coefficient in order to predict the train and test error of DNNs. Their coefficient is defined as a square root of the Jacobian of the neural network calculated wrt its input, multiplied by the covariance matrix of the Jacobian, and normalized by the covariance matrix of the input. The presence of the Jacobian in it calls for the differentiability assumption making its application to most of the neural networks with ReLU non-linearity impossible as is. The authors didn't provide any implementation of their coefficient and we were not able to find any other study reporting the reproduced results from this work.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: Anonymized code to reproduce experiments is available as a zip file, with a README file to explain how to run it.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental evaluations: 
 We consider computer vision models trained and evaluated on the same Imagenet dataset with 1,000 output categories (Imagenet-1K) publicly available at [43] . The non-linearity signatures of different studied models presented in the paper is calculated by passing batches of size 512 through the pre-trained models for the entirety of the Imagenet-1K validation set (see Appendix H for more datasets) with a total of 50,000 images. We include the following landmark architectures in our study:Alexnet [14] , four VGG models [16] , Googlenet [44] , Inception v3 [17] , five Resnet models [18] ,four Densenet models [19] , four MNASNet models [45] , four EfficientNet models [46] , five ViT models, three Swin transformer [47] and four Convnext models [48] . We include MNASNet and EfficientNet models as prominent representatives of the neural architecture search approach [49] .Such models are expected to explicitly maximize the accuracy for a given computational budget.Swin transformer and Convnext models are introduced as ViTs with traditional computer vision priors. Their presence will be useful to better grasp how such priors impact ViTs. We refer the reader to Appendix E for more practical details.We give a general outlook of the developments in computer vision over the last decade when seen through the lens of their non-linearity. In Figure 3 we present the minimum, median, and maximum values of the affinity scores calculated for the considered neural networks (see Appendix F for raw non-linearity signatures). We immediately see that until the arrival of transformers, the trend of the landmark models was to decrease their non-linearity, rather than to increase it. On a more fine-grained level, we note that pure convolution Closer look at accuracy/non-linearity trade-off Different families of vision models leverage different characteristics of their internal non-linearity to achieve better performance. To better understand this phenomenon, we now turn our attention to a more detailed analysis of the accuracy/non-linearity trade-off by looking for a statistic extracted from their non-linearity signatures that is the most predictive of their accuracy as measured by the R 2 score. Additionally, we also want to understand whether the non-linearity of DNNs can explain their performance better than the traditional characteristics such as the number of parameters, the number of giga floating point operations per second (GFLOPS), and the depth. From the results presented in Figure 4 , we observe the following. First, the information extracted from the non-linearity signatures often correlates more with the final accuracy, than the usual DNN characteristics. This is the case for Residual networks (ResNets and DenseNets), ViTs, and vision models influenced by transformers (Post-ViT). Unsurprisingly, for models based on neural architecture search (NAS-based, i.e. EfficientNets and MNASNets) the number of parameters is the most informative metric as they are specifically designed to reach the highest accuracy with the increasing model size and compute. For Pre-residual pure convolutional models (Alexnet, VGGs, Googlenet, and Inception), the spread of the non-linearity explains the accuracy increase similarly to depth. Second, we observe that all models preceding ViTs were implicitly optimizing the spread of their affinity score values to achieve better performance. After the arrival of the transformers, the observed trend is to increase either the median or the minimum values of the non-linearity. This suggests a fundamental shift in the implicit bias that the transformers carry. Distinct signature for every architecture Non-linearity signature correctly identifies the different families of neural architectures. To show this, we perform hierarchical clustering using pairwise dynamic time warping (DTW) distances [50] between the non-linearity signatures of the models from connections make the distribution of the affinity scores bimodal with one such mode centered around highly linear activation functions. This confirms in a principled way that residual connections indeed tend to enable the learning of the identity function just as suggested in the seminal work that proposed them [18] . Non-linearity signatures can also be applied to meaningfully identify training methods, such as popular nowadays self-supervised approaches, for a fixed architecture (see Appendix I).Uniqueness of the affinity score No other metric extracted from the activation functions of the considered networks exhibits a strong consistent correlation with the non-linearity signature. To validate this claim, we compare in Table 1 the Pearson correlation between the non-linearity signature and several other metrics comparing the inputs and the outputs of the activation functions. We can see that for different models the non-linearity correlates with different metrics suggesting that it captures the information that other metrics fail to capture consistently across all architectures. This becomes even more apparent when analyzing the individual correlation values (in Appendix G). Overall, the proposed affinity score and the non-linearity signatures derived from it offer a unique perspective on the developments in the ML field. 
discussions: 
 We proposed the first sound approach to measure non-linearity of activation functions in neural networks and defined their non-linearity signature based on it. We further used non-linearity signatures to provide a meaningful overview of the evolution of neural architectures proposed over the last decade with clear interpretable patterns. We showed that until the arrival of transformers, the trend in DNNs was to decrease their non-linearity, rather than to increase it. Vision transformers changed this pattern drastically. We also showcased that our measure is unique, as no other metric correlates strongly with it across all architectures.In the future, our work can be applied to study the non-linearity of the LLM models to better understand the effect of different architectural choices in them. On a higher level, our approach can also be used to identify new disruptive neural architectures by identifying those of them that leverage different internal non-linearity characteristics to obtain better performance. This capacity of identifying novel technologies is even more crucial in the age of very large models where experimenting with the building blocks of the optimized backbone comes at a very high cost.This paper presents work whose goal is to advance the field of Machine Learning and better understand the underlying behavior of Deep Neural Networks architectures. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.An important assumption of Theorem 3.3, is that the activation function that we want to analyze through ρ aff needs to be a positive definite transformation of the inputs. Fortunately, this is the case for activation functions, that we consider in this paper. Finally, we note that despite the strong correlation between the statistics extracted from the non-linearity signatures for certain DNNs' architectures,we are yet to show that explicitly optimizing affinity scores through backpropagation can have an actionable impact on DNNs performance or its other properties, such as robustness or transferability.In this section, we provide proofs of the main theoretical results from the paper.Corollary 3.2. Without loss of generality, let X, Y ∈ P 2 (R d ) be centered, and such that Y = T X,where T is a positive semi-definite linear transformation. Then, T is the OT map from X to Y .Proof. We first proof that we can consider centered distributions without loss of generality. To this end, we note thatEQUATIONimplying that splitting the 2-Wasserstein distance into two independent terms concerning the L 2 distance between the means and the 2-Wasserstein distance between the centered measures.Furthermore, if we have an OT map T ′ between X -E[X] and Y -E[Y ], thenEQUATIONis the OT map between X and Y .To prove the statement of the Corollary, we now need to apply Theorem 3.1 to the convex ϕ(x) =x T T x, where T is positive semi-definite. T aff is the OT map between N X and N Y and can be calculated in closed-formEQUATIONProof. Corollary 3.2 states that T is an OT map, andΣ(T N X ) = T Σ(X)T = Σ(Y ).Therefore, T N X = N Y , and by Theorem 3.1, T is the OT map between N X and N Y . Finally, we computeEQUATION2 (X, Y ).Proposition 3.5. Let X, Y ∈ P 2 (R d ) and N X , N Y be their normal approximations. Then,1. |W 2 (N X , N Y ) -W 2 (X, Y )| ≤ 2 Tr (Σ(X)Σ(Y )) 1 2 √ Tr[Σ(X)]+Tr[Σ(Y )] . 2. For T aff as in (4), W 2 (T aff X, Y ) ≤ √ 2 Tr [Σ(Y )] 1 2 .Proof. By Theorem 3.4, we have W 2 (N X , N Y ) ≤ W 2 (X, Y ). On the other hand,W 2 2 (X, Y ) = min γ∈ADM(X,Y ) R d ×R d ∥x -y∥ 2 dγ(x, y) ≤ R d ×R d ∥x∥ 2 + ∥y∥ 2 dγ(x, y) = Tr[Σ(X)] + Tr[Σ(Y )].Combining the above inequalities, we get|W 2 (N X , N Y ) -W 2 (X, Y )| ≤ Tr[Σ(X)] + Tr[Σ(Y )] -W 2 (N X , N Y ) . Let a = Tr[Σ(X)] + Tr[Σ(Y )], and so W 2 2 (N X , N Y ) = a -b, where b = 2 Tr (Σ(X)Σ(Y )) 1 2 .Then the RHS of can be written as√ a - √ a -b = |a -(a -b)| √ a + √ a -b ≤ b √ a ,where the inequality follows from positivity of W 2 (N X , N Y ) = √ a -b. Letting X = T aff X in the obtained bound gives 2).Many works aimed to improve the way how the non-linearity -represented by activation functionscan be defined in DNNs. As an example, a recent survey on the commonly used activation functions in deep neural networks [40] identifies over 40 activation functions with first references to sigmoid dating back to the seminal paper [36] published in late 80s. The fashion for activation functions used in deep neural networks evolved over the years in a substantial way, just as the neural architectures themselves.Saturating activations, such as sigmoid and hyperbolic tan, inspired by computational neuroscience were a number one choice up until the arrival of rectifier linear unit (ReLU) in 2010. After being the workhorse of many famous models over the years, the arrival of transformers popularized Gaussian Error Linear Unit (GELU) which is now commonly used in many large language models including GPTs.We illustrate in Figure 6 the affinity scores obtained after a single pass of the data through the following activation functions: Sigmoid, ReLU [51] , GELU [52] , ReLU6 [53] , LeakyReLU [54] with a default value of the slope, Tanh, HardTanh, SiLU [55] , and HardSwish [56] . As the non- This plot allows us to derive several important conclusions. 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. We note that in terms of 1) both GELU and ReLU may attain affinity scores that are close to 0, which is not the case for Sigmoid. For 2), we observe that the non-linearity of Sigmoid and GELU is maintained in a wide range, while for ReLU it is rather narrow. We can also see a distinct pattern of more modern activation functions, such as SiLU and HardSwish having a stronger non-linearity pattern in large subdomains. We also note that despite having a shape similar to Sigmoid, Tanh may allow for much lower affinity scores. Finally, the variations of ReLU seem to have a very similar shape with LeakyReLU being on average more linear than ReLU and ReLU6. Dimensionality reduction Manipulating 4-order tensors is computationally prohibitive and thus we need to find an appropriate lossless function r to facilitate this task. One possible choice for r may be a vectorization operator that flattens each tensor into a vector. In practice, however, such flattening still leads to very high-dimensional data representations. In our work, we propose to use averaging over the spatial dimensions to get a suitable representation of the manipulated tensors. In Figure 7 (left), we show that the affinity score is robust wrt such an averaging scheme and maintains the same values as its flattened counterpart.The non-linearity signature requires calculating the affinity score over "wide" matrices. Indeed, after the reduction step is applied to a batch of n tensors of size h × w × c, we end up with matrices of size n × c where n may be much smaller than c. This is also the case when input tensors are 2D when the batch size is smaller than the dimensionality of the embedding space. To obtain a well-defined estimate of the covariance matrix in this case, we use a known tool from the statistics literature called Ledoit-Wolfe shrinkage [57] . In Figure 7 (right), weshow that shrinkage allows us to obtain a stable estimate of the affinity scores that remain constant in all regimes.Robustness to batch size and different seeds In this section, we highlight the robustness of the non-linearity signature with respect to the batch size and the random seed used for training. To this end, we concentrate on VGG16 architecture and CIFAR10 dataset to avoid costly Imagenet retraining.In Figure 8 , we present the obtained result where the batch size was varied between 128 and 1024 with an increment of 128 (left plot) and when VGG16 model was retrained with seeds varying from 1 to 9 (right plot). The obtained results show that the affinity score is robust to these parameters suggesting that the obtained results are not subject to a strong stochasticity. Impact of training Finally, we also show how a non-linearity signature of a VGG16 model looks like at the beginning and in the end of training on Imagenet. We extract its non-linearity signature at initialization when making a feedforward pass over the whole CIFAR10 dataset and compare it to the non-linearity signature obtained in the end. In Figure 9 , we can see that at initialization the network's non-linearity signature is increasing, reaching almost a perfectly linear pattern in the last layers. Training the network enhances the non-linearity in a non-monotone way. Importantly, it also highlights that the non-linearity signature is capturing information from the training process. In Figure 10 , we portray the raw non-linearity signatures of several representative networks studied in the main paper. We use different color codes for distinct activation functions appearing repeatedly in the considered architecture (for instance, every first ReLU in a residual block of a Resnet). We also indicate the mean standard deviation of the affinity scores over batches in the title.We see that the non-linearities across ReLU activations in all of Alexnet's 8 layers remain stable. Its successor, VGG network, reveals tiny, yet observable, variations in the non-linearity propagation with increasing depth and, slightly lower overall non-linearity values. We attribute this to the decreased size of the convolutional filters (3x3 vs. 7x7). The Googlenet architecture was the first model to consider learning features at different scales in parallel within the so-called inception modules.This add more variability as affinity scores of activation in Googlenet vary between 0.6 and 0.9.Despite being almost 20 times smaller than VGG16, the accuracy of Googlenet on Imagenet remains comparable, suggesting that increasing and varying the linearity is a way to have high accuracy with a limited computational complexity compared to predecessors. This finding is further confirmed with Inception v3 that pushed the spread of the affinity score toward being more linear in some hidden layers. When comparing this behavior with Alexnet, we note just how far we are from it. Resnets achieve the same spread of values of the non-linearity but in a different, and arguably, simpler way.Indeed, the activation after the skip connection exhibits affinity scores close to 1, while the activations in the hidden layers remain much lower. Densenet, that connect each layer to all previous layers and not just to the one that precedes it, is slightly more non-linear than Resnet152, although the two bear a striking similarity: they both have an activation function that maintains the non-linearity low with increasing depth. Additionally, transition layers in Densenet act as linearizers and allow it to reset the non-linearity propagation in the network by reducing the feature map size. ViTs (Large with 16x16 and 32x32 patch sizes, and Huge with 14x14 patches) are all highly non-linear models to the degree yet unseen. Interestingly, as seen in Figure 11 the patch size affects the non-linearity propagation in a non-trivial way: for 16x16 size a model is more non-linear in the early layers, while gradually becoming more and more linear later, while 32x32 patch size leads to a plateau in the hidden layers of MLP blocks, with a steep change toward linearity only in the final layer. We hypothesize that attention modules in ViT act as a focusing lens and output the embeddings in the domain where the activation function is the most non-linear.Finally, we explore the role of increasing depth for VGG and Resnet architectures. We consider VGG11, VGG13, VGG16 and VGG19 models in the first case, and Resnet18, Resnet34, Resnet50, Resnet101 and Resnet152. The results are presented in Figure 12 and Figure 13 for VGGs and Resnets, respectively. Interestingly, VGGs do not change their non-linearity signature with increasing depth. In the case of Resnets, we can see that the separation between more linear post-residual activations becomes more distinct and approaches 1 for deeper networks. We consider the following metrics as 1) the linear CKA [38] commonly used to assess the similarity of neural representations, the average change in 2) SPARSITY and 3) ENTROPY before and after the application of the activation function as well as the 4) Frobenius NORM between the input and output of the activation functions, and the 5) R 2 score between the linear model fitted on the input and the output of the activation function. We present in Table 2 , the detailed values of Pearson correlations obtained for each architecture and all the metrics considered in this study. In Figure 14 , we show the full matrix of pairwise DTW distances [50] obtained between architectures, then used to obtain the clustering presented in the main text. Below, we compare the results obtained on CIFAR10, CIFAR100 datasets as well as when the random data tensors are passed through the network. As the number of plots for all chosen 33 models on these datasets will not allow for a meaningful visual analysis, we rather plot the differences -in terms of the DTW distance -between the non-linearity signature of the model on Imagenet dataset with respect to three other datasets. We present the obtained results in Figure 15 .We can see that the overall deviation for CIFAR10 and CIFAR100 remains lower than for Random dataset suggesting that these datasets are semantically closer to Imagenet. In this section, we show that the non-linearity signature of a network remains almost unchanged when considering other pertaining methodologies such as for instance, self-supervised ones. To this end, we use 17 Resnet50 architecture pre-trained on Imagenet within the next 3 families of learning approaches:1. SwAV [58] , DINO [59] , and MoCo [60] that belong to the family of contrastive learning methods with prototypes; 2. Resnet50 [18] , Wide Resnet50 [61] , TRex, and TRex* [62] that are supervised learning approaches;3. SCE [63] , Truncated Triplet [64] , and ReSSL [65] that perform contrastive learning using relational information.From the dendrogram presented in Figure 16 , we can observe that the DTW distances between the non-linearity signatures of all the learning methodologies described above allow us to correctly cluster them into meaningful groups. This is rather striking as the DTW distances between the different instances of the Resnet50 model are rather small in magnitude suggesting that the affinity scores still retain the fact that it is the same model being trained in many different ways.While providing a fine-grained clustering of different pre-trained models for a given fixed architecture, the average affinity scores over batches remain surprisingly concentrated as shown in Table 3 . This hints at the fact that the non-linearity signature is characteristic of architecture but can also be subtly multi-faceted when it comes to its different variations. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Experimental details are described in Section 4 and Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Standard deviations across multiple batch of data are reported.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: All experiments are carried out on a single A100 GPU.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss broader impacts in Appendix A.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: No such risks, no checkpoints released.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Torchvision contributors credited for checkpoints, and datasets as well, in Section 4.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: This paper introduces a method to quantify the non-linearity of individual activation functions in neural networks using common optimal transport results. It reveals distinct patterns in the evolution of non-linearity across vision architectures. The reviewers have raised various valid concerns about the writing or the applications of this paper. In addition, concerns about the novelty of the theoretical results have been mentioned, but those have been of secondary nature in the decision. The writing along with the insights should be improved, along with the rigor of the theoretical results.

Reviews: review: paper_summary: paper_summary: This paper introduces a novel method for quantifying the non-linearity of activation functions in neural networks, termed the "non-linearity signature." Using an affinity score derived from optimal transport theory, it measures the non-linearity of individual activation functions. It defines the non-linearity signature as a comprehensive set of these scores across all functions in a deep neural network (DNN). The study compares these signatures across a range of popular DNN architectures in computer vision, revealing clear patterns in their evolution over the past decade, notably showing a trend towards decreasing non-linearity until the disruptive impact of vision transformers. It emphasizes the uniqueness of their measure, as it does not strongly correlate with other metrics across different architectures. The approach could potentially be applied to analyze the non-linearity of newer large language models (LLMs) and identify innovative neural architectures that optimize internal non-linear characteristics for enhanced performance, crucial in the era of costly experiments with large-scale model optimizations.
strength_weakness: strength_weakness: 1. The authors should discuss more activation functions. Currently, only ReLU, Tanh, and Sigmoid are included.  While these are among the most commonly used activation functions in neural networks, many other activation functions have been introduced and shown to be effective in various contexts, like GELU. Including a more comprehensive analysis of a diverse set of activation functions would enhance the robustness and applicability of their proposed method. 
2. There is currently some research on the nonlinearity of deep neural networks that should be compared and discussed.
3.  It would be beneficial to showcase examples from domains beyond computer vision. While the paper focuses on computer vision tasks, it may not address the non-linearity signature's applicability to other domains such as NLP, speech recognition, or reinforcement learning. The findings might be less generalizable if the proposed measure does not perform equally well across diverse types of tasks and data.
questions: questions: How would the model’s nonlinearity affect different downstream tasks? Does the impact vary across datasets?
limitations: limitations: Yes, they have discussed the assumption of Theorem 3.3.
paper_summary: paper_summary: This study proposes empirical statistics about different DNN architectures in the hope to shed some light into why some architectures are better than others for some computer vision tasks. To do so, the study leverages common optimal transport results on DNN's internal representations, under some strong assumption about the distribution of those representations.
strength_weakness: strength_weakness: - I do not agree with the following statement `Without non-linear activation functions,
84 most of DNNs, no matter how deep, reduce to a linear function unable to learn complex patterns.` as to me, models such as transformers with linear attention and linear MLP blocks have no actual nonlinearity but are higher order polynomial of the input, i.e., are not linear. Could the authors provide clarifications on that statement or did I misunderstand something?

- I also disagree with the following `Activation functions were also early identified [29, 30, 31, 32] as a key to making even a shallow
86 network capable of approximating any function, however complex it may be, to arbitrary precision.` since again, Fourier series for example can approximate any function as well. Hence DNN nonlinearities are certainly not the key ingredient to function approximation in general

- Many formal results such as Theorem 3.3 are well known and have been established for years (even decades) but no reference is provided which is misleading to the reader.

- Fig 2. is also misleading since the "nonlinearity" of any activation function depends on the range of the inputs. The only case that wouldn't be true is e.g. for ones with constant second derivatives, i.e., a linear activation function.... hence again that statement is highly misleading in presenting ReLU as inherently benefiting form that property compared to others

- the statement `No other metric extracted from the activation functions of the
260 considered networks exhibits a strong consistent correlation with the non-linearity signature.` is again an overstatement as the authors only compare with a few alternatives and theorem is provided to support such a statement

- the statement `We proposed the first sound approach to measure non-linearity of activation functions in neural
270 networks` is also incorrect, see e.g. 
  - https://jmlr.org/papers/v20/18-418.html
  - https://arxiv.org/pdf/2301.09554
  - https://arxiv.org/abs/1810.09274
  all the above works have been published in peer reviewed journals/conferences
questions: questions: Please see the **Weaknesses** section
limitations: limitations: In addition to my concerns expressed above, the study does not provide any actionable insights or understanding on the "why" of different architectures performing differently beyond the proposed statistical numbers. How could one use the provided analysis to better design model architectures or for model selection? 

Also, the paper does not provide any novel theoretical results. All the major theorems and results are already widely known within the community, yet they are presented as part of the contributions. With that in mind, the paper solely leverages existing OT tools, with some underlying simplifications on the DNN's data distribution, and report computed metrics. Hence the study falls below acceptance level in my opinion and would need a major rewriting + additional novel contributions to be worth acceptance.

The writing style is also filled with unsupported claims and highly misleading statements (see the **Weaknesses** examples).
paper_summary: paper_summary: This paper proposes the affinity score, which measures the non-linearity of an activation function $\sigma(X)$ given the distribution of $X$.
The affinity score is defined based on how well the 2-Wasserstein distance $W_2(X, Y)$, where $Y=\sigma(X)$, is approximated by $W_2(N_X, N_Y)$, where $N_X$ and $N_Y$ are Gaussian approximations of the distributions of $X$ and $Y$, respectively.
Note that $W_2(N_X, N_Y)$ has a closed-form solution, and it holds that $W_2(X, Y) = W_2(N_X, N_Y)$ if the relation between $X$ and $Y$ is locally affine on the support of the given $X$.
The authors then propose to characterize a DNN model by the set of affinity scores of activation functions in the model under a given input distribution.
Experimental results suggest that the affinity scores are relatively low in transformer-based vision models, meaning that the activation functions are used in a more non-linear region compared to CNN models.
strength_weakness: strength_weakness: * It is empirically shown that the proposed score has a low correlation with existing non-linearity metrics such as R^2, but it is unclear whether the existing metrics are insufficient to analyze the DNN models in the way proposed in this paper. 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}$.
* In my opinion, one would expect the nonlinearity score to behave symmetrically at $x=0$ for activation functions like ReLU, but the proposed affinity score seems to have a lower score at negative $x$, as shown in Fig.2 or Fig.6. Is there any reasonable explanation for such a behavior of the proposed score?
questions: questions: * 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.

* This is just a comment out of interest, I would like to see if the affinity score of a neuron in the same model changes depending on the input class.
limitations: limitations: See weakness above.

Comments: title: Thank you for your reply and reconsidering the score
Thank you
Thank you
comment: We would like to thank the reviewer for reconsidering their score.

> I find that many insights are echoing previous studies

If our study echoes other known results on this topic, we would gladly discuss and contextualize them in our work if the reviewer kindly provides us with more specific pointers to such studies. There seem to be no other metrics for measuring non-linearity as we propose to do it but we remain open to discovering other works deriving similar insights with other tools. We genuinely tried to find reasonably related works for this contribution but didn't manage to find them.

> I am still unconvinced that the method allows for "actionable results"

As for the providing actionable results, can the reviewer elaborate a bit more on what is considered actionable? We believe that our study provides comparisons between models in a theoretically-grounded way. It is thus actionable as it allows to quantify a quantity of interest and use this quantity in a variety of ways (to compare models, to understand the role of different parts of the network etc). We believe that such a comparison can be done in future works as well, acting as a tool in them. We remain open to discussing it further until the end of the rebuttal period if the reviewer shares with us they understand by actionable.
We thank the reviewer for their reply. We are happy that the reviewers found the provided illustrations insightful and we will include them in the new revision as suggested.
I want to thank the authors for providing clear answers to each of my concerns. In light of those findings and while considering the other reviews I have updated my score. I still believe that some improvements could be done in terms of presentation (which would have helped my initial assessment) and I am still unconvinced that the method allows for "actionable results". Also, I find that many insights are echoing previous studies and are not surprising (which is why I raised my score to 4 and not 5).
Thank you for your responses.

From Figure 1 in the supplemental pdf, it appears that R^2 has a similar tendency to $\rho_{aff}$, e.g., the bimodality of ResNet152 and the less linear behavior of Vit Huge 14. As the author notes, it is curious that R^2 becomes negative.

Figure 4 in the supplemental pdf is interesting, showing another difference in Conv nets and ViTs. I would like the authors to include these results in the main manuscript.

I would like to keep my score at 7.