Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods
Authors: Felix Dangel
Venue: NeurIPS 2024 poster (2024)
Keywords: Tensor networks
convolutions
KFAC
einsum
Second-order optimization
Primary area: optimization
TL;DR: We introduce einsum implementations of many operations related to convolution that speed up the computation of second-order methods like KFAC

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
Felix Dangel
Vector Institute
Toronto, Canada
fdangel tueeeresttettttitttttttee..ee
Abstract
Despite their simple intuition, convolutions are more tedious to analyze than dense layers, which complicates the transfer of theoretical and algorithmic ideas to convo- lutions. 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, derive diagrams of various autodiff operations and popular support, batching, channel groups,
vature approximations with full and generalization any convolution dimension. Further, provide convolution- specific transformations based on the connectivity pattern which allow to simplify diagrams before evaluation. Finally, probe performance. Our TN implementa- tion accelerates 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.
Introduction
Convolutional neural networks [CNNs, 39] mark 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 explain with graphical illustrations [e.g. 21J. Yet, convolutions are more challenging to analyze than dense layers in multi-layer perceptrons (MLPs) or transformers [71]. One reason that they are hard express in matrix notation and-ever in index notation-compe ccomactt expressions that are convenient to work with only exist for special thyper-paaameter [e.g. 27, 2]. Many
pppre-parrmeter (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 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 of these methods require also reflects this gap.
Some functions only support special dimensions [15]. Others use less efficient workarounds ($5.1) 1) or
not provided at all (SB.4). And they are hard modify as the code either closed-source [12] or
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

from sincanv inport cony, inw.p,ndss.patars, from einopa Inport winsun; inport torch
I, K, D. (C_in, T[0l, kfo_shape (C_in K[O] x(t), C_in x[0] K[t]) final shape
a convolution lyper-paraneen paranetera (2- 2-tuplex)
input
def ifc_ ():
CoConppae Kronecker unfoldd. K, D, P, 8)
in2col."
I_unf torch nn
....emmtoaaataaaa functional
raturn ainson(3_unf X_anf, " out, out -2 3
k).
Compute Kranecker factor its tansor network.
Pii con__indes_pattene (I K[0], 8[0]. P[O]. D(0])
P12 conv_indes.pettera (I[1), K[1] stil, P[t] D[13)
Pi1, *c.in 11 12, 11 kl, 12 o2 k2, c.in. 11, 12., 11, ol. kt. 12, o2. k2." 6.in B2 c-1n- k1. k2_ wet...raaaaeeee......
sinaun(X, Pit,
11
convolutions Computa Kranecker factor via 1ta simplified tanzor network.".. x, reahape (C_is, I[0] // K[G], K[0] I[3] // K[s], k[:]) return einsun (t, X, "c_in ol k2, cuin. kl. o2. k2_"
k2 c_tn_ k2_ rwwhapa ahac_a)ppe)
Figure 1: Many convolution-related routines can be expressed as TNs and evaluated with einsum. We illustrate this for the input- based factor of KFAC for convolutions [KFC, 27], whose standard implementation (top) requires unfolding the input (high memory). The TN (middle) enables internal optimizations inside einsum (e.g. with contraction path optimizers like opt_einsum [66]). (Bottom) In many cases, the TN further simplifies due structures in the index pattern, which reduces cost.
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 diagrams. These diagrams are simpler to parse than mathematical equations and can seamlessly be (i) manipulated to take derivatives, add batching, einsum. This yields simple, modifiable implementations that benefit from automated under-the-hooed the-hood- optimizations 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:
extract sub-tensors, (ii) merged with other diagrams, and (iii) evaluated with
1. We the TN format of convolution from Hayashi al. [29] to derive diagrams and einsum formulas for autodiff and less standard routines for curvature approximations with support for all syper-parametent batching, groups, and any dimension (Table 1).
2. We present transformations based on the convolution's 's connectivity pattern re-wire and symbolically simplify TNs before evaluation (example in Figure 1).
3. We compare default and TN implementations, demonstrating optimal peak memory reduc- tion and run time improvements up to 4.5 x for a recent KFAC variant, and showcase their flexibility 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 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 ( ([))s, ; = Atz)) Parenthesized indices
mean reshapes, e.g. [a] (+.) = [A]a, with the flattened matrix A.

---

01
01 02
(2).02)
(01.02)
[Gumt o2)
(Cout,
(1)
Gost
o1
(9m Nx K2)
(9n Ra
A(W)
01
n
N1
no
O2
02
102
Arg
(Can &is k2)
(Can (2)
F2
X
Cin
W
Cout
Go
Gin
Cour
(a) Convolution
(b) Input unfolding
(c) Kernel unfolding
Figure 2: TNs of (a) 2d convolution and (b,c) connections to matrix multiplication view. The connectivity along each dimension is explicit via index pattern tensor N.
2.1 Convolution
2d convolutions process channels of 2d signals X E RCax4x/ with Cn channels of spatial dimen- sions' 11, 12 by sliding collection of Cou filter banks, arranged in kernel W E IR x X xixk x with kernel size K1, K2, over the input. The sliding operation depends on various hyper- parameters [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 E xCaxx0xx0 with spatial dimensions O1, O2. Optionally, bias from b E RR is added per channel.
One way implement convolution is via matrix multiplication [10], similar fully-connected layers. First, one extracts the overlapping patches from the input for each output, then flattens and column-stactoo them into matrix [X] E IR K X2xO,O, called the unfolded input (or im2col). Multiplying matrix view W E R Con C K2 K2Kg K2 of the kernel onto the unfolded input then yields matrix view Y of Y (the vector of ones, 10,0, copies the bias for each channel).
= W + b bl, E C
We also view convolution affine map of the flattened input BCddds into vector view of Y with a Toeptltz-struuctured matrix A A(W E (Cs00,02xClli
y A x
l0,0, RR
(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.
Tensor Multiplication
Tensor multiplication unifies outer (Kronecker), element-wise (Hadamard), and inner products and uses the onput-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 described by the index tuples S1, S2, S3 where S3 C U S2) (converting tuples to sets needed). Any product A and B can be described by the multiplication operator (2.,5) with
C ((, B)
[C|s, E(s (S.
A s. B]ss
(3)
summing over indices that are not present in the output. E.g., for two matrices A, B, their product is AB ((..)).G.))((..))) (A, B) (see HH.2), their Hadamard product A ((...)). B). and their Kronecker product A B ((i.)). (k.)).)).............)..)))).)). (G)))))))))) B). Libraries support this func- tionality via einsum, which takes string encoding of SI, S2, Sa. followed by A, B. It also accepts longer sequences A An with index tuples Si, S2,...... Sy and output index tuple SN+1.
N
AN+1 Sw An) + |AN+1]SNH
Σ
(4)
(UN., S.
WE
We prefer 11, the more common choice H, W simplify the generalization higher dimensions.

---

Table 1: Contraction expressions of operations related 2d convolution. They include batching and channel groups, which are standard features in implementations. We describe each operation by tuple of input tensors and a contraction string that uses the einops library's syntax [59] which can express index (un-)grouping. Some quantities are only correct up to a scalar factor which is suppressed for brevity. See šB for visualizations and Table B3 for more operations.
Operation
Operands
Contraction string (einops [59] convention)
Conv. (no bias)
X. ( n W
to (g cutn) 12, 11 kt, 12 02 x2, (g (_ont) c,in
(g cLous)
Unf. input (im2col)
X. m m
tu clin 12, 11 e1 x2,
12 (o_ta k2) (et 02)*
Unf. kernel (Toeplitz)
n n W
of
x2
25 (c.00b 02) [c_in 11 12)
Weight VJP Input VJP conv.)
X, n 3 Vo
ta (g c_1m) 12, 11 x1,
02 ki, 12 k2, n. (g c_out)
(g c_ous) o1
W, (1) vo
"g c_put) c.in
KFC//FAC-eepand
X. (
n
X, n
2)
to (g in) 11 12.
11
x1, 12
X2 i (g cLio.) 11 13. (e
31 01 k1.
32
X00 (e.in. x1_ 12_)"
KFAC-reduce
X, ( 3
X, (
(2)
lg e_in)
o2 x2. (n a_in.] 12,
ol. xi...
12 o2
Co. kl. x2_""
Binary and N-ary tensor multiplication commutative: We can simultaneously permute operands and their index tuples without changing the result,
*(S1.92 Sa) A, B) * (S2,S1 Ss) [B, A) *(.S....S).) A., A,..) +(..S,,.....) Aj. Ai,
They 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 understand tensor multiplications is via diagrams developed by e.g. Penrose [53]. Rank-K tensors are represented by nodes with legs labelled by the index's 's name-. O denotes vector
Ø matrix B, and e rank-3 tensor C. A Kronecker delta 6)..j 8,ح is simply line, inner multiplication, join the legs of the involved indices, c.g. the matrix multiplication diagram is : - m-c. Element-wise multiplication is similar, but with a leg sticking out. The
; - =y= . Multiplications are indicated by connections between legs. For
Hadamard and Kronecker product diagrams
(G.)
(j.1)
(3:4)
(5)
B
Note that the outer tensor product rank-4 tensor and must be reshaped (indicated by black triangles") ) into matrix. This syntax allows for extracting and embedding tensors along diagonals; e.g. taking a matrix diagonal, 4) ناا or forming a diagonal matrix, :- ت and generalizes to larger diagonal blocks (SB). In the following, we stick to the simplest case to avoid the more advanced syntax. However, it shows the expressive power of TNs and required to support common features of convolutions like channel groups (known as separable convolutions).
which represents the connectivity pattern between input, output, and kernel. P., on KI 02, is I if input locations (i1, i2) overlap with kernel positions (k1, k2) when computing output locations (or, 02) and 0 otherwise. The spatial couplings are independent along each dimension, hence P decomposes
Application to convolution: We define a binary tensor P E {0, x0.1))xx0xxXxxxxx000
tt. . og where the index pattern tensor n (0, (4, encodes the connectivity along dimension j. With that, one obtains
into P Pi f2. ( n!
On I1.2
Kr.Ka
Yom101:02 = DeC. +
+ Σ
X
11.01 No 12.02,02 Wenn
Ga=1
f1.42=1
ky
1k==1
Without bias, this translates into the diagram in Figure 2a.
We use identical shapes for all tensors. Leg otientation does not assign properties like co-/contra- co-fcontra-variance. 'Reshape can be seen as multiplication with a one- one-hot tensor, but we decided to use separate symbol to emphasize it merely serves for re- e--iterpreetatin and does not much computation.
variance.

---

99
4
W
Caee
(a) Differentiation/veight Jac.
(b) Input Jac.
(c) Weight VJP
(d) Input VIP/tr. conv.
Figure 3: TN differentiation graphical manipulation. (a) Differentiating convolution w.r.t. W is cutting out of the diagram and yields the weight Jacobian. (b) Same procedure applied to the
Jacobian w.r.t. X. (c) VJP for the weight and (d) input Jacobian (transpose convolution). Jacobians
are shaded, only their contraction with v ( highlighted.
TNs for Convolution Operations
We now demonstrate the elegance TNs for computing derivatives ($3.1), autodiff operations ($3.2). and approximate second-ordd -oder information ($3.3) by graphical manipulation. For simplicity, we exclude batching (vmap-ing like JAX [8]) and channel groups, and provide the diagrams with full support in $B. Table summarizes our derivations (with batching and groups). As warm-up, we identify the unfolded input and kernel from the matrix-multiplicee They follow by contracting the index patterns with either the input or kernel (Figures 2b and 2c),
view (Equations (1) and (2)).
X ka)
Σ Xen.i
(
(2)
kn. 62,02.22
$1.32
[A(W)] (Con (c.,,
n
W Ke
12.02,2
k1.
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 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, AN+1|Sw+:-1 St B|Ax ulsx+1/01A 1/0AA/,, w.r.t. input Aj collects all partial derivatives and addressed through indices S.+1 x 5 with S) an independent copy of Sy. Assume that A, only enters once in the tensor multiplication. Then, taking the derivative of Equation (4) w.r.t. [Ayls;
simply replaces the tensor by Kronecker delta 05,,55,
0 A AN+IISNH öA,yss
Σ
[Als ... (A) A)-11s,-1 II 8,,0 (.,((Ajttsss ".. [.(NNsx sES,
(6)
(.)
If an index i E Sj is summed, i Sn+1. we can sum the Kronecker delta 8,, effectively replacing all occurrences of by i'. instead i part of the output index, E S 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 Jw Y, this introduces structure: we re-interpret the two sub-diagrams in Figure 3a as matrices, compare with the Kronecker diagram from Equation (5) and use Figure 2b, we find [X] % Ic... for the Jacobian's ss matrix view [c.g. 16]. Figure 3b shows the input Jacobian Jx which is tensor view of A(W), expected from the matrix-vector perspective of Equation (2).
Differentiating TN is more convenient than using matrix calculus [44] as it amounts a simple graphical manipulation, does not rely a flattening convention, and therefore preserves the full index structure. The resulting TN can still translated back matrix language, desired. also simplifies the computation of higher-order derivatives (e.g. /a//)) since differentiation yields another TN and can thus be repeated. If tensor occurs more than once in a TN, the product rule applies and the derivative is sum of TNs with one occurrence removed.

---

Figure TNs of input-based Kronecker
factors for KFAC approximations of the Fisher/GGN (no batching, no groups). The unfolded input is shaded, only addi- tional contractions are highlighted. (a) l (KFC/KFAC-expand) from Grosse & Martens [27], (b) & (KFAC-reduce) from Eschenhagen et al. [23] (vectors of
kn
A1
n
1/2
X
X
+
KFAC-expand
(b) KFAC-reduce
ones effectively amount to sums).
3.2 Autodiff & Connections to Transpose Convolution
Although Jacobians useful, crucial routines for autodiff are vector-Jacobian and facobian-vecter products (VJPs, JVPs). Both are simple to realize with TNs due to access to full Jacobians. VJPs are used in backpropagation pull back tensor V () IR Ie ReCax0xx0x from the output to the input weight space. The VJP results v(x) E Caxix12 and ((W) E XCCaaxCaaKi
are
ViY) Y 01. Cr OX4, 11.
OY Y 01n.01.02
Σ Con,01,00
( Sue K.
Σ 01.20
V(V) 0on.01.02
aw el
Both are simply new TNs constructed from contracting the vector with the respective Jacobian, see Figures and 3d (VJPs are analogous). The input VJP often used to define transpose convolution [21]. In the matrix-muliplicationn perspective (Equation (2)), this operation is defined relative to a convolution with kernel W by multiplication with AWW T_ i.e. using the same connectivity pattern but mapping from the convolution's 's output to input space. The TN in Figure 3d makes this sharing explicit and cleanly defines transpose convolution."
3.3 Kronecker-factored Approximate Curvature
The Jacobian diagrams allow us to construct the TNs of second-order information like the Fisher/gen- eralized Causs-Newton (GGN) matrix and sub-tensors like its diagonal (SC). 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 [X] which requires large memory. State-of-the- libraries that provide access to KFAC [17, 51J also use this approach. Using TNs, we can often avoid expanding [X] explicitly and save memory. Here, we describe the existing KFAC approximations and their TNs (see $5.1 for their run time evaluation).
KFC :xpandd: Grosse & Martens [27] introduce a Kronecker approximation for the ker- nel's GGN, G s3 Ω @ r where r E CC XCuu XCur and the input-based factor S2 = [X] E K, K K2 (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 the concept of weight sharing, also present in convolutions. They identify two approxi- mations: KFACCexpand and KFAC-reduce. The former corresponds to KFC [27]. The latter shows similar performance in downstream tasks, but is cheaper compute. It relies the column-averaged unfolded input, i.e. the average over all patches sharing the same weights. KFAC-reduce approximates G ~ Ω @ I with I E CCa. and 52 1/10, 1 O. (1) Oa (Figure 4b; averaged over batch). For convolutions, this is arguably "more natural' approximation as it becomes exact in certain limits [23], in contrast to the expand approximation.
[X]] E K, Ka
Our approach enables to derive KFAC for transpose convo- lutions. 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 unfolded input. From this formulation we can immediately obtain KFAC through the weight sharing view of Eschenhagen al. [23]. The Kronecker factor requires unfolding the input similar 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).
KFAC for transpose convolution:
'Standalone implementations of transpose convolution require another parameter to unambiguously recon- struct the convolution's input dimension (see $D for how to compute N in this case).

---

Figure 5: TN illus- trations of index pat- tern simplifications and transformations. See $ D.3 for the math formulation.
n.I.K.S P.D.
U.
D.K.K.K......
Kernel-output swap
Dense convolution
V
nJ. 0,
N K.O.
(c) Down-sampling convolution
TN Simplifications & Implementation
Many convolutions in real-world CNNs use structured connectivity patterns that allow for simplifica- tions which we descrihe here along with implementation aspects.
4.1 Index Pattern Structure & Simplifications
The index pattern I encodes the connectivity of convolution and depends its hyper-parameters.
Along one dimension, n = n(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 DI). Its entries are
K,,
S, P)). = (o1-
,....,K and output size OII,K, = 1 Since n is binary and has size linear in 1,O, K, it is cheap to pre- compute and cache. The index pattern's symmetries allow for re-wiring TN. For instance, the symmetry of (k: D) and (o in Equation (7) and OT, , K, S. P, D) permits a kernel-output swap, exchanging the role of kernel and output dimension (Figure 5a). Rochette al. [58] used this to phrase the per-example gradient computation (Figure 3c) as convolution.
with i I.....I.o 1,.... O, (KI---KK-KK-
For many convolutions of real-world CNNs (see SE for hyper-paranee perem-pareeter study) the index pattern possesses structure that simplifies its contraction with other tensors into either smaller contractions or reshapes: Dense convolutions adjacent tiles of the input. Their index pattern's 's action can be expressed as a cheap reshape (Fig- ure 5b). Such convolutions are common in DenseNets [33], MobileNets [31, 60], ResNets [30], and Cos [42]. Inception V3 [69] has 2d mixed-dense convolutions that are dense along one dimension. Down-sampling convolutions sub-set of their input, and are used in R [30], ResNext101 [72], and WideResNet101 [73].
shared kernel size and stride, and thus process von-overlap
a larger stride than kernel size, hence only process
Their pattern contracts with tensor like that of a dense convolution with a sub-tensor (Figure 5c). 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 which to carry out the summations in 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 binary tree, but the underlying search in general at least #P- hard [14]. Fortunately, there exist heuristics to find high-quality contraction trees for TNs with hundreds tensors [32, 25, 13], implemented in packages like opt_einsut [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 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: N is sparse as only small fraction of the input contributes to an output element. For 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 []]...- Hence nnz
OK, which corresponds to sparsity of 1/1. Padding leads to more kernel elements that do not contribute to an
7

---

Figure 6: Run time ratios of TN (w/o simplifications) vs. standard im- plementation for dense convolutions of 9 CNNs. With simplifications, convo- lution and input VJP achieve median ratios slightly above 1, and the TN im- plementation faster for weight VJP, KFC & KFACCreduce. The code in Fig- ure corresponds to default, TN, and simplified TN KFC implementation.
TN versus PT (ognrittmic)
0.3
Forward
Inpur VJP
Weight VJP
KPC
KFAC-red.
output pixel, and therefore sparser n. For down-sampling and dense convolutions, we showed how IT's algebraic structure allows to simplify its contraction. However, if that is not possible, N contains explicit zeros that add unnecessary FLOPS. One way to circumvent this is to match TN with that of an operation with efficient implementation (like im2col, (transpose) convolution) using transformations like the kernel-ouppul swap or by introducing identity tensors to complete template, as done in Rochette et al. [58], Dangel [15] for per-sample gradients and im2col.
TNs offer principled approach for stochas- tic approximation via Monte-Carlo estimation to save memory and run time the cost of accu- racy. The basic idea is best explained on a matrix product C := AB Σn= [A] [B] n. with A E R B E IR N,O, To approximate the sum, we introduce a distribution over n range, then use Souum-roww-ampline
Approximate contraction & structured dropout:
[CRS, I] to form an unbiased Monte- Monte-Carlo approximation with sam- pled indices, which only requires the sub-matrices with active column-row pairs. Bernoulli-CR, samples without replacement by assigning a Bernoulli random variable Bernoull(zT,)) with prob- ability nn for column-row pair n. to included in the contraction. The Bernoulli estimator Č= SN/Z. [A), BB), with 2n ~ Bernoulli(TT). With a shared keep probability. Tu : this yields the unbiased estimator C" 1/rE m==......N A'B' where A' = AK and B' = KB with K = diag( ....... 2x) are the sub-matrices of A, B containing the active column-row pairs. CRS applies to single contraction. For TNs with multiple sums, we can apply individually. and also impose 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 in PyTorch [52]. The simplifica- tions 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 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 second-order method). This is intuitive as more frequently used routines have been optimized more aggressively.
an Nvidia Tesla T4 GPU (16 GB); see SF for all details. Due to
Impact of simplifications: While general convolutions remain unaffected (Figure F18d) when applying the transformations of 54, mixed dense, dense, and down-sampling convolutions consistently enjoy significant run time improvements (Figures F1 8a 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 's highly optimized weight VJP, also for down-sampling
"iissmm does yet support index un-grouping, so must reshape manually before and after.
8

---

103
Figure 7: Extra memory used by the standard versus our TN imple- mentation (simplifications enabled) of KFAC-reduce. Each point represents convolution from 9 CNNs, clipped below by MiB. TNs consistently use less memory than the standard imple- mentation (one exception), and often no extra memory at all. We observe memory savings up to 3 GiB.
Aux. TN [MIB]
10a
104
10
10
10a
10a
10
Auxiliary memory PT [MiB]
convolutions (Figure F21). For KFC, the median performance ratios 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 S2 and consistently outperforms the Py PyTocc implementation in the median when simplifications are enabled (Figure F23). The standard implementation unfolds the input, takes the row-average. then forms its outer product. The TN does not need to expand [X] in memory and instead averages the index pattern tensors, which reduces peak memory and run time. We observe performance ratios down to 0.22 x (speed-ups up = 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 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 52 (setting I x 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 rrecconditiones on ResNet18 and VGG19 on ImageNet- per-iteration time and peak memory on an NVIDIA A40 with 48 GiB of RAM. For SINGD, compare computing the Kronecker factors with the standard approach ( ('SINGD') via input unfolding versus TN implementation ('SINGD+TN'). Table 2 summarizes the results.
synthetic data (3, 256, 256) using a batch size of 128. We measured
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 factor of 2 from 32 GiB to 16 GiB. This enables using larger batches or more fre- quently updating the pre-conditioner, under- lining the utility of our approach for reduc- ing the computational gap between approx- imate second-order and first-order methods.
Table 2: Impact of our TN implementation on SINGD's run time and peak memory compared SGD.
Optimizer
Per iter. [s]
Peak mem. [GiB]
RooNeil
SGD
0.12 (1.0x)
(1.00)
SINGD
0.19 (1.7x)
4.5 (1.3x)
SINGD+TN
0.16 (1.3x)
3.6 (1.0x)
SGD
0.69 (1.0x)
(1.00)
VOKIO
SINGD
1.0 (1.5x)
32 (2.3%)
SINGD+TN
0.80 (1.2x)
(1.1x)
5.2 Randomized Autodiff via Approximate Contraction
CRS is an alternative to checkpointing [26] to lower memory consumption of backpropagation [50. 11, I]. Here, we focus 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 Gin- which eliminates many gradient entries as the index is part of the gradient. The randomized gradient would thus only train sub-tensor of the kernel per step. Oktay et al. [50], Chen et al. [11] I] 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 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
(b) Synthetie data
channel spatial
Normalized croo
2
0.2
0.4
0.6
0.3
0.2
0.4
0.6
0.8
Keep probability p
Keep probability p
Figure 8: Sampling spatial axes is more effective than channels both (a) real-world and (b) synthetic data. We take the untrained AlI-CNN-C [68] for CIFAR- 100 with cross-entropy loss, disable dropout, and modify the convolutions to fraction p of X when computing the weight gradient via Bernoulli-CRS. For mini-batches size 128, we compute the deterministic gradients for all kernels, then flatten and concatenate them into vector g: likewise for its proxy 9. CRS is described by (Pe.. Pia Piy), the keep rates along the channel and spatial dimensions. We compare channel and spatial sampling with same memory reduction, i.e. (p, 1, 1) and (1, /P, vp). To measure approximation quality, we use the normalized residual norm llg 90121/ ||ll|2 and report mean and standard deviation of 10 different model and batch initializations.
does not zero out the gradient for entire filters. In Figure 8 we compare the gradient approximation errors of channel and spatial sub- sampling. For the same memory reduction, spatial sub-sampling yields 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 performaanceffexibliit 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 tensor- valued functions and compiler optimizations through linear algebra and common sub-expression elimination. Phrasing convolution as e einsum, we allow it to be integrated into such frameworks, benefit from their optimizations, and complement them with our cnnvolution-speccii simplifications.
7 Conclusion
We used tensor networks (TNs), a diagrammatic representation of tensor multiplications, 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 syper-parameters, channel groups, batching, and generalization to arbitrary dimensions. All amount to simple einsum expres- sions that can easily be modified-c.g, perform stochastic backpropagation-aat and benefit from under- the-hood optimizations before evaluation. We complemented those by convolution symbolic simplifications based on structure in the connectivity pattern and showed their effectiveness to advance second-orde 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 transforma- tions (differentiation, batching, slicing, simplification) in terms of diagrams at less cognitive load without sacrificing rigour. We believe they are a powerful tool for the MI. community that will open up new algorithmic possibilities due to their simplicity & flexibility.
10

---

Acknowledgments and Disclosure of Funding
The author would like 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
[1] Adelman, M., Levy, K., Hakimi, I., and Silberstein, M. Faster neural network training with approximate tensor operations. In Advances Neural Information Processing Systems (NeurlPS), 2021.
[2] 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 neural information processing systems (NeurIPS), 2019.
[3] 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.
[4] Becker, S. and Lecun, Y. Improving the convergence of back proppaapaaationn learning with second-order methods, 1989.
Benzing, Gradient descent neurons and its link to approximate second-order optimization. In International Conference on Machine Learning (ICML), 2022.
Biamonte, and Bergholm, V. Tensor networks nutshell, 2017.
[7] Botev, A., Ritter, H., and Barber, D. Practical Gauss-Newton optimisation for deep learning. In Internationnal Conference on Machine Learning (ICML), 2017.
[8] Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. M. Leary, C., Maclaurin, D., and Wanderman-Milne, S. JAX: composable transformations of Nuthon+NumPy programs, 2018.
[9] Bridgeman, J. C. and Chubb, C. T. Hand-waving and interpretive dance: introductory course tensor networks. Journal of Physies A: Mathematical and theoretical, 2017.
[10] Chellapilla, K., Puri, S., and Simard, P. High performance convolutional neural networks for document processing. In International Workshop Frontiers Handwriting Recognition, 2006.
[11] Chen, J., Xu, K., Wang, Y., Cheng. Y., and Yao, A. DropIT: Dropping intermediate tensors for memory- efficient DNN training. In International Conference on Learning Representations (ICLR), 2023.
[12] Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Catanzaro, B., and Shelhamer, E. cudnn: Efficient primitives deep learning. 2014.
[13] cuQuantum development team, T. cuQuantum SDK: A high-perfommance library for accelerating quantum information science, 2023.
[14] Damm, C., Holzer, M., and McKenzie, P. The complexity of tensor calculus. computational complexity, 2002.
[15] Dangel, F. unfoldNd: (u=1,2.3) unfold (im2col) and fold (col2im) in pytorch, 2021.
[16] Dangel, F., Harmeling, S., and Hennig, P. Modular block-diagonal curvature approximations for feed- forward architectures. In International Conference on Artificial Intelligence and Statistics (AISTATS). 2020.
[17] Dangel, F., Kunstner, F., and Hennig, P. BackPACK: Packing more into backprop. In International Conference on Learning Representations (ICLR), 2020.
[18] 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.
[19] Dangel, F. J. Backpropagation beyond the gradient. 2023.
[20] Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace redux effortless bayesian deep lcarning. In Advances Neural Information Processing Systems (NeurIPS), 2021.
[21] Dumoulin, V. and Visin, F. A guide to convolution arithmetic for deep learning. 2016.

---

[22] Elsayed, M.. Farrahi, H. Dangel. F., and Mahmood, A. R. Revisiting scalable hessian diagonal approx- imations for applications in reinforcement learning. In International Conference Machine Learning (ICML). 2024.
[23] Eschenhagen, R., Immer, A.. Tumner, R. E., Schneider, F., and Hennig. P. Xronecker-factoree approximate curvature for modern neural network architectures. In Advances Neural Information Processing Systems (NeurIPS), 2023.
[24] Goldfarb, D., Ren, Y., and Bahamou, A. Practical quasi-newton methods for training deep neural networks, 2021.
[25] Gray, J. and Kourtis, S. Hyper-optimized tensor network contraction. Quantuon, 2021.
[26] Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentia- tion. SIAM, 2008.
[27] Grosse, R. and Martens, J. A foctecee-ffactored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning (ICML), 2016.
[28] Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A.. Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E.,, Hubinger, E., Lukošiūte, K., Nguyen, K., Joseph, N., McCandlish, S., Kaplan, J., and Bowman, S. Studying large language model generalization with influence functions, 2023.
[29] 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 (NeurlPS), 2019.
[30] 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.
[31] Howard, A. G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., and Adam, Mobilenets: Efficient convolutional neural networks for mobile vision applications. 2017.
[32] Huang, C., Zhang, F., Newman, M., Ni, X., Ding, D., Cai. J., Gao, X., Wang, T., Wu, F., Zhang, G., et al. Efficient parallclization of tensor network contraction for simulating quantum computation. Nature Computational Science, 2021.
[33] Huang, G., Liu, Z., Van Der Maaten, L., and Weinherger, K. Q. Densely connected convolutional networks. In IEEE conference computer vision and pattern recognition (CVPR), 2017.
[34] Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks, 2020.
[35] Jastrzebski. S.. Szymezak. M.. Fort, S., Arpit, D., Tabor, J.. Cho K., and Geras K. The break-even point on optimization trajectories of deep neural networks. International Conference on Learning Representations (ICLR), 2020.
[36] Krizhevsky, A., Sutskever, L,, and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Syxtems (NeurIPS), 2012.
[37] Laue, S., Mitterreiter, M., and Giesen, J. Computing higher order derivatives of matrix and tensor expressions. Advunces in Neural Information Processing Systems (NeurIPS), 2018
[38] Laue, S., Mitterreiter, M., and Giesen, J. A simple and efficient tensor calculus. In AAAI Conference on Artificial Intelligence, 2020.
[39] LeCun, Y., Boser, B., Denker, S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1989.
[40] Lin, W., Duruisseaux, V., Leok, M., Nielsen, F., Khan, M. E., and Schmidt, M. Simplifying momentum- based riemannian submanifold optimization. 2023.
[41] Lin, W.. Dangel, F., Eschenhagen, R.. Neklyudov, K., Kristiadi, A., Turner. R. E., and Makhzani, A. Structured inverse-free natural gradient descent: Memory--fficient & numerically-stable KFAC. In International Conference on Machine Learning (ICML), 2024.
[42] Liu, Z., Mao, H., Wu, C., Feichtenhoo,, C., Darrell, T., and Xie. S. A convnet for the 2020s. IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), 2022.
12

---

[43] Ma, Ye, J., and Solomonik. E. Autohoot: Automatic high-order optimization for tensors. In International Conference on Parallel Architectures and Compilation Techniques (PACT), 2020.
[44] Magnus, J. R. and Neudecker, H. Matrix Differential Calculus with Applications in Statistics and Econometrics. Probabilistics and Statisties. 1999.
[45] Martens, J. Deep learning via Hessian-free optimization. In International Conference on Machine Learning (ICML), 2010.
[46] Martens, J. New insights and perspectives on the natural gradient method, 2020.
[47]
Martens, J. and Grosse, R. Optimizing neural networks with Krotecker-factored approximate curvature. In International Conference on Machine Learning (ICML), 2015.
[48] Martens, J., Ba, and Johnson, M. Kronccker-factored curvature approximations for recurrent neural networks. International Conference on Learning Representations (ICLR), 2018.
[49] Novak, R.,. Sohl-Dickstein, J., and Schoenholz, Fast finite width neural tangent kernel. In International Conference on Machine Learning (ICML), 2022.
[50] Oktay, D., McGreivy, N., Aduol, Beatson, A., and Adams, R. P. Randomized automatic differentiation. In International Conference on Learning Representations (ICLR), 2021.
[51] Osawa, K., Ishikawa, S., Yokota, R., Li, S., and Hoefler, T. Asdl: A unified interface for gradient preconditioning in pytorch, 2023.
[52] 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., Steinee, B., Fang, L.,, Bai, J., and Chintala, S. PyTorch: An imperative style, gegh-performacc deep learning library. In Advances in Neural Information Processing Systems (NeurIPS). 2019.
[53] Penrose, Applications of negative dimensional tensors. Combinatorial Mathematics and its Applications, 1971.
[54] 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.
[55] 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.
[56] Ren, Y., Bahamou, A., and Goldfarh, Kronecker-factored quasi-newton methods for deep learning. 2022.
[S7] Rigamonti, R., Sironi, A., Lepetit, V., and Fua, P. Learning separable filters. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2013.
[58] Rochette, G., Manoel, A., and Tramel, W. Efficient per-example gradient computations in convolutional neural networks, 2019.
[59] Rogozhnikov, A. Finops: Clear and reliahle tensor manipulations with einstein-like notation. In Interna- tional Conference on Learning Representations (ICIR), 2022.
[60] Sandler, M., Howard, A., Zhu, M., Zhmoginov, A., and Chen, L.-C. Mobilenetv2: Inverted residuals and linear bottlenecks. IEEE conference on computer vision and pattern recognition (CVPR), 2018.
[61] Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep lincar neural networks, 2014.
[62] Schneider, F. Balles, I.., and Hennig, P. DeepOBS: A deep learning optimizer benchmark suite. In International Conference on Learning Representations (ICIR), 2019.
[63] Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 2002.
[64] 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.
[65] Singh, S. P., Hofmann, T., and Schölkopf, B. The hessian perspective into the nature of convolutional neural networks. 2023.
13

---

[66] Smith, D. G. A. and Gray. opt einsum A python package for optimizing contraction order for einsum-like expressions. Journal of Open Source Software (JOSS), 2018.
[67] Smith, S. W. The scientist and engineer's 's guide to digital signal processing. 1997.
[68] Springenberg, T., Dosovitskiy, A., Brox, T., and Riedmiller, M. Striving for simplicity: The all convolutional net, 2015.
[69] Szegedy, C., Vanhoucke, V. loffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In IEFE conference computer vision and pattern recognition (CVPR), 2016.
[70] Tai, C., Xiao, T., Zhang, Y., Wang, X., et al. Convolutional neural networks with low-rank regularization. 2015.
[71] Vaswami, 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.
[72] Xie, S., Girshick, R., Dollár, P., Tu, Z., and He, K. Aggregated residual transformations for deep neural networks. In IEEE conference computer vision and pattern recognition (CVPR), 2017.
[73] Zagoruyko, S. and Komodakis, N. Wide residual networks. 2016.
[74] Zhang, F. A parallel tensor network contraction algorithm and its applications quantum computation. 2020.
14

---

Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods (Supplementary Material)
A
Limitations
16
Visual Tour of Tensor Network Operations for Convolutions
16
B.1
Convolution & First-order Derivatives
16
B.2
Exact Second-order Information
16
B.3
Kronecker-factored Approximate Curvature (KFAC) for Grouped Convolutions
18
B.4 Kronecker-factored Approximate Curvature (KFAC) for Transpose Convolution
19
B.5 Further Operations & Extensive Overview
20
C Exact Second-Order Information
24
D Implementation Details
25
D.I Index Pattern Tensor Computation for Convolutions
25
D.2 Index Pattern Tensor for Standalone Transpose Conveittoo......................
25
D.3 Details on Index Pattern Simplifications
25
Convolution Layer Hyper-parameter Analysis
27
Run Time Evaluation Details (GPU)
30
F1
Protocol & Overview
30
F2
Forward Pass
32
F33
Input VJP
35
F44
Weight VJP
38
F.5
KFC Factor (KFAC-expand)
41
F.6
KFAC-reduce Factor
44
G
Memory Evaluation Details (CPU)
47
G.1
Theoretical Empirical Analysis for KFAC-reduce Factor
H Miscellaneous
50
Example: Associativity of Tensor Multiplication.
50
H.2 Example: Matrix-matrix Multiplication as Tensor Multiplication
50
15

---

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 advantageous optimization approach which opt_einsum does not account for. The second-order quantity TNs from &C and $3.3 contain such sub-expressions, for instance [X] and
1 [X] in KFAC-expand and KFAC-reduce, and ( in the GGN quantities from Figure C16. The efficiency of CSE depends on how costly the shared tensor 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 faster run times interesting to integrate CSE into the contraction path optimization and develop a heuristic to choose a contraction path, for instance based weighted sum of FLOPs and memory.
GPU. It would be
No index slicing: We mention index slicing technique to reduce peak memory of, and distribute, TN contractions. However, our implementation does not use index slicing, although there are packages like cotengra [25] with interface similar to opt__iinnee 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.
We did not experiment with index
B Visual Tour of Tensor Network Operations for Convolutions
Here, we extend the presented operations with batch axis and allow for grouped convolutions.
B.1 Convolution & First-order Derivatives
Adding batch dimension (vmap-ing): Adding batch axis to all presented operations is trivial. We only need 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 Cj input channels into G groups of size Önn == Ca/G, then perform independent convolu- tions per group, each producing Cout Ga//G output channels which are concatenated in the output. Each group uses a kernel ofof size Cunt x Oin x K₁ x K2. These kernels are stacked into a single tensor W E CC. K2.22 such that [W] (g.)..) W g. To support groups, we thus decompose the channel indices into := (Cin- g) and Cout == (Cout: g). For the forward pass this yields the grouped convolution (without bias)
Y Yig (mn))))).1.0 En X(g.a) $1.01 .8 k2.02,82
(2)
W (g. (. k1
(B8)
Figure B9a shows the batched version 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 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 E kKIxKJ be matrix of K horizontally and vertically concatenated blocks
16

---

01
01
01
(1)
1
8
02
(3
h/2
2/2
i2
k2
77.
Cin
W
Cout
Gn
out
Cout
n2
in
W
Cout
(a) Convolution (no bias)
(b) Weight Jacobian
Y (9.(9.9m))001:00
Yn m. Pan).
(c) Input Jacobian
a of P Paa My Mg
i YY (9.
Equation (B8)
Oxn dat Maaiii
1
81
02
(2)
A/2
k2
g
X
n2'
Y
W
Cout
Cout
(d) Weight VJP
(e) Input VJP V(X) n' no' E,p.,,
((W)
0 Y (pm.)). Vn Va 0g aw And M, Pig-Bomm 01:
8Y (Ym.(rrau))) va 02 0X (gg 1 12
Σ
01:02
Figure B9: TNs of the (a) forward pass, (b, c) Jacobians, and (d, e) VJPs with batch axis and channel groups. They generalize Figures 2 and 3 from the main text. For the VJPs, the Jacobians are shaded.
kn
k2
02
Gm
Gin
X
X
X
71
Cout
Caul
(YY
Y
Y
n'
(a) GGN diagonal
(b) GGN Gram matrix (empirical NTK)
Figure B10: TNs (a) the GGN diagonal and (b) the GGN Gram matrix with batching and channel groups. They extend Figures C16b and C16c from the main text.
A dli. RIXJ_ K; 1.... K. We can extract the diagonal blocks by restoring the sub-structure,
17

---

gK1
k1
h
9K2
0-2
222
02
9C,
Can
S
Y
Cour
Fout
Figure B11: TN of a GGN mini-block diagonal without batching and channel groups.
0
k1
k
ky
11
kz
2/2
02
0
0/
n
1/2
82
X
X
X
X
Cin
Can
(a) KFC/KTAC-eXpand factor
(b) KFFAC-reduce factor
Figure B12: TN diagrams of input-based factors in Kronecker approximations of the GGN for convolutions with batching and channel groups. They extend Figure from the main text.
then taking the diagonal along the K-dimensional index.
A (R.A)
(ke,
(Ki,)
We apply this procedure the GGN from Figure Cl6a. Assume we want to divide the output channel, imput channel, and spatial dimensions into Gc... Gc., : GK, , Gk2 groups A group will thus be indexed with tuple (gc.u> 9C.: 9K, (K,) and the corresponding GGN block will be of dimension Cous/ / x Cin/ Gc, x K2/GK, x K2/GKz x Coa//G... X Cn/GG. Gc. x K2/Gk, K2/G and correspond to the GGN for [W]] (Ic..:))
((9x,,)))) This process of un-grouping the output dimensions, then taking the diagonal along the group indices, is illustrated in Figure B 1. Note that we choose Gc.. Cout: GC., Cin. Gk, K1, GK2 K2, each block will
single number and hence recover the GGN diagonal from Figure Cl6b. If instead we Gc... Gc.. GK, GK2 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 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, can then apply established the KFAC approximation for convolutions without groups. For a group g we have the kerell W
X 9 (g...... and the unfolded input of its associated input channels,
[X,] [X]
[g.) (or [Xn.g] [X.]ig:)):: [[X (...))....] in the batched setting).
Applying the regular KFC approximation to the kernel of group this yields the Fisher approximation 22 @ T with r, C and
KFC/KFAC-expand for grouped convolutions:
18

---

R = I/N [ 1hnggXXXXXX,,. [[ E RaKKi where Xn.g is the input tensor for sample n and group g (remember the index structure (2) Figure B12a shows the diagram for [NRg]e
Proceeding in the same way, but using the unfolded input averaged over output locations, we obtain the Fisher approximation S2 & I with ro E IR (RC XCm and S2, 1/N(O;0,)² Σn 12, o,
KFAC-reduce for grouped convolutions:
g)) R.R., K for the kernel of group g. Figure 12b shows the diagram for (1(0002)0211)
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 = Caa to an output space y = x1x02, The convolution has haper-parameters
K1,2 S12, P1,, D12 with index patterns ( n(h, P1, E R4x00xx and n (12, K2 S2 P2, D2) 22x00x,
Transpose convolution as matrix multiplication: Transpose convolution maps Y E y into an X E X. In ML frameworks like PyTorch, its kernel W is stored as Cout Cin x K1 x K₂ tensor. The relation X W xT Y where *т denotes transpose convolution is given by Figure 3d,
Cat
K
K2
O₂
Σ Cos
11.01.1 K2.02.8g Ko Won
(22)
(B9)
=1=
Our goal is turn the express the above matrix multiplication. To do that, first define the matrix reshape X of via X 122x such that [X]
(o..)(1111)) X Next, we consider a
transposed kernel W of W with changed order of the first two indices, i.e. W BR x such that
W Cn. Eo We Cas
(B10)
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 E RCx Cau such that [ (ea. matrix multiplication
W ka we can express Equation (B9) as
X = W[
(B11)
where | E RC Ka Ka is the transpose input to the transpose convolution (note that [. * [lr!)
0,
( (eer k1, (11.b)))
Σ
Σ
(
(22) 1g.00,62
Y Yoo.o1000
(B12)
To the best of our knowledge there is no API for [-]r in existing ML frameworks. Our approach can provide a simple and efficient implementation of [.] through TN shown in Figure B13aa which corresponds to Equation (B12). As Equation (B11) is of the same form as Equation (1), it is now straightforward 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) ~ Ω & F, with I CCxC computed from backpropagated gradients, and the input-based Kronecker factor
Ω [V E Ka
(B13)
See Figure B13b for the corresponding TN.
19

---

0
02
04
3
Gost
Y
Cost
Cout
Soul
(a)
Transpose-
unfolded input
(b) KFAC-expand
(c) KFAC-reduce
Figure B13: TNs for extending KFAC transpose convolutions (no batching and groups).
For KFAC-reduce, we have G(w) l2 & r, with I RC computed from backpropagated gradients, and the input-based Kronecker factor
KFAC-reduce:
Ω = (1122 (17,, (IT,, Ro
(B14)
See Figure B13c for the corresponding TN.
With batching and groups: In the presence of groups, we have per-group kernels W, XixKs Ka as well as per-group transpose- unfolded inputs [VgIr M.rga.. -[Maa) R 12 Each group corresponds transpose convolution in itself. With batching, we have additional leading batch dimension, i.e. [Yn.gIT- Applying the same steps from above, we can define the KFAC approximation for the GGN w.r.t. the flattened per-group kernel wy of Wg.
[W) (g:))s: E R (a////xa///xK K2 and Wg E R. Calaxxo///xx x X
For KFAC-expand, we have G(w,) Zg Ig. g. with Ig E C computed from backpropa-
gated gradients, and the input-bat based Kronecker factor
N
2
o oll CafoKk Kg Ka/Caa K2
For KFAC-reduce, we have G(w,) Re S2, 2 G with F, RC//axaa/o computed from backpropa- gated gradients, and the input-based Kronecker factor
N
2(O,02) DY (1 IY (nglT))
R Ko
B.5 Further Operations & Extensive Overview
Consecutive convolutions: We chain two, or more, convolutions into a single TN diagram (Figure B 14) to obtain 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 (() CaC K1 K2 and v(x) E x 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 without summing over the batch axis which amounts 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 (xx) RCux K2 K2 in input space to obtain the JVP, with perturbations ((XI) (0.02xC.K,Kk Ka to obtain the VJP.
20

---

0/2
(7(2)
102
1/g
X
W
Coul
W
Cout
Figure B14: TN two consecutive convolutions without groups and without batch axis.
Approximate Hessian diagonals (HesScale/BLL89) Becker & Lecun [4], Elsayed al. [22] proposed approximate procedures for the Hessian diagonal which cost roughly a gradient. They can be understood modifications of the Hessian backpropagation equations from Dangel et [16].
Consider a layer with input æ, output y, and weights w inside a sequential feedforward neural network (for convolutional layer, these correspond to the flattened input, output, and kernel). To compute per-layer Hessians of loss l, each layer backpropagates its incoming Hessian V according to [16]
de Vayi:
V2e = (J+B) ( ( + Σ
dy:
(B15)
Oe
- V (Jwy) +
Σ
V ys.
By:
The scheme of [4, 22] imposes diagonal structure on the backpropagated quantity. layer receives a backpropagated diagonal d such that diag(d(v) V and backpropagates it according to Equation (B 15), but with a post-processing step to obtain diagonal backpropagated quantity,
d( = diag (Joy)) T diag(d (Jau))) + diag
Σ
Vays dy:
Ol
(B16)
www) = diag (Jwy)) diag(d (Jww))) + diag
Σ
OR
V2
Vi
Oy.
where d)) (())
V2 and (iag( V is an approximation the Hessian diagonal.
For convolutional layers, which are linear in the input and weight, the second summands are zero due Vžn V P.y... The first terms of Equation (B1 require (i) embedding diagonal vector into a matrix, (ii) applying MJPs and JMPs, and (iii) extracting the result's diagonal. Those can be expressed single TN. We show the diagrams in Figure B 15, using tensors rather than their flattened versions, that (ar, y, w. (() (v)) w)) + (X, Y, W. D ( ( D (W))
21

---

02
11
22
122
62
(V)
Cout
W
W
D
Cout
Cin
(a) HesScalc/BL89 input backpropagation
(b) HesScale/BL89 weight backpropagation
k1
02
k2
Cin
71
(Y)
Cust
Gi
(c) HesScale/BL89 input backpropagation (+ batch, (d) HesScale/BL.9 weight backpropagation (+ batch, groups)
groups)
Figure B15: TN diagrams for HesScale/BL.sg .89 [4, 22] backpropagations through convolutional layers to approximate the Hessian diagonals D (x) (W) JMPs and MJPs are shaded. (a, b) show the simple versions without batching and without channel groups. (c, d) include batching and channel groups.
22

---

Table B3: Extensive list of convolution and related operations (extension from Table in the main text). All operations consider two spatial dimensions and support batching and channel groups. Generalization other dimensions follow by introducing more spatial indices i_3, o_3, ... and kernel indices k_3,
Operation
Operands
Contruction string (einops [59] convention)
Convolution (no bias) [29]
X, m W
ta cutol
11 of xi,
c.005) c.io
c.ost) o2*
Unfolded input (in2col. [X])
X n m
ta clin 12, 11 ol lit, 12
(o.sa ki X2) (ol o2)
Unfolded kemel (Treplitz)
n()) n(9), W
ris ot lety ed nuoot cutn k2 (ruonb of o2) (euth 11 120*
Folded oulgut (col2in)
Y, Y. (
ta (g eLous) oo o2, 11 01 lat, 12 02 x2 (g (gin) 11 12"
Transpose unfolded input ([
Y, .
"w. (g cLous) o1 oz, 11 01 x1, 12 02
Kg cuin k1 x2) 11 12"
Convolution weight VJP
m() m(2) you
fon (g mukn) 12, st lit 12 k2, (g SLoot)
cLout cutn kt k2"
Convolution input (transpose convolution)
W. ( ( v00 (0) (()) v0
"gg cLout? c_in ki 22, 11 ki, 12, [g (Lout) [g cLous) c,in k2"
(g c.1a) 12"
Convolution weight sample/hatched) [58]
VJP
(per-
tw. (g cuinl 12, 11 01 ai, x2, (g cLost)
Convolution weight JVP
X m( yow
ta (g c,in) 11 12, 11 01. ki, 12 02 X2, (g cLout) c,io 21 12 n. (g cLout) 01 c2* "C cueot) outn 11 12, 11 k1. 02 k2. <g cu.n)
23
Convolution input JVP
W n yoo
(g cLoct) oz
in2col VIP
n(), n v(D)) n(), ( yoo
12) (ot
12"
in2col JVP
his
320 (o1
KFC/KFAC expand (27, 23]
X x (. x m
ts [g ccinl 11 12,
11 o1 at,, 12 02 X2, (g e.knn) 11, 12., R1_ o1 kt., 12, 02 22.
+> (c_in k1 121
(e.in. 141., X2_1"
reACcreduce [23]
X. k0.,n p,),,, ", x.n m
ta (g c_inl 11
12,
12 at, 12 02 X2, n. (g 4,3n.) 11. 12., 11. o1. I4.,, 12, 22, 12.
>> K (c ut x2) [e.
44. x2_
KFCKKFAC-eApand transpose convolution
v.n m()) Y, (
ta (g c_ous)
11 01 21,
k2, fg cLout.) 01. 02., 11, of Rd., 12, 12.
ts x cuu x:
(oLo0s EL. 12.)"
KerC-reduce for transpose convolution
Y. n Y.. (
k2. (g (Lont o1. o2., 11_ ot. KL. 12. o2_ X2. (c.ost. kL. k2.)"
11 kt,
GGN Gram/empirical NTK matrix [18,51, 49]
n m(a) s(,, x. x
(2) (())
ts (g s_in] 12, 11 kt.
k2, (g. c.ous] 02, m_ (s sLin) it, 12., i1, st_ 21, s_in)
12, 02_ k2. 6- n- (g c_out) ol- oZ >> (o m) (c_ n_)"
GGN/Fisher diagonal [17,
m X. n n s00 n m sm
(g clous]
11,
12,
u2""
GGN/Fisher diagonal (per sample/bacchee)
X, n m S0,,,.''
"2 ig clinl 12, 11 21, 12 02 X2, [g c.ous) 01 02, (g c.inl 11, 12., 11, ot, ki.
12, 02. 82, [g c.out)
01. 02.
c.005) c,in K2"
Approximate weight Hessian diagonal [4, Appoximmae input Hessian dingonal [4, 22] Approximate weight Hessian diagonal (per-
X. o n D X. ( ((
(g c_in)
(g. eLmot) 62, (ute) 19L 11_
[g 000) X1
w, n n D00 w. m0)n m
cLout)
kl
kt. o2 k2, ig ecout) 02, (e moout) cain Xt_ k2., it of kl..
x2_
c.in) 12"
x. ( Ln a DM x, m
ta (g clin) 11 12, 11 o1 125. 12 o2 X2, (g cLoot) ef 69, a (g (.iz) 1t, 12., 11. o1 ki, 12, 02 22 (g e_mmt) 1_18 RI 22*
sampletbatched)

---

1
X
8
Fex
5
(a) GGN/Fisher
(b) GGN/Fisher diagonal
(c) Gram (empirical NTK)
Figure C16: TN composition and sub-tensor extraction for second-order information. Weight MJPs from Figure 3c are shaded. (a) exact and (b) diagonal of the kernel's GGN (the same applies to structurally similar matrices like the gradient covariance [35]). (c) TN of the GGN Gram matrix.
C Exact Second-Order Information
Here, we look at computing second-order information of a loss w.r.t. to the kernel of convolution. Its computation can be phrased as backpropagation with 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 (æc, t) and its loss f(w) = t) where f = (x) E IRC is the pre- diction of CNN with convolution with flattened kernel w and flattened output y (deriva- tions carry over to a batch loss). The kernel's generalized Hauss-Newton (GGN) matrix [63] G(w) wJ ) [))(((( E Cu K,,,C.C,K,K is 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]. is the self-outer product of a backpropagated symmetric factorization S(v) (J,f) (()) O O2xC of the loss Hessian, V/((f.)) .)) ((()) During backpropagation, the convolution extracts information about G(w) (Jwy) ( Jwy..
In TN notation, this easy 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) FCou x 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 51D, or mini-block diagonals (Figure B11 [16, 3J). This also removes redundant computations compared to computing, then slicing, the matrix. The same ideas apply to the GGN Gram matrix (S(")) ((w) e RCxC (Figure C16c). It contains the GGN spectrum [18] and is related the empirical NTK for square loss [49].
K2 x C. XK K22 of G(w) (Figure C16a). The GGN often
24

---

D Implementation Details
Here we present details on the index pattern computation, and additional transformations.
D.1 Index Pattern Tensor Computation for Convolutions
Algorithm DI lists pseudo-code for the index pattern computation from the convolution hyper- parameters K, S, P, D, and the spatial input dimension I, that is (11, K,S, S, P. D). Unlike in the main text, we use O-based indexing which is more common in numerical libraries. For self-consisteney, re-state the relation of the sperr-paraeeters output dimension from [21, Relationship 15),
D,,,K,S.PD))-1 =1+ 2P K - (K- 1)(D-)
(D17)
S
Algorithm DI Computing the convolution index pattern tensor NT for a spatial dimension.
Require: Input size I E N+, kernel size K' E N+, stride S E N+ padding P E Nt. dilation D E N+
O - 1+2P K- (K(K-I)(D-)
Compute output dimension [21, Relationship 15]
x 0xxoxk
Imitialize index pattern tensor
for o = 0,..., 0 - 1, k = O,... ...... K do
b Use O-based indexing!
i < kD + oS P
> Reconstruct contributing input element & Check in bounds
#
then
11,o, t-1
end if end for return Index pattern tensor 17 E (0, (xxox
D.2 Index Pattern Tensor for Standalone Transpose Convolution
Although transpose convolution is defined w.r.t. reference convolution with hyper-parameters S,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 CCnn xO, x02 (the convolution's s output space) to IR C x/x (the convolution's 's input space). For convolution with
S> 1, we cannot infer I from O, K,S, P, D, as multiple Is map to the same (T 2P K (K (K-1)(( - 1)) mod 5 * 0 (see the floor operation in Algorithm D1). We need to either supply directly, or the remainder
A - K - (K 1)(D - S(0 1)
(often called output_ppaddin to make I unambiguous. Then, we compute
I (0 1)5 2P + K + +(K 1)(D - 1) + A.
(D18)
to get A) and call Algorithm D1 to obtain ((T(O,4), S.
D.3 Details Index Pattern Simplifications
In the following, we will assume the absence of boundary pixels that don't overlap with the kernel, that is
I - (K-1 mod 5=0,
(D19)
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- aeeem-erreeter analysis of real-world CNNs (SE), 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 into tuple index (i, k) of size 1/K K:
[1, K, K, 0, 1) [n(1, K, K, O, 1) 62,0o
Point-wise convolutions (K = S 1) are special case with pattern [N 1, 1, 0, 1)]..o,k ò
25

---

Point-wise convolutions with K = 5 = 1 are common in DenseNets [33], MobileNets [31, 60] and ResNets [30]. Inception V3 [69] has 2d 'mixed dense' convolutions that are point-wise along one spatial dimension. ConvNeXt [42] uses dense convolutions with K = = E (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 summed, all participating tensors can be pruned to remove the explicit zeros. Assume mod S = 0. Then, pruming amounts to un grogpuu i into (i', 8) of size 1/s x S, narrowing s to K entries, and grouping back into an index i of size K1/s. After pruning, the index pattern represents a dense convolution with input size KI Kijs, kernel size K, and stride K. In contraction with some tensor V,
M K,,,K,S> K,0, EE [N(K//s K, K, O, 1)
with sub-tensor [V] V........ (M....(".>)........ V.....e..)................. (:.:............................ where 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
in R8 [30], ResNext101 [72], and WideResNet101 [73]. Those convolutions discard 75 % of their input! Knowledge that an operation only consumes 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 ker- nel and output dimensions in an index pattern yields another index pattern with same input size, kernel size O K, S, P, D), and swapped stride and dilation:
n K, S, D,D. [n(I, D,P.
This transformation is easy to see from the symmetry of (k, D) and (o, in Equation (7) and O(I, K, S, P. D). 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 commensions--esentiee [58] use this to phrase the computation of per-example gradients convolution.
the TN for convolution from Figure with kernel V. Rochette et al.
SD.3 presents more properties of N based the sub-sampling interpretation of stride and dilation along the output and kernel dimensions. We also provide 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) sub- sample non-strided convolutions along the oufput dimension, ignoring all but every Sth output [21]. In other words, [T, K, P, D))i,,, = [n(I, K, 1, D(o,)+((- or, in tensor notation (F)s
denotes slicing with steps of S),
n (L,K,S,P,D) = L,,,
Transformation DS (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 Dth kernel element. In other words, [N K, S, S)
[N K + (K- - - -)) P, 1)),, D(k-1) or, in tensor notation,
nII,
[N - -K 1)
Transformation D6 (Transpose convolution as convolution) Assume Equation (D19), Consider non-strided (S = 1), non-dilated (D convolution with index pattern (TT K, K,1,P,1) and output dimension OII K, P, 1). Transposing the spatial dimensions and flipping the kernel dimension yields another index pattern with modified padding P' K P- 1. other words, for all i i= 1, I,..., K, I....,O
[N(1 K, 1, P,1. [N(O, P'L)), 1)
26

---

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, Inception V3, MobileNetV2, ResNext101). Regarding the seper-parameter, we make the following observations:
Many CNNs not use bias term. This is because the output those layers feeds directly into batch normalization layer, which is invariant under the addition of bias term.
All investigated convolutions use default dilation.
Group convolutions are rarely used. MobileNetV2 and ConyNeXt-base (Tables E4g and E4i) group convolutions that interpret each individual channel as group. ResNext101
(Table E4f) uses group convolutions that interpret a collection of channels as 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.
Inception V3 (Table E4h) uses -wo-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 fraction of the input. The index pattern can be simplified.
27

---

Table E4: Hyper-parameters of convolutions in different CNNs. For convolutions with identical hyper-parameteny we only show one instance and its multiplicity.
(a) 3c3d CIFAR-1O (3, 32, 32)
Nome (count)
shape
Output
Kerne
Stride
Padding
Dillaion
Groups
cocot.0 (1)
13, 32)
(64, 28, 25)
15.3)
1..11
(0.
(3.3)
Yes
Geteral
ea0x2.0 (1)
(64, 14, 14)
(96, 12, 12)
(3.3]
01
(6.0)
(1.3)
You
Gcteral
00nx3.1 (1)
196, 8)
(128, 5)
(3.3)
0.
(6,0)
(1.3)
Yos
Geteral
(b) 2c2d, Fashion MNIST (1, 28, 28)
tosunt)
Ingut shape
Outges shage
Stride
Padding
Gropp
Type
ovnvt (1)
(I. 53, 12)
(32, 28, 250
(5.5) (5.5)
11.11
(5. 09
(1.1) (1.3
Yes
General
c0092.3 (1)
(22, 18, 18)
(64, 34, 14)
(1.1)
(0,0)
Yes
Geteral
(c) AlI-CNN-C, CIFAR- 100 (3, 32, 32)
Name (ovunt)
Iaput shape
Output shupe
Kernel
Stride
Padding
Dilution
Groups
Type
cancl I1
[3, 14)
(96. 12, 321
13, 9)
(1.
(0. ())
[1. 1)
Yox
General
comv2.1
(96. 34, 34)
(96. 32, 32)
(3.3)
1)
(0.0)
Yex
Graznd
conv5.1
(96.
(96. 16,
03.33
(2.2)
General
convi.1
(96.
(192. 166
(1.1)
00,
Generl
co215.1
(192,
(192,
Geneal
comvt.1
(192,
(192.
(2.2)
Genera
conv7.0
(192.
(3.3)
(0,
General
[192.
coa19.1
(192,6.6)
(d) AlexNet, ImageNet (3, 256, 256)
Name (connt)
Input shage
Output shape
Kernel
Stride
Padding
DNlation
Groupe
Blas
Pype
(11.11
3 (3)
(64,
(192, 31,31)
(5.5)
(1.1)
(2.2)
Yes
General
features (1)
[192, 15)
(384, 159
(3.3)
(.
(1. 1)
Yes
General
Reaturce& ())
(384,
(256. 15)
(3.3)
(.
(1.1)
(1.1)
Yes
Guerrl
festuresID
(256, 15,
(256, 15,
(3.3)
0..1)
(..
(1.1)
Yes
Ganeral
(c) ResNet18, ImageNet (3, 256, 256)
Name (comnt)
Ingut
Ontput shope
Kernel
Stride
Padding
INtation
Groups
Blas
Type
com'1 (1)
256, 156)
(M4 128, 128)
(7.7)
(2:2)
(1. 1)
No
Grneal
Sperl. (4)
(64, 54, 64)
(64, 64, 64)
(3.3)
(1.1)
(1.
(1. I)
1
No
Geneal
Decrrr0 (1)
(64, 54, 64)
(128, 32, 32)
(3.3)
(2.2)
(1.3)
(1. 3)
1
No
Generl
2per220cmm2233 (3)
(128, 32, 32)
(128. 32, 32)
(3.3)
[1.1]
().
(L.
No
Geneal
(64, 64, 64)
(128, 32, 32)
1.
(2.2)
(9.
(L.I)
1
No
Dowp
Sayer30comrl (t)
(128, 32, 52)
(256. 16. 16)
(3.3)
(2.2)
(1.b)
(L.
1
No
Gencal
bper3Dcmm21 (3)
(258. 16. 16)
(258. 16, 16)
(3.3)
(1.1)
(1. 11
(1.1)
No
Grncal
(128. 32, 32)
(258, 16, 16)
(1. 13
(2.3)
(0. 01
(1.
No
Down
kggel ())
(25%) 16, 160
1512. 8. 91
(3. ()
(2.2)
(1.
(1.1)
No
Groeal
1penl (?)
(512. 8.
1512, 8.
(3.
(1.1]
().
(1.1)
No
Groeal
Unes4.0.dowssmmptaa (1)
(256, 16,
85)2.8.8)
(1.1)
(2,2)
(0.0)
(1.1)
No
Down
(f) Reseeex1011 ImageNet (3, 256, 256)
(count)
Iaput shape
Output shupe
Stride
Padding
Dilation
Geoupe
Type
convl (1)
(3 256, 256)
(54, 138, 178)
(7.
[2. 21
(3. 3)
(1. 11.1)
Ko
General
byyri 1.c.ne (2)
(64, 64.64)
(256. 64, 64)
[1. 1)
(1.
Ko
Danc:
Iystrl...um22..2 (3)
(256, 64)
(256. 64, 64)
[3,
[1. 11
(3.1)
(3.
32
Ku
Geneal
(256.
(256. 06,
11
Demo
(256.
(512) 04,
11
10.00
(1.
Dene
Javer2.0.omv2
Geseral
1aver20.omn3
(L
Dene
laye2 Doe2DDdlemempplee0
(256.
32,
Dows
layer2.1coms2
(512,
Gueaal
92,
laye: aywrst.dwwwmmppllcc..
(512,
(1024,
12.21
OLD
Dows
byerrs1
(1054,
13.35
11
Love4...cavl
(1024. 16, 16)
(2048.
11
11
(0.0)
(L
Dane
loye40..our2
(2043, 16. 16)
13
12.21
(.
Guseral
1ov44..0mm33.. (5)
(2048.
(1.
11. 11
1)
Dane
layes Deoommmpplee (1)
(1024, 16, 16)
(2045,
(2
06 03
1)
Dows
laves 1eeel..ooo3
(2048.
(2045, 5)
13.55
11. 11
(.1)
(1.
No
Geueral
(g) Non ImageNet (3, 256, 256)
Name (comnt)
Input
Outpat shape
Kernel
Stride
Padding
DNlattin
Groupe
Blas
Type
f.0
Cl 256, 2563
(128, 64, 64)
(9.4)
(4.4)
().0)
(1. 1)
Yes
Donse
featues. 1...cc.00.....00 (3)
(128, 64, 64)
(128, 64, 64)
(7.7)
(1.
(3.3)
1)
128
Yes
General
feanues (1)
(128, 64, 64)
(256, 12, 32)
(2.2)
(2.2)
(0.00
(1.1)
Yes
Dense
festure 3.0.Muck. (3)
(256. 12. 32)
(256, 12, 32)
(7.7)
(1.
(3.3)
(1. 1)
256
Yeu
Genzal
festures A.1 (1)
(256, 12. 32)
(512, 16. 16)
(2.2)
(2.2)
(0.
(1.1)
1
Yeu
Dane
fronms 5.0 S.0.book- (27)
(512. 16, 16)
(517, 10, 16)
(7.7)
(1.13
(.1)
(1.1)
512
You
Giner
fratures 6.1 (1)
(512. 16, 16)
(1094, R
(2.2)
(2,2)
(0.0)
(1.
1
Yeu
Dans:
kemmees....Ncca..
(1024, 8,3)
(1024, S.
(7.7)
0,11
(3,3)
(L
1024
Yes
Generl
28

---

(h) Inception V3, ImageNct (3, 299, 299)
Nome (count)
Input shape
Ourput shape
Kernel
Stride
Padding
Dilation
Groups
Blas
Type
Comrdd 1a,333 (1)
295)
(32. 1499
(3.3)
12.
10.00
(1.
General
Comrd 21,333
149,
(32. 147)
10,00
General
Cun24 (I)
147.
(64. 147)
(..1)
Con
(80,
Comcll 9s3oow
(0.0)
Misel, Shw.shhhtell..mm
Mised 3hoomcchis 1.000r
(68,
(1.1)
13
03
Dense
bhomachoss 7,0mm
(64,
(5.5)
(2.2)
Mixed boand336ddh.2 coux
(96, 15,
(3.3)
(1.
(L. 13
(..)
Geseral
Mixed branch3u3db__.
(96)
(L. 13
Generul
Mised Sobraach pool.come
25,
(31,
1
03
MixeL Se
3556,33,33)
(64,
(L.
(0.0)
(.
Deme
Mixed bruchh Looos
35,6,
(58,35,33)
D
(L.
(0.0)
(.
Dene
Mixol hrnch lel (4)
(64. 15,
(1. 1
(L 1
(8.
(.1
Dene
Mixal beacch5t Lcom
15,
(48,
(1. D
(L. I
(0.0)
(.
Dene
Mixed In braaachh....omm (1)
(588, 35)
634,, 17, 171
(3.3)
12. 21
10.
(.1)
No
General
Mixed to Sea cone (1)
(96, 39
(M% 17, 175
12. 21
(.1
No
General
Mixedd th hninchhisl oom (17)
(788, 17, 17)
(197, 17, 17)
1
(6.
0.D
Dene
Mixed 6xtrrach767_ L.op (2)
(7, 17)
(128, 17, 171
(1
IL II
(0.
(..)
No
Dense
Mised 6aorunch7k] 2.com (2)
(128, 17)
112% 17, 131
11. 71
IL II
(0,
(.1)
No
Desse mix
Mixel Rilmmddis7 Jooe (1)
(128, 17)
(182, 17, 171
(7. 1I
(1. II
(3.0)
(.1)
No
Dense mis
Mixed Sh (atnma777d8.....m coux (2)
(128. 17, 17)
(125, 17,
177
17. II
(L. II
(3.
(1
No
Desse mis
Mined da hranch7c7dhi Same (1)
(128, 17)
1990 172 123
(1.7)
(1. 1
(0. 9)
0.1
No
Denee mis
Mixol nz bchhnchfs).
Leam (4)
(768, 17)
1150. 172 121
(1. 1
(6. 0))
(.
No
Dens:
Misal fiz bc.hnad757722 2.0m (4)
(166) 17)
(150.
171
171
(1. 71
(1. II
(0, 3)
(..
No-
Dense mis
Mixed 6i (schmnnch777_ 3.cmm (2)
(160) 17, 179
1192 17. 135
CT. 1I
ULU
(3.0)
(LD1)
No-
Derse mix
Mised be hrhncchhcidhh 2.oo (4)
(160, 17. 179
1190 172 141
C7. 1I
(3.00
(1.1)
No
Derse mix
Mixed de branddi77ddh_5..o 55 (2)
(160) 17. 17)
(192) 17, 151
(1
(0.3)
(1.1)
No
Desse mix
Mixed, 6o brunhcc77k7 2.cor (4)
(192, 17. 173
(152 1T. 151
(L 1
(0.3)
(1.1)
No-
Desse mix
Misol, fir Scchnnch722_.. 2am_ah
(192, 17.
(00) 17. 17
(L I
(3.0)
(3.15
No
Dense mis
Anxlogit conc (1)
(768.
1133 533 555
(L 1
(6,0)
(3
No
Done
And ogis conel com
(178.5.5)
Ceneral
Mixol 7a bobncchhs 2.0m
(192, 17,
(S S
12. 21
(0.0)
(3.15
Ni
Genenal
Misedd TullTTueeaaaa hew.ennee77334444 4mms
(192, 17, 175
(102 K
(5.5)
21
(0.0)
(1.15
Ni
Gennn
Mixed, 7h beanchtcc con
Mixed Phoomnchic 1.c0m
(S84, K.5
Deno
Mixed_Th Thbs__hhhanddd3._.
(384, 5.51
Mixed, bhbbhhhcccccc...o 2actss_____hoooeee
(384,
Missd, The bannch3a
Mised_ Th Melna
(584,
Geama
Misei_ Sasiul__nmuuhhppee
Mixed pood.co
(1280) 5,
(192)
(0.0)
Mixal, hehnnchhs
(0.0)
Mixed, hraneh3sddi
(438,
0.1
03
brunch poot.com
(2088,
(192,
0.1
0.10
(0.00
(1.1)
(i) No1,,, ImageNet (3, 256, 256)
(count)
Inpot shape
Outgot shape
Padding
Dilation
Groups
Type
(1)
(3, 256 256)
[32, 12K, 125)
(3.1)
(3: 31
(1.1)
(1.13
No
Genernl
ficttnos (1)
(22 128, 175)
[32. 12%, 125)
(3.1)
(1. 1)
(1.1)
(1,
50.
No-
Nuuuresl.coom (3)
(32. 128. 125)
[16. 125. 125)
(1.3)
(1.3)
(0, (1.
(1
No
Dease
featine. 2 (1)
(16. 128, 125)
(96. 125. 125)
(1.3)
(1.1)
(0. (0.0)
(1.1)
No
Desse
featurs. 1.0
(96. 128. 125)
(D6, 64, 641
(3.3)
(2.2)
(L.))
(1
96
No
Genend
2.0m
(95 64, 64)
(24, 64, 641
(1.1)
(1.3)
(1
No
Dease
foetnnx Namm66002 (2)
(24.6 64, 64)
[144, 64,
(1.13
(1,
No.
Dese
3 (.)
[144, 64)
[144, 64,
(3.3)
(L.1)
(L.U)
(1,
144
No
Genend
featur:....omm 2(1)
[144, 64)
(34, 64, 6-4)
(1.33
(L.1)
13
No
Demc
fommrnc4com com:),
(144. 56)
[146,
13
144
No
Cenera
frammme...om 2
[144, 12)
12)
11,
No.
Desen
No
Desse
(192.
1192.
No
General
f
(192.
Dems
Com.o.T.om....0)
(192.
16,
Genend
faaterz.7.com
fomino (4
Genend
11.13
Dence
13
Seatures (3)
(95, 16)
1576, 16, 10T
(3. 1)
(LD
Desce
Seatures 12.com11:0122
1576, 10, 100
(3. 1)
(L
O.
General
(estuns: 12.cmm22 (2)
(576,
(36, 19T
(1. 1)
(0.0)
(L
Deas
Scatures Accmm1000
(576,
(376, (9
(3.2)
(2,2)
(L.
Genend
festres. (1)
(575, 5,
(160,3.9 90
(1.1)
(0.0)
O.D
Desse
fcatunos 15 (3)
(163. 8)
(.1
(0.0)
Desec
Scatancs 15ams (.0(3)
(063, 5)
(990, 80
(3.3)
(.1)
(L.
O.
960
Ganed
feotues 22)
(901 S, 5)
(100) 30
()
0. 11
01
11 1)
No
Dense
fromuss (1)
(WX1
(19), 8.
(.
(. 11
(0.0)
1)
No
Denen
Seatires (1)
(320, 5,5)
(1250 3,
(3.3)
(1.3)
(0,
()
No
Desst
29

---

General
Dense mix
Dense
Down
10'
TN PT (logurithnic)
100
10
Forward
Inpar VIP
Weight VIP
KPAC-red.
Forward
Input VIP
Weight VIP
KFC
KPAC-red.
Forward
Inpat VIP
Weight VIP
KFC
KPAC-red.
Forward
Input VIP
Weight VIP
KFC.
KFAC-red
Figure F17: Benchmark overview. We measure the performance ratios of our TN implementation
w.r.t. base line in PyTorch (PT). Blue boxes show the performance ratios of TN versus PT, second- color boxes show the performance ratios of TN+opt versus PT.
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 (16GB memory). We use batch size of 32 for the ImageNet architectures, and 1 128 for the others.
F1 Protocol & Overview
We compare different implementations of the same operations in PyTorch. The base line (refer- enced by 'PT') uses PyTorch's 's built-in functionalities for convolutions and related operations, such as torch. nn functional. conv2d (forward), torch. nn functional. unfold (KFC, KFAC- reduce), and PyTorch's '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 modified equation and operand list ('TN + opt'). Finally, we determine the contraction path using _et_einnumm.. contract_paaeeeee and perform the contraction with its Py PyTorch back-end (opt_ einsum.....mm.......... contract). We only measure the contraction time 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) 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 is time divided by the base line's. Ratios larger than mean that the TN implementation is slower, ratios smaller than 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 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.
30

---

(a) Mixed dense
(b) Dense
3.16
TN vensus (logarihmiic)
IN (lgarithmicc)
Inpat WJPWeight
KFAC-red.
Forward Input VJPWeight
KPC
KFAC-red.
(c) Down-sampling
(d) General
TN (logarihmic)
3.16
IN (logarihimi))
10
0.1
Forward Inpat VIPWeight VIP KFC KFAC-red.
Farward Input VJPWeight VJP
KFC KFAC
Figure F18: Impact of TN simplifications (non-simplified performance ratios shown in blue). TN simplifications improve performance on (a) mixed dense, (b) dense, and (c) down-sampling convolu- tions. (d) General convolutions not affected by TN simplifications.
31

---

F.2 Forward Pass
We compare TN and TN+opt with PyTorch's torch. mn functionnaa..... conv2d. Figure F19 visual- izes the performance ratios for different convolution categories. Table F5 contains the detailed run times and performance factors.
10a
TN PT (logurihmicc
10a
General
Dense mix
Dense
Down
Figure F19: Forward pass performance ratios of TN versus PT and TN+opt versus PT for different convolution types on GPU.
32

---

Table F5: Forward pass performance comparison on GPU.
(a) 3c3d, CIFAR- 10, input shape (128, 3. 32, 32)
Nome
[s
[9
Factor
TN + opt
FT
Factor
co0.10
1.26 1263-000 10-3
4.63 10-1
2.73%
1.23 10-3
4.64 10
2.64
c0022
1.81 10
4.52 10
4.22s
1.79 1073
4.58- 10
3.85
com3311
1.21- 10
4.11 10
2,94x
1,35 10
4.10 10"
2.83x
(b) F-MNIST 2c2d, input shape (128, 1, 28, 28)
Name
opt [s]
PT()]
Factur
conv1.1
8.21 10
2:25
10
3.65x
7.67
2.25 10
3.41x
c00.2.1
3.56
7.43
4.79x
3.24
7.83
4.14x
(c) CIFAR- 100 AII-CNN-C, input shape (128, 3, 32, 32)
Name
[N
Factor
TN + opt [s]
PT
Factor
c00-1.1
1.01 10
4.20 10
2.41x
9.45 10
4.19 10
2.25x
cmm2
1.94 10
5.09.100
6.26
1.88 10
3.10-
10
0.08
8.66 10-3
2.86
3.00
7.77
2.86
10
2.72
comv
8.58
1.75
4.91>
7.27-10
1.75
10
4.48
com5.1
1.67 10
2.91 10
5.74X
1.51 10
2.91
10-3
5.19x
6.13 10
2.24
5.08 10
2.27x
com70 10
2.58.
8:26 10
3.12s
2.51. 1611-107)
8.27
3.03
cm08.1
8.20
10
2.96 10
2.77%
3.42 10
2.97
10
1.15x
come)
7.52 10
2,35 10
3.19x
3.01 10
2.55
1.28
(d) Alexnet, input shape (32, 256, 256)
Name
TN|]
FTB]
TN opt [s]
PT
Factor
10
2:45 10-3
1-79 1198-002 10
2.44 10
1eature.3
7.43 10
2.87 10
2.79x
7.27 7.27-10 10
2.85
10
2.55x
4.68 10
1.04 10
4.52x
3.22-10-3
1.02 10
3.14x
6.15- 10
1.50
6.16-10-3
1.84 10
3.34:
Gaatuoos 10
4.41
1.31 10
9.36x
4.36
1.81- 10
3.35x
(e) ResNet18, input shape (32, 3, 256, 256)
Name
TN [+]
PT
Faciur
TN 4 [s]
PT]]
Fattor
conw
1.44 10
4.07 10
3.53x
1:44 10
4.08-100 1003
3.53x
layerl Baomw
1.05-10
1.78- 10
5.91
1.05 10
1.79-10- 10
5.87
layer2 0.cor)
10-3
1.89
6.88 10-3
1.51 10
4.54X
6.91 10-3
1.52 172-1103
4.64
1.60 10
-
4.23x
119-11--
3.80-10-1
1.37
U.oom)
2.01
laperk D.omn:
5.02
1.30
10
3.85
5.05
1.91-10
3.871
lgera....wasampph
1.10-10 10
3.79
10
2.91X
5.61 10
1.79-10
1.45%
layerá
20
2.36
2 10-6-003
2.36 1006-1000
layerd D.omr
4.47.
1.40
10
3.18x
4.51
1.40
3.21
jperi...dowasamppee
9.90 10
3.81
10
2,00X
5.16
10
3.83 10
1.35x
(i) ResNext101, input shape (32, 3, 256, 256)
Name
[N
TN apt [+]
PT
uxte]
1.45 10-3
4.07 10-3
3.57
1.44 10
4.07 -0-3
8.64
layerl
lapert.i.comr 10.com2
3.03 10
9.86 10-3
3.07x
3.03 10
9.86 10-3
3.08x
layerl 3
1.51 10
6.54 10
2.31x
10 -
6.54-10" 10-3
1.16x
10
1.29
1.29
3.33 10
4.93 10
6.75A
3.59 10
4.93
6.75
1.05 10
6.24 10
1.09%
6.84 10
6-24-10-3
1.10x
7.65
3.30
3.71 10
10-3
1ayer2.Leaom
4.50
3.27A
1.50
4.69
8.27
laser)))
1.67- 10
1:23 10
1.35
1.28 10
1.23-10"
1.04s
1.76-10 10
2.65 10
6.05
1.76-10"* 10
2.86.10-3
6.65
5.38 10
3.20 10-3
1.74X
3.42 10
3.20 10-3
1.07
layer3 Leam2
7.64
2.49. 10
3.07
10
10-8
8.07
layert
layer4 Deson?
10
2.02
10
3.99
8:06
2.02
4.00s
7.85 10
6.28
10
1.25%
6.33
10
6.28 10
1.01x
Juwrs....owssuppp..
3.34
1
3.44
10
layerk Leamd
4.76
1.36
10
3.51x
4.77
10
1.36
10
3.52
(g) ConvNeXt-base input shape (32, 3, 256, 256)
Name
TN
PT
Factor
IN sptiss
PT((]]
Factor
featuncs 0.0
4.26 10
9.88 16
4.31x
1.20 10-3
9.94 10
1.21s
Reateros (1
5.07 10
7.61
10
6.06x
5.07 10-
7.01- 10
6.66s
(cateox.2.1
7.60- 10
3.21-10 -00
2.37
3.89
3.20-10-2
Reateres 3 S0
2.36 10
3.81 10-
6.18x
2:35 2.35-002
3.81 10-8
6.17s
Reauees A.1
5.41 10
3.38 10-3
1.60x
3.52 11-3-
1.25-10
1.04x
Acatunos 0
1.11- 10
1.91
5.70
1-8--0--
5.60s
fraturos
4.54- 10
3.69 10
1.23X
3.04-10"
3.70-10"
0.93v
Reatenos 7 (
1.00 10
1.01 10
1.05x
1.02 1.02-1002
1.01-10
1.01s

---

(h) Inception 33, input shape (32, 3, 299, 299)
Name
1N [s]
PT [s]
Factor
TN-
opt [s]
PT(s]
Factor
Com 14, 2o
1.02 10
9.85 10
10.35x
1.01
10
9.79 10
10.30
Com 21, po
8.23 10"
5.14 10-3
6.30x
3.19
10
6.16 10
6.18K
Candd 20,3m
4.83-101
3.14-10-3
5.98
8.14-10
6.87
Comld come
4.86-10
1.17-10" 10
4:28s
1.72
1.17 10m3
1.48
Com kow
3.69 10
7.64 1664-100 10-
4.83x
3.65
10
7.54 10
4.77x
10
3.68%
10
1.62
Miuet_3hhraua5ss__ Looo
1.64
3.30X
1.63x
Misol, Sklearch5555 2ams
5.01 100
1:25 10-3
4.07
4.85
10
1:23 10
3.94
Mixed Sobrncchtsddl 2ooe
1.40-10"3
1:31 100--
3:38x
4:31
10
1.31
3.30
Mixed Manoo_3ttraccbaaaaa...
3.50X
3.40x
Misol ShhnaddS
1.33 10
3.26 10
4.00
7.04
10
3.27 10
2.15
Mierd Schmnchts
2.08 10
6.41-10
3.26x
1,05
10
0.40
1.61x
Misal, SchhanchS.....
188-10033
10
Mixed 5dbnnchhil com
2.18 10
6.89 10
3.12x
1.13
10"3
6.98 10
1.02X
Mixel ScbcanchSes__ Leozv
1.96-10"
6.91 10
2.84X
1.13
10"3
6.98 10-4
1.64x
1.15-10"
7.12- 10
1.07-00
10"2
7.13- 1
1.51
Mixed babnach)ssth 3chi Som
2.61
8.99
2.90x
2.36
8.00 10
2.02
Mixal Shbnnchhss
2.16-10"
1:22- 10
1.77s
1.41
10
1.22 10
1.15
Mixed 6h (htrarchicT_1
1.67 10
8.15 10
2,05
1.10
10-3
8.16 10
1.35x
Mixel @hbrma777__..om
2:14
2.66X
5.06
2.19x
Misol, Akleandh777 Jams
3.59-10"3
1.06 1
2.45
2:27
10
1.06 10
2.15
Mixed dnbronch7s7dh1 2,oe
2.17-10"
10
2,76x
10
7.88 10
2.26x
Misad #hbonehh5725l
10"3
Mixed & (c.tradhTTT_. 1.c0r
2.05 100
1.16 10-3
1.77X
1.41
10"3
1.16 10-3
1.21x
Misel_ Schaana777_2ao
8.19 10
1.12 10
2,84X
2.72
10
1.12 10-3
2.42x
Misal, Ac.randris7_
1:25-
10"3
Mixed 6 6cbronchhs2dhl 2.0m
3:25
1.10 10
2.96X
2.75
10
1.10 10
2.51x
Mixad fe h-hreehhs7chil Some
3.19-101 10-110"
1:28
2.49s
2.75
10
1.29 10
2.12x
Mixed, de.tranchxT__aa
3.78-107
1.48-107 10
2.54X
10-3
10-3
2.16X
Mixel, fenulul.ehaadd
8.57
2.66X
2:24x
And ommmemm
6.10-10-4
2.58 10
2.69s
3:20
10
2.38 10
1.34
AmsLogis come
8.06-10"1
10
0.53x
10
0.46X
4.37
Mixed 7ammme2777633444oo
1.54 10-3
8,89 10
1.73x
1.52
9.98 10
1.71x
Mixal Thlmmchhsl
1.29-10 10-3
7.43- 10
1.75x
8.76
10-1
7.43 10-1
1.18x
Misil. Phlearchlc3
Mixed, brunch333.222 2a.com
1.49 100
9.36 10
1.38%
1.26
10
9.38 10
1.34x
Mixed Th hrancct333 2hoom
1.46-100
9:37 10
1.38x
1.28
10
0.37 10
1.37x
1.67 100-1003
10-3
10-3
10-3
1.13x
Mixal_ 3hinulT.hhannhha 2aaw
3.18
9.82
3.28A
9.88
3.26
Mixal Thhrarc poul
9.54
1.76-100
1.41s
0.30
10
0.93
Miseed Tehnnchhst
1.68 100
1.08-10
1.36X
10
1.18X
Mixed, 7chearchhd
Mised_Te Towruchhas3hllaaa..
2.25 10-3
1.56 10
1.44%
1.73
10
1.56 10
1.11x
Mixol, Schoanchpod..ne
- -
1.04 10
1.20
9.35
10
1.04-00 10
0.90s
(i) MobileNetV2, input shape (32, 3, 256, 256)
TN(s]
opt [s]
Ratues00000
6.91 10-3
7:23 10
9.65x
6.52
10
7:24
9.56x
fcalunos 1eam0.0
2.28
2.06
11.11x
2.28
2.08
11.00x
fomine conel
5.84- 10
7:61- 10
7.42x
1.56
10
7.57
10
2. 2.06x
fcatunos 2.000.0
4.27 10
1.74-10
2.45s
2.02
10
1.74 10
1.16x
features com.1.
1.09
19.35x
3:31
1.00
50.
feuture.2 cons:
2.50 10
4.98 10
6.13x
1.08
10
4.99 10
2.16x
Scalunes 3amm0.0
1.78 10
7.86 10
2,26
9.63.10
7.88 10
1.22
2:30
2:30
Jonurs.3.covv1
2.30 10
6.33 10
4.63x
1.47 10
8.34 10
2.33x
Seatures 4comc11 10
1.34 10
6.62 100
15.76x
1.04 10
6.61
10
15.72x
2com:
10
2.61 16-1-10-4
4.23x
2.61
1.92%
fcaluee Smmm000
9.24 10
3.92
2.78
5.07
3.33
1.52
Seatures com. 10
5.44 10
7.87 10
6.91x
5.42 10
7.88 10
6. 6.88%
fooure.5.comv2
3.11 10
3.93X
1.9%%
Seatures ccmmm
2.49 10
9.68%
2.37
2.51
9.44X
(ommrs-T.cvs
7.40 10
2.09 10
3.68
3.20 10
2.10 10
1.58x
8.05 10
2.91-
2.775
-
2,92 10
1.51
com...1
4.14
5.53%
2/27
4.15
5.48%
7.98 10
3.07 10
2.60%
4.63
10
3.06
10
1.11x
11:0mm2
9.98 10
4.08 10
2.42X
5.07
10
4.07
10
1.39%
A.00m
10
10
2.16x
10
4.92
10
1.14A
fratuex 12.mmc1.0
4.18 3
6.04 10
6.91
4.16
10
6.05
10
0.87
featares 12.c0m22
1.16 10
5.53 10
2.10x
7:40
3.35
10
1.338
comos.1.0
2.29
7.67%
2.28 10
7.53X
features 24.cmm22
6.95
3.90
L.78X
4.10 10
3,90
L.Gx
fcatues 15amm..
1.24-10-1
3.53 10
2,62
4.56 4
10
3.53 10
1.23A
fratmrs 15 1.mnc1.n
1.48 1.98-10-3
2.72. 10
5.161
1447-1
2.73 10
5.39x
Reatires 15,c0m22
5.80
1.43x
5.44
3.80 10
0.94x
fetaros 1
1.12-10-1
9.74 10
1.15x
7.14
9.75 10
0.73x
15.0
1.25 10-3
7.31 10
1.71X
8.01
7,31 10
1.10x
34

---

F.3 Input VJP
We compare TN and TN+opt with PyTorch implementation of the input VIP via torch. autograd grad. Figure F20 visualizes the performance ratios for different convolution categories. Table F6 contains the detailed rum times and performance factors.
10a
TN PT (logarihmic)
10a
General
Dense mix
Dease
Down
Figure F20: Input VJP performance ratios of TN versus PT and TN+opt versus for different convolution types GPU.
35

---

Table F6: Input VIP performance comparison on GPU.
(a) 3c3d, CIFAR-I0, input shape (128, 3, 32, 32)
Name
[s [g]
Factor
IN + opt [s]
FTB]
Factor
com-1.
2.24 104--10-2 10-3
1.39- 10-3
1.61x
2.19 10
1.34 10"
1.63x
2o022.0
2.61 10-0
8.29-1077
3.15x
2.55 10
7.86 10-
3.26x
1.46- 10
5.04 10
2,90s
10
3.02x
(b) F-MNIST 2c2d, input shape (128. 1, 28, 28)
Name
[N
PT
Factor
IN e opt [s]
PT
Factor
co
9.47 10
4.36
10
2.17s
8.80 10
4.40-
2.02
com2.1 1
3.07 10
9.83
10
3.74x
3.62
10
9.83
10
3.09X
(c) CIFAR-100 AlI-CNN-C, input shape (128, 3, 32, 32)
Name
TN
PT
Factor
PT[]]
Factur
comv11
1.91 10
9.84 10
1.94x
1.87
10
9:37
2.00
600221
2.00 10
5.95
3.35X
2.00 110
5.92
3.37x
cmm6.11
7.82 10
5.05 10
1.55
7.77
7.77-10
5.01
10
1.55
cone
8.23 10
3.11 100
2,65x
8.17-10
3.10-
10
2.03
10
4.36
3.57
1.55-10
4.30-
3.56
co0661
4.58 10
3.95 10-3
1.38%
4.53 10-3
3.96
10-3
1.14x
c00770
2.86 10
8.32 10
3.44X
2.81
10
8.68
10
3.24X
10
2.91
2.85%
3.47
3.32
1.04x
comv
7.76 10
2.21 10
3.51x
2.90
2.61
1.11x
(d) Alexnet, input shape (32, 3, 256, 256)
Name
[N|]
FT]]
Factur
TN opt [s]
[T[
Factor
lostonac)
1.92 10
5.48- 10-3
3.50
1.92 10-2
5.54-10"3
3.46
3
1.15 10
4.16
10
2.77
1155-1022
1.20-10
2.75:
featresu
1.49
10
5:30 10-3
1.49
3.00x
featues.5
6.26 10
1.86 10
3.36x
6.26 10-3
1.86
10-3
3.37x
Scaturoe 10
4.41-10 10
1.52-10
3.35x
4.40-10-3
1.35-10
3.26
(e) ResNet18, input shape (32, 256, 256)
[N
PTB]
opt
PT
Factor
3.38
8.50
3.
3.58
10
5.49 10-3
3.96
byart B
1.08 10
2.17
4.87
1.05
10
2.11
1.89
7:03 10
3.72
10
1.89X
6.95
10
1.66-10''
1.90x
6.91 10
1.35
10
4.47x
6.90
10-3
1.51-10-5
4.56x
2.02
9:02
2.51x
1.71
10-3
7.64 7.64-10-4
2.24x
bayer 1
3.94 10
9.05
10
1.29x
3.88- 10
3.01-10-3
1.29
5.07 10
1.31
10-3
3.87x
5.07
10-
1.36 10-
3.74x
1.15
5.96
1.94X
9.54
6.40
1.49x
layen' come
2.89 10
3.08
10-3
0.94X
2.84
10
3.12 10-5
0.81x
4.50 10
1.40
10
3.21x
4.49
10
1.44-10
3.12x
5.51
5.90-10
(1) ResNext101, input shape (32, 3, 256, 256)
Name
TN
PT
Factor
TN opt [s]
PT
Factor
oand
3.38
5.52 -3
3.97
3.38
8.48 10-3
3.98
byert Baoor
6.18 10
1.90
10
1.
2. 2.80
10
1.06 10
1.46x
byerl flcom?
3.06 10
1.17
10
2.00
3.05
10
1.17- 10
2.61
layert 0.com3
1.46
6:37
22
7.39
0.58
2.40 10
1.14
10-2
2.10x
1.44 10
1.17 10-2
1.23x
2
2.75 10
1.90
10
1.40x
2.75
10
1.95- 115-1002 10
1.41x
1:04
6:43
10
8.90 10
4.78
1.88
8. 0.00
10
1.74-
1.70s
layer21 Loon2
1.51 10
4.46 10-3
3.388
1.51 10
4.45 10
3.39x
1.94
1.25
10-2
L.66X
-
1.25 10-2
1.70 10
8.33
2.11x
1.76
8.54-
2.13
Iaper3...oom3
8:21 10
6.32
10
1.30X
6.30
10
1.32-100"
1.01x
5.51-10 10
4.51 10
1,21x
5.00 1
4.52-10 10
1.11x
7:00
1.97
3.85X
7.60
1.98
3.84x
1.51 10
1.24 10-2
1.22X
1.26
1
1.24-10-2
1.24
1.02
2oom
8.26 10
5.43
10
1.52x
1221-1002
5.4-100
1.51
7:65
6.72
1.148
6.73 10
0.83x
4.61 10
5.45
10
0.84x
4.31
10
5.45-10
10
0.79x
Layen21Loorr2
4.79 10
1.34
10
3.57X
4.79
10
1.39
10
3.44x
(g) Non-NeXt-base, input shape (32, 3, 256, 256)
TN
featuest.0
5.36 10
1.70 10-3
2.99x
1.57 10-3
1.79 10
Statuos 105lo00
4.63 10
8.60
6.36
4.63
8.68
10
5.40x
fratuecs 2.1
8.85 10
5.57 10
1.05x
3.55
3.55-10-3 10
5.36- 100
0.68
Seatures 3//lock.00
2.14
10
5.09x
1.14-10-2
4/21 110
5.09x
1.27
3.34
1.05 10-2
2.16 10-3
1.05- 1
2.16 10
4.80x
4.31 10
5.50 10
0.78x
3.25-
1.50-110m
0.59

---

(h) Inception 33, input shape (32, 3, 299, 299)
Nome
TN [s]
Facter
TN
opt [s]
PT(s]
Foctor
Cuav2 Aos
1.27 122--10"
3.19 10
3.97
1.26
10
8.21 10
3.92x
Cauy2 Ze, 3x3,co0
3.15 1055-10
5.13
10
6.16x
3.16
10
6.17 10
6.10x
8.11
8.17- 1117-10m3
6.19
Comvdd 1h bel.
5.76
1.00
10
5.31x
1:37
1:08
10m3
1.25
Chav28,4a 3dl.oo
3.71 10
1.12 10
3.30X
3.71
10
1.12 10
3,30x
MixeL Sttnnsiiiii
6.89
10
10
0.97s
Mixee.9h btbosshbss Louay
1.13
5.87
1.92X
0.98
0.9Tx
Misol SoosL_Shmmnse St Jame
4.97 1.97-10*
1.59 10
3.58
4.97
10
1.59
10
3.57
Mood honnchhhcddd Z.cerr
4.23 10
1,07
3.988
10
1.00 10"
3. 3.
5.68
1.65 10
3,41X
3.413
Mixel_Sklemmi polllastppolllco
9.70-10
5.10 10
1.90
4.77
10
5.13 10
0.93x
Minol Schrachisi com
1.48-10
8.10- 10
1.52x
8.07
10
8.10 10
1.00x
1,23-10-3
8.87-
Mited M.branhixl.com
1.68 1.68-0012
1.04 10
1.61X
9.40
10
1:05 10
0.90x
Sooo_sddructodd.. Looty
1.38 1.88-10"
8.69 10
1.59x
7.84
10
9.12 10
0.96x
1.14
1.32
10
0.88s
Mixel bonchhdddi 3oom
2.52-
1.70
1.48%
2:46
1.70-
1.45
Mixal mhlnachaal.ccm
1.78 1.78-10
1.18- 10
1.50%
1,26
10
1.18- 10
1.05%
Mixed db hhhnmn777_1
1,37
8:69
1.58X
10
N.70 10
1,07X
2.13-10
8:27
2.58K
8.27
2.16%
Misol, 3hhhmmes777_33mmm Same:
2.54
1.08 10
2.58
2:22
10
1.08 10
2 0.05.
Moxd thtinddoeddd 2.cer
2:08
8:30 10
10
8.09 10
2.222
Mixed &c (chnna777_1.
Lone
1.56 10
1.03 10-3
1.52X
1.09
10"3
1.33 10
1. 1,06X
Misol fu fc..mna777_22222
8.19-10
1.14
2.79x
2.76
10
1.13 10
2 2.43x
Mised, fch.mncc77777.
1.28
2.12
Mixed 6c hondmood 2.om
3.10
1:12- 10
2.78X
2.77
10
1.12 10
2. 1.48%
Mond 6o tonccchhhhhlll S.cmr
2.96
1:33 10
2.22%
2.61
10
1.33
10
1. 1.96
Mixed, brrnnc2777_22omm
3.77-10
1.54
10
2.17x
Misol (onoo_ffllnns777333ooo
3.65
1.51 10
2,42
2.22
5.53-10
2:74 10
2.02%
3.05
10
2 10
1.50
Austoogssgoorlcoo
6227-10
2.04
10
10
0.25x
7:05
Mixel (oo.."/atroccbho 4com
1.46 1.46-10-3
1.30 10-3
1.38X
1.42
10
1.14 10-3
1.26x
1.31-1-3
7.40 10
1.77X
8.47
10
7.50- 10
L.07
Thhnnch3e3..
Mixed Thhranchx3. 22.com
1.51
9.77 10
1.50X
1.26
10
1,02 10
1 1424X
Mited. 7hhrmnchis3 2hcom
1.50-10"1
9.77 10
1.58X
1:27
10
0. 10
1 1
1.86-10-3
1.02
10-3
10
1.10x
3.32
1.02
3.24x
3.31x
Mised 1hhranch pool
1.01-10-2
5.57 10
1.813
0.18
10
5.10-10-1
1.21
Mixed Wctnncciii..coo
1.25
1.35x
0,09x
1.45
0.95x
Mixel_ (onoo_/.iocctt33al....
2.05 10
1.68
10
1.288
1.57
10
1.56
10
0.95x
Selol_7cleanchpppe ponlame
1.27.10
8.35
10
1.53x
8.32
10
8.35
10
1.00x
ii) MobileNetV2, input shape (32, 3, 256, 256)
[T[s]
TN- opt [s]
D0
8.32 10
2.08 10
4.01x
5:25
2.02
4.08
2.27
2.63
8.64
2:27
2.60
8.71x
Reatores cons1
3.22
1.17
2.75
1.02
10
1.16- 10
0.87
Felomce22.om..00 00
7.67 1067-10
2.61-10
2,9
3.68
2.61
10
1.40x
2.0m.12
2.79
8:11
2.79
8:11
2
1.48 10-3
6.38 10
2.33%
7.53
10
8.40
10
1.18x
2.86-10-3
L.05.00
2.75
1.45 10
10 1.05
1.58
3.0mm.
2.08
2.95
2.08
2.ome
1.77-10
1.04 10
1.70%
9.75
10
1.04
10
0.94x
4.00m. 10
T:63 1163-10-3
3.15
2.42%
7.02 10
3.15
10
2.42X
2
10
festure 5.0m0.0
1.20-
5,26
2.29
6.09 10
4.83
1.28
Rrommss..cmm 1
5:41-10
1.02-
5.29%
5.39 10
1.02 10
5.27K
comu
9.63
4.00
4.35
2.11
1.07
2.10
1.07
1.97x
2
7.77 7.77-10-1
2.33 10
3.34%
3.04 10
2.33
10
1.30x
8.13-
3.11
2.38
4.63
10
3.40
1.38x
2,09
5.48
2.07
5.47
2
8.66 10
3.16 10
2.74%
4.04 10
3.16
10
1.29%
feanee....00222
9.34
4.22 10
2.21x
4.74 10
4.24
10
1.12x
12amm.n
1.16 10-3
7.11
10
7.10
fromms 12.mm: in
3.84
7.81 10
4.85
3.82
7.91
4.83x
Hammess:2oome
1.08 1.08-10"
5.71 5,71-10"
1.90x
6.13 10
5.73 10
1.07X
14
1.61
5.26-
1.60
8.26
Italues.4..0022
8.14 10-4-10-
2.87 10
2.83x
3.54
2.87
1.34X
faruns. 15
8.46.10-1
6,29-10" 10-4
1.34s
5.55 10
6.08 10
0.91x
femmn. 15 in
105
3.62.10"5
4.21x
1.50
10
3661
4.17x
feanues. cous2
9.64
2.18%
4.82
4.44
fralues 17.auv2
1.28 1.28-10
7.30- 10
1.60x
6.98
10
7.32
10
0.95%
1e0
1.29-10-
1.28
1.00x
8.76
10
1.28
10
0.68x
37

---

F.4
Weight VJP
We compare TN and TN+opt with 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 rum times and performance factors.
10a
TN versus PT (logarihmmic
10
General
Figure F21: Weight VJP performance ratios of TN versus PT and TN+opt versus PT for different convolution types on GPU.
38

---

Table F7: Weight VJP performance comparison on GPU.
(a) 3c3d, CIFAR-10, 10, input shape (128, 3, 32, 32)
Name
[N
Factor
IN + opt
FTB]
Factor
comv1.
2.27 10-3
1.50 1.50-10 10-3
1.52x
2.27 10
1.50 10"
1.51x
com-2.0
3.00 10
1.12 10
2.68%
2.09 10
1.07
10
2.78x
10
5.46.10 10
2.37s
1,25 10
2.47x
(b) F-MNIST 2c2d, input shape (128. 1, 28, 28)
Name
[N
PT
Factor
IN e opt [s]
PT
Factor
co
1.08 10-8
3.81
10
2.83*
1.03
10
4.05-
2.54x
com2.1 1
4.12 10
1.02 10
4.02x
4.09
10
1.03
10"
3.98x
(c) CIFAR-100 AlI-CNN-C, input shape (128, 3, 32, 32)
Name
TN
Factor
Tp-
PT
Factur
20m111
2.43 10"
1.02 10-3
2:39X
2.42 10
1.02
10
2.37x
20092.1
3.83
6.81x
1.97 100
5.62
3.51x
2um3.11
6.30 10
4.34 10-3
2.00
8.33 5.33-10
4:21
10
1.38
com411
8.66-100 100-110
2.64 10
3.28x
8.68-10
2.68-
10
3.21x
1.60-10 10--10
4.75%
1.61-10-2
3.42-
com6.1
5.23- 10
2.80 10
1.87x
5.17-10-3
2.81
10-3
1.80x
2.0770
3.68-10
9.97
10
2.68X
2.50 10
1.04
10
2.49X
am81
9.13-10
4.33 10
2.62
20m9
8.78-107
3.34 10
2.49x
3.93
3.50
1.12x
(d) Alexnet, input shape (32, 3, 256, 256)
Name
[N|]
FT]]
Factur
TN opt [s]
[T
Factor
listonac(
1.82 10
3.31 10-3
5.50
1.92 1.22-10-2 10-2
3-33-1002 10-3
5.46
3
2.02 10
2.57-10
7.85x
114--1022
2.58
10
4.44:
featres.
1.67
5:17
1.07
featues.5
8.16 10
1.97 10-3
4.15x
6.13
6.13-10-3
1.97 10-3
3.11x
Scaturoe 10
5.80 10
1.47-10
3.94x
4.34-10-3
1.47-10
2.95
(e) ResNet18, input shape (32, 3, 256, 256)
[N
PTB]
[st
PT
Factor
3:00 10
8.23
3.05
2:99 10
7.88 10
3.82
byart B
2.34 10
9.18
7.37
1.10
10
3.22 10
3.43
3.88 10
2.97
10
1.98X
5.87
10
2.96-10
1.98x
1.17 10
1.06
10
7.021x
6.98
10-3
1.86-10-2
4.30x
1.85
7.39
10
2.51x
7.60
10-4
7.33
1.04x
3.79 10
2.71
10
1.40x
3 10-3
2.71-10-1
10-3
1.59
6.62
1.48
10-3
4.06x
4.87-11 10-
1188-10-5
3.381
1.61
5.39
2.99X
5.45
1.12X
layen'.l comm
2.85 10
2.46
10-
1.16x
2.82
10
2.46- 10-5
1.15x
layerk...oww2
4.83 10
1.72
10
2.80%
4.31
10
1.72 10
2.50x
1.00 10
1.02
5.07-10
(1) ResNext101, input shape (32, 3, 256, 256)
Name
TN
PT
Factor
TN opt [s]
PT
Factor
oand
3.00
8.22 -3
2.99
7.88 10-3
3.82
byeer Baoor
7.08 10
2.92
10
2.02x
3:75
10
2.89- 10
1.30s
byerl flcom?
6.70 10
2.53
10
2.05
6.72
10
1.91-10"1
3.51
layert 0.c0m3
3.12
8:78
1.02
1.0/2x
layor28ooor1
2.29 10
1.80
10-2
1.28x
1.77
10
1.76-10- 10--
1.01x
2
6.64 10
1.23 10
5.40x
6.63
10
1.23 103-10-2
5.59
1.82
1.30x
8.57 10
5.21
1.64K
4.55
5.24
0.87
layer21 layer2. Leom22
4.04 10
1.21 10-2
3.331
4.04 10
1.22 10
3.323
1.84
2.03
1
2.03 1.03-10-2
1.63 10
5.77
2.83
1.63
5.82 10-2-10-8
2.81
Iaper3...oom3
1.17 10
1.07 10-2
1.10X
T.19 10
1.07 10
0.67x
6.19 10
5.95 10"3
1.01x
3.85
10 1
6.01 10
0.64x
1.47
3.17
4.65X
3.14
4.67x
1.55 10
2.10 10-2
0.74x
1.33
1
2.11- 1.11-10-2
0.63x
byyert
8.07
3.13 10
0.28x
8.06 10
3.13- 10
0.28x
8:23
1.06
0.788
1.06
10
0.63x
4.96 10
5.90
10
0.84x
3.50
10
5.89
10-8
0.60x
Layen21Loorr2
6.53 10
1.40
10
4.72X
6.62
10
1.45- 10
4.50x
(g) ConyNeXt-base input shape (32, 3, 256, 256)
[N
featuest.0
5.93 10
1.90 10-3
2.98x
1.87 10-3
1.97 10
0.95x
Statuos 105lo00
2.58 10
2.33x
2.53 1.53--102
1.00
10
2.83x
fratuecs 2.1
8.20 10
4.53 10
1.83
1.32- 10
4.52 10
0.98
Seatures 30/lock.00
1.23
5.85 10
2.10x
1.22 10
5.82 10
5.30-10"
1.08 1.08
3.74
5:39
6.05 10-3
3.63 10-3
1.6GX
1 6.03
3.64 10
1.60x
4.74 10
5.26. 10
0.90
3.53-10-3
6.17.
0.66.

---

(h) Inception 33, input shape (32, 3, 299, 299)
Nome
TN [s]
Factor
TN-
opt [s]
PT(s]
Foctor
Cuav2 Aos
1.07 -000-10
1.70 10
6.31%
1.07
10
1.70 10
6. 6.29
Cuay2 Ze, 3x3,co0
6.00-10
1.15
10
5.18x
3.11
10
1.16 10-2
2 3.688
6.10-
1.34
4.55%
1 1000-10-2
2.78
Comvd 1h bel.
5.48 10
1.82-
10
3.01x
2:26
2.12 10m3
1 1.07%
Chav28,4a Bdl
5.28 10
1:29
10
4.08x
3.39
10
1.29 10-2
2.62x
MixeL Sttnssiiiii
1.16
1.41
10
10
0.05%
Mixel_9h bobosshbss Louay
4.92
1.45
3.37X
1.47
0.95x
Misol, SthnmesS655.2 Jame
9.28 9.28-10
1.28 10
7.23%
4.85
10
1.28
10
3.77
Mood honnchhhcddd Z.cmr
7.78 10
1.75
4.45x
10
1.75-10 10
2.01x
1,05
1.85
3.51X
3.21
Mixel_Sklemmi pullaad
4.52-10
9.10 10
4.97A
1.16
10
8.36 10
1.90
Minol Schrachisi comr
6.35-10
2.00 10
3.27x
1.67
10
1.08 10
0.86
6.33-10-3
1,93
1.64-1
0.88x
Mited M.brnnhixl.com
7.45- 7.46-10-1
2.33 10
3.19x
2:03
10-3
2.31 10
0.88X
Shruusddruuc Looay
7:24-10
2.15 10
3.36X
2.00
10
2.16 10-3
0.93x
Minal bamr.leaaaccca..oomm
1.10-10"2
8.34
8.34-10 10-3
1.31x
Mixel bonchhlddi 3oom
2.46-
1.12-
2.20x
2.42
1.12- 10
2.17x
Mixal fkbenchaal.ccm
5.17-10"3
2.05 10
2.52%
1.80-
10
2.07 10
0.90x
Mixed db hhhnm0777_11
4.63-10
1.58
1.46
10-3
1,61 10-3
0.91x
Mixe_ btmm..n7777777...
3.09
1.50
1,94X
2.01
1.22%
Misol, 3elhhmmm7777_3ammm Same
3.44
2:20 10
1.50
2:57
10
2.33 10
1.01
Moxd thtinddheedd 2.cor
4.06-10
1.66 10
2.48X
10
1,08 10"3
1.08x
10
1.21x
Mixed &c (chnna777_]. cone
4.99
1.98 10-3
2.50X
1.77 10-3
2.33 10
0.87X
Misol fe dchmmna77772222
4.87-10"
2.71
1.80X
3.30
10
2.76 10"3
1.20x
Mised, fechmani777_.300
4.85
2.64
- 10
1.04x
Mixed 6c hondmoodd 2.om
3.43 10
2.80 10
1.96X
2.95
10
2.90
10
1. 1.06x
Mond 6i henccchh//hlll S.cmr
3:20-10
2:78 10
1.15%
3.41 - 10
2.82 10
1.21
Mixed, brannc2777222om
5,96
3.19
10
10-3
1.188
Misod (onos_ffllnns777333ooo
6.50
3.25 10
1.09X
1.08X
6.48-10
3.45 10
1.87%
3:45
10
3.98 10
0.89
Austoggssgoorlcoo
5.34-10
2.09
2.56X
10
2.76 10
1.
Mixel (oo.."/aorocchho 4com
1.53 10--
8.48 10
1.80X
1.52
10
9.50 10
1.79X
2.08.10-3
1.62 10
1.28X
1,07
10"3
1.63 10-3
0.88x
Thhnnci3e3..
10"3
Mixed Thhranchh33 22.com
1.56 1.56-10
1.47 10
1.06X
1.20
10
1.47 10
0.82x
Mited. 75.tmnchis3 2hcom
1.65-10
1.50 10
1.11X
1:22
10
1.50- 10-3
0.82X
2.34-10
1.65
10-3
1.66 10-3
0.80x
3.55
1.23
2.00
1.26
2.46x
Mised "hhranch pool
1,84
1.48 10
1.26%
8.67
10
1.40 1.40-10 10-3
0.59x
Mixed wctnncciii..coo
3.07
3.08 10
10
10"3
0.50x
3.30
3:11
1.06X
Mixel_ (ono___..oncchh3BB....
3.56 10
3.11
10
1.15x
2.03
10
3.10
10-3
0.65x
Mixol, poolee
2.70-10-2
1.61. 10
1.68x
1.22
10-3
1.61-
10
0.78x
ii MobileNetV2, input shape (32, 3, 256, 256)
[4]
opt [s]
feawess0.
7.70 10
1.45 10
5.30
7.07
10
1.46
5.26x
Loom Locm
1.50-
2.46-10
6.48
1.50
10
2.47
6.46
festores Lomc
1.47 1.87-10"2
2.40-10
6.11s
2.61
10
2:30
10
1.10x
8.97
4.95-10-9
1.81x
4.22
10
4,95
10
0.85x
1.0m
2.14 10
2.14
2:41
7.38 10-3
1.96 10
3.76%
1.51
10
1.89
10
1.01x
festurex 0.0
3.56 10
2.24-10
1.58
1.83
10
2.19
10
0.84x
1.34
2:46
1.04 10
2.89-10
3.61x
2,59
10
2.85
10
1.01x
festures
7.46-10-2
1.01-10"1
7.41x
7:44
10
9.84 10
7.56X
3.00 10
9.40-10-4
10
fatures fntures.5.cmm000 00
1.49
7.46-10
2.00s
7.48
7.48 10
1.00
1.0mre5...0m111 1.0
4.77 10
9.10-10"
5:24x
4.75 10
8.88 10
5.34X
3.62 10
1.06-10
10
1.01
f.00re77.00.110
2.61
7.13x
2.00
3.60
7.13X
1.48 10
5.65 105
2.38
4.64
10
5.56 10
0.84%
3 0.
1.14-10
5.36-10
2.12%
5.22
10
5.34 10
0.98x
feouues
2.44-
4.31x
2.43
5.69
2.28 10
6.32 10
2.68%
6,50
10
8.82 10
0.77%
feanues. 11.c0v22
2.36 10
8.68
2.72x
7.59
10
8.00 10
0,91x
(ssuns.12ssm.h0
1.66 10-5-1003
10-0
10
1.03
fecuno.11 12mmm.12.www 12.omm.10
3.52
8.90-10
4.29'
3.50
8.19
10
4.27
feamues femmes.12.cov2
3.27 10
1.26-10
2.59%
1.10 10
1.26
10
0.87x
14smm.1
2.07
3.90
10
3.00
10
14,0ess11.coov2
1.00
1.39
0.76x
5.10
10
1.40
0.36X
lesuno. 15
1.12 10
7.19 10-4
1.56
6,21
6.21- 10
7.10 10
0.87
forons. 15 10
2.31 100
5.96-10-
3.87s
222-10
5,96
3.83x
feanres 15com
1.34
1.41-10
6.53
1.40 10
0.47%
fratues. 2
1.50 1.50-10
1.67-10
0.95
8
1.62. 10
0.54x
180
1.53 10
1.68-10
0.91x
1.04 10
1.63
10
0.64x
40

---

F.5 KFC Factor (KFAC-expand)
We compare TN and Tot with PyTorch implementation 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.
10
(logarihmic)
10"
10-1
General
Dense mix
Dense
Down
Figure F22: KFC/KFAC-expand factor performance ratios of TN versus PT and TN+opt versus PT for different convolution types on GPU.
41

---

Table F8: KFC (KFAC-expand) factor performance comparison on GPU.
(a) 3c3d, CIFAR- 10, input shape (128, 3, 32, 32)
Name
IN Bl
PT (s
Fictse
IN + opt
FTB]
Factor
20011 1.0
1.03 10"1
2.42 10-3
0.43x
1.03 10
2.52 10"
0.41x
2.092.0
6.69 10
3.83 10-3
1.75x
6.97 10
4.52 10
1.54x
3.27-107
10
1.37s
3.53
(b) F-MNIST 2c2d, input shape (128. 1, 28, 28)
Name
IN bi
PT [4
Fictor
IN e opt [s]
PT
Factor
1.22-10
0.61x
9.50 10
1.72-
0.54x
20m22.1
1.03 10"
9.34 10
1,08X
1.02
10
9.47
10"
1.08x
(c) CIFAR-100 AlI-CNN-C, input shape (128, 3, 32, 32)
Name
TN
PT
Factor
Tp-
PT[]]
Factur
20m111
1,37 10"2
1.48 10-3
0.933
2.72
10
2:11
10
x
20092.1
1.48
2.13x
1.53 10
7:14
2.15x
cum3.11
4.77 10
1.15 10
4.17x
4.58 4.56-10
1.38
10
9.30
com4.11
2.32 10
1.34 10
2.05x
2.25-10-2
1.14-
10
1.98x
7.08- 10
5.82
1,215
1.01 101-110-]
com6.1
2.84 10
1:53 10
2.18x
2.83 2.83-10"2 10-2
9.77
10-3
2.90x
200770
8.68
5.95 10
1.48X
9.30 10
6.01
10
1.56x
am81
1.08-10
9 10
1.08%
10
20m9
1.06-10"
1.49 10
0.71x
4.61
1.50
0.30x
(d) Alexnet, input shape (32, 3, 256, 256)
Name
TN
PT
Factur
TN opt [s]
PT]]]
Factor
5.45 10-2
1.35- 10-2
4.03
0.00 10-2
4.55
3
4.57 10
4.14 10-2
1.10x
5:31- 10
1178
1.42s
featres.6
8.63
1.96
10-2
8:74
Seaturess
3.76 10-2
4.10 10-2
0.02x
4.57 1.57-10-2
4.38 1.:8-10-2
1.06x
Scatencs 10
1.52 10
1.38-10"2
1.10x
1-91-10-2
1.47-10"
1.30s
(e) ResNet18, input shape (32, 3, 256, 256)
TN
PTB]
o [s]
5.25 10
2.07 10-2
2.54x
5:22 10
2.06 1
2.64s
layerl Deomn
4.36 10
2.98
10
1.000
5.57 10
3.02 10-2
1.84:
Riper
2.81 10
6.38
10
4.41x
2.78 10
1.23-10"
100
2.25x
1ppr22...om2
2:36 10
1:90 10
1.34x
3:00-10
2.01 -00
1.03
Jaer2...owwsampk..
3.68 10
9.14
4.49X
6.45
-0
0.81x
1.34 10
9.19
10
1-06x
1.40 10-2
9.17-10-3
1.53
layer3
1.90 10
1.84- 10
1.033
1:25-10-2
1,85 10
1.16x
7.00
2.80x
4.50 10
5.72
0.80x
I.om))
8.63 10
4.79 10
1.81x
9.12
10-3
4.80
10-3
1.08
Layer
2.48 10
1.53 10
1.528
2.49
10-
1.88 10
1.32x
1.19-10
5.45
(1) ResNext101, input shape (32, 3, 256, 256)
Name
TN [s]
FT
Factor
TN o [s]
PT
Factor
ban]
10
2.05 10
A.50
5.06 10
2.05 10
2.46
layeri Dicem
3:33- 10
1.85
10
1.80
1.70 1.70-100 10
2.08 10
0.82x
layerl Deamd
1.09 10
6.60
10
1.06
1.11- 10
8.12 10
1.37
layert
7.49
1.04 10
7.52
layer2
1,00 10
1.58 10
1.05x
1.04 10
1.58 10
0.68x
1.40
4.44 10
3.15
1.41 10
1.54-10"
3.18
5:10
5.30
1.40-10
4.22
10
9.30
2.02 10
4.24 10
0.69
10per2.L.om22
5.07 10
4.23 10
1.20X
5.07 10
4.24 10
1.19x
1.14-110
5.21
1
6.27 10-3
6.11- 10
3:21
1.90
6.29
2.02
2,14s
8.77 10
4.30
10
2.04X
7:17- 10
4.33 10-
1.06x
7.50 10
3:08. 10
2.47x
2.28- 10
3.08
0.74x
Loom:
2.12
1.95
1.0%X
1:05 10-2
1.99
1.03x
Iayer
8.75 1.75-1001
4.15 10
2.11x
7.18 7.16-10 10-3
4.26 10-3
1.68
Doomd
4.70
2.47 10 2
1,91x
1.71-10-2
2.17 10
1.91
layen' D.com3
7.88
7.66
1,03
6.74
7.67
0.88X
4.54-10
2.52
10
1.80x
2:03
10
10-5
0.80x
Lcom2
1.38 10
1.16
10
L.66X
1.36
10
1.17 1.17-10-2
1.16x
(g) NonyNeXt-base, input shape (32, 3, 256, 256)
TN
[s)
featuest.0
9.93 10-2
10-3
4.71x
1.18 10-3
2:11 1003
0.56x
Statuos 105lo00
4.00 10
0.30x
6.25 1.26-10-2
1.42
0.37
fratuecs 2.1
2.37 10
4.90 10
1.85
7.81-10-3
1.93 10-3
1.59:
Seatures 30/lock.00
1.61
6.90 -
0.23x
1.57-10-2
7.12 100
0.22x
0.1242
0.88
3.98 10-3
3.35 1
3.96- 1066-10-3 1
3.43 10
0.12x
Ratuuou
0.52 10
3.30-10' 10
2.06x
4.77-10-3
8.31 8.31-10-3
1.44s

---

(h) Inception 33, input shape (32, 3, 299, 299)
Nome
TN [s]
Facter
TN-
opt [s]
PT(s)
Factor
Cuav2 Ax
3.42 10
4.00
10
8.36%
3.43
10
4.11- -10
8.33x
Cuay2 Ze, 3x3,oo
1,58 10
9.29
10
1.70x
1.91
10
0.18 10-2
2.08x
1.86-10
9.66
10
1.11x
10-1
10-4-10-2
4.09
Comvd 1h bxl.
5.26 10
2.33
10
2.28x
1.74
2:41 10m3
0.72*
Chav28,4a 3d3.oo
1.01 10
6.34
10
1.58X
1.08
10
6.16 10-2
1.76x
MixeL S.tnndiiiiisooo
4.13
2.02
10
2-13-10 10"
1.194
Mixed_9t btbossh Louay
4.12
3.60
1.138
3.65
0.00x
Misol, SthnmesS655.2 Jame
4.29 4.29-10
3:27 10
1.51
4.39
10
3.29 10
1.55
Mood honnchhhcddd Z.cmr
8.37 10
7:27 10
1.18x
1.38
10-2
7:31 7:31-10"2
1 1
1.72
1.42
1.21X
1.73x
Mixel_Sklemii palllan
4.12-110"
2.05 10
2.01
2:40
10
2.02 10
1.19
Minol Schrachisi comr
5.43
4.80 10
1.11x
3:27-
10
4.88- 10
0.67x
5.40-10-3
4.87 10
10
0,87s
Mited M.brnnhiil.com
7:24 7:24-10-1
6.08 10
1.09X
4.89
10
8.88 10
0.73x
Mive__SdruuctSe Looay
7:25-10
6.67 10
1.09X
4.89
10"3
6.59
10-3
0.73x
Minal banr.leaaacccas.ommm
7.76-10-
3.28.
10
2.41x
Mixel honchhdddi 3oom
1,399
3.50
3.68X
1.41
10
7-15
10
1.97
Mixal fkbenchaal.ccm
6.56-10
5.68 10
1.16%
4.80
10
4.22
10
1.14x
Mixed db hhhnmn777_1
6.35-10-5
6.02 10
10-3
6.03
10
0.80x
Mixe_ 2ooos_
2.01
3.60 10
0.50X
3.58
0.42x
Misol, dl (el.hlmme7777_3amm aame:
1.92
3.50 10
0.58%
1.46 10
5.58 10"
0.41x
Moxd thtinddheddd 2.cor
1.94 194-110
3.54
10"3
1050-100"
0.41x
3.49
Mixed &c (chmna777_11. cone
6.59 10
4.60 10
1.43x
4.80
10-3
4.90 10
0.98x
Misol Sosl_ffeur fcl.mnal77_22222
2.59
5.14 10
0.50x
2.08
10
6.08 10-3
0.41x
Mised, fechmmai777_330m
5.32
10
0.39
Mixed 6c hondooodd 2.om
2.51- 10
5.32 10
0.47%
2.05
10
5.25 10
0.39x
Mond 6i honccchhhhhlll S.cmr
2.53
5:21 10
0.49
2.05
10
5.12- 10
0.40x
Mixed, brrnnc2777222om
3.35
7.81 10
10
7,61 10
0.38x
Misel (onoo_ffllnn77777..ooo
3.35-10
7.52 10
0.45x
7.31
0.40x
1,09
6.14: 10
1.77%
3.82
10
0.10-10
0.63x
Austoggssgorrlcoom
8.85-10
1,07 10
10
1,09 10
0.78x
Mixel www..a/ccccccdddc 4com
6.93 6.93-10
2.93 10-3
2.36X
7.04
10
2.34 10-3
2.39x
Mixol, Saanl,76iranchisl.ooo
3,27.10-3
1.82 10
1.30
2.39
10-3
1.76- 10-3
1.36x
Thhnnci3e3..
10
Mixed Thhranchix. 2a.com
2.51
2.34 10
1,07X
3.03
10
2.32 10
1.31x
Mited. 75.tmnchis3 2hcom
2:43-10
2:24 10
1.08X
2:95
10
2.16- 10
1 1.38X
Mixed waa._//honnhbx333...oo
3,70-10-3
2.57
1.44X
10-3
10-3
1.17%
2.03
1.45. 10
1.40x
1.94
1.40
1.39
Mised 'hhranch pool
2,89
1.57 10
1.85X
2:26
10
1.57 1.57-10 10
1 1.44x
Mixed Wctnncciii..coo
7,88-10"
7.68
1.01X
10
7.86
0.88x
7.88
7.68
Mixel_ (ono___.ioccchb3BB....
7.92 10
7.67 10
1.08X
6.74
10
7.57 10
0.85x
Mixol, www.:7cceannn.lllle ponlame
7.92-10-3
7.67 10
1.08x
0.74
10-3
-
107-10-3
0.88x
ii MobileNetV2, input shape (32, 3, 256, 256)
opt [s]
feawess0.
2.30 10
3.63 10
6.34x
2.24 10
3.65 10
6.13x
Loom Locu
2.84
3.71
0.78x
2.84
10
3.72
0.76x
festores Lonc
7-35-10 10
5.48-11
1:31s
2.55
10
5.50
10
0.54x
4.28 10
3.92-100
1.47s
1.50-10
2,91
10
0.51x
1.0m
3.98
1.59X
2:51
10
1.58x
5.83 10
4.67 10
1.14%
3.06
10
5.04
10
0.61x
festurex 0.0
1.63 10
1.37
1.19
7.02
10
1.39
10
0.50x
feoturs...mm
2.07
3:68
9.72 10
9.38 10
1.04x
6.47
10
9.38
10
0.69x
festures 1.0
1.15- 10
1.02-10"
1.23x
1.15
10
1.02
10
1.13x
2.82 10
2.65-10-3
1.06%
10
2.64
fetures.5.cmm000 00
1.05
7.05
1.49
3.84
10
7.08
10
0.54x
1.0m 1.0
6.88 10
1.19-10"
0.54X
6:36
30
1.19 10
0.53x
3.39
3.16-10
1.07%
0.86%
f.0mre77.00.110
3.60
3.66-10
1.00x
3.09
3.67
1.01x
1.41 10
1.28 -0-81
1.10%
7.93
10
1.28 10
0.82%
3almmxx80000000 0.
9.96.10
6.-8
1.61x
3.57
10
6.28
10
0.54x
feouues
2.88
2.77
5/25
0.46X
2.36 10
2.24
1.06%
1.55 10
2.24
10
0.69%
feanues. 21.c0v22
2.33 10
2.22-10
1.00x
1.55
10
2.24
10
0.69x
(ssuns.12ssm.h0
10
7.08-10-4
10
7.07
10
fecuns.11 12mmm.12..ww 12.omm.10 10
1.07
8.89.10
0.46%
4.04
8.90
10
0.45x
fesmues 12.cau22
3.97 10
3.84
1.03x
2.97
10
3.85
10
0.77x
14smm.1
2.66-10 8
10
2.66
10
fesnues
1.50
1.22X
9.00
10
1.23
0.73x
lesuno. 15
9.14 10
6.34
1.44s
3.58 10
6.23 10
0.54x
forons. 15 10
9.80-10
4.0
0.21x
9.83
4.03 10
0.24x
feanres 15com
2.57
1.10x
1.85
2.35
10
0.79x
fratues. 2
2.67 10
2.35
1.10s
1.85
10
2.35
10
0.79x
180
1,15 10
7.91-10
1.46x
4.79-10
7,91
10
0.01x
43

---

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.
10
TN (logarihmic)
10"
10
General
Donse mix
Dense
Down
Figure F23: KFAC-reduce factor performance ratios of TN versus PT and TN+opt versus PT for different convolution types GPU.
44

---

Table F9: KFAC-reduce factor performance comparison on GPU,
(a) 3c3d, CIFAR- 10, input shape (128, 3, 32, 32)
Name
IN Bl
PT [s
Factse
IN + opt [s]
FTB]
Foctor
200.111
5.88-10
2:25 10
0.39x
8.59 10
2.41 10
0.36x
2.092.0
1.41- 10-3
1.79 10
0.79x
1.29 10
1.75 10
0.74x
1.33-10
2.31-101
0.57%
10
2.37
(b) F-MNIST 2c2d, input shape (128. 1, 28, 28)
Name
IN bi
PT II
Fictor
IN e opt [6]
PT
Factor
1.10-10
0.68x
1.01 10
1.83-
0.55x
20mV2.1
1.88 10
2.37 - 10
0.62x
1.34
10
2.76
10
0.56x
(c) CIFAR-100 AlI-CNN-C, input shape (128, 3, 32, 32)
Name
TN
Factor
Tp-
PT[]]
Factur
20m111
1.11 10"
1.84 10-3
0.60x
1.07 10
1.80
-3
0.56x
20092.1
3.26
0.39x
3.39 100
9:24
-3
0.39x
2um3.11
3.00 10
3.56 10
0.87s
3.08- 10
3.56
10
0.87
com4.11
1.45-100 1055-100
3.09 10
0.47%
1.46-10"
3.00
10
0.477
2.59-10 10
0.48%
2.55-
5.62
com6.1
2:46-101
3.49 10-3
0.71x
2.43 1.43-10-3
3.48- 10-3
0.70x
200770
1.65
3.03 10
0.51x
1.53 10
3.02 10
0.51x
am81
1.14-10
1.46-10
0.78x
3.59
20m9
1.14-10"
1.46 10
0.79x
3.59
1.30
0.28x
(d) Alexnet, input shape (32, 3, 256, 256)
Name
TN
PT
Factur
TN + opt [s]
[T
Factor
1.80 10"
4.22 10-3
0.44
1.54 10-3
4.22 1073
0.44x
3
1.60-
10
3.88 10"3
0.41x
1.53- 10
3.89-10" 10
0.39.
featree.6
1.51-
1.43
1.63
festres.8
1.77
10
3.02 10-3
0.59x
1.78 10-3
3.02 10-3
0.67x
Sratmoo
1.56 10
1.96 1003
0.79x
1.51-10-
1.96-10-3
0.77a
(e) ResNet18, input shape (32, 3, 256, 256)
[N
PTB]
ost
PT
Factor
1.79 10
5.41 3
0.35%
1-78
10
6.40 10-3
byart B
2.26 10
5.35
0.422
2.20
10
5.32
0.41
2:23 10
1.99
10
1.12x
2.16
10
1.96- 10
1.10x
1.47 10
3.02
10-3
0.49X
1.47-10- 10-5
3.04-10-5
0.49x
1.85
7.70
10
2.40X
3.24
10-4
7.09 10-4
0.43x
bayer&flaserr
1.48 10
1.21
10-3
1.21x
1.45- 10-3
1.21-10 10-3
1.20.
1.49
1.96
10-3
0.78x
1.36 10
1.95-10- 10
0.70x
1.25
5.43
2.31X
5.44
0.49x
layen' comm
1.49 10
1.31
10-3
1.12X
1.44
10
1.33
10-5
1.085
layerk...oww2
1.60 10
1.86
10
0.86X
1.62
10
1.86 1..6-1003
0.57
9.63
5.25
2
(1) ResNext101, input shape (32, 3, 256, 256)
Name
TN
PT
Factor
TN opt
PT
Factor
oand
1.78 10
5.38
0.39
1.77
6.39 10-3
0.33x
byert Baoor
1.87 10
1.73
10
1.08
4.47
117-10 10
1-73- 10
0.26
byerl floom?
4.44 10
1.91 10
2.29
4.43
1133-100 10
1.94-10 10
2.29:
layert 0.00m3
6:08
5.57
1.21
5.57
0.22x
6.00 10
5.57
10-3
1.09x
1.21 10
6.58 10-3
0.22x
2
1.37 10
1.18 10
1.16x
1.37
10
1.18- 10
1.16:
3:02
1,26x
3-02
6.08 10
1.77
10
3.44x
7.
10
1.77 8
0.40x
1ayer21 Leom22
4.16 10
8.91
10-3
0.42X
4.16 10
9.90 10
0.42x
3.81
3.02
10
3.02
7.88
8.42
1.29
7.90
6.43-
1.23
Iaper3...oom3
1.61 10
1.78 10
D.9IX
5.
10
1.90-10"3
0.30x
3.80 10
1.17
10"3
3.21x
5.10-10 10
10-1-1005
0.42x
2.25
5.41
0.42X
2.26
5.41
0.42x
1.61 10
1.77 10-3
0.91x
6.44 1.44-10-1
-
0.90
2oom
4.21
5.28
10
0.80x
4.21-10-8
10-
0.80x
1/23
1.45
0.85X
1.44
0.63x
byerdn.dwwnamphe
1.62 10
8.91
1.82x
4.97
10
8.83-
10
0.56x
Layen21Loorr2
2.18 10
4.73
10
0.46x
2.16
10
4.72
10
0.46x
(g) ConyNeXt-base, input shape (32, 3, 256)
TN
featues0.0
1.72 10
1.02 10-3
1.69x
7.52 10
1.01 100
0.74x
Statuos 105ld00
1.88 10
0.35x
1.58 10-2
4.40 10001000
0.35x
fratuecs 2.1
3.80 10
1 10
1.91
8.14 1.14-10-
1
10
0.484
Seatures 3.
10
0.37x
6.19-10-3
2:22 10
0.37x
4.62 10-3
1.18 1
4.57 10-3
1.16 10
0.40x
1.40 10
1.10- 10
1.27s
9.28-1001
1.02-100
0.91x

---

(h) Inception 33, input shape (32, 3, 299, 299)
Nome
TN [s]
Factor
TN-
opt [s]
PT(s]
Foctor
Cuav2[_ Ax
2.39 10
2.03 10
1.18%
2.36
10
2.00 10
1.18x
Cuay2 Ze, 3x3,cov
4.42 10
1.30
10
0.34x
4.38
10
1.30 10-2
0.34x
1.30
1.30-100 10
0.93
Comvd 1h bel.
1.32 10
2:18
10
6.12x
5.53
2:18 10
0.28
Chav28,4a 3dl
2.72 10
8.00
10
0.34x
2:74
10
9.02 10-3
0.34x
MixeL Sttnndiiiissooo
10
1.57
10
10
0.29x
BaveL_9.bnnna5455 Louay
1,43
1.57 10
0.91x
0.29x
Misol, St anm
1.56 1.56-10
3:73 10
0.42x
1.35
10
10
0.36x
Mood honcchhhcddd Z.cmrr
1.57 10
1,91
0.82X
1.44
10
1.01
10
0.76x
2.65
0.56X
0.53x
Mixol_Sklemmit pallaan
1.45-10
1.50 10
0.91x
4.62
10
1.50 10
x.29
Minol Schrachisi comr
1.79.10
2.00 10
0.90x
5.48
10
2.00- 10-3
0.27x
SchmmniS65__
1.79-10-3
2.00
10
0.27s
Mited M.brnnhixl.com
1,99
2:22 10
0.90x
5.90
10
2.21
10
0.27X
Shruusddrruc Loony
1.97-19-3
2:20 10
0.89x
5.83
10
2.18 10-3
0.27X
Minal falr.eeaaaccccccoommm
281-11
3.18
10
0.87s
Mixel honchhlddi 3oom
1,30-107)
1:22
1.07
1:44
1,28 10
1.11x
Mixal fkbenchaal.ccm
1.49-10
1.64 10
0.91x
5.47
10"4
1.82- 10"3
0.34x
Mixed db hhhnmn777_1
1455-10
1.64 10
3.60
10
1,64- 10
0.34X
1.07
0.94X
0.09x
Misol, d 3el.hmmmm777733mmm Name
9.45-10
1.25 10
0.77
7.61 10
1.23 10"
0.62x
Moxd bhiiddddeedd 2.cor
1.06-10"
1:23
0.86x
10
1.23-10" 10-110"
0.621
1.10-10
0.89
Mixed &c (chnna777_1. cone
1.43-10
1.62 10-3
0.88X
5.47
10
1.54 10-3
0.33x
Misol fe 4chrmnd77772222
1.01:10"
1.34 10
0.75x
7.84 10
1.36 10-3
0.56x
Mised, fcchmmeh777_33amm
1007-1073
1.69
10
0.48x
Mixed 6c hondmoodd 2.om
8.39
1.67 10
0.51x
1-67
10
1,88
10
0.45x
Mond 6i tenccchhh/hlll S.cmr
1.01-10
1:33 10
0.76x
7.61
10
1.35 10
0.57x
Mixed, brrnnca777_22oom
1,01-10-3
1.48
10
10-3
0.61x
Misol (onos_ffllmns7777..ooo
9.63-10"1
1.77 10
0.54x
7.38
1.79
0.41x
9.97-10
6.04 10
1.65%
3.60
10
8.58- 10
0.58x
Austoogssgoorlcoo
1.05-10
1,09
10
1,09
10
0.85X
Mixel (oo.."aaorocchho 4com
1.34-10-3
1.14 10-3
1.18X
1.36
10
1.14 10-3
1.19x
Mixol, Saaul,7hbrnnchiss.oom
1003-10-
9.07 10
1.13X
5.08
10-1
9.07 10
0.58x
Thinnci3e3..
9:30
1000-10-4
Mixed Thhranchh3. 22.com
1.13-19
7.94 10
1.43X
6.89
10
7.50 10
0.87X
Mited. 75.tmnchis3 2hcom
1.07-10
8.53 10
1:25X
7.60
10
8:51- 10
0.89x
Mixed waa._//honnhtt33A....
1.16-10
9.32 10
10-4
10
0.68x
1.55
1.08
L.56
1.02
Mised "hhranch pool
1.16
6.79 10
1.71x
5.19
10
8
10
0.76x
Mixed wcwnncciii..coo
1.45
10
0.63x
1.21
1.44
Mixel_ (onol_/.iacctt33al.... Loom
1.21 10
1.43 10
0.84x
7.66
10
1.44
10
0.53x
Mixol, www.:7/eeeeehhlllleee poilame
1.21-10-3
1.44. 10
0.84x
7.80
10
1.45-
10000
0.54x
(i) MobileNetV2, input shape (32, 3, 256, 256)
TN
opt [s]
feawess0.
1.90 10
1.66-10-
1.15x
1.91
10
1.68
10
feauure Loom
2.60 10
9.57
0.27
2,70
10
9.89
0.27x
festores Lonn
1.12-10 10
3.03
3.70s
7:12 10
3.00
10
0.24x
1.80-10
1.77 10
1.02s
4.64 10
1.70
10
0.26x
2.0m
7.01
9:06-107
9:05
2.50 10
2.38 10
1.09
6.08 10
2.40
10
0.25%
festurex
1.44 10
0.19-10"1
1.57
2.96. 10
9.40
10
0.31x
2.99
1.12
3.65 10
3.38 10
1.08%
7.52
10
3.40
10
0.23x
festures 1.0
3.01 10
3.76-10
0.80x
2.99
10
3.77
10
0.79x
1.38 10
1,16-10
1.16
fasures 5.00m 00
8.51
5.17.10
1.65
2.77
10
5.34 10
0.52x
1.0mres...0m110 1.
1.38 10
3.99
0.34X
1.36 30
3.99 10
0.34x
1.68
1.36-10
1.35
f.0mre7..00.11
1.37
1.00
0.80X
8.50 10
7.05
1.22
2.52 10
7.00 10
0.38%
3 0.
8.45
4.92-10-4
1.72s
2.19. 10
4.93
0.51x
feouurs...om111
1.16-10
1.12
2:35
8.78 10
1.30-10'
0.94%
3,06 10
9.29 10
0.33x
feanues. 11.c0v22
9.89 10
9.49-10
1.04x
3.05
10
9.25
10
0.33x
(csuns.12ss..h0
10
6.32-10--
10
5.14
10
fecuno.11 1amrm.12..oo// 12.omm.10
1.51
3.23-10
0.47%
1.27
10
3.22
10
0.39x
fesmues 12.cau2
1.14 10
1.24-10
0.92X
3.94
10
1.17
10
0.34x
14ssm.1
1.61
1.61-10
1.61
fesnues
1.14
1.67%
3.67
8.80
0.54x
lesuno. 15
9.58 10
5.23 5.23-10-4
1.82s
2.74-10
5:23 10
0.52x
feruns. 15
1.41 -00
2.25-10-3
0.63x
1
2.25 10
0.61x
feanues. 15comv1
1.15
5.81-10
1.31x
4.45
8.80
0.51x
fratues.11.. 2
1.16 10
6.80
1.31x
4.56
10
8.58
10
0.51x
180
9.51 10
5.40-10
1.76x
2.50
-00
3,22
10
0.48x
46

---

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 (Cin, Jt, I2), output of shape (Cra:: O1, O2) and kernel of shape (Com, Cin K1, K2)- The analysis with 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 the computation of the averaged unfolded input [X] (W)))) := 1/(0, O2) 10,00 O, [X] which consists of Ci K K2 numbers. In the following, we will look at the extra memory on top of storing the input X, the averaged unfolded input [X] (wvg), and the result S2,
Default implementation:
The standard implementation computes [X] (avg) via the unfolded input [X] and thus requires extra storage of C K2 (20002 numbers.
TN implementation (general case): The TN implementation requires storing the averaged index patterns (+.avg) 1/o, Ee [n lo, for 1.2. These are directly computed via slight modification of Algorithm D1 and require storing 7, K2 + I2K2 numbers. In contrast to the default implementation, spatial dimensions are de-coupled and there no dependency on Cin-
TN implementation (structured case): For structured convolutions (Figure 5) we can describe the action of the index pattern tensor through reshape and narrowing operations. MI. libraries usually perform these without allocating additional memory. Hence, our symbolic simplifications completely eliminate the allocation of temporary intermediates to compute [X] (avg)
Empirical results: To demonstrate the memory reduction inside the computation of S2 we measure its peak memory with the memory-profi ier library and subtract the memory required to store X and 52. This approximates the extra internal memory requirement of an implementationn. 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 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 vis- ible for many channels with large kernel and output sizes. One example are ConxNeXt-base's f features. 1. block. 0 convolutions with K₃ = K₂ = 7, O O₂ = 64, and Cin 128 (Ta- ble E4g). For those convolutions, we observe that the default implementation requires an additional 3, 1440 MiB (x= 3 GiB!) of memory, whereas the TN implementation has zero extra memory demand (Table G10g). This consistent with our theoretical analysis in that the overhead is storing the unfolded input, which has (N 32) (Cin 128) (O, = 64) (02 64) (K1 7) (K2 7) 822, 083, 584 float32 entries, corresponding to 3,136 MiB.
47

---

Table GIO: Additional internally required memory to compute the KFAC-reduce factor (measured on CPU). The value 0 indicates that an implementation's is peak memory matches the memory consumption of its input X and result S2
(a) 3c3d, CIFAR- 10, input shape (128, 32, 32)
Name
TN [MB
TN + [MEB]
FT [MIB]
Type
20m1 1
0.0
0.0
0.0
General
c0m2
0.0
0.0
0.0
General
20003.1
0.0
0.0
0.0
General
(b) F-MNIST 2c2d, input shape (128, 1, 28, 28)
Nume
IN (
IN (MBB)
PT (MIB]
Type
cumi 1.1
0.0
0-0
0.0
Geurrll Gruurll
211
0.0
0.0
0.0
(c) CIFAR- 100 AII-CNN- C, input shape (128, 3, 32, 32)
[N
opt [MIl)]
PT [MBB]
Type
20
0.0
0.0
0.0
General
2.1
0.0
0.0
431
Geneval
20033.1
0.0
0.0
0.0
General
comi.11
0.0
0.0
Generl
cmm5.1
0.0
0.0
215
Geural
cum6.l
0.0
0.0
Geural
com70
0.0
0.0
General
combi
0.0
0.0
Drno
com19.1
(d) Alexnet, input shape (32, 3, 256, 256)
Nome
IN [MBB]
IN opt (Mt()
PT [MBB]
Type
0.0156
175
Gaeenl
3
0.0
Gaseal
Genera
0.0195
Omenl
(e) ResNet18, input shape (32, 256, 256)
Nume
IN [MBB]
TN opt [MBB]
PT [MBB]
Type
ornvl
0.0
283
Generll
byerl
0.0
287
Ganeal
31.7
0.0
71.1
Grnzal
lnpert InpertB..cmm D.com?
0.0
0:0
143
Genen
byyer pprraaaaammmmmppoo
0.0
0.0
0.0
Down
0.0
0.0
0.0
General
lapers
0.0
0.0
70.8
Genera
layer3 Aper3.0.dame Sddwsssmpll00
0.0
0.0
0.0
Down
luyer' cooor
0.0
0.0
0.0
Genna
0.0
80.3
0.0
Genenl
bepzek Saoeesiiaoo
0.0
0.0
0.0
Down
(1) ResNext101, input shape (32, 3, 256, 256)
[MiB]
apt [MiR]
[MiR]
Type
cotv1
0.0
0.0
293
General
layeri 10.0m11
0.0
0.D
0.0
Dense
luyer1 f0oom2
576
876
1150
Generl
luyeri 3
128
0.0
127
Dense
1
128
0.0
127
Dens:
byer2....o222
256
258
576
Genzal
0.0
Dense
128
0.0
19.3
Down
byper? 2omm.
0.0
575
Gonel
l.por...comr1
0.0
Dense
124
Generl
0.0
0.0
0.0
byork.1
144
lopent coom
boyoed dowesssrio
0.0
144
General
(g) ConsNeXt-base, input shape (32, 3, 256, 256)
Namse
IN [MBB]
IN (XtBB)
PT [MBB]
Type
fratures
0.0
Dena:
fratunos 1.blook0
0.0
$1400
Guann
features 2.8 featuees bllckkk
0.0
0.0
0.0
Derse
0.0
0.0
1570
Gmenl
featues-4.. Seatures 5.ees.50:00 2toecko
0.0
0.0
0.0
Desse
0.0
0.0
784
Geneni
feamues.ó, 6.1
0.0
0.0
0.0
Desse
features 7.0 7.0.Mlxck.0
0.0
0.0
392
Genzal
48

---

(h) Inception V3, input shape (32, 3, 299, 299)
[MIB]
opp [MBB]
Type
Cuav24 363
54.6
0.0
73.0
General
Cusv2d 2r, 3o3.c
86.7
56.7
739
General
Can22 Th, Sen
84.4
64.4
758
Gouzal
Can22 3t, Lalo
166
0.0
0.0
Dane
Concld da Sehcome
52.0
0.0
482
General
Mixel Sh hhbonchiss
0.0
0.0
0.D
Deno
Misedl hracche55 Low
0.0
0.0
0.0
Dane
Mised branctt55
375
General
Mixed bhtnchhhhcdd 2.comr
84.8
General
Mixel whw.shhaaebbbBBB..
125
Mixel_5h Shbeaadac pooleer
0.0
Mixal Schhnachi
Mixed, S...annSSS,_.amm
0.0
0.0
Oh.monchhxl.ccom
0.0
Mixe_. brnnca777_111
0.0
Dense
Mixe_.5 wbw.banc377772--0
0.0
Dasse nix
Mixed boanha777777.0
0.0
0.0
Dasse nux
0.0
0.0
Daree nix
0.0
0.0
Dasse nix
Mireed schrnnch7c?
0.0195
0.0
0.0
Denoe
schemnche77 2mw
0.0155
0.0
0.0
Darsc mix
Mixed fe (chaana777_3.c0m
0.0
0.0
0.0
Dense mix
Mixed, 6c (c.2nnnh00000
2.cmm
0.0
0.0
0.0
Dense mix
Mixel fcb.r...
0.0
0.0
0.0
Danse mix
Mixed_ sccbanda77__22o0
coms
0.0
0.0
0.0
Dense mix
Mixed Se Jeleanda767__300
0.0
0.0
0.0
Dasse mix
Autlagiasccmmaann
0.0
0.0
0.0
Dane:
Auntagio.cmmm Lants
0.0
0.0
0.0
Genzal
Mosed To Iabrandt333 2com
0.0
0.0
0.D
General
Mindl 7a Tsloanchhsss 4emm
0.0
0.0
0.0
Gereral
Mixed_ 7h.rmmchiii..oo
0.0
0.0
0.0
Denie
Mixed, 76 bhhrnc233331 comr
0.0
0.0
0.0
Dense
Mised. 7h Thhhanchhcc Ja.cm
0.0
0.0
0.0
Drese mix
Mised, 763mmach333 2hcom
0.0
0.0
0.0
Dzsse nix
Mixel, 7h Thbuud3222BLLam
0.0
0.0
0.0
Denie
Minol 7h 7hleadchshhll 2omm
0.0
0.0
0.0
Gonzal
Mosol_ 7hteandh pall
0.0
0.0
0.0
Dane:
Mixed Pe/mrachix
0.D
Dense
Mixed haaranhh cow
0.0
0.0
Denor
Minel Trhhanchhhhhbl Loom
0.0
0.0
0.0
Denor
Mixed_7c.rande com
0.0
(i) MobileNetV2, input shape (32, 3, 256, 256)
Name
[N
TN + (MR)]
PT [MiR]
Type
feature
0.0
0.0
531
General
features com
26.1
26.
5776
General
Seatues LoaueeLLooooo Leans]
128
0.0
63.8
Dense
feutures 0.0
0.0
0.0
0.0
Dence
features 2arres2.c0om 2.cmm 1.0
192
192
432
General
Stature 2xtr
0.0
0.0
0.0
Dene
feutures 3.cm 0.0
0.0
0.0
0.0
Dane
fraturzs Leom 1.
34.1
70.4
648
Geura
Scatince S.mtn:
71.7
0.0
71.1
Dene:
femures..com 1.0
59.5
35.7
162
General
Seatires c
0.0
0.0
0.0
Demo
feutures 5.00m 0.0
0.0
0.0
0.0
Dene
feutures 2.00m.10
0.0
0.0
215
General
Seatrres..om comt
0.0
0.0
0.0
Dence
festures 7 7. 1.0
0.0
0.0
58.3
Gammal
Seature 2omm
0.0
0.0
0.0
Done
fraturzs comm 0.
0.0
0.0
Dema
fraturzs 1.6
0.0
107
Geand
Seatres8..om comv2
0.0
0.0
1.onv:2
0.0
0.0
Deme
feature.1 12 comit1
0.0
0.0
Dema
feumers.12.oov22 fcatues.1..cm.1.0
0.0
Geseral
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
49

---

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 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 S1, S2, S3 be index tuples of the input tensors A, B, C. and S4 S (S1 U U U3) valid output index tuple of their tensor multiplication D S2. (A, B. C). We can either first multiply A with B to obtain an intermediate tensor of index structure S1.2, or with C to obtain an intermediate tensor of index structure S2,3, before carrying intermediate index structures, we divide the indices 5 (S S2 U Sa) \ S4 that are summed over into those only shared between A, B given by $122 = (S) U S2) \ (S, Ug). and those only shared among B, given by 52,3 = (S2 U Sa) (S+ S.). This yields the intermediate indices S1,2 = (S1 U S2) \ S1,2 and S2,3 (S2 U S3) S2.3. and the parenthesizations
the remaining multiplications. To construct the
54 E EEs [A] s
B S²
C Sa
EEs Als, Σ
So.a
B
Sq
C|
5sss)
(H20)
D 29.. (S...S,,
A, B), C) *(S.. (A,
(B, C))
This generalizes n-ary multiplication, allowing 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 provide 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
[C]i, [[AA,,
The index tuples are SA (i.j), SB kj, and Sc (i, 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:
(SA US \ Sc ((i.)) (1, \1,
Now we see that matrix-matrix multiplication is a case of tensor multiplication (Equation (3)),
[C]se
Σ
[A]sa B]ss (A, B)
(SAUSH))S8c
50

---

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 bullet point list of our contributions in $ 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 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
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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See $4.2 and SA.
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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
51

---

The TN simplifications we provide in $4.1 follow straightforward from the index pattern's 's structure Equation (7) and are stated rigorously in $D.3, including their assumptions.
Justification:
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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide 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 properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide implementation details in SD, experimental and hardware details in SF and G.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
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 a specific model and empirical evaluation, 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 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 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 submis- sions 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 reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing verifying the results.
5. Open access to data and code
52

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will open-source the code reproduce all our experiments, well the raw data containing the results shown the manuscript.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
public/gun Staii//ggiidd///oeeutiiiti
onPollooooooooo for more details.
While encourage the release of code and data, 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wwtppccccc//cccccccccccccccccltt /gguides/////eeeeerre /Code/uummmmmic cotbbmisssii ssisnsolicyy 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 proposed method and baselines. only 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 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See SF and 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, as supplemental material.
7. Experiment Statistical Significance
Does the report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
53

---

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 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 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 figures symmetric error bars that would yield results that are out 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 com- puter resources (type of compute workers, memory, time of execution) needed 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 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 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https /tt///eerrip../// hccppbblcc/thhcccceeeeeee
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws 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.
54

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
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)?
Question:
Answer: [NA]
Justification:
The paper does not release any data or models that have 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
55

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 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/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS 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 crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
56

---

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


Reviews:
[{"review_id": "3k4jWEd1NN", "reviewer": "Reviewer_iQrS", "summary": "### 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.", "questions": "Please see weakness", "limitations": "Please see weakness", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "### 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.", "weaknesses": "### 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer iQrS,\n\nthanks a lot for your thorough review; we appreciate the work you put into it and are glad you find the paper innovative and support our idea of making the complex yet powerful tensor network toolbox accessible to the ML community.\n\nWe would like to inform you that we conducted a new experiment with a real second-order method based on KFAC-reduce (see our [global rebuttal](https://openreview.net/forum?id=cDS8WxnMVP&noteId=dKv8oyjgLS)). Our TN implementation was able to halve the second-order method's run time overhead compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2). We hope this provides further evidence that our approach is indeed useful to advance second-order methods and close their computational gap.\n\n## **Weaknesses 1 & 2**\n  > Using einsum for convolution operations is generally not an optimal approach\n    \n  > The einsum-based approach favors specific architectures\n\n  You are right that there are additional degrees of freedom to further accelerate our TN implementation. To give some intuition, we would like to point out that one main problem is that we cannot always leverage the index pattern's sparsity: For many convolutions, we can exploit the sparsity thanks to our proposed symbolic simplifications. However, whenever they cannot be applied, we must consider $\\Pi$ as dense because PyTorch's `einsum` does not allow contracting tensors with mixed formats. This means we must sometimes carry out multiplies by zero, which is wasteful.\n  \nWe attempted to address the scenario where our simplifications cannot be applied and have ideas how to progress on this frontier. We will describe them in an updated version of the paper as follows:\n  \n  - We experimented with TACO [2], a tensor algebra compiler that can optimize contractions of mixed-format tensors. Unfortunately, TACO's Python frontend cannot yet generate GPU code, and on CPU we ran into a memory leak in the Python frontend.\n\n  - It would be interesting to benchmark our approach on specialized hardware, e.g. [cerebras' wafer-scale engine](https://www.cerebras.net/product-chip/) processor, which [supports unstructured sparsity](https://www.cerebras.net/blog/sparsity-made-easy-introducing-the-cerebras-pytorch-sparsity-library). This would allow to remove many multiplications by zero. While we do not have access to such a chip, we are excited our approach could benefit from such hardware in the future.\n  \n\n- To the best of our knowledge, we believe that contraction path optimization of mixed-format tensor networks has not been studied sufficiently. This may be due to the absence of hardware that can leverage this information to speed up contraction. We hypothesize further performance improvements might be possible by taking the operand sparsity into account when optimizing the contraction path.\n\nPlease let us know if you think there is an additional approach we could try to boost the performance of our framework.\n\n## **Weakness 3: Limited comparison with optimized convolution algorithms**\n\n> The comparison between TN and PyTorch does not specify which convolutional algorithm is used in PyTorch's implementation.\n\nYou are right that we do not mention which convolution implementation is used in the comparison. We use PyTorch's [`torch.nn.functional.conv2d`](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv2d.html), which automatically selects a 'good' implementation according to heuristics. We think this is a good baseline implementation as most engineers and researchers rely on it in practise.\n\n> Furthermore, the paper does not compare the TN approach with other highly optimized versions of convolution\n\nConsequently you do have a point that we did not look into comparing our approach with different convolution implementations. We found that PyTorch indeed has [different convolution algorithms](https://github.com/pytorch/pytorch/blob/1848cad10802db9fa0aa066d9de195958120d863/aten/src/ATen/native/cudnn/Conv.cpp#L486-L494) implemented. However, it seems like one [cannot specify manually](https://discuss.pytorch.org/t/manually-set-cudnn-convolution-algorithm/101596/2) which algorithm should be used.\n\nWe will try to look into this and identify other convolution implementations in PyTorch in order to provide a more fine-grained comparison. If you have any pointers, please let us know.\n\n## References\n\n[2] Kjolstad, F., Chou, S., Lugato, D., Kamil, S., & Amarasinghe, S. (2017). TACO: A tool to generate tensor algebra kernels. IEEE/ACM International Conference on Automated Software Engineering (ASE).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MOno9ejadq", "reviewer": "Reviewer_cYbp", "summary": "The paper presents a novel approach to simplifying and optimising TN operations for CNNs. The authors introduce an abstraction for tensor networks to efficiently handle convolutions and leverage TN simplifications to improve computational performance. The paper details the theoretical foundations, implementation strategies and results. The results suggest improvements in run-time and memory efficiency over standard methods.", "questions": "Questions 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?", "limitations": "There is a section discussing limitations, but the authors could mention more practical factors. Such as: scalability issues and how the methods perform with larger datasets (should be straightforward with synthetic data), potential biases introduced and their impact on different types of data and potential real-world scenarios where the methods might fail or underperform.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The 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", "weaknesses": "The key weakness is that the performance demonstrated is not very convincing in many cases. This is understandable as the authors mention, standard routines have been aggressively optimised whereas the TN counterparts have been explored significantly less. However this by itself is not a reason to reject given the potential further research that can be built on this work and the interest it \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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer cYbp,\n\nthanks for your strong support and detailed review! We will apply your suggested improvements to the text.\n\nWe would like to inform you that we conducted a new experiment with a real second-order method based on KFAC-reduce (see our [global rebuttal](https://openreview.net/forum?id=cDS8WxnMVP&noteId=dKv8oyjgLS)). Our TN implementation was able to halve the second-order method's run time overhead compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2). We hope this provides further evidence that our approach is indeed useful to advance second-order methods and close their computational gap.\n\n## Questions\n\n- > Why are the (non simplified) TN implementations slower in general? Do the standard implementations have further efficiency optimisations?\n\n  This observation is correct. Non-simplified TNs perform contractions using a dense version of $\\Pi$, which is sparse. This causes wasteful multiplies by zero. Our simplifications reformulate the contraction with $\\Pi$ into `reshape`s or `slice`s, which remove these zero multiplications from the TN (see e.g. the bottom right panel of Fig. 1) and thereby speed up the computation.\n\n  Whenever simplification is not possible, the main difficulty for using the index pattern's sparsity is that PyTorch's `einsum` only accepts dense tensors. We experimented with alternatives that support mixed-format tensors, and have some ideas how to make progress on this frontier in future work. If you are interested in knowing more, feel free to have a look at our [response to Reviewer iQrS](https://openreview.net/forum?id=cDS8WxnMVP&noteId=u5s1zM6z7X).\n\n- > Does figure 6 include variability wrt contraction order?\n\n  This depends on the heuristic used to find a contraction path. Generally speaking, the heuristic depends on the number of tensors in the TN. The tensor networks in Fig. 6 contain at most 6 tensors. After looking into `opt_einsum`'s internals, we believe that the TNs with at most 4 tensors do not contain variability w.r.t. contraction path schedule, while the others might:\n\n  In all our experiments, we used `opt_einsum`'s `contract_path` method with the default `'auto'` strategy to find contraction schedules. `opt_einsum` supports many strategies and according to the [documentation](https://optimized-einsum.readthedocs.io/en/stable/path_finding.html#introduction), using `'auto'` \"will select the best of these it can while aiming to keep path finding times below around 1ms\". The exact behavior depends on the number of tensors to be contracted and is described [here](https://optimized-einsum.readthedocs.io/en/stable/path_finding.html#performance-comparison): For TNs with up to 4 tensors, the optimal schedule is searched. For TNs with 5-6 tensors, a restricted search is carried out. We are unsure whether this restricted search is deterministic or stochastic, so it may be that this approach causes some extent of variability in contraction order.\n  \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\n  This is indeed an interesting question. We did not look into this and decided to rely completely on `opt_einsum`, which uses a sophisticated heuristic based on [performance comparisons](https://optimized-einsum.readthedocs.io/en/stable/path_finding.html#performance-comparison) to strike a good balance between contraction path search time and contraction time. We believe this is a good representation of what a practitioner would do.\n\n- > How many of the operations support GPU/hardware acceleration? Would there be performance gains if they were implemented?\n\n  All operations we propose are purely PyTorch and therefore support GPU execution out of the box.\n\nThanks again, and let us know if you have further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yscA6z2iqw", "reviewer": "Reviewer_vbTD", "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.", "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", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer vbTD,\n\nthanks for your support! We appreciate the time and effort you put into reviewing our work.\n\n> The authors show faster speed of KFAC computation. Does this implementation bring better convergence, training speed or results in practice?\n\nYou have a point that we did not evaluate the impact of our achieved speed-up on a real second-order method. We are happy to inform you that we will add such an experiment to the main text (please see our [global rebuttal](https://openreview.net/forum?id=cDS8WxnMVP&noteId=dKv8oyjgLS) for a detailed description). In this experiment, we were able to halve the computational overhead of a KFAC-based method compared to SGD, and dramatically reduce its peak memory (sometimes by a factor of 2) to a value very close to that of SGD. \n\nA direct consequence of these findings is that they increase training speed, or allow using hardware with less memory.\n\nOne could also experiment with more frequent pre-conditioner updates or larger batch sizes to obtain convergence in fewer steps. We believe these are interesting directions to explore in the future that so far have been disregarded due to the high computational extra cost. To produce statistically significant statements and recommendations though, we believe it will take a larger ablation that is beyond the scope of this work.\n\nThanks again, and let us know if you have further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Evaluation on a real second-order method

Dear Reviewers,

We are glad to inform you that we successfully applied our work to a real second-order method to complement the speed-ups of fundamental operations shown in the manuscript. Specifically, we took the KFAC-based SINGD optimizer [1] and benchmarked the impact of our TN implementation on memory and run time in comparison to SGD. We will add the experiment to the main text (details below). 

**TL;DR: Our TN implementation can reduce SINGD's run time overhead by 50\% and almost completely eliminate the memory overhead. Sometimes, it even reduces the optimizer's memory footprint by a factor of 2. This enables using larger batches or more frequently updating the pre-conditioner.**

We hope that this provides further evidence for the utility of our approach and demonstrates its capability to significantly reduce the computational gap between approximate second-order methods like SINGD and first-order methods like SGD.

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

---

**Details:** We applied SINGD with KFAC-reduce and diagonal pre-conditioners to ResNet18 and VGG19 on ImageNet 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 via input unfolding versus our TN implementation.

- **ResNet18 with inputs of shape `(128, 3, 256, 256)`:**

  | Optimizer | Per iteration [s] | Peak memory [GiB] |
  |-----------|-------------------|------------------------|
  | SGD | 0.12 (1.0x) | 3.6 (1.0x) |
  | SINGD | 0.19 (1.7x) | 4.5 (1.3x) |
  | **SINGD+TN (ours)** | 0.16 (1.3x) | 3.6 (1.0x) |
 
  The TN implementation  reduces SINGD's run time overhead compared to SGD by 50\%, and reduces its peak memory to that of SGD.

- **VGG19 with inputs of shape `(128, 3, 256, 256)`:**

  | Optimizer | Per iteration [s] | Peak memory [GiB] |
  |-----------|-------------------|------------------------|
  | SGD | 0.69 (1.0x) | 14 (1.0x) |
  | SINGD | 1.0 (1.5x) | 32 (2.3x) |
  | **SINGD+TN (ours)** | 0.80 (1.2x) | 16 (1.1x) |
  
  Again, our TN implementation halves SINGD's run time overhead compared to SGD. On this network, it also dramatically lowers the memory overhead, cutting down the peak memory by a factor of 2 from 32 GiB to 16 GiB.
  
## References

[1] Lin, W., Dangel, F., Eschenhagen, R., Neklyudov, K., Kristiadi, A., Turner, R. E., & Makhzani, A.. Structured inverse-free natural gradient descent: Memory-efficient & numerically-stable KFAC. ICML 2024.

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