# TreeTensor: Boost AI System on Nested Data with Constrained Tree-Like Tensor

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 35 references that are in the map.
- First posted on arXiv 9 Feb 2026.

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Agent57: Outperforming the Atari Human Benchmark](https://mlsanity.com/paper/2003.13350/) (2020)
- [Weighted QMIX: Expanding Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2006.10800/) (2020)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Tianshou: a Highly Modularized Deep Reinforcement Learning Library](https://mlsanity.com/paper/2107.14171/) (2021)
- [AlphaStar: An Evolutionary Computation Perspective](https://mlsanity.com/paper/1902.01724/) (2019)
- [TorchBeast: A PyTorch Platform for Distributed RL](https://mlsanity.com/paper/1910.03552/) (2019)

## Abstract (arXiv)

Tensor is the most basic and essential data structure of nowadays artificial intelligence (AI) system. The natural properties of Tensor, especially the memory-continuity and slice-independence, make it feasible for training system to leverage parallel computing unit like GPU to process data simultaneously in batch, spatial or temporal dimensions. However, if we look beyond perception tasks, the data in a complicated cognitive AI system usually has hierarchical structures (i.e. nested data) with various modalities. They are inconvenient and inefficient to program directly with conventional Tensor with fixed shape. To address this issue, we summarize two main computational patterns of nested data, and then propose a general nested data container: TreeTensor. Through various constraints and magic utilities of TreeTensor, one can apply arbitrary functions and operations to nested data with almost zero cost, including some famous machine learning libraries, such as Scikit-Learn, Numpy and PyTorch. Our approach utilizes a constrained tree-structure perspective to systematically model data relationships, and it can also easily be combined with other methods to extend more usages, such as asynchronous execution and variable-length data computation. Detailed examples and benchmarks show TreeTensor not only provides powerful usability in various problems, especially one of the most complicated AI systems at present: AlphaStar for StarCraftII, but also exhibits excellent runtime efficiency without any overhead. Our project is available at https://github.com/opendilab/DI-treetensor.