📊 statistics

AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking

이 논문은 표준 스코어 기반 확산 모델이 불확실한 컨디셔닝 컨텍스트로서 결측치를 채우기 위한 확률적 증강을 사용하고 디노이징 감독을 관측된 좌표로만 제한함으로써, 불완전한 테이블 데이터를 효과적으로 생성할 수 있게 하는 플러그 앤 플레이(plug-and-play) 학습 프레임워크인 AugMask를 소개한다.

원저자: Jungkyu Kim, Taeyoung Park, Kibok Lee

게시일 2026-06-03

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jungkyu Kim, Taeyoung Park, Kibok Lee

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

다음은 AugMask 논문을 일상적인 비유를 사용하여 쉬운 개념으로 풀어낸 설명입니다.

문제점: 고장 난 레시피로 요리사를 가르치기

당신이 마스터 셰프(데이터 생성 모델인 Diffusion Model)에게 특정 요리(실제와 같은 표 형식의 데이터/tabular data 생성)를 하는 법을 가르치려 한다고 상상해 보세요. 셰프는 레시피를 따르는 데 매우 뛰어나지만, 한 가지 문제가 있습니다. 당신이 주는 레시피가 불완전하다는 것입니다. 어떤 재료는 빠져 있고, 그 자리는 그냥 빈칸이거나 "NaN"(숫자가 아님)으로 표시되어 있습니다.

표준적인 셰프(AI 모델)들은 빈 공간을 보면 혼란스러워합니다. 만약 레시피에 "설탕 2컵 [빈칸]"이라고 적혀 있다면, 그들은 요리를 할 수 없습니다.

기존 방식 (Zero-Filling): 사람들은 예전에 빈칸을 "0"이나 "물" 같은 자리 표시자로 채우곤 했습니다. 하지만 이는 셰프에게 "설탕 0컵을 넣으세요"라고 말하는 것과 같습니다. 그러면 셰프는 "0"이 실제 식재료라고 학습하게 되어, 결과적으로 맛(데이터 분포)을 망치게 됩니다.
"추측하기" 방식 (Imputation): 또 다른 방법은 셰프에게 빠진 재료를 추측하게 하고, 그 추측이 완벽해지도록 셰프를 훈련시키는 것입니다. 하지만 셰프가 추측을 틀리면, 잘못된 추측을 완벽하게 만들기 위해 계속 노력하게 되어 결국 노이즈가 가득하고 혼란스러운 주방이 되고 맙니다.

해결책: AugMask

저자들은 똑똑한 부주방장 역할을 하는 새로운 훈련 전략인 AugMask를 제안합니다. 이 방식은 **장면 설정하기(Setting the Scene)**와 **채점하기(Grading the Work)**라는 두 가지 업무를 분리함으로써 문제를 해결합니다.

1. 장면 설정하기 (Stochastic Augmentation)

부주방장(가벼운 보조 모델)은 빈 공간을 단순히 비워두거나 지루한 "0"으로 채우는 대신, 그럴듯한 추측값으로 채웁니다.

반전: 부주방장은 단 하나의 추측만 내놓는 것이 아닙니다. 그것은 하나의 범위를 제공합니다.
- 비유: 레시피에 "소금"이 빠져 있다면, 부주방장은 단순히 "1 티스푼"이라고 말하지 않습니다. 대신 "0.5에서 1.5 티스푼 사이일 수 있습니다"라고 말합니다.
왜 그럴까요? 이는 메인 셰프에게 빠진 재료가 불확실함을 가르쳐 주기 위해서입니다. 메인 셰프는 전체적인 그림(다른 재료들)을 보고, 빠진 부분이 확정적인 값이 아니라 "아마도 그럴 수도 있는 것"임을 이해하는 법을 배웁니다.

2. 채점하기 (Masking the Loss)

이 부분이 가장 중요합니다. 메인 셰프가 요리를 재현하려고 할 때, 선생님(훈련 알고리즘)은 원래 그 자리에 있었던 재료들만 채점합니다.

