표형 데이터를 위한 상태공간 모델 기반 사전학습 적합 네트워크

기본 정보

• 논문 ID: 2510.14573
• 제목: State-Space Models for Tabular Prior-Data Fitted Networks
• 저자: Felix Koch, Marcel Wever, Fabian Raisch, Benjamin Tischler
• 분류: cs.LG
• 발표 시간/학회: Proceedings of the 1st ICML Workshop on Foundation Models for Structured Data, Vancouver, Canada. 2025
• 논문 링크: https://arxiv.org/abs/2510.14573

초록

표형 데이터를 위한 기초 모델의 최근 발전(예: TabPFN)은 사전학습된 Transformer 아키텍처가 높은 예측 성능으로 베이지안 추론을 근사할 수 있음을 보여주었습니다. 그러나 Transformer는 시퀀스 길이에 대해 이차 복잡도를 가지므로, 더 효율적인 시퀀스 모델의 탐색이 필요합니다. 본 연구에서는 양방향 선형시간 구조화 상태공간 모델(SSM)인 Hydra를 TabPFN의 Transformer 대체 모델로 사용할 가능성을 조사합니다. 핵심 과제는 SSM의 입력 토큰 순서에 대한 내재적 민감성입니다. 표형 데이터에서는 행의 순서가 의미론적으로 무의미하기 때문입니다. 양방향 접근 방식이 효율성을 유지하면서 대칭적 문맥 집계를 가능하게 하는 정도를 조사합니다. 실험 결과는 이 접근 방식이 순서 의존성을 감소시키면서 원본 TabPFN 모델과 경쟁력 있는 예측 성능을 달성함을 보여줍니다.

연구 배경 및 동기

1. 해결하려는 문제: 본 연구는 표형 데이터 기초 모델에서 Transformer 아키텍처의 계산 효율성 문제, 특히 O(n²) 복잡도가 대규모 데이터셋의 확장성을 제한하는 문제를 다룹니다.
2. 문제의 중요성: TabPFN은 표형 데이터의 기초 모델로서 뛰어난 성능을 보여주며 밀리초 단위로 베이지안 추론 근사를 완료할 수 있지만, Transformer 기반 아키텍처는 대규모 데이터 처리 시 메모리 및 계산 병목 현상에 직면합니다.
3. 기존 방법의 한계:
   • Transformer의 자기주의 메커니즘은 이차 복잡도를 가짐
   • Mamba로 직접 Transformer를 대체하면 입력 시퀀스 순서에 대한 민감성이 도입됨
   • 표형 데이터에서 행의 순서는 의미론적으로 무의미하며, 이는 SSM의 인과적 설계와 충돌함
4. 연구 동기: 구조화 상태공간 모델(SSM)을 Transformer의 대체 모델로 탐색하여 선형 복잡도의 효율성 이점을 유지하면서 양방향 처리 메커니즘을 통해 입력 순서에 대한 의존성을 감소시킵니다.

핵심 기여

1. Hydra 기반 TabPFN 아키텍처 제안: 양방향 구조화 상태공간 모델 Hydra를 TabPFN에 통합하여 표형 데이터의 선형 시간 복잡도 처리를 구현합니다.
2. 반복 문맥 배열(RCP) 기술 도입: 입력을 여러 번 무작위로 배열하고 예측 결과를 평균화하여 SSM의 시퀀스 순서 민감성을 추가로 감소시킵니다.
3. 현저한 확장성 향상 달성: 원본 TabPFN과 비교하여 새로운 방법은 두 자릿수 더 큰 데이터셋을 처리할 수 있습니다(2¹⁵행에서 2¹⁷행으로 확장).
4. 경쟁력 있는 예측 성능 유지: OpenML CC-18 벤치마크 테스트에서 Hydra 기반 TabPFN의 정확도는 원본 모델보다 1.1%만 낮습니다.

방법론 상세 설명

작업 정의

본 논문은 표형 분류 작업을 연구하며, 여기서:

• 입력: 학습 및 테스트 샘플을 포함한 완전한 표형 데이터셋
• 출력: 테스트 샘플에 대한 클래스 확률 예측
• 제약: 단일 순전파에서 추론을 완료해야 하며, 그래디언트 업데이트나 미세조정이 필요 없음

모델 아키텍처

1. Hydra 아키텍처 대체

• 핵심 설계: Hydra 계층 스택으로 Transformer 인코더 대체
• 양방향 처리: 준분리 가능 행렬 믹서를 활용한 양방향 상태공간 모델링
• 계층 구조: 각 Hydra 계층은 양방향 상태공간 믹싱 후 피드포워드 변환으로 구성

2. 임베딩 전략 유지

• 원본 TabPFN의 데이터 임베딩 방법 유지
• 각 입력은 특성값과 클래스 레이블의 연결로 표현
• 추론 시 모든 가능한 레이블 할당을 주변화하여 미표시 데이터 처리

