PaperlayerView exact revision on arXiv

Source-linked AI summary

EinSort: Sorting is All We Need for Tensorizing LLM

Toshiaki Koike-Akino, Jing Liu, Ye Wang

arXiv:2606.08565v1cs.LGcs.AI

TL;DR

Large foundation models are costly to deploy, and tensor-network methods have paid relatively little attention to index ordering as a way to reveal low-rank structure. EinSort adds reversible sorting permutations to tensorization, and the paper reports improved compression quality for LLM weights and KV caches.

  • Problem

    Large foundation models are expensive to deploy, and existing tensor-network approaches have given relatively little attention to index ordering as a design choice.

    Source: §1 · p. 1

  • Method

    EinSort tensorizes LLMs with adaptive index ordering, augmenting Einstein summation with reversible permutation operators.

    Source: §1 · p. 1

  • Results

    EinSort improves compression quality for LLM weight and KV-cache tasks; sorting reduces reconstruction error, and full sorting achieves nearly un-compressed KV-cache performance.

    Source: §1 · p. 1, §2.3 · p. 3, §3 · p. 4

  • Takeaways & Limitations

    Index ordering is a useful way to reduce tensor rank and can offer compression advantages on some models.

    Source: §4 · p. 4

    Permutation storage overhead is not negligible; sorting L values requires at most ⌈log2(L!)⌉ bits using factoradic/Lehmer coding.

    Source: §2.5 · p. 3

Abstract

from arXiv · show

Tensor networks provide efficient representations for compressing large neural networks. By carefully designing shapes and topologies, they can significantly reduce memory and computational costs. However, identifying implicit low-rank structures in large foundation models remains challenging due to their enormous scale and un-structured weight distributions. We propose an adaptive tensorization method that discovers inherent low-rank structure in a target tensor by index ordering. Experiments on weight and KV-cache compression demonstrate improved reconstruction quality compared to baselines.

1. Introduction

EinSort treats tensor-index ordering as an overlooked design choice that can expose latent low-rank structure in LLM weights and KV caches. The framework uses reversible permutations and is presented as improving compression quality while reducing effective tensor rank.

1. Introduction

  • Large foundation models perform well across tasks, but their parameter counts and memory requirements make deployment expensive.The introduction lists partial activation, pruning, quantization, distillation, and rank reduction among proposed compression techniques.Source: §1 · p. 1
  • Existing tensor-network approaches emphasize topology, contraction, and tensor ranks, while giving relatively little attention to index ordering.The paper argues that ordering can substantially alter a tensor’s effective rank.Source: §1 · p. 1
  • EinSort augments tensor decomposition with adaptive index ordering to compress LLM weights and KV caches.The framework uses reversible permutation operators.Source: §1 · p. 1
  • Sorting operations can expose latent low-rank structure hidden in pretrained LLM weights and KV caches.The stated contributions include theoretical and empirical evidence for rank reduction, practical low-overhead permutation schemes, and improved compression quality.Source: §1 · p. 1

2. Tensorizing LLM for compression

The section develops tensor-network compression through Einstein summation, gauge fixing, and index ordering. EinSort inserts reversible permutations into tensor decomposition to expose low-rank structure, with sorting supported by both empirical results and a restricted theoretical result.

2.1. Einstein summation as general tensor networks

  • Einstein summation represents tensor networks through shaped tensor cores, and reducing their ranks can reduce the parameters needed to represent a tensor.For a 4-mode tensor of shape (d, d, d, d), the stated parameter count falls from d^4 to 2d(R + 1)R at maximum rank R.Source: §2.1 · pp. 1–2

2.2. Gauge fixing

  • Gauge freedom gives multiple equivalent tensor factorizations, and fixing it can improve numerical stability and reduce parameters for a chosen junction.For a 4-mode TT, the passage gives a parameter count of 2d(R + 1)R − 2R^2 after the stated block-identity choice.Source: §2.2 · p. 2

2.3. Index ordering

  • Tensor sorting dramatically reduces reconstruction error, enabling near rank-1 reconstruction; although row sorting has higher error than full sorting, its error remains more than 10 times lower than without sorting.Source: §2.3 · p. 3