규칙: 만약 레시피에 원래 "밀가루 2컵"이 있었다면, 선생님은 셰프가 밀가루를 제대로 만들었는지 확인합니다.
예외 조항: 만약 레시피에 원래 "소금" 자리가 비어 있었다면, 선생님은 셰프가 소금에 대해 무엇을 추측했든 무시합니다. 선생님은 이렇게 말합니다. "소금에 대해 무엇을 추측했는지는 상관없다. 그저 밀가루를 제대로 만드는 데 집중해라."
결과: 셰프는 "추측값"을 맥락(단서)으로 사용하여 실제 재료들을 제대로 만드는 법을 배우지만, 자신의 추측이 틀렸다고 해서 벌점을 받지는 않습니다. 이를 통해 셰프는 추측값을 암기하려 하지 않고, 실제 데이터 간의 관계를 배우기 시작합니다.

왜 이것이 작동하는가: "불확실성 페널티"

논문은 이 방식이 왜 작동하는지를 **"떨리는 손"**이라는 간단한 비유를 통해 수학적으로 설명합니다.

결정론적 추측 (나쁜 예): 만약 부주방장이 단 하나의 확신에 찬 추측(예: "정확히 1.0 티스푼입니다")을 준다면, 메인 셰프는 "좋아, 나는 이것과 정확히 일치시켜야 해"라고 생각합니다. 만약 그 추측이 조금이라도 틀리면, 셰프는 혼란에 빠집니다.
확률적 추측 (좋은 예): 만약 부주방장이 "0.5에서 1.5 사이입니다"라고 말한다면, 메인 셰프는 "이것은 흔들리고 불확실한 단서구나"라고 깨닫습니다.
페널티: 수학적으로 볼 때, 단서가 흔들릴 때(불확실성이 높을 때), 메인 셰프는 자연스럽게 그 단서에 너무 많이 의존하지 않게 됩니다. 이는 마치 학생이 시험 문제에서 흐릿한 힌트를 보고, 그것이 틀릴 수도 있다는 것을 알기에 그 힌트를 무시하는 것과 같습니다. 이는 모델이 잘못된 추측에 "갇히는" 것을 방지합니다.

결과: 더 나은 셰프

저자들은 다양한 양의 결측치(10%에서 90%까지)가 포함된 여러 데이터셋(성인 소득 데이터, 쇼핑 습관 등)에 대해 AugMask를 테스트했습니다.

결과: 이 "장면 설정 후 추측 무시하기" 전략을 사용함으로써, 표준 AI 모델(보통 결측치를 처리하는 데 어려움을 겪는 모델들)은 결측치 처리를 위해 특별히 설계된 모델들보다 더 뛰어난 성능을 보였습니다.
핵심 요약: 불완전한 레시피를 위해 특별하고 복잡한 주방을 만들 필요는 없습니다. 빈칸을 "불확실한 추측"으로 채우고, 마지막 요리를 채점할 때 그 추측은 무시하라고 셰프에게 알려주는 스마트한 방법만 있으면 됩니다.

한 문장 요약

AugMask는 AI 모델이 결측치를 처리할 때, 맥락을 제공하기 위해 빈 공간을 "불확실한 추측"으로 채우되, 채점할 때는 실제 데이터에 대해서만 엄격하게 평가함으로써 모델이 자신의 추측 때문에 혼란에 빠지는 것을 방지하도록 가르칩니다.

기술 요약: AugMask

문제 정의

점수 기반 확산 모델(Score-based diffusion models)은 강력한 딥 생성 모델로 부상했으나, 결측치가 포함된 테이블 데이터에 적용할 때 상당한 어려움에 직면합니다. 표준 확산 아키텍처는 완전히 지정된 입력을 가정하는 반면, 실제 세계의 테이블 데이터에는 종종 결측치(NaNs)가 포함되어 있습니다.

이 간극을 메우기 위한 기존 접근 방식은 크게 두 가지 범주로 나뉘며, 두 방식 모두 한계가 있습니다:

상수 플레이스홀더 임퓨테이션(Constant Placeholder Imputation): 학습 전 결측치를 고정된 상수(예: 0, 컬럼 평균)로 대체하는 방식입니다. 이는 모델이 유효한 0과 결측치를 동일하게 취급하도록 강제하여 희소성 편향(sparsity bias)을 유발하고, 예측이 결측 수준에 의존하게 만듭니다. 또한, 손실 함수를 관측된 좌표에 대해서만 마스킹하더라도, 결정론적인 플레이스홀더가 입력에 남아 있어 결측성과 상관관계가 있는 가짜 단서(spurious cue)로 작용합니다.
자기 임퓨테이션(Self-Imputation, 예: DiffPuter): 학습 중에 기대-최대화(EM) 접근 방식을 통해 결측치를 반복적으로 업데이트합니다. 이는 고정된 플레이스홀더 문제는 피할 수 있지만, 모델이 자신의 진화하는 예측치를 바탕으로 학습하게 되는 "움직이는 타겟(moving target)" 최적화 문제를 야기합니다. 이는 노이즈가 섞인 그래디언트, 비정상적 동역학(non-stationary dynamics), 그리고 지나치게 분산된 완성(completions)을 초 de할 수 있습니다. 확산 목적 함수는 점 단위의 정확도보다 다양성을 우선시하기 때문입니다.

핵심 과제는 디노이저(denoiser)에게 안정적인 수치 입력을 제공하면서도, 채워진 결측치를 정답(ground-truth) 타겟으로 취급하지 않음으로써 상수 플레이스홀더의 편향과 자기 임퓨테이션의 불안정성을 모두 피하는 것입니다.

방법론: AugMask

저자들은 AugMask를 제안합니다. 이는 결측을 인지하지 못하는 백본(missing-unaware backbones)을 불완전한 데이터에 적응시키기 위해 **컨디셔닝(conditioning)**과 **감독(supervision)**을 분리하는 플러그 앤 플레이 방식의 학습 프레임워크입니다. 이 방법은 두 가지 뚜렷한 단계로 작동합니다.

1. 확률적 조건부 증강 (입력 구성)

결정론적인 플레이스홀더를 사용하는 대신, AugMask는 컨텍스트 의존적 분포 $T(\cdot | x_{obs}, m)$ 에서 결측 항목을 샘플링하여 완전한 수치 입력을 구성합니다.

메커니즘: 각 결측 특성 $j$ 에 대해, 관측된 특성들과 결측 마스크에 조건화된 분포로부터 값 $Z_j$ 를 추출합니다.
구현: 저자들은 각 특성별 LightGBM 모델을 사용하는 경량화된 팩터화(factorized) 접근 방식을 활용합니다. 연속형 특성의 경우, 헤테로스케다스틱 가우시안(heteroskedastic Gaussian)으로부터 샘플링하기 위해 조건부 평균 회귀 모델과 로그-분산 회귀 모델을 적합시킵니다. 범주형 특성의 경우, 레이블 스무딩(label smoothing)이 적용된 조건부 분류기를 사용합니다.
역할: 이 샘플링된 값들은 디노이저를 위한 불확실한 컨디셔닝 컨텍스트(uncertain conditioning context) 역할을 수행하며, 정답 타겟이 아닙니다. 이러한 확률성은 결측의 내재적 불확실성을 인코딩합니다.

2. 관측값 전용 감독 (손실 마스킹)

디노이징 손실은 원래 관측되었던 좌표에 대해서만 적용됩니다.

목적: 모델은 관측된 좌표에 대한 재구성 오차 $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$ 를 최소화합니다. 여기서 $x_{A,0}$ 는 증강된 입력입니다.
효과: 증강된(이전에 결측되었던) 항목들은 컨텍스트를 제공하기 위해 노이즈가 섞인 입력 $x_{A,t}$ 를 통해 모델에 들어오지만, 모델은 이를 재구성하지 못하더라도 페널티를 받지 않습니다. 이는 모델이 잠재적으로 부정확한 완성치에 과도하게 의존하는 것을 방지합니다.

이론적 근거

논문은 이 학습 규칙을 **라오-블랙웰화된 목적 함수(Rao–Blackwellized objective)**와 연결합니다. 가능한 완성치들에 대해 주변화(marginalizing)함으로써, 저자들은 이 목적 함수가 **분산 가중 민감도 페널티(variance-weighted sensitivity penalty)**를 산출함을 보여줍니다.

