Machine-learning surrogates in computational science often require not only accurate predictions but also expectations, variances, conserved integrals, and normalization constants. Conventional neural networks, including Kolmogorov--Arnold networks (KANs), generally evaluate these quantities through Monte Carlo sampling or multidimensional quadrature, making repeated inference costly, particularly in high dimensions. We introduce KANs with Tractable Operators for Uncertainty and Constraints (TOUCKANs), whose finite-degree polynomial outer functions yield a decomposable sum--product representation. Under factorizing input measures, expectations and higher moments then reduce algebraically to deterministic recursions over one-dimensional leaf integrals. Independent and shared inner-block variants form a parameter-efficiency hierarchy, and a universal-approximation proof for its most restricted member extends to the complete class. The same deterministic operators allow conserved integrals and normalization to be imposed throughout training. Across uniform and Gaussian benchmarks through 50 dimensions, deterministic mean and variance errors remained below the median fixed-budget Monte Carlo errors in every case. TOUCKANs also conserved Cahn--Hilliard mass to machine precision and imposed normalization and orthogonality through deterministic moments, enabling joint computation of the lowest four Schrödinger eigenpairs in two, three, and five dimensions. TOUCKANs therefore connect KAN approximation with tractable probabilistic circuits, making statistical operators and integral constraints directly available within an expressive neural surrogate.
Figures & tables
Figure 1: TOUCKAN architecture ( 7 ), shown in the parameter-efficient diagonal form ( 11 ). Every input dimension xd is processed by the same shared network ψ:R→RQ and specialized by the per-dimension diagonal scaling λd , producing the bounded univariate branches Ψd,q=tanh(λd,qψq(xd)) . The first combination layer forms the normalized sums ξq=D1∑d=1DΨd,q∈(−1,1) , each outer block evaluates a Chebyshev expansion Φq(ξq)=∑p=0PCq,pTp(ξq) , and the final combination layer sums all branches to produce f .
Specialization and inner block Ψd(xd)
Params
rel. L2
E[f] rel. error
Var[f] rel. error
Independent tanh(ψd(xd)) [ 32 ]
41,796
(8.8±5.5)×10−4
(7.2±6.5)×10−4
(1.9±2.4)×10−4
Full-head tanh(Wdh(xd)+bd)
19,072
(9.5±2.9)×10−4
(8.8±3.0)×10−4
(2.3±1.9)×10−4
Diagonal tanh(diag(λd)ψ(xd)) [ 34 ]
3,533
(3.6±1.2)×10−3
(2.9±1.4)×10−3
(2.4±1.7)×10−4
Rank-one tanh(λdψ(xd)) [ 33 ]
2,733
(7.7±0.6)×10−2
(1.7±1.1)×10−2
(4.0±1.9)×10−3
Table 1: Parameter efficiency of the four TOUCKAN inner-block specializations on the Lüthen benchmark at D=20 . The table reports trainable parameters, relative L2 test error, and relative errors of the analytically propagated mean and variance against exact references. All variants share the same outer block and moment recursion. Error columns report the mean ± standard deviation over five random seeds, and bold marks the best mean in each column.
Surrogate
E[f] rel. error
Var[f] rel. error
Benchmark
D
rel. L2
TOUCKAN
MC (median)
TOUCKAN
MC (median)
Lüthen
5
1.0×10−6
5.64×10−10
1.34×10−3
1.13×10−8
2.10×10−3
12
2.7×10−6
1.30×10−9
3.07×10−4
9.27×10−9
2.09×10−3
25
4.5×10−6
1.77×10−9
2.53×10−4
3.81×10−8
2.35×10−3
50
8.9×10−6
1.63×10−10
1.14×10−4
1.66×10−8
1.92×10−3
Bratley
5
2.9×10−4
7.22×10−7
1.03×10−3
1.80×10−6
1.74×10−3
Table 2: Relative errors of the expectation and variance for the three benchmarks across D∈{5,12,25,50} . The surrogate column reports the test relative L2 error. TOUCKAN columns report one deterministic propagation through the trained surrogate. MC (median) is the median error over 200 independent runs of the true integrand, each with N=2×105 samples ( I.2 ). TOUCKAN is more accurate than the median Monte Carlo estimate in every case. For D≤25 , it outperforms all 200 Monte Carlo mean replications and all but at most one variance replication. At D=50 , it outperforms at least 98% of the mean and 89% of the variance replications. The variance entries marked with ‡ are limited by floating-point roundoff in the second-moment recursion rather than by the surrogate.
Figure 2: Analytical TOUCKAN moments and Monte Carlo convergence at D=50 . Rows correspond to the Lüthen, Bratley, and Oakley–O’Hagan benchmarks, while columns show the mean and raw second moment. Blue dots show 20 independent Monte Carlo estimates of the true integrand at each sample size up to N=2×106 , and the blue shaded region marks their ±1σ envelope. The red line is the exact value, and the black dashed line is the deterministic TOUCKAN estimate. The Monte Carlo envelope contracts at the expected N−1/2 rate but remains visibly dispersed at the largest budget. The TOUCKAN estimate overlaps the exact value in all six panels without additional sampling during moment evaluation.
Figure 3: Cost and accuracy of TOUCKAN propagation and Monte Carlo across D∈{5,12,25,50} . Rows denote benchmarks and columns report wall time, mean error, and raw-second-moment error. MC (1 run) uses the time of one N=2×105 run and the median error across 200 independent runs. MC (avg. of 200) reports the total time and error obtained by averaging those runs. TOUCKAN reports one warm evaluation of the compiled deterministic recursions, excluding one-time JIT compilation. The reported timings do not include surrogate training. One TOUCKAN evaluation is up to 2,600 times faster than the 200-run Monte Carlo average. For the mean and raw second moment, TOUCKAN matches or improves the averaged Monte Carlo estimator’s accuracy except for Bratley at D=50 . Timings were measured on one NVIDIA L40S GPU.
Figure 4: Cahn–Hilliard equation with the exact-mass reprojection. (a)–(c): reference solution, learned surrogate, and absolute approximation error on a logarithmic color scale for the independent TOUCKAN of Table 3 . (d), (e): training loss and relative L2 error for the five parameter-matched models of Table 3 . The dashed line marks activation of the exact-mass reprojection for the constrained models. The MLP remains unconstrained throughout. (f): 512-node quadrature mass drift ∣Q[f](t)−Q0∣ over time. Every model with the exact-mass reprojection holds this quadrature mass at machine precision, while the unconstrained MLP baseline drifts at the level of its field error.
Model
Params
rel. L2
L∞
maxt∣Q[f](t)−Q0∣
MLP (unconstrained)
14,928
2.4×10−4
5.2×10−4
2.3×10−4
TNN (constrained) [ 26 ]
14,753
4.1×10−3
6.7×10−2
2.9×10−15
TOUCKAN (independent)
14,840
1.2×10−4
7.3×10−4
5.2×10−15
TOUCKAN (full-head)
15,328
3.6×10−4
2.2×10−3
5.6×10−14
TOUCKAN (diagonal)
14,912
1.4×10−3
4.0×10−3
1.0×10−14
Table 3: Field accuracy and mass conservation for the Cahn–Hilliard benchmark. Relative L2 and L∞ errors are measured against the Fourier-spectral reference, while the mass column reports the worst-case drift maxt∣Q[f](t)−Q0∣ under the 512-node quadrature used by the reprojection. All models have approximately 15,000 trainable parameters and receive the same total second-order optimization budget. The constrained models conserve the quadrature mass to machine precision, and for the TOUCKANs independent rules with 1024 – 8192 nodes confirm the same precision for the continuum integral ( F ). For comparison on the same problem, the specialized backward-compatible PINN (bc-PINN) method [ 48 ] reports a relative error of 1.86×10−2 . Bold marks the lowest field errors; all constrained mass drifts are at machine precision.
Figure 5: Schrödinger eigenvalue problem in D=2 , 3 , and 5 (columns). Top row: maximum error in the Gram entries ⟨ψa∣ψb⟩ versus number of function evaluations, against converged tensor-product references ( 2002 , 1203 , and 485 nodes): four independent Monte Carlo runs (red, up to 108 samples), tensor-product Gauss–Hermite quadrature with nD nodes (blue), and the exact cross-moment recursion (star), which reaches 1.5×10−14 , 1.2×10−11 , and 7.5×10−8 with 256 , 384 , and 640 univariate evaluations, respectively (the D=5 value is limited by the convergence of the 485 -node reference itself). The horizontal axis counts integrand evaluations: the TOUCKAN-based recursion additionally performs O(Q2DP4) cross-moment arithmetic outside this count, and its full evaluation, including that arithmetic, completes in 1 – 3 ms per case. Middle row: convergence of the four trainable eigenvalues Ea to the exact spectrum (dotted lines) from a common initialization below the spectrum. Bottom row: training convergence, showing the relative L2 error of the eigenfunctions, the relative eigenvalue error averaged over the four states, and the off-diagonal Gram residual max∣G−I∣ . The dashed line marks the switch from first- to second-order optimization (S2). Together, these results show that deterministic Gram evaluation avoids both the Monte Carlo sampling floor and tensor-product growth while all four eigenpairs converge simultaneously.
D=2
D=3
D=5
Trainable parameters
10,963
13,075
17,299
Relative error E0
2.2×10−8
4.0×10−7
5.6×10−5
Relative error E1
2.3×10−8
1.3×10−8
2.6×10−5
Relative error E2
7.7×10−9
5.6×10−8
2.8×10−5
Relative error E3
1.4×10−8
8.1×10−8
1.3×10−4
Eigenfunction rel. L2 error
2.2×10−4
2.5×10−4
1.6×10−2
Table 4: Accuracy of the lowest four eigenpairs of the coupled oscillator in D=2 , 3 , and 5 . Eigenvalue errors are reported for each state, and the eigenfunction row reports the worst case over the four states. Normalization is evaluated using independent tensor-product Gauss–Hermite rules with 2002 , 1203 , and 485 nodes and reports the largest residual over the four states. Orthogonality is the maximum absolute off-diagonal entry of the normalized Gram matrix at the final training iteration, evaluated by the deterministic recursion at M=128 . The D=5 normalization value is limited by the reference quadrature rather than the cross-moment recursion. Training times were measured on one NVIDIA H200 GPU.
Figure 6: Lowest four eigenstates of the coupled three-dimensional oscillator on the x3=0 plane: closed-form reference (left), learned surrogate (center), and absolute error on a logarithmic color scale (right). Each row is labeled by its learned eigenvalue Ea . On this slice, the E2 and E3 states have similar projected profiles, while the Gram residuals in Table 4 quantify their full-dimensional distinction.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Sensitivity of the computed moments to the one-dimensional quadrature order M . (a) Lüthen function at D=50 : relative error of the recursion mean and variance against the exact closed-form references. (b) Schrödinger Gram matrix at D=5 : relative error against an independent 485 tensor-product Gauss–Hermite integration of the same surrogate. Dashed vertical lines mark the production orders ( M=200 and M=128 ) used for the examples in the main paper.
Figure 8: Continuum mass conservation of the Cahn–Hilliard surrogate under quadrature refinement, independent of the 512 -node rule used for enforcement. (a) Worst-case mass drift maxt∣∫fdx−Q0∣ of the constrained TOUCKAN as the evaluation grid is refined from N=8 to N=8192 midpoint nodes. The decay is spectral but non-monotone, the signature of periodic-quadrature aliasing: the error at N nodes is the sum of the integrand’s Fourier modes with k=N,2N,… . It reaches a converged plateau near 5×10−15 by N≈128 , approximately one order above machine epsilon (dotted). The production rule N=512 (dashed) lies far inside the converged regime. (b) Mass drift ∣∫fdx−Q0∣ over time t , evaluated on four independent midpoint rules ( N=64,128,512,8192 ). The under-resolved N=64 rule stays above the floor, while N=128 , the production rule N=512 , and N=8192 coincide at the double-precision floor at every time instance, so conservation holds pointwise in t and independently of the evaluation rule.
Surrogate
E[f] rel. error
Var[f] rel. error
Benchmark
D
rel. L2
TOUCKAN
MC (median)
RQMC (median)
TOUCKAN
MC (median)
RQMC (median)
Lüthen
5
1.0×10−6
5.64×10−10
1.34×10−3
2.51×10−9
1.13×10−8
2.10×10−3
3.81×10−6
12
2.7×10−6
1.30×10−9
3.07×10−4
8.46×10−10
9.27×10−9
2.09×10−3
3.81×10−6
25
4.5×10−6
1.77×10−9
2.53×10−4
8.28×10−10
3.81×10−8
2.35×10−3
3.81×10−6
50
8.9×10−6
1.63×10−10
1.14×10−4
5.46×10−10
1.66×10−8
1.92×10−3
3.97×10−6
Bratley
5
2.9×10−4
7.22×10−7
1.03×10−3
9.65×10−10
1.80×10−6
1.74×10−3
3.80×10−6
Appendix
Table 5: Relative errors against the exact mean and variance for three benchmarks across D∈{5,12,25,50} . The surrogate column reports the test relative L2 error. TOUCKAN evaluates moments of the trained surrogate by deterministic propagation. MC and RQMC instead evaluate the true integrand and report the median error over 200 independent replications. Each MC run uses N=2×105 pseudo-random samples, while each RQMC run uses a complete Owen-scrambled Sobol’ net with N=218=262,144 points. Thus, the TOUCKAN errors include the discrepancy between the surrogate and the target, whereas the sampling errors arise from direct integration of the target. TOUCKAN has lower errors than the median MC estimate in every case and than the median RQMC estimate in most cases. The variance entries marked with ‡ are limited by floating-point roundoff in the second-moment recursion rather than by the surrogate.
E[f] rel. error
Var[f] rel. error
Benchmark
D
N=2×105
N=218
N=2×105
N=218
Lüthen
5
5.31×10−7
2.51×10−9
6.44×10−6
3.81×10−6
12
1.74×10−7
8.46×10−10
2.56×10−5
3.81×10−6
25
7.79×10−8
8.28×10−10
3.30×10−5
3.81×10−6
50
4.62×10−8
5.46×10−10
5.26×10−5
3.97×10−6
Bratley
5
2.22×10−6
9.65×10−10
1.33×10−5
3.80×10−6
Appendix
Table 6: Median relative errors of RQMC over 200 scramblings at the matched budget N=2×105 and at N=218 , for the true integrands of Table 5 .
D
5
12
20
25
50
E[f]
−2.80818
−13.66773
−22.02031
−30.58880
−70.99621
Var[f]
6.39955
7.55664
34.18982
26.33687
20.07893
Appendix
Table 7: Exact reference moments of the Lüthen function ( 79 ) for every reported input dimension, in physical units.
D
λmin(G)
λmax(G)
κ(G)
off-diag. residual
2
3.97×10−4
8.96×10−2
2.26×102
2.73×10−9
3
4.58×10−5
1.42×10−3
3.11×101
4.88×10−8
5
6.99×10−9
2.23×10−7
3.20×101
5.22×10−6
Appendix
Table 8: Conditioning of the Gram matrix G at the trained parameters for the Schrödinger eigenstate runs ( Ns=4 states, evaluated by the exact recursion at the production order M=128 ). The condition number κ(G)=λmax/λmin stays at order 101 – 102 with λmin(G)>0 at the trained parameters, so the direct solve solve(G,K) is stable without regularization. Successful completion of every run indicates no observed singular solve, but these final-parameter values do not quantify the conditioning during optimization. The last column reports the off-diagonal Gram residual at the final training iteration (Figure 5 ).
Figure 9: Lowest four eigenstates of the coupled two-dimensional oscillator: closed-form reference (left), learned surrogate (center), and absolute error on a logarithmic color scale (right). Each row is labeled by its learned eigenvalue Ea .
Figure 10: Lowest four eigenstates of the coupled five-dimensional oscillator on the x3=x4=x5=0 plane: closed-form reference (left), learned surrogate (center), and absolute error on a logarithmic color scale (right). Each row is labeled by its learned eigenvalue Ea . On this slice, the E1 and E2 states have similar projected profiles, while the Gram residuals in Table 4 quantify their full-dimensional distinction.