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

Authors: felix  dangel
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:18:50

Abstract: 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.5 x while removing the standard implementation's memory overhead, and enables new hardware-efficient tensor dropout for approximate backpropagation.

Main paper: 
 introduction: 
 Convolutional neural networks [CNNs, 39] mark a milestone in the development of deep learning architectures as their 'sliding window' approach represents an important inductive bias for vision tasks. Their intuition is simple to explain with graphical illustrations [e.g. 21 ]. Yet, convolutions are more challenging to analyze than dense layers in multi-layer perceptrons (MLPs) or transformers [71] . One reason is that they are hard to express in matrix notation and-even in index notation-compact expressions that are convenient to work with only exist for special hyper-parameters [e.g. 27, 2] . Many hyper-parameters (stride, padding, . . . ) and additional features like channel groups [36] introduce even more complexity that is inherited by related routines, e.g. for autodiff. We observe a delay of analytic and algorithmic developments between MLPs vs. CNNs, e.g. The software support for less standard routines some of these methods require also reflects this gap. Some functions only support special dimensions [15] . Others use less efficient workarounds ( §5.1) or are not provided at all ( §B.4). And they are hard to modify as the code is either closed-source [12] or 38th Conference on Neural Information Processing Systems (NeurIPS 2024).from einconv import c on v _i nd e x_ p at t er n ; from einops import einsum ; import torch I , K , D , P , S = ... # convolution hyper -parameters (2 -tuples ) X = torch . rand ( C_in , I [0] , I [1] ) # input to convolution kfc_shape = ( C_in * K [0] * K [1] , C_in * K [0] * K [1] ) # final shape def kfc_im2col () :""" Compute Kronecker factor via im2col . """ X_unf = torch . nn . functional . unfold (X , K , D , P , S ) return einsum ( X_unf , X_unf , " i out , j out -> i j " ) def kfc_tn () : """ Compute Kronecker factor via its tensor network . """ Pi1 = co n v_ i nd e x_ p at te r n ( I [0] , K [0] , S [0] , P [0] , D [0]) Pi2 = co n v_ i nd e x_ p at te r n ( I [1] , K [1] , S [1] , P [1] , D [1] ) return einsum (X , Pi1 , Pi2 , X , Pi1 , Pi2 , " c_in i1 i2 , i1 o1 k1 , i2 o2 k2 , c_in_ i1_ i2_ , i1_ o1_ k1_ i2_ o2_ k2_ " + " -> c_in k1 k2 c_in_ k1_ k2_ " ) . reshape ( kfc_shape ) def kfc _s imp lif ied _tn () : # for dense convolutions """ Compute Kronecker factor via its simplified tensor network . """ X = X . reshape ( C_in , I [0] // K [0] , K [0] , I [1] // K [1] , K [1] ) return einsum (X , X , " c_in o1 k1 o2 k2 , c_in_ o1_ k1_ o2_ k2_ " + " -> c_in k1 k2 c_in_ k1_ k2_ " ) . reshape ( kfc_shape )Figure 1 : Many convolution-related routines can be expressed as TNs and evaluated with einsum.We illustrate this for the input-based factor of KFAC for convolutions [KFC, 27] , whose standard implementation (top) requires unfolding the input (high memory). The TN (middle) enables internal optimizations inside einsum (e.g. with contraction path optimizers like opt_einsum [66] ). (Bottom) In many cases, the TN further simplifies due to structures in the index pattern, which reduces cost.X X o1 o2 cin k1 k2 c ′ in k ′ 1 k ′ 2 Π (1)Π (1) Π (2) Π (2) X XEQUATIONΠ (2) Π (2) X Xi1 i2 i ′ 2 i ′ 1 o1 o2 k1 k ′ 1 k2 k ′ 2 cin c ′in written in a low-level language. This complicates the advance of existing, and the exploration of new, algorithmic ideas for convolutions.Here, we seek to reduce this complexity gap by viewing convolutions as tensor networks [TNs, 53, 6, 9] which express the underlying tensor multiplications as diagrams. These diagrams are simpler to parse than mathematical equations and can seamlessly be (i) manipulated to take derivatives, add batching, or extract sub-tensors, (ii) merged with other diagrams, and (iii) evaluated with einsum. This yields simple, modifiable implementations that benefit from automated under-the-hoodoptimizations for efficient TN contraction developed by the quantum simulation community [e.g. 66, 25, 74, 13] , like finding a high-quality contraction order or distributing computations:1. We use the TN format of convolution from Hayashi et al. [29] to derive diagrams and einsum formulas for autodiff and less standard routines for curvature approximations with support for all hyper-parameters, batching, groups, and any dimension (Table 1 ).2. We present transformations based on the convolution's connectivity pattern to re-wire and symbolically simplify TNs before evaluation (example in Figure 1 ).3. We compare default and TN implementations, demonstrating optimal peak memory reduction and run time improvements up to 4.5 x for a recent KFAC variant, and showcase their flexibility to impose hardware-efficient dropout for randomized backpropagation.Our work not only provides simpler perspectives and implementations that facilitate the exploration of algorithmic ideas for convolutions, but also directly advances second-order methods like KFAC: It enables more frequent pre-conditioner updates, using larger batches without going out of memory, and extending KFAC to transpose convolution. These improvements are important for second-order optimization and other applications like Laplace approximations [20] and influence functions [28] . 
