Title: Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods
Venue: NeurIPS 2024
Accepted: yes
Keywords: Tensor networks

convolutions

KFAC

einsum

Second-order optimization
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.

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

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

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

Approximate Hessian diagonal: 1989 vs. 2024

Hessian rank: 2021 vs. 2023

Gradient descent learning dynamics: 2014 vs. 2023

Neural tangent kernel (NTK): 2018 vs. 2019

Kronecker-factored quasi-Newton methods: 2021 vs. 2022

Kronecker-factored curvature (KFAC, KFRA, KFLR): (2015, 2017, 2017) vs. (2016, 2020, 2020)

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

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:

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

We present transformations based on the convolution's connectivity pattern to re-wire and symbolically simplify TNs before evaluation (example in Figure 1).

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

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

Π

(1)

Π

(2)

X

W

i

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

J

X

K

(

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 unfolding

2.1 Convolution

2d convolutions process channels of 2d signals X ∈ R C in × I 1 × I 2 with C in channels of spatial dimensions 1 I 1 , I 2 by sliding a collection of C out filter banks, arranged in a kernel W ∈ R C out × C in × K 1 × K 2 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 C out × O 1 × O 2 with spatial dimensions O 1 , O 2 . Optionally, a bias from b ∈ R C out 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 J X K ∈ R C in K 1 K 2 × O 1 O 2 , called the unfolded input (or im2col ). Multiplying a matrix view W ∈ R C out × C in K 1 K 2 of the kernel onto the unfolded input then yields a matrix view Y of Y (the vector of ones, 1 O 1 O 2 , copies the bias for each channel),

We can also view convolution as an affine map of the flattened input x ∈ R C in I 1 I 2 into a vector view y of Y with a Toeplitz-structured matrix A ( W ) ∈ R C out O 1 O 2 × C in I 1 I 2 ,

This perspective is uncommon in code, but used in theoretical works [e.g. 65] as it highlights the similarity between convolutions and dense layers.

2.2 Tensor Multiplication

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 ∗ ( S 1 ,S 2 ,S 3 ) with

summing 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 ,

Binary and N -ary tensor multiplication are commutative: We can simultaneously permute operands and their index tuples without changing the result,

They are also associative, i.e. we can multiply operands in any order. However, the notation becomes involved as it requires additional set arithmetic to detect summable indices (see §H.1 for an example).

2.3 Tensor Networks & Convolution

A simpler way to understand tensor multiplications is via diagrams developed by e.g. Penrose [53]. RankK tensors are represented by nodes with K legs labelled by the index's name 2 . a i denotes a vector a , B i j a matrix B , and C i j k a rank-3 tensor C . A Kronecker delta [ δ ] i,j = δ i,j is simply a line, δ i j = I i j = j i . Multiplications are indicated by connections between legs. For inner multiplication, we join the legs of the involved indices, e.g. the matrix multiplication diagram is AB i k = A B i k j . Element-wise multiplication is similar, but with a leg sticking out. The Hadamard and Kronecker product diagrams are

Note 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 } I 1 × O 1 × K 1 × I 2 × O 2 × K 2 which represents the connectivity pattern between input, output, and kernel. P i 1 ,o 1 ,k 1 ,i 2 ,o 2 ,k 2 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 i 1 ,o 1 ,k 1 ,i 2 ,o 2 ,k 2 = Π (1) i 1 ,o 1 ,k 1 Π (2) i 2 ,o 2 ,k 2 where the index pattern tensor Π ( j ) ∈ { 0 , 1 } I j × O j × K j encodes the connectivity along dimension j . With that, one obtains

Without bias, this translates into the diagram in Figure 2a.

Π

(1)

Π

(2)

X

W

I

i

1

i

2

k

1

k

2

c

in

c

′

out

o

1

o

2

c

′

in

k

′

2

k

′

1

c

out

(a) Differentiation/weight Jac.

Π

(1)

Π

(2)

W

k

1

k

2

c

′

in

c

out

i

′

1

i

′

2

o

1

o

2

Π

(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

1

o

2

c

out

(d) Input VJP/tr. conv.

3 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),

3.1 Tensor Network Differentiation

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 [ JA j A N +1 ] S N +1 ,S ′ j = ∂ [ A N +1 ] S N +1 / ∂ [ A j ] 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 δ S j ,S ′ j ,

If 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 JWY , 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 J X K ⊤ ⊗ I C out for the Jacobian's matrix view [e.g. 16]. Figure 3b shows the input Jacobian JXY 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)

Π

