1. 문제 제기 (Problem)

최근 인공지능 (AI) 시스템, 특히 지각 (Perception) 을 넘어 의사결정 (Decision) 및 인지 (Cognitive) 영역으로 확장되면서 데이터의 복잡도가 급격히 증가하고 있습니다.

• 기존 텐서의 한계: 기존 딥러닝 프레임워크 (PyTorch, TensorFlow 등) 의 핵심인 '텐서 (Tensor)'는 고정된 모양 (Fixed Shape) 의 정규화된 데이터 (이미지, 비디오 프레임 등) 에 최적화되어 있습니다. GPU 와 같은 병렬 처리 장치를 효율적으로 활용할 수 있게 해줍니다.
• 중첩 데이터 (Nested Data) 의 부재: AlphaStar(스타크래프트 II) 와 같은 복잡한 AI 시스템이나 강화학습 (DRL) 환경에서는 계층적 구조 (Hierarchical Structure) 를 가진 중첩 데이터 (딕셔너리, 리스트, 다양한 모달리티가 섞인 데이터) 가 필수적입니다.
• 현재의 비효율성: 이러한 중첩 데이터를 처리하기 위해 기존에 파이썬의 list 와 dict 를 직접 사용하거나, dm-tree, jaxlibtree 같은 라이브러리를 활용하는 경우, 데이터 구조를 매번 수동으로 정리해야 하거나 연산 시 오버헤드가 발생하며 코드가 복잡해지고 유지보수가 어렵습니다. 또한, RaggedTensor 나 nestedTensor 는 가변 길이 데이터에 초점을 맞추고 있지만, 텐서 연산의 재구현으로 인해 확장성과 성능에 부담을 줍니다.

2. 방법론 (Methodology)

저자들은 중첩 데이터를 효율적으로 처리하기 위해 TreeTensor라는 새로운 컨테이너를 제안했습니다. 이는 제한된 트리 구조 (Constrained Tree-Like Tensor) 를 기반으로 하며, 다음과 같은 핵심 메커니즘을 가집니다.

2.1. 트리 중첩 구조 정의 (Tree-Nested Structure)

• 데이터는 **값 노드 (Value Node)**와 **트리 노드 (Tree Node)**로 구성됩니다.
• 트리 노드는 하위 트리를 나타내는 키 - 값 쌍의 집합으로 정의되며, 전체 데이터는 하나의 루트 트리 노드로 표현됩니다.
• 이 구조를 통해 데이터 간의 논리적 연결 관계를 트리 속성으로 표현할 수 있습니다.

2.2. Treelize (함수 확장)

• 기존 라이브러리 (PyTorch, NumPy, Scikit-Learn 등) 의 함수를 중첩 데이터 구조에 자동으로 적용할 수 있도록 하는 핵심 기능입니다.
• 단항 함수 (Unary Function): 트리 내의 모든 값 노드에 동일한 함수를 재귀적으로 적용합니다.
• 다항 함수 (Multivariate Function): 두 개 이상의 트리 간 연산 시, 구조가 일치하는 노드끼리 연산을 수행합니다.
• 비침습적 확장 (Non-intrusive Extension): 기존 API 를 수정하지 않고, treelize 도구를 통해 클래스, 함수, 연산자, 속성 등을 트리 구조에 대응되는 형태로 변환할 수 있습니다.

2.3. 불일치 정책 (Mismatch Policies)

서로 다른 구조를 가진 트리 간 연산 시 발생하는 키 (Key) 불일치 문제를 해결하기 위해 4 가지 정책을 제공합니다:

1. Strict (기본값): 모든 키가 1:1 로 일치해야 함.
2. Inner: 교집합 (Intersection) 키만 연산.
3. Outer: 합집합 (Union) 키를 사용하며, 존재하지 않는 값은 기본값으로 채움.
4. Left: 첫 번째 트리의 키를 기준으로 함.

2.4. 제약 조건 시스템 (Constraint System)

데이터의 타입, 모양 (Shape), 저장 위치 (CUDA) 등에 대한 제약 조건을 트리 구조에 부여하여 성능 최적화와 유효성 검사를 가능하게 합니다.