2.4. Theoretical analysis

  • For sorted i.i.d. uniform values reshaped into a matrix, the matrix is asymptotically approximated in probability by one of rank at most 3.The passage cautions that this result does not guarantee low-rank factorization for arbitrary random matrices.Source: §2.4 · p. 3

2.5. EinSort: Index ordered tensor networks

  • EinSort applies a reversible permutation before tensor decomposition and its inverse afterward, extending einsum-based tensorization to search for low-rank structure.Non-negative tensorization projects values before decomposition and stores the original sign in one additional bit.Source: §2.5 · p. 3, Figure 3 · p. 3
  • Sorting’s permutation storage can be substantial: full sorting a (1024, 1024) tensor requires roughly 18.6 bits per entry.The section also proposes non-negative tensorization, retaining one sign bit to recover negative values.Source: §2.5 · p. 3

3. Experiments

The experiments evaluate EinSort for KV-cache compression and compare reductions, tensor-network topologies, and decoding runtime. EinSort achieves better perplexity than regular SVD without sorting, while compression trades off against decoding throughput.

3. Experiments

  • For Qwen3 models with eight heads, reducing KV cache is evaluated on 64-token WT2 segments followed by decoding up to 64 new tokens; shared-permutation storage is about 1.4 bits per weight.The passage also points to evaluations on other models and GSM8K, TextVQA, and LIBERO tasks.Source: §3 · p. 4
  • EinSort achieves better PPL than regular SVD without sorting for compressed KV caches, while full sorting approaches un-compressed performance.The reported comparisons cover Qwen3-0.6B and Gemma3-4B-it; Gemma3 has higher tolerance but worse base PPL than Qwen3.Source: §3 · p. 4
  • Product reduction performs best among the tested reductions; gauge fixing and non-negative tensorization improve accuracy, with the latter costing one additional bit.The best exponent is near 1 for product reduction and around 0.5 for mean reduction.Source: §3 · p. 4
  • With square-root mapping, product reduction, non-negative factorization, and gauge fixing, CP is slightly better than other tested topologies, while TR is worst.TT/TR have around 40% decoding speed down on both tested devices, TD about 70% loss, and CP is much slower due to iterative Khatri–Rao products.Source: §3 · p. 4, Table 2 · p. 4, Table 3 · p. 4

4. Conclusion

EinSort uses index ordering in tensor decomposition to discover implicit low-rank structure for LLM compression, and the paper reports sorting as a way to reduce tensor rank.

4. Conclusion

  • EinSort applies index ordering in tensor decomposition to discover implicit low-rank structure for LLM compression.The conclusion also reports that sorting reduces tensor rank and that a simple sorting mechanism showed potential advantages on some models.Source: §4 · p. 4

Impact Statement

The impact statement declines to highlight a specific societal consequence; related-work passages situate LLM compression across low-rank, KV-cache, adaptation, sorting, and tensor-network approaches.

Impact Statement

  • The authors state that they do not identify a specific societal consequence requiring emphasis.The statement describes potential societal consequences as numerous.Source: Impact Statement · p. 5

A.1. Model compression

  • LLM compression methods include quantization, pruning, distillation, and rank reduction.The passage also discusses partial activation methods.Source: §A.1 · p. 10

A.2. Low-rank compression

  • Low-rank compression builds on weight redundancy and uses matrix or tensor factorization, with possible instability or accuracy degradation.The passage traces approaches from SVD and convolutional-filter approximations to Tucker and CP decompositions.Source: §A.2 · p. 10

A.3. KV cache reduction

  • KV-cache reduction targets memory that grows with sequence length, batch size, layers, and KV hidden dimension, using eviction, quantization, or representation compression.Eviction-based methods may discard information needed for long-range reasoning.Source: §A.3 · p. 10

A.4. Low-rank adaptation

  • Low-rank adaptation includes LoRA’s addition of a product of two low-rank matrices and variants using tensor decompositions and other mappings.The passage names CP, Tucker, nonlinear, Hadamard, Kronecker, and quantum tensor-network variants.Source: §A.4 · p. 10

