GNN (20) SE(3)-Transformer：三維空間的等變性

⏪ 上篇傳送門

[embed]GNN (19) 從發現到突破上限 ⏪ 上篇傳送門medium.com

前篇筆記建立了等變性理論的完整數學框架：所有保持排列對稱性的線性層，可以被 Bell 數精確地刻畫；k 階等變網路的表達能力等價於 k-WL 測試，嚴格超越訊息傳遞框架。

等變性這個概念，最初是從圖結構的原生屬性啟發的：圖沒有自然的節點編號順序。一個正確的圖學習模型，輸出結果不能因為節點的編號方式不同而改變。因此等變性不是一個設計偏好，而是圖資料本身的性質所要求的必要條件。

這個思路可以推廣到更廣泛的場景：任何時候，當問題本身具有某種已知的對稱性，把這個對稱性直接編碼進模型架構，就比讓模型從資料中重新學習它更為根本。一個忽略已知對稱性的模型，是在用有限的訓練資料重新發現物理學家幾百年前早已確立的定律。這是根本性的浪費，也是根本性的錯誤。

在藥物設計、蛋白質結構預測、分子動力學模擬等場景中，這個問題變得尤為具體。圖的節點是三維空間中有真實座標的原子，模型必須對旋轉（rotation）與平移（translation）保持正確的回應。一個苯環無論怎麼旋轉，它的毒性不變（不變性）；預測每個原子受到的力時，把分子旋轉 90 度，力的向量也必須跟著旋轉 90 度（等變性）。這是物理定律，不是設計選擇。標準的訊息傳遞 GNN 無法保證這個性質，因為它學到的是訓練資料的特定擺放方向，而非真正的物理規律。

Thomas et al.（2018）的 Tensor Field Networks（TFN）是第一個系統性解決這個問題的框架；Fuchs et al.（2020）的 SE(3)-Transformer 在 TFN 的基礎上引入了注意力機制，把 GAT 的動態權重與 TFN 的幾何等變性整合在同一個架構中。

什麼是 SE(3)？

SE(3)（Special Euclidean group in 3D，三維特殊歐幾里德群）是描述三維空間中所有旋轉與平移操作的數學結構。「Special」指的是只包含保持手性（handedness）的旋轉（即行列式為 +1 的旋轉矩陣，構成 SO(3) 群），不包含鏡像翻轉。

對一個三維點雲（point cloud）或分子圖施加一個 SE(3) 變換，意思是：對所有點的座標同時施加同一個旋轉矩陣 R ∈ SO(3) 與平移向量 t ∈ ℝ³。

SE(3) 等變性（SE(3)-equivariance）的要求是：對輸入施加任意的 SE(3) 變換之後，模型的輸出以對應的方式跟著變換。具體地，根據輸出的類型，「對應的方式」有所不同－－這正是後續要介紹的「特徵類型」概念的核心。

為何平移等變性相對容易保證

在訊息傳遞 GNN 的框架中，平移等變性（translation equivariance）相對容易保證：只要模型在計算時只使用節點之間的相對位置（relative position）x_j — x_i，而非絕對座標，平移等變性就自動成立。把所有點加上同一個平移向量 t，相對位置 (x_j + t) — (x_i + t) = x_j — x_i 不變。TFN 與 SE(3)-Transformer 都嚴格遵守這個原則。

真正困難的是旋轉等變性（rotation equivariance）。旋轉操作改變了相對位置向量的方向，而為了確保模型對旋轉的回應方式是「正確且可控的」，需要精心設計特徵的表示方式與計算方式。

旋轉的幾何語言

TFN 與 SE(3)-Transformer 的核心設計，建立在「型態特徵（type-ℓ features）」這一種特殊的特徵表示方式。要理解這個設計，需要先意識到「不同種類的物理量，在旋轉下有不同的變換行為」。

標量、向量、與高階張量

在物理學中，不同類型的量在旋轉下有截然不同的行為：

• 純量（scalar）：旋轉下保持不變。例如把分子旋轉 90 度，一個原子的質量、電荷、溫度都不會變。純量對應型態 ℓ = 0，是一個 1 維向量（2×0+1 = 1）。
• 向量（vector）：旋轉下跟著旋轉矩陣 R 變換。例如把系統旋轉 90 度，一個粒子的速度向量也跟著旋轉 90 度。向量對應型態 ℓ = 1，是一個 3 維向量（2×1+1 = 3）。
• 高階張量（higher-order tensor）：旋轉下按更複雜的方式變換。例如描述分子鍵角的二階張量。型態 ℓ = 2 是一個 5 維向量（2×2+1 = 5），型態 ℓ 是一個 (2ℓ+1) 維向量。

