はじめに

こんにちはザワッチです。
Googleから出された表形式データを入力として出力に多クラス分類するモデルである、TabFMを使ってみました。

モデル概要

ここに詳しいリサーチブログが載っています。

学習手法
事前学習は、損失を計算して重みを更新している普通のニューラルネット学習。
今回は、分類（Classification）と回帰（Regression）がタスクにあるので、別々にチェックポイントが作られている。

分類：

回帰：

学習データセット
数億件の合成データセットのみで学習。
これらは構造的因果モデル（SCM）を使って動的に生成され、因果構造や特徴量間の関係といった、表形式タスクに典型的な帰納バイアスを与えている。

合成データが選ばれた理由として、多様で高品質なオープンソースの表データが乏しいことと、実データのプライバシー・ライセンス問題を避けるためで、実在のデータは学習に使われていないとのこと。

実際問題、因果推論できるような綺麗な学習データに使われているようなデータが各社持っていないと思うので、これでうまく汎化して推論できるのか気になるところ。結局何らかの前処理は必要そう。

アーキテクチャ

行と列に交互にアテンションをかけるハイブリッド設計で、TabPFN と TabICL の長所を組み合わせたものです。3つのステップで構成される。

まず列アテンション：
ColEmbeddingの中で、Transformer機構を呼び出して、その中で、マルチヘッドアテンションで列方向にかけることで、１特徴量で全サンプル間の関連性を計算する。

class SetTransformer(nn.Module):
  def __init__(self, num_blocks, d_model, nhead, dim_ff, num_inds,
               activation="swiglu"):
    super().__init__()
    self.blocks = nn.ModuleList([
        InducedSelfAttentionBlock(d_model, nhead, dim_ff, num_inds, activation)
        for _ in range(num_blocks)
    ])

  def forward(self, src, attn_mask=None):
    for blk in self.blocks:
      src = blk(src, attn_mask=attn_mask)
    return src
    
class InducedSelfAttentionBlock(nn.Module):
  def __init__(self, d_model, nhead, dim_ff, num_inds, activation="swiglu"):
    super().__init__()
    self.ind_vectors = nn.Parameter(torch.zeros(num_inds, d_model))
    self.mab1 = MultiheadAttentionBlock(d_model, nhead, dim_ff, activation)
    self.mab2 = MultiheadAttentionBlock(d_model, nhead, dim_ff, activation)

  def forward(self, src, attn_mask=None):
    ind = self.ind_vectors.unsqueeze(0).expand(src.shape[0], -1, -1)
    hidden = self.mab1(ind, src, src, attn_mask=attn_mask)
    return self.mab2(src, hidden, hidden)
    
class ColEmbedding(nn.Module):
  def __init__(self, d_model, num_blocks, nhead, dim_ff, num_inds):
    super().__init__()
    self.tf_col = SetTransformer(num_blocks, d_model, nhead, dim_ff, num_inds)
    self.out_w = nn.Linear(d_model, d_model)
    self.ln_w = RMSNorm(d_model)
    self.col_chunk_size = None  # chunk the independent column axis (B*HC)

  def _stage(self, src, mask):
    return self.ln_w(self.out_w(self.tf_col(src, attn_mask=mask)))

  def forward(self, x, train_size):  # x: [B,T,HC,E]
    b, t, hc, e = x.shape
    src = x.permute(0, 2, 1, 3).reshape(b * hc, t, e)  # [B*HC, T, E]
    ts = train_size.repeat_interleave(hc)  # [B*HC]
    mask = (torch.arange(t, device=x.device)[None, :] < ts[:, None])[:, None, None, :]
    cc = self.col_chunk_size
    if cc is None or src.shape[0] <= cc:
      out = self._stage(src, mask)
    else:
      out = torch.cat([self._stage(src[s:s + cc], mask[s:s + cc])
                       for s in range(0, src.shape[0], cc)], dim=0)
    return out.reshape(b, hc, t, e).permute(0, 2, 1, 3)

