这一课只做一件事:把 KITE 从三个输入走到预测分布的完整数据流走一遍,记清"哪个模块吃什么、吐什么、改的是路径上的哪一段"。你已有 Flow Matching 与 KITE 高层认知,所以不重复入门,直接进入张量级走查。读完后你应能不看论文画出这张图。
KITE 要建模的是协变量条件下的概率预测,即图 1(c) 的范式:
$$p_\theta\!\left(Y_{\text{endo}}\mid X_{\text{endo}},\,X_{\text{exo}},\,Y_{\text{exo}}\right).$$
注意三个输入的语义分裂——这是理解整个数据流的钥匙:
| 输入 | 形状 | 语义 | 谁读它 |
|---|---|---|---|
| $X_{\text{endo}}$ | $N\times T$ | 历史内生(要预测的东西的过去) | 仅 HCM |
| $X_{\text{exo}}$ | $D\times T$ | 历史外生(已知过去协变量) | KGC |
| $Y_{\text{exo}}$ | $D\times F$ | 已知未来外生 | KGC(训练+推理都用) |
| $Y_{\text{endo}}$ | $N\times F$ | 预测目标(训练可见,推理要生成) | 训练路径构造;推理要生成它 |
FM 把生成看成从源 $Y_0\sim p_0$ 沿直线传输到目标 $Y_{\text{endo}}\sim p_1$,学一个向量场 $v_\theta$ 去逼近这条直线的速度:
$$Y_s = s\,Y_{\text{endo}} + (1-s)\,Y_0,\quad \hat v = Y_{\text{endo}}-Y_0,\quad s\in[0,1].$$
训练时 $Y_{\text{endo}}$ 已知,所以 $Y_s$ 和 $\hat v$ 都能算出来,网络回归 $\hat v$。推理时从 $Y_0$ 出发,用学到的 $v_\theta$ 做 ODE 积分到 $s=1$。KITE 改的就是 $Y_0$ 和 $v_\theta$ 的条件通路,FM 直线路径本身没动。
标准 FM 源是 $\mathcal N(0,I)$。HCM 把它换成历史条件流形,且只用内生历史:
$$\mu_{\text{hist}}=f_\phi(X_{\text{endo}})\in\mathbb R^{N\times F},\quad \sigma_{\text{hist}}=\mathrm{Softplus}(g_\psi(X_{\text{endo}}))+\sigma_{\min},$$ $$\delta_{\text{hist}}=\alpha\,\frac{Mz}{\|Mz\|}+(1-\alpha)\,\epsilon,\quad Y_0=\mu_{\text{hist}}+\sigma_{\text{hist}}\,\delta_{\text{hist}}.$$ 其中 $M\in\mathbb R^{(N\times F)\times r}$ 是可学低秩流形基,$z,\epsilon\sim\mathcal N(0,I)$,$\alpha$ 可学。配套的 coverage constraint $L_{CC}$ 防止 $\sigma$ 塌缩、强制源支撑盖住目标。
FM 每一步的状态 $Y_s$ 要"听"外生条件。KGC 不把先验拼在输入上,而是改注意力的查询投影。统计先验 $S=\{s_{ij}\}\in\mathbb R^{N\times D}$ 给每对(内生 $i$,外生 $j$)一个关联强度:
$$\mathrm{Attn}(q_i,k_j)=q_i\,(W_1+s_{ij}\,W_2)\,k_j^\top.$$
实现上拆两条投影支路($A_b$ 数据驱动、$A_g$ 先验注入),再与归一化先验 $\tilde S=\mathrm{Norm}(|S|)$ 融合并 log-gating:
$$A_b=\frac{(Y_s W_{1,q})(X_{\text{exo}}W_{1,k})^\top}{\sqrt d},\quad A_g=\frac{(Y_s W_{2,q})(X_{\text{exo}}W_{2,k})^\top}{\sqrt d},$$ $$A=\mathrm{Softmax}\!\big(A_b+\tilde S\odot A_g+\log(\tilde S+\delta)\big).$$
注意维度:query 来自 $Y_s$(内生侧,$N$),key 来自 $X_{\text{exo}}$(外生侧,$D$),所以 $A\in\mathbb R^{N\times D}$ 是一张"内生×外生"依赖图。最后把外生信息加权注回状态:
$$\tilde Y_s = A\,X_{\text{exo}}\,W_v + Y_s\in\mathbb R^{N\times F},\qquad W_v\in\mathbb R^{T\times F}.$$
训练时以概率 $p_{\text{con}}$ 把条件 $c$ 换成空条件 $\varnothing$:
$$\tilde c=\begin{cases}\varnothing & \text{w.p. }p_{\text{con}}\\ c & \text{w.p. }1-p_{\text{con}}\end{cases},\qquad L_{\text{FM}}=\big|v_\theta(Y_s,s,\tilde c)-(Y_{\text{endo}}-Y_0)\big|.$$
于是同一个网络既会 $v_\theta(\cdot,c)$(协变量条件)又会 $v_\theta(\cdot,\varnothing)$(协变量无关)。推理时线性外推,$\gamma\ge 0$ 是旋钮:
$$\hat v_s=(1+\gamma)\,v_\theta(Y_s,s,c)-\gamma\,v_\theta(Y_s,s,\varnothing).$$
$\gamma=0$ 退回纯条件生成;$\gamma>0$ 把轨迹从"无外生平均行为"推开、更贴当前协变量驱动。经验最优 $\gamma\in\{1.2,1.4\}$。最后从 $Y_0$ 出发做 $N$ 步 Euler ODE:$Y_{\text{curr}}\gets Y_{\text{curr}}+\hat v\,\Delta s$。
| 瓶颈 | 模块 | 读什么 | 改路径哪段 |
|---|---|---|---|
| 拓扑落差(源太远) | HCM | $X_{\text{endo}}$ | 源 $Y_0$ |
| 伪相关放大 | KGC | $X_{\text{exo}},Y_{\text{exo}},S$ | 条件注入 → $\tilde Y_s$ |
| 条件强度不可控 | CFG | (不额外读输入) | 速度场外推 $\hat v$ |