A.5. Activation-aware compression

  • Activation-aware compression work applies activation statistics, error-propagation factors, gradient information, or activation-dependent scaling across decomposition, quantization, and pruning.The cited examples include ASVD, AA-SVD, OBD-LLM, GPTQ, AWQ, SparseGPT, SparseLLM, and Wanda.Source: §A.5 · p. 11

A.6. Permutation/sorting networks

  • Permutation and sorting networks include permutation-invariant set models and differentiable sorting relaxations such as NeuralSort, SoftSort, and Gumbel–Sinkhorn.These relaxations support gradient-based optimization of orderings.Source: §A.6 · p. 11

A.7. Tensor network

  • Tensor networks represent multidimensional arrays through networks of lower-dimensional tensors, including MPS, TTN, MERA, and PEPS.Tensorization can provide an efficient parameterization of DNN architectures.Source: §A.7 · p. 11

B. Einstein summation

Einstein summation represents tensor networks through equations and core shapes, while the sorting experiments report lower tensorization error and rank requirements.

B. Einstein summation

  • Einstein summation expresses tensor networks by specifying an equation and core shapes, and the described module supports fitting cores with ALS or gradient optimization.The implementation uses optimized contraction expressions and can initialize cores with tensorly.Source: B. Einstein summation · pp. 11–13
  • 3.4 × 10^10 to 4.4 × 10^7 FLOPs reduces the stated cost of contracting a four-mode TT with site dimension 32 and rank 20.The passage attributes the nearly three-orders-of-magnitude reduction to optimized contraction order.Source: B. Einstein summation · p. 13
  • 2.7 × 10^12 to 6.9 × 10^6 FLOPs reduces the stated cost of reconstructing a weight and multiplying it by a batched input through optimized contraction order.The passage reports more than five orders of magnitude speedup, even relative to weight materialization alone.Source: B. Einstein summation · p. 13

C. Rank reduction with tensor sorting

  • Group sorting ignores element order within groups, while sequential sorting proceeds round-robin across axes; both reduce the memory needed to represent permutations.The passage gives the group-size-dependent memory expression and states that sequential sorting requires ⌈log2(d!)⌉m/d bits for m-mode tensors with site dimension d.Source: C. Rank reduction with tensor sorting · p. 14
  • Sorting yields substantially lower reconstruction error than unsorted SVD, with full sorting strongest and row-wise, group, and sequential sorting still over 10 times lower in error.The reported trend is similar across the LLM models considered; the figure plots tensorization error against rank for Qwen3-0.6B v-projection weights.Source: C. Rank reduction with tensor sorting · p. 14, Figure 6 · p. 14

D.1. Proof of Lemma 2.1

The appendices analyze low-rank structure after sorting under several distributional settings, establish invariance under independent permutations, and compare reduction strategies.

D.1. Proof of Lemma 2.1

  • Sorted i.i.d. uniform samples reshaped into a matrix are asymptotically approximated by a matrix of rank at most three, with entrywise stochastic error O_p(L^-1/2).The proof uses the order-statistic distribution and bounds on its mean and variance.Source: §D.1 · p. 15

D.2. Random projection to make uniformity

  • After a sufficiently mixing unitary transform and marginal Gaussianization, sorted and folded matrices are expected to exhibit asymptotic low-rank structure beyond the uniform model.The argument is heuristic: transformed variables are generally not exactly independent, and finite-dimensional convergence does not establish joint uniformity.Source: §D.2 · p. 16

D.3. Row sorting

  • For N ≫ 1, a row-wise sorted matrix is asymptotically rank one under the stated assumptions on i.i.d. entries and their distribution.The rank-one form is 1_M q^T, where q is defined from the distribution’s quantile function.Source: §D.3 · pp. 16–17

D.4. Independent permutation

  • Independent row and column permutations preserve a matrix’s singular values and rank.The appendix contrasts this invariance with row-wise sorting, which can produce rank-one structure.Source: §D.4 · pp. 18–19

D.5. Reduction impact

  • Product reduction has a constant target-dependent margin, while mean reduction’s margin increases with target magnitude; under uniform variables, the mean margin decays more slowly as reduction size grows.Sliced sorting with reduction can save memory but may degrade sorting accuracy.Source: §D.5 · pp. 19–20

E. EinSort

