Convolutions Through the Lens of Tensor Networks

Dear Reviewer 4fmq,

thanks for your detailed review, especially for taking the time to formulate all of your questions. We would like to address them in the following. Please let us know if you have follow-up questions or think we did not address all of them. We are happy to discuss and provide more details.

Difference to using auto-diff & Hayashi et al.: By deriving the Jacobians, our goal was not to replace standard auto-diff. Rather, we wanted to demonstrate that using the TN formulation of convolutions yields white-box implementations, not only of many standard operations like vector-Jacobian products, but also less traditional operations used to approximate second-order information. We argue that this improves flexibility and performance:

• KFAC-reduce is a great example for a non-standard operation that is interesting for second-order methods, but whose implementation in existing ML frameworks suffers from limitations. We showed that a TN implementation's performance improves over the default implementation (New: We have added memory benchmarks to address a request from Reviewer D3NB and showed that our TN version also has consistently lower peak memory with savings up to 3 GiB).

• Our second main application, randomized autodiff, demonstrates the flexibility of a TN implementation: While ML frameworks support efficient multiplication with the Jacobian, randomizing the vector-Jacobian product is hard because its implementation is a black box. Our TN VJP is a white box that can be randomized more easily.

The applications we investigate in the main text are representatives of operations that are challenging to realize with existing frameworks. There are many others, e.g. approximating Hessian diagonals (Appendix B4), or diagonals and mini-block diagonals of GGN/Fisher matrices (Appendix B2 & C). Not only do we mention those in the introduction, but also do we provide their tensor diagrams in the appendix (see Appendix A and Table B2 for visual and pseudo-code overviews). By doing so, we make them accessible to the community, and allow other works to benefit from our proposed TN simplifications.

In summary, we believe that the focus of our work is different from Hayashi et al. in that it is concerned with investigating and improving computational performance, as well as significantly extending the amount of operations beyond the forward pass.

Support for different convolution types: We believe that our framework is fully-compatible with structured convolutions like those mentioned in [1]. You mentioned separable convolutions which we believe are referred to as 'convolutions with channel groups' in our paper. Our formulation supports channel groups, and therefore separable convolutions. To improve clarity, we have added a half sentence in the paper which mentions the alternative name.

Our derivations also carry over to other structured convolutions mentioned in Hayashi et al: From their Figure 1, we support

• depthwise separable
• bottleneck/Tucker-2
• inverted bottleneck
• flattened
• CP
• low-rank filter

The only difference is that the kernel used by those convolutions is factorized, i.e. a tensor network itself. We can simply substitute this tensor network into our derived expressions and proceed as usual, e.g. apply our symbolic simplifications. We tried to explain this compatibility in Section 6. Let us know if you are satisfied with this clarification. In case you would like us to experiment with other factorized convolutions, it would be great if you could narrow down the scope of experiments we could provide to convince you further.

Dear Reviewer Z3Mj,

thanks for your strong support! We agree that tensor networks have a great potential for machine learning and hope we can contribute to making them even more popular with this work.

Update on results: In response to Reviewer D3NB, we have extended the run time evaluation by a theoretical and empirical analysis of KFAC-reduce's memory consumption (added as Appendix G). We showed that our TN implementation is not only consistently faster (up to 4.5x), but also has much lower peak memory (with memory savings up to 3GiB for the CNNs we benchmarked in the main text).

Clarification on Section 2.2: You are right that the notation in Section 2.2 is complex and may be better explained by an example. For a matrix multiplication C = einsum("ij,jk->ik", A, B) you can think of the index tuples as S_A = ("i", "j"), S_B = ("j", "k") and S_C = ("i", "k"). The set operations in Equations 3, 4 serve to identify which indices are summed out, i.e. are not part of the output index tuple: (S_A ∪ S_B) \ S_C = ("j").

Our goal was to rigorously write down what einsum does and to convince the reader that drawing a tensor diagram is more intuitive and requires less cognitive load than using Equations 3, 4. We have added the above example to Appendix H2 and referenced it in the main text. Let us know if this improves your understanding of the presentation.

Please also let us know if you have any follow-up questions. We would be happy to discuss further.

Dear Reviewer XFQf,

Thank you for your feedback.

Contribution beyond new perspective: We would like to tackle your point that the paper's only contribution is presenting a new perspective. Specifically, we believe that our paper's experiments elaborate on concrete applications and advance the state of the art:

We can accelerate the pre-conditioner computation of a popular second-order method for deep learning by up to 4.5x. The run time improvements enable more frequent pre-conditioner updates which are so far not very popular due to the increased cost. As requested by Reviewer D3NB, we extended the run time results and measured the memory consumption, showing that our TN approach has consistently smaller peak memory with savings up to 3 GiB for KFAC-reduce. The authors of [1] mention KFAC-reduce's memory overhead caused by im2col as important limitation. By reducing the peak memory, our approach enables operating the algorithm at larger batch sizes. This is an important improvement because second-order methods seem to require larger batch sizes than first-order methods (see for instance [2, 3]).

