FlexAct: Why Learn when you can Pick?

TL;DR

FlexAct introduces a differentiable discrete selection mechanism for activation functions using Gumbel-Softmax, enabling neural networks to adaptively choose the most suitable activation during training, improving accuracy and flexibility.

Contribution

We propose a novel framework that allows neural networks to learn discrete activation functions dynamically, combining the Gumbel-Softmax trick with task-specific adaptation.

Findings

Consistently selects optimal activation functions on synthetic datasets

Enhances predictive accuracy and architectural flexibility

Bridges theoretical advances with practical utility

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.