Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods

Despite their simple intuition, convolutions are more tedious to analyze than dense layers, which complicates the transfer of theoretical and algorithmic ideas to convolutions. We simplify convolutions by viewing them as tensor networks (TNs) that allow reasoning about the underlying tensor multiplications by drawing diagrams, manipulating them to perform function transformations like differentiation, and efficiently evaluating them with einsum. To demonstrate their simplicity and expressiveness, we derive diagrams of various autodiff operations and popular curvature approximations with full hyper-parameter support, batching, channel groups, and generalization to any convolution dimension. Further, we provide convolution-specific transformations based on the connectivity pattern which allow to simplify diagrams before evaluation. Finally, we probe performance. Our TN implementation accelerates a recently-proposed KFAC variant up to 4.5x while removing the standard implementation's memory overhead, and enables new hardware-efficient tensor dropout for approximate backpropagation.

Paper

Similar papers

Peer review

Reviewer vbTD7/10 · confidence 3/52024-07-10

Summary

This paper provides a tensor network (TN) view of studying convolutional networks, including forward passes, first-order and second order derivatives. Based on these notations, the authors show how to efficiently implement some algorithms such as KFAC and structured dropout. Experimental results show that these implementations have faster computation while being memory efficient.

Strengths

1. The paper is well written with nice tensor network graphs. While many papers in the tensor decomposition field suffer from notorious indexes, the notation in this paper is consistent and clear. 2. Although the TN graphs for matrix multiplications and derivatives themselves are not new, this paper seems to be the first to provide a comprehensive study of these representations and computations in CNNs. Based on these representations, the authors show how to efficiently implement the KFAC and structured dropout algorithm, which is important for optimization. 3. Experiments show the computational and memory efficiency of the proposed TN implementations.

Weaknesses

The authors claimed some potential impacts of the proposed representation, including theoretical and practical benefits. However, they only present applications on the KFAC and structured dropout. Moreover, the derivations seem to require expertise in TNs. Therefore, I feel the potential usefulness of the proposed representation is still unclear.

Questions

The authors show faster speed of KFAC computation. Does this implementation bring better convergence, training speed or results in practice? Is there any result for the trained network?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

None

Reviewer cYbp7/10 · confidence 3/52024-07-12

Summary

The paper presents a novel approach to simplifying and optimising TN operations for CNNs. The authors introduce an abstraction for tensor networks to efficiently handle convolutions and leverage TN simplifications to improve computational performance. The paper details the theoretical foundations, implementation strategies and results. The results suggest improvements in run-time and memory efficiency over standard methods.

Strengths

The key strengths are: - The paper includes rigorous mathematical formulations and proofs to support the proposed methods. - Provides thorough details on implementation, aiding reproducibility and practical adoption. - Applicable to a wide range of convolutional operations, including those with various hyper-parameters, batching, and channel groups. The generality enhances utility and flexibility - Good use of visualisations and diagrams to clarify and explain the networks - Interest to the research community with strong potential for future extensions

Weaknesses

The key weakness is that the performance demonstrated is not very convincing in many cases. This is understandable as the authors mention, standard routines have been aggressively optimised whereas the TN counterparts have been explored significantly less. However this by itself is not a reason to reject given the potential further research that can be built on this work and the interest it Minor things: - Line 57 does not appear to have the reference to 2.2 in the right place - Two pages are spent on the preliminaries. I see the value of having this to set the context for the later sections but I wonder, if some of those were moved to the appendix, what new content could be added to the main method and results? - The NeurIPS format states "Place one line space before the figure caption and one line space after the figure." Some of the Figure captions are placed next to the figure in a two-column setup which I am not sure if satisfies the style rules - Standard errors, rather than standard deviation, might be more informative in Figure 8

Questions

Questions for authors: - Why are the (non simplified) TN implementations slower in general? Do the standard implementations have further efficiency optimisations? - Are there any practical ways to estimate a priori whether a selected TN has high contraction order variability? Note: to know the approximate variability range, not necessarily the exact values of each order - Does figure 6 include variability wrt contraction order? - How many of the operations support GPU/hardware acceleration? Would there be performance gains if they were implemented?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

4

Limitations

There is a section discussing limitations, but the authors could mention more practical factors. Such as: scalability issues and how the methods perform with larger datasets (should be straightforward with synthetic data), potential biases introduced and their impact on different types of data and potential real-world scenarios where the methods might fail or underperform.

Reviewer iQrS5/10 · confidence 4/52024-07-17

Summary

