TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer

기술 요약: TabTreeFormer

문제 정의

Transformer는 정형 데이터 생성에서 인상적인 결과를 보여주었으나, 정형 데이터의 본질적인 특성을 보존하는 데 필수적인 도메인 특화 유도 편향(inductive biases)이 부족한 경우가 많습니다. 구체적으로, 표준 Transformer는 다음과 같은 문제로 어려움을 겪습니다:

1. 비매끄러움(Non-smoothness) 및 낮은 상관관계: 정형 데이터는 이산적 특징(discrete features)과 이산적 특징 및 연속적 특징 사이의 비매끄러운 관계를 포함하는 경우가 많습니다. 신경망은 일반적으로 매끄러운 저주파 함수를 학습하는 반면, 트리 기반 모델은 조각마다 일정한(piecewise constant) 함수를 포착하는 데 뛰어나며, 회전 불변성이 없는 특성 덕분에 상관관계가 낮은 특징에 대해서도 강건합니다.
2. 다봉 분포(Multimodal Distributions): 기존 생성 모델은 연속적 특징(여러 개의 모드나 피크를 가진 특징)에서 다봉 분포를 정확하게 포착하지 못하는 경우가 많습니다.
3. 확장성 및 효율성: Transformer는 이차 복잡도(quadratic complexity)로 인해 확장성이 떨어집니다. 또한, 정형 데이터를 토큰화할 때 큰 어휘 사전(vocabulary)과 긴 시퀀스 길이를 요구하여 효율성을 저하시킵니다.
4. 개인정보 보호 및 암기: Transformer는 특히 작은 규모의 정형 데이터셋으로 훈련될 때 훈련 데이터를 암기하는 경향이 있어 개인정보 보호 위험을 초래합니다.

방법론: TabTreeFormer

저자들은 정형 데이터 특유의 유도 편향을 주입하기 위해 트리 기반 모델과 Transformer, 그리고 새로운 토크나이저를 결합한 하이브리드 아키텍처인 TabTreeFormer를 제안합니다.

1. 트리 기반 모델 통합

유도 편향을 주입하기 위해, 모델은 사전 학습된 트리 기반 모델(예: LightGBM)을 포함합니다.

• 리프 인덱스 프롬프트(Leaf Indices as Prompts): 각 행에 대해 모델은 $T$개의 트리로부터 리프 인덱스를 추출합니다. 이 인덱스들은 입력 토큰 시퀀스의 앞에 추가됩니다.
• 유도 편향 전이: 이 리프 인덱스들은 비매끄럽고 회전 불변성이 없는 구조를 인코딩하며, 이를 통해 트리 모델의 유도 편향을 Transformer로 효과적으로 전달합니다. 추론 시, 이 인덱스들은 현실적인 데이터를 생성하도록 가이드하는 프롬프트 역할을 합니다.
• 클러스터링: 리프 인덱스는 생성의 다중 동시 조건(클러스터) 역할을 수행하여, 단일 컬럼을 넘어선 조건부 생성 개념을 확장합니다.

2. 이중 양자화 토크나이저(Dual-Quantization Tokenizer)

다봉 분포를 처리하고 어휘 사전 크기와 시퀀스 길이를 줄이기 위해, 연속적 특징을 위한 새로운 이중 양자화 토크나이저가 도입되었습니다.

• 클러스터 양자화(K-Means): 첫 번째 양자화기는 K-Means 클러스터링을 사용하여 다봉 분포를 모델링하고, 각 값에 "Bin ID"를 할당합니다.
• 분위수 양자화(Quantile Quantization): 두 번째 양자화기는 분위수 기반 접근 방식을 사용하여 정밀한 수치 값을 포착하고, "Quantile ID"를 할당합니다.
• 토큰 표현: 각 수치 값은 두 개의 이산 토큰(Bin ID + Quantile ID)으로 표현되는 반면, 범주형 특징은 단일 레이블 인코딩 토큰을 사용합니다. 이 방식은 표준 토큰화 방법과 비교하여 시퀀스 길이와 어휘 사전을 크게 줄여줍니다.

3. Transformer 아키텍처 및 훈련

핵심 생성기는 자기회귀(auto-regressive) Transformer(Distill-GPT2 기반)입니다.

