math-skill 2.1.0 → 3.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (234) hide show
  1. package/README.en-US.md +159 -200
  2. package/README.md +159 -200
  3. package/agents/math-critic.en.md +2 -2
  4. package/agents/math-critic.md +8 -8
  5. package/commands/ask.md +3 -10
  6. package/design-patterns/attention/equivariant-attention.en.md +72 -0
  7. package/design-patterns/attention/equivariant-attention.md +72 -0
  8. package/design-patterns/attention/geometry-aware-attention.en.md +73 -0
  9. package/design-patterns/attention/geometry-aware-attention.md +72 -0
  10. package/design-patterns/attention/information-bottleneck-attention.en.md +72 -0
  11. package/design-patterns/attention/information-bottleneck-attention.md +72 -0
  12. package/design-patterns/attention/projection-attention.en.md +68 -0
  13. package/design-patterns/attention/projection-attention.md +68 -0
  14. package/design-patterns/attention/spectral-attention.en.md +71 -0
  15. package/design-patterns/attention/spectral-attention.md +71 -0
  16. package/design-patterns/compression/leverage-score-selection.en.md +71 -0
  17. package/design-patterns/compression/leverage-score-selection.md +71 -0
  18. package/design-patterns/compression/low-rank-kv-cache.en.md +60 -0
  19. package/design-patterns/compression/low-rank-kv-cache.md +60 -0
  20. package/design-patterns/compression/spectral-token-pruning.en.md +61 -0
  21. package/design-patterns/compression/spectral-token-pruning.md +61 -0
  22. package/design-patterns/compression/topology-preserving-compression.en.md +64 -0
  23. package/design-patterns/compression/topology-preserving-compression.md +64 -0
  24. package/design-patterns/loss/constraint-penalty.en.md +67 -0
  25. package/design-patterns/loss/constraint-penalty.md +78 -0
  26. package/design-patterns/loss/contrastive-loss.en.md +59 -0
  27. package/design-patterns/loss/contrastive-loss.md +68 -0
  28. package/design-patterns/loss/information-bottleneck-loss.en.md +69 -0
  29. package/design-patterns/loss/information-bottleneck-loss.md +74 -0
  30. package/design-patterns/loss/orthogonality-loss.en.md +56 -0
  31. package/design-patterns/loss/orthogonality-loss.md +60 -0
  32. package/design-patterns/loss/variational-loss.en.md +61 -0
  33. package/design-patterns/loss/variational-loss.md +71 -0
  34. package/design-patterns/representation/equivariant-split.en.md +87 -0
  35. package/design-patterns/representation/equivariant-split.md +89 -0
  36. package/design-patterns/representation/manifold-representation.en.md +85 -0
  37. package/design-patterns/representation/manifold-representation.md +87 -0
  38. package/design-patterns/representation/shared-private-decomposition.en.md +84 -0
  39. package/design-patterns/representation/shared-private-decomposition.md +86 -0
  40. package/design-patterns/representation/subspace-alignment.en.md +90 -0
  41. package/design-patterns/representation/subspace-alignment.md +92 -0
  42. package/design-patterns/routing/graph-routing.en.md +79 -0
  43. package/design-patterns/routing/graph-routing.md +81 -0
  44. package/design-patterns/routing/moe-routing.en.md +85 -0
  45. package/design-patterns/routing/moe-routing.md +87 -0
  46. package/design-patterns/routing/optimal-transport-routing.en.md +68 -0
  47. package/design-patterns/routing/optimal-transport-routing.md +72 -0
  48. package/design-patterns/routing/spectral-clustering-routing.en.md +88 -0
  49. package/design-patterns/routing/spectral-clustering-routing.md +90 -0
  50. package/knowledge-base/differential-geometry/connection.en.md +51 -0
  51. package/knowledge-base/differential-geometry/connection.md +51 -0
  52. package/knowledge-base/differential-geometry/curvature.en.md +49 -0
  53. package/knowledge-base/differential-geometry/curvature.md +49 -0
  54. package/knowledge-base/differential-geometry/geodesic.en.md +51 -0
  55. package/knowledge-base/differential-geometry/geodesic.md +51 -0
  56. package/knowledge-base/differential-geometry/manifold.en.md +47 -0
  57. package/knowledge-base/differential-geometry/manifold.md +47 -0
  58. package/knowledge-base/differential-geometry/metric-tensor.en.md +50 -0
  59. package/knowledge-base/differential-geometry/metric-tensor.md +50 -0
  60. package/knowledge-base/differential-geometry/tangent-space.en.md +48 -0
  61. package/knowledge-base/differential-geometry/tangent-space.md +48 -0
  62. package/knowledge-base/information-geometry/fisher-metric.en.md +56 -0
  63. package/knowledge-base/information-geometry/fisher-metric.md +56 -0
  64. package/knowledge-base/information-geometry/natural-gradient.en.md +54 -0
  65. package/knowledge-base/information-geometry/natural-gradient.md +54 -0
  66. package/knowledge-base/lie-theory/equivariance.en.md +54 -0
  67. package/knowledge-base/lie-theory/equivariance.md +54 -0
  68. package/knowledge-base/lie-theory/group-action.en.md +51 -0
  69. package/knowledge-base/lie-theory/group-action.md +51 -0
  70. package/knowledge-base/lie-theory/lie-algebra.en.md +51 -0
  71. package/knowledge-base/lie-theory/lie-algebra.md +51 -0
  72. package/knowledge-base/lie-theory/lie-group.en.md +51 -0
  73. package/knowledge-base/lie-theory/lie-group.md +51 -0
  74. package/knowledge-base/lie-theory/representation.en.md +52 -0
  75. package/knowledge-base/lie-theory/representation.md +52 -0
  76. package/knowledge-base/matrix-analysis/low-rank-approximation.en.md +49 -0
  77. package/knowledge-base/matrix-analysis/low-rank-approximation.md +49 -0
  78. package/knowledge-base/matrix-analysis/matrix-perturbation.en.md +49 -0
  79. package/knowledge-base/matrix-analysis/matrix-perturbation.md +49 -0
  80. package/knowledge-base/matrix-analysis/positive-semidefinite.en.md +50 -0
  81. package/knowledge-base/matrix-analysis/positive-semidefinite.md +50 -0
  82. package/knowledge-base/matrix-analysis/projection.en.md +47 -0
  83. package/knowledge-base/matrix-analysis/projection.md +47 -0
  84. package/knowledge-base/matrix-analysis/spectral-decomposition.en.md +48 -0
  85. package/knowledge-base/matrix-analysis/spectral-decomposition.md +48 -0
  86. package/knowledge-base/optimization/constrained-optimization.en.md +53 -0
  87. package/knowledge-base/optimization/constrained-optimization.md +53 -0
  88. package/knowledge-base/optimization/convex-optimization.en.md +50 -0
  89. package/knowledge-base/optimization/convex-optimization.md +50 -0
  90. package/knowledge-base/optimization/lagrangian-duality.en.md +50 -0
  91. package/knowledge-base/optimization/lagrangian-duality.md +50 -0
  92. package/knowledge-base/optimization/proximal-method.en.md +52 -0
  93. package/knowledge-base/optimization/proximal-method.md +52 -0
  94. package/knowledge-base/optimization/riemannian-optimization.en.md +51 -0
  95. package/knowledge-base/optimization/riemannian-optimization.md +51 -0
  96. package/knowledge-base/overview.en.md +29 -218
  97. package/knowledge-base/overview.md +29 -220
  98. package/knowledge-base/probability/concentration-inequality.en.md +45 -0
  99. package/knowledge-base/probability/concentration-inequality.md +45 -0
  100. package/knowledge-base/probability/entropy.en.md +47 -0
  101. package/knowledge-base/probability/entropy.md +47 -0
  102. package/knowledge-base/probability/fisher-information.en.md +48 -0
  103. package/knowledge-base/probability/fisher-information.md +48 -0
  104. package/knowledge-base/probability/information-bottleneck.en.md +49 -0
  105. package/knowledge-base/probability/information-bottleneck.md +49 -0
  106. package/knowledge-base/probability/kl-divergence.en.md +51 -0
  107. package/knowledge-base/probability/kl-divergence.md +51 -0
  108. package/knowledge-base/topology/euler-characteristic.en.md +54 -0
  109. package/knowledge-base/topology/euler-characteristic.md +54 -0
  110. package/knowledge-base/topology/fundamental-group.en.md +52 -0
  111. package/knowledge-base/topology/fundamental-group.md +52 -0
  112. package/knowledge-base/topology/persistent-homology.en.md +54 -0
  113. package/knowledge-base/topology/persistent-homology.md +54 -0
  114. package/lenses/algorithmic.en.md +49 -0
  115. package/lenses/algorithmic.md +49 -0
  116. package/lenses/axiomatization.en.md +48 -0
  117. package/lenses/axiomatization.md +48 -0
  118. package/lenses/categorical.en.md +48 -0
  119. package/lenses/categorical.md +48 -0
  120. package/lenses/causal.en.md +49 -0
  121. package/lenses/causal.md +49 -0
  122. package/lenses/duality.en.md +47 -0
  123. package/lenses/duality.md +47 -0
  124. package/lenses/game.en.md +50 -0
  125. package/lenses/game.md +50 -0
  126. package/lenses/geometric.en.md +53 -0
  127. package/lenses/geometric.md +53 -0
  128. package/lenses/local-to-global.en.md +52 -0
  129. package/lenses/local-to-global.md +52 -0
  130. package/lenses/perturbation.en.md +51 -0
  131. package/lenses/perturbation.md +51 -0
  132. package/lenses/probabilistic.en.md +48 -0
  133. package/lenses/probabilistic.md +48 -0
  134. package/lenses/projection.en.md +50 -0
  135. package/lenses/projection.md +50 -0
  136. package/lenses/spectral.en.md +52 -0
  137. package/lenses/spectral.md +52 -0
  138. package/lenses/symmetry.en.md +49 -0
  139. package/lenses/symmetry.md +49 -0
  140. package/lenses/topological.en.md +48 -0
  141. package/lenses/topological.md +48 -0
  142. package/lenses/variational.en.md +48 -0
  143. package/lenses/variational.md +48 -0
  144. package/package.json +9 -3
  145. package/references/books/abstract-algebra.md +5 -5
  146. package/references/books/algebraic-geometry-rising-sea.md +5 -5
  147. package/references/books/differential-geometry.md +3 -3
  148. package/references/books/matrix-analysis.md +4 -4
  149. package/references/books/micro-lie-theory.md +3 -3
  150. package/references/books/optimization-ml.md +1 -1
  151. package/references/books/smooth-manifolds.md +3 -3
  152. package/references/gpu-friendly-math.en.md +8 -8
  153. package/references/gpu-friendly-math.md +2 -2
  154. package/references/inspiration.en.md +1 -1
  155. package/references/inspiration.md +2 -2
  156. package/skills/math-research-activator/SKILL.en.md +161 -106
  157. package/skills/math-research-activator/SKILL.md +162 -110
  158. package/skills/math-research-activator/original-texts.en.md +9 -9
  159. package/skills/math-research-activator/original-texts.md +11 -11
  160. package/commands/abstraction.md +0 -11
  161. package/commands/algorithmic-thinking.md +0 -11
  162. package/commands/axiomatization.md +0 -11
  163. package/commands/causal-inference.md +0 -11
  164. package/commands/discrete-combinatorial.md +0 -11
  165. package/commands/game-theory.md +0 -11
  166. package/commands/induction-analogy.md +0 -11
  167. package/commands/information-theory.md +0 -11
  168. package/commands/logic-deduction.md +0 -11
  169. package/commands/modeling.md +0 -11
  170. package/commands/optimization.md +0 -11
  171. package/commands/probability-statistics.md +0 -11
  172. package/commands/symmetry-invariance.md +0 -11
  173. package/commands/topological-thinking.md +0 -11
  174. package/commands/transformation.md +0 -11
  175. package/skills/abstraction/SKILL.en.md +0 -117
  176. package/skills/abstraction/SKILL.md +0 -121
  177. package/skills/abstraction/original-texts.en.md +0 -163
  178. package/skills/abstraction/original-texts.md +0 -175
  179. package/skills/algorithmic-thinking/SKILL.en.md +0 -132
  180. package/skills/algorithmic-thinking/SKILL.md +0 -138
  181. package/skills/algorithmic-thinking/original-texts.en.md +0 -253
  182. package/skills/algorithmic-thinking/original-texts.md +0 -256
  183. package/skills/axiomatization/SKILL.en.md +0 -144
  184. package/skills/axiomatization/SKILL.md +0 -151
  185. package/skills/axiomatization/original-texts.en.md +0 -154
  186. package/skills/axiomatization/original-texts.md +0 -160
  187. package/skills/causal-inference/SKILL.en.md +0 -147
  188. package/skills/causal-inference/SKILL.md +0 -151
  189. package/skills/causal-inference/original-texts.en.md +0 -136
  190. package/skills/causal-inference/original-texts.md +0 -167
  191. package/skills/discrete-combinatorial/SKILL.en.md +0 -124
  192. package/skills/discrete-combinatorial/SKILL.md +0 -131
  193. package/skills/discrete-combinatorial/original-texts.en.md +0 -184
  194. package/skills/discrete-combinatorial/original-texts.md +0 -185
  195. package/skills/game-theory/SKILL.en.md +0 -117
  196. package/skills/game-theory/SKILL.md +0 -123
  197. package/skills/game-theory/original-texts.en.md +0 -131
  198. package/skills/game-theory/original-texts.md +0 -131
  199. package/skills/induction-analogy/SKILL.en.md +0 -145
  200. package/skills/induction-analogy/SKILL.md +0 -152
  201. package/skills/induction-analogy/original-texts.en.md +0 -140
  202. package/skills/induction-analogy/original-texts.md +0 -154
  203. package/skills/information-theory/SKILL.en.md +0 -134
  204. package/skills/information-theory/SKILL.md +0 -140
  205. package/skills/information-theory/original-texts.en.md +0 -127
  206. package/skills/information-theory/original-texts.md +0 -127
  207. package/skills/logic-deduction/SKILL.en.md +0 -130
  208. package/skills/logic-deduction/SKILL.md +0 -135
  209. package/skills/logic-deduction/original-texts.en.md +0 -160
  210. package/skills/logic-deduction/original-texts.md +0 -173
  211. package/skills/modeling/SKILL.en.md +0 -135
  212. package/skills/modeling/SKILL.md +0 -139
  213. package/skills/modeling/original-texts.en.md +0 -162
  214. package/skills/modeling/original-texts.md +0 -165
  215. package/skills/optimization/SKILL.en.md +0 -129
  216. package/skills/optimization/SKILL.md +0 -135
  217. package/skills/optimization/original-texts.en.md +0 -167
  218. package/skills/optimization/original-texts.md +0 -168
  219. package/skills/probability-statistics/SKILL.en.md +0 -146
  220. package/skills/probability-statistics/SKILL.md +0 -151
  221. package/skills/probability-statistics/original-texts.en.md +0 -191
  222. package/skills/probability-statistics/original-texts.md +0 -193
  223. package/skills/symmetry-invariance/SKILL.en.md +0 -135
  224. package/skills/symmetry-invariance/SKILL.md +0 -139
  225. package/skills/symmetry-invariance/original-texts.en.md +0 -206
  226. package/skills/symmetry-invariance/original-texts.md +0 -221
  227. package/skills/topological-thinking/SKILL.en.md +0 -124
  228. package/skills/topological-thinking/SKILL.md +0 -128
  229. package/skills/topological-thinking/original-texts.en.md +0 -134
  230. package/skills/topological-thinking/original-texts.md +0 -164
  231. package/skills/transformation/SKILL.en.md +0 -120
  232. package/skills/transformation/SKILL.md +0 -124
  233. package/skills/transformation/original-texts.en.md +0 -204
  234. package/skills/transformation/original-texts.md +0 -216
