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

왜 중요한가

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

핵심 기여

L2C2 프레임워크 제안

표 형식 데이터 클리닝을 '사전 분포 정렬' 문제로 정의하고, 심층 강화학습(PPO)을 사용하여 최적의 클리닝 연산자 시퀀스를 탐색하는 최초의 프레임워크를 개발했다.

TFMAwareReward 보상 함수 설계

단순 정확도뿐만 아니라 데이터 보존율과 분포 왜곡을 고려한 다목적 보상 함수를 설계하여, 행 삭제를 남발하는 퇴행적 전략을 방지하고 인컨텍스트 러닝의 효율을 극대화했다.

매개변수화된 액션 공간 도입

단순히 연산자를 선택하는 것을 넘어 KNN의 K값이나 이상치 임계값 등 세부 파라미터를 에이전트가 직접 결정하도록 하여 데이터셋별 맞춤형 정제 성능을 확보했다.

도메인 간 전이 학습 입증

특정 데이터셋에서 학습된 클리닝 정책이 처음 보는 다른 도메인의 데이터셋에서도 유효하게 작동하며, 처음부터 학습하는 것보다 60% 이상 빠른 수렴 속도를 보임을 확인했다.

핵심 아이디어 이해하기

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와 같은 별도의 보정 지표와 보상 재설정이 필요하다. 또한 단일 표 데이터셋을 기준으로 하며, 멀티 테이블 스키마에 대한 확장은 향후 과제로 남아 있다.

실무 활용

TabPFN과 같은 최신 표 형식 AI 모델을 실무에 도입할 때, 데이터 정제 과정을 자동화하고 모델 성능을 극대화하는 데 즉시 활용 가능하다.

데이터 과학자의 수동 클리닝 시간을 단축하기 위한 자동 파이프라인 구축
의료나 금융처럼 모델의 예측 정확도뿐만 아니라 신뢰도(Calibration)가 중요한 도메인의 데이터 정제
데이터 양이 적어 파운데이션 모델의 제로샷 성능에 의존해야 하는 소규모 데이터셋 분석 프로젝트

코드 공개 여부: 공개

코드 저장소 보기

키워드

TFM(표 형식 파운데이션 모델)RL(강화학습)Data Cleaning(데이터 클리닝)Prior Alignment(사전 분포 정렬)TabPFN(탭피에프엔)

용어 해설

표 형식 파운데이션 모델(Tabular Foundation Model): — 대규모 합성 데이터나 다양한 데이터셋에서 사전 학습되어, 별도의 추가 학습 없이도 새로운 표 형식 데이터에 대해 즉각적인 예측을 수행할 수 있는 AI 모델이다. 적은 양의 데이터로도 높은 성능을 내지만 입력 데이터의 분포 변화에 민감하다는 특징이 있다.
사전 분포 정렬(Prior Alignment): — 실제 입력 데이터의 통계적 분포를 모델이 학습 시 가정했던 가상의 데이터 분포(Prior)와 일치시키는 과정이다. 데이터 클리닝을 통해 이 간극을 줄임으로써 모델의 예측 정확도와 신뢰도를 동시에 회복하는 것이 목적이다.
인컨텍스트 러닝(In-Context Learning): — 모델의 가중치를 업데이트하지 않고 입력 프롬프트에 포함된 예시 데이터만을 활용하여 새로운 작업을 수행하는 방식이다. 표 형식 모델에서는 입력된 행 데이터들이 모델의 예측을 위한 문맥(Context) 역할을 수행한다.
기대 보정 오차(Expected Calibration Error): — 모델이 예측한 확률(신뢰도)과 실제 정확도 사이의 차이를 측정하는 지표이다. 이 수치가 낮을수록 모델이 자신의 예측이 맞을 확률을 더 정확하게 판단하고 있음을 의미하며, 고위험 의사결정 분야에서 매우 중요하다.

코드 예제

python

class DataQualityObserver:
    def compute_state(self, dataset):
        # 9-dimensional state vector
        r_miss = mean_missing_rate(dataset)
        w1 = mean_wasserstein_drift(dataset)
        skew = mean_absolute_skewness(dataset)
        kurt = mean_absolute_kurtosis(dataset)
        delta_bal = class_balance_ratio(dataset)
        r_ret = row_retention_ratio(dataset)
        # action history flags
        h_imp, h_out, h_scl = self.get_action_history()
        return [r_miss, w1, skew, kurt, delta_bal, r_ret, h_imp, h_out, h_scl]

데이터셋의 품질 상태를 9차원 벡터로 변환하여 RL 에이전트의 관측값으로 제공하는 로직

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드