把 TabFM 的 forward pass 完整展开｜P11

00 · READING CONTRACT

先固定符号和发布配置

类默认值不等于发布模型。下面所有主路径都按官方 v1.0.0 hardcoded config；分类与回归只在目标编码和输出头上不同。

符号	含义	v1 值 / shape
B	一个批次中的表格/ensemble 实例数	动态
T	一张表进入模型的总行数；前 `train_size[b]` 行是 context	动态
H	padding 后 feature 宽度；`d[b]` 是真实有效列数	动态
G	每个 anchor 列读取的 source slot 数	3，offset = [0,1,3]
F	Fourier 频率数	32；sin/cos 后 64
E	cell / row block embedding dim	256
C	共享 CLS token 数	8
D	ICL row representation dim	C·E = 2048
K	分类最大类数	10；回归输出 1

Column blocks

3 blocks × 2 次

每个 ColumnEmbedding 是 4-head SetTransformer，256 inducing points，FFN dim 1024；两次实例不共享参数。

Row blocks

3 blocks × 2 次

8-head、RoPE base 100000、FFN dim 1024；第一次保留全序列，第二次只保留 CLS。

ICL blocks

24-layer Transformer

d_model 2048、8 heads、无 RoPE；分类 head 输出 10 logits，回归 head 输出 1 scalar。

01 · USER INPUT

模型看到的 categorical 不是字符串，也不是 embedding ID

sklearn wrapper 先把 DataFrame 编译成 float tensor。核心网络只接收 X, y, train_size, cat_mask, d。

INPUTDataFrame / ndarray

数值、字符串类别、缺失值；训练行与测试行由 wrapper 组织。

CATEGORICALOrdinal encoding

每列类别映射到浮点序号；不是 model 内部的 per-column embedding table。

CONTINUOUS PIPELINEScale / normalize

数值与 ordinal category 进入共享的表格预处理/ensemble 变换；cat_mask保留类型身份。

BATCH COMPILERPad + ensemble

得到 X[B,T,H]、cat_mask[B,H]、d[B]、train_size[B]。

CORE MODELfloat scalar cells

NaN 再保险地变为 -100.0 sentinel，并 cast 到 compute dtype。

边界：ordinal 编码人为给类别编号；TabFM 并没有把编号的距离当作可靠度量，而是用类别专属高频 Fourier bank 尝试 decorrelate 相邻编号。但这不等价于严格的类别置换不变性。wrapper 的 feature/category ensemble 是工程缓解，不是数学保证。

02 · CELL EMBEDDER

一个 scalar 如何变成 256 维 token

CellEmbedder 先构造 overlapping feature group，再按每个 source slot 的类型选择 numeric 或 categorical Fourier 分支，最后沿 slot 求和。

直接答案：category 会经过 Fourier。它不走 numeric 的频率和 linear projection，而是走 fourier_frequencies_cat[G,F] 与 in_linear_cat[2F,E]。numeric/categorical 互不共享；但每套参数在所有同类型列之间共享。

2.1 Feature grouping：列数不变，每个 anchor 聚合 3 个 source slots

\[j_{h,g}=(h+(2^g-1))\bmod d_b,\qquad g\in\{0,1,2\}\]

所以 offset 是 [0,1,3]。输入 X ∈ R^{B×T×H} 变成 Xg ∈ R^{B×T×H×G}。当每张表有效列数不同，模运算用 d[b] 而不是 padded width H；无效 anchor 最后清零。它不是连续窗口，也不把 H 压到 H/G。

2.2 Numeric branch

\[p^{num}_{bthgf}=X^g_{bthg}\,\Omega^{num}_{gf},\quad \phi^{num}=\big[\sin p^{num};\cos p^{num}\big]\in\mathbb R^{64}\]

\[z^{num}_{bthg}=W^{num}\phi^{num}_{bthg}+b^{num}\in\mathbb R^{256}\]

