G6.4 實作：座標、力與鏡像的三種檢查

本篇重用1 篇

G6.3鏡子裡的分子，也一定要得到相同答案嗎？

同一批原子，分別檢查三種答案

G6 已經說清楚幾何模型不能只測「轉動後好像一樣」。我們要知道輸出的型態：energy 是 scalar，force 是向量，position 是座標。

這次使用三維點與 scalar node features，EGNN 先更新 features 與 coordinates，再由 invariant readout 產生 energy。訓練 target 是一個平滑的 pair-distance energy，不代表真實 molecular force field。

準備與執行

本課的實作使用 Python 3.10 以上與 PyTorch 2.x，CPU 就可以計算。請在自己的 Python 環境安裝 PyTorch，再下載 gdl_labs.py。環境目錄與下載的第三方套件不要加入網站的 Git repository。

python -m pip install torch
python gdl_labs.py --unit 6 --steps 40

檔案包含可重用的模型與 Trainer。模型、loss 與資料來源由外部指定；換資料時保留對稱一致的 tensor interface，不需要修改 training loop。

from gdl_labs import Trainer

trainer = Trainer(model, loss, lr=0.003)
history = trainer.fit(batch_fn, steps=40)
# batch_fn() returns ((model_input_1, model_input_2, ...), target)

先看介面

from gdl_labs import EGNN
import torch

model = EGNN(in_dim=2, hidden=24, layers=2)
x = torch.randn(4, 5, 3)
h = torch.randn(4, 5, 2)  # scalar node features
energy = model(x, h)      # B,1
new_x, features = model.encode(x, h)
forces = model.forces(x, h)

空間維度由 x.shape[-1] 決定，不繫結二維或三維。Scalar features 的維度則由 in_dim 指定；這兩個維度扮演不同角色。

合法變換與 expected output

程式抽一個 orthogonal $Q$ ，把 determinant 修成 $1$ ，再加共同平移 $b$ 與原子重排 $p$ ：

moved_x = x[:, p] @ Q.T + b
moved_h = h[:, p]
# Energy expected: energy
# Position expected: new_x[:, p] @ Q.T + b
# Force expected: forces[:, p] @ Q.T

請對照三個 expected 的平移差異。程式逐項 assert defect 小於浮點 tolerance，也檢查總 force；後者來自 invariant energy 的 gradient，不是一般 EGNN coordinate update 自動保證。

Force training 的 gradient

predicted_forces = model.forces(x, h, create_graph=True)
loss = (predicted_forces - target_forces).square().mean()
loss.backward()

Force loss 要再對 model parameters 微分，因此保留 coordinate gradient 的計算 graph。只做測試時可用預設 create_graph=False，但計算 force 仍必須在能對 coordinate 求 gradient 的區域裡；函數已用 torch.enable_grad() 處理外層 evaluation context。

如果 dataset 有 padding，請使用 mask 計算 energy；目前方便的 forces helper 接收完整有效點，若要 force training with mask，需把 mask 傳給 energy computation 並排除假點的 loss。

鏡像不是旋轉測試的另一個名稱

程式另外鏡射一個有四種角色的 tetrahedron：所有 pairwise distances 相同，signed volume 由 $1$ 變成 $-1$ 。Distance-based energy head 仍同值，符合它的 E(d)-invariant 假設。

若改成預測 handedness label，就不能仍用這個 invariant energy head。請先改 target action 與 feature construction，再問模型為什麼沒有學到。

作業

對同一批資料列出 permutation、translation、proper rotation、reflection 四組測試。每組報告 energy、position、force 的 expected action 與 absolute max defect，不要只列一個總數。

加分部分：比較直接的 coordinate displacement 與 energy-gradient force 的總和。說明為什麼前者可能 equivariant，卻沒有成對抵消或 conservative guarantee。

三種 expected outputs 分開測試：energy 不變，position 旋轉加平移，force 只旋轉與重排。

課堂提問

為什麼 coordinate update 的 displacement 可以總和不為零，而 energy gradient 得到的 total force 卻應接近零？

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

前者只按幾何 action 構造向量，scalar edge weights 不必成對對稱；後者來自對共同 translation invariant 的 energy，方向導數強制所有 forces 的總和為零。

動手看一次

分開測試平移、旋轉與鏡射；再故意不旋轉向量。請對照三種輸出的 expected action，而不是只看某一個總 defect。

快速回顧一下

補充接下來的擴充規劃

第一階段之後，我們會接著處理五組問題：

G7 · Tensor Features： scalar 與 vector 不夠時，如何表示角度與更高階的幾何 features？
G8 · Graph Structure and Spectral Methods： 局部訊息不足時，如何補上 graph frequency、遠方資訊與結構？
G9 · Surfaces and Local Frames： 沿曲面移動時，不同位置的方向要怎麼比較？
G10 · Approximate Symmetry and Canonicalization： 對稱只近似成立時，如何保留有用的限制？
G11 · Geometric Generative Models： score、velocity 與取樣過程要怎麼遵守幾何 action？

目前依這些問題規劃為 7、8、7、5、7 篇，共 34 篇擴充。Graph frequency 與 filtering、局部 kernel 與完整網路、ODE 與 SDE 的取樣檢查分開講；問題與偏差量測則放在一起。這是後續的分篇規劃，正文尚未完成。