Pith. sign in

REVIEW 4 major objections 6 minor 24 references

Gumbel-Softmax routing lets a network select its activation on the fly.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-03 11:23 UTC pith:2F5KS5VY

load-bearing objection Honest but thin: the regularizer does the selection work, Table 1 contradicts the headline claim, and the 'deep' experiments never happen—worth a reading group as a cautionary example, not a citable result. the 4 major comments →

arxiv 2601.06441 v2 pith:2F5KS5VY submitted 2026-01-10 cs.LG

FlexAct: Why Learn when you can Pick?

Ramnath Kumar , Kyle Ritscher , Junmin Judy , Lawrence Liu , Cho-Jui Hsieh This is my paper
classification cs.LG
keywords activation functionsGumbel-Softmaxdiscrete activation selectiongradient-norm regularizationlearned nonlinearitysynthetic regressionadaptive architectures
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper claims that a single neural network can learn to choose its own activation function from a small fixed set during training, instead of committing to one function in advance. It proposes Flex-Act, which uses the Gumbel-Softmax trick to make the discrete choice differentiable and learns routing probabilities jointly with the weights. To prevent the selection from biasing toward unbounded activations like ReLU, the paper adds a regularizer based on the gradient norm of each candidate function. On synthetic regression tasks where the true activation is known, Flex-Act recovers the correct function and matches or beats every fixed-activation baseline in all settings. If true, this would remove the need for activation hyperparameter sweeps and let one architecture adapt to different tasks.

Core claim

The central discovery is that discrete activation selection can be trained end-to-end with Gumbel-Softmax, and that without correction the selection is biased toward unbounded functions, but with a gradient-norm regularizer the model correctly selects Sigmoid or Tanh when those are the ground truth. In the paper's controlled single-layer experiments, the regularized model converges to the correct activation in every one of the five cases tested (ReLU, Sigmoid, Tanh, LeakyReLU, Identity), with mean squared error matching the best fixed-activation baseline in each column.

What carries the argument

The mechanism is a penultimate-layer router that computes a convex combination of candidate activations using Gumbel-Softmax, which allows gradients to flow through the discrete choice. The bias correction computes per-input gradient norms g_i = ||∇_x φ_i(h)|| for each candidate, converts their negative batch averages into pseudo-probabilities via softmax(-ḡ_i / λ), and adds a KL divergence between these pseudo-probabilities and the router's logits to the task loss. This regularizer counteracts the tendency for large-magnitude, unbounded activations to dominate gradient flow.

Load-bearing premise

The load-bearing premise is that the softmax of negative gradient norms (Equation 2) reliably indicates which activation function is functionally best for the task, rather than merely reflecting magnitude effects.

What would settle it

Run Flex-Act on a synthetic task where the ground-truth activation is Sigmoid but a different activation has a much larger gradient norm throughout training; if the router still picks the high-norm activation, the regularizer is not really selecting for functional suitability. Alternatively, replace the gradient-norm prior with random pseudo-probabilities and check whether selection accuracy collapses.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Removes the need for manual activation sweeps: one network can be trained once and select the right nonlinearity per task.
  • Enables the same architecture to be deployed across tasks with different functional regimes without redesign.
  • Provides interpretability: the learned selection probabilities reveal which activation each layer prefers.
  • Suggests a path toward layer-wise activation selection in deep networks, though the paper notes that naive extension would face gradient alignment and stability problems.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The paper does not separate whether the router succeeds because of the gradient-norm prior or despite it; a clean test would compare against a version using the task loss only, or random pseudo-probabilities.
  • The evidence is limited to a single-layer synthetic regression with one informative feature; whether the selection signal survives in deep networks or real data distributions is untested.
  • Because the gradient-norm prior is what forces correct selection of Sigmoid/Tanh, the method may be trading one inductive bias (fixed activation choice) for another (a fixed routing prior).

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

4 major / 6 minor

Summary. The manuscript proposes Flex-Act, a method that uses Gumbel-Softmax routing to select an activation function from a fixed set at a network layer, and introduces a gradient-norm-based KL regularizer that biases selection toward low-gradient-norm activations. The evaluation is a synthetic regression task with one informative input and five ground-truth activations (ReLU, Sigmoid, Tanh, LeakyReLU, Identity), comparing Flex-Act against fixed-activation baselines. The paper claims that Flex-Act 'consistently selects the most suitable activation function' and 'matches or outperforms all baselines in every setting without any additional computational overhead' (Section 4.1.1), and frames this as enabling plug-and-play generalization.