对象	shape	训练状态	共享规则
Ω_num	[G,F] = [3,32]	JAX 中 `nnx.Param`，初始化 N(0,σ²)，σ=1，可学习	按 group slot 区分；所有 numeric 列共享
sin/cos	[B,T,H,G,64]	无参数；Torch 用 float32 计算角度后 cast 回 compute dtype	逐 scalar/slot
W_num	[64,256]（框架存储方向可能转置）	可学习 linear；有 bias 对象，v1 主体 config `use_bias=False` 影响 transformer，cell linear checkpoint 仍按实现映射	所有 numeric 列及 slot 共用同一 projection

2.3 Categorical branch：同构，不同参数

\[p^{cat}_{bthgf}=X^g_{bthg}\,\Omega^{cat}_{gf},\quad z^{cat}_{bthg}=W^{cat}[\sin p^{cat};\cos p^{cat}]+b^{cat}\]

cat_mask[B,H] 用完全相同的 grouping 变成 cat_mask_grouped[B,1,H,G]。因此类型路由发生在 source slot，不是 anchor 列：同一个 anchor token 的三个 slot 可以分别走 numeric 与 categorical 头，然后求和。

\[e_{bth}=\sum_{g=1}^{G}\left(m_{bhg}\,z^{cat}_{bthg}+(1-m_{bhg})\,z^{num}_{bthg}\right)\in\mathbb R^{E}\]

JAX training truth

频率是可学习参数

jax/model.py 用两个 nnx.Param。训练会同时更新 Ω_num、Ω_cat、W_num、W_cat。

PyTorch inference port

频率注册为 buffer

官方 Torch port 用 register_buffer 承接已训练 checkpoint，目标是 inference/parity，不表示原始训练中的频率固定。

2.4 第一次 y 注入：只有 context rows

分类把 class id 送入 Embedding(K,E)；回归把 scalar y 送入 MLP(1→6→E, GELU)。然后把同一个 row-level y embedding 广播加到该训练行的每个 feature token：

\[\tilde e_{bth}=e_{bth}+\mathbf 1[t<n_b]q(y_{bt}),\qquad \tilde E\in\mathbb R^{B\times T\times H\times256}\]

测试行不注入 y。注意 ICL 前还会进行第二次、参数独立的 y 编码；两次注入不是共享同一个 encoder。

为什么这样设计

Fourier lift 让 shared linear head 表达非线性 scalar→token 映射。
类型专属 bank 避免 metric numeric 与 ordinal category 的频率偏好直接冲突。
按类型跨列共享，允许任意 schema zero-shot。
grouping 在 attention 前注入廉价的局部跨列上下文。

优点 / 代价

优：参数不随列数或类别 cardinality 增长。
优：categorical 无需固定全局 vocabulary。
代价：ordinal relabeling 会改变 Fourier phase。
代价：构造 [B,T,H,G,E] 中间量，显存高；Torch 用 row chunk 缓解。

可替代设计

数值：piecewise-linear bins、spline、RBF、quantile tokenization。
类别：hash/lookup embedding、DeepSets over category identity、置换等变编码。
共享：hypernetwork 按列统计产生 tokenizer。
group：随机/学习图、全列 attention、无 grouping。

03 · COLUMN EMBEDDING

Column Attention 沿行轴 T，不是沿列轴 H

源码先 transpose 为 [B,H,T,E]，再 reshape 成 [B·H,T,E]。因此每一列独立读取该表的 context rows。

\[X_c=\operatorname{reshape}(\operatorname{transpose}(X,[B,H,T,E]),[BH,T,E])\]

Axis

每个 feature 一条 row sequence

batch 维实际是 B·H，attention sequence 是 T。不同列在这个模块内不互相 attention；跨列交互留给 RowInteraction。

Mask

keys 只允许 context rows

mask[b,h,1,1,t] = (t < train_size[b])。训练与测试 query 都只能从 labeled context 的 column distribution 取信息。

内部是带 256 inducing points 的 SetTransformer，共 3 blocks。可把单 block 抽象为 induced set attention：先由 inducing queries 汇总 context rows，再让每个 row query 读取 inducing representation。最后 Linear(E,E) → RMSNorm，reshape 回 [B,T,H,E]。

