Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FlexAct: Why Learn when you can Pick?

Ramnath Kumar, Kyle Ritscher, Junmin Judy, Lawrence Liu, Cho-Jui Hsieh

TL;DR

Activation functions are often fixed, hindering performance across scales and tasks. Flex-Act introduces layer-wise discrete activation routing using the Gumbel-Softmax trick to select among a fixed set of activations during training, with a gradient-norm-based regularizer to prevent bias toward unbounded functions. The approach yields end-to-end differentiable learning of per-layer activations, with analytical consideration of gradient propagation for discrete choices. In synthetic regression experiments, Flex-Act consistently recovers the correct activation and improves training stability, highlighting potential for robust, modular neural architectures that adapt their non-linear cores to task demands.

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.

FlexAct: Why Learn when you can Pick?

TL;DR

Activation functions are often fixed, hindering performance across scales and tasks. Flex-Act introduces layer-wise discrete activation routing using the Gumbel-Softmax trick to select among a fixed set of activations during training, with a gradient-norm-based regularizer to prevent bias toward unbounded functions. The approach yields end-to-end differentiable learning of per-layer activations, with analytical consideration of gradient propagation for discrete choices. In synthetic regression experiments, Flex-Act consistently recovers the correct activation and improves training stability, highlighting potential for robust, modular neural architectures that adapt their non-linear cores to task demands.

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.
Paper Structure (29 sections, 20 equations, 3 figures, 1 table)

This paper contains 29 sections, 20 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Motivation.
  3. Overview and Evaluation.
  4. Contributions.
  5. Related Work
  6. Traditional Fixed Activation Functions.
  7. Parameterized Activation Functions.
  8. General Parametric Activation Function.
  9. Background
  10. Neural Network Preliminaries.
  11. Activation Function Design.
  12. ReLU
  13. Sigmoid
  14. Tanh
  15. Leaky ReLU
  16. ...and 14 more sections

Figures (3)

  • Figure 1: Gumbel-Softmax activation selection probabilities over training epochs for various ground-truth activation functions. The model consistently converges to the appropriate non-linearity across tasks.
  • Figure 2:
  • Figure 3: Ablation analysis comparing our proposed method (Flex-Act) against fixed-function models across five ground truth transformations: ReLU, Sigmoid, Tanh, Leaky-ReLU, and Identity. The plots show predicted output values against input features. Flex-Act consistently approximates the true functional forms more faithfully, even in the absence of explicit architectural inductive bias.