KAN layers as dense GEMMs in JAX: a small Neural QMC experiment (FermiKAN)

[Samu-Physicist]

Update (Sep 29, 2026): After a closer review, I found problems in the design described below. The 4D ε-manifold embedding was unnecessary (real solid harmonics have no singularity at the origin), and the electron-nucleus cusp is not exactly satisfied for molecular orbitals. I have revised the README and the Hashnode article accordingly and withdrawn the interpretability claims for now. I'm fixing the code and will post updated results in this thread.

Hi everyone! I’m an R&D researcher at a chemical manufacturer. I usually rely on commercial quantum chemistry solvers, but Neural QMC caught my eye—theoretically, it offers highly accurate quantum mechanics evaluation at $\mathcal{O}(N^4)$ computational cost!

I started experimenting with DeepMind’s FermiNet and successfully reached chemical accuracy for H2 and LiH systems. That got me thinking: what if we replaced the MLPs with KANs (Kolmogorov-Arnold Networks) acting as a learnable physical basis set on single-electron orbitals? This could potentially demonstrate a "Glass-Box" AI, highly interpretable and compatible with LLMs (Neuro-Symbolic approach).

Thanks to JAX's incredible expressiveness, implementing these complex basis function expansions was elegant. By replacing branching B-splines with orthogonal polynomials and utilizing jnp.einsum, I was able to compile the KAN architecture into GEMMs. I built this entirely as a weekend project, pairing with Gemini as my coding partner to push the limits of what I could do with JAX.

Here are the results from my H2 PoC:

• The Good: The Good: The model uses about 330x fewer parameters than the default FermiNet configuration for H2 (which is much larger than H2 requires).
• The Bad (The Wall): I fell into a dependency hell with the K-FAC optimizer and my KAN implementation, so I had to fall back to Adam.
• The Trade-offs: The accuracy is currently about ~0.01 Hartree shy of the original FermiNet. In an apples-to-apples comparison (using Adam), it took 4x the iterations to converge.

This is a personal, non-peer-reviewed project. I've written a detailed breakdown of the architecture, the math, and my struggles here:

👉 Link to GitHub Repo

👉 Link to Hashnode Article

I’d love to hear your thoughts! I’m especially hoping to collaborate with JAX/XLA wizards (I’d really appreciate some advice on overcoming the K-FAC integration issue) and domain experts!