TabSwift: An Efficient Tabular Foundation Model with Row-Wise Attention

기술 요약: TABSWIFT

문제 정의
TabPFN과 같은 표 형식 파운데이션 모델(Tabular Foundation Models, TFMs)은 테스트 시 파라미터 업데이트 없이 인컨텍스트 학습(In-Context Learning, ICL)을 활용함으로써 소-중규모 데이터셋에서 경쟁력 있는 성능을 입증해 왔습니다. 그러나 최근 TFM의 발전은 특징별 이질성(feature-wise heterogeneity)을 포착하기 위해 점점 더 복잡한 아키텍처(예: 행/열 교차 어텐션)에 의존하는 경향이 있습니다. 이러한 복잡한 설계는 추론 비용과 메모리 사용량을 크게 증가시켜, 지연 시간에 민서한 실제 배포 환경에서의 활용을 제한합니다. 경량화된 행 방향 어텐션 전용 백본(예: TabPFN v1)은 효율적이지만, 고도로 이질적인 작업에서는 성능의 한계에 부딪힐 위험이 있습니다. 본 연구가 다루는 핵심 질문은, 현대적인 안정화 및 사전 학습 기술으로 강화된 경량 행 방향 어텐션 전용 백본이 우수한 정확도-효율성 트레이드오프를 제공하면서 경쟁력 있는 파운데이션 모델 성능을 달나할 수 있는가 하는 점입니다.

방법론
저자들은 TABSWIFT를 제안합니다. 이는 효율적인 행 방향 어텐션 전용 구조를 유지하면서 세 가지 핵심 구성 요소를 도입한 경량 표 형식 파운데이션 모델입니다.

1. 안정화된 행 방향 백본 (Stabilized Row-Wise Backbone):

   • 아키텍처: 모델은 각 테이블 행을 하나의 토큰으로 취급하며, TabPFN v2나 TabICL에서 발견되는 비용이 많이 드는 열 방향 어텐션 패스를 피하고 오직 행(인스턴스) 간의 셀프 어텐션만을 적용합니다.
   • 게이트 어텐션 (Gated Attention): 사전 학습 중 최적화 안정성을 개선하기 위해, 저자들은 요소별 게이트 어텐션 메커니즘을 통합했습니다. 이는 프로젝션 전 스케일드 닷 프로덕트 어텐션(SDPA) 출력에 곱셈 게이트를 적용하여, 무시할 만한 오버헤드로 입력 의존적 희소성과 비선형성을 도입합니다.
   • 학습 가능한 레지스터 토큰 (Learnable Register Tokens): 입력 시퀀스 앞에 작은 규모의 학습 가능한 레지스터 토큰 세트를 추가합니다. 이 토큰들은 글로벌 컨텍스트와 태스크 수준의 정보를 집계하는 잠재 슬롯 역할을 하며, 모델이 레이어를 거치며 데이터셋 수준의 표현을 유지하고 정교화하도록 돕습니다.

2. 통합 사전 학습 공식 (Unified Pretraining Formulation):

   • TABSWIFT는 단일 사전 학습된 모델 내에서 분류와 회귀를 모두 지원하는 통합된 목적 함수로 학습됩니다.
   • 사전 학습 파이프라인은 대규모 합성 표 형식 작업 풀(20,000 스텝, 스텝당 512개 작업)을 활용합니다. 각 합성 작업에 대해, 타겟은 이산화된 분류 타겟과 정규화된 회귀 타겟을 모두 제공하도록 처리되어, 각 작업 유형에 대한 별도의 경량 MLP 헤드를 동시에 최적화할 수 있게 합니다.