preliminaries: 
 We briefly review 2d convolution ( §2.1), tensor multiplication and einsum ( §2.2), then introduce the graphical TN notation and apply it to convolution ( §2.3). Bold lower-case (a), upper-case (A), and upper-case sans-serif (A) symbols indicate vectors, matrices, and tensors. Entries follow the same convention but use regular font weight; [•] denotes slicing ( [A] i,j = A i,j ). Parenthesized indices mean reshapes, e.g. [a] (i,j) = [A] i,j with a the flattened matrix A.2Π (1)Π (2) X Wi 1 i 2 k 1 k 2 c in c out o 1 o 2 Y = o 1 o 2 c out (a) Convolution Π (1) Π (2) X i 1 i 2 (o 1 , o 2 ) (c in , k 1 , k 2 ) c in k 1 k 2 o 1 o 2 X (c in , k 1 , k 2 ) (o 1 , o 2 ) = (b) Input unfolding Π (1) Π (2) W k 1 k 2 (c out , o 1 , o 2 ) (c in , i 1 , i 2 ) c in c out i 1 i 2 o 1 o 2 A(W) (c in , i 1 , i 2 ) (c out , o 1 , o 2 ) = (c) Kernel unfoldingFigure 2 : TNs of (a) 2d convolution and (b,c) connections to its matrix multiplication view. The connectivity along each dimension is explicit via an index pattern tensor Π.2d convolutions process channels of 2d signals X ∈ R Cin×I1×I2 with C in channels of spatial dimensions1 I 1 , I 2 by sliding a collection of C out filter banks, arranged in a kernel W ∈ R Cout×Cin×K1×K2 with kernel size K 1 , K 2 , over the input. The sliding operation depends on various hyperparameters [padding, stride, . . . , see 21] . At each step, the filters are contracted with the overlapping area, yielding the channel values of a pixel in the output Y ∈ R Cout×O1×O2 with spatial dimensions O 1 , O 2 . Optionally, a bias from b ∈ R Cout is added per channel.One way to implement convolution is via matrix multiplication [10] , similar to fully-connected layers. First, one extracts the overlapping patches from the input for each output, then flattens and column-stacks them into a matrix X ∈ R CinK1K2×O1O2 , called the unfolded input (or im2col). Multiplying a matrix view W ∈ R Cout×CinK1K2 of the kernel onto the unfolded input then yields a matrix view Y of Y (the vector of ones, 1 O1O2 , copies the bias for each channel),EQUATIONWe can also view convolution as an affine map of the flattened input x ∈ R CinI1I2 into a vector view y of Y with a Toeplitz-structured matrix A(W) ∈ R CoutO1O2×CinI1I2 ,EQUATION)This perspective is uncommon in code, but used in theoretical works [e.g. 65] as it highlights the similarity between convolutions and dense layers.Tensor multiplication unifies outer (Kronecker), element-wise (Hadamard), and inner products and uses the input-output index relation to infer the multiplication type. We start with the binary case, then generalize to more inputs: Consider A, B, C whose index names are described by the index tuples S 1 , S 2 , S 3 where S 3 ⊆ (S 1 ∪ S 2 ) (converting tuples to sets if needed). Any product of A and B can be described by the multiplication operator * (S1,S2,S3) withEQUATIONsumming over indices that are not present in the output. E.g., for two matrices A, B, their product is AB = * ((i,j),(j,k),(i,k)) (A, B) (see §H.2), their Hadamard product A⊙B = * ((i,j),(i,j),(i,j)) (A, B), and their Kronecker product A ⊗ B = * ((i,j),(k,l),((i,k),(j,l))) (A, B). Libraries support this functionality via einsum, which takes a string encoding of S 1 , S 2 , S 3 , followed by A, B. It also accepts longer sequences A 1 , . . . , A N with index tuples S 1 , S 2 , . . . , S N and output index tuple S N +1 ,EQUATIONTable 1 : Contraction expressions of operations related to 2d convolution. They include batching and channel groups, which are standard features in implementations. We describe each operation by a tuple of input tensors and a contraction string that uses the einops library's syntax [59] which can express index (un-)grouping. Some quantities are only correct up to a scalar factor which is suppressed for brevity. See §B for visualizations and Table B3 for more operations.A simpler way to understand tensor multiplications is via diagrams developed by e.g. Penrose [53] .Rank-K tensors are represented by nodes with K legs labelled by the index's name 2 . a i denotes a vector a, B . Element-wise multiplication is similar, but with a leg sticking out. The Hadamard and Kronecker product diagrams areEQUATIONNote that the outer tensor product is a rank-4 tensor and must be reshaped (indicated by black triangles 3 ) into a matrix. This syntax allows for extracting and embedding tensors along diagonals; e.g. taking a matrix diagonal, diag(A) i = A i , or forming a diagonal matrix, diag(a) i i = a i i ; and generalizes to larger diagonal blocks ( §B). In the following, we stick to the simplest case to avoid the more advanced syntax. However, it shows the expressive power of TNs and is required to support common features of convolutions like channel groups (known as separable convolutions).Application to convolution: We define a binary tensor P ∈ {0, 1} I1×O1×K1×I2×O2×K2 which represents the connectivity pattern between input, output, and kernel. P i1,o1,k1,i2,o2,k2 is 1 if input locations (i 1 , i 2 ) overlap with kernel positions (k 1 , k 2 ) when computing output locations (o 1 , o 2 ) and 0 otherwise. The spatial couplings are independent along each dimension, hence P decomposes into P i1,o1,k1,i2,o2,k2 = Π Π (1) Π (2) Wk 1 k 2 c ′ in c out i ′ 1 i ′ 2 o 1 o 2 (b) Input Jac.Π (1) Π (2) X V (Y) i 1 i 2 c ′ in k ′ 1 k ′ 2 o 1 o 2 c ′ out (c) Weight VJP Π (1) Π (2) V (Y) W k 1 k 2 c ′ in i ′ 1 i ′ 2 o 1Although a transpose convolution is defined w.r.t. a reference convolution with hyper-parameters K, S, P, D, most libraries offer standalone implementations of transpose convolution. We describe the transpose convolution by its associated convolution, that is as a mapping from R Cout×O1×O2 (the convolution's output space) to R Cin×I1×I2 (the convolution's input space). For convolution with S > 1, we cannot infer I from O, K, S, P, D, as multiple Is map to the same O if (I + 2P -K -(K -1)(D -1)) mod S ̸ = 0 (see the floor operation in Algorithm D1). We need to either supply I directly, or the remainderA = I + 2P -K -(K -1)(D -1) -S(O -1)(often called output_padding) to make I unambiguous. Then, we computeI = (O -1)S -2P + K + (K -1)(D -1) + A . (D18)to get I(O, A) and call Algorithm D1 to obtain Π(I(O, A), K, S, P, D). 
tns for convolution operations: 
 We now demonstrate the elegance of TNs for computing derivatives ( §3.1), autodiff operations ( §3.2), and approximate second-order information ( §3.3) by graphical manipulation. For simplicity, we exclude batching (vmap-ing like in JAX [8] ) and channel groups, and provide the diagrams with full support in §B. Table 1 summarizes our derivations (with batching and groups). As a warm-up, we identify the unfolded input and kernel from the matrix-multiplication view (Equations ( 1) and ( 2)). They follow by contracting the index patterns with either the input or kernel (Figures 2b and 2c ),[ X ] (cin,k1,k2),(o1,o1) = i1,i2X cin,i1,i2 Πi2,o2,k2 ,EQUATIONi2,o2,k2 W cout,cin,k1,k2 .Derivatives play a crucial role in theoretical and practical ML. First, we show that differentiating a TN diagram amounts to a simple graphical manipulation. Then, we derive the Jacobians of convolution.Consider an arbitrary TN represented by the tensor multiplication from Equation (4) . The Jacobian tensor [J Aj A N +1 ] S N +1 ,S ′ j = ∂[A N +1 ] S N +1/∂[Aj ] S ′ j w.r.t. an input A j collects all partial derivatives and is addressed through indices S n+1 × S ′ j with S ′ j an independent copy of S j . Assume that A j only enters once in the tensor multiplication. Then, taking the derivative of Equation (4) w.r.t. [A j ] S ′ j simply replaces the tensor by a Kronecker delta δ Sj ,S ′ j ,EQUATIONIf an index i ∈ S j is summed, i ̸ ∈ S n+1 , we can sum the Kronecker delta δ i,i ′ , effectively replacing all occurrences of i by i ′ . If instead i is part of the output index, i ∈ S n+1 , the Kronecker delta remains part of the Jacobian and imposes structure. Figure 3a illustrates this process in diagrams for differentiating a convolution w.r.t. its kernel. Equation ( 6) amounts to cutting out the argument of differentiation and assigning new indices to the resulting open legs. For the weight Jacobian J W Y, this introduces structure: If we re-interpret the two sub-diagrams in Figure 3a as matrices, compare with the Kronecker diagram from Equation (5) and use Figure 2b , we find X ⊤ ⊗ I Cout for the Jacobian's matrix view [e.g. 16]. Figure 3b shows the input Jacobian J X Y which is a tensor view of A(W), as expected from the matrix-vector perspective of Equation (2) .Differentiating a TN is more convenient than using matrix calculus [44] as it amounts to a simple graphical manipulation, does not rely on a flattening convention, and therefore preserves the full index structure. The resulting TN can still be translated back to matrix language, if desired. It also simplifies the computation of higher-order derivatives (e.g. ∂ 2 Y /∂W∂X), since differentiation yields another TN and can thus be repeated. If a tensor occurs more than once in a TN, the product rule applies and the derivative is a sum of TNs with one occurrence removed.ΠΠ (1)Π (2) Π (2) X X i 1 i 2 i ′ 2 i ′ 1 o 1 o 2 k 1 k ′ 1 k 2 k ′ 2 c in c ′ in (a) KFAC-expand 1 1 Π (1) Π (1)1 1 The unfolded input is shaded, only additional contractions are highlighted. (a) Ω (KFC/KFAC-expand) from Grosse & Martens [27] , (b) Ω (KFAC-reduce) from Eschenhagen et al. [23] (vectors of ones effectively amount to sums).Π (2) Π (2) X X i1 i2 i ′ 2 i ′ 1 o1 o2 o ′ 1 o ′ 2 k1 k ′ 1 k2 k ′ 2 cin c ′ in (b) KFAC-reduceAlthough Jacobians are useful, crucial routines for autodiff are vector-Jacobian and Jacobian-vector products (VJPs, JVPs). Both are simple to realize with TNs due to access to full Jacobians. VJPs are used in backpropagation to pull back a tensor V (Y) ∈ R Cout×O1×O2 from the output to the input or weight space. The VJP results V (X) ∈ R Cin×I1×I2 and V (W) ∈ R Cout×Cin×K1×K2 are V (X) c ′ in ,i ′ 1 ,i ′ Π(I, K, S, P, D) = Π(I, K, D, P, S) i o k i k o (a) Kernel-output swap Π(I, K, K, 0, 1) = Π i o k i o k (b) Dense convolution V Π(I, K, S > K, 0, 1) = Ṽ Π( IK /S, K, K, 0, 1) i i ′ o k o k (c) Down-sampling convolution 
