Learnable Activation Functions

Learnable / adaptive activation functions.

Implements:

  • RationalActivation — the Rational Activation Function (RAF) from Transformers with Learnable Activation Functions (Fang et al., EACL 2023, arXiv:2208.14111). A RAF is a learnable ratio of two low-degree polynomials (a Padé approximant), with five “version” variants (A/B/C/D/N) controlling the denominator form. The default configuration matches the paper: degree (5, 4), version "A" (the “safe” per-term absolute-value denominator), initialized by a least-squares fit to GELU on [-3, 3].

    Note:

    Despite being requested as “Rectified Activation Function”, the paper defines RAF = Rational Activation Function. This implementation follows the paper.

  • SwishTrainable — Swish with a learnable beta parameter.

  • Maxout — the Goodfellow et al. (2013) maxout unit.

class src.model.activation_function.learnable.Maxout(*args: Any, **kwargs: Any)[source]

Bases: Module

Maxout unit: max_{i=0..k-1} (W_i x + b_i).

Computes k linear projections of the input and takes the elementwise maximum. A universal approximator of any continuous function given enough pieces. Increases parameter count by a factor of k.

Parameters:

  • dim – Input/output feature dimension.

  • num_pieces – Number of linear pieces k to maximize over. Default: 2.

proj

Linear layer producing dim * num_pieces outputs.

Reference: Goodfellow et al. (2013), arXiv:1302.4389; Lederer §2.3.4.

forward(x: torch.Tensor) → torch.Tensor[source]

Apply Maxout.

Parameters:

x – Input tensor of shape (..., dim).

Returns:

Tensor of shape (..., dim).

class src.model.activation_function.learnable.RationalActivation(*args: Any, **kwargs: Any)[source]

Bases: Module

Rational Activation Function (RAF): a learnable Padé P(x)/Q(x).

\[F(x) = \frac{P(x)}{Q(x)} = \frac{\sum_{j=0}^{m} a_j x^{j}}{1 + R(x)}\]

where the denominator term R(x) depends on version:

  • "A" (default, “safe”): R(x) = \sum_{k} |b_k x^{k+1}| (abs per term)

  • "B": R(x) = |\sum_{k} b_k x^{k+1}| (abs of the whole sum)

  • "C": R(x) = |b_0 + b_1 x + \dots + b_n x^{n}| with floor 0.1

  • "D": like "B" but with uniform multiplicative noise on the denominator weights during training only.

  • "N": R(x) = \sum_{k} b_k x^{k+1} (no abs; can have poles)

Versions A/B/D keep Q(x) >= 1 (no division by zero); C keeps it >= 0.1; N is unsafe. The paper uses version "A".

Parameters:

  • degrees – (m, n) numerator and denominator polynomial degrees. Default: (5, 4) (the paper default).

  • version – Denominator form. One of "A" (default), "B", "C", "D", "N".

  • approx_func – Initialization target. One of "gelu" (default), "relu", "leaky_relu", "leaky_relu_0.1", "sigmoid", "tanh", "swish", "silu", "identity".

  • trainable – If False, freeze the rational parameters. Default: True.

  • input_scaling – If True, apply per-token min-max scaling of the input to [-3, 3] before the rational (the RAFT preprocessing that keeps inputs inside the fitted range). Default: False.

  • noise_eps – Std of the version-D multiplicative noise. Default: 0.1.

numerator

Learnable parameter of shape (m + 1,).

denominator

Learnable parameter of shape (n,).

Reference:

Fang et al. (2023). Transformers with Learnable Activation Functions. arXiv:2208.14111.

forward(x: torch.Tensor) → torch.Tensor[source]

Apply the rational activation elementwise.

Parameters:

x – Input tensor of arbitrary shape.

Returns:

Tensor of same shape.

class src.model.activation_function.learnable.SwishTrainable(*args: Any, **kwargs: Any)[source]

Bases: Module

Swish with a learnable beta: x * sigmoid(beta * x).

When beta -> 0 the function approaches the linear/2 map; large beta recovers ReLU. Shared-scalar beta by default (per-channel via dim).

Parameters:

  • dim – If given, beta is per-channel of shape (dim,). If None, beta is a single shared scalar. Default: None.

  • beta_init – Initial value of beta. Default: 1.0 (== SiLU).

beta

Learnable Swish slope (scalar or (dim,)).

Reference: Ramachandran et al. (2017), arXiv:1710.05941; survey §6.

forward(x: torch.Tensor) → torch.Tensor[source]

Apply trainable Swish elementwise.

Parameters:

x – Input tensor of arbitrary shape.

Returns:

Tensor of same shape.