3. 반복 문맥 배열(RCP)

알고리즘 흐름은 다음과 같습니다:

입력: 배열 횟수 r, 문맥 D, 테스트 샘플 xtest
출력: 예측된 클래스 값
초기화: 빈 리스트 outputs ← []
for i = 1 to r do
    D의 행 섞기: Dp ← shuffle(D)
    xtest를 Dp에 연결: Din ← Dp ∪ xtest
    예측: outputs[i] ← PFN.predict(Din)
end for
outputs의 평균값 반환

기술적 혁신 포인트

1. 양방향성이 순서 민감성 해결: 단방향 Mamba와 비교하여 Hydra의 양방향 처리는 문맥 정보를 대칭적으로 집계하여 입력 순서에 대한 의존성을 감소시킵니다.
2. 선형 복잡도: 준분리 가능 행렬 곱셈기를 통해 O(n) 복잡도를 구현하며, Transformer의 O(n²)과 비교하여 현저한 이점이 있습니다.
3. RCP 전략: 여러 번의 무작위 배열과 결과 평균화를 통해 순서 민감성을 추가로 낮추는 혁신적 방법으로, 표형 데이터의 특성을 위한 맞춤형 설계입니다.

실험 설정

데이터셋

• 주요 데이터셋: OpenML CC-18 벤치마크 테스트 스위트
• 필터링 조건: ≤2000행, ≤100특성, ≤10클래스
• 최종 데이터셋: 30개의 다중 클래스 분류 데이터셋
• 데이터 분할: 각 데이터셋을 16회 무작위로 학습/테스트 세트로 분할

평가 지표

1. 정확도(Accuracy): 분류 정확률
2. AUC OvO: One-vs-One 다중 클래스 AUC
3. KL 발산: 다양한 입력 배열 하에서 예측 분포의 차이를 측정하여 순서 민감성 평가
4. 추론 시간: 다양한 입력 규모에서의 계산 시간
5. 메모리 사용: 처리 가능한 최대 데이터셋 규모

비교 방법

• Transformer 기반 TabPFN: 원본 기준 모델
• Mamba 기반 TabPFN: 단방향 SSM 대체 방안
• Hydra 기반 TabPFN: 본 논문에서 제안한 양방향 SSM 방안

구현 세부사항

• 학습 하드웨어: Nvidia A40 GPU (48GB)
• 테스트 하드웨어: NVIDIA H100 80GB
• 학습 시간: Transformer 48시간, Mamba 52시간, Hydra 134시간
• 주요 하이퍼파라미터:
  • 학습률: 0.0001
  • SSM 계층 수: 24층(Transformer의 2배)
  • 임베딩 차원: 1024

실험 결과

주요 결과

1. 확장성 비교

• Transformer 한계: 2¹⁵행(80GB 메모리 제한)
• Hydra 한계: 2¹⁷행(PyTorch 32비트 인덱스 제한, 하드웨어 제한 아님)
• 성능 향상: 처리 가능한 데이터 규모 100배 증가

2. 예측 성능 비교

• Hydra vs Transformer: 정확도 평균 차이 -1.1%, AUC 차이 -1.1%
• Hydra vs Mamba: Hydra 정확도 평균 3.6% 높음
• 분산 분석: Hydra는 Mamba보다 낮은 성능 분산을 보임

3. 순서 민감성 분석

KL 발산으로 측정:

• RCP 횟수 증가에 따라 KL 발산이 현저히 감소
• Hydra는 Mamba보다 낮은 순서 민감성을 보임
• RCP 전략은 이상 배열의 영향을 효과적으로 감소

소거 실험

RCP 횟수의 영향

• 정확도: RCP 횟수 증가에 따라 향상되지만 개선 폭은 상대적으로 작음
• KL 발산: 현저히 감소하여 순서 의존성 감소를 나타냄
• 계산 비용: 추론 시간이 r배 선형 증가

아키텍처 비교

• 단방향 vs 양방향: Hydra의 양방향 메커니즘이 Mamba의 단방향 처리보다 명확히 우수
• 계층 수 설정: Mamba 논문의 권장사항을 따라 Transformer 계층 수의 2배 사용

실험 발견

1. 양방향성의 중요성: 양방향 처리는 표형 데이터의 무순서 특성에 매우 중요
2. 효율성과 성능의 균형: 경쟁력 있는 성능을 유지하면서 현저한 효율성 향상 달성
3. RCP의 효과성: 다중 배열 평균 전략이 순서 민감성을 효과적으로 감소
4. 하드웨어 제한 극복: Transformer의 대규모 데이터 메모리 제한을 성공적으로 극복

관련 연구

표형 기초 모델

• TabPFN: 표형 데이터의 획기적인 Transformer 모델
• TabFlex: 선형 주의를 사용한 확장 방안
• Mambular: Mamba 기반의 표형 심층학습 모델