• 서수 임베딩(Ordinal Embeddings): 분위수 토큰은 서수 관계(ID가 가까울수록 값이 가까움)를 가지므로, 저자들은 일반적인 위치 또는 학습된 임베딩 대신 비선형 단조 스케일드 시그모이드 함수를 사용하는 함수 기반 서수 임베딩을 설계했습니다.
• 서수 교차 엔트로피 손실(Ordinal Cross-Entropy Loss, OCEL): 표준 교차 엔트로피는 모든 클래스를 동일하게 구분합니다. 저자들은 분위수 토큰의 서수적 특성을 존중하여, 예측값이 타겟 클래스에서 멀어질수록 더 큰 벌점을 부여하는 가중치 적용 OCEL을 제안합니다.
• 개인정보 보호: 암기를 방지하기 위해 모델은 강력한 입력 토큰 마스킹을 적용하고 검증 손실(validation loss) 기반의 조기 종료 메커니즘을 활용합니다. 훈련 세트는 두 개의 하위 집합으로 나뉘며, 각각 별도의 네트워크가 훈련됩니다. 생성 시에는 과적합 없이 분포 커버리지를 확보하기 위해 두 네트워크에서 샘플을 추출합니다.

주요 기여

1. 트리 사전(Tree Priors)의 최초 통합: 저자들이 아는 바에 따르면, 이는 생성 작업을 위해 트리 기반 모델을 사용하여 정형 특화 유도 편향(비매끄러움, 회전 불변성 없음)을 Transformer에 도입한 첫 번째 연구입니다.
2. 이중 양자화 토크나이저: 고정밀도 및 압축된 표현을 유지하면서 다봉 연속 분포를 모델링할 수 있는 새로운 토크나이저를 제안하였으며, 이를 서수 인식 임베딩 및 손실 함수와 결합했습니다.
3. 포괄적인 평가: 모델은 9개의 다양한 데이터셋에 대해 8개의 최신 기술(SOTA) 베이스라인(GAN, VAE, 확산 모델 및 기타 Transformer 포함)과 함께 평가되었으며, 유용성(utility), 충실도(fidelity), 개인정보 보호(privacy), 효율성(efficiency) 측면에서 우수한 균형을 입증했습니다.

실험 결과

9개의 데이터셋(예: Adult, Bank, Boston, Credit)과 8개의 베이스라인을 대상으로 평가되었습니다:

• 유용성(Utility): TabTreeFormer는 다운스트림 머신러 러닝 효용성(MLE) 작업(로지스틱 회귀, 랜덤 포레스트, XGBoost)에서 일관되게 베이스라인을 능가했습니다. 가장 우수한 변체인 TabTreeFormer-NM(마스킹 없음)은 데이터 유용성을 우선시할 때 최상의 베이스라인 대비 최대 44%의 성능 향상을 달성했습니다.
• 충실도(Fidelity): 모델은 베이스라인과 대등한 "Shape" 점수(주변 분포 유사도)를 기록했으며, 특히 다봉 분포와 특징 간 관계를 포착하는 데 있어 우수한 "Trend" 점수를 보였습니다.
• 개인정보 보호(Privacy): 마스킹을 적용한 변체(TabTreeFormer-S 및 L)는 강력한 개인정보 보호 탄력성을 보여주었습니다. 높은 암기 위험을 보인 여러 베이스라인(예: TabuLa, Forest Diffusion)과 달리 유의미한 개인정보 유출이 감지되지 않았습니다.
• 효율성(Efficiency): 하이브리드 구조임에도 불구하고, TabTreeFormer는 80M 개 이상의 파라미터를 가진 베이스라인들과 비교하여 훨씬 작은 모델 크기(5M–40M 파라미터)로 더 빠른 생성 시간과 대등한 성능을 달성했습니다.

의의

본 논문은 TabTreeFormer가 정형 예측에서의 트리 기반 모델의 성공과 생성 모델링에서의 유도 편향 저활용 사이의 근본적인 간극을 해결한다고 주장합니다. 트리 기반 사전 지식과 Transformer 아키텍처를 효과적으로 결합하고, 서수 데이터에 특화된 토큰화 및 손실 함수를 도입함으로써, TabTreeFormer는 고품질의 합성 정형 데이터를 생성하기 위한 강력한 솔루션을 제공합니다. 이는 도메인 특화 지식을 통합하는 것이 특히 데이터 유용성이 중요한 시나리오에서 생성 성능을 크게 향상시킬 수 있음을 입증하며, 동시에 경쟁력 있는 효율성과 개인정보 보호 보장을 유지합니다.