tn4ml: Tensor Network Training and Customization for Machine Learning

TL;DR

The paper introduces tn4ml, a Python library that structures Tensor Network models into end-to-end machine learning pipelines, covering data embedding, TN architecture and initialization, optimization, and evaluation. It focuses on one-dimensional TNs (e.g., MPS, MPO, SMPO) and offers multiple embedding schemes, including Product State and Entangled State (with Patch Embedding FRQI), along with various initialization and optimization strategies implemented via JAX-based tooling. The authors demonstrate the framework on two use cases—binary breast cancer classification and MNIST-based anomaly detection—highlighting how hyperparameters like bond dimension and spacing influence accuracy and computational efficiency on CPU/GPU hardware. They provide an open-source, MIT-licensed tool designed to help practitioners build explainable TN-based models and to guide hyperparameter choices for TN pipelines.

Abstract

Tensor Networks have emerged as a prominent alternative to neural networks for addressing Machine Learning challenges in foundational sciences, paving the way for their applications to real-life problems. This paper introduces tn4ml, a novel library designed to seamlessly integrate Tensor Networks into optimization pipelines for Machine Learning tasks. Inspired by existing Machine Learning frameworks, the library offers a user-friendly structure with modules for data embedding, objective function definition, and model training using diverse optimization strategies. We demonstrate its versatility through two examples: supervised learning on tabular data and unsupervised learning on an image dataset. Additionally, we analyze how customizing the parts of the Machine Learning pipeline for Tensor Networks influences performance metrics.

Figures (10)

• Figure 1: Graphical representation of MPS with upper real $i_k$ and virtual indices $D_k$, and MPO with additional lower real indices $j_k$. Dimensions of real and virtual indices are chosen following different criteria. Commonly upper indices follow embedding dimension, while lower and virtual indices are treated as hyperparameters.
• Figure 2: Graphical representation of SMPO with upper real $i_k$, virtual indices $D_k$, and lower real indices $j_k$. The spacing parameter $S$ determines the number of lower indices.
• Figure 3: ML pipeline for TNs consists of (1) data embedding procedure; (2) choice of TN architecture and initialization; (3) model optimization with choice of objective function and training strategy and (4) model evaluation, where P indicates a parametrized TN model.
• Figure 4: Graphical notation of a product state for Product State Embedding where each local feature map $\phi_i(x_i)$ has dimension $d$, and no virtual indices between tensors.
• Figure 5: The objective function for classification is obtained by contracting embedded data $\Phi(x)$ with the MPS model $P$, resulting in another MPS, which is further contracted to produce a real-valued vector $\vec{y}$ with size equal to the number of classes $C$.