Learnable activation functions

Overview

A standard multi-layer perceptron applies a fixed nonlinearity to each neuron after a learned linear map 2404.19756 §1. The learnable variant moves the trainable part onto the edges: each connection between two neurons carries its own univariate function, and each neuron only sums what arrives 2404.19756 §2.2. The invariant every variant shares is this split between learnable edge functions and parameter-free summing nodes.

The motivation is twofold. The authors want better accuracy per parameter than MLPs on smooth functions, and they want a network whose parts can be read and simplified 2404.19756 §1. Both claims rest on the same construction, described below.

Variants

Only one variant has a card in this concept so far.

Spline edges (Kolmogorov-Arnold networks)

Each edge applies a base term plus a B-spline term, w_b * silu(x) + w_s * spline(x), with cubic splines on a grid that can be refined during training 2404.19756 §2.2. Refining the grid is done by a least-squares fit from the coarse spline, without restarting training 2404.19756 §2.4. Pruning removes weak nodes and edges, and selected edges can be replaced with named functions such as sine or exponential 2404.19756 §2.5.1. Card: KAN: Kolmogorov-Arnold Networks.

Reported results: on fifteen special functions, the best KAN is lower in RMSE than the best MLP in every row of the table 2404.19756 tab. 1. On the Feynman equations, KANs and MLPs are comparable on average, and the MLP has the lowest loss in several rows 2404.19756 tab. 2. On a Poisson problem, a small KAN reaches a lower error with far fewer parameters than a larger MLP 2404.19756 §3.4. On knot signatures, a KAN beats an MLP baseline quoted from another group 2404.19756 tab. 3.

When to use which

With one card, the only comparison is KAN against MLP. The table lists what the paper supports, not general advice.

Implementation sketch

The forward pass has the same shape as an MLP, with the matrix product replaced by a sum over per-edge functions 2404.19756 §2.2. The edge function is the part that changes between variants.

# Forward pass: one learnable function per input-output pair  [arxiv:2404.19756 §2.2]
# Shapes [n_0, n_1, ..., n_L]; x has shape (batch, n_0)

def forward(x, layers):
    for layer in layers:
        # each edge applies its own univariate function  [arxiv:2404.19756 §2.2]
        edge_outputs = layer.edge_functions(x)   # (batch, n_in, n_out)
        # nodes only add the incoming edge outputs  [arxiv:2404.19756 §2.2]
        x = edge_outputs.sum(axis=1)             # (batch, n_out)
    return x

# Spline edge: base term plus trainable spline  [arxiv:2404.19756 §2.2]
# Cubic (k = 3) B-splines on a grid of G intervals

def edge(x, w_b, w_s, coefs):
    # residual-style base term, with silu as the default  [arxiv:2404.19756 §2.2]
    base = w_b * silu(x)
    # trainable combination of B-spline basis functions  [arxiv:2404.19756 §2.2]
    spline = sum(c * B(x) for c, B in zip(coefs, basis))
    return base + w_s * spline

# Grid extension: refit a finer spline to the coarse one  [arxiv:2404.19756 §2.4]

def extend_grid(coarse_coefs, coarse_basis, fine_basis, samples):
    # values of the coarse spline on sampled inputs  [arxiv:2404.19756 §2.4]
    targets = coarse_basis(samples) @ coarse_coefs
    # least-squares fit onto the finer grid  [arxiv:2404.19756 §2.4]
    return least_squares(fine_basis(samples), targets)

Where variants differ: the basis (splines here; other papers may use other families), the initialisation, and the sparsity penalty. The card's Paper vs code section lists where the released code departs from the paper's stated choices.

Contradictions

None known. Only one card covers this concept. The authors' own numbers are not checked against an independent reproduction here. The paper also does not state the optimiser or grid for its PDE experiment 2404.19756 §3.4.

Papers

• KAN: Kolmogorov-Arnold Networks: replaces MLP weights and fixed node activations with learnable spline edges, with pruning and symbolic fitting for interpretability.