conv_einsum: Efficient Tensor Convolutions

• conv_einsum is a generalized framework for representing and evaluating multilinear operations with explicit convolution indices in tensor networks.
• It couples Einstein-style notation with a meta-algorithm that selects FLOPs-minimizing evaluation paths for optimal performance.
• The framework unifies various tensor decompositions (CP, Tucker, TT, etc.) for convolutional layers, leading to significant runtime and memory gains.

conv_einsum is a generalized, einsum-like framework for representing and efficiently evaluating multilinear operations that include convolutions, especially those arising in tensorial neural networks (TNNs). It augments Einstein-style indexed expressions with explicit convolutional indices and couples that notation to a meta-algorithm that chooses FLOPs-minimizing pairwise evaluation paths. In this formulation, ordinary tensor contraction, tensor batch product, tensor outer product, and tensor convolution coexist inside a single symbolic program, allowing tensorized convolutional layers to be written as one expression rather than as an ad hoc sequence of primitive kernels (Rabbani et al., 2024).

1. Conceptual basis and scope

The framework was introduced to address a specific gap in the einsum ecosystem. Standard einsum can evaluate simple multilinear operations such as contractions, but existing implementations cannot process multi-way convolutions. This matters particularly for TNNs, where a dense layer or convolutional kernel is reshaped into a higher-order tensor and then factorized by schemes such as CP, Tucker, Tensor Train (TT), Tensor Ring (TR), Block-Term (BT), or Hierarchical Tucker (HT). After factorization, a forward pass becomes a sequence of multilinear operations whose evaluation order can greatly affect FLOPs and memory use (Rabbani et al., 2024).

The framework therefore treats convolutional tensor programs as indexed symbolic objects rather than as fixed layer primitives. In this view, a tensorized convolutional layer is not merely “a convolution plus a factorization”; it is a structured multilinear expression whose computational cost depends on how intermediate tensors are formed. This is closely aligned with tensor-network reasoning, where one distinguishes the mathematical expression from the contraction path used to evaluate it. Related work made this perspective explicit by writing convolutions as tensor networks with binary connectivity tensors or selector tensors, but conv_einsum moves the convolutional structure into the symbolic notation itself and then optimizes the resulting execution path (Dangel, 2023).

2. Notation and formal semantics

A conv_einsum expression has the form

$$\text{"input}_1,\text{ input}_2,\ldots \to \text{output} \mid \text{conv\_indices"}$$

where the indices to the right of | are treated as convolution dimensions rather than as ordinary equality-matched Einstein indices. The paper states the index roles as follows: indices appearing in multiple inputs but not in the output are contraction indices; indices appearing in multiple inputs and also in the output are batch-product or shared indices; indices appearing in one input and the output are outer-product or preserved indices; indices listed after the pipe are convolution dimensions (Rabbani et al., 2024).

A standard tensor convolution example is

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$

where $x$ is the convolution index and $b,c,d,e$ are preserved. A standard 2D convolutional layer is written as

$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$

Here $b$ is batch, $s$ is the input-channel contraction index, $t$ is the output-channel index, and $h,w$ are convolution dimensions. An important semantic distinction from ordinary einsum is that a convolution label can be used even when the corresponding input sizes differ, because convolution does not require exact dimension equality in the way contraction does.

The appendix formalizes the primitive multilinear operations underlying the language. For mode-$(k,\ell)$ contraction between tensors $$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$0 and $$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$1,

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$2

For mode-$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$3 convolution, the corresponding fibers are convolved rather than equality-contracted:

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$4

The same appendix also gives tensor batch product and tensor outer product, so conv_einsum strictly subsumes ordinary einsum-style tensor contraction. If no convolution indices are listed after |, the notation reduces to standard Einstein-style multilinear evaluation (Rabbani et al., 2024).

3. Execution model and path optimization

The implementation strategy separates two-input evaluation from many-input path planning. For a two-input conv_einsum, the paper argues that the operation can be reduced to GPU-optimized primitives such as einsum, conv1d, and conv2d, together with reshapes and permutations. The implementation classifies index roles into convolution, contraction, outer product, batch product, and self-contraction. Self-contraction can be removed in preprocessing by summation; grouped convolution handles the batch-product-like case; multiple same-type non-convolution indices can be merged into one compound mode and then restored after evaluation (Rabbani et al., 2024).

