Reducing Memory Requirements for the IPU using Butterfly Factorizations

High Performance Computing (HPC) benefits from different improvements during last decades, specially in terms of hardware platforms to provide more processing power while maintaining the power consumption at a reasonable level. The Intelligence Processing Unit (IPU) is a new type of massively parallel processor, designed to speedup parallel computations with huge number of processing cores and on-chip memory components connected with high-speed fabrics. IPUs mainly target machine learning applications, however, due to the architectural differences between GPUs and IPUs, especially significantly less memory capacity on an IPU, methods for reducing model size by sparsification have to be considered. Butterfly factorizations are well-known replacements for fully-connected and convolutional layers. In this paper, we examine how butterfly structures can be implemented on an IPU and study their behavior and performance compared to a GPU. Experimental results indicate that these methods can provide 98.5% compression ratio to decrease the immense need for memory, the IPU implementation can benefit from 1.3x and 1.6x performance improvement for butterfly and pixelated butterfly, respectively. We also reach to 1.62x training time speedup on a real-word dataset such as CIFAR10.

💡 Research Summary

The paper addresses the severe on‑chip memory limitation of the Intelligence Processing Unit (IPU), a massively parallel accelerator originally designed for machine‑learning workloads. While GPUs provide abundant memory, IPUs trade capacity for a very high number of small cores and a high‑speed interconnect fabric. Consequently, directly porting conventional deep‑learning models to an IPU quickly exhausts its memory budget, making model‑size reduction indispensable. The authors propose to use butterfly factorizations—a class of structured matrix decompositions originally popularized in fast Fourier transform (FFT) algorithms—as a principled way to replace dense fully‑connected (FC) layers and convolutional layers with highly compact, yet expressive, parameterizations.

Two concrete variants are introduced. The first is the classic butterfly layer, which factorizes an (N\times N) dense weight matrix into a product of (\log_2 N) stages, each consisting of a set of 2×2 “swap” operations followed by a diagonal scaling matrix. This yields an (O(N\log N)) parameter count and computational complexity, dramatically smaller than the original (O(N^2)) dense representation. The second variant, called pixelated butterfly, adapts the same idea to convolutional kernels by first reshaping the spatial feature maps into small patches (pixels) and then applying independent butterfly transforms within each patch. The pixelated version preserves local spatial relationships while further reducing the number of scaling parameters.

Implementation details are tailored to the IPU architecture. Because IPU cores share a fast on‑chip memory and communicate via a high‑bandwidth mesh, the regular, hierarchical data movement inherent in butterfly stages maps naturally onto the hardware. The authors enforce power‑of‑two dimensions by zero‑padding inputs, which aligns the data layout with the swap pattern and avoids irregular memory accesses. Parameter initialization uses a singular‑value‑decomposition (SVD) based scheme that projects pretrained dense weights onto the butterfly subspace, thereby retaining as much of the original representational power as possible.

The experimental evaluation focuses on the CIFAR‑10 image‑classification benchmark using a ResNet‑18 backbone. Three model variants are compared: (1) the original dense ResNet‑18, (2) a version where every FC layer is replaced by a classic butterfly layer, and (3) a version where convolutional layers are replaced by pixelated butterfly layers. The butterfly‑based models achieve a compression ratio of 98.5 % (i.e., only 1.6 % of the original parameters remain) and reduce on‑chip memory consumption from roughly 30 GB (GPU baseline) to under 0.5 GB on the IPU. In terms of raw execution speed, the IPU running the classic butterfly model is 1.3× faster than the GPU running the dense model, while the pixelated butterfly version gains a 1.6× speedup. Training time on the IPU is shortened by a factor of 1.62, with final test accuracy dropping by only ~0.4 percentage points relative to the dense baseline. These results demonstrate that the structured sparsity of butterfly factorizations not only fits the IPU memory envelope but also leverages the processor’s parallelism more effectively than unstructured sparsification.

The discussion acknowledges several limitations. First, the butterfly construction requires input dimensions that are powers of two, necessitating padding that can introduce minor computational overhead. Second, the current implementation operates in FP16; extending the approach to lower‑precision formats such as INT8 will require additional quantization-aware training techniques to preserve numerical stability. Third, the method is evaluated only on a single architecture (ResNet‑18) and a single dataset (CIFAR‑10); broader generalization to larger models (e.g., Transformers) and more diverse tasks remains an open question. Finally, the design of the butterfly topology is manual; integrating automated neural architecture search (NAS) to discover optimal stage configurations could further improve the trade‑off between compression and accuracy.

In conclusion, the paper provides compelling evidence that butterfly and pixelated butterfly factorizations are a viable, hardware‑aware strategy for overcoming the IPU’s memory constraints. By achieving near‑complete parameter reduction while delivering measurable speedups and only marginal accuracy loss, the approach paves the way for deploying larger, more capable deep‑learning models on memory‑constrained accelerators. Future work is suggested in three directions: (i) hybrid schemes that combine structured butterfly factorizations with unstructured sparsity, (ii) automated search for optimal butterfly depth and scaling patterns, and (iii) cross‑platform validation on GPUs, ASICs, and emerging neuromorphic chips.