@@ -0,0 +1,72 @@
1
+ # Equivariant Attention
2
+ > **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as ✅ verified / ⚠️ retrofittable (needs validation) / ❌ infeasible. Unmarked claims are theoretically possible but require engineering validation.
3
+
4
+ ## Applicable Problems
5
+ When the input possesses an **explicit symmetry group $G$ action** (rotation, translation, permutation, reflection, etc.), and the desired model output should be **equivariant** (covariant) rather than invariant under the same transformations, equivariant constraints must be directly encoded into the attention mechanism. Typical scenarios include: 3D point clouds / molecules ($E(3)$ rigid-body group), image classification ($D_n$ rotation/reflection group), set data ($S_n$ permutation group), and multi-view / multi-sensor fusion.
6
+
7
+ ## Mathematical Inspiration
8
+ - Lenses: [symmetry, categorical (unified framework for group actions)]
9
+ - Knowledge: [`../../knowledge-base/probability/concentration-inequality.md` (sample efficiency gains under equivariant constraints -- data equivalence along orbits), `../../knowledge-base/probability/entropy.md` (equivariant constraints reduce output distribution entropy, yielding stronger inductive bias)]
10
+
11
+ ## Required Mathematical Knowledge
12
+ - **Group Representation Theory Basics**: Linear representation of a group $G$, $\rho: G \to GL(V)$, and irreducible representation decomposition
13
+ - **Equivariant Map Definition**: $f(g \cdot x) = \rho_{\text{out}}(g) \cdot f(x)$ for all $g \in G$
14
+ - **Orbit-Stabilizer Theorem** (see `references/books/abstract-algebra.md` Ch.5): $|orbit| = |G|/|stab|$, giving the parameter sharing multiplier
15
+ - **Schur's Lemma**: An equivariant linear map between irreducible representations is either zero or a scalar multiple
16
+
17
+ ## AI Module Form
18
+
19
+ **Core Idea**: Replace the standard attention $Q, K, V$ with **steerable features**, ensuring that attention weights are invariant under group actions and that outputs are equivariant under group actions.
20
+
21
+ **Scheme A: Permutation-Equivariant Attention ($S_n$ group, set data)**:
22
+ ```python
23
+ # DeepSets / Set Transformer style
24
+ # Attention weights are permutation-invariant: pi(Q)pi(K)^T = QK^T (permutations cancel)
25
+ # Output is permutation-equivariant: pi(softmax(QK^T) V) = softmax(QK^T) pi(V)
26
+ Q, K, V = W_q(X), W_k(X), W_v(X) # pointwise linear transform
27
+ scores = Q @ K.T / sqrt(d) # permutation-invariant
28
+ attn = softmax(scores) # permutation-invariant
29
+ output = attn @ V # permutation-equivariant (V is equivariant)
30
+ ```
31
+
32
+ **Scheme B: $E(3)$-Equivariant Attention (3D point clouds / molecules)**:
33
+ ```python
34
+ # Decompose features into scalars + vectors + higher-order tensors (spherical harmonic basis)
35
+ # Attention weights computed using only scalar features (rotation-invariant)
36
+ scalar_Q = scalar_proj(X_scalar) # scalars only -> rotation-invariant
37
+ scalar_K = scalar_proj(X_scalar)
38
+ scores = scalar_Q @ scalar_K.T / sqrt(d_s) # rotation-invariant attention weights
39
+
40
+ # V contains equivariant features (scalars + vectors), weighted by invariant weights
41
+ output_scalar = softmax(scores) @ V_scalar # scalar -> invariant
42
+ output_vector = softmax(scores) @ V_vector # vector -> equivariant (rotation-covariant)
43
+ ```
44
+
45
+ **Scheme C: $D_n$-Equivariant Attention (image rotation / reflection)**:
46
+ ```python
47
+ # G-CNN style: for each group element g in D_n, transform input with rho(g) then compute attention
48
+ # Weights shared along orbits (same W_q/W_k/W_v), aggregate outputs over all group elements
49
+ output = mean(softmax((rho(g)@X@W_q) @ (rho(g)@X@W_k).T/sqrt(d)) @ (rho(g)@X@W_v) for g in D_n)
50
+ ```
51
+
52
+ ## Implementable Architectures
53
+ - **SE(3)-Transformer / Equiformer**: Spherical harmonic features + equivariant attention for molecular property prediction and protein structure
54
+ - **Set Transformer**: $S_n$ permutation-equivariant attention + Induced Set Attention (low-rank inducing points for complexity reduction)
55
+ - **G-CNN Attention**: $D_n$ rotation/reflection equivariance for remote sensing imagery and medical imaging
56
+
57
+ ## GPU Feasibility
58
+ - **Dimension 1 Tensorization**: Group actions implemented as $\rho(g)$ matrix multiplications; equivariant features stored as batched tensors
59
+ - **Dimension 2 GEMM-mappability**: $\rho(g) X$ and $Q K^T$ are both GEMM operations; $|G|$ group elements map to batched GEMM
60
+ - **Dimension 3 Complexity**: $|G|$-fold computation overhead; acceptable for small groups ($|D_4|=8$), infeasible for large groups ($|S_n|=n!$). Remedy: use generators + Cayley graph propagation instead of full group enumeration
61
+ - **Dimension 4 Memory**: Requires storing $|G|$ copies of intermediate features; mitigated by chunking + gradient checkpointing
62
+ - **Dimension 5 Low Precision**: Orthogonal representation matrices are numerically stable under bf16
63
+ - **Dimension 6 Parallelism**: $|G|$ group elements are naturally parallelizable (along the batch dimension)
64
+ - **Dimension 7 Sparsity**: Permutation $\rho(g)$ is extremely sparse and can be encoded as gather indices
65
+ - **Dimension 8 Operator Fusion**: Group action + linear duality can be fused into a single batched GEMM
66
+
67
+ ## Paper Phrasing
68
+ "We propose an equivariant attention mechanism that constrains attention weights to be group invariants and attention outputs to be group equivariants, directly encoding the inductive bias of symmetry group $G$ into the model architecture without additional data augmentation, achieving a $|G|/|stab|$-fold improvement in parameter efficiency."
69
+
70
+ ## Risks
71
+ - **Cost of Incorrect Group Selection**: If the data does not possess the assumed symmetry (e.g., molecules lacking full $E(3)$ symmetry), equivariant constraints will impair expressiveness. The symmetry assumption must be validated first, or "approximate equivariance" (soft equivariance) should be used.
72
+ - **Memory Explosion of High-Order Representations**: The $L$-th order spherical harmonic representation of $SO(3)$ has dimension $(2L+1)^2$, causing storage and computation to grow rapidly for high-order features ($L \geq 3$). In practice, truncation to $L \leq 2$ is standard.
@@ -0,0 +1,72 @@
1
+ # 等变注意力 / Equivariant Attention
2
+ > **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「✅ 已验证 / ⚠️ 可改造需验证 / ❌ 不可行」。未标注的视为理论可行,需工程验证。
3
+
4
+ ## 适用问题
5
+ 当输入具有**明确的对称群 $G$ 作用**(旋转、平移、置换、反射等),且期望模型输出在相同变换下**协变**(equivariant)而非不变时,需要将等变约束直接编入注意力机制。典型场景:3D 点云/分子($E(3)$ 刚体群)、图像分类($D_n$ 旋转/镜像群)、集合数据($S_n$ 置换群)、多视角/多传感器融合。
6
+
7
+ ## 数学思想来源
8
+ - 透镜:[symmetry(对称透镜), categorical(范畴化透镜 — 群作用的统一框架)]
9
+ - 知识:[`../../knowledge-base/probability/concentration-inequality.md`(等变约束下的样本效率提升 — 轨道上的数据等价性), `../../knowledge-base/probability/entropy.md`(等变约束降低输出分布熵 → 更强的归纳偏置)]
10
+
11
+ ## 需要的数学知识
12
+ - **群表示论基础**:群 $G$ 的线性表示 $\rho: G \to GL(V)$,不可约表示分解
13
+ - **等变映射定义**:$f(g \cdot x) = \rho_{\text{out}}(g) \cdot f(x)$,对所有 $g \in G$
14
+ - **轨道-稳定子定理**(参见 `references/books/abstract-algebra.md` Ch.5):$|orbit| = |G|/|stab|$,参数共享倍率
15
+ - **Schur 引理**:不可约表示间的等变线性映射要么为零要么为标量乘
16
+
17
+ ## AI 模块形式
18
+
19
+ **核心思路**:将标准注意力的 $Q, K, V$ 替换为**等变特征**(steerable features),确保注意力权重在群作用下不变,输出在群作用下等变。
20
+
21
+ **方案 A:置换等变注意力($S_n$ 群,集合数据)**:
22
+ ```python
23
+ # DeepSets / Set Transformer 风格
24
+ # 注意力权重对置换不变:π(Q)π(K)^T = QK^T(置换抵消)
25
+ # 输出对置换等变:π(softmax(QK^T) V) = softmax(QK^T) π(V)
26
+ Q, K, V = W_q(X), W_k(X), W_v(X) # 逐点线性变换
27
+ scores = Q @ K.T / sqrt(d) # 置换不变
28
+ attn = softmax(scores) # 置换不变
29
+ output = attn @ V # 置换等变(V 是等变的)
30
+ ```
31
+
32
+ **方案 B:$E(3)$ 等变注意力(3D 点云/分子)**:
33
+ ```python
34
+ # 将特征分解为标量 + 向量 + 高阶张量(球谐基)
35
+ # 注意力权重只用标量特征计算(旋转不变)
36
+ scalar_Q = scalar_proj(X_scalar) # 仅标量 → 旋转不变
37
+ scalar_K = scalar_proj(X_scalar)
38
+ scores = scalar_Q @ scalar_K.T / sqrt(d_s) # 旋转不变的注意力权重
39
+
40
+ # V 包含等变特征(标量 + 向量),用不变的权重加权
41
+ output_scalar = softmax(scores) @ V_scalar # 标量 → 不变
42
+ output_vector = softmax(scores) @ V_vector # 向量 → 等变(旋转协变)
43
+ ```
44
+
45
+ **方案 C:$D_n$ 等变注意力(图像旋转/镜像)**:
46
+ ```python
47
+ # G-CNN 风格:对每个群元素 g ∈ D_n,用 ρ(g) 变换输入后计算注意力
48
+ # 权重沿轨道共享(同一组 W_q/W_k/W_v),聚合所有群元素输出
49
+ output = mean(softmax((rho(g)@X@W_q) @ (rho(g)@X@W_k).T/sqrt(d)) @ (rho(g)@X@W_v) for g in D_n)
50
+ ```
51
+
52
+ ## 可实现结构
53
+ - **SE(3)-Transformer / Equiformer**:球谐特征 + 等变注意力,用于分子性质预测和蛋白质结构
54
+ - **Set Transformer**:$S_n$ 置换等变注意力 + Induced Set Attention(低秩诱导点降复杂度)
55
+ - **G-CNN 注意力**:$D_n$ 旋转/镜像等变,用于遥感图像和医学影像
56
+
57
+ ## GPU 可行性
58
+ - **维度 1 张量化 ✅**:群作用实现为 $\rho(g)$ 矩阵乘,等变特征为 batched 张量
59
+ - **维度 2 GEMM 可映射 ✅**:$\rho(g) X$ 和 $Q K^T$ 均为 GEMM;$|G|$ 个群元素 → batched GEMM
60
+ - **维度 3 复杂度 ⚠️**:$|G|$ 倍计算量,小群($|D_4|=8$)可接受,大群($|S_n|=n!$)不可。改造:用生成元 + Cayley 图传播替代整群枚举
61
+ - **维度 4 显存 ⚠️**:需存储 $|G|$ 份中间特征,可分块 + gradient checkpointing
62
+ - **维度 5 低精度 ✅**:正交表示矩阵在 bf16 下数值稳定
63
+ - **维度 6 并行 ✅**:$|G|$ 个群元素天然并行(batch 维)
64
+ - **维度 7 稀疏 ⚠️**:置换 $\rho(g)$ 极度稀疏,可编码为 gather 索引
65
+ - **维度 8 算子融合 ✅**:群作用 + 线性变换可融合为单次 batched GEMM
66
+
67
+ ## 论文表述方式
68
+ "我们提出等变注意力机制,通过将注意力权重约束为群不变量、注意力输出约束为群等变量,在不增加数据增强的情况下将对称群 $G$ 的归纳偏置直接编码进模型结构,参数效率提升 $|G|/|stab|$ 倍。"
69
+
70
+ ## 风险
71
+ - **群选择错误的代价**:若数据不具有假定的对称性(如分子不具有完全 $E(3)$ 对称性),等变约束会伤害表达能力。需先验证对称性假设,或使用"近似等变"(soft equivariance)。
72
+ - **高阶表示的显存爆炸**:$SO(3)$ 的 $L$ 阶球谐表示维度为 $(2L+1)^2$,高阶特征 ($L \geq 3$) 的存储和计算量急剧增长。实践中通常截断到 $L \leq 2$。
@@ -0,0 +1,73 @@
1
+ # Geometry-Aware Attention
2
+ > **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as ✅ verified / ⚠️ retrofittable (needs validation) / ❌ infeasible. Unmarked claims are theoretically possible but require engineering validation.
3
+
4
+ ## Applicable Problems
5
+ When **known geometric relationships** exist between tokens/keys (spatial distance, manifold geodesic distance, hierarchical structure, temporal interval), the position-agnostic inner product of standard attention cannot exploit this structural information. Geometry-aware attention **directly injects geometric priors into attention weight computation**, enabling the model to naturally respect the metric structure of the underlying space. Typical scenarios include: 3D scene understanding, molecular conformation, temporal forecasting, hierarchical text structure (paragraph-sentence-word), and knowledge graphs.
6
+
7
+ ## Mathematical Inspiration
8
+ - Lenses: [symmetry (metric invariance), duality (coordinate-system-independent representation)]
9
+ - Knowledge: [`../../knowledge-base/information-geometry/fisher-metric.md` (geometric metrics on distribution spaces), `../../knowledge-base/probability/concentration-inequality.md` (concentration behavior under geometric constraints)]
10
+
11
+ ## Required Mathematical Knowledge
12
+ - **Metric Spaces and Distance Functions**: Euclidean distance, geodesic distance, tree distance, Wasserstein distance
13
+ - **RBF / Kernel Methods**: $k(x, y) = \exp(-d(x,y)^2 / 2\sigma^2)$, converting distance to similarity
14
+ - **Geometric Interpretation of Positional Encoding**: RoPE = action of the rotation group $SO(2)$ (relative distance encoded as rotation angle difference); ALiBi = exponentially decaying distance bias
15
+
16
+ ## AI Module Form
17
+
18
+ **Core Idea**: Explicitly introduce a geometric distance term into attention scores so that attention to distant tokens naturally decays:
19
+
20
+ $$\text{GeoAttn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}} + \text{GeoBias}(i, j)\right) V$$
21
+
22
+ **Scheme A: Distance-Biased Attention (General Metric Spaces)**:
23
+ ```python
24
+ # d_ij is the geometric distance between token i and j (precomputed or computed online)
25
+ # Learnable distance bias function
26
+ distance_bias = MLP(d_ij) # or simply -alpha * d_ij
27
+ scores = (Q @ K.T) / sqrt(d) + distance_bias
28
+ attn = softmax(scores) @ V
29
+ ```
30
+
31
+ **Scheme B: Relative Positional Encoding (Geometric Generalization of RoPE / ALiBi)**:
32
+ ```python
33
+ # Generalize RoPE: position -> group element g_i, relative position -> g_i g_j^{-1}
34
+ # Learnable relative position bias on arbitrary metric spaces
35
+ def geo_attention(Q, K, V, positions):
36
+ rel_pos = pairwise_difference(positions) # (n, n, coord_dim)
37
+ geo_bias = geo_encoder(rel_pos) # (n, n) learnable mapping
38
+ scores = (Q @ K.T) / sqrt(d) + geo_bias
39
+ return softmax(scores) @ V
40
+ ```
41
+
42
+ **Scheme C: Manifold-Aware Attention (Non-Euclidean Spaces)**:
43
+ ```python
44
+ # When tokens reside on a non-Euclidean manifold (hyperbolic space, sphere),
45
+ # use geodesic distance instead of Euclidean distance
46
+ def manifold_attention(Q, K, V, manifold):
47
+ geodist = manifold.geodesic_distance_matrix(positions) # (n, n)
48
+ geo_kernel = exp(-geodist^2 / (2 * sigma^2)) # RBF kernel
49
+ scores = (Q @ K.T) / sqrt(d) + log(geo_kernel + eps)
50
+ return softmax(scores) @ V
51
+ ```
52
+
53
+ ## Implementable Architectures
54
+ - **RoPE Extensions**: From $SO(2)$ rotations to higher-dimensional rotation groups $SO(2k)$, encoding multi-dimensional positional information (2D image patches, 3D voxels)
55
+ - **Hierarchical Positional Bias**: Use LCA (Lowest Common Ancestor) depth as distance in tree structures, suitable for hierarchical geometric of documents and code
56
+ - **Molecular Conformation Attention**: 3D atomic coordinates -> distance matrix -> geometric bias, used in molecular GNNs and protein structure prediction
57
+
58
+ ## GPU Feasibility
59
+ - **Dimension 1 Tensorization**: Distance matrices and bias matrices are dense tensors with element-wise operations
60
+ - **Dimension 2 GEMM-mappability**: The main body $Q K^T$ is standard GEMM; geometric bias is additive and does not interrupt GEMM
61
+ - **Dimension 3 Complexity**: Pairwise distance matrix construction and storage are $O(n^2)$; however, block-wise computation + online softmax can avoid materializing the full matrix
62
+ - **Dimension 4 Memory**: The $n \times n$ distance matrix occupies GPU memory; mitigated by block/streaming computation (compatible with FlashAttention)
63
+ - **Dimension 5 Low Precision**: Distance computation and bias addition are stable under bf16; the exp in RBF requires overflow caution (clamp distances)
64
+ - **Dimension 6 Parallelism**: Distance computation and attention can be pipelined; pairwise distances can be computed in parallel blocks
65
+ - **Dimension 7 Sparsity**: Distant biases tend toward $-\infty$ (approaching zero after softmax), naturally inducing structured sparsity (local attention windows)
66
+ - **Dimension 8 Operator Fusion**: Geometric bias can be fused into FlashAttention's online softmax loop (added after $QK^T$, before softmax)
67
+
68
+ ## Paper Phrasing
69
+ "We propose geometry-aware attention, which explicitly injects learnable geometric distance bias terms into attention scores, enabling the model to naturally respect the metric structure of the input space. This achieves exponential decay of attention to distant tokens without increasing parameter count, while maintaining compatibility with FlashAttention."
70
+
71
+ ## Risks
72
+ - **Geometric Prior Conflicting with Data**: If geometric distance is inconsistent with semantic relevance (e.g., distant but semantically related tokens in text), strong geometric bias will impair model expressiveness. The bias should be learnable and overridable by content-based attention.
73
+ - **Non-Differentiability of Distance Computation**: Certain distances (graph shortest path, tree LCA depth) are non-differentiable or computationally expensive. Differentiable relaxations (softmin instead of min) or precomputed distance matrices are needed.
@@ -0,0 +1,72 @@
1
+ # 几何感知注意力 / Geometry-Aware Attention
2
+ > **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「✅ 已验证 / ⚠️ 可改造需验证 / ❌ 不可行」。未标注的视为理论可行,需工程验证。
3
+
4
+ ## 适用问题
5
+ 当 token/key 之间存在**已知的几何关系**(空间距离、流形测地距、层级结构、时序间隔)时,标准注意力的位置无关内积无法利用这些结构信息。几何感知注意力将**几何先验直接注入注意力权重计算**,使模型天然尊重底空间的度量结构。典型场景:3D 场景理解、分子构象、时序预测、层级文本结构(段落-句子-词)、知识图谱。
6
+
7
+ ## 数学思想来源
8
+ - 透镜:[symmetry(对称透镜 — 度量不变性), duality(对偶透镜 — 坐标系无关表达)]
9
+ - 知识:[`../../knowledge-base/information-geometry/fisher-metric.md`(分布空间的几何度量), `../../knowledge-base/probability/concentration-inequality.md`(几何约束下的浓度行为)]
10
+
11
+ ## 需要的数学知识
12
+ - **度量空间与距离函数**:欧氏距离、测地距、树距离、Wasserstein 距离
13
+ - **RBF / 核方法**:$k(x, y) = \exp(-d(x,y)^2 / 2\sigma^2)$,距离→相似度的转换
14
+ - **位置编码的几何解释**:RoPE = 旋转群 $SO(2)$ 的作用(相对距离编码为旋转角度差);ALiBi = 指数衰减的距离偏置
15
+
16
+ ## AI 模块形式
17
+
18
+ **核心思路**:在注意力分数中显式引入几何距离项,使远距离 token 的注意力天然衰减:
19
+
20
+ $$\text{GeoAttn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}} + \text{GeoBias}(i, j)\right) V$$
21
+
22
+ **方案 A:距离偏置注意力(通用度量空间)**:
23
+ ```python
24
+ # d_ij 为 token i 和 j 之间的几何距离(预计算或在线计算)
25
+ # 可学习距离偏置函数
26
+ distance_bias = MLP(d_ij) # 或简单的 -α * d_ij
27
+ scores = (Q @ K.T) / sqrt(d) + distance_bias
28
+ attn = softmax(scores) @ V
29
+ ```
30
+
31
+ **方案 B:相对位置编码(RoPE / ALiBi 的几何推广)**:
32
+ ```python
33
+ # 推广 RoPE:位置 → 群元素 g_i,相对位置 → g_i g_j^{-1}
34
+ # 可学习任意度量空间上的相对位置偏置
35
+ def geo_attention(Q, K, V, positions):
36
+ rel_pos = pairwise_difference(positions) # (n, n, coord_dim)
37
+ geo_bias = geo_encoder(rel_pos) # (n, n) 可学习映射
38
+ scores = (Q @ K.T) / sqrt(d) + geo_bias
39
+ return softmax(scores) @ V
40
+ ```
41
+
42
+ **方案 C:流形感知注意力(非欧空间)**:
43
+ ```python
44
+ # token 位于非欧流形(双曲空间、球面)时,用测地距替代欧氏距
45
+ def manifold_attention(Q, K, V, manifold):
46
+ geodist = manifold.geodesic_distance_matrix(positions) # (n, n)
47
+ geo_kernel = exp(-geodist^2 / (2 * sigma^2)) # RBF 核
48
+ scores = (Q @ K.T) / sqrt(d) + log(geo_kernel + eps)
49
+ return softmax(scores) @ V
50
+ ```
51
+
52
+ ## 可实现结构
53
+ - **RoPE 扩展**:从 $SO(2)$ 旋转到更高维旋转群 $SO(2k)$,编码多维位置信息(2D 图像 patch、3D voxel)
54
+ - **层级位置偏置**:树结构中用 LCA(最近公共祖先)深度作为距离,适合文档/代码的层级建模
55
+ - **分子构象注意力**:3D 原子坐标 → 距离矩阵 → 几何偏置,用于分子 GNN 和蛋白质结构预测
56
+
57
+ ## GPU 可行性
58
+ - **维度 1 张量化 ✅**:距离矩阵和偏置矩阵均为稠密张量,逐元素运算
59
+ - **维度 2 GEMM 可映射 ✅**:主体 $Q K^T$ 为标准 GEMM;几何偏置为加法,不阻断 GEMM
60
+ - **维度 3 复杂度 ⚠️**:成对距离矩阵 $O(n^2)$ 构建和存储;但可用分块计算 + online softmax 避免物化完整矩阵
61
+ - **维度 4 显存 ⚠️**:$n \times n$ 距离矩阵占用显存;可分块/流式计算(与 FlashAttention 兼容)
62
+ - **维度 5 低精度 ✅**:距离计算和偏置加法在 bf16 下稳定;RBF 的 exp 需注意上溢(clamp 距离)
63
+ - **维度 6 并行 ✅**:距离计算和注意力可流水线并行,成对距离可分块并行
64
+ - **维度 7 稀疏 ✅**:远距离偏置趋于 $-\infty$(softmax 后趋零),天然诱导结构化稀疏(局部注意力窗口)
65
+ - **维度 8 算子融合 ✅**:几何偏置可融入 FlashAttention 的 online softmax 循环中(加在 $QK^T$ 之后、softmax 之前)
66
+
67
+ ## 论文表述方式
68
+ "我们提出几何感知注意力,通过在注意力分数中显式注入可学习的几何距离偏置项,使模型天然尊重输入空间的度量结构,在不增加参数量的情况下将远距离 token 的注意力指数衰减,同时保持与 FlashAttention 的兼容性。"
69
+
70
+ ## 风险
71
+ - **几何先验与数据冲突**:若几何距离与语义相关性不一致(如文本中远距离但语义相关的 token),强几何偏置会伤害模型表达。需让偏置可学习且可被内容注意力覆盖。
72
+ - **距离计算的非可微性**:某些距离(图最短路径、树 LCA 深度)不可微或计算复杂。需使用可微松弛(softmin 替代 min)或预计算距离矩阵。
@@ -0,0 +1,72 @@
1
+ # Information Bottleneck Attention
2
+ > **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as ✅ verified / ⚠️ retrofittable (needs validation) / ❌ infeasible. Unmarked claims are theoretically possible but require engineering validation.
3
+
4
+ ## Applicable Problems
5
+ When the attention mechanism needs to **selectively transmit useful information while suppressing redundant/noisy information**, information bottleneck theory can guide the learning of attention weights -- maximizing the mutual information $I(Z;Y)$ of the attention distribution with respect to the target $Y$, while minimizing the mutual information $I(X;Z)$ with respect to the input $X$. Typical scenarios include: long-document summarization (filtering large numbers of irrelevant tokens), multimodal alignment (cross-modal noise suppression), and interpretability (attention weights as visualization of information flow).
6
+
7
+ ## Mathematical Inspiration
8
+ - Lenses: [categorical (information-theoretic framework unifying attention design), variational (constrained variational and Lagrangian duality)]
9
+ - Knowledge: [`../../knowledge-base/probability/information-bottleneck.md` (IB objective and variational lower bound), `../../knowledge-base/probability/kl-divergence.md` (implementation of KL regularization), `../../knowledge-base/probability/entropy.md` (mutual information estimation)]
10
+
11
+ ## Required Mathematical Knowledge
12
+ - **Information Bottleneck Objective**: $\min I(X;Z) - \beta I(Z;Y)$, balancing compression and prediction
13
+ - **Variational Information Bottleneck (VIB)**: Replacing intractable mutual information with variational lower bounds
14
+ - **Correspondence Between Mutual Information and Attention**: Softmax attention weights $\alpha_{ij}$ can be interpreted as information channel allocations from key $j$ to query $i$
15
+
16
+ ## AI Module Form
17
+
18
+ **Core Idea**: Treat attention as an information bottleneck -- attention weights determine "how much information flows from values to output," while KL regularization constrains the information throughput.
19
+
20
+ **Scheme A: KL-Regularized Attention (Simplest IB Attention)**:
21
+ ```python
22
+ # KL divergence between attention distribution q(z|x) and uniform prior as IB regularizer
23
+ scores = (Q @ K.T) / sqrt(d)
24
+ attn = softmax(scores)
25
+ uniform_prior = ones_like(attn) / n
26
+ kl_reg = (attn * (log(attn + eps) - log(uniform_prior))).sum(dim=-1).mean()
27
+ loss = task_loss + beta * kl_reg
28
+ output = attn @ V
29
+ ```
30
+
31
+ **Scheme B: Variational Information Bottleneck Attention (VIB-Attention)**:
32
+ ```python
33
+ # Introduce stochastic bottleneck Z ~ q(Z|attention) to limit information throughput
34
+ scores = (Q @ K.T) / sqrt(d)
35
+ attn = softmax(scores)
36
+ mu_z, log_var_z = attn, linear(attn) # mean and learnable variance
37
+ z = softmax(mu_z + exp(0.5 * log_var_z) * randn_like(mu_z)) # reparameterization + softmax
38
+ kl = 0.5 * (mu_z^2 + exp(log_var_z) - log_var_z - 1).sum(-1).mean() # KL vs prior
39
+ output = z @ V
40
+ loss = task_loss + beta * kl
41
+ ```
42
+
43
+ **Scheme C: Mutual Information Maximization Attention (DIM Style)**:
44
+ ```python
45
+ # Directly maximize I(Z;Y) via InfoNCE, with KL constraining I(X;Z)
46
+ Z = softmax((Q @ K.T) / sqrt(d)) @ V
47
+ info_nce = infonce_loss(Z, target_embedding, negatives, tau=0.1)
48
+ kl_bottleneck = estimate_kl(X, Z) # MINE/NWJ estimator
49
+ loss = -info_nce + beta * kl_bottleneck
50
+ ```
51
+
52
+ ## Implementable Architectures
53
+ - **IB-Sparse Attention**: KL regularization serves as a soft constraint on attention sparsity -- larger KL relative to uniform prior leads to more concentrated attention, yielding implicit Top-K selection
54
+ - **IB Interpretation of Dropout**: Dropout is a form of stochastic information bottleneck -- randomly blocking information channels forces the model to learn robust representations. The dropout rate corresponds to the $\beta$ parameter
55
+ - **Multi-Head Information Allocation**: Different heads learn different information bottlenecks (different $\beta$); some heads transmit global information, others only local information
56
+
57
+ ## GPU Feasibility
58
+ - **Dimension 1 Tensorization**: KL regularization involves element-wise operations; VIB reparameterization sampling is element-wise
59
+ - **Dimension 2 GEMM-mappability**: The main body $QK^T$ and $attn \cdot V$ are standard GEMM; regularization introduces no new GEMM operations
60
+ - **Dimension 3 Complexity**: KL regularization is $O(n)$ per token, adding no asymptotic complexity
61
+ - **Dimension 4 Memory**: VIB requires additional $\mu_z$ and $\log\sigma_z$, approximately doubling attention weight memory
62
+ - **Dimension 5 Low Precision**: log/exp in KL computations are stable under bf16 (standard log-softmax tricks)
63
+ - **Dimension 6 Parallelism**: Regularization can be computed in parallel with forward propagation, introducing no serial dependencies
64
+ - **Dimension 7 Sparsity**: KL regularization implicitly induces attention sparsity, amenable to block-sparse acceleration
65
+ - **Dimension 8 Operator Fusion**: KL can be fused into the softmax kernel (FusedSoftmaxKL)
66
+
67
+ ## Paper Phrasing
68
+ "We propose information bottleneck attention, which models attention as an information bottleneck variational problem. By maximizing the mutual information between output and target while minimizing redundant information transmission from the input, it achieves theoretically optimal information selection and demonstrates enhanced sparsity and interpretability in experiments."
69
+
70
+ ## Risks
71
+ - **High Variance of Mutual Information Estimators**: The MINE/NWJ/InfoNCE estimators in Scheme C exhibit high variance in high-dimensional spaces, potentially causing training instability. It is recommended to first validate the basic effect of IB attention with Scheme A (KL regularization) before attempting the full IB objective.
72
+ - **Difficulty of $\beta$ Tuning**: $\beta$ controls the compression-prediction trade-off, and the optimal $\beta$ varies significantly across tasks. Too small a $\beta$ degenerates to standard attention; too large a $\beta$ causes underfitting. Adaptive $\beta$ scheduling or information plane monitoring is recommended.
@@ -0,0 +1,72 @@
1
+ # 信息瓶颈注意力 / Information Bottleneck Attention
2
+ > **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「✅ 已验证 / ⚠️ 可改造需验证 / ❌ 不可行」。未标注的视为理论可行,需工程验证。
3
+
4
+ ## 适用问题
5
+ 当注意力机制需要**选择性传递有用信息、抑制冗余/噪声信息**时,用信息瓶颈理论指导注意力权重的学习——使注意力分布最大化关于目标 $Y$ 的互信息 $I(Z;Y)$,同时最小化关于输入 $X$ 的互信息 $I(X;Z)$。典型场景:长文档摘要(大量无关 token 需被过滤)、多模态对齐(跨模态噪声抑制)、可解释性(注意力权重作为信息流的可视化)。
6
+
7
+ ## 数学思想来源
8
+ - 透镜:[categorical(范畴化透镜 — 信息论框架统一注意力设计), variational(变分透镜 — 约束优化与拉格朗日对偶)]
9
+ - 知识:[`../../knowledge-base/probability/information-bottleneck.md`(IB 目标函数与变分下界), `../../knowledge-base/probability/kl-divergence.md`(KL 正则项的实现), `../../knowledge-base/probability/entropy.md`(互信息估计)]
10
+
11
+ ## 需要的数学知识
12
+ - **信息瓶颈 IB 目标**:$\min I(X;Z) - \beta I(Z;Y)$,压缩-预测权衡
13
+ - **变分信息瓶颈 VIB**:用变分下界替代难以估计的互信息
14
+ - **互信息与注意力的对应**:softmax 注意力权重 $\alpha_{ij}$ 可解释为从 key $j$ 到 query $i$ 的信息信道分配
15
+
16
+ ## AI 模块形式
17
+
18
+ **核心思路**:将注意力视为信息瓶颈——注意力权重决定"让多少信息从 value 流向输出",同时用 KL 正则约束信息量不过大。
19
+
20
+ **方案 A:KL 正则化注意力(最简 IB 注意力)**:
21
+ ```python
22
+ # 注意力分布 q(z|x) 与均匀先验的 KL 散度作为 IB 正则
23
+ scores = (Q @ K.T) / sqrt(d)
24
+ attn = softmax(scores)
25
+ uniform_prior = ones_like(attn) / n
26
+ kl_reg = (attn * (log(attn + eps) - log(uniform_prior))).sum(dim=-1).mean()
27
+ loss = task_loss + beta * kl_reg
28
+ output = attn @ V
29
+ ```
30
+
31
+ **方案 B:变分信息瓶颈注意力(VIB-Attention)**:
32
+ ```python
33
+ # 引入随机瓶颈 Z ~ q(Z|attention),限制信息通过量
34
+ scores = (Q @ K.T) / sqrt(d)
35
+ attn = softmax(scores)
36
+ mu_z, log_var_z = attn, linear(attn) # 均值与可学习方差
37
+ z = softmax(mu_z + exp(0.5 * log_var_z) * randn_like(mu_z)) # 重参数化 + softmax
38
+ kl = 0.5 * (mu_z^2 + exp(log_var_z) - log_var_z - 1).sum(-1).mean() # KL vs 先验
39
+ output = z @ V
40
+ loss = task_loss + beta * kl
41
+ ```
42
+
43
+ **方案 C:互信息最大化注意力(DIM 风格)**:
44
+ ```python
45
+ # 用 InfoNCE 直接最大化 I(Z;Y),配合 KL 约束 I(X;Z)
46
+ Z = softmax((Q @ K.T) / sqrt(d)) @ V
47
+ info_nce = infonce_loss(Z, target_embedding, negatives, tau=0.1)
48
+ kl_bottleneck = estimate_kl(X, Z) # MINE/NWJ 估计器
49
+ loss = -info_nce + beta * kl_bottleneck
50
+ ```
51
+
52
+ ## 可实现结构
53
+ - **IB-Sparse Attention**:将 KL 正则化为注意力稀疏性的软约束——KL 对均匀先越大 → 注意力越集中 → 隐式 Top-K 选择
54
+ - **Dropout 的 IB 解释**:Dropout 是一种随机信息瓶颈——随机阻断部分信息通道,迫使模型学习鲁棒表示。Dropout rate 对应 $\beta$ 参数
55
+ - **Multi-Head 信息分配**:不同 head 学习不同的信息瓶颈(不同 $\beta$),有的 head 传递全局信息,有的只传递局部信息
56
+
57
+ ## GPU 可行性
58
+ - **维度 1 张量化 ✅**:KL 正则项为逐元素运算,VIB 重参数化采样为逐元素
59
+ - **维度 2 GEMM 可映射 ✅**:主体 $QK^T$ 和 $attn \cdot V$ 为标准 GEMM,正则项不引入新 GEMM
60
+ - **维度 3 复杂度 ✅**:KL 正则 $O(n)$ 逐 token,不增加渐近复杂度
61
+ - **维度 4 显存 ⚠️**:VIB 需额外 $\mu_z$ 和 $\log\sigma_z$,约 2x 注意力权重显存
62
+ - **维度 5 低精度 ✅**:KL 中 log/exp 在 bf16 下稳定(标准 log-softmax 技巧)
63
+ - **维度 6 并行 ✅**:正则项可与前向传播并行,不引入串行依赖
64
+ - **维度 7 稀疏 ✅**:KL 正则隐式诱导注意力稀疏性,可利用 block-sparse 加速
65
+ - **维度 8 算子融合 ✅**:KL 可融入 softmax kernel(FusedSoftmaxKL)
66
+
67
+ ## 论文表述方式
68
+ "我们提出信息瓶颈注意力机制,通过将注意力建模为信息瓶颈优化问题,在最大化输出与目标的互信息的同时最小化输入信息的冗余传递,实现理论上最优的信息选择,并在实验中展现出更强的稀疏性和可解释性。"
69
+
70
+ ## 风险
71
+ - **互信息估计的方差问题**:方案 C 中的 MINE/NWJ/InfoNCE 估计器在高维空间中方差大,可能导致训练不稳定。建议先用方案 A(KL 正则)验证 IB 注意力的基本效果,再尝试完整 IB 目标。
72
+ - **$\beta$ 调参困难**:$\beta$ 控制压缩-预测权衡,不同任务最优 $\beta$ 差异大。$\beta$ 过小退化为标准注意力,$\beta$ 过大导致欠拟合。建议自适应 $\beta$ 调度或信息平面监控。
@@ -0,0 +1,68 @@
1
+ # Projection Attention
2
+ > **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as ✅ verified / ⚠️ retrofittable (needs validation) / ❌ infeasible. Unmarked claims are theoretically possible but require engineering validation.
3
+
4
+ ## Applicable Problems
5
+ When the token/key space dimensionality is too high, the standard attention $Q K^T$ inner product tends to become uniform in high-dimensional spaces ("attention collapse"). In such cases, key/query vectors must first be projected onto a **subspace with superior geometric structure** before computing attention. Typical scenarios include: KV-Cache compression for long-context LLMs, attention sparsification in high-dimensional embedding spaces, and attention alignment of heterogeneous features in multimodal fusion.
6
+
7
+ ## Mathematical Inspiration
8
+ - Lenses: [projection, spectral, probabilistic]
9
+ - Knowledge: [`../../knowledge-base/probability/concentration-inequality.md` (high-dimensional concentration inequalities provide theoretical explanation for attention collapse), `../../knowledge-base/probability/entropy.md` (entropy of attention distributions as a quality metric)]
10
+
11
+ ## Required Mathematical Knowledge
12
+ - **Johnson-Lindenstrauss Lemma**: High-dimensional point sets can be projected onto an $O(\log n / \epsilon^2)$-dimensional subspace while approximately preserving distances
13
+ - **Random Projections and Subspace Embeddings**: Sparse projection matrices (e.g., CountSketch, SRHT) approximately preserve inner products
14
+ - **SVD / PCA Truncation**: Optimal low-rank subspace projection that maximizes retained variance
15
+
16
+ ## AI Module Form
17
+
18
+ **Core Idea**: Project $Q, K \in \mathbb{R}^{n \times d}$ onto an $r$-dimensional subspace ($r \ll d$) before computing attention:
19
+
20
+ $$\text{ProjAttn}(Q, K, V) = \text{softmax}\left(\frac{Q P_Q (K P_K)^T}{\sqrt{r}}\right) V$$
21
+
22
+ where $P_Q, P_K \in \mathbb{R}^{d \times r}$ are projection matrices.
23
+
24
+ **Three Projection Strategies**:
25
+
26
+ 1. **Learnable Projection** ($P$ as trainable parameters):
27
+ ```
28
+ P_Q = Linear(d, r, bias=False) # r << d
29
+ P_K = Linear(d, r, bias=False)
30
+ scores = (Q @ P_Q) @ (K @ P_K).T / sqrt(r)
31
+ attn = softmax(scores) @ V
32
+ ```
33
+
34
+ 2. **Random Fixed Projection** (guaranteed by the JL Lemma):
35
+ ```
36
+ P = random_gaussian(d, r) / sqrt(r) # fixed, not trained
37
+ scores = (Q @ P) @ (K @ P).T / sqrt(r)
38
+ ```
39
+
40
+ 3. **Data-Adaptive Projection** (online PCA):
41
+ ```
42
+ # Maintain running covariance of K, take top-r eigenvectors
43
+ C = running_mean(K^T @ K) # d x d
44
+ eigenvecs = top_r_eigenvectors(C) # d x r
45
+ scores = (Q @ eigenvecs) @ (K @ eigenvecs).T / sqrt(r)
46
+ ```
47
+
48
+ ## Implementable Architectures
49
+ - **Multi-Head Projection**: Each head uses a different $P_h \in \mathbb{R}^{d_h \times r}$, with total computation $O(n \cdot d \cdot r + n^2 \cdot r)$. When $r \ll d$, this saves the $O(n^2 d)$ cost of $Q K^T$
50
+ - **Key-cache Compression**: The projected $K' = K P_K \in \mathbb{R}^{n \times r}$ replaces the original $K$ in storage, reducing Key-cache memory by about $d/r$. Full KV-Cache compression requires a separate V compression/reconstruction mechanism
51
+ - **Hierarchical Projection**: Shallow layers use small $r$ (coarse filtering), while deep layers use large $r$ (fine ranking)
52
+
53
+ ## GPU Feasibility
54
+ - **Dimension 1 Tensorization**: Projection = matrix multiplication, attention = matrix multiplication chain; entirely tensor operations
55
+ - **Dimension 2 GEMM-mappability**: $Q P_Q$ and $K P_K$ are both standard GEMM operations, fully utilizing Tensor Cores
56
+ - **Dimension 3 Complexity**: Projection cost $O(ndr)$ is far below attention cost $O(n^2 d)$, and after projection $r \ll d$ reduces attention to $O(n^2 r)$
57
+ - **Dimension 4 Memory**: Key-cache compressed by $d/r$; V-cache must be handled separately
58
+ - **Dimension 5 Low Precision**: Projection matrices are orthogonal or near-orthogonal, yielding numerical stability; bf16 is acceptable
59
+ - **Dimension 6 Parallelism**: Projection can be pipelined with attention; Multi-Head is naturally parallel across heads
60
+ - **Dimension 7 Sparsity**: Projection matrices are inherently dense; sparse projections (CountSketch) may introduce gather/scatter operations
61
+ - **Dimension 8 Operator Fusion ⚠️ Retrofittable, needs kernel-level validation**: Projection may be integrated into a FlashAttention-style kernel, but requires kernel-level verification
62
+
63
+ ## Paper Phrasing
64
+ "We decompose attention computation into two steps -- low-dimensional subspace projection and projected-space attention -- achieving a $d/r$-fold Key-cache compression (full KV-Cache compression requires separate V compression/reconstruction) while preserving attention quality and guaranteeing Johnson-Lindenstrauss distance preservation."
65
+
66
+ ## Risks
67
+ - **Projection Direction Degeneracy**: Learnable projections may collapse onto a few directions (deterioration of the condition number of $P^T P$), causing attention distributions to degenerate. Orthogonal regularization $\|P^T P - I\|_F^2$ is required.
68
+ - **Irreversible Information Loss**: Projection discards $(d-r)$ dimensions of information; if the task relies on features in these dimensions, performance will degrade. It is recommended to combine with residual connections (weighted combination of original attention and projected attention).
@@ -0,0 +1,68 @@
1
+ # 投影注意力 / Projection Attention
2
+ > **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「✅ 已验证 / ⚠️ 可改造需验证 / ❌ 不可行」。未标注的视为理论可行,需工程验证。
3
+
4
+ ## 适用问题
5
+ 当输入的 token/key 空间维度过高,标准注意力的 $Q K^T$ 内积在高维空间中趋于均匀("注意力坍缩")时,需要先将 key/query 投影到一个**几何结构更优的子空间**再做注意力计算。典型场景:长上下文 LLM 的 KV-Cache 压缩、高维嵌入空间的注意力稀疏化、多模态融合中异构特征的注意力对齐。
6
+
7
+ ## 数学思想来源
8
+ - 透镜:[projection(投影与分解), spectral(谱分解), probabilistic(概率)]
9
+ - 知识:[`../../knowledge-base/probability/concentration-inequality.md`(高维浓度不等式 → 注意力坍缩的理论解释), `../../knowledge-base/probability/entropy.md`(注意力分布的熵作为质量指标)]
10
+
11
+ ## 需要的数学知识
12
+ - **Johnson-Lindenstrauss 引理**:高维点集可投影到 $O(\log n / \epsilon^2)$ 维子空间而近似保持距离
13
+ - **随机投影与子空间嵌入**:稀疏投影矩阵(如 CountSketch、SRHT)近似保持内积
14
+ - **SVD / PCA 截断**:最优低秩子空间投影,最大化保留方差
15
+
16
+ ## AI 模块形式
17
+
18
+ **核心思路**:将 $Q, K \in \mathbb{R}^{n \times d}$ 先投影到 $r$ 维子空间($r \ll d$),再计算注意力:
19
+
20
+ $$\text{ProjAttn}(Q, K, V) = \text{softmax}\left(\frac{Q P_Q (K P_K)^T}{\sqrt{r}}\right) V$$
21
+
22
+ 其中 $P_Q, P_K \in \mathbb{R}^{d \times r}$ 为投影矩阵。
23
+
24
+ **三种投影策略**:
25
+
26
+ 1. **可学习投影**($P$ 为训练参数):
27
+ ```
28
+ P_Q = Linear(d, r, bias=False) # r << d
29
+ P_K = Linear(d, r, bias=False)
30
+ scores = (Q @ P_Q) @ (K @ P_K).T / sqrt(r)
31
+ attn = softmax(scores) @ V
32
+ ```
33
+
34
+ 2. **随机固定投影**(JL 引理保证):
35
+ ```
36
+ P = random_gaussian(d, r) / sqrt(r) # 固定不训练
37
+ scores = (Q @ P) @ (K @ P).T / sqrt(r)
38
+ ```
39
+
40
+ 3. **数据自适应投影**(在线 PCA):
41
+ ```
42
+ # 维护 K 的 running covariance,取 top-r 特征向量
43
+ C = running_mean(K^T @ K) # d x d
44
+ eigenvecs = top_r_eigenvectors(C) # d x r
45
+ scores = (Q @ eigenvecs) @ (K @ eigenvecs).T / sqrt(r)
46
+ ```
47
+
48
+ ## 可实现结构
49
+ - **Multi-Head 投影**:每个 head 使用不同的 $P_h \in \mathbb{R}^{d_h \times r}$,总计算量 $O(n \cdot d \cdot r + n^2 \cdot r)$,当 $r \ll d$ 时节省 $Q K^T$ 的 $O(n^2 d)$ 开销
50
+ - **KV-Cache 压缩**:投影后的 $K' = K P_K \in \mathbb{R}^{n \times r}$ 替代原始 $K$ 存储,Key 部分显存节省约 $d/r$ 倍(完整 KV-Cache 压缩需额外设计 V 的压缩/重构机制)
51
+ - **分层投影**:浅层用小 $r$(粗筛),深层用大 $r$(精排)
52
+
53
+ ## GPU 可行性
54
+ - **维度 1 张量化 ✅**:投影 = 矩阵乘,注意力 = 矩阵乘链,全程张量运算
55
+ - **维度 2 GEMM 可映射 ✅**:$Q P_Q$ 和 $K P_K$ 均为标准 GEMM,吃满 Tensor Core
56
+ - **维度 3 复杂度 ✅**:投影 $O(ndr)$ 远低于注意力 $O(n^2 d)$,且投影后 $r \ll d$ 使注意力 $O(n^2 r)$
57
+ - **维度 4 显存 ✅**:Key-cache 压缩 $d/r$ 倍(V-cache 需独立处理),直接降低显存峰值
58
+ - **维度 5 低精度 ✅**:投影矩阵正交/近正交时数值稳定,bf16 可接受
59
+ - **维度 6 并行 ✅**:投影可与注意力流水线并行,Multi-Head 天然跨头并行
60
+ - **维度 7 稀疏 ⚠️**:投影矩阵本身稠密;若用稀疏投影(CountSketch),可能引入 gather/scatter
61
+ - **维度 8 算子融合 ⚠️可改造需验证**:投影可融入 FlashAttention 的 online softmax 循环中
62
+
63
+ ## 论文表述方式
64
+ "我们将注意力计算分解为低维子空间投影与投影空间注意力两步,在保持注意力质量的同时将 Key-cache 压缩 $d/r$ 倍(V-cache 需独立处理),并保证 Johnson-Lindenstrauss 距离保持性。"
65
+
66
+ ## 风险
67
+ - **投影方向退化**:可学习投影可能坍缩到少数方向($P^T P$ 条件数恶化),导致注意力分布退化。需加正交正则化 $\|P^T P - I\|_F^2$。
68
+ - **信息丢失不可逆**:投影丢弃了 $(d-r)$ 维信息,若任务恰需这些维度的特征则性能下降。建议配合残差连接(原始注意力 + 投影注意力的加权组合)。