Tensor network kernel machines (TNKMs) combine nonlinear features with low-rank tensor representations to model complex relationships efficiently. However, selecting an appropriate tensor rank remains challenging, as conventional rank-search methods require training multiple candidate models, increasing computational costs. In this work, we focus on a canonical polyadic (CP) tensor network parameterization and propose GB-CPKM, a gradient boosting framework that builds a higher-capacity predictor by sequentially adding fixed-rank learners. This approach provides a new perspective on rank selection by incrementally increasing model capacity, with ensemble size selected based on validation performance rather than through explicit rank search. Experiments on real-world regression benchmarks show that GB-CPKM trains up to 6x faster and uses up to 2x less peak memory than a conventional rank-selection method based on cross-validation, while achieving comparable predictive performance. These results highlight the potential of gradient boosting to improve the computational efficiency and scalability of tensor-network learning.
Figures & tables
Figure 1: Comparison of the training procedures for CP tensor network kernel machines (CPKMs) using gradient boosting (GB) and rank selection via validation (CV). (a) GB-CPKM incrementally builds an ensemble by training fixed-rank R^ base learners and concatenating their factor matrices along the rank dimension to represent the combined model. The ensemble size M is selected based on validation performance. (b) In contrast, CV-CPKM independently trains candidate models with different CP ranks and selects the rank that achieves the lowest validation error. In the diagrams, each base learner model is represented by D factor matrices of size I×R^ , where I is the number of rows and R^ is the CP rank. Different colors distinguish the trainable factor matrices. Factors in gray are fixed and are not updated.
Figure 2: Ablation study of base-learner rank in GB-CPKM on four real-world datasets, with each column corresponding to a dataset. The top row shows cumulative training time (in seconds), while the bottom row reports validation MSE on a logarithmic scale. Across the evaluated datasets, higher-rank base learners achieve better predictive performance while also reducing cumulative training time. However, rank selection should also account for memory requirements, which may constrain the use of higher-rank models in practice.
Figure 3: Comparison of GB-CPKM (red) and CV-CPKM (blue) on four real-world datasets, with each column corresponding to a dataset. From top to bottom, the rows show cumulative training time (in seconds), test MSE, and peak memory consumption (in MiB), plotted against the number of considered candidate CP ranks. Solid lines indicate the mean across 10 runs, and shaded regions represent ± 1 standard deviation. GB-CPKM consistently requires less training time and peak memory than CV-CPKM, while achieving comparable predictive performance. The test MSE of GB-CPKM is similar to that of CV-CPKM, with overlapping variability bands in some cases, and approaches comparable values across the evaluated datasets. These results highlight the potential of GB-CPKM to reduce computational costs without substantially compromising predictive accuracy.
MSE (test) ↓
Time (sec.) ↓
Memory (MiB) ↓
Data
N
D
CV-CPKM
GB-CPKM
GB-CPKM(H)
CV-CPKM
GB-CPKM
CV-CPKM
GB-CPKM
Energy
768
8
0.0026 ±0.00
0.0043±0.00
0.0038±0.00
0.14±0.02
0.02 ±0.01
534
444
Concrete
1030
8
0.1618 ±0.05
0.1889±0.04
0.1801±0.05
0.18±0.00
0.06 ±0.01
572
461
Kin8nm
8192
8
0.1077 ±0.01
0.1721±0.03
0.1209±0.01
1.28±0.03
0.33 ±0.02
865
551
Naval
11934
16
0.0000 ±0.00
0.0074±0.00
0.0010±0.00
1.78±0.06
0.36 ±0.02
1090
630
Precipitation
628474
3
0.9419 ±0.04
0.9454±0.04
0.9448±0.04
88.59±11.84
8.35 ±1.05
2885
1156
Table 1: Performance comparison of GB-CPKM and CV-CPKM across regression benchmarks. We report test mean squared error (MSE), training time (seconds), and peak memory consumption (MiB). Each experiment is repeated over 10 runs, with results reported as mean ± 1 standard deviation. N denotes the number of samples, and D denotes the input dimensionality. Overall, GB-CPKM trains up to 6× faster and uses up to 2× less peak memory, on average, than CV-CPKM, while achieving comparable test MSE. The hybrid variant, GB-CPKM(H), further improves predictive performance on most datasets by performing an additional full ALS sweep on the combined model.
Tensor networks represent high-dimensional data and quantum states as interconnected low-rank tensors to enable efficient computation, learning, and simulation.
Tensor Network (TN) kernel machines accelerate supervised learning by using low-rank TNs to reduce computation and memory. However, most TN-based methods are deterministic with manually tuned ranks and feature dimensions, while the existing Bayesian TN kernel machine is restricted to the CANDECOMP/PARAFAC decomposition (CPD), whose single rank is shared by all modes. We extend this framework to the tensor train (TT) and propose Bayesian Tensor Train Kernel Machines (BTT-KM), whose hierarchical ARD-type priors on the TT cores infer the TT ranks and feature dimensions from data. Since each TT core is coupled to the others only through its neighbors, its closed-form posterior update requires propagating the contributions of all other cores along the chain. Mean-field variational inference yields a Bayesian alternating least squares algorithm with the same computational complexity as its deterministic counterpart and a closed-form predictive distribution. A synthetic experiment illustrates that BTT-KM recovers the true TT ranks and feature dimensions from an over-parameterized initialization.
Afra Kilic, Kim Batselier
Delft Center for Systems and Control, Delft University of Technology, Delft, 2628 CN, The Netherlands
We study rank selection for low-rank tensor regression under random covariates design. Under a Gaussian random-design model and some mild conditions, we derive population expressions for the expected training-testing discrepancy (optimism) for both CP and Tucker decomposition. We further demonstrate that the optimism is minimized at the true tensor rank for both CP and Tucker regression. This yields a prediction-oriented rank-selection rule that aligns with cross-validation and extends naturally to tensor-model averaging. We also discuss conditions under which under- or over-ranked models may appear preferable, thereby clarifying the scope of the method. Finally, we showcase its practical utility on a real-world image regression task and extend its application to tensor-based compression of neural network, highlighting its potential for model selection in deep learning.
Haoming Shi, Eric C. Chi, Hengrui Luo
Department of Statistics, Rice University · School of Statistics, University of Minnesota
Developing nonlinear models that are both expressive and computationally efficient remains a challenge in machine learning and nonlinear system identification. Tensor network kernel machines (TNKM) address this challenge by combining nonlinear feature representations with compact low-rank tensor-network parameterizations. However, practical and extensible software frameworks for developing TNKM models remain limited. In this work, we introduce "tnkm", an open-source Python library for constructing and training TNKM models using JAX. The library provides a unified interface for combining different feature maps, tensor-network architectures, and optimization strategies, including alternating least squares and gradient-based methods. We demonstrate the capabilities of "tnkm" on nonlinear benchmark problems, showing that the implemented models achieve competitive prediction accuracy while retaining compact parameterizations and efficient training. The proposed framework facilitates reproducible development and application of tensor-network-based learning methods.
Albert Saiapin, Kim Batselier
Delft Center for Systems and Control, TU Delft, Netherlands