Paper: Convolutions Through the Lens of Tensor Networks (ICLR-2024)

Reviewer comment (Novelty):
From current draft, I am a bit confused on the difference between this work and (Hayashi et al. 2019). I suppose this paper proposes to compute some first-order and second-order information from tensor network as well. Then what is the difference between using standard auto-differentiation packages and the proposed method based on tensor network? It seems that in experiments, the authors only compare the proposed method with standard PyTorch implementations, but not tensor network combined with auto-differentiation packages. Some explanations may be needed here.

Author response:
**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.

Reviewer comment (Reproducibility):
I also wonder how is the index tensor \Pi stored in real applications. Since it should be a very sparse tensor, do we have to use some sparse formats? How will it affect the computation time? The authors may need to add more details here.

Author response:
**Leveraging sparsity of the index pattern:** You are right that the index pattern tensor is sparse. Right now, we store this tensor in dense (boolean) format because (1) it is relatively small and (2) currently we cannot leverage its sparsity to speed up the contraction:

(1) Take for example the first layer of ConvNeXt-base which has large input sizes ($I_{1,2} = 256$), output sizes ($O_{1,2} = 64$), and kernel sizes ($K_{1,2} = 4$). Storing that layer's index pattern as dense boolean tensor requires $256 \cdot 64 \cdot 4 = 65,536$ bits, or 8 kiB.

(2) Ideally, we would like to not only store the index pattern in sparse format, but also leverage its sparsity in the contraction with other tensors. However, the `einsum` implementation in PyTorch that we rely on requires all tensors to be dense (to the best of our knowledge, this limitation applies to all other popular ML frameworks).

We tried experimenting with TACO [2], a tensor algebra compiler which is capable to compile tensor expressions involving sparse and dense tensors. The Python front-end only supports CPU code generation and we eventually had to discard it because we kept encountering memory leaks that seem to be caused by its Python front-end.

For convolutions with special structure, our symbolic simplifications from Section 4 leverage the sparsity of $\mathbf{\mathsf{\Pi}}$ in that they re-express its contraction through cheap operations such as `reshape`s and `narrow`s. In the general case, however, we are unaware of a framework that can leverage the sparsity. Such a framework would certainly further improve our approach.

Reviewer comment (Experiments):
I also note that most experiments are performed with simple convolution operations, while there are also many different types of convolutions (e.g., separate convolution). It would be better if the authors can also compare with these operations to demonstrate the flexibility of tensor network.

Author response:
**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.

Reviewer comment (Presentation):
Minor: formatting issues. Some captions in the appendix seem to be overlap with the page head.

Author response:
**Minor:** We have uploaded a version that fixes the formatting issues of the tables in the appendix. Thanks for pointing them out.

Reviewer comment (Writing):
In Section 2.2, I cannot fully follow how to use the set operation with the index tuples to model the tensor contractions. Could you give a more intuitive explanation or examples?

Author response:
**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.

Reviewer comment (Novelty):
The main concern is contribution. The authors point out the advantages of this perspective rather than developing a framework in a novel way. From this point of view, the contribution seems limited.

Author response:
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.

Reviewer comment (Novelty):
The authors' central contribution only lies in some derivation based on this perspective, such as automatic differentiation.

Author response:
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.

Reviewer comment (Novelty):
I'm not sure whether the proposal of perspective is a contribution and means much to the community or not. It would be helpful to provide some explanation about this point.

Author response:
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.

Reviewer comment (Presentation):
The advantages of using tensor networks to represent CNNs are not thoroughly discussed in this article. The paper primarily focuses on the advantages in the context of KFC, leading me to believe that it is primarily aimed at addressing memory consumption issues within KFC. Therefore, it might be more appropriate to modify the paper's topic and title to "Accelerating KFC with Tensor Network (TN) Methods.”

Author response:
**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).

Reviewer comment (Theory):
The paper provides a comprehensive guide on how to use Tensor Networks to represent CNNs, and offers detailed operations for various CNNs. However, in terms of the advantages of using Tensor Networks to represent CNNs, the paper lacks further analysis and discussion beyond a brief analysis in the context of KFC. For instance, once CNNs are represented in the form of TN, could this representation also be benefit to other second-order analysis and optimization methods, such as the Approximate Hessian diagonal, KBFGS, and Hessian rank mentioned in the Introduction? If this is possible, I would prefer to see the authors provide a more in-depth discussion.

Author response:
**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).