U6.6 實作：MeanFlow 的導數、落點與半群，要分開檢查

本篇重用3 篇

先選一條真的知道答案的 flow

U6.2 的公式很短，但程式可能把兩個時間交換、少算空間方向，或讓 target gradient 流回網路。若只看生成樣本像不像資料，這些錯誤很難定位。因此這一篇沿用 U5 的 Gaussian 實驗方式，換成能直接計算任意時間 map 的 transport。

下載完整 Python 程式（NumPy） · 下載可獨立執行的 notebook

python -m pip install numpy
python dma_fast_maps.py --unit 6 --steps 2500 --seeds 0 1 2 --out u6-results

這裡真的更新一個小 MLP；oracle map 只作對答案。程式內的 JVP 由 tanh 網路的解析導數計算，並先與 centered finite difference 核對；若改成 PyTorch，使用 torch.func.jvp 的方向完全相同。

從 marginals 算出 map 與 velocity

取獨立的 $X_0\sim\mathcal N(0,1)$ 、 $X_1\sim\mathcal N(2,0.6^2)$ ，沿用 $I_t=(1-t)X_0+tX_1$ 。因此 marginal mean 與 variance 是

m_t=2t,\qquad q_t=(1-t)^2+0.36t^2.

對同一個標準化座標 $z$ ，令 $x_t=m_t+\sqrt{q_t}z$ 。這樣每個時刻的 law 都是所需 marginal；微分後得到

v_t(x)=2+\frac{q'_t}{2q_t}(x-m_t).

它也等於 reference velocity 的條件平均。從位置 $x$ 找回 $z=(x-m_t)/\sqrt{q_t}$ ，再送到時間 $r$ ，便得

\psi_{t\to r}(x) =m_r+\sqrt{\frac{q_r}{q_t}}(x-m_t).

這是精確 map，不是 Euler approximation。平均速度是 $(\psi_{t\to r}(x)-x)/(r-t)$ ；程式 mean_velocity 用等價而數值穩定的表示，並在 $r=t$ 回傳 velocity 極限。

先驗證導數方向，再開始訓練

網路輸入順序固定為 $(x,r,t)$ 。沿同一條 trajectory 改變起始時間時，位置以 velocity 移動，目的時間不變，方向必須是 $(v,0,1)$ 。verify_identities() 在訓練前檢查四件事：JVP 與有限差分一致、精確 MeanFlow identity 成立、對角線回到 $v_t$ 、精確 map 滿足半群。

若只把 $t$ 擾動、位置不動，計算的就是 $\partial_tu$ ，漏掉 $J_xu\,v$ 。這不是同一個微分問題，不能用更久的訓練補救。

從資料造 input，再組成固定 target

每個 batch 抽 $(X_0,X_1)$ 與 $t\le r$ ，部分時間對令 $r=t$ ，其餘保留長段。自己合成的 $x=(1-t)X_0+tX_1$ 有一個條件速度 $\tilde v=X_1-X_0$ 。

核心 target 與 U6.2 完全對齊：

features = meanflow_features(x, r, t)
direction = np.column_stack((velocity/3, np.zeros_like(t), np.ones_like(t)))
target = velocity + (r-t)*net.jvp(features, direction)

features 把位置除以 $3$ 以調整輸入尺度，所以擾動方向的位置分量也要除以 $3$ ；這是 chain rule。兩處 velocity 使用同一份條件 label。接著

pred, hidden = net.forward(features)
net.update(features, hidden, 2*(pred-target)/batch)

只有 pred 進入 backprop，target 與 JVP branch 都當成已算好的數值。這讓 conditional velocity 的 affine expected-update 論證適用。示範程式保留這個原始結構，沒有再加非線性 tangent normalization。

關掉導數修正，網路會少學到什麼？

程式同時訓練 MeanFlow 與 no_correction。兩者使用相同 network、初始 seed、資料與更新數；後者將所有時間對的 target 都改成 $\tilde v$ 。如此回歸會傾向學輸入點的 instantaneous velocity，而沒有要求整段平均。

先在不跑程式的情況下，利用解析 $u$ 與 $v$ 比較一步位移，寫下哪些起點可能偏多、偏少。執行後再看結果；一個 finite-budget 的對照不保證某條線對所有 seeds 都較低。

課堂提問Q1

如果增加 sampling steps 後 empirical $W_2$ 變小，是否已證明模型的半群更準？

先想一想，再展開看整理後的答案

沒有。分布距離只比較最後一批輸出，可能看不出逐點 map 的差異。半群 residual 則比較同一個輸入的直接與分段落點；兩者可以一起改善，也可以有不同變化。

更直接的反例是恆等 map：它的半群 residual 為零，但在這個 $\mathcal N(0,1)$ 到 $\mathcal N(2,0.36)$ 的任務中，output law 明顯錯誤。必須同時對照解析 map 與資料 reference。

把三個數字放在同一張結果表

u6-results.csv 對每個模型、seed 與 calls 數列出：

數字	比較什麼？	它沒有單獨保證什麼？
`map_rmse`	同一起點的模型結果與精確 map	整體多樣性的一般結論
`semigroup_rmse`	直接 $0\to1$ 與 $0\to1/2\to1$	正確的資料動力學
`empirical_w2`	生成樣本與獨立資料 reference	特定逐點 coupling 正確

semigroup_rmse 固定比較這一種切法，因此同一 seed／模型在不同 calls 列會重複；它不是每種 schedule 的綜合 residual。每組使用 $4096$ samples，oracle 多步的 map 與半群差距應只剩浮點量級，empirical $W_2$ 仍有抽樣地板。

實際執行 $2500$ updates、seeds $0,1,2$ ，一步 empirical $W_2$ 的三種子範圍為 MeanFlow $0.159$ – $0.175$ 、移除修正 $0.528$ – $0.542$ ，oracle $0.0181$ – $0.0199$ 。這是目前這個小網路與訓練預算的結果：修正項有實際影響，學出的 map 也仍有可量差距。三種子範圍不是信賴區間，不能替代更大規模的實驗。

下一輪實驗，一次只改一項

將 JVP direction 改成只有時間分量，重新執行，分開看 identity check 與 learned-map 結果，確認錯誤先在哪個層級出現。
固定模型，比較 $1,2,4,8$ calls，畫三種指標；若多步變差，檢查中途 state 的分布與各段 map error，不要先假設程式一定有 bug。
將長段 target 改成兩短跳合成，加入 $d=0$ 的 local FM，實作 U6.3 的 Shortcut；比較相同總網路求值預算。
另以 Gaussian 平移 generator 實作 U6.4 的 score-difference 更新，先用解析 student score 檢查正負號，再換成 auxiliary model，分開量估計誤差。

互動 demo：先用解析答案對照。 每個數字由目前參數重新計算；實際訓練結果另見實作程式。

消化一下

參考文獻

主要理論參考：The Principles of Diffusion Models，§11.5 的平均速度與 sampling；解析 Gaussian transport、NumPy JVP 與評估程式是本課程實驗。正文的符號沿用本課程；方法實作細節再對照下列原始論文。