型態 ℓ 特徵的旋轉變換由 Wigner-D 矩陣（Wigner-D matrix）D^ℓ(R) 描述。Wigner-D 矩陣是 SO(3) 群的不可約表示（irreducible representation） — — 把它理解為「如何精確描述型態 ℓ 的物理量在旋轉下的變換方式」的數學規則。對於型態 0 的特徵，D⁰(R) = 1（純量不變）；對於型態 1 的特徵，D¹(R) = R（向量跟著旋轉矩陣變換）。

為何需要混合不同型態的特徵

在真實的分子計算中，不同型態的特徵往往需要互相作用。例如：兩個原子的相對位置向量（型態 1）與鍵能（型態 0）相乘，可以得到描述鍵方向能量的型態 1 特徵。這個「型態 0 × 型態 1 → 型態 1」的組合，在數學上由 Clebsch-Gordan 係數（Clebsch-Gordan coefficients）精確地描述。

Clebsch-Gordan 係數在物理學中是量子力學角動量疊加的標準工具 — — 它告訴你兩個不同型態的特徵相乘後，能夠產生哪些型態的輸出，以及各個輸出型態的係數是多少。TFN 與 SE(3)-Transformer 都以 Clebsch-Gordan 係數作為型態間特徵混合的數學基礎。

一般來說，型態 ℓ₁ 與型態 ℓ₂ 的特徵相乘，可以產生型態 ℓ 的輸出，其中 ℓ 的範圍是 |ℓ₁ — ℓ₂| ≤ ℓ ≤ ℓ₁ + ℓ₂。這個規則保證了特徵混合操作的等變性。只要每一步都遵守 Clebsch-Gordan 的規則，整個計算流程就自動維持 SE(3) 等變。

TFN：等變卷積的先驅

Thomas et al.（2018）的 Tensor Field Networks 是第一個在三維點雲上實現嚴格 SE(3) 等變的神經網路框架，也是 SE(3)-Transformer 的直接技術前身。

以球諧函數作為等變濾波器

TFN 的核心問題是：如何設計一個對旋轉保持等變的卷積濾波器？

在標準 CNN 中，濾波器是一組可學習的數值，對任意方向的輸入都有固定的回應。這在 2D 圖像上可行，因為圖像有自然的座標系；但在三維點雲中，濾波器必須對所有可能的旋轉方向都給出一致的等變回應，這對任意形狀的濾波器幾乎不可能。

TFN 的解法是把濾波器分解為兩個部分的乘積：

• 徑向函數（radial function）φ(r)：只依賴兩個節點之間的距離 r = ||x_j — x_i||，是一個可學習的函數（通常用 MLP 實現）。距離是旋轉不變量，意即旋轉不改變兩點之間的距離，因此這個部分可以自由學習而不破壞等變性。
• 球諧函數（spherical harmonics）Y^ℓ_m(x̂)：只依賴兩個節點之間的相對方向 x̂ = (x_j — x_i)/||x_j — x_i||（單位向量），是固定的、不可學習的基底函數。球諧函數是定義在球面上的函數族，是 SO(3) 群在球面上的自然正交基底。它們在旋轉下按 Wigner-D 矩陣精確地變換。

把這兩個部分組合起來：

[embed]

其中 ℓ_i 是輸入特徵的型態，ℓ_f 是輸出特徵的型態。濾波器 W 的角度方向完全由球諧函數決定（固定且等變），只有徑向部分是可學習的。這樣，無論輸入點雲如何旋轉，濾波器的等變性都由球諧函數的數學性質保證。

訊息傳遞

有了等變濾波器之後，TFN 的前向傳播可以寫成訊息傳遞的形式。對於節點 i，其型態 ℓ 輸出特徵的計算為：

第一項是 自交互項（self-interaction）：節點 i 自身的型態 ℓ 特徵乘以可學習的純量權重 w^{ℓℓ}，直接傳遞到下一層，相當於 ResNet 的跳躍連接（skip connection）確保資訊不會在深層網路中消失。