(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

Π

(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

c

in

c

′

in

(b) KFAC-reduce

Figure 4: TNs of input-based Kronecker factors for KFAC approximations of the Fisher/GGN (no batching, no groups). 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).

3.2 Autodiff & Connections to Transpose Convolution

Although 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 C out × O 1 × O 2 from the output to the input or weight space. The VJP results V ( X ) ∈ R C in × I 1 × I 2 and V ( W ) ∈ R C out × C in × K 1 × K 2 are

Both are simply new TNs constructed from contracting the vector with the respective Jacobian, see Figures 3c and 3d (VJPs are analogous). The input VJP is often used to define transpose convolution [21]. In the matrix-multiplication perspective (Equation (2)), this operation is defined relative to a convolution with kernel W by multiplication with A ( W ) ⊤ , i.e. using the same connectivity pattern but mapping from the convolution's output to input space. The TN in Figure 3d makes this sharing explicit and cleanly defines transpose convolution. 4

3.3 Kronecker-factored Approximate Curvature

The Jacobian diagrams allow us to construct the TNs of second-order information like the Fisher/generalized Gauss-Newton (GGN) matrix and sub-tensors like its diagonal (§C). Here, we focus on the popular Kronecker-factored approximation of the GGN [47, 27, 23, 48] whose input-based Kronecker factor relies on the unfolded input J X K which requires large memory. State-of-the-art libraries that provide access to KFAC [17, 51] also use this approach. Using TNs, we can often avoid expanding J X K explicitly and save memory. Here, we describe the existing KFAC approximations and their TNs (see §5.1 for their run time evaluation).

KFC (KFAC-expand): Grosse & Martens [27] introduce a Kronecker approximation for the kernel's GGN, G ≈ Ω ⊗ Γ where Γ ∈ R C out × C out and the input-based factor Ω = J X KJ X K ⊤ ∈ R C in K 1 K 2 × C in K 1 K 2 (Figure 4a), the unfolded input's self-inner product (averaged over a batch).

KFAC-reduce: Eschenhagen et al. [23] generalized KFAC to graph neural networks and transformers based on the concept of weight sharing, also present in convolutions. They identify two approximations: KFAC-expand and KFAC-reduce. The former corresponds to KFC [27]. The latter shows similar performance in downstream tasks, but is cheaper to compute. It relies on the column-averaged unfolded input, i.e. the average over all patches sharing the same weights. KFAC-reduce approximates G ≈ ˆ Ω ⊗ ˆ Γ with ˆ Γ ∈ R C out × C out and ˆ Ω = 1 / ( O 1 O 2 ) 2 1 ⊤ O 1 O 2 J X K ( 1 ⊤ O 1 O 2 J X K ) ⊤ ∈ R C in K 1 K 2 × C in K 1 K 2 (Figure 4b; averaged over a batch). For convolutions, this is arguably a 'more natural' approximation as it becomes exact in certain limits [23], in contrast to the expand approximation.

KFAC for transpose convolution: Our approach enables us to derive KFAC for transpose convolutions. We are not aware of previous works doing so. This seems surprising because, similar to §2.1, transpose convolution can be seen as matrix multiplication between the kernel and an unfolded input. From this formulation we can immediately obtain KFAC through the weight sharing view of Eschenhagen et al. [23]. The Kronecker factor requires unfolding the input similar to im2col , but for transpose convolutions. This operation is currently not provided by ML libraries. We can overcome this limitation, express the unfolding operation as TN, and-for the first time-establish KFAC (expand and reduce) for transpose convolutions (see §B.4 for details).

Π

(

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)

=

˜

V

Π

(

IK

/

S

, K, K,

0

,

1)

i

i

′

o

k

o

k

Figure 5: TN illus- trations of index pat- tern simplifications and transformations. See § D.3 for the math formulation.

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

4.1 Index Pattern Structure & Simplifications

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 are

with i = 1 , . . . , I, o = 1 , . . . , O, k = 1 , . . . , K and output size O ( I, K, S, P, D ) = 1 + ⌊ ( I +2 P -( 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 ˜ V (Figure 5c). §5.1 shows that those simplifications accelerate computations.

4.2 Practical Benefits of the TN Abstraction & Limitations for Convolutions

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.

3

1

0

.

3

TN versus PT (logarithmic)

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 product C := 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 is ˜ C := ∑ N n =1 z n / π 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.

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

5.1 Run Time Evaluation

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

10

0

10

1

10

2

10

3

10

0

10

1

10

2

10

3

Auxiliary memory PT [MiB]

Aux. mem. TN [MiB]

convolutions (Figure F21). For KFC, the median performance ratios are well below 1 for dense, mixed dense & sub-sampling convolutions (Figure F22).

KFAC-reduce: 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 J X K 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.

Downstream improvements with KFAC-reduce: 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.

5.2 Randomized Autodiff via Approximate Contraction

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

(a) Real-world data

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

2

4

Keep probability

p

Normalized error

channel

spatial

(b) Synthetic data

0

0

.

2

0

.

4

0

.

6

0

.

8

1

0

2

4

Keep probability

p

Normalized error

channel

spatial

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.

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

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

Acknowledgments and Disclosure of Funding

The author would like to thank Luca Thiede, Andres Fernandez Rodríguez, and Kirill Neklyudov for providing feedback to the manuscript. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.

References

Adelman, M., Levy, K., Hakimi, I., and Silberstein, M. Faster neural network training with approximate tensor operations. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R. R., and Wang, R. On exact computation with an infinitely wide neural net. Advances in neural information processing systems (NeurIPS) , 2019.

Bahamou, A., Goldfarb, D., and Ren, Y. A mini-block fisher method for deep neural networks. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2023.

Becker, S. and Lecun, Y. Improving the convergence of back-propagation learning with second-order methods. 1989.

Benzing, F. Gradient descent on neurons and its link to approximate second-order optimization. In International Conference on Machine Learning (ICML) , 2022.

Biamonte, J. and Bergholm, V. Tensor networks in a nutshell, 2017.

Botev, A., Ritter, H., and Barber, D. Practical Gauss-Newton optimisation for deep learning. In International Conference on Machine Learning (ICML) , 2017.

Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., and Wanderman-Milne, S. JAX: composable transformations of Python+NumPy programs, 2018.

Bridgeman, J. C. and Chubb, C. T. Hand-waving and interpretive dance: an introductory course on tensor networks. Journal of Physics A: Mathematical and theoretical , 2017.

Chellapilla, K., Puri, S., and Simard, P. High performance convolutional neural networks for document processing. In International Workshop on Frontiers in Handwriting Recognition , 2006.

Chen, J., Xu, K., Wang, Y., Cheng, Y., and Yao, A. DropIT: Dropping intermediate tensors for memoryefficient DNN training. In International Conference on Learning Representations (ICLR) , 2023.

Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Catanzaro, B., and Shelhamer, E. cudnn: Efficient primitives for deep learning. 2014.

cuQuantum development team, T. cuQuantum SDK: A high-performance library for accelerating quantum information science, 2023.

Damm, C., Holzer, M., and McKenzie, P. The complexity of tensor calculus. computational complexity , 2002.

Dangel, F. unfoldNd: (n=1,2,3)-dimensional unfold (im2col) and fold (col2im) in pytorch, 2021.

Dangel, F., Harmeling, S., and Hennig, P. Modular block-diagonal curvature approximations for feedforward architectures. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2020.

Dangel, F., Kunstner, F., and Hennig, P. BackPACK: Packing more into backprop. In International Conference on Learning Representations (ICLR) , 2020.

Dangel, F., Tatzel, L., and Hennig, P. ViViT: Curvature access through the generalized gauss-newton's low-rank structure. Transactions on Machine Learning Research (TMLR) , 2022.

Dangel, F. J. Backpropagation beyond the gradient. 2023.

Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace redux effortless bayesian deep learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Dumoulin, V. and Visin, F. A guide to convolution arithmetic for deep learning. 2016.

Elsayed, M., Farrahi, H., Dangel, F., and Mahmood, A. R. Revisiting scalable hessian diagonal approximations for applications in reinforcement learning. In International Conference on Machine Learning (ICML) , 2024.

Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., and Hennig, P. Kronecker-factored approximate curvature for modern neural network architectures. In Advances in Neural Information Processing Systems (NeurIPS) , 2023.

Goldfarb, D., Ren, Y., and Bahamou, A. Practical quasi-newton methods for training deep neural networks, 2021.

Gray, J. and Kourtis, S. Hyper-optimized tensor network contraction. Quantum , 2021.

Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

Grosse, R. and Martens, J. A kronecker-factored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning (ICML) , 2016.

Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., Hubinger, E., Lukoši¯ ut˙ e, K., Nguyen, K., Joseph, N., McCandlish, S., Kaplan, J., and Bowman, S. R. Studying large language model generalization with influence functions, 2023.

Hayashi, K., Yamaguchi, T., Sugawara, Y., and Maeda, S.-i. Exploring unexplored tensor network decompositions for convolutional neural networks. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In IEEE conference on computer vision and pattern recognition (CVPR) , 2016.

Howard, A. G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., and Adam, H. Mobilenets: Efficient convolutional neural networks for mobile vision applications. 2017.

Huang, C., Zhang, F., Newman, M., Ni, X., Ding, D., Cai, J., Gao, X., Wang, T., Wu, F., Zhang, G., et al. Efficient parallelization of tensor network contraction for simulating quantum computation. Nature Computational Science , 2021.

Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. In IEEE conference on computer vision and pattern recognition (CVPR) , 2017.

Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks, 2020.

Jastrzebski, S., Szymczak, M., Fort, S., Arpit, D., Tabor, J., Cho*, K., and Geras*, K. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations (ICLR) , 2020.

Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems (NeurIPS) , 2012.

Laue, S., Mitterreiter, M., and Giesen, J. Computing higher order derivatives of matrix and tensor expressions. Advances in Neural Information Processing Systems (NeurIPS) , 2018.

Laue, S., Mitterreiter, M., and Giesen, J. A simple and efficient tensor calculus. In AAAI Conference on Artificial Intelligence , 2020.

LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D. Backpropagation applied to handwritten zip code recognition. Neural Computation , 1989.

Lin, W., Duruisseaux, V., Leok, M., Nielsen, F., Khan, M. E., and Schmidt, M. Simplifying momentumbased riemannian submanifold optimization. 2023.

Lin, W., Dangel, F., Eschenhagen, R., Neklyudov, K., Kristiadi, A., Turner, R. E., and Makhzani, A. Structured inverse-free natural gradient descent: Memory-efficient & numerically-stable KFAC. In International Conference on Machine Learning (ICML) , 2024.

Liu, Z., Mao, H., Wu, C.-Y., Feichtenhofer, C., Darrell, T., and Xie, S. A convnet for the 2020s. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

Ma, L., Ye, J., and Solomonik, E. Autohoot: Automatic high-order optimization for tensors. In International Conference on Parallel Architectures and Compilation Techniques (PACT) , 2020.

Magnus, J. R. and Neudecker, H. Matrix Differential Calculus with Applications in Statistics and Econometrics . Probabilistics and Statistics. 1999.

Martens, J. Deep learning via Hessian-free optimization. In International Conference on Machine Learning (ICML) , 2010.

Martens, J. New insights and perspectives on the natural gradient method, 2020.

Martens, J. and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning (ICML) , 2015.

Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations (ICLR) , 2018.

Novak, R., Sohl-Dickstein, J., and Schoenholz, S. S. Fast finite width neural tangent kernel. In International Conference on Machine Learning (ICML) , 2022.

Oktay, D., McGreivy, N., Aduol, J., Beatson, A., and Adams, R. P. Randomized automatic differentiation. In International Conference on Learning Representations (ICLR) , 2021.

Osawa, K., Ishikawa, S., Yokota, R., Li, S., and Hoefler, T. Asdl: A unified interface for gradient preconditioning in pytorch, 2023.

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS) . 2019.

Penrose, R. Applications of negative dimensional tensors. Combinatorial Mathematics and its Applications , 1971.

Petersen, F., Sutter, T., Borgelt, C., Huh, D., Kuehne, H., Sun, Y., and Deussen, O. ISAAC newton: Input-based approximate curvature for newton's method. In International Conference on Learning Representations (ICLR) , 2023.

Pinson, H., Lenaerts, J., and Ginis, V . Linear cnns discover the statistical structure of the dataset using only the most dominant frequencies. In International Conference on Machine Learning (ICML) , 2023.

Ren, Y., Bahamou, A., and Goldfarb, D. Kronecker-factored quasi-newton methods for deep learning, 2022.

Rigamonti, R., Sironi, A., Lepetit, V., and Fua, P. Learning separable filters. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2013.

Rochette, G., Manoel, A., and Tramel, E. W. Efficient per-example gradient computations in convolutional neural networks, 2019.

Rogozhnikov, A. Einops: Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations (ICLR) , 2022.

Sandler, M., Howard, A., Zhu, M., Zhmoginov, A., and Chen, L.-C. Mobilenetv2: Inverted residuals and linear bottlenecks. In IEEE conference on computer vision and pattern recognition (CVPR) , 2018.

Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2014.

Schneider, F., Balles, L., and Hennig, P. DeepOBS: A deep learning optimizer benchmark suite. In International Conference on Learning Representations (ICLR) , 2019.

Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation , 2002.

Singh, S. P., Bachmann, G., and Hofmann, T. Analytic insights into structure and rank of neural network hessian maps. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Singh, S. P., Hofmann, T., and Schölkopf, B. The hessian perspective into the nature of convolutional neural networks. 2023.

Smith, D. G. A. and Gray, J. opt_einsum - A python package for optimizing contraction order for einsum-like expressions. Journal of Open Source Software (JOSS) , 2018.

Smith, S. W. The scientist and engineer's guide to digital signal processing. 1997.

Springenberg, J. T., Dosovitskiy, A., Brox, T., and Riedmiller, M. Striving for simplicity: The all convolutional net, 2015.

Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In IEEE conference on computer vision and pattern recognition (CVPR) , 2016.

Tai, C., Xiao, T., Zhang, Y., Wang, X., et al. Convolutional neural networks with low-rank regularization. 2015.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Xie, S., Girshick, R., Dollár, P., Tu, Z., and He, K. Aggregated residual transformations for deep neural networks. In IEEE conference on computer vision and pattern recognition (CVPR) , 2017.

Zagoruyko, S. and Komodakis, N. Wide residual networks. 2016.

Zhang, F. A parallel tensor network contraction algorithm and its applications in quantum computation. 2020.

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

A Limitations

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 J X K and 1 ⊤ O 1 O 2 J X K 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: Wemention index slicing as a technique to reduce peak memory of, and distribute, TNcontractions. 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.

B Visual Tour of Tensor Network Operations for Convolutions

Here, we extend the presented operations with a batch axis and allow for grouped convolutions.

B.1 Convolution & First-order Derivatives

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 ˜ C in := C in / G , then perform independent convolutions per group, each producing ˜ C out := C out / G output channels which are concatenated in the output. Each group uses a kernel W g of size ˜ C out × ˜ C in × K 1 × K 2 . These kernels are stacked into a single tensor W ∈ R C out , ˜ C in ,K 1 ,K 2 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)

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.

B.2 Exact Second-order Information

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

v

Π

(1)

Π

(2)

X

W

i

1

i

2

k

1

k

2

o

1

o

2

n

˜

c

in

˜

c

out

g

(a) Convolution (no bias)

Y

n,

(

g,

˜

c

out

)

,o

1

,o

2

Equation (B8)

Π

(1)

Π

(2)

X

i

1

i

2

k

′

1

k

′

2

o

1

o

2

˜

c

′

out

˜

c

out

n

˜

c

in

g

g

′

(b) Weight Jacobian

∂

Y

n,

(

g,

˜

c

out

)

,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

′

˜

c

out

g

g

′

∂

Y

n,

(

g,

˜

c

out

)

,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,

˜

c

out

,

o

1

,o

2

V

n,

(

g,

˜

c

out

)

,o

1

,o

2

∂

Y

n,

(

g,

˜

c

out

)

,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

′

˜

c

out

(e) Input VJP

V

(

X

)

n

′

,

(

g

′

,

˜

c

′

in

)

,i

′

1

,i

′

2

=

∑

n,g,

˜

c

out

,

o

1

,o

2

V

n,

(

g,

˜

c

out

)

,o

1

,o

2

∂

Y

n,

(

g,

˜

c

out

)

,o

1

,o

2

∂

X

n

′

,

(

g

′

,

˜

c

′

in

)

,,i

′

1

,i

′

2

d

e

l

t

a

δ

Π

(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

n

g

˜

c

in

˜

c

out

(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

˜

c

in

˜

c

out

(b) GGN Gram matrix (empirical NTK)

A ( k 1 ,k 2 ) ∈ R I × J , k i = 1 . . . , K . We can extract the diagonal blocks by restoring the sub-structure,

Π

(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

˜

k

1

˜

k

′

1

g

K

1

˜

k

2

˜

k

′

2

g

K

2

˜

c

in

˜

c

′

in

g

C

in

˜

c

out

˜

c

′

out

g

C

out

Π

(1)

Π

(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

n

˜

c

in

˜

c

′

in

g

(a) KFC/KFAC-expand factor

1

1

Π

(1)

Π

(1)

1

1

Π

(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

˜

c

in

˜

c

′

in

g

n

(b) KFAC-reduce factor

then taking the diagonal along the K -dimensional index,

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 C out , G C in , G K 1 , G K 2 groups. A group will thus be indexed with a tuple ( g C out , g C in , g K 1 , g K 2 ) and the corresponding GGN block will be of dimension C out /G C out × C in /G C in × K 1 /G K 1 × K 2 /G K 2 × C out /G C out × C in /G C in × K 1 /G K 1 × K 2 /G K 2 and 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 choose G C out = C out , G C in = C in , G K 1 = K 1 , G K 2 = K 2 , each block will be a single number and hence we recover the GGN diagonal from Figure C16b. If instead we G C out = G C in G K 1 G K 2 = 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).

B.3 Kronecker-factored Approximate Curvature (KFAC) for Grouped Convolutions

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, J X g K = J X K ( g, :) , : , : = J [ X ] ( g, :) , : , : K (or J X n,g K = J X n K ( g, :) , : , : = J [ X ] n, ( g, :) , : , : K 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 ˜ C out × ˜ C out and

Ω g = 1 / N ∑ N n =1 J X n,g KJ X n,g K ⊤ ∈ R ˜ C in K 1 K 2 × ˜ C in K 1 K 2 where X n,g is the input tensor for sample n and group g (remember the index structure X n, ( g, ˜ c in ) ,i 1 ,i 2 ). 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 ˜ C out × ˜ C out and ˆ Ω g = 1 / N ( O 1 O 2 ) 2 ∑ N n =1 1 ⊤ O 1 O 2 J X n,g K ( 1 ⊤ O 1 O 2 J X n,g K ) ⊤ ∈ R ˜ C in K 1 K 2 × ˜ C in K 1 K 2 for the kernel of group g . Figure B12b shows the diagram for { N ( O 1 O 2 ) 2 ˆ Ω g } G g =1 .

B.4 Kronecker-factored Approximate Curvature (KFAC) for Transpose Convolution

Here we derive the KFAC approximation for transpose convolutions.

We describe transpose convolution in terms of its associated convolution from an input space X = R C in × I 1 × I 2 to an output space Y = R C out × O 1 × O 2 . 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 I 1 × O 1 × K 1 and Π (2) = Π ( I 2 , K 2 , S 2 , P 2 , D 2 ) ∈ R I 2 × O 2 × K 2 .

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 as C out × C in × K 1 × K 2 tensor. The relation X = ˜ W ⋆ T Y where ⋆ T denotes transpose convolution is given by Figure 3d,

Our goal is to turn the express the above as matrix multiplication. To do that, we first define the matrix reshape X of X via X ∈ R C in × I 1 I 2 such that [ X ] c in , ( i 1 ,i 2 ) = X c in ,i 1 ,i 2 . Next, we consider a transposed kernel W of ˜ W with changed order of the first two indices, i.e. W ∈ R C in × C out × K 1 × K 2 such that

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 reshape W of W via W ∈ R C in × C out K 1 K 2 such that [ W ] c in , ( c out ,k 1 ,k 2 ) = W c in ,c out ,k 1 ,k 2 , we can express Equation (B9) as matrix multiplication

J K where J Y K T ∈ R C out K 1 K 2 × I 1 I 2 is the transpose-unfolded input to the transpose convolution (note that J · K = J · K T!)

̸

To the best of our knowledge there is no API for J · K T in existing ML frameworks. Our approach can provide a simple and efficient implementation of J · K 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.

KFAC-expand: 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 C in × C in computed from backpropagated gradients, and the input-based Kronecker factor

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-reduce

KFAC-reduce: For KFAC-reduce, we have G ( w ) ≈ ˆ Ω ⊗ ˆ Γ , with ˆ Γ ∈ R C in × C in computed from backpropagated gradients, and the input-based Kronecker factor

See Figure B13c for the corresponding TN.

With batching and groups: In the presence of G groups, we have per-group kernels ˜ W g = [ ˜ W ] ( g, :) , : , : , : ∈ R C out / G × C in / G × K 1 × K 2 and W g ∈ R C in / G × C out / G × K 1 × K 2 , as well as per-group transposeunfolded inputs J Y g K T = J Y K T ( g, :) , : , : = J [ Y ] ( g, :) , : , : K T ∈ R C out / G K 1 K 2 × I 1 I 2 . Each group corresponds to a transpose convolution in itself. With batching, we have an additional leading batch dimension, i.e. J Y n,g K 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 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

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

B.5 Further Operations & Extensive Overview

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 C out × C in × K 1 × K 2 and V ( X ) ∈ R C in × I 1 × I 2 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.

Batched convolution weight VJP: 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 C in × K 1 × K 2 in input space to obtain the JVP, or with perturbations V ( J X K ) ∈ R O 1 O 2 × C in K 1 K 2 to obtain the VJP.

Π

(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

′

2

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

Consider 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]

The 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,

where 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 ∇ 2 x 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 is ( x , y , w , d ( x ) , d ( y ) , d ( w ) ) → ( X , Y , W , D ( X ) , D ( Y ) , D ( W ) ) .

Π

(1)

Π

(1)

Π

(2)

Π

(2)

D

(

Y

)

W

Π

(1)

Π

(2)

W

k

1

k

2

k

′

1

k

′

2

o

1

o

2

i

1

i

2

c

out

c

in

(a) HesScale/BL89 input backpropagation

Π

(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

c

out

c

in

(b) HesScale/BL89 weight backpropagation

Π

(1)

Π

(1)

Π

(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

˜

c

˜

c

in

Π

(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

˜

c

out

g

˜

c

out in

Π

(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

′

out

c

′

in

k

′

1

k

′

2

c

out

c

in

k

1

k

2

c

(a) GGN/Fisher

Π

(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

out

c

in

k

1

k

2

c

(b) GGN/Fisher 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

in

c

out

c

c

′

(c) Gram (empirical NTK)

C Exact Second-Order Information

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 C out C in K 1 K 2 × C out C in K 1 K 2 is 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 factorization S ( y ) = ( J y f ) ⊤ S ( f ) ∈ R C out O 1 O 2 × 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 C out × C in × K 1 × K 2 × C out × C in × K 1 × K 2 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].

D Implementation Details

Here we present details on the index pattern computation, and additional transformations.

D.1 Index Pattern Tensor Computation for Convolutions

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],

Algorithm D1 Computing the convolution index pattern tensor Π for a spatial dimension.

Require:

D.2 Index Pattern Tensor for Standalone Transpose Convolution

̸

Although 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 C out × O 1 × O 2 (the convolution's output space) to R C in × I 1 × I 2 (the convolution's input space). For convolution with S > 1 , we cannot infer I from O,K,S,P,D , as multiple I s map to the same O if ( I +2 P -K -( K -1)( D -1)) mod S = 0 (see the floor operation in Algorithm D1). We need to either supply I directly, or the remainder

(often called output_padding ) to make I unambiguous. Then, we compute

to get I ( O,A ) and call Algorithm D1 to obtain Π ( I ( O,A ) , K, S, P, D ) .

D.3 Details on Index Pattern Simplifications

In the following, we will assume the absence of boundary pixels that don't overlap with the kernel, that is

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 ( ˜ i, ˜ k ) of size I / 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 ˜ i 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 ,

with sub-tensor [ ˜ V ] ..., ˜ i,... = [[ V ] ..., ( i ′ ,s ) ,... ] ..., (: , : K ) ,... where : K means narrowing to K elements.

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.

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:

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):

Transformation D4 (Sub-sampling interpretation of stride) Strided convolutions ( S > 1 ) subsample non-strided convolutions along the output dimension, ignoring all but every S th output [21]. In other words, [ Π ( I, K, S, P, D )] i,o,k = [ Π ( I, K, 1 , P, D )] i, 1+ S ( o -1) ,k or, in tensor notation ( [ · ] :: S denotes slicing with steps of S ),

Transformation D5 (Sub-sampling interpretation of dilation) Dilated convolutions ( D > 1 ) with kernel size K sub-sample the kernel of a non-dilated convolution of kernel size K + ( D -1)( K -1) , ignoring all but every D th kernel element. In other words, [ Π ( I, K, S, P, D )] i,o,k = [ Π ( I, K +( K -1)( D -1) , S, P, 1)] i,o, 1+ D ( k -1) or, in tensor notation,

Transformation D6 (Transpose convolution as convolution) Assume Equation (D19) . Consider a non-strided ( S = 1 ), non-dilated ( D = 1 ) convolution with index pattern Π ( I, K, 1 , P, 1) and output dimension O ( I, K, 1 , P, 1) . Transposing the spatial dimensions and flipping the kernel dimension yields another index pattern with modified padding P ′ = K -P -1 . In other words, for all i = 1 , . . . , I , k = 1 , . . . , K , o = 1 , . . . , O

E Convolution Layer Hyper-parameter Analysis

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

(c) All-CNN-C, CIFAR-100 (3, 32, 32)

(d) AlexNet, ImageNet (3, 256, 256)

(e) ResNet18, ImageNet (3, 256, 256)

(f) ResNext101_32x8d, ImageNet (3, 256, 256)

(g) ConvNeXt-base, ImageNet (3, 256, 256)

Forward

Input VJP

Weight VJP

KFC

KFAC-red.

Forward

Input VJP

Weight VJP

KFC

KFAC-red.

Forward

Input VJP

Weight VJP

KFC

KFAC-red.

Forward

Input VJP

Weight VJP

KFC

KFAC-red.

10

-

1

10

0

10

1

General

Dense mix

Dense

Down

TN versus PT (logarithmic)

F Run Time Evaluation Details (GPU)

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.

F.1 Protocol & Overview

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.

(a) Mixed dense

Forward

Input VJPWeight VJP

KFC

KFAC-red.

1

3

.

16

TN versus PT (logarithmic)

(b) Dense

Forward

Input VJPWeight VJP

KFC

KFAC-red.

1

TN versus PT (logarithmic)

(c) Down-sampling

Forward

Input VJPWeight VJP

KFC

KFAC-red.

1

3

.

16

TN versus PT (logarithmic)

(d) General

Forward

Input VJPWeight VJP

KFC

KFAC-red.

0

.

1

1

10

TN versus PT (logarithmic)

F.2 Forward Pass

We compare TN and TN+opt with PyTorch's torch.nn.functional.conv2d . Figure F19 visualizes the performance ratios for different convolution categories. Table F5 contains the detailed run times and performance factors.

General

Dense mix

Dense

Down

10

0

10

1

TN versus PT (logarithmic)

(b) F-MNIST 2c2d, input shape (128, 1, 28, 28)

(c) CIFAR-100 All-CNN-C, input shape (128, 3, 32, 32)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

Type

General

Dense

Type

General

Down

Type

General

Dense own

Type

Dense

General

Type

General

Dense

Dense mix

Type

General

Dense

F.3 Input VJP

We compare TN and TN+opt with a PyTorch implementation of the input VJP via torch.autograd.grad . Figure F20 visualizes the performance ratios for different convolution categories. Table F6 contains the detailed run times and performance factors.

General

Dense mix

Dense

Down

10

0

10

1

TN versus PT (logarithmic)

(c) CIFAR-100 All-CNN-C, input shape (128, 3, 32, 32)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

(f) ResNext101, input shape (32, 3, 256, 256)

(g) ConvNeXt-base, input shape (32, 3, 256, 256)

Type

General

Dense

Type

General

Down

Type

General

Dense own

Type

Dense

General

Type

General

Dense

Dense mix

Type

General

Dense

F.4 Weight VJP

We compare TN and TN+opt with a PyTorch implementation of the weight VJP via torch.autograd.grad . Figure F21 visualizes the performance ratios for different convolution categories. Table F7 contains the detailed run times and performance factors.

General

Dense mix

Dense

Down

10

0

10

1

TN versus PT (logarithmic)

(c) CIFAR-100 All-CNN-C, input shape (128, 3, 32, 32)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

(f) ResNext101, input shape (32, 3, 256, 256)

Type

General

Dense

Type

General

Down

Type

General

Dense own

Type

Dense

General

Type

General

Dense

Dense mix

Type

General

Dense

F.5 KFC Factor (KFAC-expand)

We compare TN and TN+opt with a PyTorch implementation of the input-based KFC factor based on torch.nn.functional.unfold . Figure F22 visualizes the performance ratios for different convolution categories. Table F8 contains the detailed run times and performance factors.

General

Dense mix

Dense

Down

10

-

1

10

0

10

1

TN versus PT (logarithmic)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

(f) ResNext101, input shape (32, 3, 256, 256)

(g) ConvNeXt-base, input shape (32, 3, 256, 256)

Type

General

Dense

Type

General

Down

Type

General

Dense own

Type

Dense

General

Type

General

Dense

Dense mix

Type

General

Dense

F.6 KFAC-reduce Factor

We compare TN and TN+opt with a PyTorch implementation of the input-based KFAC-reduce factor based on torch.nn.functional.unfold . Figure F23 visualizes the performance ratios for different convolution categories. Table F9 contains the detailed run times and performance factors.

General

Dense mix

Dense

Down

10

-

0

.

5

10

0

10

0

.

5

TN versus PT (logarithmic)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

(f) ResNext101, input shape (32, 3, 256, 256)

Type

General

Dense

Type

General

Down

Type

General

Dense own

Type

Dense

General

Type

General

Dense

Dense mix

Type

General

Dense

G Memory Evaluation Details (CPU)

Here, we investigate the peak memory consumption of our proposed TN implementations.

G.1 Theoretical & Empirical Analysis for KFAC-reduce Factor

We assume a two-dimensional convolution with input X of shape ( C in , I 1 , I 2 ) , output of shape ( C out , O 1 , O 2 ) and kernel of shape ( C out , C in , K 1 , K 2 ) . The analysis with a batch dimension is analogous; hence we suppress it here to de-clutter the notation.

The main difference between the default and our proposed TN implementation of ˆ Ω from §3.3 lies in the computation of the averaged unfolded input J X K ( avg ) := 1 / ( O 1 O 2 ) 1 ⊤ O 1 O 2 J X K which consists of C in K 1 K 2 numbers. In the following, we will look at the extra memory on top of storing the input X , the averaged unfolded input J X K ( avg ) , and the result ˆ Ω .

Default implementation: The standard implementation computes J X K ( avg ) via the unfolded input J X K and thus requires extra storage of C in K 1 K 2 O 1 O 2 numbers.

TN implementation (general case): The TN implementation requires storing the averaged index patterns Π ( i, avg ) := 1 / O i ∑ O i o =1 [ Π ( i ) ] : ,o, : for i = 1 , 2 . These are directly computed via a slight modification of Algorithm D1 and require storing I 1 K 1 + I 2 K 2 numbers. In contrast to the default implementation, spatial dimensions are de-coupled and there is no dependency on C in .

TN implementation (structured case): For structured convolutions (Figure 5) we can describe the action of the index pattern tensor through reshape and narrowing operations. ML libraries usually perform these without allocating additional memory. Hence, our symbolic simplifications completely eliminate the allocation of temporary intermediates to compute J X K ( avg ) .

Empirical results: To demonstrate the memory reduction inside the computation of ˆ Ω we measure its peak memory with the memory-profiler library and subtract the memory required to store X and ˆ Ω . This approximates the extra internal memory requirement of an implementation. With the setup of §F we report the minimum additional memory over 50 independent runs in Table G10. We consistently observe that the TN implementation has lower peak memory, which is further reduced by our symbolic simplifications (see for example the effect on ResNext101's dense and down-sampling convolutions in Table G10f).

Our theoretical analysis from above suggests that the peak memory difference becomes most visible for many channels with large kernel and output sizes. One example are ConxNeXt-base's features.1.0.block.0 convolutions with K 1 = K 2 = 7 , O 1 = O 2 = 64 , and C in = 128 (Table E4g). For those convolutions, we observe that the default implementation requires an additional 3,140 MiB ( ≈ 3 GiB!) of memory, whereas the TN implementation has zero extra memory demand (Table G10g). This is consistent with our theoretical analysis in that the overhead is storing the unfolded input, which has ( N = 32) · ( C in = 128) · ( O 1 = 64) · ( O 2 = 64) · ( K 1 = 7) · ( K 2 = 7) = 822 , 083 , 584 float32 entries, corresponding to 3,136 MiB.

(c) CIFAR-100 All-CNN-C, input shape (128, 3, 32, 32)

(d) Alexnet, input shape (32, 3, 256, 256)

(e) ResNet18, input shape (32, 3, 256, 256)

(f) ResNext101, input shape (32, 3, 256, 256)

H Miscellaneous

H.1 Example: Associativity of Tensor Multiplication

Here, we demonstrate associativity of tensor multiplication through an example. The technical challenge is that an index can only be summed once there are no remaining tensors sharing it. Therefore, we must carry indices that are summed in later multiplications in the intermediate results, which requires some set arithmetic on the index sets.

Let S 1 , S 2 , S 3 be index tuples of the input tensors A , B , C , and S 4 ⊆ ( S 1 ∪ S 2 ∪ S 3 ) a valid output index tuple of their tensor multiplication D = ∗ ( S 1 ,S 2 ,S 3 ,S 4 ) ( A , B , C ) . We can either first multiply A with B to obtain an intermediate tensor of index structure S 1 , 2 , or B with C to obtain an intermediate tensor of index structure S 2 , 3 , before carrying out the remaining multiplications. To construct the intermediate index structures, we divide the indices ˜ S = ( S 1 ∪ S 2 ∪ S 3 ) \ S 4 that are summed over into those only shared between A , B given by ˜ S 1 , 2 = ( S 1 ∪ S 2 ) \ ( S 4 ∪ S 3 ) , and those only shared among B , C given by ˜ S 2 , 3 = ( S 2 ∪ S 3 ) \ ( S 4 ∪ S 1 ) . This yields the intermediate indices S 1 , 2 = ( S 1 ∪ S 2 ) \ ˜ S 1 , 2 and S 2 , 3 = ( S 2 ∪ S 3 ) \ ˜ S 2 , 3 , and the parenthesizations

This generalizes to n -ary multiplication, allowing to break it down into smaller multiplications. However, the index notation and set arithmetic from Equation (H20) quickly becomes impractical.

H.2 Example: Matrix-matrix Multiplication as Tensor Multiplication

Here we provide a small self-contained example that demonstrates Equation (3) for matrix-matrix multiplication.

Consider two matrices A , B which are compatible for multiplication and let C = AB . In index notation, we have

The index tuples are S A = ( i, j ) , S B = ( j, k ) , and S C = ( i, k ) . Next, we evaluate which indices are summed over. Since the order of those indices does not matter, we can interpret the tuples as sets and use set arithmetic:

Now we see that matrix-matrix multiplication is a case of tensor multiplication (Equation (3)),

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We provide a bullet point list of our contributions in §1, each of which references the part of the paper that outlines the contribution.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See §4.2 and §A.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The TN simplifications we provide in §4.1 follow straightforward from the index pattern's structure Equation (7) and are stated rigorously in §D.3, including their assumptions.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide implementation details in §D, experimental and hardware details in §F and G.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will open-source the code to reproduce all our experiments, as well as the raw data containing the results shown in the manuscript.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: See §F and G for details on the experimental setting.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Our run time plots contain box plots with medians and quartiles reported over different convolutions, and the randomized backpropagation results show mean and standard deviations for 10 different model and batch initializations.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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.

8. 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).

9. 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 ?

Answer: [Yes]

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

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

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper does not release any data or models that have a high risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the papers introducing the neural network architectures and data sets used in our experiments.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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.

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

STRENGTHS:
1. The paper is well written with nice tensor network graphs. While many papers in the tensor decomposition field suffer from notorious indexes, the notation in this paper is consistent and clear.

2. Although the TN graphs for matrix multiplications and derivatives themselves are not new, this paper seems to be the first to provide a comprehensive study of these representations and computations in CNNs. Based on these representations, the authors show how to efficiently implement the KFAC and structured dropout algorithm, which is important for optimization.

3. Experiments show the computational and memory efficiency of the proposed TN implementations.

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

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

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe key strengths are: \n- The paper includes rigorous mathematical formulations and proofs to support the proposed methods.\n- Provides thorough details on implementation, aiding reproducibility and practical adoption.\n- Applicable to a wide range of convolutional operations, including those with various hyper-parameters, batching, and channel groups. The generality enhances utility and flexibility\n- Good use of visualisations and diagrams to clarify and explain the networks\n- Interest to the research community with strong potential for future extensions\n\nWEAKNESSES:\nThe 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 \n\nMinor things:\n- Line 57 does not appear to have the reference to 2.2 in the right place\n- 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? \n- 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\n- Standard errors, rather than standard deviation, might be more informative in Figure 8\n\nQUESTIONS:\nQuestions for authors:\n- Why are the (non simplified) TN implementations slower in general? Do the standard implementations have further efficiency optimisations? \n- 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\n- Does figure 6 include variability wrt contraction order?\n- How many of the operations support GPU/hardware acceleration? Would there be performance gains if they were implemented?\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n### Summary of \"Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods\"\n\nThe 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:\n\n1. **Recast of Convolution Operation into Tensor Networks**:\n   - 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.\n\n2. **Tensor Networks for Machine Learning**:\n   - 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.\n\n3. **Implementation with Einsum and Opt_Einsum**:\n   - 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).\n\n4. **Insights into Second-Order Optimization**:\n   - 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.\n\nIn 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.\n\nSTRENGTHS:\n### Strengths of \"Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods\"\n\n1. **Innovative Perspective**:\n   - The paper introduces an innovative perspective by recasting convolution operations into tensor networks (TNs). This novel approach simplifies the analysis and implementation of convolutions, making complex operations more understandable and accessible.\n\n2. **Advances in Second-Order Optimization**:\n   - The paper shows that the tensor network view of convolutions can greatly benefit second-order optimization methods, particularly Kronecker-factored Approximate Curvature (KFAC). This provides new insights and practical improvements in the computation of curvature approximations.\n\n3. **Comprehensive Diagrams and Transformations**:\n   - The inclusion of detailed diagrams and specific transformations based on the connectivity pattern of convolutions enhances the clarity and utility of the proposed methods. These visual aids make it easier to understand and implement the tensor network approach.\n\n4. **Flexibility and Modifiability**:\n   - The tensor network framework allows for easy modifications and supports various convolutional configurations, including different hyper-parameters, batching, channel groups, and convolution dimensions. This flexibility makes the approach broadly applicable.\n\n5. **Experimental Validation**:\n   - The paper provides robust experimental validation, demonstrating the practical benefits of the tensor network approach in various convolutional operations and configurations. The significant speed-ups and memory savings observed in the experiments highlight the effectiveness of the proposed methods.\n\nWEAKNESSES:\n### Weaknesses of \"Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods\"\n\nWhile 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:\n\n1. **Suboptimal Use of Einsum for Convolutions**:\n   - 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.\n\n2. **Performance Trade-offs for Specific Architectures**:\n   - 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.\n\n3. **Limited Comparison with Optimized Convolution Algorithms**:\n   - 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.\n\nIn 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.\n\nQUESTIONS:\nPlease see weakness\n\nLIMITATIONS:\nPlease see weakness", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}