해석: 목적 함수의 국소 근사(local approximation)에서 2차 항은 증강의 조건부 분산에 비례합니다. 높은 조건부 분산은 특정 완성치에 대한 모델의 민감도를 억제하여, 불확실한 채움에 대한 의존을 효과적으로 방지합니다. 낮은 분산은 해당 완성치가 컨텍스트로서 더 직접적으로 사용될 수 있게 합니다.
AugFull과의 대비: 저자들은 AugMask를 "AugFull"(증강된 좌표를 포함한 모든 좌표에 손실을 적용하는 방식)과 비교합니다. 그들은 AugFull이 증강 규칙이 잘못 지정될 경우 편향을 유발할 수 있음을 입증했습니다. 즉, 모델이 잠재적으로 틀린 의사 타겟(pseudo-targets)에 맞추도록 강제되기 때문입니다. AugMask는 완성을 엄격히 컨텍스트로 취급함으로써 이를 피합니다.

주요 기여

불완전한 테이블 확산을 위한 학습 원칙: 컨디셔닝(확률적 증강을 통한)과 감독(관측값 전용 손실 마스킹을 통한)을 분리하는 AugMask를 도입했습니다.
주변화 관점: 확률적 증강이 분산 가중 민감도 페널티를 유도하여, 불확실한 완성치에 대한 과도한 의존을 방지하도록 모델을 규제한다는 이론적 분석을 제시했습니다.
플러그 앤 플레이 레시피: 구조적 변경 없이도 기존의 결측 미인지(missing-unaware) 점수 기반 확산 모델이 불완전한 데이터에서 학습할 수 있게 하는 경량 프레임워크를 제공합니다.

실험 결과

저자들은 다양한 데이터셋(Adult, Default, Shoppers, Magic, Beijing, News)과 다양한 결측 레짐(MCAR 및 MAR) 및 0.1에서 0.9 사이의 결측 비율에 대해 AugMask를 평가했습니다.

성능: AugMask는 표준 결측 미인지 확산 모델(예: CDTD, TabDiff)이 특화된 결측 인지 베이스라인(예: MIWAE, MissDiff, DiffPuter, ForestDiff)보다 우수한 성능을 내도록 합니다.
생성 품질: AugMask로 학습된 모델은 우수한 샘플 수준 충실도( $\alpha$ -Precision, $\beta$ -Recall)와 특성 수준 통계량(Trend, Shape)을 달성했습니다.
머신러닝 유용성: AugMask 모델이 생성한 합성 데이터로 학습된 다운스트림 태스크(분류 및 회귀)는 베이스라인에 비해 향상된 성능을 보였습니다.
강건성: AugMask의 이점은 결측 비율이 높아질수록 더욱 두드러졌습니다. 높은 결측률(예: 90%)에서도 AugMask는 다른 방법론에서 관찰되는 심각한 성능 저하를 피했습니다.
민감도 분석: 이 방법은 증강의 잘못된 지정(misspecification)에 대해 강건함을 입증했습니다. 증강 규칙이 변형되더라도(예: 확률성을 줄이거나 편향을 도입함), AugMask는 AugFull보다 성능 우위를 유지했습니다. 이는 완성을 타겟이 아닌 컨텍스트로 취급하는 것이 핵심 요소임을 확인시켜 줍니다.

의의 및 주장

논문은 AugMask가 완전히 지정된 확산 백본과 불완전한 실제 데이터 사이의 불일치에 대한 원칙적인 해결책을 제공한다고 주장합니다. 확률적 완성을 타겟이 아닌 불확실한 컨텍스트로 취급함으로써, 이 방법은 다음을 달성합니다:

상수 플레이스홀더가 도입하는 희소성 편향과 인위적인 패턴을 방지합니다.
자기 임퓨테이션 방식의 최적화 불안정성과 움직이는 타겟 역학을 방지합니다.
증강의 분산을 활용하여 모델을 규제함으로써, 불확실한 입력에 대한 민감도를 낮춥니다.

저자들은 AugMask를 복잡한 아키텍처 수정이나 특화된 손실 함수 없이도 기존의 고성능 결측 미인지 모델을 결측 데이터 환경에 배치할 수 있게 하는 일반적인 전략으로 포지셔닝합니다. 본 연구는 생성에 초점을 맞추고 있으며, 이를 잠재 확산 파이프라인(latent diffusion pipelines)이나 임퓨테이션 특화 목적 함수로 확장하는 것은 향후 연구 과제로 남겨두었습니다.