EinSort combines reversible index permutations with tensor decomposition, and extends this approach with non-negative tensorization, gauge fixing, activation-aware objectives, and adaptive compression options. The section describes permutation-storage trade-offs and reports sorting and gauge-fixing results across Qwen3 and Phi3-mini.

E. EinSort

  • Shared permutations reorder tensor entries before decomposition; slicing reduces permutation storage, reaching about 0.4 bits per weight at d=32.Non-negative tensorization decomposes non-negative values and stores one additional sign bit for reconstruction.Source: E. EinSort · pp. 21–22

F. Nonlinear mapping and reduction

G. Gauge fixing

  • Repeated LU gauge fixing makes TT cores block identity and reduces parameters from 2dr(1 + r) to 2dr(1 + r) − 3r^2.SVD canonicalization makes all cores except the orthogonality center orthogonal and is described as improving numerical stability.Source: G. Gauge fixing · pp. 23–24

H. Activation-aware tensorization

  • Sorted tensor decomposition can incorporate activation and gradient statistics through preconditioners, extending activation-aware tensorization objectives to EinSort.Source: H. Activation-aware tensorization · p. 27

I. Quantization-aware tensorization

  • Quantization-aware tensorization quantizes or prunes tensor cores while using a straight-through estimator to pass gradients through einsum contraction.Source: §I · p. 27

J. Test-time adaptation

  • Test-time adaptation can select tensor ranks from online activations or update tensor cores using an online preconditioner.A small scoring network may be trained on synthetic or held-out calibration data; an alternative updates cores locally using an MPO and DMRG.Source: J. Test-time adaptation · p. 28

K. Activation compression

  • EinSort can reduce the memory needed for intermediate activations retained across layers for gradient updates, and reverse permutations let activations contract with decomposed weights.Source: K. Activation compression · pp. 28–29

L. Connection with quantum gates

The appendices connect reversible index permutations to logic gates and explore alternative ways to represent, learn, and order them. They also describe how folding and axis choices affect tensor decomposition and permutation storage.

L. Connection with quantum gates

  • Wπ has singular values [11.5, 0.2, 0.0, 0.0], making it nearly rank 2, compared with [10.6, 4.1, 1.6, 0.6] for unsorted W0.The row-dependent column permutation is reversible and can be implemented by a CNOT chain.Source: L. Connection with quantum gates · p. 29
  • Allowing approximately 50% permutation cancellation through don’t-care optimization reduces the gate count from 1976 to 744.Source: L. Connection with quantum gates · p. 30

M. Algebraic polynomial permutation

  • Quadratic polynomial permutations use π(i) = a + bi + ci^2 mod n with stated bijectivity conditions, but do not cover all permutations unless n is prime.For arbitrary n, the passage describes using polynomial permutations over a finite field instead.Source: M. Algebraic polynomial permutation · p. 30

N. Learned permutation operations

  • Permutation and tensorization choices can be learned by gradient updates or, for non-differentiable hyperparameters, reinforcement learning.The proposed choices include nonlinear mappings, tensor ranks, shapes, and topologies.Source: N. Learned permutation operations · p. 30

O. Folding and ordering

  • Sequentially sorting pairs of axes reduces permutation storage from 18.6 to 17.1 bits for d=32; sorting each of four axes sequentially reduces it to 14.7 bits.Axis swapping, folding, sorting, reduction, and decomposition choices can be adjusted independently to minimize reconstruction error.Source: O. Folding and ordering · p. 31

P. Benefits and limitations

EinSort offers flexible tensor-network design and broad compression applicability, with potential rank-reduction benefits from sorting. Its practical use is limited by permutation memory, unavailable optimized CUDA kernels, and the difficulty of choosing designs and computing large factorizations.

Benefits

  • Einsum-based formulation supports designing almost arbitrary tensor-network topologies.Sorting also has a theoretical basis for potentially reducing tensor ranks.Source: P. Benefits and limitations · p. 31
  • EinSort can adjust slicing order, mode selection, folding combination, nonlinear mapping, and reduction operations.These options provide flexibility in configuring compression.Source: P. Benefits and limitations · p. 31
  • EinSort can be applied to weight, KV-cache, and forward-activation compression.The listed activation use is compression for backward-gradient calculation.Source: P. Benefits and limitations · p. 31