第二項是 鄰居聚合項：對節點 j ≠ i 的所有型態 k 的特徵，用等變濾波器 W^{ℓk}(x_j — x_i) 做加權求和。濾波器 W^{ℓk} 把型態 k 的輸入特徵轉換成型態 ℓ 的輸出，並且這個轉換對旋轉保持等變。

模型建構

TFN 的每個節點維護一個按型態分層的特徵字典。型態 0 的特徵（純量）、型態 1 的特徵（向量）、型態 2 的特徵（高階張量）分別存放在不同的「槽」中。每個槽有若干個通道（channels），類似於 CNN 的特徵圖深度。

每一層的計算，是把來自不同節點、不同型態的特徵，透過等變濾波器混合後輸出到下一層的對應型態槽中。整個計算流程完全保持 SE(3) 等變性，因為每一步都嚴格遵守 Clebsch-Gordan 的型態混合規則。

TFN 的根本限制在於它的 聚合權重是靜態的。每對節點 (i, j) 之間的聚合權重完全由它們的相對位置 x_j — x_i 決定，與節點的特徵內容無關。這與 GCN 的各向同性聚合有著類似的問題：模型無法根據鄰居的語義內容動態調整對不同鄰居的關注程度。SE(3)-Transformer 正是為了解決這個限制而設計的。

SE(3)-Transformer：注意力遇上幾何等變性

