표 형식 파운데이션 모델을 위한 사전 분포 정렬 기반 데이터 클리닝

표 형식 파운데이션 모델(TFM)은 적은 데이터로도 강력한 성능을 내지만, 입력 데이터에 결측치나 이상치가 있으면 성능이 급격히 저하된다. 이 논문은 강화학습을 통해 데이터 클리닝 순서를 최적화하여 모델이 학습한 가상 분포와 실제 데이터를 일치시킴으로써 정확도를 높이고 모델의 확신 편향 문제를 해결한다.

TabPFN과 같은 표 형식 파운데이션 모델은 수백만 개의 가상 데이터셋을 생성하여 학습된다. 이때 모델은 입력 데이터가 대략적으로 깨끗한 가우시안 분포를 따를 것이라고 가정하는 '내부 사전 분포(Prior)'를 가지게 된다. 하지만 실제 데이터는 결측치, 중복, 이상치가 가득하여 이 사전 분포와 큰 간극(Prior Mismatch)이 발생하며, 이는 모델의 예측 정확도를 떨어뜨릴 뿐만 아니라 모델이 틀린 답을 내면서도 강하게 확신하게 만드는 신뢰도 문제를 야기한다.

이 문제를 해결하기 위해 데이터 클리닝을 단순한 전처리가 아닌, 모델의 사전 분포에 데이터를 맞추는 과정으로 재정의했다. 클리닝 연산자는 적용 순서에 따라 데이터 분포를 계속 변화시키기 때문에 정적인 규칙으로는 최적의 결과를 얻기 어렵다. 따라서 각 단계의 데이터 상태를 9차원 벡터로 관찰하고, 다음 연산자를 결정하는 순차적 의사결정 문제로 접근하여 강화학습 에이전트가 최적의 경로를 찾도록 했다.

특히 인컨텍스트 러닝 메커니즘에서 데이터의 양(행의 수)이 줄어들면 예측의 불확실성이 비선형적으로 증가한다는 점에 주목했다. 이를 위해 행을 삭제할 때 단순 선형 페널티가 아닌 제곱 형태의 페널티를 부여하여, 모델이 가능한 한 데이터를 보존하면서도 분포를 정렬할 수 있는 정교한 클리닝 전략을 학습하게 유도했다.

데이터 클리닝 과정을 유한 시간 지평(Finite-horizon) 마르코프 의사결정 과정(MDP)으로 모델링했다. 상태 공간 S는 결측률, Wasserstein 거리 기반의 분포 드리프트, 왜곡도(Skewness), 첨도(Kurtosis), 클래스 균형, 행 보존율 및 적용된 액션 이력을 포함하는 9차원 벡터로 구성된다.

액션 공간 A는 결측치 대체(Mean, Median, KNN), 이상치 제거(IQR, Z-score), 스케일링(Min-Max, Z-score)의 세 가지 가족으로 나뉘며, 각 액션은 연속적이거나 이산적인 세부 파라미터를 포함한다. 에이전트는 PPO(Proximal Policy Optimization) 알고리즘을 사용하여 이 복잡한 액션 공간에서 보상을 극대화하는 정책을 학습한다.

보상 함수는 본 논문의 핵심으로, TFMAwareReward를 제안했다. 이는 [TabPFN 정확도] + [행 보존율의 제곱] + [데이터 품질 점수] - [분포 왜곡 페널티]의 조합으로 계산된다. 특히 행 보존율에 제곱(α=2)을 적용하여, 데이터가 적어질수록 예측 불확실성이 급증하는 인컨텍스트 러닝의 특성을 수학적으로 반영했다.

10개의 OpenML 벤치마크 데이터셋을 대상으로 실험한 결과, 제안된 TFMAwareReward를 사용한 방식이 기존 Random Forest 기반 보상 방식보다 4개 데이터셋에서 구조적으로 더 우수한 클리닝 파이프라인을 선택했으며, 나머지 6개에서도 동등한 성능을 보였다. 평균 정확도는 0.851로 대조군(0.843) 대비 유의미한 향상을 기록했다.

신뢰도 측면에서, 사전 분포 정렬 클리닝을 거친 데이터는 정제하지 않은 경우나 표준 전처리를 거친 경우보다 기대 보정 오차(ECE)가 모든 오류 유형(결측치, 이상치, 중복 등)에서 감소했다. 특히 중복 데이터가 포함된 경우 ECE 개선 효과가 가장 뚜렷하게 나타났다.

전이 학습 실험에서는 Ionosphere 데이터셋으로 사전 학습된 정책을 Phoneme, Adult, Bank 데이터셋에 적용했을 때, 단 2,000단계의 미세 조정만으로도 처음부터 5,000단계 동안 학습한 모델의 성능을 압도했다. 이는 클리닝 정책이 데이터셋에 특화된 지식을 넘어, 일반적인 데이터 정제 원리를 학습했음을 시사한다.

L2C2 아키텍처는 DataQualityObserver, Parameterized Action Module, TFMAwareReward Function의 세 가지 핵심 컴포넌트로 구성된다. 상태 벡터의 Wasserstein 거리 계산은 1차원 주변 분포(Marginal distribution)에 대해 정렬 기반의 폐쇄형(Closed-form) 계산을 수행하여 연산 효율성을 확보했다.

보상 함수 설계에서 기존의 단순 정확도 보상이 왜 행 삭제와 같은 퇴행적 전략(Degenerate strategy)으로 이어지는지 실험적으로 증명했다. 예를 들어 결측치가 있는 행을 모두 삭제하면 데이터셋은 깨끗해지지만, 남은 데이터가 너무 적어 모델의 사후 분포(Posterior)가 수렴하지 못하게 된다. 이를 방지하기 위해 Bernstein-von Mises 정리에 기반한 O(1/√n) 불확실성 스케일링을 보상 설계에 반영했다.

학습 안정성을 위해 두 가지 가드레일을 적용했다. 첫째, 동일한 액션 가족을 연속으로 적용하는 것을 제한하는 '반복 방지 가드', 둘째, 행 수가 10개 미만으로 떨어질 경우 에피소드를 강제 종료하는 '최소 행 수 체크'를 통해 강화학습의 발산을 막았다.

현재 프레임워크는 분류(Classification) 작업에 최적화되어 있으며, 회귀(Regression) 작업에 적용하기 위해서는 CRPS와 같은 별도의 보정 지표와 보상 재설정이 필요하다. 또한 단일 표 데이터셋을 기준으로 하며, 멀티 테이블 스키마에 대한 확장은 향후 과제로 남아 있다.