tn simplifications & implementation: 
 Many convolutions in real-world CNNs use structured connectivity patterns that allow for simplifications which we describe here along with implementation aspects.The index pattern Π encodes the connectivity of a convolution and depends on its hyper-parameters.Along one dimension, Π = Π(I, K, S, P, D) with input size I, kernel size K, stride S, padding P , and dilation D. We provide pseudo-code for computing Π in §D which is easy to implement efficiently with standard functions from any numerical library (Algorithm D1). Its entries areEQUATIONwith i = 1, . . . , I, o = 1, . . . , O, k = 1, . . . , K and output size O(I, K, S, P, D) = 1 + ⌊ (I+2P -(K+(K-1)(D-1))) /S⌋. Since Π is binary and has size linear in I, O, K, it is cheap to precompute and cache. The index pattern's symmetries allow for re-wiring a TN. For instance, the symmetry of (k, D) and (o, S) in Equation (7) and O(I, K, S, P, D) permits a kernel-output swap, exchanging the role of kernel and output dimension (Figure 5a ). Rochette et al. [58] used this to phrase the per-example gradient computation (Figure 3c ) as convolution.For many convolutions of real-world CNNs (see §E for a hyper-parameter study) the index pattern possesses structure that simplifies its contraction with other tensors into either smaller contractions or reshapes: Dense convolutions use a shared kernel size and stride, and thus process non-overlapping adjacent tiles of the input. Their index pattern's action can be expressed as a cheap reshape (Figure 5b ). Such convolutions are common in DenseNets [33] , MobileNets [31, 60] , ResNets [30] ,and ConvNeXts [42] . InceptionV3 [69] has 2d mixed-dense convolutions that are dense along one dimension. Down-sampling convolutions use a larger stride than kernel size, hence only process a sub-set of their input, and are used in ResNet18 [30] , ResNext101 [72] , and WideResNet101 [73] .Their pattern contracts with a tensor V like that of a dense convolution with a sub-tensor Ṽ (Figure 5c ). §5.1 shows that those simplifications accelerate computations.Contraction order optimization: There exist various orders in which to carry out the summations in a TN and their performance can vary by orders of magnitude. One extreme approach is to carry out all summations via nested for-loops. This so-called Feynman path integral algorithm requires little memory, but many FLOPS since it does not re-cycle intermediate results. The other extreme is sequential pair-wise contraction. This builds up intermediate results and can greatly reduce FLOPS. The schedule is represented by a binary tree, but the underlying search is in general at least #Phard [14] . Fortunately, there exist heuristics to find high-quality contraction trees for TNs with hundreds of tensors [32, 25, 13] , implemented in packages like opt_einsum [66] .Index slicing: A common problem with high-quality schedules is that intermediates exceed memory. Dynamic slicing [32] (e.g. cotengra [25] ) is a simple method to decompose a contraction until it becomes feasible by breaking it up into smaller identical sub-tasks whose aggregation adds a small overhead. This enables peak memory reduction and distribution.Sparsity: Π is sparse as only a small fraction of the input contributes to an output element. For a convolution with stride S < K and default parameters (P = 0, D = 1), for fixed output and kernel indices k, o, there is exactly one non-zero entry in [Π] :,o,k . Hence nnz(Π) = OK, which corresponds to a sparsity of 1 /I. Padding leads to more kernel elements that do not contribute to an Forward Input VJP Weight VJP KFC KFAC-red. output pixel, and therefore a sparser Π. For down-sampling and dense convolutions, we showed how Π's algebraic structure allows to simplify its contraction. However, if that is not possible, Π contains explicit zeros that add unnecessary FLOPS. One way to circumvent this is to match a TN with that of an operation with efficient implementation (like im2col, (transpose) convolution) using transformations like the kernel-output swap or by introducing identity tensors to complete a template, as done in Rochette et al. [58] , Dangel [15] for per-sample gradients and im2col.Approximate contraction & structured dropout: TNs offer a principled approach for stochastic approximation via Monte-Carlo estimation to save memory and run time at the cost of accuracy. The basic idea is best explained on a matrix productC := AB = N n=1 [A] :,n [B] n,: with A ∈ R I×N , B ∈ R N,O .To approximate the sum, we introduce a distribution over n's range, then use column-row-sampling [CRS, 1] to form an unbiased Monte-Carlo approximation with sampled indices, which only requires the sub-matrices with active column-row pairs. Bernoulli-CRS samples without replacement by assigning a Bernoulli random variable Bernoulli(π n ) with probability π n for column-row pair n to be included in the contraction. The Bernoulli estimator isC := N n=1 zn /πn [A] n,: [B] n,: with z n ∼ Bernoulli(π n ).With a shared keep probability, π n := p, this yields the unbiased estimator C ′ = 1 /p n=1,...,N A ′ B ′ where A ′ = AK and B ′ = KB with K = diag(z 1 , . . . , z N ) are the sub-matrices of A, B containing the active column-row pairs. CRS applies to a single contraction. For TNs with multiple sums, we can apply it individually, and also impose a distribution over the result indices, which computes a (scaled) sub-tensor. 
