# FlexAct: Why Learn when you can Pick?

arXiv: https://arxiv.org/abs/2601.06441 · HTML: https://mlsanity.com/paper/2601.06441/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 42 references that are in the map.
- First posted on arXiv 10 Jan 2026.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [KAN: Kolmogorov-Arnold Networks](https://mlsanity.com/paper/2404.19756/) (2024)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [The Shattered Gradients Problem: If resnets are the answer, then what is the question?](https://mlsanity.com/paper/1702.08591/) (2017)
- [Activation Functions in Deep Learning: A Comprehensive Survey and Benchmark](https://mlsanity.com/paper/2109.14545/) (2021)
- [Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice](https://mlsanity.com/paper/1711.04735/) (2017)
- [Activation Functions: Comparison of trends in Practice and Research for Deep Learning](https://mlsanity.com/paper/1811.03378/) (2018)
- [Rational neural networks](https://mlsanity.com/paper/2004.01902/) (2020)
- [GANS for Sequences of Discrete Elements with the Gumbel-softmax Distribution](https://mlsanity.com/paper/1611.04051/) (2016)
- [Learning Activation Functions: A new paradigm for understanding Neural Networks](https://mlsanity.com/paper/1906.09529/) (2019)
- [Characterizing the Spectrum of the NTK via a Power Series Expansion](https://mlsanity.com/paper/2211.07844/) (2022)
- [Learning Combinations of Activation Functions](https://mlsanity.com/paper/1801.09403/) (2018)
- [Comparison of non-linear activation functions for deep neural networks on MNIST classification task](https://mlsanity.com/paper/1804.02763/) (2018)
- [SPLASH: Learnable Activation Functions for Improving Accuracy and Adversarial Robustness](https://mlsanity.com/paper/2006.08947/) (2020)
- [End-to-end learnable EEG channel selection for deep neural networks with Gumbel-softmax](https://mlsanity.com/paper/2102.09050/) (2021)
- [Adaptively Customizing Activation Functions for Various Layers](https://mlsanity.com/paper/2112.09442/) (2021)
- [APALU: A Trainable, Adaptive Activation Function for Deep Learning Networks](https://mlsanity.com/paper/2402.08244/) (2024)

## Abstract (arXiv)

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.