TensorKrowch: Smooth integration of tensor networks in machine learning

TL;DR

The paper addresses the challenge of bringing tensor networks into mainstream machine learning practice by introducing TensorKrowch, a PyTorch‑based library that lets researchers build, train, and integrate tensor‑network layers. It defines Nodes/Edges as building blocks to compose arbitrary tensor networks and contract them end‑to‑end within DL models. Key contributions include efficient memory and contraction optimizations (trace, memory modes, and shared memory), built‑in architectures such as $MPS$, $MPO$, $TTN$, and $PEPS$, and tensorization tools for integrating with pretrained models. The work facilitates rapid prototyping and broader adoption of tensor networks in ML, supported by documentation and community‑driven contribution guidelines.

Abstract

Tensor networks are factorizations of high-dimensional tensors into networks of smaller tensors. They have applications in physics and mathematics, and recently have been proposed as promising machine learning architectures. To ease the integration of tensor networks in machine learning pipelines, we introduce TensorKrowch, an open source Python library built on top of PyTorch. Providing a user-friendly interface, TensorKrowch allows users to construct any tensor network, train it, and integrate it as a layer in more intricate deep learning models. In this paper, we describe the main functionality and basic usage of TensorKrowch, and provide technical details on its building blocks and the optimizations performed to achieve efficient operation.

Figures (2)

• Figure 1: Comparison of the maximum memory usage of one contraction of the built-in MPSLayer tensor network, when the model is traced or not, using different bond dimensions. Contraction is performed in a training regime: 1) an example data tensor is passed through the model, 2) gradients are computed via backpropagation, and 3) parameters are updated according to the gradients. All the contractions are computed in CPU using a batch size of 500, with both memory modes auto_stack and auto_unbind set to True, both contraction arguments inline_input and inline_mats also set to True, and the following arguments for the MPSLayer model: n_features=1000, in_dim=2, out_dim=10. All the experiments were run on an Intel Xeon CPU E5-2620 v4 with 256GB of RAM.
• Figure 2: Comparison of running times of one contraction of the built-in MPSLayer tensor network, using different bond dimensions. Contraction is performed in different regimes: training/inference, parallel/inline algorithm, CPU/GPU execution, and using different combinations for the options auto_stack and auto_unbind. For training, 1) an example data tensor is passed through the model, 2) gradients are computed via backpropagation and 3) parameters are updated according to the gradients. For inference, only the example data tensor is passed to perform one contraction of the model. Parallel and inline refer to the two possible algorithms that can be used to contract MPSLayer, specified by the argument inline_mats. When inline_mats = True (inline), the matrices of the MPS are contracted sequentially; when inline_mats = False (parallel), matrices are iteratively stacked, contracted in parallel and unbound until the chain of matrices is reduced to a single one. The argument inline_input is always set to False, indicating that the contraction of the MPS nodes with the respective input nodes is performed in parallel. Solid lines represent CPU times, dashed lines represent GPU times. All the contractions are computed using a batch size of 100 and the following arguments for the MPSLayer model: n_features=100, in_dim=2, out_dim=10. All the experiments were run on an Intel Xeon CPU E5-2620 v4 with 256GB of RAM and an NVIDIA GeForce RTX 3090.