Fuchs et al. (2020）的 SE(3)-Transformer 把 TFN 的幾何等變性與 Transformer 的自注意力機制整合在一起，讓模型能夠在保持 SE(3) 等變性的同時，根據鄰居的特徵內容動態調整聚合權重。

不變的注意力係數，等變的 Value

讓我們先回顧 GAT 的注意力設計。注意力係數 α_ij 是一個純量，它決定節點 i 對鄰居 j 的關注程度；訊息 V(h_j) 是節點 j 的特徵向量，被加權求和後更新節點 i 的表示。

在 SE(3)-Transformer 中，這個設計被分成兩個不同的對稱性要求：

• 注意力係數 α_ij 必須是 SE(3) 不變的（invariant）：不管輸入點雲如何旋轉，節點 i 對節點 j 的注意力係數應該保持不變。「這個鄰居重不重要」這個判斷，不應該因為整個系統的旋轉方向而改變。
• 訊息 V(x_j — x_i, f_j) 必須是 SE(3) 等變的（equivariant）：它攜帶了方向性的幾何資訊，必須隨著旋轉以正確的方式跟著變換。

這個「不變的注意力 × 等變的訊息」的組合，是整個架構設計的核心思想。注意力係數只是加權求和的純量係數，純量與等變向量的乘積仍然是等變的，因此整體的注意力聚合結果保持 SE(3) 等變性。

前向傳播

SE(3)-Transformer 的節點更新公式為：

與 TFN 的公式相比，差異只在第二項多了注意力係數 α_ij。

W^{ℓk}_V(x_j — x_i) 是訊息的濾波器（value filter），與 TFN 的等變濾波器結構相同（球諧函數 × 徑向函數），負責把型態 k 的輸入特徵轉換成型態 ℓ 的輸出特徵。

如何計算 SE(3) 不變的注意力係數

注意力係數 α_ij 的計算方式，是 SE(3)-Transformer 設計中最精妙的部分。

標準的點積注意力（dot-product attention）計算 q_i · k_j，其中 q 和 k 是特徵向量的線性投影。但在 SE(3) 等變的設定下，q_i 和 k_j 本身都是等變的型態 ℓ 特徵向量。我們好奇的是，對它們做點積能否保持不變性？

關鍵性質在於：兩個在旋轉下以相同 Wigner-D 矩陣變換的向量，它們的點積是旋轉不變的。這是 Wigner-D 矩陣的正交性保證的：若 q → D^ℓ(R) q 且 k → D^ℓ(R) k，則 q^T (D^ℓ(R))^T D^ℓ(R) k = q^T k，因為 Wigner-D 矩陣是正交矩陣。

SE(3)-Transformer 利用這個性質設計查詢與鍵：

查詢 qi 只依賴節點 i 自身的特徵，鍵 k{ij} 依賴節點 j 的特徵以及相對位置（透過等變濾波器 W^{ℓk}_K）。兩者都是等變的型態特徵向量，它們的點積就是 SE(3) 不變的注意力分數，再透過 softmax 正規化得到注意力係數：

注意力的角度調製作用

SE(3)-Transformer 的注意力機制還解決了 TFN 的另一個技術限制：角度約束問題（angular constraint problem）。

在 TFN 的等變濾波器 W^{ℓk}(x) = φ(r) · Y^ℓ(x̂) 中，角度方向完全由球諧函數 Y^ℓ(x̂) 決定，是固定的、不可學習的。這意味著濾波器在角度方向上的自由度被完全鎖定。模型只能學習「從哪個距離聚合資訊」（徑向函數），無法學習「從哪個方向聚合資訊」（角度方向）。

注意力係數 α_ij 的引入，等同於給每個邊 (i, j) 乘以一個資料相依的純量調製（scalar modulation）。當 α_ij 因角度方向的不同而有所差異時，模型實際上是在對 TFN 濾波器的角度方向做動態加權。在不違反等變性的前提下，增加了模型在角度方向上的表達自由度。

自交互層：等變的跳躍連接

SE(3)-Transformer 的自交互層（self-interaction layer）是一個在單個節點內部、不涉及鄰居資訊的線性操作，其功能類似於 CNN 的 1×1 卷積或 ResNet 的跳躍連接。

在注意力聚合層中，節點 i 不對自身進行注意力計算（避免自迴路）。自交互層補償了這個設計：把節點 i 在層 L 的自身特徵，直接傳遞到層 L+1，與鄰居聚合的結果相加。

SE(3)-Transformer 提出了兩種自交互設計：

• 線性自交互（linear self-interaction）：對同一型態 ℓ 的不同通道做可學習的線性混合，不同點共享相同的混合權重。這是最基本的跳躍連接形式。
• 注意力自交互（attentive self-interaction）：混合係數不再是固定的全局權重，而是由節點自身特徵通過 MLP 動態計算得到的注意力權重。這讓不同的節點可以學習不同的通道混合方式，增加了表達能力。

兩種自交互都保持 SE(3) 不變性，因為它們只在同一型態的通道之間做混合，不改變特徵的型態結構。注意力自交互的係數是由型態特徵的內積計算得到的純量，同樣是 SE(3) 不變的。

GAT 的啟示

GAT 與 SE(3)-Transformer 的共同結構

兩個模型都有注意力係數 α_ij（決定鄰居 j 對節點 i 的重要性）與 V 訊息（鄰居 j 貢獻的特徵內容），都對鄰域做注意力加權聚合，都有自身特徵的跳躍連接。

GAT 與 SE(3)-Transformer 的差異

• 特徵的型態結構：GAT 的節點特徵是普通的實數向量，沒有幾何型態的概念。SE(3)-Transformer 的節點特徵是按型態分層的字典－－型態 0 的純量、型態 1 的向量、型態 2 的高階張量分別存放，每個型態有獨立的通道。
• 注意力係數的計算：GAT 的注意力係數由節點特徵的線性變換加 LeakyReLU 計算，不考慮幾何位置。SE(3)-Transformer 的注意力係數由等變型態特徵的點積計算，保證 SE(3) 不變性，同時利用相對位置資訊 x_j — x_i 通過等變濾波器。
• V 訊息的計算：GAT 的 V 訊息是節點特徵的線性投影，與幾何位置無關。SE(3)-Transformer 的 V 訊息由 TFN 等變濾波器（球諧函數 × 徑向函數）計算，同時依賴節點特徵與相對位置，並且在旋轉下以正確的型態方式變換。
• 可學習參數的結構：GAT 每個注意力頭只需要學習投影矩陣 W_Q、W_K、W_V 與注意力向量 a。SE(3)-Transformer 需要為每個型態對 (ℓ_in, ℓ_out) 學習一個徑向函數（通常是 MLP），同時 Clebsch-Gordan 係數是固定的數學常數，不需要學習。

模型表現如何？

N 體模擬（N-Body Simulation）

這個實驗直接測試等變性本身的價值。五個粒子各帶有正電或負電荷，任務是預測 500 個時間步之後每個粒子的位置與速度。這是一個嚴格的等變任務：旋轉輸入，輸出的位置與速度預測也應該跟著旋轉。

實驗結果清楚地說明了等變性的必要性：非等變的 Set Transformer 在輸入旋轉時，預測的粒子軌跡發生了顯著偏移；SE(3)-Transformer 的預測在任意旋轉下保持一致。SE(3)-Transformer 的均方誤差（MSE）比非等變基線低了接近一個數量級，驗證了等變性不只是數學上的優雅性質，而是真實帶來了預測精度的提升。

三維物體分類（ScanObjectNN）

ScanObjectNN 是一個真實世界的三維物體點雲分類資料集，包含 2,902 個物體的 15 個類別。這個任務的特殊性在於：現實中物體有重力方向的偏好（「朝上」的方向是有意義的），並非完全旋轉不變。

實驗顯示，完全 SO(3) 不變的模型（忽略重力方向）在未做資料增強時性能損失明顯，而 SE(3)-Transformer +z（額外提供 z 軸方向作為對稱性破缺的提示）能夠學習到應該尊重哪些對稱性，在使用資料增強訓練後展現出最穩健的泛化性能。

分子性質預測（QM9）

QM9 是分子機器學習的標準 Benchmark，包含 134,000 個分子，任務是預測多個量子化學性質。SE(3)-Transformer 在多個回歸任務上達到了與當時最先進的等變模型（Cormorant、TFN）相當的性能，在部分任務上略遜於 LieConv，但優於不具備等變性的基線模型。

更重要的是，SE(3)-Transformer 的等變誤差（equivariance error，衡量模型輸出對旋轉操作的一致性）接近浮點計算精度的極限，說明架構設計確實在數值上嚴格保持了 SE(3) 等變性。

SE(3)-Transformer 把等變性理論推進了一步：從抽象的排列對稱性，到具體的三維幾何對稱性。型態特徵的分層結構（type-ℓ features）是等變性理論在三維空間中的具體實現。每個型態對應 SO(3) 的一個不可約表示，Wigner-D 矩陣與 Clebsch-Gordan 係數是保持旋轉等變性的數學工具，球諧函數是等變濾波器的角度基底。

在這個意義上，TFN 與 SE(3)-Transformer 可以被理解為：把 Maron et al. 的等變性框架（找出所有保持對稱性的線性層）應用到 SO(3) 群而非置換群（permutation group）的具體實現。

等變性理論的探索到此告一段落。無論是上一篇筆記的高階等變網路（突破 WL 上界），還是本篇筆記的 SE(3)-Transformer（保證三維幾何等變性），它們都嘗試在「訊息傳遞框架之外」尋找更強大架構。

然而，如同上一篇筆記中提到，訊息傳遞框架在絕大多數應用場景中仍然佔據主導地位。因此，在實作中，我們更常遇到的問題，是即使一個模型在理論上具備 1-WL 的表達能力，仍然會面臨的兩項常見失效模式：

• Over-smoothing：為什麼堆疊更多層反而讓性能下降？
• Over-squashing：長距離資訊傳播在哪裡遭遇瓶頸？

這兩個問題將由下一篇筆記展開討論。

⏩ 下篇傳送門

[embed]GNN (21) Over-Smoothing：不再被社會同化 在幾乎所有的 GNN 實作中，當你嘗試把網路堆疊更多層時，性能不升反降，超過某個層數後便急劇崩潰。這個現象早在 GCN 的原始論文（Kipf & Welling, 2017）就已被注意到，當時論文僅使用 2 層…medium.com

Reference

1. 謝秉翰（2026）。圖神經網路概論〔課程〕。國立臺灣大學。
2. Fuchs, F. B., Worrall, D. E., Fischer, V., & Welling, M. (2020). SE(3)-Transformers: 3D roto-translation equivariant attention networks. In Advances in Neural Information Processing Systems (NeurIPS), 33.
3. Maron, H., Ben-Hamu, H., Shamir, N., & Lipman, Y. (2019). Invariant and equivariant graph networks. In Proceedings of the 7th International Conference on Learning Representations (ICLR).
4. Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., & Riley, P. (2018). Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. arXiv preprint arXiv:1802.08219.
5. Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., & Bengio, Y. (2018). Graph attention networks. In Proceedings of the 6th International Conference on Learning Representations (ICLR).
6. Xu, K., Hu, W., Leskovec, J., & Jegelka, S. (2019). How powerful are graph neural networks? In Proceedings of the 7th International Conference on Learning Representations (ICLR).
7. Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R., & Smola, A. (2017). Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), 30.
8. Anthropic. (2026). Claude Sonnet 4.6 [Large language model]. https://claude.ai