Our faster and more memory-efficient KFAC implementation for convolutions is not only relevant to the optimization community. KFAC is also a popular curvature approximation for Bayesian neural networks to construct Laplace approximations [4].

Please let us know if you have any follow-up questions. We would be happy to discuss them!

References:

[1] Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., & Hennig, P. (NeurIPS 2023): Kronecker-factored approximate curvature for modern neural network architectures.

[2] Martens, J., & Grosse, R. (ICML 2015): Optimizing neural networks with Kronecker-factored approximate curvature.

[3] Grosse, R., & Martens, J. (ICML 2016). A kronecker-factored approximate Fisher matrix for convolution layers.

[4] Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., & Hennig, P. (NeurIPS 2021). Laplace redux - effortless bayesian deep learning.

Dear Reviewer D3NB,

thanks for your thorough review and detailed questions. We are happy to add more details to our answers below and discuss follow-up questions; please let us know if you have any.

Memory consumption (Q1): You have a point that we do not support our claim of reduced peak memory for KFAC-reduce with evidence. To alleviate this concern, we added a theoretical analysis and empirical evaluation in Appendix G:

For KFAC-reduce, the main difference between default and TN implementation is the computation of the averaged unfolded input $[[\mathbf{\mathsf{X}} ]]^{(\text{avg})} := \frac{1}{(O_1 O_2)} 1_{O_1 O_2}^{\top} [[ \mathbf{\mathsf{X}} ]]$ which consists of $C_{\text{in}} K_1 K_2$ numbers.

• The default implementation needs to build up the unfolded input. This requires extra storage of $C_{\text{in}} K_1 K_2 O_1 O_2$ numbers.

• Our proposed TN implementation uses the averaged index patterns, directly computed via a modification of Algorithm D1. They consist of $I_1 K_1 + I_2 K_2$ numbers. In contrast to the default implementation, spatial dimensions are de-coupled and there is no dependency on $C_{\text{in}}$.

• For structured convolutions, we can use our proposed simplifications to express the pattern tensor action through reshapes and narrows that often do not require additional memory. This completely eliminates the need to store additional tensors.

To empirically demonstrate this memory reduction inside the computation of the KFAC-reduce factor, we measured its peak memory and subtracted the memory of storing the input $\mathbf{\mathsf{X}}$ and the result $\hat{\mathbf{\Omega}}$. This serves as a proxy for the extra memory that is temporarily required. We consistently observe that the default implementation requires more extra memory than the TN implementation, whose demand is further reduced for structured convolutions if we enable our simplifications.

For example, our TN implementation uses 3 GiB less memory on ConvNeXt-base's features.1.0.block.0 convolutions. This is a considerable fraction of the 16-32 GiB available on most contemporary GPUs.

These results further substantiate our claims as they complement our demonstrated run time improvements. Our work leads to direct improvements of KFC in that it not only enables more frequent pre-conditioner updates (currently, updating every $10-100$ steps is popular), but also extends the regime of feasible batch sizes such an optimizer can operate in. This is important as second-order methods seem to require larger batch sizes than first-order methods (e.g. [1, 2]).

Applications other than KFC (Q2): We take your point that our performance evaluation focuses on KFC. This is because it makes a great example for a non-standard operation that is interesting for the design of new algorithms, but whose implementation in existing ML frameworks suffers from limitations. Our second application, randomized autodiff, is another---yet orthogonal---instance that demonstrates the flexibility of a TN implementation: While ML frameworks support efficient multiplication with the Jacobian, randomizing the vector-Jacobian product is hard because its implementation is a black box. By deriving the VJP as a TN, we provide a white box implementation which can be randomized more easily.

The applications we provide in the main text are representatives of operations that are challenging to realize with existing frameworks, e.g. approximating Hessian diagonals (Appendix B4), or diagonals and mini-block diagonals of GGN/Fisher matrices (Appendix B2 & C). Note that we not only mention them in the introduction, but also provide their tensor diagrams in the appendix (see Appendix A and Table B2 for visual and pseudo-code overviews). By doing so, we make them accessible to the community, and allow other works to benefit from our proposed TN simplifications. Please let us know if you would like us to provide more details on one of these operations in the appendix, or add the discussion of a new quantity that is currently not addressed.

Our faster and more memory-efficient KFAC implementation for convolutions is not only relevant for optimization. KFAC is also a popular curvature approximation for Bayesian applications with neural networks, such as constructing Laplace approximations (e.g. [3] and references within).

References:

[1] Martens, J., & Grosse, R. (ICML 2015): Optimizing neural networks with Kronecker-factored approximate curvature.

[2] Grosse, R., & Martens, J. (ICML 2016). A kronecker-factored approximate Fisher matrix for convolution layers.

[3] Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., & Hennig, P. (NeurIPS 2021). Laplace redux - effortless bayesian deep learning.