The main algorithmic contribution is the extension of netcon-style tensor-network path search to expressions containing convolution. For multi-input expressions, the framework searches over binary evaluation trees and scores each candidate merge using a convolution-aware tnn-cost rather than a contraction-only cost. The paper gives an example API,

$b$1

and reports naive and optimized FLOP counts together with the resulting path. This places conv_einsum in the same general category as contraction-order optimizers, but with convolution treated as a first-class multilinear operation rather than as a preprocessing step (Rabbani et al., 2024).

The cost model in the appendix covers contraction, batch product, outer product, and convolution. For outer product, the cost is

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$5

and for convolution the paper gives the direct-evaluation scaling

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$6

The framework also introduces a training-aware objective. If an atomic operation is $$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$7, then the training cost is taken as

$$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$8

where $$T = \mathrm{conv\_einsum}("xbc,\,ade \to xbcde \mid x", T_1, T_2),$$9 and $x$0 are the backward computations for the two inputs. This shifts optimization away from forward FLOPs alone and toward total training work.

A significant limitation is stated explicitly. If a convolution index appears in exactly two inputs, the implementation supports non-commutative convolution, treating the larger mode as features and the smaller as filters. If a convolution index appears in more than two inputs, the current implementation supports only commutative convolution with circular padding. This makes conv_einsum broader than ordinary einsum, but not a complete symbolic model of all multi-input deep-learning convolution semantics (Rabbani et al., 2024).

4. Tensorial neural networks and decomposed convolutional layers

The strongest use case in the paper is tensorized convolutional layers. conv_einsum provides a single symbolic language in which CP, reshaped CP, Tucker, reshaped Tucker, TT, TR, BT, and HT layers can be written directly.

For standard CP decomposition, a kernel $x$1 is factorized into $x$2, $x$3, $x$4, and $x$5, and kernel reconstruction is written as

$x$6

The corresponding convolutional layer is

$x$7

For reshaped CP with $x$8, input and output channels are factorized and the forward layer becomes

$x$9

For Tucker, the standard layer is

$b,c,d,e$0

and reshaped Tucker adds a core tensor $b,c,d,e$1 and corresponding rank factors. The appendix further shows analogous formulas for reshaped TT, TR, BT, and HT, making the framework decomposition-agnostic at the symbolic level (Rabbani et al., 2024).

The paper’s theorems target exactly these tensorized cases. For reshaped CP convolutional kernels, under assumptions such as

$b,c,d,e$2

together with inequalities of the form

$b,c,d,e$3

for some $b,c,d,e$4, and $b,c,d,e$5, the paper states that there exists a pairwise evaluation path cheaper than naive left-to-right. The proof compares

$b,c,d,e$6

with

$b,c,d,e$7

and concludes

$b,c,d,e$8

An analogous statement is given for reshaped Tucker. The computational interpretation is straightforward: if large spatial dimensions $b,c,d,e$9 are introduced too early, every subsequent intermediate pays that cost; if cheaper contractions among factor tensors are performed first, the expensive spatial convolution can be delayed (Rabbani et al., 2024).

5. Empirical behavior and measured gains

The experiments were run in PyTorch on an NVIDIA GeForce RTX 2080Ti and covered three settings: video classification with a two-stream ConvNet and ResNet-101 backbone on UCF-101, automatic speech recognition with Conformer on LibriSpeech, and image classification with ResNet-34 on CIFAR-10 and ImageNet. For reshaped CP experiments, the paper uses $$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$0, reshaping only input and output channel modes (Rabbani et al., 2024).

