# Nimble GNN Embedding with Tensor-Train Decomposition

arXiv: https://arxiv.org/abs/2206.10581 · HTML: https://mlsanity.com/paper/2206.10581/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 30 references that are in the map.
- First posted on arXiv 21 Jun 2022.

## Builds on (most cited first, 20 of 20)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://mlsanity.com/paper/1710.01878/) (2017)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [On the Expressive Power of Deep Learning: A Tensor Analysis](https://mlsanity.com/paper/1509.05009/) (2015)
- [Provable Benefit of Orthogonal Initialization in Optimizing Deep Linear Networks](https://mlsanity.com/paper/2001.05992/) (2020)
- [Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems](https://mlsanity.com/paper/1909.02107/) (2019)
- [Distributed Hierarchical GPU Parameter Server for Massive Scale Deep Learning Ads Systems](https://mlsanity.com/paper/2003.05622/) (2020)
- [TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models](https://mlsanity.com/paper/2101.11714/) (2021)
- [Ultimate tensorization: compressing convolutional and FC layers alike](https://mlsanity.com/paper/1611.03214/) (2016)
- [Hash Embeddings for Efficient Word Representations](https://mlsanity.com/paper/1709.03933/) (2017)
- [Wide Compression: Tensor Ring Nets](https://mlsanity.com/paper/1802.09052/) (2018)
- [Tensorized Embedding Layers for Efficient Model Compression](https://mlsanity.com/paper/1901.10787/) (2019)
- [Getting deep recommenders fit: Bloom embeddings for sparse binary input/output networks](https://mlsanity.com/paper/1706.03993/) (2017)
- [Training with Multi-Layer Embeddings for Model Reduction](https://mlsanity.com/paper/2006.05623/) (2020)
- [Position-based Hash Embeddings For Scaling Graph Neural Networks](https://mlsanity.com/paper/2109.00101/) (2021)

## Built on by (most cited first, 6 of 6)

- [Tensor Networks Meet Neural Networks: A Survey and Future Perspectives](https://mlsanity.com/paper/2302.09019/) (2023)
- [Characterizing the Accuracy -- Efficiency Trade-off of Low-rank Decomposition in Language Models](https://mlsanity.com/paper/2405.06626/) (2024)
- [A Thorough Performance Benchmarking on Lightweight Embedding-based Recommender Systems](https://mlsanity.com/paper/2406.17335/) (2024)
- [GraphScale: A Framework to Enable Machine Learning over Billion-node Graphs](https://mlsanity.com/paper/2407.15452/) (2024)
- [TensorCodec: Compact Lossy Compression of Tensors without Strong Data Assumptions](https://mlsanity.com/paper/2309.10310/) (2023)
- [Coarse-to-Fine Lightweight Meta-Embedding for ID-Based Recommendation](https://mlsanity.com/paper/2501.11870/) (2025)

## Abstract (arXiv)

This paper describes a new method for representing embedding tables of graph neural networks (GNNs) more compactly via tensor-train (TT) decomposition. We consider the scenario where (a) the graph data that lack node features, thereby requiring the learning of embeddings during training; and (b) we wish to exploit GPU platforms, where smaller tables are needed to reduce host-to-GPU communication even for large-memory GPUs. The use of TT enables a compact parameterization of the embedding, rendering it small enough to fit entirely on modern GPUs even for massive graphs. When combined with judicious schemes for initialization and hierarchical graph partitioning, this approach can reduce the size of node embedding vectors by 1,659 times to 81,362 times on large publicly available benchmark datasets, achieving comparable or better accuracy and significant speedups on multi-GPU systems. In some cases, our model without explicit node features on input can even match the accuracy of models that use node features.