Significance. The idea of discrete activation selection is a legitimate and potentially useful alternative to parameterized activations, and the synthetic ground-truth setup is a sensible way to test identifiability. The α=0 ablation and selection-probability plots are transparent and helpful. However, the paper's empirical and conceptual support is not sufficient: Table 1 contradicts the headline performance claim; the regularizer in Eq. (2) encodes a hand-designed prior that favors low-gradient-norm activations and is independent of the task loss, so the 'discovery' of Sigmoid/Tanh is partly manufactured; and the stated contributions on gradient derivations and deep-setting gains are not delivered by the manuscript. The paper honestly acknowledges the heuristic nature of the regularizer and the single-layer scope in Section 5, but those acknowledgments undercut the broad claims in the abstract and conclusion.

major comments (4)
  1. [Section 4.1.1, Table 1] The claim that Flex-Act 'matches or outperforms all baselines in every setting' is contradicted by Table 1. For the ReLU, Sigmoid, Tanh, and LeakyReLU ground truths, Flex-Act (α=0.3) has MSE 0.0001, 0.0011, 0.0001, and 0.0001, respectively, while the fixed activation matching the ground truth has 0.0000 in every case. Flex-Act is never the best performer except for Identity, and for Sigmoid it is an order of magnitude worse than the fixed Sigmoid baseline. At best it is second-best. The claim of 'no additional computational overhead' is also unsupported: computing the gradient norms in Eq. (1)–(2) for all candidate activations requires additional backward computations.
  2. [Section 3.3, Eq. (2)] The pseudo-probabilities p̃_i = softmax(-ḡ_i / λ) depend only on the input distribution and current weights, not on the target y or the task loss. Minimizing the KL term in Eq. (3)–(4) therefore always pushes the router toward the candidate with the smallest gradient norm, which is typically a saturating function such as Sigmoid or Tanh. In the Sigmoid/Tanh experiments this coincides with the ground truth, so Figures 1–2 do not demonstrate task-driven discovery. Table 1 shows that α=0 fails to select these activations, confirming that the regularizer, not the task loss, is doing the work. Section 5 admits the regularizer is heuristic. This undercuts the internal validity of the 'consistent selection' claim.
  3. [Section 3.2 and Contributions] The second contribution bullet promises 'analytically derive the backpropagation equations for networks employing discrete activation routing,' but Section 3.2 only restates the standard Gumbel-Softmax reparameterization and the convex combination; no gradient derivation is supplied. The third bullet promises 'empirical gains in deep settings,' but Section 4 contains only single-layer synthetic regression; no deep network is trained. Both stated contributions are unsupported by the manuscript.
  4. [Section 5 and Conclusion] The Discussion states that Flex-Act is 'primarily designed to be only added at the penultimate layer,' that naive extension 'would face issues with gradient alignment, exploding gradients, and unstable training,' and that 'our current regularizer is heuristic and derived from empirical intuition.' These limitations directly contradict the abstract's claim that Flex-Act 'dynamically learns the optimal activation function' and the Discussion's opening claim of 'purely data-driven' discovery. As written, the paper's conclusions substantially exceed its evidence.
minor comments (6)
  1. [Table 1 vs. Figure 3] The reported MSEs for the same Flex-Act configuration differ between Table 1 and Figure 3 captions; for example, Tanh is 0.0001 in Table 1 but 0.0006 in Figure 3, and LeakyReLU is 0.0001 in Table 1 but 0.0000 in Figure 3. The numbers should be reconciled.
  2. [Eq. (4)] The symbol N in Eq. (4) is not defined. If it is the batch size, it should say so; if it is the number of layers, the single-layer setting should be clarified.
  3. [Figure 2] Figure 2 does not state which ground-truth activation is being plotted. The caption should specify the setting, otherwise the effect of α on selection probabilities is ambiguous.
  4. [Typos and formatting] There are several typographical errors: 'penulimate' (Sections 3.0.5, 4.1), 'hypermater' (Related Work), and 'T anh' in Figure 1. Also, the text alternates between 'each layer' and 'the penultimate layer' when describing where Flex-Act is applied; this should be made consistent.
  5. [References] Several references are incomplete: He et al. 2015a and 2015b point to the same URL, Devlin (2018) and Vaswani (2017) lack full author lists and venue information, and a number of entries are missing page numbers or conference details.
  6. [Table 1 caption] The caption states that the best result is bolded and the second-best is marked with †, but Flex-Act rows carry † in columns where it is tied with the fixed model (e.g., Identity) or is not clearly second-best. The formatting rules should be applied consistently.