class TabFM(nn.Module):
  def __init__(self, *, embed_dim=8, max_classes=3, col_num_blocks=2,
               col_nhead=2, col_num_inds=4, row_num_blocks=2, row_nhead=2,
               row_num_cls=2, icl_num_blocks=2, icl_nhead=2, ff_factor=2,
               feature_group_size=3, num_freq=32, decoder_hidden=None,
               is_classifier=True):
    super().__init__()
    self.max_classes = max_classes
    self.is_classifier = is_classifier
    ff = embed_dim * ff_factor
    icl_dim = embed_dim * row_num_cls
    self.cell_embedder = CellEmbedder(embed_dim, max_classes, feature_group_size,
                                      num_freq, is_classifier)
    self.col_embedder = ColEmbedding(embed_dim, col_num_blocks, col_nhead, ff, col_num_inds)
    self.col_embedder_2 = ColEmbedding(embed_dim, col_num_blocks, col_nhead, ff, col_num_inds)
    self.row_interactor = RowInteraction(embed_dim, row_num_blocks, row_nhead, ff,
                                         row_num_cls, output_full=True)
    self.row_interactor_2 = RowInteraction(embed_dim, row_num_blocks, row_nhead, ff,
                                           row_num_cls, output_full=False)
    self.cls_tokens = nn.Parameter(torch.zeros(row_num_cls, embed_dim))
    self.icl_predictor = ICLearning(icl_dim, icl_num_blocks, icl_nhead, max_classes,
                                    icl_dim * ff_factor,
                                    decoder_hidden or icl_dim * 2, is_classifier)

    # Enable activation chunking by default (see the module constants above).
    # The per-block knobs stay settable (e.g. to None) for benchmarking.
    for module in self.modules():
      if hasattr(module, "row_chunk_size"):
        module.row_chunk_size = _ROW_CHUNK_SIZE
      if hasattr(module, "col_chunk_size"):
        module.col_chunk_size = _COL_CHUNK_SIZE
      if hasattr(module, "ffn_chunk_size"):
        module.ffn_chunk_size = _FFN_CHUNK_SIZE

  def forward(self, x, y, train_size, cat_mask=None, d=None):
    # Mirror the JAX model's entry: replace NaN with the -100 sentinel and cast
    # to the compute dtype (JAX: `jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)`).
    # NaN is already imputed in the shared preprocessing, so nan_to_num is a
    # no-op in the normal flow, but it keeps the model robust + JAX-faithful.
    x = torch.nan_to_num(x, nan=-100.0).to(self.cls_tokens.dtype)
    emb = self.cell_embedder(x, y, train_size, cat_mask, d=d)
    emb = self.col_embedder(emb, train_size)
    b, t, _, e = emb.shape
    cls = self.cls_tokens.expand(b, t, -1, -1)
    emb = torch.cat([cls, emb], dim=2)
    emb = self.row_interactor(emb, d=d)
    emb = self.col_embedder_2(emb, train_size)
    reps = self.row_interactor_2(emb, d=d)
    return self.icl_predictor(reps, y, train_size)

次に行アテンション：
RoPEで特徴量の位置を区別しながら、行の中の特徴量どうしを見比べさせる。

RoPEで特徴量の位置を区別しながら、行の中の特徴量どうしを見比べさせて、後にCLSベクトルを抜き出してその行の要約としている。（self.out_ln(out if self.output_full else out[:, : self.num_cls, :])）

次に行圧縮：CLSトークンと RoPE を用いた行レベルのアテンションで、各行を1つの密なベクトルに要約する。
RoPEで特徴量の位置を区別しながら、行の中の特徴量どうしを見比べさせて、後にCLSベクトルを抜き出してその行の要約としている。（self.out_ln(out if self.output_full else out[:, : self.num_cls, :])）

