[24′ NeurIPS] ReMasker: Imputing Tabular Data with Masked Autoencoding

This entry is part 3 of 3 in the series Tabular Model

1. 들어가며

이번 글에서는 Meta에서 2024년에 발표한 ReMasker 논문을 리뷰합니다.

ReMasker는 테이블형 데이터셋을 대상으로 하는 결측치 대체(Imputation) 모델입니다. 데이터 임퓨테이션은 누락된 데이터 값을 추정하여 채우는 작업으로, 데이터를 분석하고 모델링하는 데 있어 중요한 전처리 단계입니다. 이미지 생성 분야의 경우, 이미지의 일부분을 마스킹하고 이를 채우는 이미지 인페인팅 기법이 유사한 문제로 자주 언급됩니다. 마찬가지로 테이블형 데이터에서도 결측치를 채워 데이터 완전성을 확보하는 임퓨테이션이 필요합니다.

최근 이미지 생성 분야에서는 디퓨전 모델이 주목받고 있으며, 이 모델을 사전 학습하여 인페인팅에 활용하거나 조건부 생성 모델로 변형하여 사용하는 방법이 연구되고 있습니다. 그러나 테이블형 데이터는 이미지와는 다른 특성을 가지므로, 이를 효과적으로 처리할 수 있는 전용 모델이 필요합니다.

이전에 연구된 TabSyn은 테이블형 데이터셋을 생성하는 디퓨전 모델을 제안하며, RePaint 방식의 임퓨테이션을 적용한 예입니다. 하지만 TabSyn은 주로 데이터셋 생성에 중점을 두었기 때문에 임퓨테이션 성능 면에서 추가적인 개선이 필요했습니다.

이번에 리뷰하는 ReMasker는 요즘 생성 모델의 트렌드인 디퓨전 모델을 사용하지 않고 MAE(Masked Autoencoder) 방식을 적용하여 테이블형 데이터 임퓨테이션에서 높은 성능을 보여줍니다.

2. 제안 방법

이번 장에서는 ReMasker 모델이 테이블형 데이터의 결측치 대체(Imputation) 문제를 어떻게 해결하는지 구체적으로 살펴보겠습니다.

2-1. Tabular Imputation Task

먼저, 테이블형 데이터의 결측치 대체 문제를 정의해 보겠습니다. 테이블형 데이터셋은 여러 개의 특성(Feature)으로 이루어져 있고, 일부 특성 값이 누락된 경우도 발생합니다. 임퓨테이션 문제는 이러한 누락된 데이터를 추정해 채우는 작업으로, 모델은 랜덤하게 마스킹된 결측값들을 예측하는 방식으로 학습합니다.

ReMasker에서는 기존 데이터셋에 존재하는 결측값뿐만 아니라 추가적인 마스크(Mask)를 더 적용하는데, 이를 ReMask라고 명명합니다. 마스크 비중을 높이는 이유는 모델이 더 많은 결측값을 복원하도록 학습시키고, 임퓨테이션 능력을 강화하기 위함입니다.

2-2. 모델 구조 (Architecture)

그림1. ReMasker Architecture
그림1. ReMasker Architecture

위 그림은 ReMasker 모델의 구조를 보여줍니다.

ReMasker 모델은 크게 학습(Fitting)추론(Imputation) 두 단계로 나누어집니다. 학습 단계에서는 마스킹된 데이터가 입력되고, 이를 복원(Reconstruction)하는 방식으로 모델이 학습됩니다. 추론 단계에서는 학습된 모델이 마스킹된 값을 복원하며 예측을 수행합니다.

그림2. Architecture 상세 분석
그림2. Architecture 상세 분석

2-3. 인코딩 (Encoding)

각 단계별로 좀 더 자세히 살펴보겠습니다. 먼저, 입력 데이터는 마스킹(M)된 테이블형 데이터셋입니다. 각 사각형은 테이블의 각 특성 값을 나타냅니다. ReMasker의 인코더와 디코더는 Transformer 구조를 기반으로 하기 때문에, 데이터를 인코딩하고 디코딩할 때 선형 변환(Linear Projection)을 통해 각 특성의 차원을 맞춰줍니다.

이후, 데이터 순서가 중요한 테이블형 데이터 특성에 맞게 위치 인코딩(Positional Encoding)이 추가됩니다. 이제 마스킹된 부분을 제외한 데이터가 Self-Attention 연산을 수행하며 특성 간의 관계를 학습합니다. 이렇게 인코딩된 벡터는 특성 간 상관관계를 반영한 새로운 벡터로 변환됩니다.

2-4. 디코딩 (Decoding)

이제 디코딩 과정입니다. 디코딩 단계에서는 복원해야 할 마스킹 값을 예측하는 것이 주된 목적입니다. 앞서 제거했던 마스킹 위치를 복원하여 모델이 원래 크기의 벡터를 생성하도록 합니다. 이 벡터는 디코딩용 Transformer 블록을 거쳐 Self-Attention 연산을 통해 주변 특성 값과의 관계를 반영하게 됩니다.

최종적으로 선형 변환을 통해 원래의 데이터 형태로 변환된 뒤, MSE(Mean Squared Error) 손실 함수를 사용하여 실제값과 유사해지도록 학습됩니다.