\[Z^{col}=\operatorname{RMSNorm}\left(W_o\,\operatorname{SetTransformer}(X_c;M_{train})\right)\in\mathbb R^{B\times T\times H\times E}\]

为什么不用普通 T×T self-attention：inducing points 把 set attention 的主要 row 交互从 O(T²) 变成近似 O(TI)，I=256；代价是 information bottleneck，并且 column stage 依赖 context set 的统计表达能力。

04 · CLS + ROW INTERACTION

跨特征交互两次，第二次才做 row compression

8 个全局共享的 learnable CLS vectors 被复制到每一张表、每一行。Row attention 的 sequence 轴是 C+H。

\[X_r=[c_1,\ldots,c_C,z_1,\ldots,z_H]\in\mathbb R^{B\times T\times(C+H)\times E}\]

PREPEND8 CLS tokens

[B,T,H,256] → [B,T,8+H,256]。CLS 参数跨全部行/表共享。

ROW #1Full sequence

reshape [B·T,8+H,256]；3-layer 8-head Transformer + RoPE；保留 CLS 与 feature tokens。

COLUMN #2Re-contextualize

再次沿 T 做一个独立 SetTransformer。输入列轴现在包含 8 个 CLS channel 与 H 个 feature channel。

ROW #2CLS-only output

再次跨 8+H 交互，但只取前 8 个 token；feature outputs 丢弃。

FLATTENRow state

[B,T,8,256] → [B,T,2048]，供 dataset-wise ICL。

有效列 mask 只允许 key 位置 < C+d[b]，padding feature 不被读取。RoPE 作用在 row 内 token 序列的位置轴，即 CLS + feature position；这给模型列顺序信息，也意味着核心 row block 本身不是严格列置换不变。

设计目的

交替 column/row，让“列分布统计”和“同一行内特征关系”反复交换信息。
多个 CLS 提供比单 CLS 更宽的汇聚瓶颈。
固定 2048 维 row state 隔离 variable-H schema 与 ICL Transformer。

优点 / 代价

优：处理变长列数，最终 row dim 固定。
优：第二轮 column 可传播第一次 row interaction 的结果。
代价：8-token bottleneck 可能丢失细粒度 feature 信息。
代价：RoPE 与 grouping 都引入列顺序敏感性。

可替代设计

attention pooling / Perceiver latent array。
mean/sum pooling + DeepSets，换取严格不变性。
动态 routing / mixture of CLS experts。
feature graph message passing 或 cross-attention decoder。

05 · ICL TRANSFORMER + HEAD

在 row 级别做 dataset-wise learning

到这里每一行已压成 2048 维 token。24-layer ICL Transformer 沿 T 读取 labeled context rows，并为所有行输出预测。

5.1 第二次 y 注入

分类：one-hot K 维再 linear 到 2048；回归：MLP 1→4096→2048。只给训练行相加：

