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

Abstract: Tensor is the most basic and essential data structure of nowadays 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.

• The paper introduces TreeTensor, a unified tree-like tensor structure that efficiently handles deeply nested, variable-length data in AI systems.
• It leverages treelize operations and seamless API extensions to generalize tensor computations across heterogeneous tree nodes.
• Benchmark results demonstrate significant gains in runtime efficiency and code simplicity, benefiting advanced reinforcement learning and multimodal applications.

TreeTensor: A Constrained Tree-Like Tensor for Nested Data in Machine Learning Systems

Motivation and Context

The increasing complexity of cognitive AI systems, particularly in domains such as reinforcement learning and multimodal perception-decision architectures, has introduced substantial challenges in data representation and manipulation. Traditional tensor constructs, as used by frameworks like PyTorch and TensorFlow, are inherently limited to regular, fixed-shaped arrays, hindering efficient and expressive modeling of deeply nested, heterogeneous, or variable-length data structures. This limitation is particularly pressing in advanced systems like AlphaStar, which require handling intricate, hierarchical data spanning spatial, scalar, vector, and variable-length entity observations.

TreeTensor addresses this gap by providing a unified, highly extensible container with a tree-nested structure that encapsulates both the logical relationships within complex data and the operational flexibility needed for scalable, efficient computation in modern AI pipelines. Unlike prior solutions (e.g., dm-tree, jax-libtree, torchbeast, RaggedTensor, nestedTensor), TreeTensor systematically generalizes nested data handling, enables seamless extension of existing APIs, and incorporates constraint mechanisms for robust data validation.

Figure 1: Comparison of transition between native PyTorch and TreeTensor-augmented PyTorch, illustrating seamless handling of nested data and its application in AlphaStar's complex data structuring.

Architecture of TreeTensor

TreeTensor is implemented as an extensible Python-based library, leveraging Cython for performance-critical components. The architecture is layered: TreeStorage encapsulates the base tree structure, TreeValue provides node-level value representation, and various utilities (including treelize and FastTreeValue) facilitate efficient operations and API extension.

Figure 2: Overview of TreeTensor's engineering design and main classes, showing integration points for extending library functions and inheriting features.

The central construct is a tree node, which can be either a value node $n^v$ (representing actual data values) or a tree node $n^t$ (representing nested subtrees). TreeTensor exposes operations for applying unary and multivariate functions across all nodes, supporting the composability and vectorization familiar from standard tensor operations but generalized for arbitrarily deep and heterogeneous trees.

Figure 3: Structure of tree nodes in TreeTensor, highlighting the distinction between value (square) and tree (circular) nodes.

Treelize and API Extension

The treelize operation is foundational, enabling users to lift any function, method, operator, or property onto TreeTensor objects. This generalization is achieved by recursively traversing the tree and applying the target operation to each value node, maintaining the structural integrity of the tree.

Unary operations (e.g., $\mathrm{sigmoid}(x)$, $\mathrm{split}(x, n)$) are automatically broadcast across nodes, while multivariate operations (e.g., $x + y$, $x \cdot y - z$) are jointly applied to aligned nodes from multiple trees.

Figure 4: Application of a unary function across all nodes of a TreeTensor, preserving structure and transforming values.

Figure 5: Application of a multivariate function across three TreeTensor instances, yielding a new tree with transformed values.

TreeTensor provides seamless API compatibility with PyTorch, Numpy, Scikit-Learn and other libraries. This is implemented via automatic treelize decorators, allowing existing library classes and functions to accept and operate on tree-nested arguments without intrusive modifications.

Figure 6: Extension of Numpy functions to TreeTensor, generating a tree of ndarrays and supporting attribute/method access.

Structural and Functional Utilities

TreeTensor introduces structural utilities such as subside and rise, which dynamically adjust the position of external data structures (lists, dicts) within the tree, supporting flattening and reconstruction as needed for integration into batch processing, multi-agent reinforcement learning, and variable-length data scenarios.

Figure 7: Illustration of subside and rise operations, transforming the storage layout between external containers and internal value nodes.

Functional utilities (map, filter, mask, reduce) are implemented natively, providing high-level recursive operations familiar from functional programming paradigms, but extended to the nested data context.

Constraint System for Robust Data Validation

To improve reliability and operational safety, TreeTensor incorporates a constraint system. Constraints can be inheritable (applied recursively to all leaf nodes, e.g., "all Tensors must be float32 on CUDA:0") or non-inheritable (imposed at specific nodes, e.g., "positive/negative margin for metric learning batch"). Constraints are constructed in parallel with the tree structure and validated automatically or manually during runtime operations, ensuring data integrity even in mutable contexts.

Figure 8: Construction of a constraint tree accompanying TreeTensor, demonstrating constraint inheritance and validation logic.

Figure 9: Example of constrained TreeTensor in multi-agent RL for long-horizon decision-making, with shape and type constraints across agents and time steps.

Figure 10: Sampling and optimization process for constrained TreeTensor, supporting efficient batch formation and loss computation.

Efficiency and Scalability

TreeTensor is engineered for runtime efficiency, employing Cython for core routines and explicit memory management where feasible. Comparative benchmarks with Tianshou Batch demonstrate TreeTensor's superior performance in key operations (stack, split, cat, init, set), especially as data size and complexity increase.

Figure 11: Speed performance comparison for stack and split operations between TreeTensor and Tianshou Batch, showing logarithmic dominance as element count grows.

Strong numerical results reported in the paper include:

• Up to an order of magnitude reduction in code complexity, extensibility, and maintainability metrics when implementing advanced RL algorithms, including AlphaStar collate functions, MuZero, WQMIX, and TREX.
• Near-native runtime overhead, with TreeTensor implementations showing equal or faster execution compared to native batch structures in real-world scenarios.

Policy-Based Handling of Structure Mismatches

TreeTensor provides four policies for resolving mismatches in structure during multivariate operations: Strict (default, requiring exact alignment), Inner (intersection), Outer (union), and Left (first-tree alignment). These policies enable flexible batch processing and variance reduction in scenarios with dynamic, multimodal data.

Limitations

TreeTensor does not maintain parent-pointer references, enforcing a loosely coupled structure. This precludes locality-maintenance operations akin to segment trees and may impact certain advanced tree-based algorithms. Risky constraints (e.g., those that may be invalidated by in-place modification in PyTorch) require manual validation and careful design to avoid performance degradation.

Practical and Theoretical Implications

Practically, TreeTensor enables concise, readable, and extensible implementations of complex, nested data handling in reinforcement learning, perception-decision pipelines, and multimodal AI systems. Theoretically, TreeTensor provides a framework for systematic modeling of hierarchical data relationships, constraint propagation, and asynchronous/variable-length computation.

TreeTensor's design opens avenues for scalable integration with future machine learning libraries, automatic differentiation frameworks, and distributed reinforcement learning platforms, particularly as tasks and architectures continue to diversify in complexity and data structure.

Conclusion

TreeTensor represents a significant advancement in the systematic handling of nested, hierarchical, and variable-length data structures within AI systems. Its design balances extensibility, efficiency, and robustness, facilitating scalable implementations in research and industry. The constraint system offers rigorous guarantees for data integrity, while treelize operations enable seamless extension of existing library functions. The architecture is well positioned to serve as a foundational component for future developments in multimodal machine learning, complex agent-based environments, and scalable RL systems.