Limitations

  • Choosing the best tensor topology, folding shapes, and index ordering is not straightforward, while large tensor factorization is computationally challenging.The paper also notes a connection between sorting and entanglement operations in quantum tensor networks.Source: P. Benefits and limitations · pp. 31–32
  • Permutation memory is difficult to reduce, and compatible high-performance CUDA kernels for contraction and index reordering are unavailable.These limitations affect memory overhead and kernel optimization.Source: P. Benefits and limitations · p. 32

Q. Experiments setup

Experiments evaluate EinSort on LLM tasks using the SparseLLM setting and code base, decomposing transformer linear layers. The setup covers several model sizes and measures WikiText-2 perplexity, with KV-cache sorting memory adjustable through slicing.

Q. Experiments setup

  • Experiments use the SparseLLM setting and code base, decompose transformer linear layers, and run on NVIDIA A40 or A100 GPUs.EinSort is implemented in PyTorch, with HuggingFace Transformers used for models and datasets.Source: Q. Experiments setup · p. 32
  • LLM experiments cover Qwen3, Gemma3, and Phi3 models at different sizes, measuring perplexity on raw-WikiText2.The setup refers to Appendices S and T for model and dataset details.Source: Q. Experiments setup · p. 32
  • KV-cache sorting memory can be controlled by adjusting slicing, with reported per-parameter overheads of 1.44 bits for Qwen3-0.6B/1.7B, 3.14 bits for Gemma3-4B, and 0.35 bits for Phi3-mini.The setup uses the head axis for reduction and 64-stacked head dimensions for sorting.Source: Q. Experiments setup · p. 32

R.1. LLM model variants

Across LLM, VLM, and robot-manipulation evaluations, EinSort improves compression performance over the stated baselines. The evaluated model families include language, vision-language, and vision-language-action systems.

R.1. LLM model variants

  • For Phi3-mini KV-cache compression, EinSort significantly improves tensorization accuracy toward full-sort SVD while using 0.35 bits of permutation memory.The full-sort case requires approximately 11.1 bits for permutation memory.Source: §R.1 · p. 32

R.2. LLM math reasoning

  • EinSort significantly improves Phi-4-mini GSM8K accuracy over regular SVD compression.The figure plots average accuracy with a Wilson 99% confidence interval.Source: §R.2 · p. 33, Figure 15 · p. 33

R.3. VLM visual reasoning

  • At 20% compression, EinSort keeps high accuracy on Qwen3-VL TextVQA evaluations, while the other weight-compression baselines degrade steeply.The comparison covers Qwen3-VL-4B and 8B against plain SVD, ASVD, SVD-LLM, and LatentLLM.Source: §R.3 · p. 33, Figure 16 · p. 33

R.4. VLA robot manipulation

  • EinSort significantly outperforms state-of-the-art weight-decomposition methods across four LIBERO benchmarks for X-VLA, π0.5, and VLA-JEPA.Success rates are evaluated over 800 rollouts; occasional results above uncompressed models are described as potentially due to regularization, with gaps within Wilson confidence intervals.Source: §R.4 · p. 34, Table 5 · p. 34, Table 6 · p. 34

S. Large foundation models

  • The evaluated VLA models use Florence-2, PaliGemma, and Qwen3-VL backbones, respectively, and the appendix describes their architectures and model families.The listed model families include X-VLA, π0.5, and VLA-JEPA.Source: S. Large foundation models · pp. 35–36

T. Datasets

The dataset appendix describes language modeling, mathematical reasoning, visual question answering, and simulated manipulation benchmarks. The following appendix section lists libraries used in the work.

T. Datasets

  • The datasets cover WikiText-2 language modeling, GSM8K mathematical reasoning, TextVQA image-text question answering, and LIBERO simulated long-horizon manipulation.The passages provide dataset sizes, task descriptions, and licensing information.Source: T. Datasets · pp. 36–37

U. Libraries

  • The libraries appendix lists dependencies including RevKit, Galois, LeRobot, PyTorch, opt_einsum, einops, TensorLy, Transformers, datasets, SciPy, and PyEDA.The entries include version and license information for several libraries.Source: U. Libraries · p. 37

Related papers

Loading 2606.08565v1…