[Paper Review] Tetra-AML: Automatic Machine Learning via Tensor Networks

A. Naumov, Ar. Melnikov+6 more|arXiv (Cornell University)|Mar 28, 2023

Advanced Neural Network Applications6 citations

TL;DR

Tetra-AML introduces an automated machine learning framework that leverages tensor networks and a novel black-box optimizer, TetraOpt, to perform neural architecture search (NAS) and hyperparameter optimization (HPO) efficiently. It achieves state-of-the-art performance on CIFAR-10 with 3.2% accuracy loss while reducing ResNet-18 memory usage by 14.5× through combined quantization, pruning, and tensor network compression.

ABSTRACT

Neural networks have revolutionized many aspects of society but in the era of huge models with billions of parameters, optimizing and deploying them for commercial applications can require significant computational and financial resources. To address these challenges, we introduce the Tetra-AML toolbox, which automates neural architecture search and hyperparameter optimization via a custom-developed black-box Tensor train Optimization algorithm, TetraOpt. The toolbox also provides model compression through quantization and pruning, augmented by compression using tensor networks. Here, we analyze a unified benchmark for optimizing neural networks in computer vision tasks and show the superior performance of our approach compared to Bayesian optimization on the CIFAR-10 dataset. We also demonstrate the compression of ResNet-18 neural networks, where we use 14.5 times less memory while losing just 3.2% of accuracy. The presented framework is generic, not limited by computer vision problems, supports hardware acceleration (such as with GPUs and TPUs) and can be further extended to quantum hardware and to hybrid quantum machine learning models.

Motivation & Objective

• Address the rising computational and deployment costs of large deep neural networks in commercial applications.
• Reduce the time and resource burden of manual neural architecture and hyperparameter tuning.
• Enable efficient model compression for deployment on memory-constrained devices without significant accuracy loss.
• Extend the framework to support quantum and hybrid quantum-classical machine learning workloads.
• Provide a generic, hardware-accelerated toolbox for automated, end-to-end model optimization across diverse ML tasks.

Proposed method

• Employ a custom black-box optimization algorithm, TetraOpt, based on tensor trains to perform NAS and HPO in high-dimensional search spaces.
• Utilize tensor network decomposition to compress convolutional layers by representing weight tensors as low-rank tensor trains, reducing parameter count and memory footprint.
• Integrate standard compression techniques—pruning and quantization—alongside tensor network compression for synergistic efficiency gains.
• Support hardware acceleration via GPUs and TPUs, and extend TetraOpt to a quantum-optimized variant for future use in quantum machine learning.
• Design a modular pipeline where users define a dataset and search space, and the system automatically trains, optimizes, and compresses models.
• Apply post-training compression to already trained models, preserving accuracy while minimizing inference costs.

Figure 1: General scheme of Tetra-AML . Both Neural Architecture Search (NAS) and Hyperparameter Optimization (HPO) are performed via TetraOpt (Terra Quantum’s black-box optimizer based on Tensor Trains). Then the model is compressed via tensor network methods, quantization and pruning.

Experimental results

Research questions

• RQ1 Can tensor network-based optimization outperform Bayesian optimization in neural architecture and hyperparameter search for computer vision tasks?
• RQ2 To what extent can tensor network compression reduce model size and inference cost while maintaining high accuracy?
• RQ3 How effective is the combination of pruning, quantization, and tensor network compression in reducing memory usage for deep neural networks?
• RQ4 Can the TetraOpt optimizer be extended to support quantum machine learning and hybrid quantum-classical model optimization?
• RQ5 What is the trade-off between compression ratio and accuracy loss in compressed ResNet-18 models on CIFAR-10?

Key findings

• TetraOpt outperforms Bayesian optimization in terms of validation accuracy on the CIFAR-10 dataset when given the same number of model evaluations.
• ResNet-18 models compressed using Tetra-AML achieve a 14.5× reduction in memory usage while incurring only a 3.2% drop in accuracy.
• Tensor network compression enables exponential reduction in memory and runtime by decomposing large weight tensors into low-rank tensor trains.
• The framework supports hardware acceleration via GPUs and TPUs, enabling efficient training and inference on standard and high-performance hardware.
• The integration of pruning, quantization, and tensor network compression leads to significantly improved model efficiency compared to individual techniques.
• The framework is extensible to quantum machine learning, with TetraOpt's quantum-optimized variant capable of improving hyperparameter and architecture search for quantum and hybrid models.

Figure 2: Validation accuracy dependence on the number of neural network models (architectures) runs for TetraOpt , Bayesian, and Random Search algorithms. TetraOpt achieves 93.7% accuracy, while Bayesian optimization and Random Search find architectures only with 93.5% and 92.8% accuracy, respectiv