[Paper Review] TensorNetwork for Machine Learning

Stavros Efthymiou, Jack D. Hidary+1 more|arXiv (Cornell University)|Jun 7, 2019

Quantum many-body systemsPhysics and Astronomy31 references62 citations

TL;DR

The paper demonstrates image classification using matrix product state tensor networks via the TensorNetwork library, achieving 98% MNIST and 88% Fashion-MNIST test accuracy with GPU-accelerated automatic-gradient training.

ABSTRACT

We demonstrate the use of tensor networks for image classification with the TensorNetwork open source library. We explain in detail the encoding of image data into a matrix product state form, and describe how to contract the network in a way that is parallelizable and well-suited to automatic gradients for optimization. Applying the technique to the MNIST and Fashion-MNIST datasets we find out-of-the-box performance of 98% and 88% accuracy, respectively, using the same tensor network architecture. The TensorNetwork library allows us to seamlessly move from CPU to GPU hardware, and we see a factor of more than 10 improvement in computational speed using a GPU.

Motivation & Objective

• Show how tensor networks can be applied to image classification.
• Encode image data into a matrix product state form and train with automatic gradients.
• Demonstrate performance and speedups on MNIST and Fashion-MNIST datasets.
• Provide open-source code and TensorNetwork integration with TensorFlow for practitioners.

Proposed method

• Encode each image pixel into a two-dimensional local feature map to form a data tensor.
• Represent the classifier as an MPS tensor with a label index and compute inner products with the encoded data to obtain f^(l)(x).
• Train using multi-class cross-entropy with softmax over labels and backpropagation via automatic differentiation.
• Discuss contraction orders and computational cost, favoring a parallelizable contraction strategy.
• Utilize TensorFlow backend to enable automatic gradients and Adam optimization for training.
• Compare CPU vs GPU performance and assess dependence on bond dimension chi (chi >= ~10).

Experimental results

Research questions

• RQ1 Can a matrix product state tensor network classify images effectively on MNIST and Fashion-MNIST?
• RQ2 What is the impact of bond dimension chi on accuracy and training cost?
• RQ3 How does TensorNetwork with TensorFlow enable gradient-based optimization for tensor networks?
• RQ4 What are the practical speedups when moving from CPU to GPU for this approach?
• RQ5 How do contraction orders affect computational efficiency and parallelism?

Key findings

• MNIST test accuracy ~98% on the full 60k train, 10k test split using the MPS network.
• Fashion-MNIST test accuracy ~88% under the same architecture and settings.
• GPU training yields about 10x faster wall-clock time per epoch compared to CPU for the same codebase.
• Final accuracies show little dependence on bond dimension chi for chi ≳ 10.
• Cross-entropy loss with softmax performs comparably to mean-squared error in terms of final accuracy.
• Training with automatic gradients (TensorFlow) is feasible and effective for MPS-based classification.