Activations

The activations of torch.nn.functional, as expression functions over a tile or a scalar. An activation applied where a result is stored fuses into that store: no tile holds the pre-activation, and the activation costs arithmetic only.

@autotune(R in [16])
kernel bias_silu(X: tensor<f32>[M, 256], Bias: tensor<f32>[1, 256], Y: tensor<f32>[M, 256]) {
  let r = ctx.block.x * R
  Y[r :+ R, :] = silu(X[r :+ R, :] + Bias[0 :+ 1, :])
}

The bias row broadcasts over the R rows, and the sum, the activation and the store are one sweep. silu reads its argument twice in its definition, x / (1.0 + exp(-x)); the compiler stages the slice once all the same.

Rectifiers: relu(x), leaky_relu(x, 0.1), elu(x, 1), selu(x), plotted from the library's source.

• per element, in registers

RectifiersGELU and its kinSquashingSmoothShrinkage

Computed from the definitions in phobos-lang/src/stdlib/activation.ph, in double precision; on the device exp and tanh are approximate.

Rectifiers

relu

fn relu(x)expression

max(x, 0).

Source, std/activation.ph:7

fn relu(x) = max(x, 0.0)

relu6

fn relu6(x)expression

relu, capped at 6.

Source, std/activation.ph:10

fn relu6(x) = clamp(x, 0.0, 6.0)

leaky_relu

fn leaky_relu(x, slope)expression

x for a positive x, slope * x otherwise. slope in [0, 1].

Source, std/activation.ph:13

fn leaky_relu(x, slope) = max(x, x * slope)

hardtanh

fn hardtanh(x)expression

x limited to [-1, 1].

Source, std/activation.ph:16

fn hardtanh(x) = clamp(x, -1.0, 1.0)

elu

fn elu(x, alpha)expression

x for a positive x, alpha * (e^x - 1) otherwise. alpha > 0.

Source, std/activation.ph:29

fn elu(x, alpha) = max(x, 0.0) + min(alpha * (exp(x) - 1.0), 0.0)

selu

fn selu(x)expression

elu scaled to be self-normalizing, with the constants of Klambauer et al.

Source, std/activation.ph:32

fn selu(x) = 1.0507009873554805 * elu(x, 1.6732632423543772)

celu

fn celu(x, alpha)expression

x for a positive x, alpha * (e^(x / alpha) - 1) otherwise. alpha > 0.

Source, std/activation.ph:35

fn celu(x, alpha) = max(x, 0.0) + min(alpha * (exp(x / alpha) - 1.0), 0.0)

squared_relu

fn squared_relu(x)expression

relu(x)^2, Primer's.

Source, std/activation.ph:76

fn squared_relu(x) = square(max(x, 0.0))

Thresholds and shrinkage

An activation shaped like a condition is a where, which is a value like any other inside an expression function, so it fuses the same way.

threshold

fn threshold(x, th, value)expression

x where it is above th, value elsewhere.

Source, std/activation.ph:19

fn threshold(x, th, value) = where(x > th, x, value)

hardshrink, softshrink, tanhshrink

fn hardshrink(x, l)fn softshrink(x, l)fn tanhshrink(x)expression

Shrinkage, as sparse coding uses it: hardshrink zeroes [-l, l] and keeps x outside it, softshrink zeroes it and moves the rest l toward 0, and tanhshrink is x - tanh(x). l >= 0.

Source, std/activation.ph:24

fn hardshrink(x, l) = where(abs(x) > l, x, 0.0)
fn softshrink(x, l) = x - clamp(x, -l, l)
fn tanhshrink(x) = x - tanh(x)

Sigmoid and its relatives

sigmoid

fn sigmoid(x)expression

The logistic function, 1 / (1 + e^-x).

Source, std/activation.ph:38

fn sigmoid(x) = 1.0 / (1.0 + exp(-x))

log_sigmoid

fn log_sigmoid(x)expression

log(sigmoid(x)), without the overflow of computing it so.

Source, std/activation.ph:41

fn log_sigmoid(x) = min(x, 0.0) - log(1.0 + exp(-abs(x)))

hardsigmoid

fn hardsigmoid(x)expression

A piecewise-linear sigmoid: 0 below -3, 1 above 3.

Source, std/activation.ph:44

fn hardsigmoid(x) = clamp(x / 6.0 + 0.5, 0.0, 1.0)

softsign

fn softsign(x)expression

x / (1 + |x|).

Source, std/activation.ph:73

fn softsign(x) = x / (1.0 + abs(x))

GELU and SiLU

gelu

fn gelu(x)expression

x * Phi(x), Phi the standard normal's CDF, through erf: PyTorch's gelu with approximate="none".

Source, std/activation.ph:57

fn gelu(x) = 0.5 * x * (1.0 + erf(x * 0.7071067811865476))

gelu_tanh

fn gelu_tanh(x)expression

gelu by its tanh approximation, GPT-2's: PyTorch's approximate="tanh".

Source, std/activation.ph:60

fn gelu_tanh(x) = 0.5 * x * (1.0 + tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)))

quick_gelu

fn quick_gelu(x)expression

x * sigmoid(1.702 * x), CLIP's cheaper approximation of gelu.

Source, std/activation.ph:63

fn quick_gelu(x) = x / (1.0 + exp(-1.702 * x))

silu

fn silu(x)expression

x * sigmoid(x), also called swish.

Source, std/activation.ph:47

fn silu(x) = x / (1.0 + exp(-x))

swish

fn swish(x, beta)expression

x * sigmoid(beta * x).

Source, std/activation.ph:50

fn swish(x, beta) = x / (1.0 + exp(-beta * x))

hardswish

fn hardswish(x)expression

x * hardsigmoid(x), MobileNetV3's.

Source, std/activation.ph:53

fn hardswish(x) = x * hardsigmoid(x)

mish

fn mish(x)expression

x * tanh(softplus(x)).

Source, std/activation.ph:70

fn mish(x) = x * tanh(softplus(x))

softplus

fn softplus(x)expression

log(1 + e^x), in the form that neither overflows for a large x nor loses a small one.

Source, std/activation.ph:67

fn softplus(x) = max(x, 0.0) + log(1.0 + exp(-abs(x)))

Gated linear units

A transformer’s MLP projects its input twice, a gate and an up branch, and multiplies the gate’s activation into the other. The product is one more operation in the same sweep:

@autotune(R in [8])
kernel swiglu_rows(G: tensor<f32>[M, 512], U: tensor<f32>[M, 512], H: tensor<f16>[M, 512]) {
  let r = ctx.block.x * R
  H[r :+ R, :] = swiglu(G[r :+ R, :], U[r :+ R, :])
}

swiglu, geglu, geglu_tanh, reglu

fn swiglu(gate, up)fn geglu(gate, up)fn geglu_tanh(gate, up)fn reglu(gate, up)expression

Gated linear units: the gate's activation times the linear branch up, as a transformer's MLP joins its two projections. The names follow Shazeer's "GLU Variants Improve Transformer".

Source, std/activation.ph:81

fn swiglu(gate, up) = silu(gate) * up
fn geglu(gate, up) = gelu(gate) * up
fn geglu_tanh(gate, up) = gelu_tanh(gate) * up
fn reglu(gate, up) = relu(gate) * up