For ImageNet classification with an RCP ResNet-34 and checkpointing, Table 1 reports the following minutes per epoch. At compression ratio 5%, conv_einsum requires 21.8 train minutes and 1.06 test minutes, compared with 32.4 and 1.42 for naive evaluation with checkpointing. At 10%, the corresponding numbers are 23.4 and 1.47 versus 35.7 and 1.77. At 20%, they are 28.9 and 1.95 versus 39.9 and 2.35. At 50%, they are 34.1 and 2.54 versus 50.1 and 3.54. At 100%, they are 41.6 and 3.08 versus 65.6 and 4.08. The improvement therefore persists across the full compression range and grows in absolute terms at larger model sizes (Rabbani et al., 2024).

Table 2 reports FLOP reductions for CP convolutional layers in ResNet-34 at CR = 100%, batch size 128.

Layer group	FLOP speedup
conv1	$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$1
conv2_x	$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$2
conv3_x	$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$3
conv4_x	$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$4
conv5_x	$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$5

The largest gains occur in deeper layers, where feature and channel dimensions make naive left-to-right scheduling especially costly. Figures 3a–d show consistent runtime gains on CIFAR-10 and LibriSpeech, with and without checkpointing. Figure 1 shows a stronger memory effect on UCF-101: naive evaluation often fails due to memory overflow except at the highest compression levels, whereas conv_einsum remains usable across many settings. The paper’s interpretation is that optimal paths produce uniformly smaller intermediate products and therefore improve both runtime and memory efficiency (Rabbani et al., 2024).

The appendix also reports that conv_einsum outperforms naive implementations across multiple tensor decomposition families, including RCP, RTR, RTT, and RTK. Accuracy is not presented as the main contribution; the paper states that the central result is computational and memory efficiency rather than an intrinsic accuracy gain from the notation or path search itself.

6. Broader context, limitations, and significance

conv_einsum sits at the intersection of three strands of work: Einstein-style tensor programming, tensor-network path optimization, and convolution-as-structured-contraction. Earlier tensor-network tools such as ncon() separate symbolic network specification from explicit contraction sequence, making path control a first-class part of the interface rather than an implicit consequence of notation (Pfeifer et al., 2014). More recent work on tensor-network contraction planning argues that contraction paths should be represented as symbolic trees rather than as flat execution traces, emphasizing the partial-order structure of contraction. This suggests that the planning layer of conv_einsum belongs naturally to a symbolic intermediate-representation problem, not only to a kernel-dispatch problem (Sanchez-Ramirez et al., 2024).

A second neighboring line of work rewrites convolution itself as an einsum-like tensor network. One formulation factorizes convolution into ordinary dense tensor multiplication plus binary connectivity tensors $$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$6 that encode stride, padding, dilation, batching, channel groups, and arbitrary convolution dimension. In 2D, the core contraction is

$$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$7

which can be evaluated with einsum and then transformed for autodiff and second-order methods (Dangel, 2023). A closely related approach introduces dummy binary tensors $$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$8 and $$Y = \mathrm{conv\_einsum}("bshw,\,tshw \to bthw \mid hw", X, W).$$9 so that exact convolution becomes a contraction over selector tensors. In that formulation, standard convolution is written as

$b$0

making clear that plain einsum alone is insufficient unless one explicitly materializes structured selector tensors (Hayashi et al., 2019). conv_einsum can be read as an attempt to internalize this distinction: instead of requiring auxiliary selector tensors in the front-end notation, it marks convolution dimensions explicitly and optimizes the resulting multilinear program.

This also clarifies a common misconception. conv_einsum is not merely a wrapper around conv2d, nor is it simply “einsum with a convolution flag.” It is a symbolic representation plus a path-optimization layer for tensor programs that mix contractions and convolutions. At the same time, it is not a complete universal convolution compiler. The paper’s own implementation restricts multi-way non-commutative convolution and currently supports only commutative convolution with circular padding when a convolution index appears in more than two inputs (Rabbani et al., 2024).

Its broader significance lies in making tensorized convolutional layers legible to tensor-network methods. Once a layer is written as a conv_einsum expression, the same conceptual tools apply that have long applied to ordinary einsum expressions: path optimization, symbolic decomposition, cost modeling, and checkpointing. This suggests that the lasting contribution of conv_einsum is not only a library routine, but a representational layer that brings convolutional multilinear programs into the same optimization space as tensor contractions.