• 상속 제약 (Inheritance Constraint): 하위 노드 전체에 적용되는 조건 (예: "모든 노드가 float32 타입").
• 비상속 제약 (Non-inheritance Constraint): 특정 노드 간의 관계나 조건 (예: "anchor 노드와 positive 노드의 shape 가 일치해야 함").
• 제약 트리 (Constraint Tree): 데이터 트리와 1:1 로 매핑되어, 데이터 수정 시 자동으로 제약 조건을 검증하거나 연산 전/후에 유효성을 검사합니다.

2.5. 성능 최적화

• Python 기반이지만, 핵심 컴포넌트는 Cython으로 작성되어 정적 라이브러리로 컴파일됩니다.
• 동적 타입 결정 비용 감소, 인라인 (Inlining) 처리, 불필요한 캡슐화 계층 제거 등을 통해 오버헤드를 최소화합니다.

3. 주요 기여 (Key Contributions)

1. 범용 중첩 데이터 컨테이너 제안: 고정된 모양의 텐서와 파이썬의 비정형 리스트/딕셔너리 사이의 간극을 메우는 TreeTensor 를 개발했습니다.
2. 제로 비용 (Zero Cost) 연산: Scikit-Learn, NumPy, PyTorch 등 기존 머신러닝 라이브러리의 함수를 중첩 데이터에 거의 오버헤드 없이 적용할 수 있는 treelize 메커니즘을 구현했습니다.
3. 유연한 불일치 처리 및 제약 시스템: 다양한 데이터 구조 간의 연산을 위한 정책과, 데이터 무결성을 보장하고 특수 연산을 최적화할 수 있는 제약 조건 시스템을 설계했습니다.
4. 실제 복잡한 AI 시스템 적용: AlphaStar(스타크래프트 II) 와 같은 초고난도 강화학습 프로젝트의 데이터 처리 파이프라인을 TreeTensor 로 최적화하여 코드의 복잡성을 획기적으로 줄이고 가독성을 높였습니다.

4. 실험 결과 (Results)

• 코드 품질 향상: AlphaStar 의 데이터 정합 (collate) 함수, MuZero, WQMIX, TREX 등 다양한 강화학습 알고리즘을 TreeTensor 로 재구현한 결과, 코드 라인 수와 순환 복잡도 (Cyclomatic Complexity) 가 크게 감소했습니다. (예: AlphaStar collate 함수는 177 줄에서 66 줄로 감소, 복잡도는 28.0 에서 7.6 으로 감소).
• 실행 효율성: Tianshou 의 Batch 라이브러리 및 기타 유사 라이브러리와 비교한 벤치마크에서, 대부분의 연산 (stack, cat, split, init 등) 에서 TreeTensor 가 더 빠른 성능을 보였습니다. 특히 데이터 크기와 차원이 커질수록 성능 차이가 두드러졌습니다.
  • 예: split 연산에서 TreeTensor 는 62µs, Tianshou Batch 는 653µs 로 약 10 배 이상의 속도 향상을 보였습니다.
• 유연성: 가변 길이 데이터 처리 및 비동기 실행 등 다양한 시나리오에 쉽게 확장 가능함이 입증되었습니다.

5. 의의 및 결론 (Significance)

• AI 시스템 개발의 간소화: 복잡한 중첩 데이터를 다루는 AI 시스템 (특히 강화학습, 멀티모달 학습) 의 개발 난이도를 낮추고, 코드의 가독성과 유지보수성을 크게 향상시킵니다.
• 성능과 확장성의 동시 달성: 기존 트리 기반 라이브러리가 가진 성능 저하 문제를 해결하면서도, 기존 텐서 라이브러리의 병렬 처리 능력을 유지합니다.
• 미래 지향성: TreeTensor 는 단순한 데이터 구조를 넘어, 다양한 AI 알고리즘과 라이브러리를 통합하는 플랫폼으로서의 잠재력을 보여주며, 더 복잡하고 일반적인 AI 문제를 해결하는 데 필수적인 도구로 자리 잡을 수 있습니다.

이 논문은 데이터의 복잡성이 증가하는 현대 AI 환경에서, 구조화된 중첩 데이터를 효율적이고 직관적으로 처리할 수 있는 새로운 패러다임을 제시했다는 점에서 중요한 의의를 가집니다.