Circularity Check

1 steps flagged

Sigmoid/Tanh 'discovery' is largely induced by the gradient-norm KL target (Eq. 2–4), so the central selection claim partially reduces to the regularizer's functional form.

specific steps
  1. fitted input called prediction [Section 3.3, Eqs. (1)–(4); Section 4.1.1, Figure 2; Section 5 Discussion]
    "In preliminary experiments, we observed that the model consistently favored ReLU and Leaky ReLU where the activations have large unbounded outputs regardless of data-specific structure... We convert the negative average gradient norms into pseudo-probabilities using a softmax function: p̃_i = exp(−ḡ_i/λ)/Σ_j exp(−ḡ_j/λ), where ḡ_i is the average gradient norm for activation function φ_i over the batch... The total loss is then a combination of the primary task loss ... and the Kullback-Leibler divergence ... L = L_task + αL_KL."

    The KL term in Eq. (4) explicitly drives the router probabilities p_i toward p̃_i, which by Eq. (2) is a deterministic function of negative average gradient norms and does not depend on the labels y or the task loss. Because bounded/saturating activations such as Sigmoid and Tanh have smaller |∇_x φ| than ReLU in the active regime, the regularizer by construction assigns them higher pseudo-probability. Thus the α=0.3 'discovery' of Sigmoid/Tanh in Figure 2 is the regularizer's own target, not an emergent data-driven result. The paper itself shows α=0 fails on these cases and that the same synthetic tasks motivated the heuristic; the Discussion even admits 'our current regularizer is heuristic.' Hence the central selection claim is partially manufactured by the regularizer's functional form

full rationale

The main circularity is in the bias-correction mechanism. Eq. (2) defines pseudo-probabilities from gradient norms only, and Eq. (3)–(4) force the Gumbel-Softmax router to match those pseudo-probabilities. Therefore, when the ground truth is Sigmoid or Tanh, the regularized model's 'selection' of those activations is substantially predetermined by the choice of softmax(-ḡ/λ), which favors low-gradient-norm functions. The paper's own ablation (α=0) confirms that without this term the model does not select Sigmoid/Tanh, so the corrected selection is not an emergent property of the learning task. This is a fitted-input-called-prediction pattern: the regularizer was introduced after observing the ReLU bias on these very tasks, and the evaluation then presents the regularized outcome as data-driven discovery. I did not find load-bearing self-citation or imported uniqueness theorems; those patterns are absent. However, because the task loss does contribute for ReLU/LeakyReLU/Identity and the framework itself is not vacuous, the circularity is partial rather than total: score 6.

Axiom & Free-Parameter Ledger

6 free parameters · 4 axioms · 0 invented entities

The central claim rests on a heuristic regularizer (the gradient-norm pseudo-probabilities) applied to a toy synthetic task, plus standard Gumbel-Softmax machinery. There are no new physical or mathematical entities; the free parameters are mostly unspecified or chosen by hand.

free parameters (6)
  • alpha (KL regularization weight) = 0.3 (ablation at 0.0)
    Chosen to correct the observed ReLU bias; no tuning procedure reported, and selection results depend strongly on it.
  • lambda (gradient-norm scaling) = not specified
    Equation (2) requires lambda to convert gradient norms to pseudo-probabilities, but no value or schedule is given.
  • tau (Gumbel-Softmax temperature) = not specified
    Annealing schedule not described; controls hardness of selection and affects optimization.
  • k (ground-truth scale) = 5
    The authors state it is 'arbitrarily set to 5 for the sake of simplicity' (Section 4.1).
  • LeakyReLU slope = not specified
    Defined as a 'small positive constant' with no value given (Section 3.0.4).
  • network width/architecture = not specified
    Experiments use a single-layer network, but width, optimizer, and initialization are not reported.