### Summary of "Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods" The paper "Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods" presents a compelling and innovative recast of convolution operations into tensor networks, offering significant benefits for both theoretical understanding and practical applications in machine learning. Here’s a motivated summary with the key points: 1. **Recast of Convolution Operation into Tensor Networks**: - The authors provide a nice review and recast of convolutions by viewing them as tensor networks (TNs). This perspective simplifies the analysis of convolutions, allowing for intuitive reasoning about underlying tensor multiplications through easily interpretable diagrams. This approach makes it possible to perform function transformations like differentiation more effectively and efficiently. 2. **Tensor Networks for Machine Learning**: - This paper is very good material for studying tensor networks within the context of machine learning. By representing convolutions as TNs, the authors bridge the gap between complex convolutional operations and the more straightforward tensor manipulations, making advanced theoretical concepts accessible to practitioners and researchers in machine learning. 3. **Implementation with Einsum and Opt_Einsum**: - The authors implement convolutions using the einsum and opt_einsum functions in PyTorch/Python, providing a detailed comparison of the efficiency of these implementations. The tensor network approach not only accelerates the computation but also reduces memory overhead significantly. The paper demonstrates that tensor network implementations can achieve substantial performance improvements, particularly for less standard routines like Kronecker-factored Approximate Curvature (KFAC). 4. **Insights into Second-Order Optimization**: - The implementation of KFAC benefits significantly from the tensor network view of convolutions. The TN perspective provides new insights into second-order optimization methods, allowing for more efficient computation of curvature approximations. The paper shows that the tensor network approach can lead to more frequent pre-conditioner updates, larger batch sizes without memory issues, and extends KFAC to transpose convolutions, demonstrating its broad applicability to second-order optimization in neural networks. In summary, this paper provides a well-rounded and innovative approach to understanding and optimizing convolutions through tensor networks. It not only offers theoretical insights but also demonstrates practical performance benefits, making it a valuable resource for researchers and practitioners in machine learning.

Strengths

### Strengths of "Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods" 1. **Innovative Perspective**: - The paper introduces an innovative perspective by recasting convolution operations into tensor networks (TNs). This novel approach simplifies the analysis and implementation of convolutions, making complex operations more understandable and accessible. 2. **Advances in Second-Order Optimization**: - The paper shows that the tensor network view of convolutions can greatly benefit second-order optimization methods, particularly Kronecker-factored Approximate Curvature (KFAC). This provides new insights and practical improvements in the computation of curvature approximations. 3. **Comprehensive Diagrams and Transformations**: - The inclusion of detailed diagrams and specific transformations based on the connectivity pattern of convolutions enhances the clarity and utility of the proposed methods. These visual aids make it easier to understand and implement the tensor network approach. 4. **Flexibility and Modifiability**: - The tensor network framework allows for easy modifications and supports various convolutional configurations, including different hyper-parameters, batching, channel groups, and convolution dimensions. This flexibility makes the approach broadly applicable. 5. **Experimental Validation**: - The paper provides robust experimental validation, demonstrating the practical benefits of the tensor network approach in various convolutional operations and configurations. The significant speed-ups and memory savings observed in the experiments highlight the effectiveness of the proposed methods.

Weaknesses

### Weaknesses of "Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods" While the paper "Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods" offers numerous strengths, it also has some notable weaknesses: 1. **Suboptimal Use of Einsum for Convolutions**: - Using einsum for convolution operations is generally not an optimal approach. The paper itself acknowledges that an optimal method may not exist, as indicated by the mixed results in Tables F4-7, where the performance comparison between the default PyTorch (PT) implementation and the tensor network (TN) approach shows that TN is not consistently faster. This suggests that while TNs offer theoretical benefits, their practical efficiency can vary significantly. 2. **Performance Trade-offs for Specific Architectures**: - The einsum-based approach favors specific architectures but makes some of the frequently used models slower. This raises questions about the general applicability of the proposed method. There is a lack of intuition or explanation provided for why certain architectures benefit from the TN approach while others do not. Understanding these trade-offs is crucial for evaluating the practical utility of the method across different use cases. 3. **Limited Comparison with Optimized Convolution Algorithms**: - The comparison between TN and PyTorch does not specify which convolutional algorithm is used in PyTorch's implementation. Furthermore, the paper does not compare the TN approach with other highly optimized versions of convolution, such as Winograd or Fourier transform-based methods. These optimized algorithms are known to significantly improve convolution performance, and a comparison with them would provide a more comprehensive assessment of the TN approach's efficiency. In summary, while the paper presents a novel and theoretically sound method for implementing convolutions using tensor networks, its practical applicability is limited by the suboptimal performance of einsum for convolutions, unexplained performance trade-offs across different architectures, and a lack of comparison with other optimized convolution algorithms. Addressing these weaknesses could enhance the robustness and utility of the proposed approach.

Questions

Please see weakness

Rating

5

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Please see weakness

Authorsrebuttal2024-08-12

We would once again like to thank all reviewers for their time and effort and their constructive feedback which helped us further improve the manuscript. If there are any remaining questions, please let us know before the end of the discussion phase.

Reviewer vbTD2024-08-13

Thanks for the authors' response. I will maintain the score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC