기술 요약: 표 형식 파운데이션 모델을 위한 엔드 투 엔드 압축 (End-to-End Compression for Tabular Foundation Models)

문제 정의

디코더 전용 트랜스포머 아키텍처를 활용한 인컨텍스트 러닝(ICL)을 사용하는 표 형식 파운데이션 모델(TFM)은 최근 그래디언트 부스팅 결정 트리(GBDT)의 지배력에 도전하고 있습니다. 그러나 이러한 모델들은 근본적인 스케일링 병목 현상에 직면해 있습니다. 즉, 추론 복잡도가 훈련 인스턴스 수 $N$에 대해 이차적($O(N^2)$) 또는 선형적($O(N)$, KV 캐싱 사용 시)으로 증가합니다. 이러한 조건부 컨텍스트 크기에 대한 이차적 의존성은 메모리와 연산 오버헤드를 과도하게 발생시켜, TFM이 소규모 또는 중규모 데이터셋에 국한되도록 만듭니다. 결과적으로, 모델은 정보가 풍부한 대규모 데이터셋에서 성능이 저하되는 경향이 있으며, 거대한 컨텍스트 윈도우를 요구하는 실제 시나리오를 처리하는 데 어려움을 겪습니다.

방법론: TACO

이 병목 현상을 해결하기 위해 저자들은 예측 전 학습 데이터셋을 잠재 공간(latent space)으로 압축하도록 설계된 엔드 투 엔드 모듈형 아키텍처인 TACO(Tabular Compression)를 제안합니다. 이 프레임워크는 두 개의 공동 학습된 구성 요소로 이루어집니다:

1. 압축기 (Compressor, $g$): 원래의 훈련 데이터셋 $D_{train}$ ($N$개의 행 보유)을 $K$개의 행을 가진 조밀하고 낮은 랭크의 표현인 $D_{mini-train}$으로 매핑하는 트랜스포머 기반 모듈입니다 ($K \ll N$). 압축기는 데이터의 원형적 패턴(prototypical patterns)을 포착하는 응축된 임베딩을 학습하도록 훈련됩니다. 이는 데이터의 일부로부터 초기화된 "더미(dummy)" 테이블을 사용하여, 전체 데이터셋을 대표하도록 반복적으로 정제됩니다.
2. 예측기 (Predictor, $f$): 인컨텍스트 예측을 수행하는 표준 표 형식 트랜스포머(TabPFNv2 아키텍처 기반)입니다. 예측기는 전체 $N$-행 데이터셋에 어텐션을 수행하는 대신, 압축된 컨텍스트 $D_{mini-train}$과 테스트 배치 특징(features)에 조건화됩니다.

학습 전략:
압축기와 예측기는 엔드 투 엔드 방식으로 공동 메타 학습(meta-learned)됩니다. 목적 함수는 데이터셋의 사전 분포(prior)에 대해 예측 손실을 최소화하는 것이며, 이는 압축기가 예측기의 성능을 극대화하는 표현을 생성하도록 유도합니다. 학습 파이프라인은 다음을 포함합니다:

• 합성 데이터 사전 훈련: 구조적 인과 모델(SCM)을 사용하여 다양한 합성 데이터셋을 생성합니다.
• 커리큘럼 학습: 훈련 중에 시퀀스 길이를 점진적으로 증가시킵니다 (1k에서 60k 행까지).
• 실제 데이터 적응: 자기 지도 학습(행/특징 부트스트래핑)을 사용하여 실제 표 형식 데이터셋에 미세 조정합니다.
• 멀티 레이트(Multi-Rate) 훈련: 다양한 리소스 제약 조건 하에서 작동할 수 있도록 여러 압축률($r \in \{1, 2, 4, 8, 16\}\%$)을 동시에 처리하도록 모델을 훈련합니다.

확장 메커니즘:
단일 순방향 패스(forward pass)의 메모리 용량을 초과하는 데이터셋(예: 100만 행 이상)의 경우, TACO는 청킹 및 스티칭(chunking-and-stitching) 전략을 사용합니다. 데이터셋을 서로 겹치지 않는 청크로 분할하여 각각 독립적으로 압축한 뒤, 생성된 요약본들을 연결하여 예측기를 위한 글로벌 컨텍스트를 형성합니다.

주요 기여

1. 새로운 아키텍처: 표 형식 파운데이션 모델의 훈련 컨텍스트를 압축하는 최초의 엔드 투 엔드 방식인 TACO를 도입하여, 예측 품질을 원시 조건부 컨텍스트 크기로부터 분리했습니다.
2. 효율성 이득: TACO가 최신 예측기 전용 트랜스포머(POT)와 비교했을 때, 1% 압축률에서 성능을 유지하면서도 추론 시간을 최대 94배 가속화하고 98%의 메모리 절감을 달 achieves함을 실증적으로 보여주었습니다.
3. 확장성: 메모리 제약으로 인해 표준 TFM이 불가능한 작업(예: 약 150만 행의 MetroPT-3 데이터셋)을 수용할 수 있는 검증된 전략을 제시합니다.
4. 공동 최적화: 압축기와 예측기를 공동 훈련하는 것이 예측기를 고정(freezing)하는 것보다 우수하다는 증거를 제시하며, 이를 통해 예측기가 압축된 입력에 맞춰 자신의 표현 공간을 적응할 수 있게 합니다.

실험 결과

저자들은 TabArena 벤치마크와 추가적인 대규모 데이터셋(TabFSBench, TableShift, MetroPT-3)에서 TACO를 평가했습니다:

• 성능 대비 속도: 1% 압축률에서 TACO는 압축되지 않은 POT 베이스라인에 비해 후속 예측에서 94배의 속도 향상을 달성하고 메모리 소비를 97.6% 줄였습니다. 성능 저하는 전체 컨텍스트 베이스라인과 비교했을 때 통계적으로 무시할 수 있는 수준입니다 (ROC-AUC 차이가 0에 가까움).
• 메모리 효율성: KV 캐싱을 사용하는 시나리오에서 POT는 32,000개 이상의 행이 있는 데이터셋에서 빈번하게 OOM(Out-of-Memory) 오류를 겪는 반면, TACO는 실행 가능합니다. TACO는 KV 캐싱 설정에서 VRAM 사용량을 최대 92%까지 줄입니다.
• 베이스라인 비교: TACO는 모든 압축률에 걸쳐 무작위 샘플링(random sampling) 및 k-최근접 이웃(kNN) 서브샘플링과 같은 단순 압축 베이스라인보다 성능이 뛰어납니다.
• 대규모 데이터 처리: 약 150만 행의 MetroPT-3 데이터셋에서 TACO는 훈련 세트를 약 1,200개의 행으로 압축했으며, 베이스라인 및 XGBoost보다 우수한 정밀도(precision)와 AUPRC를 달성했습니다. 반면 POT와 TabPFNv2는 실행에 실패했습니다.
• 강건성: TableShift 벤치마크에서 TACO는 메모리 제한으로 인해 표준 TFM이 실패한 15개 데이터셋 중 12개에서 성공적으로 작동했으며, 분포 변화(distribution shift) 하에서도 경쟁력 있는 정확도를 유지했습니다.

의의 및 주장

본 논문은 TACO가 엄격한 추론 시간 예산 하에서도 예측력을 유지하는 데이터 표현을 학습함으로써, 표 형식의 인컨텍스트 추론을 근본적으로 재구성한다고 주장합니다. 저자들은 이 방법이 원칙적인 정확도-지연 시간(accuracy-latency) 트레이드오프를 가능하게 하여, 인컨텍스트 표 형식 예측기가 기존의 트랜스포머 기반 접근법으로는 접근할 수 없었던 대규모 실제 데이터셋으로 확장될 수 있게 한다고 단언합니다.

이 연구는 컨텍스트를 지능적으로 압축한다면 어텐션 메커니즘의 이차 복잡도가 극복 불가능한 장벽이 아님을 강조합니다. 컨텍스트 크기를 $N/K$ 배만큼 줄임으로써, TACO는 메모리와 연산 비용의 스케일링을 선형화하여 고성능 파운데이션 모델이 대규모 산업용 데이터를 다루는 데 있어 실용적으로 기능하게 만듭니다. 저자들은 TACO가 확장성을 진전시켰으나, 향후 연구에서는 결측값에 대한 다양한 사전 분포, 분포 변화, 그리고 회귀 및 시계열 작업으로의 확장을 다루어야 한다고 결론지었습니다.