Optimized Architectures for Kolmogorov-Arnold Networks

This paper proposes a principled, end-to-end differentiable framework that combines overprovisioned architectures with sparsification, deep supervision, and depth selection to learn compact, interpretable Kolmogorov-Arnold networks without sacrificing accuracy, thereby resolving the tension between model expressiveness and interpretability in scientific machine learning.

The Gist

Imagine you are trying to teach a robot to understand the laws of physics or predict the weather. You have two main choices for how to build its brain:

1. The "Black Box" Brain: You give it a massive, complex neural network. It's incredibly smart and accurate, but it's like a giant, tangled ball of yarn. You can't see how it figured out the answer, only that it did. Scientists hate this because they need to understand the "why," not just the "what."
2. The "Transparent" Brain (KANs): Recently, a new type of brain called a Kolmogorov–Arnold Network (KAN) was invented. Instead of just having fixed weights (like a standard calculator), KANs learn little, simple mathematical curves for every connection. This makes them transparent; you can look at the brain and say, "Ah, I see, it learned that $x$ squared plus sine of $y$ equals the answer." It's like looking at a clear glass engine instead of a black box.

The Problem:
The catch with KANs is that to make them smart enough to solve hard problems, you usually have to build them huge. You give them thousands of connections. While the math is transparent, a brain with 10,000 transparent connections is still too messy for a human to understand. It's like having a library where every book is written in plain English, but there are 10 million books. You still can't find the one story you need.

The Solution: "Grow Big, Then Trim"
The authors of this paper propose a clever strategy: Overprovision, then Sparsify.

Think of it like sculpting a statue out of a giant block of marble.

1. Overprovisioning (The Big Block): Instead of trying to carve the perfect statue immediately, they start with a massive, over-sized block of marble. They give the KAN way more connections and layers than it probably needs.
2. The Sculpting Tools (The New Architecture): They equip the KAN with three special tools to carve away the excess:
   • Edge Gates (The Chisel): These are tiny switches on every single connection. During training, the network learns to flip the switch to "OFF" for connections that aren't doing any work. It's like pruning a bonsai tree, cutting off dead branches so the tree grows a beautiful, compact shape.
   • Forward Connections (The Elevator): Imagine a building where every floor has an elevator that goes straight to the roof. This lets the network skip unnecessary middle layers if the answer is simple. It helps the network decide, "Do I need to go deep, or can I solve this right now?"
   • Exit Gates (The Early Exit): Imagine a hallway with doors on every floor. Usually, you have to walk all the way to the end. But these doors let the network say, "I have the answer on the 2nd floor; I don't need to go to the 10th." This allows the network to choose its own depth.

The "Smart Scale" (Minimum Description Length)
How does the network know how much to cut? They use a principle called Minimum Description Length (MDL).
Think of this as a strict budget for the network's "backpack."

• The backpack needs to carry the answer (Accuracy).
• But the backpack also has a weight limit (Complexity).
• The network is penalized if its backpack is too heavy (too many connections).
• The goal is to find the lightest backpack that still holds the answer perfectly.

What They Found:
They tested this on everything from simple math puzzles to predicting chaotic weather patterns and the strength of concrete.

• Just cutting branches (Sparsification) wasn't enough. If you just cut connections but don't let the network choose its depth, it often gets confused and loses accuracy.
• The Magic Combo: When they combined the chisels (cutting edges) with the elevators and exits (choosing depth), the results were amazing.
  • The networks became tiny (sometimes 90% smaller than the original).
  • They stayed super accurate (often even better than the big models).
  • They became easy to read. The final models were so simple that a human could actually look at the math and understand the logic.

The Takeaway:
This paper shows that we don't have to choose between "Smart but confusing" and "Simple but dumb." By starting with a massive, flexible brain and teaching it to prune itself down to the essentials, we can create AI that is both a genius and a clear, understandable teacher. It turns the "Black Box" into a "Glass House" that is small enough to walk through.

Technical Summary

1. Problem Statement

The Tension Between Accuracy and Interpretability:
Deep learning has revolutionized scientific modeling but often sacrifices interpretability for accuracy. Complex architectures (e.g., DenseNet blocks, skip connections) improve performance but make models "black boxes."
The Limitation of Current KANs:
Kolmogorov–Arnold Networks (KANs) offer a promising alternative by learning univariate activation functions on edges rather than fixed weights, theoretically offering better interpretability. However, standard KANs face the same trade-off: overprovisioning (making the network large) improves expressiveness but destroys interpretability.
The Gap:
Existing methods to prune KANs (e.g., post-hoc pruning) are discrete and separate from training, making them expensive and suboptimal. There is a lack of a differentiable, end-to-end framework that jointly optimizes the network's activations, structure (sparsity), and depth to find compact, accurate, and interpretable models without manual architecture search.

2. Methodology

The authors propose a framework that combines overprovisioned architectures with differentiable sparsification and depth selection, guided by a Minimum Description Length (MDL) objective.

A. Architectural Components

The proposed architecture integrates three key mechanisms into a standard KAN:

1. Differentiable Edge Gates (E):
   • Based on $\ell_0$ regularization (Louizos et al.), using a continuous relaxation (Gumbel-Softmax/stretched sigmoid) to learn binary gates for every activation function (edge).
   • Allows the network to learn which edges to keep (1) or prune (0) during training via gradient descent.
2. DenseNet-style Forward Connections (F):
   • Connects the input and all previous layer outputs to the current layer.
   • Provides deep supervision (gradients flow directly to early layers) and allows early features to bypass the "trunk," effectively enabling the network to skip unnecessary depth.
3. Learnable Exit Gates (X):
   • Implements a Multi-Exit architecture where every layer has an output head.
   • Uses a categorical Gumbel-Softmax relaxation to learn a probability distribution over which exit (layer depth) to use for the final prediction.
   • This provides explicit depth selection, allowing the model to dynamically choose the optimal network depth for a given input or task.

B. The Learning Objective (MDL)

The training objective is based on the Minimum Description Length (MDL) principle, which balances data fit against model complexity:
$$\mathcal{L}_{MDL} = \mathcal{L}_{data} + \mathcal{L}_{model}$$

• Data Loss ($\mathcal{L}_{data}$): Mean Squared Error (MSE).
• Model Complexity ($\mathcal{L}_{model}$): Approximated using a BIC-style penalty proportional to the number of open gates (edges and exits).
  • For exit gates, the complexity is the expected description length, weighted by the probability of using a specific exit path.
  • This formulation encourages the model to find the simplest architecture that sufficiently explains the data.

C. Training Procedure

• Warmup: The trunk trains first to stabilize spline activations.
• Annealing: The temperature parameters for the Gumbel-Softmax relaxations (for both edge gates and exit gates) are annealed from high (exploratory) to low (deterministic) during training.
• Inference: Gates are deterministically thresholded (e.g., $E[\tilde{z}] > 0.5$) to produce a fully discrete, interpretable architecture.

3. Key Contributions

1. Differentiable Joint Optimization: The paper introduces the first end-to-end differentiable method to jointly optimize KAN activations, sparsity (edge gates), and depth (exit gates) without discrete search.
2. Synergy of Mechanisms: It demonstrates that edge-level sparsification alone is insufficient (it reduces edge count but not compositional depth). The combination of sparsification with depth selection (via Forward Connections or Multi-Exits) is required to achieve both high accuracy and high parsimony.
3. MDL-Guided Architecture Search: By framing architecture search as a differentiable optimization problem under an MDL objective, the authors provide a principled way to trade off accuracy and complexity.
4. Comprehensive Benchmarking: The study evaluates a $2 \times 2 \times 2$ factorial design (presence/absence of E, F, X) across:
   • Symbolic regression benchmarks (Nguyen-10).
   • Dynamical systems (Ikeda map, 3-species ecosystem).
   • Real-world datasets (Concrete strength, Superconductor critical temperature).

4. Results

The experiments reveal distinct performance patterns across the different architectural conditions:

• Edge Gates (E) Alone: Often reduces the number of edges but harms accuracy. The model struggles to find the correct compositional depth, leading to underfitting or poor symbolic fidelity.
• Depth Selection (F or X) + Sparsification (E):
  • EX (Edge + Exit): Consistently achieved the highest Pareto hypervolume (best trade-off between accuracy and size) across most benchmarks. It discovered models that were significantly smaller (e.g., 16 edges vs. 48 for the Ikeda map) with equal or better accuracy.
  • EFX (Edge + Forward + Exit): The most expressive architecture. While not always the absolute best on every single metric, it provided a robust "safe default" that outperformed baselines and single-mechanism variants.
  • FX (Forward + Exit, no Edge): Often added too many parameters (overprovisioning) without pruning, leading to larger models.
• Specific Findings:
  • Symbolic Regression: Conditions with depth selection (EX, EFX) found models with fewer edges and lower RMSE than the baseline.
  • Dynamical Systems: The Ikeda map saw massive reductions in model size (66% reduction) with no accuracy loss. The ecosystem model was more sensitive to over-regularization but still benefited from depth selection.
  • Real-World Data: On the Concrete dataset, the EFX model at $\beta=0.01$ achieved a 0.8% improvement in RMSE while reducing model size by 82% (from 351 edges to 64).

5. Significance and Conclusion

• Resolving the Interpretability-Accuracy Trade-off: The paper proves that scientific machine learning models can be made both more expressive and more interpretable simultaneously. By learning the architecture during training, the resulting KANs are compact and reveal the underlying functional structure of the data.
• Principled Architecture Search: It shifts architecture search from a discrete, hyperparameter-heavy process to a continuous, gradient-based optimization, making it scalable and integrable with standard deep learning pipelines.
• Practical Recommendation: The authors suggest using the EFX architecture (Edge gates + Forward connections + Exit gates) as a default for scientific modeling tasks where the optimal architecture is unknown a priori. This combination provides the necessary "slack" to explore complex functions while the MDL objective and differentiable gates automatically prune the model to its essential, interpretable core.

In summary, this work provides a robust, differentiable framework for "growing" and "pruning" KANs simultaneously, offering a new pathway for creating scientific AI models that are not only accurate but also transparent and parsimonious.