技術要約: TabTreeFormer

問題提起

Transformerは表形式データ（tabular data）の生成において目覚ましい成果を示しているが、表形式データ固有の特性を保持するために不可欠な、ドメイン特有の帰納バイアス（inductive biases）が欠如していることが多い。具体的には、標準的なTransformerは以下の問題に直面する：

1. 非平滑性と低相関: 表形式データには離散的な特徴量や、離散的な特徴量と連続的な特徴量の間の非平滑な関係が含まれることが多い。ニューラルネットワークは通常、平滑で低周波な関数を学習する傾向があるが、決定木モデルは区分定数関数（piecewise constant functions）のキャプチャに長けており、非回転不変性（non-rotational invariance）により低相関な特徴に対しても堅牢である。
2. マルチモーダル分布: 既存の生成モデルは、連続的な特徴量におけるマルチモーダルな分布（複数のモードやピークを持つ特徴量）を正確に捉えることができない場合が多い。
3. スケーラビリティと効率性: Transformerは二次的な計算複雑性を伴うため、スケーラビリティに乏しい。さらに、表形式データをトークン化する際に、大きな語彙数や長いシーケンス長を必要とすることが多く、効率を低下させる。
4. プライバシーと記憶（Memorization）: Transformerは、特に小規模な表形式データセットで学習する場合、訓練データを記憶しやすい傾向があり、プライバシーのリスクをもたらす。

手法: TabTreeFormer

著者らは、表形式特有の帰納バイアスを注入するために、ツリーベースのモデルとTransformer、および新しいトークナイザを統合したハイブリッドアーキテクチャであるTabTreeFormerを提案している。

1. ツリーベースモデルの統合

帰納バイアスを注入するため、モデルは事前学習済みのツリーベースモデル（例：LightGBM）を組み込んでいる。

• リーフインデックスをプロンプトとして利用: データセットの各行に対して、$T$個の決定木のリーフインデックスを抽出する。これらのインデックスは、入力トークンシーケンスの先頭に付加される。
• 帰納バイアスの転移: これらのリーフインデックスは、非平滑かつ非回転不変な構造をエンコードしており、これによりツリーモデルの帰納バイアスをTransformerへと効果的に転移させる。推論時、これらのインデックスは現実的なデータを生成するためのガイドとして機能するプロンプトとなる。
• クラスタリング: リーフインデックスは、生成のための複数の同時並行的な条件（クラスター）として機能し、単一のカラムを超えた条件付き生成の概念を拡張する。

2. 二重量子化トークナイザ（Dual-Quantization Tokenizer）

マルチモーダル分布を扱い、語彙数とシーケンス長を削減するために、連続的な特徴量に対して新しい二重量子化トークナイザを導入する。

• クラスター量子化 (K-Means): 第一の量子化器は、マルチモーダルな分布をモデル化するためにK-Meansクラスタリングを使用し、各値に「Bin ID」を割り当てる。
• 分位点量子化 (Quantile Quantization): 第二の量子化器は、精密な数値を捉えるために分位点ベースのアプローチを採用し、「Quantile ID」を割り当てる。
• トークン表現: 各数値は2つの離散トークン（Bin ID + Quantile ID）によって表現され、カテゴリカル特徴量は単一のラベルエンコードされたトークンを使用する。この手法は、標準的なトークン化手法と比較して、シーケンス長と語彙サイズを大幅に削減する。

3. Transformerアーキテクチャと学習

コアとなる生成器は、自己回帰型Transformer（Distill-GPT2に基づく）である。

• 順序埋め込み (Ordinal Embeddings): 分位点トークンは順序関係（近いIDは近い値を示す）を持つため、著者らは標準的な位置エンコーディングや学習可能な埋め込みではなく、非線形な単調スケーリング・シグモイド関数を用いた関数ベースの順序埋め込みを設計した。
• 順序交差エントロピー損失 (OCEL): 標準的な交差エントロピーはすべてのクラスを等しく区別するが、著者らは、分位点トークンの順序性を尊重し、ターゲットクラスに近い予測へのペナルティを遠い予測よりも小さくする、重み付けされたOCELを提案している。
• プライバシー保護: 記憶を防ぐため、モデルは入力トークンの強力なマスキングを行い、検証損失に基づく早期終了メカニズムを利用する。訓練セットは2つのサブセットに分割され、それぞれに別のネットワークが学習される。生成時には、過学習を避けつつ分布の網羅性を確保するために、両方のネットワークからサンプルが抽出される。

主な貢献

1. ツリー事前分布の初の統合: 著者らの知る限り、生成タスクにおいて、ツリーベースのモデルを用いて表形式特有の帰納バイアス（非平滑性、非回転不変性）をTransformerに導入した最初の研究である。
2. 二重量子化トークナイザ: 高い精度とコンパクトな表現を維持しながら、マルチモーダルな連続分布をモデル化できる新しいトークナイザであり、順序を考慮した埋め込みおよび損失関数と組み合わされている。
3. 包括的な評価: モデルは9つの多様なデータセットに対して8つの最先端のベースライン（GAN、VAE、拡散モデル、および他のTransformerを含む）と比較評価されており、有用性（utility）、忠実度（fidelity）、プライバシー、および効率性の優れたバランスを実証している。

実験結果

9つのデータセット（Adult, Bank, Boston, Creditなど）および8つのベースラインを用いて評価を行った：

• 有用性 (Utility): TabTreeFormerは、ダウンストリームの機械学習有効性（MLE）タスク（ロジスティック回帰、ランダムフォレスト、XGBoost）において、一貫してベースラインを上回っている。最も優れたバリアント（マスクなしのTabTreeFormer-NM）は、データの有用性を優先する場合、最良のベースラインに対して最大44%の性能向上を達成した。
• 忠実度 (Fidelity): 本モデルは、ベースラインと同等の「Shape」スコア（周辺分布の類似性）を達成し、特にマルチモーダル分布や特徴量間の関係を捉える能力において、より優れた「Trend」スコア（相関の忠実性）を示した。
• プライバシー: マスキングを施したバリアント（TabTreeFormer-SおよびL）は、強力なプライバシー耐性を示し、高い記憶のリスクを示したいくつかのベースライン（TabuLa、Forest Diffusionなど）とは異なり、データセット全体を通じて重大なプライバシー漏洩は検出されなかった。
• 効率性: ハイブリッドな性質にもかかわらず、TabTreeFormerは、80M以上のパラメータを持つベースラインと比較して、より速い生成時間と、大幅に小さいモデルサイズ（5M–40Mパラメータ）を実現している。

意義

本論文は、TabTreeFormerが、表形式の予測における決定木モデルの成功と、生成モデリングにおけるそれらの帰納バイアスの活用不足との間の根本的なギャップを埋めるものであると主張している。ツリーベースの事前分布をTransformerアーキテクチャと効果的に組み合わせ、順序データのための特殊なトークン化および損失関数を導入することで、TabTreeFormerは高品質な合成表形式データの生成に対する堅牢なソリューションを提供する。これは、ドメイン固有の知識を組み込むことが、特にデータの有用性が極めて重要なシナリオにおいて、競争力のある効率性とプライバシー保証を維持しながら、生成性能を大幅に向上させ得ることを示している。