Dear Reviewer vbTD,

thanks for your support! We appreciate the time and effort you put into reviewing our work.

The authors show faster speed of KFAC computation. Does this implementation bring better convergence, training speed or results in practice?

You have a point that we did not evaluate the impact of our achieved speed-up on a real second-order method. We are happy to inform you that we will add such an experiment to the main text (please see our global rebuttal for a detailed description). In this experiment, we were able to halve the computational overhead of a KFAC-based method compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2) to a value very close to that of SGD.

A direct consequence of these findings is that they increase training speed, or allow using hardware with less memory.

One could also experiment with more frequent pre-conditioner updates or larger batch sizes to obtain convergence in fewer steps. We believe these are interesting directions to explore in the future that so far have been disregarded due to the high computational extra cost. To produce statistically significant statements and recommendations though, we believe it will take a larger ablation that is beyond the scope of this work.

Thanks again, and let us know if you have further questions.

Dear Reviewer cYbp,

thanks for your strong support and detailed review! We will apply your suggested improvements to the text.

We would like to inform you that we conducted a new experiment with a real second-order method based on KFAC-reduce (see our global rebuttal). Our TN implementation was able to halve the second-order method's run time overhead compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2). We hope this provides further evidence that our approach is indeed useful to advance second-order methods and close their computational gap.

Questions

• Why are the (non simplified) TN implementations slower in general? Do the standard implementations have further efficiency optimisations?

  This observation is correct. Non-simplified TNs perform contractions using a dense version of $\Pi$, which is sparse. This causes wasteful multiplies by zero. Our simplifications reformulate the contraction with $\Pi$ into reshapes or slices, which remove these zero multiplications from the TN (see e.g. the bottom right panel of Fig. 1) and thereby speed up the computation.

  Whenever simplification is not possible, the main difficulty for using the index pattern's sparsity is that PyTorch's einsum only accepts dense tensors. We experimented with alternatives that support mixed-format tensors, and have some ideas how to make progress on this frontier in future work. If you are interested in knowing more, feel free to have a look at our response to Reviewer iQrS.

• Does figure 6 include variability wrt contraction order?

  This depends on the heuristic used to find a contraction path. Generally speaking, the heuristic depends on the number of tensors in the TN. The tensor networks in Fig. 6 contain at most 6 tensors. After looking into opt_einsum's internals, we believe that the TNs with at most 4 tensors do not contain variability w.r.t. contraction path schedule, while the others might:

  In all our experiments, we used opt_einsum's contract_path method with the default 'auto' strategy to find contraction schedules. opt_einsum supports many strategies and according to the documentation, using 'auto' "will select the best of these it can while aiming to keep path finding times below around 1ms". The exact behavior depends on the number of tensors to be contracted and is described here: For TNs with up to 4 tensors, the optimal schedule is searched. For TNs with 5-6 tensors, a restricted search is carried out. We are unsure whether this restricted search is deterministic or stochastic, so it may be that this approach causes some extent of variability in contraction order.

• 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

  This is indeed an interesting question. We did not look into this and decided to rely completely on opt_einsum, which uses a sophisticated heuristic based on performance comparisons to strike a good balance between contraction path search time and contraction time. We believe this is a good representation of what a practitioner would do.

• How many of the operations support GPU/hardware acceleration? Would there be performance gains if they were implemented?

  All operations we propose are purely PyTorch and therefore support GPU execution out of the box.

Thanks again, and let us know if you have further questions.

Dear Reviewer iQrS,

thanks a lot for your thorough review; we appreciate the work you put into it and are glad you find the paper innovative and support our idea of making the complex yet powerful tensor network toolbox accessible to the ML community.

We would like to inform you that we conducted a new experiment with a real second-order method based on KFAC-reduce (see our global rebuttal). Our TN implementation was able to halve the second-order method's run time overhead compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2). We hope this provides further evidence that our approach is indeed useful to advance second-order methods and close their computational gap.

Weaknesses 1 & 2

Using einsum for convolution operations is generally not an optimal approach

The einsum-based approach favors specific architectures

You are right that there are additional degrees of freedom to further accelerate our TN implementation. To give some intuition, we would like to point out that one main problem is that we cannot always leverage the index pattern's sparsity: For many convolutions, we can exploit the sparsity thanks to our proposed symbolic simplifications. However, whenever they cannot be applied, we must consider $\Pi$ as dense because PyTorch's einsum does not allow contracting tensors with mixed formats. This means we must sometimes carry out multiplies by zero, which is wasteful.

We attempted to address the scenario where our simplifications cannot be applied and have ideas how to progress on this frontier. We will describe them in an updated version of the paper as follows:

• We experimented with TACO [2], a tensor algebra compiler that can optimize contractions of mixed-format tensors. Unfortunately, TACO's Python frontend cannot yet generate GPU code, and on CPU we ran into a memory leak in the Python frontend.

• It would be interesting to benchmark our approach on specialized hardware, e.g. cerebras' wafer-scale engine processor, which supports unstructured sparsity. This would allow to remove many multiplications by zero. While we do not have access to such a chip, we are excited our approach could benefit from such hardware in the future.

• To the best of our knowledge, we believe that contraction path optimization of mixed-format tensor networks has not been studied sufficiently. This may be due to the absence of hardware that can leverage this information to speed up contraction. We hypothesize further performance improvements might be possible by taking the operand sparsity into account when optimizing the contraction path.

Please let us know if you think there is an additional approach we could try to boost the performance of our framework.

Weakness 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.

You are right that we do not mention which convolution implementation is used in the comparison. We use PyTorch's torch.nn.functional.conv2d, which automatically selects a 'good' implementation according to heuristics. We think this is a good baseline implementation as most engineers and researchers rely on it in practise.

Furthermore, the paper does not compare the TN approach with other highly optimized versions of convolution

Consequently you do have a point that we did not look into comparing our approach with different convolution implementations. We found that PyTorch indeed has different convolution algorithms implemented. However, it seems like one cannot specify manually which algorithm should be used.

We will try to look into this and identify other convolution implementations in PyTorch in order to provide a more fine-grained comparison. If you have any pointers, please let us know.

References

[2] Kjolstad, F., Chou, S., Lugato, D., Kamil, S., & Amarasinghe, S. (2017). TACO: A tool to generate tensor algebra kernels. IEEE/ACM International Conference on Automated Software Engineering (ASE).