experiments: 
 Here, we demonstrate computational benefits of TNs for less standard routines of second-order methods and showcase their flexibility to perform stochastic autodiff in novel ways.We implement the presented TNs' contraction strings and operands 5 in PyTorch [52] . The simplifications from §4 can be applied on top and yield a modified einsum expression. To find a contraction schedule, we use opt_einsum [66] with default settings. We extract the unique convolutions of 9 architectures for ImageNet and smaller data sets, then compare some operations from Table 1 with their standard implementation on an Nvidia Tesla T4 GPU (16 GB); see §F for all details. Due to space constraints, we highlight important insights here and provide references to the corresponding material in the appendix. In general, the performance gap between standard and TN implementation decreases the less common an operation is (Figure F17 ); from forward pass (inference & training), to VJPs (training), to KFAC (training with a second-order method). This is intuitive as more frequently used routines have been optimized more aggressively.Impact of simplifications: While general convolutions remain unaffected (Figure F18d ) when applying the transformations of §4, mixed dense, dense, and down-sampling convolutions consistently enjoy significant run time improvements (Figures F18a to F18c). As an example, we show the performance comparison for dense convolutions in Figure 6 : The performance ratio's median between TN and standard forward and input VJP is close to 1, that is both require almost the same time. In the median, the TN even outperforms PyTorch's highly optimized weight VJP, also for down-sampling Figure 7 : Extra memory used by the standard versus our TN implementation (simplifications enabled) of KFAC-reduce. Each point represents a convolution from 9 CNNs, clipped below by 1 MiB. TNs consistently use less memory than the standard implementation (one exception), and often no extra memory at all. We observe memory savings up to 3 GiB. convolutions (Figure F21 ). For KFC, the median performance ratios are well below 1 for dense, mixed dense & sub-sampling convolutions (Figure F22 ).For all convolution types, the TN implementation achieves its largest improvements for Ω and consistently outperforms the PyTorch implementation in the median when simplifications are enabled (Figure F23 ). The standard implementation unfolds the input, takes the row-average, then forms its outer product. The TN does not need to expand X in memory and instead averages the index pattern tensors, which reduces peak memory and run time. We observe performance ratios down to 0.22 x (speed-ups up to ≈ 4.5 x, Table F9 ) and consistently lower memory consumption with savings up to 3 GiB (Figure 7 ). Hence, our approach not only significantly reduces the overhead of 2nd-order optimizers based on KFAC-reduce, but also allows them to operate on larger batches without exceeding memory (Eschenhagen et al. [23] specifically mention memory as important limitation of their method). Other examples for KFAC algorithms where computing the input-based Kronecker factor adds significant time and memory overhead are that of Petersen et al. [54] , Benzing [5] which only use Ω (setting Γ ∝ I), or Lin et al. [41, 40] which remove matrix inversion.To demonstrate the speed-ups of KFAC-reduce in practical algorithms, we apply our work to the SINGD optimizer [41] and benchmark the impact of our TN implementation on its memory and run time in comparison to SGD without momentum. Concretely, we investigate SINGD with KFAC-reduce and diagonal pre-conditioners on ResNet18 and VGG19 on ImageNet-like synthetic data (3, 256, 256) using a batch size of 128. We measured per-iteration time and peak memory on an NVIDIA A40 with 48 GiB of RAM. For SINGD, we compare computing the Kronecker factors with the standard approach ('SINGD') via input unfolding versus our TN implementation ('SINGD+TN'). Table 2 summarizes the results. On both nets, our TN implementation halves SINGD's run time, and almost completely eliminates the memory, overhead compared to SGD. On VGG19, it dramatically lowers the memory overhead, cutting it down by a factor of 2 from 32 GiB to 16 GiB. This enables using larger batches or more frequently updating the pre-conditioner, underlining the utility of our approach for reducing the computational gap between approximate second-order and first-order methods.CRS is an alternative to checkpointing [26] to lower memory consumption of backpropagation [50, 11, 1] . Here, we focus on unbiased gradient approximations by applying the exact forward pass, but CRS when computing the weight VJP, which requires storing a sub-tensor of X. For convolutions, the approaches of existing works are limited by the supported functionality of ML libraries. Adelman et al. [1] restrict to sampling X along c in , which eliminates many gradient entries as the index is part of the gradient. The randomized gradient would thus only train a sub-tensor of the kernel per step. Oktay et al. [50] , Chen et al. [11] apply unstructured dropout to X, store it in sparse form, and restore the sparsified tensor during the backward pass. This reduces memory, but not computation.Our TN implementation is more flexible and can, for example, tackle spatial dimensions with CRS. This reduces memory to the same extent, but also run time due to fewer contractions. Importantly, it Figure 8 : Sampling spatial axes is more effective than channels on both (a) real-world and (b) synthetic data. We take the untrained All-CNN-C [68] for CIFAR-100 with cross-entropy loss, disable dropout, and modify the convolutions to use a fraction p of X when computing the weight gradient via Bernoulli-CRS. For mini-batches of size 128, we compute the deterministic gradients for all kernels, then flatten and concatenate them into a vector g; likewise for its proxy ĝ. CRS is described by (p cin , p i1 , p i2 ), the keep rates along the channel and spatial dimensions. We compare channel and spatial sampling with same memory reduction, i.e. (p, 1, 1) and (1, √ p, √ p). To measure approximation quality, we use the normalized residual norm ∥g-ĝ∥ 2/∥g∥ 2 and report mean and standard deviation of 10 different model and batch initializations.does not zero out the gradient for entire filters. In Figure 8 we compare the gradient approximation errors of channel and spatial sub-sampling. For the same memory reduction, spatial sub-sampling yields a smaller approximation error on both real & synthetic data. E.g., instead of keeping 75 % of channels, we achieve the same approximation quality using only 35 % of pixels. 
related work: 
 Structured convolutions: We use the TN formulation of convolution from Hayashi et al. [29] who focus on connecting kernel factorizations to existing (depth-wise separable [31, 60] , factored [69] , bottleneck [30] , flattened/CP decomposed, low-rank filter [67, 57, 70] ) convolutions and explore new factorizations. Our work focuses on operations related to convolutions, diagram manipulations, the index pattern structure, and computational performance/flexibility. Structured convolutions integrate seamlessly with our framework by replacing the kernel with its factorized TN.Higher-order autodiff: ML frameworks focus on differentiating scalar-valued objectives once.Recent works [37, 38, 43] developed a tensor calculus to compute higher-order derivatives of tensorvalued functions and compiler optimizations through linear algebra and common sub-expression elimination. Phrasing convolution as einsum, we allow it to be integrated into such frameworks, benefit from their optimizations, and complement them with our convolution-specific simplifications. 