3. 적응형 레이어별 조기 종료 메커니즘 (Adaptive Layer-Wise Early-Exit Mechanism):

   • 추론 지연 시간을 더욱 줄이기 위해, 모델은 쿼리당 적응형 조기 종료 전략을 채택합니다.
   • 메커니즘: 마지막 $E$개의 트랜스포머 레이어에 경량 예측 헤드와 "종료 헤드(exit heads)"를 부착합니다. 종료 헤드는 현재 쿼리 표현과 레지스터 토큰의 풀링된 요약을 결려하여 신뢰도 점수를 예측합니다.
   • 정책: 추론 중에 모델은 레이어를 순차적으로 처리합니다. 특정 쿼리에 대한 예측 신뢰도 점수가 임계값 $\tau$를 초과하면, 모델은 조기에 종료하고 예측값을 출력합니다. 이를 통해 "쉬운" 샘플은 더 적은 레이어로 처리하고, "어려운" 샘플은 전체 깊이를 통과하게 함으로써 "언제든(anytime)" 추론이 가능하게 합니다.

주요 기여

• TABSWIFT 모델: 게이트 어텐션으로 안정화되고 레지스터 토큰으로 증강된 행 방향 어텐션 전용 백본이, 훨씬 더 무겁고 복잡한 TFM(예: TabPFN v2, TabICL)과 경쟁할 만한 성능을 달성하면서도 현저히 낮은 추론 비용을 유지함을 입증했습니다.
• 통합 아키텍처: 분류 및 회귀 작업을 모두 처리할 수 있는 단일 사전 학습 모델을 도입하여 배포를 단순화했습니다.
• 적응형 추론: 신뢰도 예측 기반의 조기 종료 전략을 제안하여, 예측 성능의 손실을 최소화하면서도 평균 추론 시간을 실질적으로 단축했습니다.
• 오픈 소스: 코드가 커뮤니티에 공개되었습니다.

실험 결과
저자들은 300개의 이진, 다중 클래스 및 회귀 데이터셋으로 구성된 대규모 컬렉션인 TALENT 벤치마크에서 TABSWIFT를 평가했습니다.

• 성능: TABSWIFT는 분류(AUC) 및 회귀(RMSE) 작업 모두에서 가장 강력한 베이스라인(TabPFN v2, TabICL, LimiX)과 대등한 평균 순위를 달성했습니다. 통계적 검정(Wilcoxon-Holm) 결과, 상위 성능 모델들과 유의미한 차이가 없는 경우가 많았습니다.
• 효율성: 추론 시간 측정 결과, TABSWIFT는 복잡한 TFM보다 현저히 낮은 지연 시간으로 작동합니다. 그림 1은 TABSWIFT가 높은 성능과 낮은 지연 시간이라는 이상적인 지점에 더 가깝게 위치하며, 훨씬 더 나은 정확도-효율성 트레이드오프를 달성함을 보여줍니다.
• 절제 연구 (Ablation Studies): 실험을 통해 게이트 어텐션과 레지스터 토큰 모두가 베이스라인 행 방향 어텐션 모델 대비 성능 향상에 기여함을 확인했습니다.
• 조기 종료 분석: 종료 임계값 $\tau$를 변화시킴에 따라 유리한 파레토 프런티어(Pareto frontier)가 나타납니다. 레지스터 토큰에 조건화된 학습된 종료 헤드는 엔트로피 기반의 중단 베이스라인보다 우수한 성능을 보였습니다. 임베딩 공간 시각화는 모델이 모호하거나 중첩된 샘플을 더 깊은 계산이 필요한 것으로 올바르게 식별하고, 확신이 있는 샘플은 조기에 종료함을 보여줍니다.

의의 및 주장
본 논문은 TABSWIFT가 엄격한 추론 예산이 요구되는 실제 배포 환경에 적합한 효율적이고 언제든 사용 가능한(anytime) 표 형식 인컨텍스트 학습을 가능하게 한다고 주장합니다. 미니멀한 행 방향 설계를 재검토하고 현대적인 안정화 기술을 적용함으로써, 저자들은 최첨단 결과를 얻기 위해 반드시 계산 비용이 많이 드는 교차 어텐션 패턴을 채택할 필요는 없다고 주장합니다. 이 연구는 TABSWIFT가 모델 선택이나 앙상블을 위한 강력하고 효율적인 구성 요소로서, 지연 시간에 민감한 실제 표 형식 애플리케이션의 서빙을 위한 견고한 솔루션을 제공함을 시사합니다.