\[r'_{bt}=r_{bt}+\mathbf 1[t<n_b]g(y_{bt})\]

5.2 最关键的 attention mask

\[M_{b,1,1,j}=\mathbf 1[j<n_b]\]

mask 只约束 key/value 位置，不约束 query。结果是：训练 row 和测试 row 都可读取所有 training rows；任何 query 都不能把 test rows 当作 key。因此 test rows 彼此不可见，也不会泄漏测试标签。这个结构更像“并行查询同一 context memory”，不是普通 causal next-token mask。

5.3 Prediction head

\[h=\operatorname{RMSNorm}(\operatorname{ICL}_{24}(r';M)),\qquad \hat y=\operatorname{MLP}(h)\]

Classification

2048 → 4096 → 10

输出 [B,T,10] logits；wrapper 再依据任务真实 class 数裁切/ensemble/softmax，并逆映射标签。

Regression

2048 → 4096 → 1

输出 [B,T,1]；wrapper 还原 target preprocessing 与 ensemble 聚合。

ICL 无 RoPE：官方 v1 明确 use_rope=False。在 mask 允许的 context 集上，模型不依赖人为 row position encoding；这更接近表格行作为集合。替代方案包括 DeepSets/SetTransformer、cross-attention（test queries → train memory）、kernel/nearest-neighbor memory 或带显式 task token 的 encoder-decoder。

06 · PREFILL / DECODE

缓存改变计算组织，不改变信息边界

普通 forward 把 train+test 一次送入；服务路径把 context 编码一次，随后多个 test batches 重用 Column inducing representations 与 ICL K/V。

阶段	输入	缓存 / 计算	输出
prefill	`x_ctx[B,Tc,H]`, `y_ctx[B,Tc]`	补到 128 倍数；两次 Column 每层 inducing hidden；ICL 每层 K/V；记录有效 train_size	context logits + `{col1,col2,icl}`
decode	`x_test[B,Tq,H]`	y 全设为 -100；重算 row-independent cell/row stages；Column 跳过 context 汇总并读取 cached inducing hidden；ICL query 读取 cached K/V	`[B,Tq,K_or_1]`，裁掉 padding

Torch 的 ICL cache 可做 per-tensor symmetric int8 quantization；prefill_train_size 保持全精度。Column cache 与 ICL cache 是不同层级的复用：前者缓存每列的 context-set summary，后者缓存 row-level context K/V。

Sentinel

-100.0

NaN、padding、decode 时未知 y 共用协议；mask 决定其是否可参与信息流。

Padding

block size 128

prefill/decode 均补齐，最终按原始 T 截断。它是运行时效率策略，不是模型结构维度。

Chunking

4096 / 16 / 8192

Torch 默认按 rows、column instances、FFN tokens 分块；注释声明结果精确等价，用于 40GB GPU 内存安全。

07 · COMPLETE GRAPH

完整 forward 计算图

下面把普通 forward 的数据流、主 shape 与 attention 轴放在一张图里。n=train_size，v1: G=3,F=32,E=256,C=8,D=2048。

TabFM v1.0.0 · ordinary forward · classification shown

Raw table / DataFrame
  │ categorical: per-column ordinal encode → float scalar
  │ numeric/category preprocessing + ensemble + padding
  ▼
X [B,T,H]    y [B,T]    cat_mask [B,H]    d [B]    train_size n [B]
  │ nan_to_num(-100), cast
  ▼
FEATURE GROUPING  offsets [0,1,3], modulo d[b]
Xg [B,T,H,G=3]
  │
  ├─ numeric source slot ─ Ω_num [3,32] ─ sin/cos → [64] ─ W_num [64,256] ─┐
  │                                                                       │ select per source slot
  └─ category source slot ─ Ω_cat [3,32] ─ sin/cos → [64] ─ W_cat [64,256] ─┘ by grouped cat_mask
                                                                            │ sum over G
                                                                            ▼
CELL TOKENS [B,T,H,256]
  │ + y_cell(y)[B,T,256] broadcast over H, only t < n
  ▼
COLUMN EMBEDDING #1
  │ transpose/reshape → [B·H,T,256]
  │ SetTransformer ×3; attention sequence = ROWS T; keys only t < n; inducing I=256
  │ Linear + RMSNorm; reshape back
  ▼
[B,T,H,256]
  │ prepend shared CLS [8,256] to every row
  ▼
[B,T,8+H,256]
  │
  ▼ ROW INTERACTION #1
  │ reshape → [B·T,8+H,256]
  │ Transformer ×3; attention sequence = CLS+FEATURES; 8 heads + RoPE
  │ invalid padded features masked; output FULL sequence
  ▼
[B,T,8+H,256]
  │
  ▼ COLUMN EMBEDDING #2  (independent parameters)
  │ attention sequence = ROWS T for each CLS/feature channel
  ▼
[B,T,8+H,256]
  │
  ▼ ROW INTERACTION #2  (independent parameters)
  │ same row-wise transformer geometry; output ONLY first 8 CLS
  ▼
[B,T,8,256] ─ flatten ─→ ROW REPRESENTATION R [B,T,2048]
  │ + y_icl(y)[B,T,2048], only t < n (second, independent y encoder)
  ▼
ICL TRANSFORMER ×24  [B,T,2048]
  │ attention sequence = ROWS T; NO RoPE
  │ every query may use keys j only when j < n
  │ test rows cannot read other test rows
  ▼
RMSNorm [B,T,2048]
  │ MLP 2048 → 4096 → 10
  ▼
LOGITS [B,T,10]  → wrapper: real-class slice / ensemble / probability / inverse transform

逐阶段 shape ledger

#	模块 / 函数	输入	输出	作用
1	CellEmbedder.feature_grouping	[B,T,H]	[B,T,H,3]	每个 anchor 聚合指数偏移 source slots
2	Fourier + typed linear	[B,T,H,3]	[B,T,H,3,256]	每 slot scalar 非线性 lifting，按类型路由
3	slot sum + y #1	[B,T,H,3,256]	[B,T,H,256]	构造 cell token；context label 广播
4	ColEmbedding #1	[B,T,H,256]	[B,T,H,256]	按列读取 context row distribution
5	prepend CLS	[B,T,H,256]	[B,T,8+H,256]	建立固定容量 row aggregation slots
6	RowInteraction #1	[B,T,8+H,256]	同 shape	行内 feature/CLS 交互，保留全序列
7	ColEmbedding #2	[B,T,8+H,256]	同 shape	把 row-aware token 再做跨 context row 校准
8	RowInteraction #2	[B,T,8+H,256]	[B,T,2048]	只保留 8 CLS 并 flatten，压缩一行
9	y #2 + ICLearning	[B,T,2048]	[B,T,2048]	以 labeled rows 为 context 完成 dataset-wise adaptation
10	decoder	[B,T,2048]	[B,T,10] / [B,T,1]	分类 logits / 回归值

08 · EVIDENCE & BOUNDARY

每个判断对应到哪里

审计固定在官方仓库 main@b8a8b090c66d1b9e7af278003461582219996b6a。行号对这个 commit 有效。

JAX authority

`tabfm/src/jax/model.py`

YEmbeddingScheme 1479–1486
CellEmbedder 1489–1782；频率参数 1565–1584；类型路由 1699–1741
ColEmbedding 1785–1953
RowInteraction 1957 起
ICLearning 2073–2259；mask 2215–2232
TabFM 2263 起；普通/prefill/decode 在文件后段

Runtime cross-check

`tabfm/src/pytorch/model.py`

CellEmbedder 354–428；Torch frequency buffers 362–365
ColEmbedding 431–499；transpose 显示 attention 轴
RowInteraction 502–535
ICLearning 631–679
TabFM.forward 728–742
prefill/decode 744–817

Interface

`classifier_and_regressor.py`

CategoricalOrdinalEncoder、continuous preprocessing、ensemble、padding、真实 train_size/d/cat_mask 组织均在此。核心模型里的类别值不是字符串或 lookup index tensor，而是已编译 float scalar + type mask。

Released config

`jax/tabfm_v1_0_0.py`

40–65 行固定 v1 配置：E=256、G=3、F=32、C=8、Column/Row 各 3 blocks、ICL 24 blocks、SwiGLU、learned Fourier path 与 post-embedding y injection。

验证：官方 pytorch/model_test.py 在完整 JAX/PyTorch/dev extras 下结果为 6 passed, 2 subtests passed；另用小配置 hook 验证 [2,5,4] → cell [2,5,4,16] → CLS+row [2,5,6,16] → row state [2,5,32] → logits [2,5,5]。仓库未修改。测试证明实现 parity 与 shape，不证明模型性能、论文主张或每个设计选择的 ablation 优越性。

源码能确认“它具体怎么算”；源码本身不能确认“为什么这是最优设计”。本文的设计目的来自实现结构和官方代码注释；优缺点与替代设计属于工程分析，已与代码事实分开陈述。

TABFM TECHNICAL SERIES

技术细节 01一个数，为什么先绕进频率空间？技术细节 02不减一列，为什么每个 token 却读了三列？技术细节 03 · 当前TabFM 完整 forward pass