axioms (4)
  • standard math Gumbel-Softmax provides unbiased/low-variance gradient estimates for discrete selection.
    Standard result from Jang et al. (2017), used without proof.
  • ad hoc to paper Average gradient norms of candidate activations are a valid signal of functional suitability.
    This is the basis of Eq. (2)-(4); the authors themselves call the regularizer 'heuristic and derived from empirical intuition' (Section 5).
  • domain assumption Single-layer synthetic regression with one informative feature is representative of activation selection in deep networks.
    The experiments only cover this toy setting, yet the paper claims relevance to deep architectures.
  • domain assumption A predefined set of five activations is sufficient to demonstrate 'optimal' selection.
    The candidate set excludes GELU, Swish, and many parameterized forms; the claim of optimality is restricted to this menu.

pith-pipeline@v1.3.0-alltime-deepseek · 11835 in / 12013 out tokens · 115523 ms · 2026-08-03T11:23:18.490523+00:00 · methodology

0 comments
read the original abstract

Learning activation functions has emerged as a promising direction in deep learning, allowing networks to adapt activation mechanisms to task-specific demands. In this work, we introduce a novel framework that employs the Gumbel-Softmax trick to enable discrete yet differentiable selection among a predefined set of activation functions during training. Our method dynamically learns the optimal activation function independently of the input, thereby enhancing both predictive accuracy and architectural flexibility. Experiments on synthetic datasets show that our model consistently selects the most suitable activation function, underscoring its effectiveness. These results connect theoretical advances with practical utility, paving the way for more adaptive and modular neural architectures in complex learning scenarios.

Figures

Figures reproduced from arXiv: 2601.06441 by Cho-Jui Hsieh, Junmin Judy, Kyle Ritscher, Lawrence Liu, Ramnath Kumar.

