결측값 없는 트랜스포머 기반 테이블 데이터 학습

초록

본 논문은 결측값을 별도로 보정하거나 초기화하지 않고, 두 개의 마스크를 활용해 트랜스포머에 직접 적용함으로써 테이블형 데이터의 분류 성능을 크게 향상시키는 Imputation‑Free Incremental Attention Learning (IFIAL) 방법을 제안한다. 특징을 겹치는 고정 크기 파티션으로 나누어 점진적으로 학습함으로써 메모리 사용을 최적화하고, 17개의 데이터셋에서 11개의 최신 방법을 능가한다.

상세 분석

IFIAL은 기존 결측값 처리 방식(삭제, 평균·중앙값 대체, MICE 등)과 달리 결측값 자체를 모델 입력에 포함시키되, 트랜스포머의 어텐션 연산 단계에서만 제외한다는 점이 핵심이다. 이를 위해 두 종류의 어텐션 마스크를 설계한다. 첫 번째 마스크 M₁은 열 방향으로 결측 피처에 해당하는 열 전체에 −∞ 값을 부여해 소프트맥스 연산 시 0으로 변환한다. 두 번째 마스크 M₂는 행 방향으로 결측 피처에 해당하는 행을 0으로 만들어 해당 피처의 출력이 완전히 차단되도록 한다. 이렇게 하면 Q·Kᵀ 행렬에 결측값이 포함된 위치는 어텐션 스코어 계산에서 실질적으로 무시된다.

또한 IFIAL은 Feature‑Tokenized Transformer(FTT)를 기반으로 하며, 피처 이름과 범주형 값을 사전학습된 언어 모델 임베딩으로, 연속형 값은 선형 투영을 통해 임베딩한다. 이러한 임베딩은 이후 gated transformer 레이어에서 미세조정된다.

특징 파티셔닝 전략은 전체 피처 집합을 크기 k 인 고정 파티션으로 나누고, 인접 파티션 간에 ⌈k/2⌉ 만큼 겹치게 설계한다. 겹침을 통해 이전 파티션에서 학습된 가중치를 새로운 파티션에 자연스럽게 전달할 수 있다. 파티션 수 P 는 식 P = 1 + ⌈(d−k)/(k−⌈k/2⌉)⌉ 로 정의되며, d는 전체 피처 수이다. 이 방식은 O(n²) 복잡도를 갖는 전통적인 트랜스포머의 메모리 요구량을 파티션당 O(k²)로 낮추어 대규모 피처 공간에서도 GPU 메모리 초과 없이 학습이 가능하도록 만든다.

실험에서는 17개의 이질적인 테이블 데이터셋(의료, 금융, 생물학 등)과 4가지 결측 메커니즘(MCAR, MAR, MNAR, 혼합) 및 다양한 결측 비율(10%~70%)을 적용하였다. IFIAL은 평균 순위에서 11개의 베이스라인(전통 ML, XGBoost, Random Forest, LightGBM, DeepImpute, GAIN, MissForest 등)을 모두 앞섰으며, 특히 결측 비율이 높을수록 성능 격차가 확대되는 경향을 보였다. 파티션 크기를 전체 피처의 절반(k = d/2)으로 설정했을 때 정확도와 학습 시간 사이의 최적 균형을 달성했으며, 더 작은 파티션은 메모리 절감은 되지만 성능 저하가 발생하고, 큰 파티션은 메모리 초과 위험이 있었다.

또한, 마스크 기반 접근법이 결측값 자체를 완전히 무시하면서도 관측값만을 활용해 충분히 풍부한 표현을 학습한다는 점을 시각화된 어텐션 히트맵을 통해 검증하였다. 이는 기존 방법이 결측값을 임의값으로 채워 넣음으로써 발생할 수 있는 피처 중요도 왜곡을 방지한다는 의미이다.

한계점으로는 현재 구현이 완전히 순차적 파티션 학습에 의존하므로, 파티션 순서가 성능에 미치는 영향을 완전히 해소하지 못했다는 점과, 매우 높은 차원의 범주형 피처(수천 개)에서는 임베딩 차원 축소가 추가로 필요할 수 있다는 점을 언급한다. 향후 연구에서는 파티션 순서 최적화와 동적 파티션 크기 조정, 그리고 마스크와 결합된 사전학습된 결측 인코더를 탐색할 계획이다.