표형 데이터에 효율적인 상태공간 모델 기반 TabPFN

초록

본 논문은 TabPFN의 Transformer 백본을 선형 시간 복잡도를 갖는 구조화 상태공간 모델(SSM)인 Hydra로 교체하고, 행 순서에 대한 민감성을 완화하기 위해 양방향 처리와 Repeated Context Permutations(RCP)를 도입한다. 실험 결과, Hydra‑ 기반 TabPFN은 메모리·시간 효율이 크게 향상되면서도 Transformer 대비 예측 정확도와 AUC 차이가 1 % 이하로 유지됨을 보였다.

상세 분석

TabPFN은 대규모 합성 데이터로 사전 학습된 Transformer를 이용해 테이블 전체를 한 번에 입력하고, 베이지안 추론을 근사함으로써 몇 밀리초 안에 고성능 예측을 제공한다. 그러나 Transformer의 자기‑Attention 연산은 입력 행 수 N에 대해 O(N²) 메모리·시간 복잡도를 가지므로, 수천 행을 초과하는 데이터셋에서는 GPU 메모리 한계에 부딪힌다. 이를 해결하고자 저자들은 선형 시간 SSM인 Mamba와 그 양방향 확장인 Hydra를 검토한다. Mamba는 상태 변수와 선택적 업데이트 메커니즘을 통해 O(N) 복잡도를 제공하지만, 기본적으로 causal(단방향) 구조이기 때문에 입력 행 순서에 크게 의존한다. 테이블 데이터는 행 순서가 의미가 없으므로, 이러한 순서 민감성은 성능 저하와 불안정성을 초래한다.

Hydra는 quasi‑separable matrix mixer를 활용해 전·후 방향으로 상태를 혼합함으로써 양방향 컨텍스트를 동시에 고려한다. 이는 순서 의존성을 완화하고, Transformer가 제공하는 전역 상호작용을 어느 정도 재현한다. 저자들은 Hydra를 TabPFN에 적용하기 위해 (1) Transformer Encoder를 Hydra 레이어 스택으로 교체하고, (2) 기존 임베딩 방식을 그대로 유지하며, (3) 합성 작업 분포에 맞춰 전체 모델을 재학습하는 파이프라인을 설계했다.

또한, 순서 의존성을 추가로 감소시키기 위해 Repeated Context Permutations(RCP)를 도입했다. RCP는 입력 테이블을 r번 무작위로 섞은 뒤 각각에 대해 예측을 수행하고, 결과 확률을 평균한다. 이 과정은 추론 시간을 r배 증가시키지만, KL‑다이버전스가 현저히 감소하고, 특히 불리한 행 순서에 의해 발생하는 오류를 평균화함으로써 정확도가 소폭 상승한다는 실험적 증거를 제시한다.

실험은 OpenML CC‑18 벤치마크(30개 데이터셋, 각 16번 랜덤 분할)에서 수행되었다. 결과는 다음과 같다. 첫째, Hydra는 입력 크기가 2¹⁷(≈130 k) 행까지 메모리 초과 없이 처리 가능했으며, Transformer는 2¹⁵(≈33 k) 행에서 한계에 봉착했다. 둘째, 정확도와 AUC 면에서 Hydra는 평균 1.1 % 차이로 Transformer와 거의 동등한 성능을 보였으며, Mamba보다 변동성이 적었다. 셋째, RCP를 5번 적용했을 때 KL‑다이버전스가 크게 감소하고, 정확도는 약 0.5 % 정도 향상되었다(하지만 AUC에는 유의미한 변화가 없었다).

이러한 결과는 (i) 선형‑시간 SSM이 대규모 테이블에 대한 실용적인 대안이 될 수 있음을, (ii) 양방향 설계가 순서 민감성을 완화하면서도 전역 컨텍스트를 효과적으로 통합한다는 점을, (iii) 무작위 순열 평균화가 모델의 안정성을 높이는 간단하지만 유용한 기법임을 시사한다. 향후 연구에서는 10 k 행을 넘어서는 초대형 테이블에 대한 확장, 최적 행 순서 탐색, 그리고 FlashAttention 등 기존 Transformer 가속기법과의 혼합 사용을 검토할 여지가 있다.