class Encoder(nn.Module):
  def __init__(self, num_blocks, d_model, nhead, dim_ff, activation="swiglu",
               rope_base=100000.0):
    super().__init__()
    # One RoPE per Encoder (mirrors JAX `tf_row.rope.freqs`), shared by all blocks.
    self.rope = RoPE(d_model // nhead, rope_base) if rope_base is not None else None
    self.blocks = nn.ModuleList([
        MultiheadAttentionBlock(d_model, nhead, dim_ff, activation, rope_base)
        for _ in range(num_blocks)
    ])

  def forward(self, x, attn_mask=None):
    for blk in self.blocks:
      x = blk(x, attn_mask=attn_mask, rope=self.rope)
    return x

class RowInteraction(nn.Module):
  def __init__(self, d_model, num_blocks, nhead, dim_ff, num_cls,
               rope_base=100000.0, output_full=True):
    super().__init__()
    self.tf_row = Encoder(num_blocks, d_model, nhead, dim_ff, rope_base=rope_base)
    self.out_ln = RMSNorm(d_model)
    self.num_cls = num_cls
    self.output_full = output_full
    self.row_chunk_size = None  # chunk the independent row axis (B*T)

  def _stage(self, src, mask=None):
    out = self.tf_row(src, attn_mask=mask)
    return self.out_ln(out if self.output_full else out[:, : self.num_cls, :])

  def forward(self, x, d=None):  # x: [B,T,HC,E]
    b, t, hc, e = x.shape
    src = x.reshape(b * t, hc, e)
    # Mask cross-column attention to the valid columns (CLS + d real features);
    # padded columns (>= d + num_cls) must not be attended to. Matches JAX.
    mask = None
    if d is not None:
      d_padded = d.to(torch.long) + self.num_cls  # [B]
      valid = torch.arange(hc, device=x.device)[None, :] < d_padded[:, None]  # [B, HC]
      mask = valid.repeat_interleave(t, dim=0)[:, None, None, :]  # [B*T, 1, 1, HC]
    rc = self.row_chunk_size
    if rc is None or src.shape[0] <= rc:
      out = self._stage(src, mask)
    else:
      out = torch.cat([self._stage(src[s:s + rc],
                                   None if mask is None else mask[s:s + rc])
                       for s in range(0, src.shape[0], rc)], dim=0)
    if self.output_full:
      return out.reshape(b, t, hc, e)
    return out.reshape(b, t, -1)

最後にICL Transformer：
24ブロックの因果Transformerが圧縮済みの行ベクトル列を処理し、学習行を文脈として扱いながらテスト行の予測を出力する。
分類なら decoder が max_classes 個のロジットを、回帰なら1個の値を出力する。

class ICLearning(nn.Module):
  def __init__(self, d_model, num_blocks, nhead, max_classes, dim_ff,
               decoder_hidden, is_classifier=True):
    super().__init__()
    self.tf_icl = Encoder(num_blocks, d_model, nhead, dim_ff, rope_base=None)  # ICL has no RoPE
    self.ln = RMSNorm(d_model)
    self.is_classifier = is_classifier
    if is_classifier:  # one-hot y-encode; decode to per-class logits
      self.y_encoder = OneHotAndLinear(max_classes, d_model)
      self.decoder = MLP(d_model, [decoder_hidden], max_classes)
    else:  # MLP y-encode the scalar target; decode to a single value
      self.y_encoder = MLP(1, [decoder_hidden], d_model)
      self.decoder = MLP(d_model, [decoder_hidden], 1)

  def forward(self, reps, y, train_size):  # reps: [B,T,d_model]
    b, t, _ = reps.shape
    tm = (torch.arange(t, device=reps.device)[None, :] < train_size[:, None])
    if self.is_classifier:
      y_enc = self.y_encoder(y)
    else:
      y_enc = self.y_encoder(y[..., None].to(reps.dtype))
    r = reps + y_enc * tm[..., None]
    mask = tm[:, None, None, :]
    out = self.tf_icl(r, attn_mask=mask)
    return self.decoder(self.ln(out))

TabFM.forward で、列→行→列→行、と手で交互に積んでいます。列用と行用のブロックが2セットずつ用意されているのがわかります。col_embedder と col_embedder_2、row_interactor と row_interactor_2 の2つずつです。

例えば、とある購買顧客データから「月額料金が高い」「利用月数が短い」から、この行は"高額なのに短期"という組み合わせ情報を得られて、（1回目の列x行アテンション）２回目にその情報量をもとに、サンプル全体を見渡し（列アテンション）、どれだけ珍しいかを掴む。その後の行圧縮で、意味あるベクトル表現に変えて、ICLでラベル予測をする的な。

  def forward(self, x, y, train_size, cat_mask=None, d=None):
    # Mirror the JAX model's entry: replace NaN with the -100 sentinel and cast
    # to the compute dtype (JAX: `jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)`).
    # NaN is already imputed in the shared preprocessing, so nan_to_num is a
    # no-op in the normal flow, but it keeps the model robust + JAX-faithful.
    x = torch.nan_to_num(x, nan=-100.0).to(self.cls_tokens.dtype)
    emb = self.cell_embedder(x, y, train_size, cat_mask, d=d)
    emb = self.col_embedder(emb, train_size)
    b, t, _, e = emb.shape
    cls = self.cls_tokens.expand(b, t, -1, -1)
    emb = torch.cat([cls, emb], dim=2)
    emb = self.row_interactor(emb, d=d)
    emb = self.col_embedder_2(emb, train_size)
    reps = self.row_interactor_2(emb, d=d)
    return self.icl_predictor(reps, y, train_size)

Pytorch：

JAX：

入力
数値列・カテゴリ列が混在した表形式データ（pandas DataFrame または numpy 配列）を受け付ける。

具体的にはラベル付きの学習行（X_train, y_train）と予測対象のテスト行（X_test）を、まとめて1つのプロンプトとしてモデルに渡す。

画像・音声・動画・生テキスト、グラフや系列などの非表形式データには対応していない。

出力
2種類のタスクに対応。

分類では、2値・マルチクラス（最大10クラスまで）に対応し、クラス確率を出力する。

回帰では、連続値のターゲットを予測する。

ソースコードは Apache 2.0 ですが、モデル重みは TabFM 非商用ライセンス v1.0 で、商用利用はできない。また、Googleの公式サポート製品ではありませんとのこと

特筆すべき特徴

We’ve seen a massive shift in how people handle time-series forecasting since we launched TimesFM. Now, we’re bringing that same "zero-shot" logic to tabular data.
We introduce TabFM, a new foundation model for tabular data to simplify classification and regression workflows.

リサーチブログの先頭にあるように、
ゼロショットで表形式データ（tabular data）を分類や回帰ができることが強みぽい。背景には、LLMがin-context learningのゼロショットで結構いい精度出せたことを応用している。
従来の表形式データの予測をICL問題として捉えることで、手動によるモデル学習、ハイパーパラメータ調整、複雑な特徴量エンジニアリングの必要性を排除していることがかなりでかい。
一方で、こちらで用意した学習データの全行をコンテキストを渡す形になってしまうので、メモリ使用量が同じ類のモデルより、バカ食う。
層化サンプリング（クラス比率を保った抽出）で行数を抑えつつ、TabFMは最大500特徴量までを想定して最適化されているので、テーブルデータのカラム最大500まではいけそう

使ってみた

HuggingFaceに公開されているので、こちらをダウンロード。

サンプルで使用するデータセット
与信・購買系のデータセットである、UCIの Default of Credit Card Clients（Taiwan, 2005）

台湾のクレジットカード顧客30,000件のデータで、与信限度額・属性（性別・学歴・婚姻・年齢）に加え、2005年4〜9月の返済状況・請求額・支払額の履歴（各6ヶ月分）を含みます。説明変数は23。
各月は別特徴量となる。

取得は ucimlrepo から直接ロードするのがお手軽。

from ucimlrepo import fetch_ucirepo
ds = fetch_ucirepo(id=350)   # Default of Credit Card Clients
X = ds.data.features          # 23の説明変数
y = ds.data.targets           # 翌月デフォルトか否か

何を予測するか
翌月にそのカード利用者が債務不履行（デフォルト）するか / しないかの2値予測。
手がかりは与信限度額、性別・学歴・婚姻・年齢、そして過去6ヶ月分の返済状況・請求額・支払額の履歴。

分析してみた結果：１つ目

より不均衡に強いROC-AUCで見ても、
TabFM 0.779 > LightGBM 0.763 > ロジスティック回帰 0.719で首位。正解率だけでなく順位付けの質でも上回る

中位のリスク層はどちらにも靡くためモデルでも判別するのは難しそう

さらに分析

TabFMは特徴量を横断して全サンプルの傾向を読む、と聞き、モデルが何を手がかりにしているかを2枚のヒートマップで見てみ
た。
名前は似ているが、測っているものは正反対

まず、相関ヒートマップでデータそのものの性質を見る。

各列のペアが値として一緒に動くかを見た図で、ここでは、たとえば請求額が直近でも5ヶ月前でもほぼ同じ＝顧客は毎月おおよそ同じ額を使う安定した消費行動をしている、というデータの素の構造がわかる。

返済状況・請求額・支払額がそれぞれ濃い赤のブロックになり、6列あっても実質1列分の情報しかない冗長さが垣間見える。

毎月おおよそ同じ額を使う、安定した消費パターンで、6列も持っているのに情報としてはほぼ1列分の冗長さしかなく、そして実際に相互作用ヒートマップでも請求額の組み合わせは薄い（TabFMもそこからは新しい情報を引き出せていない）。

これは、データ上は請求額も返済状況も似た冗長構造を持つが、TabFMは返済状況の組み合わせだけを予測に使い、請求額の組み合わせはほぼ使っていないと読める。

そのほか

とりあえず、非デフォルトで予測しても77.9%
Mac16GBでアンサンブルの方もやってみようとしたが、重すぎアンド遅過ぎて断念

まとめ

久しぶりに記事を書いて疲れました。
おわり