conclusion: 
 We used tensor networks (TNs), a diagrammatic representation of tensor multiplications, to simplify convolutions and many related operations. We derived the diagrams of autodiff and less standard routines for curvature approximations like KFAC with support for all hyper-parameters, channel groups, batching, and generalization to arbitrary dimensions. All amount to simple einsum expressions that can easily be modified-e.g. to perform stochastic backpropagation-and benefit from under-the-hood optimizations before evaluation. We complemented those by convolution-specific symbolic simplifications based on structure in the connectivity pattern and showed their effectiveness to advance second-order methods. Our TN implementation accelerates the computation of KFAC up to 4.5 x and uses significantly less memory. Beyond performance improvements, the simplifying perspective also allowed us to formulate KFAC for transpose convolution. More broadly, our work underlines the elegance of TNs for reasoning about tensor multiplications and function transformations (differentiation, batching, slicing, simplification) in terms of diagrams at less cognitive load without sacrificing rigour. We believe they are a powerful tool for the ML community that will open up new algorithmic possibilities due to their simplicity & flexibility.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: All results were obtained on a single GPU to be comparable in terms of run time. See §F for the details.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: This work aims to provide a simplifying perspective and implementations of otherwise hard-to-access operations for convolutions to facilitate the exploration of algorithmic ideas and advance existing second-order methods. We don't see any direct negative societal impacts.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Here we comment on limitations on our approach.No common sub-expression elimination (CSE): Our implementation relies on opt_einsum which focuses on contraction order optimization. This optimization is efficient when all operands are different. However, with multiple occurrences of operands, computing shared sub-expressions might be an advantageous optimization approach which opt_einsum does not account for. The second-order quantity TNs from §C and §3.3 contain such sub-expressions, for instance X and 1 ⊤ O1O2 X in KFAC-expand and KFAC-reduce, and S (W) in the GGN quantities from Figure C16 . The efficiency of CSE depends on how costly the shared tensor is to compute. For instance, computing S (W) is expensive and therefore CSE is the more suitable optimization technique. For the input-based Kronecker factors which require the unfolded input, either contraction path optimization or CSE might be better. This is because the optimal contraction order may not correspond to 2x input unfolding and exhibit more parallelism which may lead to faster run times on a GPU. It would be interesting to integrate CSE into the contraction path optimization and develop a heuristic to choose a contraction path, for instance based on a weighted sum of FLOPs and memory.No index slicing: We mention index slicing as a technique to reduce peak memory of, and distribute, TN contractions. However, our implementation does not use index slicing, although there are packages like cotengra [25] with an interface similar to opt_einsum. We did not experiment with index slicing as our benchmark uses a single GPU and did not encounter out-of-memory errors. Still, we mention this technique, as, in combination with CSE, it could automatically reduce peak memory of the GGN quantities from Figure C16 which suffer from high memory requirements.Here, we extend the presented operations with a batch axis and allow for grouped convolutions.Adding a batch dimension (vmap-ing): Adding a batch axis to all presented operations is trivial. We only need to add an additional leg to the batched tensors, and connect these legs via element-wise or inner multiplication, depending on whether the result tensor is batched or not.Grouped convolutions: Grouped convolutions were originally proposed by Krizhevsky et al. [36] and allow for parallelizing, distributing, and reducing the parameters of the convolution operation. They split C in input channels into G groups of size Cin := Cin /G, then perform independent convolutions per group, each producing Cout := Cout /G output channels which are concatenated in the output. Each group uses a kernel W g of size Cout × Cin × K 1 × K 2 . These kernels are stacked into a single tensor W ∈ R Cout, Cin,K1,K2 such that [W] (g,:),:,:,: = W g . To support groups, we thus decompose the channel indices into c in := (c in , g) and c out := (c out , g). For the forward pass this yields the grouped convolution (without bias)Y (g,cout),o1,o2 = i1,i2,cin,k1,k2 X (g,cin),i1,i2 Π (1) i1,o1,k1 Π (2)i2,o2,k2 W (g,cout),cin,k1,k2 .(B8)Figure B9a shows the batched version of Equation (B8) as TN. Applying the differentiation rule from §3 leads to the Jacobians and VJPs shown in the remaining panels of Figure B9 .In Figure B12 we show the TNs for the GGN diagonal and the GGN Gram matrix (empirical NTK matrix) from Figure C16 extended by channel groups and a batch axis.Diagonal block extraction: Combined with index un-grouping, diagonal extraction generalizes to larger blocks: Let A ∈ R KI×KJ be a matrix of K horizontally and vertically concatenated blocksΠ (1) Π (2) X W i 1 i 2 k 1 k 2 o 1 o 2 n cin cout g (a) Convolution (no bias) Y n,(g,cout),o 1 ,o 2 Equation (B8) Π (1) Π (2) X v d e l t a i 1 i 2 k ′ 1 k ′ 2 o 1 o 2 c′ out cout n cin g g ′ (b) Weight Jacobian ∂Y n,(g,cout),o 1 ,o 2 ∂W (g ′ ,c ′ out ),c ′ in ,k ′ 1 ,k ′ 2 Π (1)Π (2) δ W Πk 1 k 2 c′ in i ′ 1 i ′ 2 o 1 o 2 n n ′ cout g g ′ (c) Input Jacobian ∂Y n,(g,cout),o 1 ,o 2 ∂X n ′ ,(g ′ ,c ′ in ),i ′ 1 ,i ′ 2 Π (1) Π (2) X V i 1 i 2 k ′ 1 k ′ 2 o 1 o 2 n c′ in c′ out g ′ (d) Weight VJP V (W) (g ′ ,c ′ out ),c ′ in ,k ′ 1 ,k ′ 2 = n,g,cout, o 1 ,o 2 V n,(g,cout),o 1 ,o 2 ∂Y n,(g,cout),o 1 ,o 2 ∂W (g ′ ,c ′ out ),c ′ in ,k ′ 1 ,k ′ 2 Π (1) Π (2) V (Y) W k 1 k 2 c′ in i ′ 1 i ′ 2 o 1 o 2 n ′ g ′ cout (e) Input VJP V (X) n ′ ,(g ′ ,c ′ in ),i ′ 1 ,i ′ 2 = n,g,cout, o 1 ,o 2 V n,(g,cout),o 1 ,o 2 ∂Y n,(g,cout),o 1 ,o 2 ∂X n ′ ,(g ′ ,c ′ in ),,i ′ 1 ,i ′ 2Π (2) X S (Y) A (k1,k2) ∈ R I×J , k i = 1 . . . , K. We can extract the diagonal blocks by restoring the sub-structure,Π (1) Π (2) X S (Y) o 1 o 2 o ′ 1 o ′ 2 i 1 i 2 i ′ 1 i ′ 2 k 1 k 2 c n g cin cout (a) GGN diagonal Π (1) Π (2) X S (Y) Π (1) Π (2) X S (Y) o 1 o 2 o ′ 1 o ′ 2 i 1 i 2 i ′ 1 i ′ 2 k 1 k 2 c c ′ n n ′ g cin cout (b) GGN Gram matrix (empirical NTK)Π (1) Π (2) X S (Y) Π (1) Π (2) X S (Y) o 1 o 2 o ′ 1 o ′ 2 i 1 i 2 i ′ 1 i ′ 2 c k1 k′ 1 g K1 k2 k′ 2 g K2 cin c′ in g Cin cout c′ out g Cout Figure B11: TN of a GGN mini-block diagonal without batching and channel groups.EQUATIONΠ (1) Π (1)1 1 then taking the diagonal along the K-dimensional index,Π (2) Π (2) X X i 1 i 2 i ′ 2 i ′ 1 o 1 o 2 o ′ 1 o ′ 2 k 1 k ′ 1 k 2 k ′ 2 cin c′ in g n (b) KFAC-reduce factorA (k,k) k i j = A (k, j) (k, i) k i j .We can apply this procedure to the GGN from Figure C16a . Assume we want to divide the output channel, input channel, and spatial dimensions into G Cout , G Cin , G K1 , G K2 groups. A group will thus be indexed with a tuple (g Cout , g Cin , g K1 , g K2 ) and the corresponding GGN block will be of dimensionC out /G Cout × C in /G Cin × K 1 /G K1 × K 2 /G K2 × C out /G Cout × C in /G Cin × K 1 /G K1 × K 2 /G K2and correspond to the GGN for [W] (g C out ,:),(g C in ,:),(g K 1 ,:),(g K 2 ,:) . This process of un-grouping the output dimensions, then taking the diagonal along the group indices, is illustrated in Figure B11 . Note that if we chooseG Cout = C out , G Cin = C in , G K1 = K 1 , G K2 = K 2 ,each block will be a single number and hence we recover the GGN diagonal from Figure C16b . If instead we G Cout = G Cin G K1 G K2 = 1, we obtain the full GGN from Figure C16a . The outlined schemes allows to extract mini-blocks of arbitrary size along the diagonal (subject to the total dimension).We were unable to find a definition of KFAC for grouped convolutions. Hence, we derive it here and present the TN diagrams. We use the perspective that grouped convolutions are independent convolutions over channel groups which are then concatenated. For each of those convolutions, we can then apply established the KFAC approximation for convolutions without groups. For a group g we have the kernel W g = [W] (g,:),:,:,: and the unfolded input of its associated input channels, X g = X (g,:),:,: = [X] (g,:),:,: (or X n,g = X n (g,:),:,: = [X] n,(g,:),:,: in the batched setting). KFC/KFAC-expand for grouped convolutions: Applying the regular KFC approximation to the kernel of group g, this yields the Fisher approximation Ω g ⊗ Γ g with Γ g ∈ R Cout× Cout andΩ g = 1 /N N n=1 X n,g X n,g ⊤ ∈ R CinK1K2× CinK1K2where X n,g is the input tensor for sample n and group g (remember the index structure X n,(g,cin),i1,i2 ). Figure B12a shows the diagram for {N Ω g } G g=1 .KFAC-reduce for grouped convolutions: Proceeding in the same way, but using the unfolded input averaged over output locations, we obtain the Fisher approximationΩg ⊗ Γg with Γg ∈ R Cout× Cout and Ωg = 1 /N(O1O2) 2 N n=1 1 ⊤ O1O2 X n,g (1 ⊤ O1O2 X n,g ) ⊤ ∈ R CinK1K2× CinK1K2for the kernel of group g. Figure B12b shows the diagram for{N (O 1 O 2 ) 2Ωg } G g=1 .Here we derive the KFAC approximation for transpose convolutions.We describe transpose convolution in terms of its associated convolution from an input spaceX = R Cin×I1×I2 to an output space Y = R Cout×O1×O2 . The convolution has hyper-parameters K 1,2 , S 1,2 , P 1,2 , D 1,2 with index patterns Π (1) = Π(I 1 , K 1 , S 1 , P 1 , D 1 ) ∈ R I1×O1×K1 and Π (2) = Π(I 2 , K 2 , S 2 , P 2 , D 2 ) ∈ R I2×O2×K2 .Transpose convolution as matrix multiplication: Transpose convolution maps a Y ∈ Y into an X ∈ X . In ML frameworks like PyTorch, its kernel W is stored asC out × C in × K 1 × K 2 tensor. The relation X = W ⋆ T Ywhere ⋆ T denotes transpose convolution is given by Figure 3d ,EQUATIONOur goal is to turn the express the above as matrix multiplication. To do that, we first define the matrix reshapeX of X via X ∈ R Cin×I1I2 such that [X] cin,(i1,i2) = X cin,i1,i2. Next, we consider a transposed kernel W of W with changed order of the first two indices, i.e. W ∈ R Cin×Cout×K1×K2 such thatEQUATION)This transposition is necessary to convert the kernel's layout in the ML framework to a layout that admits Equation (B9) to be expressed as matrix multiplication. Using a matrix reshapeW of W via W ∈ R Cin×CoutK1K2 such that [W ] cin,(cout,k1,k2) = W cin,cout,k1,k2, we can express Equation (B9) as matrix multiplicationX = W Y T (B11)where Y T ∈ R CoutK1K2×I1I2 is the transpose-unfolded input to the transpose convolution (note that• ̸ = • T !) [ Y T ] (cout,k1,k2),(i1,i2) = O1 o1=1 O2 o2=1 Π (1) i1,o1,k1 Π (2) i2,o2,k2 Y cout,o1,o2 . (B12)To the best of our knowledge there is no API for • T in existing ML frameworks. Our approach can provide a simple and efficient implementation of • through the TN shown in Figure B13a which corresponds to Equation (B12). As Equation (B11) is of the same form as Equation ( 1), it is now straightforward to write down the KFAC approximations for transpose convolution.We will define the KFAC-expand approximation for the GGN w.r.t. the flattened kernel w of W . Note that, in practise, this approximation must be properly transformed back to the layout W of the ML framework. We have G(w) ≈ Ω ⊗ Γ, with Γ ∈ R Cin×Cin computed from backpropagated gradients, and the input-based Kronecker factorΩ = Y T Y ⊤ T ∈ R CoutK1K2×CoutK1K2 . (B13)See Figure B13b for the corresponding TN.Π (1) Π (2) Y o 1 o 2 c out k 1 k 2 i 1 i 2 (a) Transpose- unfolded input Π (1)Π ( 1)Π (2) Π (2) Y Y o 1 o 2 o ′ 1 o ′ 2 i 1 i 2 k 1 k ′ 1 k 2 k ′ 2 c out c ′ out (b) KFAC-expand Π (1) Π (1)1 1Π (2) Π (2) 1 1 Y Y o 1 o 2 o ′ 1 o ′ 2 i 1 i ′ 1 i 2 i ′ 2 k 1 k ′ 1 k 2 k ′ 2 c out c ′ out (c) KFAC-reduceFigure B13 : TNs for extending KFAC to transpose convolutions (no batching and groups).For KFAC-reduce, we have G(w) ≈ Ω ⊗ Γ, with Γ ∈ R Cin×Cin computed from backpropagated gradients, and the input-based Kronecker factorΩ = 1 (I 1 I 2 ) 2 1 ⊤ I1I2 Y T 1 ⊤ I1I2 Y T ⊤ ∈ R CoutK1K2×CoutK1K2 . (B14)See Figure B13c for the corresponding TN.With batching and groups: In the presence of G groups, we have per-group kernels Wg = [ W] (g,:),:,:,: ∈ R C out/G× C in/G×K1×K2 and W g ∈ R C in/G× C out/G×K 1×K2 , as well as per-group transposeunfolded inputsY g T = Y T (g,:),:,: = [Y] (g,:),:,: T ∈ R C out/GK 1K2×I1 I2. Each group corresponds to a transpose convolution in itself. With batching, we have an additional leading batch dimension, i.e. Y n,g T . Applying the same steps from above, we can define the KFAC approximation for the GGN w.r.t. the flattened per-group kernel w g of W g .For KFAC-expand, we haveG(w g ) ≈ Ω g ⊗ Γ g , with Γ g ∈ R C in/G× C in/Gcomputed from backpropagated gradients, and the input-based Kronecker factorΩ g = 1 N N n=1 Y n,g T Y n,g ⊤ T ∈ R C out/GK 1K2 × C out/GK 1K2 .For KFAC-reduce, we have G(w g ) ≈ Ωg ⊗ Γg , with Γg ∈ R C in/G× C in/G computed from backpropagated gradients, and the input-based Kronecker factorΩg = 1 N (O 1 O 2 ) 2 N n=1 1 ⊤ I1I2 Y n,g T 1 ⊤ I1I2 Y n,g T ⊤ ∈ R C out/GK 1K2 × C out/GK 1K2 .Consecutive convolutions: We can chain two, or more, convolutions into a single TN diagram (Figure B14 ) to obtain a deep linear CNN [65] similar to deep linear networks which are popular for analytical studies.Convolution weight/input JVPs: In the main text, we derived the Jacobians of convolution ( §3.1) which can be used to derive the JVPs. A JVP propagates perturbations V (W) ∈ R Cout×Cin×K1×K2 and V (X) ∈ R Cin×I1×I2 in the input space to perturbations in the output space by contracting the perturbation with the Jacobian. See Table B3 for the general einsum expressions.To obtain per-sample gradients, the weight VJP must be carried out without summing over the batch axis which amounts to keeping the batch index in the output index tuple.VJPs and JVPs of im2col: With the TN differentiation technique described in §3.1 we can compute the Jacobian of the unfolding operation, then contract it with perturbations V (X) ∈ R Cin×K1×K2 in input space to obtain the JVP, or with perturbations V ( X ) ∈ R O1O2×CinK1K2 to obtain the VJP. Approximate Hessian diagonals (HesScale/BL89): Becker & Lecun [4] , Elsayed et al. [22] proposed approximate procedures for the Hessian diagonal which cost roughly a gradient. They can be understood as modifications of the Hessian backpropagation equations from Dangel et al. [16] .Π (1) Π ′ (1) Π (2) Π ′ (2) X W W ′ c in i 1 i 2 k 1 k 2 k ′ 1 k ′ 2 c out o 1 o 2 c ′ out o ′ 1 o ′ 2Consider a layer with input x, output y, and weights w inside a sequential feedforward neural network (for a convolutional layer, these correspond to the flattened input, output, and kernel). To compute per-layer Hessians of a loss ℓ, each layer backpropagates its incoming Hessian ∇ 2 y ℓ according to [16]EQUATIONThe scheme of [4, 22] imposes diagonal structure on the backpropagated quantity. A layer receives a backpropagated diagonal d (y) such that diag(d (y) ) ≈ ∇ 2 y ℓ, and backpropagates it according to Equation (B15), but with a post-processing step to obtain a diagonal backpropagated quantity,EQUATIONwhere diag(d (x) ) ≈ ∇ 2 x ℓ and diag(d (w) ) ≈ ∇ 2 w ℓ is an approximation to the Hessian diagonal. For convolutional layers, which are linear in the input and weight, the second summands are zero due to ∇ 2x y i = 0 = ∇ 2 w y i . The first terms of Equation (B16) require (i) embedding a diagonal vector into a matrix, (ii) applying MJPs and JMPs, and (iii) extracting the result's diagonal. Those can be expressed as a single TN. We show the diagrams in Figure B15 , using tensors rather than their flattened versions, that isEQUATIONΠ (2) Π (2) D (Y) W Π (1) Π (2) W k 1 k 2 k ′ 1 k ′ 2 o 1 o 2 i 1 i 2 n g cout cin (c) HesScale/BL89 input backpropagation (+ batch, groups) Π (1) Π (1) Π (2) Π (2) D (Y) X Π (1) Π (2) X i 1 i ′ 1 i 2 i ′ 2 o 1 o 2 k 1 k 2 n cout g cin (d) HesScale/BL89 weight backpropagation (+ batch, groups)Figure B15 : TN diagrams for HesScale/BL89 [4, 22] backpropagations through convolutional layers to approximate the Hessian diagonals D (X) , D (W) . JMPs and MJPs are shaded. (a, b) show the simple versions without batching and without channel groups. (c, d) include batching and channel groups.Table B3 : Extensive list of convolution and related operations (extension from Table 1 in the main text). All operations consider two spatial dimensions and support batching and channel groups. Generalization to other dimensions follow by introducing more spatial indices i_3, o_3, ... and kernel indices k_3, ....Convolution (no bias) [29] X, Π (1) , Π (2) , W"n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, (g c_out) c_in k1 k2 -> n (g c_out) o1 o2"Unfolded input (im2col, X ) X, Π (1) , Π"n c_in i1 i2, i1 o1 k1, i2 o2 k2 -> n (c_in k1 k2) (o1 o2)"Unfolded kernel (Toeplitz) Π (1) , Π (2) , WEQUATION"n (g c_out) o1 o2, i1 o1 k1, i2 o2 k2 -> n (g c_in) i1 i2"EQUATION"n (g c_out) o1 o2, i1 o1 k1, i2 o2 k2 -> n (g c_in k1 k2) i1 i2"Convolution weight VJP X, Π (1) , Π (2) , V (Y) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2 -> c_out c_in k1 k2"Convolution input VJP (transpose convolution) W, Π (1) , Π (2) , V (Y) "(g c_out) c_in k1 k2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2 -> n (g c_in) i1 i2"Convolution weight VJP (persample/batched) [58] X, Π (1) , Π (2) , V (Y) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2 -> n (g c_out) c_in k1 k2" Convolution weight JVP X, Π (1) , Π (2) , V (W) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, (g c_out) c_in k1 k2 -> n (g c_out) o1 o2"Convolution input JVP W, Π (1) , Π (2) , V (X) "(g c_out) c_in i1 i2, i1 o1 k1, i2 o2 k2, n (g c_in) i1 i2 -> n (g c_out) o1 o2"im2col VJP Π (1) , Π (2) , V ( X )"i1 o1 k1, i2 o2 k2, n (c_in k1 k2) (o1 o2) -> n c_in i1 i2" im2col JVP Π (1) , Π (2) , V (X)"i1 o1 k1, i2 o2 k2, n c_in i1 i2 -> n (c_in k1 k2) (o1 o2)"KFC/KFAC-expand [27, 23] X, Π (1) , Π (2) , X, Π (1) , Π (2) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_in_) i1_ i2_, i1_ o1 k1_, i2_ o2 k2_ -> g (c_in k1 k2) (c_in_ k1_ k2_)"X, Π (1) , Π (2) , X, Π (1) , Π (2) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_in_) i1_ i2_, i1_ o1_ k1_, i2_ o2_ k2_ -> g (c k1 k2) (c_ k1_ k2_)"KFC/KFAC-expand for transpose convolution Y, Π (1) , Π (2) , Y, Π (1) , Π (2) "n (g c_out) o1 o2, i1 o1 k1, i2 o2 k2, n (g c_out_) o1_ o2_, i1_ o1 k1_, i2_ o2 k2_ -> g (c_out k1 k2) (c_out_ k1_ k2_)"KFAC-reduce for transpose convolution Y, Π (1) , Π (2) , Y, Π (1) , Π"n (g c_out) o1 o2, i1 o1 k1, i2 o2 k2, n (g c_out_) o1_ o2_, i1_ o1_ k1_, i2_ o2_ k2_ -> g (c_out k1 k2) (c_out_ k1_ k2_)"GGN Gram/empirical NTK matrix [18, 51, 49] X, Π (1) , Π (2) , S (Y) , X, Π (1) , Π (2) , S (Y) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, c n (g c_out) o1 o2, n_ (g c_in) i1_ i2_, i1_ o1_ k1, i2_ o2_ k2, c_ n_ (g c_out) o1_ o2_ -> (c n) (c_ n_)"GGN/Fisher diagonal [17, 51] X, Π (1) , Π (2) , S (Y) , X, Π (1) , Π (2) , S (Y) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, c n (g c_out) o1 o2, n (g c_in) i1_ i2_, i1_ o1_ k1, i2_ o2_ k2, c n (g c_out) o1_ o2_ -> (g c_out) c_in k1 k2"GGN/Fisher diagonal (per-sample/batched) X, Π (1) , Π (2) , S (Y) , X, Π (1) , Π (2) , S (Y) "n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, c n (g c_out) o1 o2, n (g c_in) i1_ i2_, i1_ o1_ k1, i2_ o2_ k2, c n (g c_out) o1_ o2_ -> n (g c_out) c_in k1 k2"Approximate weight Hessian diagonal [4, 22] X, Π (1) , Π (2) , D (Y) , X, Π (1) , Π"n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2, n (g c_in) i1_ i2_, i1_ o1 k1, i2_ o2 k2 -> (g c_out) c_in k1 k2"Approximate input Hessian diagonal [4, 22] W, Π (1) , Π (2) , D (Y) , W, Π (1) , Π"(g c_out) c_in k1 k2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2, (g c_out) c_in k1_ k2_, i1 o1 k1_, i2 o2 k2_ -> n (g c_in) i1 i2"Approximate weight Hessian diagonal (persample/batched) X, Π (1) , Π (2) , D (Y) , X, Π (1) , Π"n (g c_in) i1 i2, i1 o1 k1, i2 o2 k2, n (g c_out) o1 o2, n (g c_in) i1_ i2_, i1_ o1 k1, i2_ o2 k2 -> n (g c_out) c_in k1 k2" Π (1) Π (2) X S (Y) Π (1) Π (2) X S (Y) o1 o2 o ′ 1 o ′ 2 i1 i2 i ′ 1 i ′ 2 c ′ out c ′ in k ′ 1 k ′ 2 cout cin k1 k2 c (a) GGN/Fisher Π (1) Π (2) X S (Y) Π (1) Π (2) X S (Y) o1 o2 o ′ 1 o ′ 2 i1 i2 i ′ 1 i ′ 2 cout cin k1 k2 c (b) GGN/Fisher diagonal Π (1) Π (2) X S (Y) Π (1) Π (2) X S (Y) o1 o2 o ′ 1 o ′ 2 i1 i2 i ′ 1 i ′ 2 k1 k2 cin cout c c ′ (c) Gram (empirical NTK)Here, we look at computing second-order information of a loss w.r.t. to the kernel of a convolution. Its computation can be phrased as backpropagation with a final extraction step [19] which contains less standard operations like Jacobian-matrix products (JMPs) and sub-tensor extraction. TNs can express this extraction step in a single diagram.Consider a datum (x, t) and its loss ℓ(w) = ℓ(f , t) where f := f w (x) ∈ R C is the prediction of a CNN with a convolution with flattened kernel w and flattened output y (derivations carry over to a batch loss). The kernel's generalized Gauss-Newton (GGN) matrix [63] G(w) = (J w f ) ⊤ ∇ 2 f ℓ(J w f ) ∈ R CoutCinK1K2×CoutCinK1K2is a positive semi-definite Hessian proxy preferred by many applications [e.g. 20, 45] and coincides with the Fisher information matrix for many common losses [46] . It is the self-outer product of a backpropagated symmetric factorizationS (y) = (J y f ) ⊤ S (f ) ∈ R CoutO1O2×C of the loss Hessian, ∇ 2 f ℓ(f , y) = S (f ) (S (f ) ) ⊤ . During backpropagation, the convolution extracts information about G(w) = (J w y) ⊤ S (y) (S (y) ) ⊤ J w y.In TN notation, this is easy to express without flattening: We simply compose two VJP diagrams from Figure 3c with an extra leg (MJP) and add the outer-product contraction to obtain the tensor version G(W) ∈ R Cout×Cin×K1×K2×Cout×Cin×K1×K2 of G(w) (Figure C16a ). The GGN is often further approximated by sub-tensors as it is too large. These slicing operations are also easy to integrate into the diagrams, e.g. to extract diagonal elements (Figure C16b [17, 51] ), or mini-block diagonals (Figure B11 [16, 3] ). This also removes redundant computations compared to computing, then slicing, the matrix. The same ideas apply to the GGN Gram matrix (S (w) ) ⊤ S (w) ∈ R C×C (Figure C16c ). It contains the GGN spectrum [18] and is related to the empirical NTK for square loss [49] .Here we present details on the index pattern computation, and additional transformations.Algorithm D1 lists pseudo-code for the index pattern computation from the convolution hyperparameters K, S, P, D, and the spatial input dimension I, that is Π(I, K, S, P, D). Unlike in the main text, we use 0-based indexing which is more common in numerical libraries. For self-consistency, we re-state the relation of the hyper-parameters to output dimension from [21, Relationship 15] ,EQUATION)Algorithm D1 Computing the convolution index pattern tensor Π for a spatial dimension.Require:Input size I ∈ N + , kernel size K ∈ N + , stride S ∈ N + , padding P ∈ N + 0 , dilation D ∈ N + O ← 1 + I+2P -K-(K-1)(D-1) S ▷ Compute output dimension [21, Relationship 15] Π ← 0I×O×K ▷ Initialize index pattern tensor for o = 0, . . . , O -1, k = 0, . . . , K -1 do ▷ Use 0-based indexing! i ← kD + oS -P ▷ Reconstruct contributing input element if 0 ≤ i ≤ I -1 then ▷ Check in bounds Π i,o,k ← 1 end if end for return Index pattern tensor Π ∈ {0, 1} I×O×K D.In the following, we will assume the absence of boundary pixels that don't overlap with the kernel, that isEQUATION)where the floor operation in O(I, K, S, P, D) is obsolete. This can always be assured by narrowing X before a convolution. Based on our hyper-parameter analysis of real-world CNNs ( §E), we identify:Transformation D1 (Dense convolutions) Assume Equation (D19). For K = S with default padding and dilation (P = 0, D = 1), patches are adjacent non-overlapping tiles, accessible by un-grouping the input index i into a tuple index ( ĩ, k) of size I /K × K:[Π(I, K, K, 0, 1)] i,o,k = [Π(I, K, K, 0, 1)] ( ĩ, k),o,k = δ ĩ,o δ k,k .Point-wise convolutions (K = S = 1) are a special case with pattern [Π(I, 1, 1, 0, 1)] i,o,k = δ i,o .Point-wise convolutions with K = S = 1 are common in DenseNets [33] , MobileNets [31, 60] and ResNets [30] . InceptionV3 [69] has 2d 'mixed dense' convolutions that are point-wise along one spatial dimension. ConvNeXt [42] uses dense convolutions with K = S ∈ {2, 4}.Transformation D2 (Down-sampling convolutions) For S > K with default padding and dilation (P = 0, D = 1), some elements do not overlap with the kernel. If the input dimension i is summed, all participating tensors can be pruned to remove the explicit zeros. Assume I mod S = 0. Then, pruning amounts to un-grouping i into (i ′ , s) of size I /S × S, narrowing s to K entries, and grouping back into an index ĩ of size KI /S. After pruning, the index pattern represents a dense convolution with input size KI /S, kernel size K, and stride K. In a contraction with some tensor V, Transformation D2 converts down-sampling convolutions to dense convolutions, which can be further simplified with Transformation D1. We find down-sampling convolutions with S = 2 > K = 1 in ResNet18 [30] , ResNext101 [72] , and WideResNet101 [73] . Those convolutions discard 75 % of their input! Knowledge that an operation only consumes a fraction of its input could be used to eliminate those 'dead' computations in preceding operations, reducing FLOPS and memory.I i=1 [V] ...,i,... [Π(I, K, S > K, 0, 1)] i,o,k = I /S ĩ=1 [ Ṽ] ...,Transformation D3 (Kernel-output dimension swap) Assume Equation (D19). Transposing kernel and output dimensions in an index pattern yields another index pattern with same input size, kernel size O(I, K, S, P, D), and swapped stride and dilation:[Π(I, K, S, P, D)] i,o,k = [Π(I, O, D, P, S)] i,k,o .This transformation is easy to see from the symmetry of (k, D) and (o, S) in Equation ( 7) and O(I, K, S, P, D). It converts index pattern contractions over output into kernel dimensions, like in convolutions. An example is the weight VJP from Figure 3c , which-after swapping kernel and output dimensions-resembles the TN for convolution from Figure 2 with kernel V. Rochette et al. [58] use this to phrase the computation of per-example gradients as convolution. §D.3 presents more properties of Π based on the sub-sampling interpretation of stride and dilation along the output and kernel dimensions. We also provide a transformation for swapping input and output dimensions, relating convolution and transpose convolution as described in [21] .For completeness, we state additional index pattern tensor properties here (using 1-based indexing): TransformationHere we give an overview of and characterize convolutions in popular architectures (see Table E4 ). We include moderately deep CNNs on Fashion MNIST, CIFAR-10, and CIFAR-100 from the DeepOBS benchmark [62] , and deep CNNs on ImageNet (AlexNet, ResNet18, InceptionV3, MobileNetV2, ResNext101). Regarding the hyper-parameters, we make the following observations:• Many CNNs do not use a bias term. This is because the output of those layers feeds directly into a batch normalization layer, which is invariant under the addition of a bias term. • All investigated convolutions use default dilation.• Group convolutions are rarely used. MobileNetV2 and ConvNeXt-base (Tables E4g and E4i ) use group convolutions that interpret each individual channel as a group. ResNext101 (Table E4f ) uses group convolutions that interpret a collection of channels as a group.ConvNeXt-base (Table E4g ) uses dense convolutions with P = 0 and S = K ∈ {2, 4}. • Many networks use dense convolutions, that is convolutions with unit kernel size (K = 1), unit stride (S = 1), and no padding (P = 0). These convolutions have a trivial index pattern and can therefore be simplified. • InceptionV3 (Table E4h ) uses two-dimensional convolutions with one trivial dimension ('mixed dense') with unit kernel size, unit stride, and no padding along one direction. For this spatial dimension, the index pattern can be simplified. • ResNet18 (Table E4e ) and ResNext101 (Table E4f ) use convolutions with S > K for down-sampling whose kernel only overlaps with a fraction of the input. The index pattern can be simplified. EQUATIONDense mix AuxLogits.conv0.conv (1) (768, 5,EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONDense mix Mixed_7b.branch3x3_2b.conv (4) ( Here we provide all details on the run time evaluation from the main text. We consider the convolutions from the CNNs from §E. Experiments were carried out on an Nvidia Tesla T4 (16 GB memory). We use a batch size of 32 for the ImageNet architectures, and 128 for the others.We compare different implementations of the same operations in PyTorch. The base line (referenced by 'PT') uses PyTorch's built-in functionalities for convolutions and related operations, such as torch.nn.functional.conv2d (forward), torch.nn.functional.unfold (KFC, KFACreduce), and PyTorch's built-in automatic differentiation torch.autograd.grad (VJPs).Our TN implementation (referenced by 'TN') sets up operands and the string-valued equation for each routine. Optionally, we can apply the simplifications from §4 as a post-processing step before contraction, which yields a modified equation and operand list ('TN + opt'). Finally, we determine the contraction path using opt_einsum.contract_path and perform the contraction with its PyTorch back-end (opt_einsum.contract). We only measure the contraction time as in practical settings, the contraction path search would be executed once, then cached. We also exclude final operations to obtain the correct shape or scale (flattening, reshaping, scaling by constant) in all implementations (including the base line).For each operation and each convolution layer, we perform 50 independent repetitions and report the minimum time in tables. To summarize those tables, we extract the performance ratios, that is the TN implementation's run time divided by the base line's. Ratios larger than 1 mean that the TN implementation is slower, ratios smaller than 1 indicate that it is faster than the base line. We collect those ratios for the different convolution types (general, mixed dense, dense, sub-sampling) and display them separately using box plots. Each operation has two boxes, corresponding to the un-simplified (TN), and the simplified (TN + opt) implementation. For the box plots, we use matplotlib's default settings (a box extends from the first quartile to the third quartile of the data, with a line at the median; whiskers extend from the box by 1.5x the inter-quartile range; flier points are those past the end of the whiskers). Figure F17 summarizes the entire GPU benchmark. Figure F18 shows the same information with each convolution type as an individual plot. We compare TN and TN+opt with PyTorch's torch.nn.functional.conv2d. Figure F19 visualizes the performance ratios for different convolution categories. 3.81 • 10 -3 6.17 2.62 • 10 -3 0.17 5.17 Table G10 : Additional internally required memory to compute the KFAC-reduce factor (measured on CPU). The value 0 indicates that an implementation's peak memory matches the memory consumption of its input X and result Ω.(a) 3c3d, CIFAR-10, input shape (128, 3, 32, 32)• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: We have read the Code of Ethics and believe that our work conforms to it.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We prefer I1, I2 over the more common choice H, W to simplify the generalization to higher dimensions.We use identical shapes for all tensors. Leg orientation does not assign properties like co-/contra-variance.Reshape can be seen as multiplication with a one-hot tensor, but we decided to use a separate symbol to emphasize it merely serves for re-interpretation and does not cause much computation.Standalone implementations of transpose convolution require another parameter to unambiguously reconstruct the convolution's input dimension (see §D for how to compute Π in this case).einsum does not yet support index un-grouping, so we must reshape manually before and after. 


Decision: The reviewers all found the paper well-written, well-presented, and the proposed framing of convolutions as tensor networks compelling. Despite some concerns about practical performance, they believe the paper will be of interest to the community. The authors addressed the questions in their rebuttal.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: ### 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: questions: Please see weakness
limitations: limitations: Please see weakness
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: None

Comments: comment: Thanks for the authors' response. I will maintain the score.
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.