2-5. 종합

ReMasker는 테이블형 데이터셋을 위한 MAE(Masked Autoencoder)로, BERT 방식과 유사하게 설계되었습니다. Transformer의 Self-Attention 연산을 통해 데이터의 단기 및 장기 상관관계를 학습하고, 마스킹된 값을 복원하는 방식을 채택한 것입니다.

기존 TabSyn과 달리, ReMasker는 임퓨테이션 문제에 초점을 맞춰 설계되었기 때문에 결측값 대체 성능에서 더 높은 성능을 보입니다.

3. 실험

이제 ReMasker 모델이 실제로 얼마나 효과적인지 실험을 통해 확인해 보겠습니다. 먼저 기존 방법들과의 성능 비교를 살펴본 후, 다양한 구성 요소가 모델 성능에 미치는 영향을 분석한 Ablation Study 실험 결과를 검토합니다.

3-1. 성능 비교 (Performance)

ReMasker의 성능을 평가하기 위해 12개의 Regression Task 벤치마크 데이터셋을 사용하였습니다. 각 모델의 성능을 비교하기 위해 세 가지 평가 지표를 활용했으며, 그 결과는 아래와 같습니다.

그림3. 성능 비교 실험 결과
그림3. 성능 비교 실험 결과

사용된 평가 지표는 다음과 같습니다.

  1. RMSE (Root Mean Squared Error) : 복원된 값이 실제값(Ground Truth)과 얼마나 유사한지를 나타내는 지표입니다. RMSE 값이 낮을수록 더 정확하게 복원했음을 의미합니다.
  2. WD (Wasserstein Distance) : 복원된 값의 분포가 원래 데이터 분포와 얼마나 유사한지를 평가합니다. 이 값이 작을수록 기존 데이터와 분포가 비슷함을 의미합니다.
  3. AUROC (Area Under the Receiver Operating Characteristic Curve) : 임퓨테이션된 데이터셋으로 로지스틱 회귀 모델을 학습하여 평가한 성능 지표입니다. AUROC 값이 높을수록 결측치를 정확하게 복원하여 모델이 잘 학습된 결과를 의미합니다.

이 결과에서 ReMasker는 대부분의 평가 지표에서 기존 모델을 상회하는 우수한 성능을 보여줍니다. 이는 ReMasker가 특히 임퓨테이션에 최적화되어 있음을 보여주는 지표입니다.

3-2. Ablation Study

이번에는 ReMasker 모델의 구성 요소가 성능에 미치는 영향을 분석한 Ablation Study 결과를 살펴보겠습니다. 특히 Backbone Network를 Transformer, Linear, CNN으로 각각 변경하며 비교한 실험이 인상적입니다.

그림4. Backbone에 따른 성능 비교 실험 결과
그림4. Backbone에 따른 성능 비교 실험 결과

이 실험에서는 Transformer를 Backbone으로 사용했을 때 가장 뛰어난 성능을 보였습니다. 이는 Transformer가 각 특성 간의 상관관계를 잘 학습할 수 있는 구조이기 때문이며, 특히 마스킹을 효과적으로 처리할 수 있는 장점이 성능 차이로 이어진 것으로 보입니다. Linear나 CNN을 Backbone으로 사용했을 때에는 마스킹 처리와 연관 학습에서 Transformer만큼의 효율을 보이지 못했습니다.

4. 마치며

이번 글에서는 Meta에서 발표한 ReMasker 모델의 구조와 성능에 대해 알아보았습니다. ReMasker는 테이블형 데이터의 결측치를 대체하기 위해 설계된 모델로, MAE(Masked Autoencoder)를 기반으로 하여 결측치 복원 성능을 높였습니다. 기존의 Tabular Imputation 연구에서 주로 사용된 Diffusion Model을 대신하여 Transformer 기반 구조를 채택한 것이 특징적입니다.

ReMasker의 성능 실험 결과는, 테이블형 데이터의 임퓨테이션 문제에서 기존 모델을 능가하는 성과를 보여줍니다. 특히 Transformer를 Backbone으로 사용하여 각 특성 간의 상관관계를 학습하는 방식이 결측치를 정교하게 복원하는 데 큰 역할을 했습니다. 실험 결과에서도 볼 수 있듯이, ReMasker는 다양한 데이터셋과 지표에서 안정적이고 우수한 성능을 입증하였습니다.

ReMasker의 접근 방식은 테이블형 데이터셋의 특성을 고려한 최적화된 임퓨테이션 모델을 제시하며, 특히 결측 데이터가 많은 실무 환경에서 효율적이고 실용적인 솔루션이 될 수 있습니다. 앞으로도 테이블형 데이터의 특성에 맞는 임퓨테이션 모델이 발전하면서, 보다 정교한 데이터 분석과 예측이 가능해질 것으로 기대됩니다.

5. 참고 자료

  1. ReMasker 논문
  2. Diffusion Model에 관한 모든 것

Series Navigation<< [21′ AAAI] TabNet : Attentive Interpretable Tabular Learning
0 0 votes
Article Rating
Subscribe
Notify of
guest

0 Comments
Inline Feedbacks
View all comments
0
Would love your thoughts, please comment.x
()
x
Scroll to Top