Research map / Architectures and efficiency / CNN design
Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Views convolutions as tensor networks to derive einsum expressions for autodiff, KFAC-type curvature approximations (including for transpose convolution), and structure-based simplifications, yielding up to 4.5x faster KFAC-reduce with lower memory.
Problem
Convolutions are hard to analyze and implement for less standard routines (higher-order autodiff, Kronecker-factored curvature, transpose convolution) across all hyper-parameters, batching, groups and dimensions, and existing implementations are inefficient or missing.
Approach
Represent convolution as a tensor network with index pattern tensors encoding connectivity, derive Jacobians, VJPs and KFAC factors by diagrammatic manipulation, express them as einsum strings, and apply convolution-specific index-pattern simplifications (dense, down-sampling, kernel-output swap) before contraction.
Key results
- TN implementation of KFAC-reduce achieves speed-ups up to ~4.5x (performance ratio down to 0.22x) with memory savings up to 3 GiB across convolutions of 9 CNNs.
- On VGG19 with SINGD, the TN implementation reduces peak memory from 32 GiB to 16 GiB and halves per-iteration time relative to standard SINGD (0.80 s vs 1.0 s).
- Random CRS-based spatial sub-sampling reaches the same gradient approximation quality as channel sub-sampling using only 35% of pixels instead of keeping 75% of channels.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the TN implementation instead of input unfolding to compute SINGD's KFAC-reduce factors | SINGD on ResNet18 · ImageNet-like synthetic data (3, 256, 256), batch size 128 | peak memory (GiB): 4.5 -> 3.6 | Helped |
| Use the TN implementation instead of input unfolding to compute SINGD's KFAC-reduce factors | SINGD on VGG19 · ImageNet-like synthetic data (3, 256, 256), batch size 128 | peak memory (GiB): 32 -> 16 | Helped |
| Use the TN implementation instead of input unfolding to compute SINGD's KFAC-reduce factors | SINGD on ResNet18 · ImageNet-like synthetic data (3, 256, 256), batch size 128 | per-iteration time (s): 0.19 -> 0.16 | Helped |
| Use the TN implementation instead of input unfolding to compute SINGD's KFAC-reduce factors | SINGD on VGG19 · ImageNet-like synthetic data (3, 256, 256), batch size 128 | per-iteration time (s): 1.0 -> 0.80 | Helped |
| Avoid explicitly unfolding the input when computing KFAC-reduce | KFAC-reduce convolution implementation · Convolutions across 9 CNNs | runtime ratio (TN vs. standard implementation): 0.22\,x | Helped |
| Subsample spatial pixels rather than input channels for randomized weight-VJP computation | TN randomized autodiff · Real and synthetic data | gradient approximation error: 75\,\% of channels -> 35\,\% of pixels | Helped |
Surprising findings
- The TN implementation of weight VJP in the median outperforms PyTorch's highly optimized weight VJP, including for down-sampling convolutions.
Limitations
- No common sub-expression elimination; opt_einsum only optimizes contraction order
- No index slicing implemented; single-GPU benchmark without OOM encountered
- Index pattern tensors contain explicit zeros for general convolutions, leaving unsimplified cases costly
Experimental setup
Datasets: ImageNet-like synthetic data (3x256x256)
Benchmarks: convolutions of 9 CNN architectures (ImageNet and smaller datasets)
Metrics: run time ratio peak memory per-iteration time
Baselines: standard PyTorch convolution implementation standard KFAC-reduce implementation SINGD SGD
Builds on: KFAC KFC KFAC-reduce opt_einsum column-row sampling
Closest papers by meaning
- Exploring Unexplored Tensor Network Decompositions for Convolutional Neural Networks (NeurIPS 2019 · similarity 0.90)
- T-Net: Parametrizing Fully Convolutional Nets With a Single High-Order Tensor (CVPR 2019 · similarity 0.89)
- Deep Tensor Convolution on Multicores (ICML 2017 · similarity 0.88)
- Distribution-Aware Tensor Decomposition for Compression of Convolutional Neural Networks (NeurIPS 2025 · similarity 0.88)
- Coarse-To-Fine Tensor Trains for Compact Visual Representations (ICML 2024 · similarity 0.87)
- Tensorion: A Tensor-Aware Generalization of the Muon Optimizer (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.