Source-linked AI summary
Scalability Analysis of Distributed Kolmogorov-Arnold Network Training on High-Performance Computing Systems
Guangneng Chen, David Garcia Selfa, Pablo Quesada Barriuso
TL;DR
This paper addresses the limited characterization of KAN behavior under multi-GPU, multi-node data-parallel training. It empirically studies DDP KAN scaling across four dimensions, finding 74.7% parallel efficiency and 5.97× speedup at 8 GPUs alongside stable distributed weak scaling and topology-related communication variation.
Problem
KAN’s behavior under replicated multi-GPU, multi-node data-parallel training had not been systematically characterized, including whether its per-edge learnable activations alter communication or memory scaling.
Source: §1 · pp. 2–3Method
The paper empirically evaluates DDP-based KAN training across strong scaling, weak scaling, communication overhead, and model-size scaling.
Source: §1 · p. 3Results
At 8 GPUs, KAN training reaches 74.7% parallel efficiency and 5.97× speedup, with stable weak scaling across distributed configurations and non-monotonic 1.3%–6.1% communication overhead.
Source: §6 · p. 17Takeaways & Limitations
The results support training KAN with standard data-parallel DDP without pathologies beyond those expected for conventional architectures, within the study’s stated scope.
Source: §6 · p. 17Mechanistic explanations were not directly confirmed through ablation, and matched-architecture controls were unavailable to distinguish general DDP behavior from KAN-specific effects.
Source: §5.5 · p. 16
Abstract
from arXiv · showhide
Kolmogorov-Arnold Networks (KANs) replace the fixed activation functions and linear weights of Multi-Layer Perceptrons (MLPs) with learnable univariate functions on network edges, offering improved interpretability and, in some settings, competitive parameter efficiency. While the approximation properties of KANs have received considerable attention, their behavior under distributed, multi-GPU training has not been systematically characterized. This paper presents an empirical scalability study of data-parallel KAN training on multi-node, multi-GPU high-performance computing (HPC) infrastructure, evaluated along four dimensions: strong scaling, weak scaling, communication overhead, and model-size scaling. Experiments were conducted on the FinisTerrae III supercomputer using up to 8 NVIDIA A100 GPUs across 4 nodes with PyTorch Distributed Data Parallel (DDP). KAN training reaches 74.7% parallel efficiency at 8 GPUs with a 5.97x speedup, consistent with conventional deep learning workloads. Weak scaling shows an initial single-to-multi-GPU throughput drop followed by strong stability. Communication overhead follows a non-monotonic pattern (1.3%-6.1%), driven primarily by All-Reduce algorithm selection and inter-node latency rather than KAN's edge-wise gradient structure. The parameter-to-memory ratio improves with model size even as training time scales unfavorably. These results indicate that operator-level and data-parallel optimizations for KAN are complementary. We provide deployment guidelines for GPU topology and model-size selection, and discuss the limitations of a synthetic-regression evaluation.
1. Introduction
KANs move learnable univariate functions from MLP-style node activations to network edges, motivating questions about their computational cost and distributed-training behavior. This study empirically evaluates KAN scaling under DDP and examines whether its architectural differences produce distinct communication or memory behavior.
1. Introduction
- KAN’s learnable spline-parameterized univariate functions sit on network edges, while nodes sum their inputs; MLPs instead use fixed node activations and learnable edge weights.The construction is motivated by the Kolmogorov-Arnold representation theorem [2, 3].Source: §1 · pp. 1–2, Figure 1 · p. 2
- KAN’s interpretability and parameter-efficiency interest [1] [4] comes with the computational cost of evaluating and differentiating B-spline bases at every edge.Source: §1 · p. 2
- Prior KAN parallelization targets computation within a single training run, leaving multi-GPU, multi-node data-parallel replication as a separate question.MatrixKAN [5] accelerates spline evaluation, while Polar and Poluektov [6] address Newton-Kaczmarz fitting; neither studies replicated DDP training.Source: §1 · pp. 2–3
- It was unclear whether KAN’s per-edge learnable activations would yield distinct communication or memory scaling, or whether parameter count and ordinary DDP costs would dominate.Source: §1 · p. 3
- The study characterizes strong and weak scaling, communication overhead, and model-size scaling for DDP-based KAN training across configurations from one GPU to eight GPUs on four nodes.It tests whether observed patterns are consistent with general DDP behavior or differ for KAN.Source: §1 · p. 3
2. Related Work
This study applies distributed-training scaling methodology to KAN, complementing prior work that parallelizes computation within a single KAN run. It positions function-approximation architectures as relevant background and frames DDP scalability as a distinct systems-level question.
2.1. Data-Parallel Distributed Training
- Data-parallel training is a widely used distributed strategy, with prior work addressing scaling, stragglers, communication reduction, and All-Reduce design.The study evaluates standard synchronous, non-quantized DDP using PyTorch’s implementation [10], alongside topology-dependent NCCL behavior.Source: §2.1 · pp. 3–4
2.2. Parallelization Approaches for KAN
- Prior KAN approaches parallelize within a single-device run, whereas this study examines replicated multi-GPU, multi-node DDP training.MatrixKAN [5] accelerates spline evaluation; Polar and Poluektov [6] parallelize Newton-Kaczmarz fitting. The directions are additive: reducing per-GPU compute can make communication a larger share of step time.Source: §2.2 · p. 4
2.3. Function-Approximation Architectures
- KAN belongs to function-approximation architectures using structured basis functions rather than affine transformations followed by pointwise nonlinearities.Spline theory underlies KAN’s edge functions and relates to the memory and computation patterns examined in the study.Source: §2.3 · p. 4
2.4. Positioning of This Work
- The paper contributes an empirical scalability characterization of reference, unoptimized KAN under DDP on multi-node HPC hardware, applying established distributed-training measures to KAN.It does not accelerate KAN’s internal spline computation, a distinct problem addressed by MatrixKAN [5] and related work [6].Source: §2.4 · pp. 4–5
3. Methodology
The methodology evaluates DDP-based KAN training across four scalability dimensions, using controlled experiments on FinisTerrae III. It specifies the model and synthetic task, performance metrics, hardware topology, and configurations used for each experiment.
3. Methodology
- The evaluation varies one experimental factor at a time, reserves compute resources to avoid co-located-job interference, and warms up the system before timing.Source: §3 · p. 5
3.1. Performance Metrics
- Speedup, parallel efficiency, throughput, communication overhead, and parameter-to-memory ratio quantify scaling, utilization, synchronization time, and parameter storage density.The parameter-to-memory ratio is parameter count divided by peak GPU memory usage, not a general measure of memory-utilization efficiency.Source: §3.1 · p. 5
3.2. Experimental Platform
- Experiments ran on FinisTerrae III with A100 GPUs, PCIe intra-node links, and 200 Gb/s HDR InfiniBand inter-node links.Some configurations used one GPU per node to isolate inter-node communication; 4G4N and 8G4N then differed in whether GPU pairs shared a node.Source: §3.2 · pp. 5–6
3.3. Model and Data
- The model was a one-hidden-layer reference PyKAN KAN with input dimension 2, scalar output, cubic splines, and grid size 10; experiments used noisy synthetic regression on f(x, y) = x² + y².The synthetic task holds dataset size, noise, and per-sample compute cost controllable; results are not validated on real scientific or engineering datasets.Source: §3.3 · p. 6
3.4. Experimental Configurations
- Strong scaling fixed 100,000 samples, while weak scaling fixed 25,000 samples per GPU; model scaling varied width from 32 to 256 under 2G4N.Communication experiments isolated intra-node, inter-node, and hybrid patterns. Each configuration ran once for 40 epochs, so reported epoch variability does not estimate run-to-run variability.Source: §3.4 · p. 6
4. Experimental Results
The results section covers four evaluation dimensions: strong scaling, weak scaling, communication overhead, and model-size scaling.
4.1. Strong Scaling Performance Analysis
With a fixed workload of 100,000 samples, KAN training scales to 8 GPUs with 74.7% parallel efficiency and 5.97× speedup. Throughput and topology results indicate that the scaling pattern reflects growing communication overhead rather than a KAN-specific penalty.
4.1.1. Speedup and Parallel Efficiency
- 5.97× speedup and 74.7% parallel efficiency at 8 GPUs, with efficiency declining monotonically as GPU count grows.The reported efficiencies are 100.0%, 79.9%, 78.8%, 77.7%, 77.2%, and 74.7% across the 1, 2 (intra-node), 2 (inter-node), 4 (2-node), 4 (4-node), and 8-GPU configurations.Source: Figure 2 · p. 7, Table 1 · p. 7, §4.1.1 · p. 8
- The 74.7%–100% efficiency range is broadly consistent with conventional MLP and CNN DDP overhead, with no observed additional penalty from KAN's spline-based parameters [7] [10].Per-edge spline evaluation and gradients are local and independent, so they are not expected to add synchronization beyond DDP's gradient All-Reduce.Source: §4.1.1 · p. 8
- At fixed GPU count, node distribution has a secondary effect: 2G2N trails 2G1N by roughly one efficiency point, while 4G4N and 4G2N are nearly identical.The observed topology pattern is traced to All-Reduce algorithm selection rather than a KAN-specific effect.Source: §4.1.1 · p. 8
4.1.2. Throughput Analysis
- 26,618 samples/s total throughput at 8 GPUs, up from 4,442 samples/s on 1 GPU; per-GPU throughput falls 25.1%.The 5.99× total-throughput increase closely tracks measured speedup, while each additional GPU contributes less incremental throughput as communication overhead grows.Source: Figure 3 · p. 8, §4.1.2 · p. 9
- At 8 GPUs, scaling remains practically useful, with diminishing but still positive returns expected for sufficiently large workloads; the Amdahl-equivalent non-parallelizable fraction is approximately 4.9%.That fraction summarizes communication, synchronization, and residual load imbalance, rather than a literal serial code fraction.Source: §4.1.1 · p. 8, §4.1.2 · p. 9
4.2. Weak Scaling Stability Analysis
Weak scaling keeps each GPU's workload at 25,000 samples as dataset size and GPU count grow proportionally. Distributed configurations show stable per-GPU throughput after an initial drop from the single-GPU baseline.
4.2.1. Per-GPU Throughput Stability
- 3,440–3,463 samples/s/GPU across distributed configurations, versus 4,390 samples/s/GPU on one GPU; CV falls from 9.7% overall to under 0.3% among distributed runs.The distributed configurations span 2 to 8 GPUs.Source: Figure 4 · p. 9, Table 2 · p. 9, §4.2.1 · p. 10
- The approximately 21% single-to-multi-GPU throughput drop reflects fixed DDP costs, while throughput remains stable from 2 to 8 GPUs.Distributed runs take approximately 290 seconds across configurations, indicating that added communication is amortized by the proportionally larger dataset.Source: §4.2.1 · p. 10
- Once fixed distributed-training overhead is paid, per-GPU efficiency does not further degrade as workers increase, provided per-GPU workload stays approximately constant.The paper treats this as general DDP behavior under the tested workload, not as distinguishing weak-scaling behavior specific to KAN.Source: §4.2.1 · p. 10
4.3. Communication Overhead Analysis
Communication overhead varies non-monotonically with GPU placement, and the authors attribute the pattern primarily to All-Reduce behavior and network topology rather than KAN’s gradient structure.
4.3. Communication Overhead Analysis
- 1.3% (2G1N), 6.1% (2G2N), 1.4% (4G4N), 3.0% (4G2N), and 4.6% (8G4N): overhead follows a non-monotonic pattern attributed to All-Reduce and topology.Table 3 and Figure 5 report the topology breakdown; the proposed explanation is algorithm selection and network topology, not a KAN-specific gradient effect.Source: Figure 5 · p. 10, Table 3 · p. 11, §4.3 · p. 11
- The proposed mechanisms differ by placement: two-rank inter-node exchange is latency-sensitive, four-rank rings can pipeline, and hybrid tiers add coordination cost.The two-rank explanation is inferred rather than confirmed by NCCL_DEBUG traces; the four-rank and hybrid explanations are also presented as plausible mechanisms.Source: §4.3 · p. 11
- Efficiency exceeded 95% in four of five configurations; favor colocating two GPUs and avoiding mixed intra-/inter-node tiers when scaling beyond two GPUs.The 2G2N configuration is the outlier; 4G4N outperforms 4G2N despite using exclusively inter-node links.Source: §4.3 · p. 11
- The authors interpret the topology effects as likely general DDP behavior, but note that this was not tested with a matched-architecture ablation.The interpretation is bounded to models with comparable gradient tensor sizes; KAN’s per-edge spline coefficients affect that tensor size.Source: §4.3 · pp. 11–12
4.4. Model Scaling Characteristics Analysis
Across four model widths in a fixed 2G4N configuration, training time grew approximately linearly with width while the parameter-to-memory ratio improved; loss gains were modest relative to the per-parameter throughput cost.
4.4.1. Performance-Parameter Trade-off
- 81.5 s (Small) to 577.0 s (Extra-Large): training time rose 7.08× as width increased 8×, approximately linearly.The authors connect the trend to per-edge B-spline evaluation cost and a linearly growing edge count; explaining the roughly 11% gap from exact linear growth would require kernel-level profiling.Source: Figure 6 · p. 12, §4.4.1 · pp. 12–13, Table 4
- Final training loss improves from 0.100 (Small) to 0.080 (Extra-Large), modest relative to the 56× decrease in per-parameter throughput.Total throughput also falls from 24,624 to 3,469 samples/s across those sizes.Source: §4.4.1 · p. 13
4.4.2. Parameter-to-Memory Ratio Trends
- The parameter-to-memory ratio rises from 74,268 parameters/GB (Small) to 118,066 parameters/GB (Extra-Large), despite peak memory growing from 0.04 GB to 0.18 GB.The authors hypothesize that fixed per-process memory costs are amortized over larger models, but did not isolate those costs experimentally.Source: Figure 6 · p. 12, §4.4.2 · p. 13
5. Comprehensive Performance Evaluation and Discussion
No GPU configuration dominates all four performance dimensions; the preferred deployment depends on whether throughput or efficiency matters more. The discussion proposes complementary communication and computation optimizations while bounding its guidance to the tested workload and platform.
5.1. Multi-Dimensional Performance Synthesis
- Four-GPU configurations balance scalability and communication efficiency, while eight GPUs maximize throughput and one GPU maximizes compute efficiency; no configuration dominates all four axes.Figure 7 is a qualitative synthesis of normalized scores, not a single quantitative composite metric.Source: Figure 7 · p. 14, §5.1 · p. 14
5.2. Optimization Strategies
- Hierarchical All-Reduce and overlapping gradient communication with backward computation are proposed to address topology penalties and coordination overhead.The proposed hierarchy aggregates within a node over PCIe before crossing nodes over InfiniBand; the suggestions are not evaluated here.Source: §5.2 · p. 14
- Matrix-batched spline evaluation could reduce per-GPU compute time within DDP, making communication-aware topology selection more important in relative terms.MatrixKAN is discussed as complementary but was not implemented or evaluated in this study.Source: §5.2 · p. 15
5.3. Deployment Guidance
- For two GPUs, colocate them; at four GPUs, prefer one GPU per node, and use eight GPUs when throughput outweighs added communication overhead.The guidance is restricted to two-dimensional regression, models up to 21.5K parameters, A100 GPUs, and HDR InfiniBand.Source: §5.3 · p. 15
- Choose model width incrementally against accuracy and training-time budgets rather than maximizing memory use; the parameter-to-memory ratio improves with width.The guidance reflects the regression task studied, where per-parameter throughput falls sharply while accuracy gains are comparatively modest.Source: §5.3 · p. 15
5.4. Complementarity with Operator-Level Optimization
- Operator-level acceleration and topology-aware DDP address different training-stack costs and could be combined, but their joint effect was not measured.MatrixKAN reduces per-pass compute time; this study characterizes communication as replicated training spans GPUs and nodes.Source: §5.4 · p. 15
5.5. Limitations and Future Directions
- The findings are bounded by synthetic two-dimensional regression, models up to 21.5K parameters, and a single cluster of at most eight GPUs.Real scientific or engineering workloads, larger models, and other platforms remain unvalidated; the reported overhead percentages are platform-specific.Source: §5.5 · pp. 15–16
- The study uses reference recursive spline evaluation, and its proposed DDP/NCCL mechanisms lack direct matched-architecture ablation.Training-time values are upper bounds for an optimized implementation; whether qualitative trends persist is unconfirmed. A KAN–MLP control is identified as a priority for follow-up.Source: §5.5 · pp. 16–17
6. Conclusions
The study measures data-parallel KAN scalability across four dimensions and finds stable distributed scaling alongside topology-dependent communication overhead. Its conclusions support standard DDP within the tested scope, while emphasizing that real workloads, larger models, and matched-architecture controls remain to be tested.
6. Conclusions
- At eight GPUs, KAN reaches 74.7% parallel efficiency and 5.97× speedup; weak-scaling variation is 9.7% overall and 0.3% among distributed configurations.Communication overhead is 1.3%–6.1%; larger models improve the parameter-to-memory ratio from 74K to 118K parameters/GB while training time grows with width.Source: §6 · p. 17
- The study distinguishes general DDP behavior from effects more plausibly tied to KAN's per-edge spline parameterization, and frames its system-level findings as complementary to MatrixKAN.The two approaches address different training-stack layers: operator-level computation and multi-device communication.Source: §6 · p. 17
- The reported results support training KAN with standard DDP without pathologies beyond those expected for conventional architectures, within the study's tested scope.The scope is synthetic two-dimensional regression, models up to 21.5K parameters, and one HPC platform; real workloads and matched-architecture controls remain necessary for broader generalization.Source: §6 · p. 17
- Experimental scripts, logs, and analysis code are publicly available.The paper also states that its experiment and figure-generation code is available online.Source: §6 · p. 17, Data availability · p. 18