Speedrunning Tabular Foundation Model Pretraining

기술 요약: 테이블형 파운데이션 모델(Tabular Foundation Model) 사전 학습의 스피드런

문제점
테이블형 파운데이션 모델(TFM)을 사전 학습하는 것은 계산 비용이 많이 드는 병목 현상이며, 이는 새로운 아키텍처, 사전 지식(prior), 최적화 전략의 반복 주기를 늦춥니다. 대규모 언어 모델(LLM)의 사전 학습을 가속화하기 위한 상당한 노력이 기울여진 반면, TFM을 위한 동등한 커뮤니티 주도 프레임워크는 존재하지 않습니다. 높은 사전 학습 비용은 빠른 프로토타이핑과 속도 향상 기술의 축적을 제한합니다.

방법론
저자들은 TabPFN v2를 최소한의 교육용으로 재구현한 nanoTabPFN을 중심으로 한 공개적인 "스피드run" 경진대회를 소개합니다. 이 경진대회는 기여자가 단일 파일로 구성된 훈련 스크립트를 수정하여 특정 다운스트림 성능 목표에 도달하는 데 걸리는 실제 벽시계 시간(wallclock time)을 최소화하는 고정된 프로토콜을 수립합니다.

• 목표: 목표는 서브샘플링된 TabArena 데이터셋(34개 분류 작업)에 대해 랜덤 포레스트(Random Forest) 베이스라인보다 높은 평균 검증 ROC AUC를 달시키는 신경망을 사전 학습하는 것입니다.
• 제약 조건: 모든 실행은 단일 NVIDIA L40S GPU를 사용해야 합니다. 보고된 시간은 클러스터 노이즈를 억제하기 위해 여러 번의 검증 실행 결과의 중앙값입니다. 기록은 사전 학습된 가중치를 로드해서는 안 되며, 고정된 평가 파이프라인을 사용해야 합니다.
• 진행 방식: 스피드런은 누적적입니다. 즉, 각 새로운 기록은 이전의 기술을 기반으로 구축되어, 특정 최적화의 영향을 격리하여 보여줍니다.

주요 기여 및 기술
본 논문은 새로운 모델 제품군을 제안하는 것이 아니라, 언어 모델링에서 다양한 기술이 테이블형 데이터로 어떻게 전이되는지를 보여주는 재현 가능한 프로토콜과 일련의 기록들을 제안합니다. 베이스라인(74.32분)에서 현재 최고 기록(0.92분)까지의 진행 과정에는 다음과 같은 주요 기술적 개입이 포함되었습니다:

1. 옵티마이저 교체: 2D 가중치 행렬에 대해 스케줄 프리(schedule-free) AdamW 옵티마이저를 Muon 옵티마이저로 교체하여 샘플 효율성을 크게 개선했습니다.
2. 아키텍처 및 정밀도:
   • Scaled Dot-Product Attention (SDPA) 구현.
   • Post-norm에서 Pre-norm 아키텍처로 전환.
   • 혼합 정밀도 훈련(bf16) 및 TF32 행렬 곱셈 채택.
   • 하이퍼파라미터 조정(학습률을 $10^{-3}$으로 상향) 및 어텐션 헤드를 줄이면서 임베딩 차원 확장 (192 $\to$ 256).
3. 처리량(Throughput) 최적화: Muon을 위한 Newton–Schulz 반복(iteration) 배치 처리 및 트랜스포머 인코더 순전파(forward pass) 컴파일.
4. 아키텍처 수정:
   • 잔차 감쇠(Residual Decay): 초기 트랜스포머 블록으로 들어오는 입력을 지수적으로 가중치 감소 시킴 ($0.95^i$).
   • RMSNorm: LayerNorm을 더 낮은 정밀도의 RMSNorm으로 교체.
   • 생각하는 행(Thinking Rows): 입력에 16개(이후 24개)의 학습 가능한 "생각하는 행"을 추가하여 데이터 행이 이를 참조하게 함으로써 성능을 크게 향상시킴.
   • 반복적 특징 그룹화(Repeated Feature Grouping): 유사한 분포를 가진 특징들의 표현 붕괴를 방지하기 위해 특정 오프셋(예: $j, j+1, j+3$)과 함께 열을 공동 임베딩함.
5. 고급 훈련 전략:
   • 최신 가중치 평균화(Latest Weight Averaging, LAWA): 평가를 위해서만 마지막 10개의 체크포인트를 평균함.
   • 가중치 감쇠(Weight Decay): AdamW와 Muon 파라미터 모두에 가중치 감쇠 적용.
   • 디코더 수정: 디코더에 타겟 임베딩 대신 특징 임베딩의 평균을 전달.
   • 오토리서치(Autoresearch): 최종 기록은 최적의 구성을 발견하기 위해 LLM 기반의 하이퍼파라미터 및 아키텍처 탐색 루프(autoresearch)를 활용했습니다.

결과
현재 최고 기록은 단일 NVIDIA L40S GPU에서 목표 다운스트림 성능을 0.92분 만에 달성했습니다. 이는 다음을 의미합니다:

• 74.32분인 베이스라인 대비 벽시계 시간 기준 81배의 속도 향상.
• 필요한 합성 데이터셋 수의 22배 감소 (80,576개에서 3,648개로).
• 전체 베이스라인 예산(256,000개 데이터셋) 동안 훈련했을 때, 최적화된 모델은 전처리나 앙상블 없이도 TabPFN v2.5(0.8286)에 근접하는 0.82의 평균 ROC AUC에 도달했습니다.

의의
본 논문은 주요 기여가 스피드런 형식 그 자체라고 주장하며, 이는 커뮤니티가 사전 학습 개선 사항을 추가, 검증 및 중첩할 수 있는 간단하고 개방된 프로토콜을 제공합니다. 합성 사전 지식(synthetic prior)과 평가 파이프라인을 고정함으로써, 리더보드는 어떤 실질적인 최적화가 TFM에 효과적인지를 격리하여 보여줍니다. 저자들은 현재의 최고 기록이 합성 사전 지식을 고정된 상태로 유지하고 있음을 언급하며, 이는 해당 사전 지식 자체가 향 향후 개선을 위한 "손대지 않은 방향"임을 시사합니다. 이 연구는 언어 모델링에서 성공적인 기술들(Muon, SDPA, 특정 아키텍처 미세 조정 등)이 테이블형 파운데이션 모델로 효과적으로 전이되어 훈련 비용을 획기적으로 줄일 수 있음을 입증합니다.