상태공간 모델

• Mamba: 선택적 상태공간 모델로 선형 복잡도 구현
• Hydra: 양방향 SSM 확장으로 비인과 모델링 지원
• S4: 구조화 상태공간 시퀀스 모델의 기초 연구

효율성 최적화 방법

• FlashAttention: IO 최적화를 통한 Transformer 메모리 요구 감소
• Linear Attention: 선형 복잡도의 주의 메커니즘 대체 방안

결론 및 논의

주요 결론

1. Hydra는 TabPFN의 확장성 문제를 성공적으로 해결하여 처리 능력을 두 자릿수 향상
2. 양방향 SSM은 단방향 SSM보다 표형 데이터의 무순서 특성에 더 적합
3. RCP 전략은 SSM 순서 민감성을 감소시키는 효과적인 방법
4. 선형 복잡도를 유지하면서 Transformer와 경쟁력 있는 성능 달성

한계

1. 재학습 필요: 아키텍처 차이로 인해 전체 모델 재학습 필요
2. 문맥 제한: 실험은 여전히 1000행 이내로 제한되어 대규모 시나리오를 충분히 탐색하지 못함
3. RCP 오버헤드: 다중 배열이 r배의 추론 시간 증가
4. 순서 최적화: 최적 배열 전략에 대한 심층 연구 부족

향후 방향

1. 대규모 검증: >10k행의 데이터셋에서 SSM 기반 TabPFN 테스트
2. 최적 배열: SSM을 위한 최적 행 배열 전략 연구
3. 아키텍처 최적화: 더 효율적인 양방향 SSM 아키텍처 탐색
4. 이론적 분석: 양방향성이 표형 데이터 모델링에 미치는 이론적 기초 심층 이해

심층 평가

장점

1. 문제 정의의 명확성: TabPFN의 핵심 병목을 정확히 파악하고 맞춤형 해결책 제시
2. 기술 선택의 합리성: Hydra의 양방향 특성이 표형 데이터의 무순서 특성과 잘 부합
3. 완전한 실험 설계: 성능, 효율성, 순서 민감성 등 다차원적 평가 포함
4. 강한 결과 설득력: 성능 유지하면서 현저한 확장성 향상 달성
5. 높은 방법 실용성: RCP 전략은 단순하고 효과적이며 구현 및 배포 용이

부족한 점

1. 제한된 혁신 정도: 주로 기존 기술의 조합 응용으로 근본적 혁신 부족
2. 불충분한 이론적 분석: 양방향성이 순서 민감성 문제를 해결하는 이유에 대한 심층 이론적 설명 부족
3. 제한된 실험 규모: 여전히 상대적으로 작은 데이터셋으로 제한되어 대규모 처리 능력을 충분히 입증하지 못함
4. 불완전한 비교: 다른 선형 복잡도 방법(예: Linear Attention)과의 직접 비교 부족
5. 불충분한 하이퍼파라미터 분석: 높은 학습 비용으로 인해 충분한 하이퍼파라미터 최적화 미실시

영향력

1. 학술적 기여: 표형 기초 모델의 효율성 최적화에 새로운 사고와 실증적 증거 제공
2. 실용적 가치: 실제 응용에서의 확장성 문제 해결로 높은 실용 가치 보유
3. 영감 제공: SSM의 구조화 데이터 모델링 잠재력을 보여주어 관련 연구 영감 제공 가능
4. 재현성: 코드 공개 가능하고 실험 설정이 상세하여 우수한 재현성 보유

적용 시나리오

1. 대규모 표형 분류: 특히 많은 샘플을 처리해야 하는 표형 분류 작업에 적합
2. 실시간 추론 시나리오: 선형 복잡도로 추론 속도에 엄격한 요구사항이 있는 응용에 적합
3. 자원 제한 환경: Transformer보다 적은 메모리와 계산 자원 필요
4. 소수 샘플 학습: TabPFN의 소수 샘플 시나리오에서의 이점 유지

참고문헌

주요 참고문헌 포함:

1. Hollmann et al. (2023) - TabPFN 원본 논문
2. Gu & Dao (2023) - Mamba 아키텍처
3. Hwang et al. (2024) - Hydra 양방향 SSM
4. Dao et al. (2022) - FlashAttention 최적화 기술
5. Zeng et al. (2024) - TabFlex 선형 주의 방법

본 논문은 표형 기초 모델의 확장성 문제 해결에 가치 있는 기여를 하였으며, 양방향 SSM과 반복 배열 전략을 교묘하게 결합하여 효율성과 성능의 요구를 성공적으로 균형 있게 조정했습니다. 이론적 혁신 측면에서 부족함이 있지만, 실용적 가치와 향후 연구에 대한 영감 제공 의미는 인정할 만합니다.