KAN 2.0: Kolmogorov-Arnold Networks Meet Science

In plain words

Most neural networks learn weights on their connections and apply one fixed bend (a nonlinearity) at each neuron. Kolmogorov-Arnold Networks (KANs) move the learning onto the connections instead: each connection carries its own small function of one number. Because each function is one-dimensional, you can plot it and, when it looks like a familiar shape, replace it with a simple formula. That makes a KAN easier to read than a typical network.

This paper argues that a scientist needs more from a model than readability. They want to put what is already known into the network, and to pull out which inputs matter, how the parts group together, and which formula fits. The authors extend the existing pykan library with multiplication nodes, a compiler that turns a formula into an equivalent network, and tools for each of these jobs.

They test the tools on physics problems: conserved quantities, Lagrangians (a compact way to write how a system moves), a hidden symmetry of the Schwarzschild black hole, and constitutive laws (how a material's stress responds to deformation). The results are demonstrations on small problems, and the authors report them themselves.

TL;DR

KAN 2.0 is a toolbox built on Kolmogorov-Arnold Networks for scientific use. Its main additions are multiplication nodes (MultKAN), a compiler from symbolic formulas to KANs, and tools that rank inputs, detect modular and symmetric structure, and extract formulas from trained networks. Use it if you want to fit a small KAN to physics-style data and read a formula out of it. The results are author-reported and mostly qualitative. Most training settings are not in the paper, so read the official code's defaults before reproducing anything.

Problem

• Interpretability in KANs was tied mostly to extracting symbolic formulas. The authors argue that formulas are not always necessary or feasible: in chemistry and biology the systems are often too complex, and modular structure and key features may be enough §1.
• Embedding prior scientific knowledge into a KAN is, in the authors' words, an overlooked task §1.
• A product of two variables is not a primitive in a KAN. It has to be built from additions and univariate functions, for example through exp and log for positive numbers §2.
• The earlier KAN importance score, the L1 norm of each edge function, looks only at local information. The paper shows cases where it ranks inputs wrongly §4.1 fig. 6.
• Recovering a symbolic formula from a dense KAN is hard without help, so the paper gives three tricks §4.3 fig. 9.
• Assumptions: the examples are small-scale and physics-related. The modularity tests use second-order derivatives and need smooth functions §4.2.2. The authors say KANs are more interpretable than MLPs at small scale, and that interpretability falls as networks grow §7.

Method

KANs place a learnable one-variable function on each edge; a node sums what arrives §2 eq. 2. The paper adds multiplication nodes, three ways to inject known structure §3, three ways to extract structure §4, and applies them to physics tasks §5.

1. MultKAN (multiplication nodes). Each layer has addition nodes and multiplication nodes. Every KAN layer first produces subnodes. Addition nodes copy one subnode. Multiplication nodes multiply k subnodes; the paper sets k=2 for simplicity and notes that any k ≥ 2 works §2 eq. 4 eq. 5. The multiplication layer has no trainable parameters, so sparsity penalties still apply §2.

# MultKAN layer  [§2] [eq. 4] [eq. 5]
# Input:  nodes, addition part then multiplication part
# Output: next-layer nodes, same two parts

def multkan_layer(x, n_next_a, n_next_m, arity=2):
    # Learnable edge functions produce the subnodes.  [§2]
    subnodes = kan_layer(x)               # width n_next_a + arity * n_next_m
    # Addition nodes copy their subnode unchanged.  [eq. 4]
    additions = subnodes[:n_next_a]
    # Multiplication node m takes the product of its arity subnodes.  [eq. 4]
    products = []
    for m in range(n_next_m):
        start = n_next_a + m * arity
        products.append(prod(subnodes[start : start + arity]))
    # Additions first, products after.  [eq. 4] [eq. 5]
    return concat(additions, products)

2. Inject known structure §3.

   • Auxiliary inputs (augment_input): add a computed feature, such as β = v/c or γ = 1/√(1−β²), as an extra input. The paper shows a relativistic-mass example where this cuts the edge count §3.1 fig. 3 eq. 6.
   • Modules (module(start_layer, chain)): keep connections inside user-specified groups and remove connections between groups §3.2 eq. 7.
   • Formula to KAN (kanpiler(input_variables, output_expression)): parse the formula into a tree, add identity dummy edges and subnodes so the tree fits KAN layers, and return a model in symbolic mode §3.3 eq. 8 fig. 5. Expansion (expand_width, expand_depth) adds zero-valued functions, and perturb makes them trainable §3.3.

3. Extract structure §4.

   • Feature importance. Node and edge attribution scores are computed backwards from the outputs. Each edge score is the subnode's score times the edge's output spread divided by the subnode's spread §4.1 eq. 9. prune_input keeps the inputs that score highest §4.1.
   • Anatomical modularity (auto_swap): assign neurons positions and swap them to shorten connections, so clusters become visible §4.2.1 fig. 7.
   • Functional modularity (test_separability, test_general_separability, test_symmetry, plot_tree): test additive or multiplicative separability eq. 10 eq. 11, generalized separability eq. 12 eq. 13, and generalized symmetry eq. 14 eq. 15. Apply the tests recursively to build a tree of groups §4.2.2 fig. 8.

# Attribution scores for a trained KAN  [§4.1] [eq. 9]
# E[layer][i][j]: std of the edge output, subnode i from input node j
# N[layer][i]:    std of subnode i
# The paper's E, N are standard deviations over the data  [§4.1]

def feature_attribution(E, N, depth, n_outputs):
    # Every output unit starts with score 1.  [§4.1]
    scores = identity(n_outputs)            # (n_outputs, n_outputs)
    for layer in reversed(range(depth)):
        # A multiplication node passes its score to its subnodes.  [§4.1]
        sub_scores = copy_to_subnodes(scores, layer)
        # Edge score: subnode score x edge spread / subnode spread  [eq. 9]
        edges = sub_scores[:, :, None] * E[layer]
        edges = edges / (N[layer][:, None] + 1e-4)   # 1e-4: code guard
        # Node score: sum of the edge scores that feed it.  [eq. 9]
        scores = edges.sum(axis=1)
    # Per-output rows; feature_score averages over output units.  [code]
    return scores.mean(axis=0)

# Symmetry test: does f depend on z only through the direction of
# grad_y f?  [§4.2.2] [eq. 14] [eq. 15]
# y = first input group, z = the other inputs

def is_symmetric(f, y_idx, z_idx, points, threshold):
    # Gradient of f with respect to y.  [§4.2.2]
    grad_y = gradient(f, points, wrt=y_idx)
    # Unit direction of that gradient.  [eq. 15]
    direction = grad_y / norm(grad_y)
    # The direction must not change along z.  [eq. 15]
    dependence = jacobian(direction, points, wrt=z_idx)
    # threshold is not stated in the paper; see Hyperparameters.
    return max(abs(dependence)) < threshold

4. Symbolic extraction, three tricks §4.3 fig. 9.

   • A: detect and reuse modules. Train a large KAN, probe its tree, then build the detected module structure into a second KAN and symbolify it §4.3.
   • B: sparse initialization (sparse_init=True): start with a sparse connection pattern. The paper says symbolic formulas usually correspond to sparse KANs, so a sparse start suits them §4.3 fig. 5.
   • C: hypothesis testing with checkpoints. Each training or pruning step saves a version. rewind returns to an earlier version, so alternative hypotheses can be tried and compared §4.3.

5. Applications §5. Each uses the tools above on a physics task.

   • Conserved quantities §5.1 fig. 10: use the condition that a conserved H satisfies f·∇H = 0. Parameterize H by a KAN and minimize the mean squared dot product over samples drawn from [−1, 1]^d (loss written inline, not numbered).
   • Lagrangians §5.2 fig. 11: a Lagrangian network predicts acceleration from the Euler-Lagrange form eq. 18. The kinetic term is initialized from the compiled formula T = ½mq̇² to keep the mass matrix positive. Symbolic regression is then applied per edge.
   • Hidden symmetry of the Schwarzschild black hole §5.3 fig. 12: an MLP or KAN learns the time shift Δt(r) that removes the metric's time-radius coupling (Gullstrand-Painlevé transformation, given inline). The KAN is later refined to very low loss with a 1000-interval grid, using the MLP as a teacher §5.3.
   • Constitutive laws §5.4 fig. 13: learn the stress P as a function of deformation F for a Neo-Hookean material, compiling the linear law first and then growing, perturbing, pruning and symbolifying the KAN §5.4 eq. 21 eq. 22.

Key equations

• MultKAN multiplication layer §2 eq. 4: \mathbf{M}_l(\mathbf{z}_l) = \mathrm{concatenate}\big(\mathbf{z}_l[:n^a_{l+1}],\; \mathbf{z}_l[n^a_{l+1}::2] \odot \mathbf{z}_l[n^a_{l+1}+1::2]\big)
• Attribution recursion §4.1 eq. 9 (the paper writes E and N with mixed index sets; the code uses the edge's own std, see Pitfalls): B_{l-1,i,j} = A_{l,j}\,\frac{E_{l,j}}{N_{l+1,j}}, \qquad A_{l-1,i} = \sum_j B_{l-1,i,j}
• Generalized symmetry condition §4.2.2 eq. 15: \nabla_{\mathbf{z}}\Big(\widehat{\nabla_{\mathbf{y}} f}\Big) = 0
• Lagrangian acceleration, as printed but with the brace closed here §5.2 eq. 18 (see Pitfalls): \ddot{\mathbf{q}} = (\nabla_{\dot{\mathbf{q}}}\nabla^T_{\dot{\mathbf{q}}}\mathcal{L})^{-1}\big[\nabla_{\mathbf{q}}\mathcal{L} - (\nabla_{\mathbf{q}}\nabla^T_{\dot{\mathbf{q}}}\dot{\mathbf{q}})\big]
• Linear constitutive law, small deformation §5.4 eq. 21: \mathbf{P}_l = \mu(\mathbf{F}+\mathbf{F}^T-2\mathbf{I}) + \lambda(\mathrm{Tr}(\mathbf{F})-3)\mathbf{I}
• Neo-Hookean law §5.4 eq. 22: \mathbf{P} = \mu(\mathbf{F}\mathbf{F}^T-\mathbf{I}) + \lambda\log(\det\mathbf{F})\,\mathbf{I}

Hyperparameters

The paper states almost no training settings. Rows marked "code default" come from the official pykan code at code_commit, not from the paper. Check them before using them as the paper's values.

Results

Author-reported. Setups are small toy or physics problems, each compared with the paper's own baseline in that section. The timing claim does not state hardware §7.

Paper vs code

Official repository: KindXiaoming/pykan. The paper gives this link in its introduction §1. The LICENSE file names Ziming Liu, the first author LICENSE. Commit ecde4ec3274d3bef1ad737479cf126aed38ab530. License: MIT. Read-only review. Nothing was executed, and I did not re-run the physics notebooks or check the paper's numbers.

• Attribution aggregation. The paper sets every output unit to score 1 and writes the scores with standard deviations §4.1 eq. 9. The code keeps one row per output unit and feature_score averages those rows over outputs MultKAN.py MultKAN.py MultKAN.py. The paper does not say how outputs are combined. The code adds 1e-4 to the denominator MultKAN.py.
• Subnode and node affine terms. The code applies a scale and bias to subnodes and nodes MultKAN.py. The paper has no such terms. By default they are fixed at 1 and 0 (affine_trainable=False) MultKAN.py MultKAN.py. So they have no effect unless enabled.
• Speed mode. The paper says model.speed() is one line that turns on the efficiency changes §7. In code, speed() sets symbolic_enabled=False, and the forward pass then drops every symbolic function MultKAN.py MultKAN.py. A model with active symbolic functions gives different outputs after speed().
• Fitting defaults. The paper gives no training settings. The code uses MSE with LBFGS and no regularization by default (lamb=0). With lamb=0 it stops saving activations, and during fit it turns off the symbolic branch when no symbolic function is active MultKAN.py MultKAN.py MultKAN.py.
• Symmetry test. The paper's condition is that the derivative of the unit gradient direction with respect to z vanishes §4.2.2 eq. 15. The code normalizes with a 1e-6 guard, multiplies by the input standard deviation, takes the median absolute value over points, and compares the maximum against 1e-3 hypothesis.py hypothesis.py. The paper states none of these choices.
• Multiplication layer. Consistent with the paper's stride-2 product eq. 4. The code uses consecutive subnode groups and supports other arities and mixed arities MultKAN.py MultKAN.py.
• Sparse initialization. The paper says the sparse pattern suits symbolic formulas §4.3. The code places inputs and outputs on evenly spaced positions and connects each to its nearest counterpart utils.py KANLayer.py. The paper does not define the rule.
• Expansion. The paper says the expansion methods add zero functions §3.3. The code's expand_depth sets the mask to zero, as described MultKAN.py. perturb's magnitude default is 1.0, and the paper gives no value MultKAN.py.
• expand_width signature. The Physics 4A tutorial calls expand_width(0,1,perturb=True) Physics_4A_constitutive_laws_P11.ipynb. The method at this commit takes no perturb argument MultKAN.py. The tutorial and the method disagree.
• Tree thresholds. The paper names no thresholds §4.2.2. The code's separability threshold differs between plot_tree/MultKAN.tree (1e-1) and get_tree_node (1e-2) hypothesis.py hypothesis.py MultKAN.py.
• Version. The paper says the release is pykan 0.2.x §1. This commit's setup.py says 0.2.8 setup.py.
• Tutorial values. The conserved-quantity notebook prunes with edge_th=5e-2 Physics_2B_conservation_law_2D.ipynb. The paper states no pruning threshold. Use notebook values only as examples.

Pitfalls

• Attribution scores. The paper's formula writes E with two indices where it defines an edge std with three §4.1 eq. 9. The code's edge score is the edge's own output std over the subnode's std, plus a 1e-4 guard MultKAN.py. Implement the code's version, and check the sign and scale on a toy case first.
• Multi-output models. Feature scores average over outputs MultKAN.py. For per-output importance, use the per-output rows from attribute MultKAN.py.
• speed() drops symbolic functions (see Paper vs code). Call it only for models with no active symbolic functions, or after checking the outputs.
• Fit defaults hide regularization. With the default lamb=0, neither the L1 nor the entropy penalty is applied MultKAN.py MultKAN.py. The paper says sparse regularization "can be directly applied" §2, but gives no strength. The LBFGS history and strong-Wolfe line search are fixed in code MultKAN.py.
• Symbolic extraction is noise-sensitive. The paper reports that symbolic regression is sensitive to noise and to compositional forms it cannot search for. Its fix was to fix the outer function by hand (fix_symbolic) §5.2.
• Seed dependence. Different seeds give different modules and different solutions, including a sub-optimal one in the auxiliary-input example §3.1 fig. 3. Hidden symmetry: roughly 1 in 3 seeds found a globally smooth solution; the rest did not §5.3.
• Domain walls. A spline KAN on a 1→1 problem learned a solution with a sign flip and a loss spike at the singular point §5.3 fig. 12. Regularization or seed choice removes it, per the authors.
• Prior knowledge can hurt. In the P12 example, the run with the compiled linear law ended at loss 7e-3 and the run without prior reached 6e-9. The authors suspect the linear law sits near a poor local minimum, and warn against generalizing from one case §5.4.
• Second-order derivatives. Modularity tests rely on second derivatives and need smooth functions. The paper gives a family of functions where small changes make the second derivatives blow up §4.2.2.
• Tree and symmetry results depend on thresholds. The paper says results can depend on seed and threshold §4.2.2. Thresholds are not stated, and the code's defaults differ by place (see Hyperparameters).
• Lagrangian stability. The mass matrix must stay positive, so the kinetic term is initialized from a compiled formula §5.2. Random initialization was reported as much less stable, with no numbers given §5.2.
• rewind semantics. The paper says rewind('0.1') renames version 0.1 to 1.1 §3.3. From reading the code, rewind saves the calling object's current state under <round>.<state> and then loads the requested checkpoint MultKAN.py MultKAN.py. Which state gets renamed therefore depends on which object you call it on. I did not run this.
• Scale and efficiency. KANs still trail MLPs in efficiency at large scale, and interpretability falls as models grow §7. The memory claim (from O(LN²G) to O(LNG)) is not checked here.
• Printed equation. The Lagrangian acceleration eq. 18 is typeset with an unbalanced brace in its last term. Check it against the Euler-Lagrange equation before implementing.
• Unverified claims. The claim that the compiler handles all 120 Feynman equations, and the paper's cited MSE-versus-grid scaling §3.3 [liu2024kan], are not verified here.

Lineage

• Builds on: KAN (Liu et al., 2024). This paper keeps its KAN layers and spline edges and adds multiplication nodes, a new importance score, and the tools above §2 §4.1.
• Compared against: MLPs (hidden symmetry, functional tree, anatomical modularity) §4.2 §5.3; the L1-norm edge score §4.1; the MLP-based Lagrangian neural network, where this paper uses a KAN instead of the MLP §5.2. No card in this wiki covers the Lagrangian neural network.
• Superseded by: none known from wiki content.