Research map / Architectures and efficiency / RNNs and normalization
From Alexnet to Transformers: Measuring the Non-linearity of Deep Neural Networks with Affine Optimal Transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces the affinity score, a closed-form optimal-transport-based measure of activation non-linearity, and aggregates it into per-network "non-linearity signatures" to compare ImageNet vision architectures from AlexNet to transformers.
Problem
Existing theory relates DNN expressive power to depth or width and cannot compare architecture families; there is no widely accepted measure of how non-linear a network or its activation functions are.
Approach
Measure activation non-linearity as the normalized 2-Wasserstein gap between the output and its best affine OT map from the input, using closed-form affine OT between Gaussian approximations (Gelbrich bound); collect these affinity scores over all activations of a network into a non-linearity signature.
Key results
- Affinity-score statistics correlate with ImageNet-1K accuracy; e.g., the minimum affinity score gives R2 of 0.974 for pre-residual convnets and the number of parameters is most predictive for NAS-based models (R2 0.955).
- Vision transformers break the trend of decreasing non-linearity seen in earlier landmark architectures, becoming the most non-linear among landmark models.
- Residual connections produce a bimodal affinity-score distribution, with one mode centered on highly linear activations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase Glorot weight magnitude from U(-s,s) to U(-6s,6s) for ReLU (ablation) | 3-layer randomly initialized MLP · 2D coordinate-grid inputs | ReLU affinity score: 0.9(±0.04) -> 0.82(±0.07) | Helped |
| Increase Glorot weight magnitude from U(-s,s) to U(-6s,6s) for Tanh (ablation) | 3-layer randomly initialized MLP · 2D coordinate-grid inputs | Tanh affinity score: 0.96(±0.02) -> 0.73(±0.05) | Helped |
| Use s′ = s − 0.05 and increase weight magnitude to U(-6s′,6s′) (ablation) | 3-layer randomly initialized MLP with ReLU · 2D coordinate-grid inputs | ReLU affinity score: 0.9(±0.05) -> 0.65(±0.1) | Helped |
| Use s′ = s − 0.05 and increase weight magnitude to U(-6s′,6s′) (ablation) | 3-layer randomly initialized MLP with Tanh · 2D coordinate-grid inputs | Tanh affinity score: 0.96(±0.01) -> 0.74(±0.05) | Helped |
| Add residual connections (ablation) | ResNet20 · ImageNet-1K | distribution of affinity scores | Mixed |
Surprising findings
- No 'perfect' non-linearity exists in the 1D formulation: the affinity score approaches 0 only when both terms vanish, i.e., when the activation receives only negative inputs.
- Changing the Glorot initialization scale slightly makes ReLU networks behave much like tanh networks, contradicting the neural redshift claim that ReLU has a universal simplicity bias.
Limitations
- Affinity score relies on Gaussian approximations of activation inputs and outputs, so it captures only second-order (mean and covariance) structure.
- Analysis is limited to ImageNet computer vision models using pretrained checkpoints.
- Correlations with accuracy are computed over small families of models, so they may not generalize.
Experimental setup
Datasets: ImageNet-1K validation set
Benchmarks: ImageNet-1K
Metrics: affinity score Pearson correlation R2 score dynamic time warping distance
Baselines: R2 score of linear regression mapping Non-linearity coefficient (Ru et al.) CKA depth number of parameters GFLOPS
Builds on: Gelbrich bound neural redshift Kantorovich optimal transport
Closest papers by meaning
- Nonlinear Advantage: Trained Networks Might Not Be As Complex as You Think (ICML 2023 · similarity 0.88)
- From Hard to Soft: Understanding Deep Network Nonlinearities via Vector Quantization and Statistical Inference (ICLR 2019 · similarity 0.87)
- On the Hardness of Training Deep Neural Networks Discretely (AAAI 2025 · similarity 0.86)
- Curved Representation Space of Vision Transformers (AAAI 2024 · similarity 0.85)
- Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit (ICLR 2024 · similarity 0.85)
- What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.