Tensor Trains Accelerate Quantum-Inspired Homogenization on TPUs, GPUs, CPUs, Enabling 10x Faster Processing of High-resolution Datasets

The increasing resolution of modern imaging techniques generates complex datasets that strain the capabilities of conventional computational modelling, particularly in fields like materials science and engineering. Sascha H. Hauck, Matthias Kabel, and Nicolas R. Gauger, from Fraunhofer ITWM and the University of Kaiserslautern-Landau, investigate how a novel approach utilising Tensor Trains can overcome these limitations. Their work focuses on accelerating a technique called homogenization, which predicts the effective properties of complex materials, by leveraging the power of modern hardware accelerators. By benchmarking fundamental Tensor Train operations on CPUs, GPUs, and TPUs, the team demonstrates significant speed-ups, up to ten times faster than traditional CPU-based methods, and reveals that both GPUs and TPUs offer comparable performance, opening new possibilities for tackling previously intractable industrial-scale datasets and advancing the field of multi-scale modelling.

Recognizing the limitations of traditional techniques when applied to very large datasets, the team focused on reducing both computational cost and memory requirements. This research pioneers the transfer of core TT operations to Graphics Processing Units (GPUs) and Tensor Processing Units (TPUs), enabling efficient handling of complex data. The study demonstrates that these quantum-inspired algorithms can effectively solve homogenization problems, which determine the effective properties of materials.

Researchers meticulously benchmarked core TT operations on each accelerator, measuring execution times and efficiency. This detailed analysis revealed that TPUs generally outperform GPUs and CPUs for tensor operations, particularly when processing multiple operations simultaneously. The team also investigated Polar-based TT compression, a technique to further reduce memory usage, and assessed its performance on the accelerators. Results show that Polar decomposition generally outperforms Singular Value Decomposition in terms of speed, although the optimal choice depends on the specific application and desired balance between speed and accuracy.

The Roofline model revealed that memory bandwidth often limits computational speed. By applying this accelerated algorithm to a realistic scenario, scientists demonstrated the feasibility of analysing previously intractable datasets. The results show that both GPUs and TPUs achieve comparable performance, with speed-ups of up to ten times compared to traditional CPU-based implementations. This work addresses the computational challenges posed by these datasets, specifically for homogenization techniques used to predict material behaviour. The team investigated the performance of fundamental TT operations on modern hardware accelerators using the JAX framework, comparing CPUs, GPUs, and TPUs.

Benchmarking revealed that adapting the SFFT-based algorithm for use on accelerators achieves speed-ups of up to ten times relative to CPU implementation, enabling the treatment of previously infeasible dataset sizes. Results demonstrate that both GPUs and TPUs achieve comparable performance in realistic scenarios, despite the relative immaturity of the TPU ecosystem, highlighting the potential of both architectures for accelerating quantum-inspired techniques. The team provides the first systematic benchmarking of TT algebra on TPU hardware, alongside a direct comparison with GPU performance. Scientists have successfully adapted a Superfast-Fourier Transform (SFFT)-based homogenization algorithm, originally designed for central processing units, to take advantage of modern hardware accelerators, including graphics processing units and tensor processing units. By leveraging the memory efficiency of Tensor Train decompositions, the team achieved speed-ups of up to ten times compared to traditional CPU-based implementations. While acknowledging that the performance of the algorithm depends on the underlying data exhibiting suitable structure to achieve a low Tensor Train rank, this work represents a crucial step towards enabling detailed analysis of complex materials and processes. Future research will likely focus on extending this accelerated approach to even more complex models and exploring its application to a wider range of scientific and engineering problems.