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

TL;DR

TreeTensor introduces a constrained, tree-like nested data container to address complex hierarchical data in AI systems, enabling tensor-like parallelism while preserving structural relationships. The core ideas are the treelize operation, which extends arbitrary Python functions to TreeTensor via a mapping from $f$ to $F$ and supports class-level extension, and a mismatch-robust policy framework plus a propagating constraint tree to enforce data properties. The authors present detailed definitions of tree nodes, value nodes, and the constraint machinery, plus practical extensions to PyTorch, NumPy, and Scikit-Learn, with empirical DRL benchmarks (including AlphaStar-like workloads) showing improved usability and competitive runtime performance. Overall, TreeTensor aims to simplify development for perception-decision AI systems, improve performance, and facilitate integration with existing libraries through non-intrusive treelization and constraint-based extensions.

Abstract

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.

Figures (11)

• Figure 1: Up: Comparison and transition between native PyTorch and TreeTensor-augmented PyTorch. Users can handle more complicated AI task on nested using TreeTensor, with the consistent usage experience and zero switch cost. Down: Transition in algorithm and data when progressing towards more complex and general AI system. This example is from related information of AlphaStar arulkumaran_alphastar_2019 and shows the concrete data structures and algorithms combinations.
• Figure 2: TreeTensor's Overview. The names appearing in this figure are the names appearing in the engineering design. TreeTensor is a Python-based application that using Cython for speedup. TreeStorage is the lowest layer, and it is used to encapsulate the tree structure. TreeValue class, tree calculating utilities, and the treelize tool can then be created. Current library functions, classes, and APIs can be enhanced and merged with the FastTreeValue class and additional integrated functions, which can totally inherit the original features.
• Figure 3: Tree Nodes $n^t$ and $n^t_x$'s Structure. Square nodes indicate value nodes, while circular nodes represent tree nodes.
• Figure 4: Operate Function $p$ on All the Nodes Of Tree $n_1$. The structure of $n^{\prime}_1$ is identical to that of $n_1$, but its values are handled by function $p$.
• Figure 5: Operate Function $h$ Among Trees $n_1$, $n_2$ and $n_3$. The structure of $n^{\prime}$ is identical to the original trees, but its values are handled by function $h$.