Figure 1
Figure 1. Figure 1: Gumbel-Softmax activation selection probabilities over training epochs for various ground-truth [PITH_FULL_IMAGE:figures/full_fig_p010_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Effect of KL-based regularization on activation selection. Without regularization, the model biases [PITH_FULL_IMAGE:figures/full_fig_p011_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Ablation analysis comparing our proposed method ( [PITH_FULL_IMAGE:figures/full_fig_p015_3.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

24 extracted references · 22 linked inside Pith

  1. [8]

    Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun

    URLhttps://arxiv.org/abs/1906.09529. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015a. URLhttps://arxiv.org/abs/1502.01852. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on image...

    Pith/arXiv arXiv 1906

  2. [9]

    Eric Jang, Shixiang Gu, and Ben Poole

    URLhttps://arxiv.org/abs/2112.09442. Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax,

    Pith/arXiv arXiv

  3. [11]

    URL https://arxiv.org/abs/2001.08361. Matt J. Kusner and José Miguel Hernández-Lobato. Gans for sequences of discrete elements with the gumbel-softmax distribution,

    Pith/arXiv arXiv 2001

  4. [12]

    Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljačić, Thomas Y

    URLhttps://arxiv.org/abs/1611.04051. Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljačić, Thomas Y. Hou, and Max Tegmark. Kan: Kolmogorov-arnold networks,

    Pith/arXiv arXiv

  5. [14]

    URLhttp://dx.doi.org/10.1109/ICPR.2018.8545362

    doi: 10.1109/ icpr.2018.8545362. URLhttp://dx.doi.org/10.1109/ICPR.2018.8545362. Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on Machine Learning (ICML-10), pp. 807–814,

    arXiv 2018

  6. [16]

    Comparison of non-linear activation functions for deep neural networks on mnist classifica- tion task.arXiv preprint arXiv:1804.02763,

    Dabal Pedamonti. Comparison of non-linear activation functions for deep neural networks on mnist classifica- tion task.arXiv preprint arXiv:1804.02763,

    Pith/arXiv arXiv

  7. [17]

    12 Preprint, Under Review Prajit Ramachandran, Barret Zoph, and Quoc V

    URLhttps://arxiv.org/abs/1711.04735. 12 Preprint, Under Review Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions,

    Pith/arXiv arXiv

  8. [18]

    URL https://arxiv.org/abs/1710.05941. David E. Rumelhart and James L. McClelland.Learning Internal Representations by Error Propagation, pp. 318–362

    Pith/arXiv arXiv

  9. [20]

    URLhttps://dx.doi.org/10.1088/1741-2552/ac115d

    doi: 10.1088/ 1741-2552/ac115d. URLhttps://dx.doi.org/10.1088/1741-2552/ac115d. Barathi Subramanian, Rathinaraja Jeyaraj, Rakhmonov Akhrorjon Akhmadjon Ugli, and Jeonghong Kim. Apalu: A trainable, adaptive activation function for deep learning networks.arXiv preprint arXiv:2402.08244,

    arXiv

  10. [21]

    Mohammadamin Tavakoli, Forest Agostinelli, and Pierre Baldi

    URLhttps://arxiv.org/abs/1602.07261. Mohammadamin Tavakoli, Forest Agostinelli, and Pierre Baldi. Splash: Learnable activation functions for improving accuracy and adversarial robustness,

    Pith/arXiv arXiv

  11. [22]

    A Vaswani

    URLhttps://arxiv.org/abs/2006.08947. A Vaswani. Attention is all you need.Advances in Neural Information Processing Systems,

    Pith/arXiv arXiv 2006

  12. [23]

    Hongyi Zhang, Yann N

    URLhttps://arxiv.org/abs/1611.03530. Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normaliza- tion,

    Pith/arXiv arXiv

  13. [24]

    URLhttps://arxiv.org/abs/1901.09321. 13 Preprint, Under Review A Appendix To assess the inductive flexibility of learned nonlinearities, we perform an ablation analysis comparing our approach (Flex-Act) with standard fixed activation models. The five subfigures in Figure 3 illustrate the predicted output values versus input features for each activation se...

    Pith/arXiv arXiv 1901

  14. [2010]

    Activation functions: Comparison of trends in practice and research for deep learning.arXiv preprint arXiv:1811.03378,

    Chigozie Nwankpa, Winifred Ijomah, Anthony Gachagan, and Stephen Marshall. Activation functions: Comparison of trends in practice and research for deep learning.arXiv preprint arXiv:1811.03378,

    Pith/arXiv arXiv

  15. [2014]

    Thomas Strypsteen and Alexander Bertrand

    URLhttps://arxiv.org/abs/1312.6120. Thomas Strypsteen and Alexander Bertrand. End-to-end learnable eeg channel selection for deep neural networks with gumbel-softmax.Journal of Neural Engineering, 18(4):0460a9, jul

    Pith/arXiv arXiv

  16. [2015]

    Bert: Pre-training of deep bidirectional transformers for language understanding.arXiv preprint arXiv:1810.04805,

    11 Preprint, Under Review Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding.arXiv preprint arXiv:1810.04805,

    Pith/arXiv arXiv

  17. [2016]

    David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams

    URL https: //arxiv.org/abs/1607.06450. David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question?,

    Pith/arXiv arXiv

  18. [2017]

    Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B

    URL https://arxiv.org/abs/1611.01144. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models,

    Pith/arXiv arXiv

  19. [2018]

    Yoshua Bengio, Nicholas Léonard, and Aaron Courville

    URL https://arxiv.org/abs/1702.08591. Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation.arXiv preprint arXiv:1308.3432,

    Pith/arXiv arXiv

  20. [2019]

    Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E

    URLhttps://arxiv.org/abs/ 1803.08375. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization,

    Pith/arXiv arXiv

  21. [2020]

    Djork-Arné Clevert

    URL https: //arxiv.org/abs/2004.01902. Djork-Arné Clevert. Fast and accurate deep network learning by exponential linear units (elus).arXiv preprint arXiv:1511.07289,

    Pith/arXiv arXiv 2004

  22. [2021]

    org/abs/2010.11929

    URLhttps://arxiv. org/abs/2010.11929. Shiv Ram Dubey, Satish Kumar Singh, and Bidyut Baran Chaudhuri. Activation functions in deep learning: A comprehensive survey and benchmark,

    Pith/arXiv arXiv 2010

  23. [2022]

    Xavier Glorot and Yoshua Bengio

    URLhttps://arxiv.org/abs/2109.14545. Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. InProceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249–256. JMLR Workshop and Conference Proceedings,

    Pith/arXiv arXiv

  24. [2024]

    Andrew L Maas, Awni Y Hannun, Andrew Y Ng, et al

    URLhttps://arxiv.org/abs/2404.19756. Andrew L Maas, Awni Y Hannun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. InProc. icml, volume 30, pp

    Pith/arXiv arXiv