math-skill 2.1.0 → 3.0.1
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.
- package/README.en-US.md +165 -198
- package/README.md +165 -198
- package/agents/math-critic.en.md +28 -26
- package/agents/math-critic.md +33 -33
- package/commands/ask.en.md +10 -0
- package/commands/ask.md +3 -10
- package/design-patterns/attention/equivariant-attention.en.md +72 -0
- package/design-patterns/attention/equivariant-attention.md +72 -0
- package/design-patterns/attention/geometry-aware-attention.en.md +73 -0
- package/design-patterns/attention/geometry-aware-attention.md +72 -0
- package/design-patterns/attention/information-bottleneck-attention.en.md +72 -0
- package/design-patterns/attention/information-bottleneck-attention.md +72 -0
- package/design-patterns/attention/projection-attention.en.md +68 -0
- package/design-patterns/attention/projection-attention.md +68 -0
- package/design-patterns/attention/spectral-attention.en.md +71 -0
- package/design-patterns/attention/spectral-attention.md +71 -0
- package/design-patterns/compression/leverage-score-selection.en.md +71 -0
- package/design-patterns/compression/leverage-score-selection.md +71 -0
- package/design-patterns/compression/low-rank-kv-cache.en.md +60 -0
- package/design-patterns/compression/low-rank-kv-cache.md +60 -0
- package/design-patterns/compression/spectral-token-pruning.en.md +61 -0
- package/design-patterns/compression/spectral-token-pruning.md +61 -0
- package/design-patterns/compression/topology-preserving-compression.en.md +64 -0
- package/design-patterns/compression/topology-preserving-compression.md +64 -0
- package/design-patterns/loss/constraint-penalty.en.md +67 -0
- package/design-patterns/loss/constraint-penalty.md +78 -0
- package/design-patterns/loss/contrastive-loss.en.md +59 -0
- package/design-patterns/loss/contrastive-loss.md +68 -0
- package/design-patterns/loss/information-bottleneck-loss.en.md +69 -0
- package/design-patterns/loss/information-bottleneck-loss.md +74 -0
- package/design-patterns/loss/orthogonality-loss.en.md +56 -0
- package/design-patterns/loss/orthogonality-loss.md +60 -0
- package/design-patterns/loss/variational-loss.en.md +61 -0
- package/design-patterns/loss/variational-loss.md +71 -0
- package/design-patterns/representation/equivariant-split.en.md +87 -0
- package/design-patterns/representation/equivariant-split.md +89 -0
- package/design-patterns/representation/manifold-representation.en.md +85 -0
- package/design-patterns/representation/manifold-representation.md +87 -0
- package/design-patterns/representation/shared-private-decomposition.en.md +84 -0
- package/design-patterns/representation/shared-private-decomposition.md +86 -0
- package/design-patterns/representation/subspace-alignment.en.md +90 -0
- package/design-patterns/representation/subspace-alignment.md +92 -0
- package/design-patterns/routing/graph-routing.en.md +79 -0
- package/design-patterns/routing/graph-routing.md +81 -0
- package/design-patterns/routing/moe-routing.en.md +85 -0
- package/design-patterns/routing/moe-routing.md +87 -0
- package/design-patterns/routing/optimal-transport-routing.en.md +68 -0
- package/design-patterns/routing/optimal-transport-routing.md +72 -0
- package/design-patterns/routing/spectral-clustering-routing.en.md +88 -0
- package/design-patterns/routing/spectral-clustering-routing.md +90 -0
- package/knowledge-base/differential-geometry/connection.en.md +51 -0
- package/knowledge-base/differential-geometry/connection.md +51 -0
- package/knowledge-base/differential-geometry/curvature.en.md +49 -0
- package/knowledge-base/differential-geometry/curvature.md +49 -0
- package/knowledge-base/differential-geometry/geodesic.en.md +51 -0
- package/knowledge-base/differential-geometry/geodesic.md +51 -0
- package/knowledge-base/differential-geometry/manifold.en.md +47 -0
- package/knowledge-base/differential-geometry/manifold.md +47 -0
- package/knowledge-base/differential-geometry/metric-tensor.en.md +50 -0
- package/knowledge-base/differential-geometry/metric-tensor.md +50 -0
- package/knowledge-base/differential-geometry/tangent-space.en.md +48 -0
- package/knowledge-base/differential-geometry/tangent-space.md +48 -0
- package/knowledge-base/information-geometry/fisher-metric.en.md +56 -0
- package/knowledge-base/information-geometry/fisher-metric.md +56 -0
- package/knowledge-base/information-geometry/natural-gradient.en.md +54 -0
- package/knowledge-base/information-geometry/natural-gradient.md +54 -0
- package/knowledge-base/lie-theory/equivariance.en.md +54 -0
- package/knowledge-base/lie-theory/equivariance.md +54 -0
- package/knowledge-base/lie-theory/group-action.en.md +51 -0
- package/knowledge-base/lie-theory/group-action.md +51 -0
- package/knowledge-base/lie-theory/lie-algebra.en.md +51 -0
- package/knowledge-base/lie-theory/lie-algebra.md +51 -0
- package/knowledge-base/lie-theory/lie-group.en.md +51 -0
- package/knowledge-base/lie-theory/lie-group.md +51 -0
- package/knowledge-base/lie-theory/representation.en.md +52 -0
- package/knowledge-base/lie-theory/representation.md +52 -0
- package/knowledge-base/matrix-analysis/low-rank-approximation.en.md +49 -0
- package/knowledge-base/matrix-analysis/low-rank-approximation.md +49 -0
- package/knowledge-base/matrix-analysis/matrix-perturbation.en.md +49 -0
- package/knowledge-base/matrix-analysis/matrix-perturbation.md +49 -0
- package/knowledge-base/matrix-analysis/positive-semidefinite.en.md +50 -0
- package/knowledge-base/matrix-analysis/positive-semidefinite.md +50 -0
- package/knowledge-base/matrix-analysis/projection.en.md +47 -0
- package/knowledge-base/matrix-analysis/projection.md +47 -0
- package/knowledge-base/matrix-analysis/spectral-decomposition.en.md +48 -0
- package/knowledge-base/matrix-analysis/spectral-decomposition.md +48 -0
- package/knowledge-base/optimization/constrained-optimization.en.md +53 -0
- package/knowledge-base/optimization/constrained-optimization.md +53 -0
- package/knowledge-base/optimization/convex-optimization.en.md +50 -0
- package/knowledge-base/optimization/convex-optimization.md +50 -0
- package/knowledge-base/optimization/lagrangian-duality.en.md +50 -0
- package/knowledge-base/optimization/lagrangian-duality.md +50 -0
- package/knowledge-base/optimization/proximal-method.en.md +52 -0
- package/knowledge-base/optimization/proximal-method.md +52 -0
- package/knowledge-base/optimization/riemannian-optimization.en.md +51 -0
- package/knowledge-base/optimization/riemannian-optimization.md +51 -0
- package/knowledge-base/overview.en.md +29 -218
- package/knowledge-base/overview.md +29 -220
- package/knowledge-base/probability/concentration-inequality.en.md +45 -0
- package/knowledge-base/probability/concentration-inequality.md +45 -0
- package/knowledge-base/probability/entropy.en.md +47 -0
- package/knowledge-base/probability/entropy.md +47 -0
- package/knowledge-base/probability/fisher-information.en.md +48 -0
- package/knowledge-base/probability/fisher-information.md +48 -0
- package/knowledge-base/probability/information-bottleneck.en.md +49 -0
- package/knowledge-base/probability/information-bottleneck.md +49 -0
- package/knowledge-base/probability/kl-divergence.en.md +51 -0
- package/knowledge-base/probability/kl-divergence.md +51 -0
- package/knowledge-base/topology/euler-characteristic.en.md +54 -0
- package/knowledge-base/topology/euler-characteristic.md +54 -0
- package/knowledge-base/topology/fundamental-group.en.md +52 -0
- package/knowledge-base/topology/fundamental-group.md +52 -0
- package/knowledge-base/topology/persistent-homology.en.md +54 -0
- package/knowledge-base/topology/persistent-homology.md +54 -0
- package/lenses/algorithmic.en.md +49 -0
- package/lenses/algorithmic.md +49 -0
- package/lenses/axiomatization.en.md +48 -0
- package/lenses/axiomatization.md +48 -0
- package/lenses/categorical.en.md +48 -0
- package/lenses/categorical.md +48 -0
- package/lenses/causal.en.md +49 -0
- package/lenses/causal.md +49 -0
- package/lenses/duality.en.md +47 -0
- package/lenses/duality.md +47 -0
- package/lenses/game.en.md +50 -0
- package/lenses/game.md +50 -0
- package/lenses/geometric.en.md +53 -0
- package/lenses/geometric.md +53 -0
- package/lenses/local-to-global.en.md +52 -0
- package/lenses/local-to-global.md +52 -0
- package/lenses/perturbation.en.md +51 -0
- package/lenses/perturbation.md +51 -0
- package/lenses/probabilistic.en.md +48 -0
- package/lenses/probabilistic.md +48 -0
- package/lenses/projection.en.md +50 -0
- package/lenses/projection.md +50 -0
- package/lenses/spectral.en.md +52 -0
- package/lenses/spectral.md +52 -0
- package/lenses/symmetry.en.md +49 -0
- package/lenses/symmetry.md +49 -0
- package/lenses/topological.en.md +48 -0
- package/lenses/topological.md +48 -0
- package/lenses/variational.en.md +48 -0
- package/lenses/variational.md +48 -0
- package/package.json +9 -3
- package/references/agentic-workflow.en.md +1 -1
- package/references/agentic-workflow.md +1 -1
- package/references/books/abstract-algebra.en.md +124 -0
- package/references/books/abstract-algebra.md +22 -22
- package/references/books/algebraic-geometry-rising-sea.en.md +171 -0
- package/references/books/algebraic-geometry-rising-sea.md +11 -11
- package/references/books/differential-geometry.en.md +140 -0
- package/references/books/differential-geometry.md +19 -19
- package/references/books/matrix-analysis.en.md +146 -0
- package/references/books/matrix-analysis.md +10 -10
- package/references/books/micro-lie-theory.en.md +116 -0
- package/references/books/micro-lie-theory.md +17 -17
- package/references/books/optimization-ml.en.md +163 -0
- package/references/books/optimization-ml.md +7 -8
- package/references/books/smooth-manifolds.en.md +105 -0
- package/references/books/smooth-manifolds.md +14 -14
- package/references/gpu-friendly-math.en.md +23 -10
- package/references/gpu-friendly-math.md +17 -4
- package/references/inspiration.en.md +1 -1
- package/references/inspiration.md +2 -2
- package/skills/math-research-activator/SKILL.en.md +195 -114
- package/skills/math-research-activator/SKILL.md +194 -119
- package/skills/math-research-activator/original-texts.en.md +9 -9
- package/skills/math-research-activator/original-texts.md +11 -11
- package/commands/abstraction.md +0 -11
- package/commands/algorithmic-thinking.md +0 -11
- package/commands/axiomatization.md +0 -11
- package/commands/causal-inference.md +0 -11
- package/commands/discrete-combinatorial.md +0 -11
- package/commands/game-theory.md +0 -11
- package/commands/induction-analogy.md +0 -11
- package/commands/information-theory.md +0 -11
- package/commands/logic-deduction.md +0 -11
- package/commands/modeling.md +0 -11
- package/commands/optimization.md +0 -11
- package/commands/probability-statistics.md +0 -11
- package/commands/symmetry-invariance.md +0 -11
- package/commands/topological-thinking.md +0 -11
- package/commands/transformation.md +0 -11
- package/skills/abstraction/SKILL.en.md +0 -117
- package/skills/abstraction/SKILL.md +0 -121
- package/skills/abstraction/original-texts.en.md +0 -163
- package/skills/abstraction/original-texts.md +0 -175
- package/skills/algorithmic-thinking/SKILL.en.md +0 -132
- package/skills/algorithmic-thinking/SKILL.md +0 -138
- package/skills/algorithmic-thinking/original-texts.en.md +0 -253
- package/skills/algorithmic-thinking/original-texts.md +0 -256
- package/skills/axiomatization/SKILL.en.md +0 -144
- package/skills/axiomatization/SKILL.md +0 -151
- package/skills/axiomatization/original-texts.en.md +0 -154
- package/skills/axiomatization/original-texts.md +0 -160
- package/skills/causal-inference/SKILL.en.md +0 -147
- package/skills/causal-inference/SKILL.md +0 -151
- package/skills/causal-inference/original-texts.en.md +0 -136
- package/skills/causal-inference/original-texts.md +0 -167
- package/skills/discrete-combinatorial/SKILL.en.md +0 -124
- package/skills/discrete-combinatorial/SKILL.md +0 -131
- package/skills/discrete-combinatorial/original-texts.en.md +0 -184
- package/skills/discrete-combinatorial/original-texts.md +0 -185
- package/skills/game-theory/SKILL.en.md +0 -117
- package/skills/game-theory/SKILL.md +0 -123
- package/skills/game-theory/original-texts.en.md +0 -131
- package/skills/game-theory/original-texts.md +0 -131
- package/skills/induction-analogy/SKILL.en.md +0 -145
- package/skills/induction-analogy/SKILL.md +0 -152
- package/skills/induction-analogy/original-texts.en.md +0 -140
- package/skills/induction-analogy/original-texts.md +0 -154
- package/skills/information-theory/SKILL.en.md +0 -134
- package/skills/information-theory/SKILL.md +0 -140
- package/skills/information-theory/original-texts.en.md +0 -127
- package/skills/information-theory/original-texts.md +0 -127
- package/skills/logic-deduction/SKILL.en.md +0 -130
- package/skills/logic-deduction/SKILL.md +0 -135
- package/skills/logic-deduction/original-texts.en.md +0 -160
- package/skills/logic-deduction/original-texts.md +0 -173
- package/skills/modeling/SKILL.en.md +0 -135
- package/skills/modeling/SKILL.md +0 -139
- package/skills/modeling/original-texts.en.md +0 -162
- package/skills/modeling/original-texts.md +0 -165
- package/skills/optimization/SKILL.en.md +0 -129
- package/skills/optimization/SKILL.md +0 -135
- package/skills/optimization/original-texts.en.md +0 -167
- package/skills/optimization/original-texts.md +0 -168
- package/skills/probability-statistics/SKILL.en.md +0 -146
- package/skills/probability-statistics/SKILL.md +0 -151
- package/skills/probability-statistics/original-texts.en.md +0 -191
- package/skills/probability-statistics/original-texts.md +0 -193
- package/skills/symmetry-invariance/SKILL.en.md +0 -135
- package/skills/symmetry-invariance/SKILL.md +0 -139
- package/skills/symmetry-invariance/original-texts.en.md +0 -206
- package/skills/symmetry-invariance/original-texts.md +0 -221
- package/skills/topological-thinking/SKILL.en.md +0 -124
- package/skills/topological-thinking/SKILL.md +0 -128
- package/skills/topological-thinking/original-texts.en.md +0 -134
- package/skills/topological-thinking/original-texts.md +0 -164
- package/skills/transformation/SKILL.en.md +0 -120
- package/skills/transformation/SKILL.md +0 -124
- package/skills/transformation/original-texts.en.md +0 -204
- package/skills/transformation/original-texts.md +0 -216
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
# Graph Routing
|
|
2
|
+
> **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as [v] verified / [~] retrofittable (needs validation) / [x] infeasible. Unmarked claims are theoretically possible but require engineering validation.
|
|
3
|
+
|
|
4
|
+
## Applicable Problems
|
|
5
|
+
Use when there exists a known or learnable topological structure among modules/experts. Typical scenarios:
|
|
6
|
+
(1) Hierarchical MoE -- experts organized in a tree structure, routing proceeds along tree edges;
|
|
7
|
+
(2) Pipeline/serial routing -- inputs pass through multiple processing stages in DAG order;
|
|
8
|
+
(3) Spatial/temporal correlation routing -- tokens at adjacent positions tend to be routed to similar experts (spatial continuity);
|
|
9
|
+
(4) Knowledge-graph-guided expert selection -- experts organized according to a concept graph.
|
|
10
|
+
Core requirement: **leverage structural priors to constrain routing decisions and reduce the search space**.
|
|
11
|
+
|
|
12
|
+
## Mathematical Inspiration
|
|
13
|
+
- Lenses: lenses/geometric.md (graph Laplacian, spectral graph theory), lenses/probabilistic.md (message passing, information flow)
|
|
14
|
+
- Knowledge: knowledge-base/matrix-analysis/projection.md (adjacency matrix, spectral decomposition),
|
|
15
|
+
knowledge-base/optimization/lagrangian-duality.md (variational on graphs, diffusion processes)
|
|
16
|
+
|
|
17
|
+
## Required Mathematical Background
|
|
18
|
+
- **Graph Laplacian**: L = D - A (combinatorial) or L_sym = D^{-1/2} L D^{-1/2} (normalized)
|
|
19
|
+
Eigendecomposition L = U Lambda U^T provides the graph frequency-domain basis; low-frequency components correspond to smooth signals
|
|
20
|
+
- **Graph Diffusion / Random Walk**: P = D^{-1} A is the transition matrix, P^t describes the distribution after t steps
|
|
21
|
+
PageRank: pi = alpha * P^T * pi + (1 - alpha) * v, balancing graph structure and prior preferences
|
|
22
|
+
- **Graph Neural Network Message Passing**:
|
|
23
|
+
h_i^{(l+1)} = sigma(sum_{j in N(i)} W^{(l)} h_j^{(l)} / sqrt(d_i * d_j))
|
|
24
|
+
Equivalent to a single sparse matrix multiplication L_sym * H * W
|
|
25
|
+
- **Min-Cut Spectral Clustering**: min cut(A, B) s.t. vol(A) = vol(B) => approximate solution given by the Fiedler vector of L
|
|
26
|
+
|
|
27
|
+
## AI Module Form
|
|
28
|
+
```
|
|
29
|
+
Module: GraphRouter
|
|
30
|
+
Input: token representations X in R^{N x d}, expert graph G = (V, E) with |V| = K
|
|
31
|
+
|
|
32
|
+
Method 1 - Graph Diffusion Routing (precomputed):
|
|
33
|
+
A in R^{K x K} // expert adjacency matrix (predefined or learnable)
|
|
34
|
+
P = softmax(A / tau, dim=-1) // normalized transition probabilities
|
|
35
|
+
P_t = matrix_power(P, t) // t-step diffusion, t = 2 ~ 5
|
|
36
|
+
// routing score = initial score * diffusion matrix
|
|
37
|
+
score_init = X @ W_gate // N x K, standard gate
|
|
38
|
+
score_final = score_init @ P_t // N x K, graph-diffusion smoothed
|
|
39
|
+
// A single GEMM (N x K) @ (K x K) incorporates graph structure
|
|
40
|
+
|
|
41
|
+
Method 2 - GNN Routing (learnable graph structure):
|
|
42
|
+
H_0 = expert_embeddings // K x d
|
|
43
|
+
H_1 = ReLU(L_norm @ H_0 @ W_1) // 1-layer GCN
|
|
44
|
+
H_2 = L_norm @ H_1 @ W_2 // 2-layer GCN
|
|
45
|
+
score = X @ H_2^T // N x K routing scores
|
|
46
|
+
// Graph structure updated end-to-end via learnable parameterization of A
|
|
47
|
+
|
|
48
|
+
Method 3 - Hierarchical Tree Routing (O(log K) complexity):
|
|
49
|
+
// Experts organized as a binary tree; each internal node is a binary classifier
|
|
50
|
+
for level in range(depth): // depth = log2(K)
|
|
51
|
+
direction = sigmoid(X @ w_level + b_level) // left/right subtree selection
|
|
52
|
+
path_prob *= direction // accumulate path probability
|
|
53
|
+
// Total computation: O(N * d * log K) vs. O(N * d * K) for standard MoE
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## Implementable Structures
|
|
57
|
+
- **Sparse adjacency matrix**: Use torch.sparse to store A; sparse matmul replaces dense operations
|
|
58
|
+
- **Precomputed diffusion kernel**: P_t fixed during early training, periodically recomputed (once per epoch)
|
|
59
|
+
- **Graph structure learning**: A = softmax(MLP(E_i + E_j)) parameterizes edge weights for end-to-end learning
|
|
60
|
+
- **Hierarchical tree implementation**: Represented as a complete binary tree array with level-wise vectorization
|
|
61
|
+
|
|
62
|
+
## GPU Feasibility
|
|
63
|
+
- **Tensorization**: GCN layers perform sparse matrix times dense matrix (SpMM), supported by both PyTorch and cuSPARSE
|
|
64
|
+
- **GEMM-mappable**: score_init @ P_t in Method 1 is standard GEMM (N x K) @ (K x K)
|
|
65
|
+
- **Complexity**: Method 1 O(N * K^2) diffusion + O(N * K * d) gate; Method 3 O(N * d * log K) significantly better than O(N * d * K)
|
|
66
|
+
- **Memory & KV-Cache**: A matrix stored sparsely at K x K; hierarchical tree parameters d x log K are negligible
|
|
67
|
+
- **Low-precision stability**: Probability matrix P and softmax in fp16 require attention to normalization precision
|
|
68
|
+
- **Parallelism & Communication**: GNN message passing can be batched in parallel; hierarchical tree nodes at the same level can be evaluated independently in parallel
|
|
69
|
+
- **Sparse structure**: Graph adjacency matrices are naturally sparse (degree << K); SpMM reduces from O(K^2) to O(K * avg_deg)
|
|
70
|
+
- **Operator fusion**: L_norm @ H @ W in GCN can be fused into a single sparse GEMM
|
|
71
|
+
|
|
72
|
+
## Paper-Worthy Formulation
|
|
73
|
+
"Leveraging hierarchical tree/graph topology among experts, we compress routing decisions from a flat O(N * K) search to O(N * log K) tree traversal or O(N * K * avg_deg) graph diffusion. Fiedler spectral analysis reveals that the graph's algebraic connectivity lambda_2 directly governs the diversity-coherence trade-off in routing."
|
|
74
|
+
|
|
75
|
+
## Risks
|
|
76
|
+
- Incorrect graph structure priors can misguide routing toward suboptimal experts
|
|
77
|
+
- Over-smoothing in graph diffusion (excessively large t) causes all tokens to route to the same expert
|
|
78
|
+
- Learnable graph structure increases parameter count and overfitting risk
|
|
79
|
+
- Binary classification errors in hierarchical trees accumulate layer by layer, degrading performance at greater depths
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# Graph Routing(图路由)
|
|
2
|
+
> **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「[v] 已验证 / [~] 可改造需验证 / [x] 不可行」。未标注的视为理论可行,需工程验证。
|
|
3
|
+
|
|
4
|
+
## 适用问题
|
|
5
|
+
当模块/专家之间存在已知的或可学习的拓扑结构时使用。典型场景:
|
|
6
|
+
(1) 层次化 MoE——专家按树形结构组织,路由沿树的边进行;
|
|
7
|
+
(2) Pipeline/串行路由——输入按 DAG 顺序经过多个处理阶段;
|
|
8
|
+
(3) 空间/时间相关路由——相邻位置的 token 倾向路由到相近专家(空间连续性);
|
|
9
|
+
(4) 知识图谱引导的专家选择——专家按概念图谱组织。
|
|
10
|
+
核心诉求:**利用结构化先验约束路由决策,减少搜索空间**。
|
|
11
|
+
|
|
12
|
+
## 数学思想来源
|
|
13
|
+
- 透镜:lenses/geometric.md(图拉普拉斯、谱图论)、lenses/probabilistic.md(消息传递、信息流)
|
|
14
|
+
- 知识:knowledge-base/matrix-analysis/projection.md(邻接矩阵、谱分解)、
|
|
15
|
+
knowledge-base/optimization/lagrangian-duality.md(图上优化、扩散过程)
|
|
16
|
+
|
|
17
|
+
## 需要的数学知识
|
|
18
|
+
- **图拉普拉斯**:L = D - A(组合)或 L_sym = D^{-1/2} L D^{-1/2}(归一化)
|
|
19
|
+
特征分解 L = U Λ U^T 给出图的频域基,低频分量对应平滑信号
|
|
20
|
+
- **图扩散/Random Walk**:P = D^{-1}A 为转移矩阵,P^t 描述 t 步后的分布
|
|
21
|
+
PageRank: π = α·P^T·π + (1-α)·v,平衡图结构与先验偏好
|
|
22
|
+
- **Graph Neural Network 消息传递**:
|
|
23
|
+
h_i^{(l+1)} = σ(Σ_{j∈N(i)} W^{(l)} h_j^{(l)} / √(d_i·d_j))
|
|
24
|
+
等价于一次稀疏矩阵乘法 L_sym · H · W
|
|
25
|
+
- **Min-Cut 谱聚类**:min cut(A,B) s.t. vol(A)=vol(B) → 近似解为 L 的 Fiedler 向量
|
|
26
|
+
|
|
27
|
+
## AI 模块形式
|
|
28
|
+
```
|
|
29
|
+
模块:GraphRouter
|
|
30
|
+
输入:token 表示 X ∈ R^{N×d},专家图 G = (V, E) 其中 |V|=K
|
|
31
|
+
|
|
32
|
+
方法1 - 图扩散路由(预计算):
|
|
33
|
+
A ∈ R^{K×K} // 专家邻接矩阵(预定义或可学习)
|
|
34
|
+
P = softmax(A / τ, dim=-1) // 归一化转移概率
|
|
35
|
+
P_t = matrix_power(P, t) // t 步扩散,t=2~5
|
|
36
|
+
// 路由分数 = 初始分数 × 扩散矩阵
|
|
37
|
+
score_init = X @ W_gate // N×K,标准 gate
|
|
38
|
+
score_final = score_init @ P_t // N×K,图扩散平滑
|
|
39
|
+
// 一次 GEMM (N×K)@(K×K) 即可融入图结构
|
|
40
|
+
|
|
41
|
+
方法2 - GNN 路由(可学习图结构):
|
|
42
|
+
H_0 = expert_embeddings // K×d
|
|
43
|
+
H_1 = ReLU(L_norm @ H_0 @ W_1) // 1层 GCN
|
|
44
|
+
H_2 = L_norm @ H_1 @ W_2 // 2层 GCN
|
|
45
|
+
score = X @ H_2^T // N×K 路由分数
|
|
46
|
+
// 图结构通过 A 的可学习参数化端到端更新
|
|
47
|
+
|
|
48
|
+
方法3 - 层次树路由(O(log K) 复杂度):
|
|
49
|
+
// 专家组织为二叉树,每个内部节点为二分类器
|
|
50
|
+
for level in range(depth): // depth = log₂(K)
|
|
51
|
+
direction = sigmoid(X @ w_level + b_level) // 左/右子树选择
|
|
52
|
+
path_prob *= direction // 路径概率累积
|
|
53
|
+
// 总计算量:O(N·d·log K) vs 标准 MoE 的 O(N·d·K)
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## 可实现结构
|
|
57
|
+
- **稀疏邻接矩阵**:用 torch.sparse 存储 A,稀疏 matmul 替代 dense
|
|
58
|
+
- **预计算扩散核**:P_t 在训练初期固定,周期性重计算(每 epoch 一次)
|
|
59
|
+
- **图结构学习**:A = softmax(MLP(E_i ⊕ E_j)) 参数化边权重,端到端学习
|
|
60
|
+
- **层次树实现**:用完全二叉树的数组表示,level-wise 向量化
|
|
61
|
+
|
|
62
|
+
## GPU 可行性
|
|
63
|
+
- **张量化**:GCN 层为稀疏矩阵×稠密矩阵 (SpMM),PyTorch 和 cuSPARSE 均支持
|
|
64
|
+
- **GEMM 可映射**:方法1 的 score_init@P_t 为标准 GEMM (N×K)@(K×K)
|
|
65
|
+
- **复杂度**:方法1 O(N·K²) 扩散 + O(N·K·d) gate;方法3 O(N·d·log K) 显著优于 O(N·d·K)
|
|
66
|
+
- **显存与 KV-Cache**:A 矩阵 K×K 稀疏存储;层次树参数 d×log K 极小
|
|
67
|
+
- **低精度稳定**:概率矩阵 P 和 softmax 在 fp16 下需注意归一化精度
|
|
68
|
+
- **并行与通信**:GNN 消息传递可批量并行;层次树同层节点独立可并行判断
|
|
69
|
+
- **稀疏结构**:图邻接矩阵天然稀疏(度 << K),SpMM 加速比 O(K²) 到 O(K·avg_deg)
|
|
70
|
+
- **算子融合**:GCN 的 L_norm@H@W 可融合为单次稀疏 GEMM
|
|
71
|
+
|
|
72
|
+
## 论文表述方式
|
|
73
|
+
"利用专家间的层次树/图拓扑结构,将路由决策从 O(N·K) 的平铺搜索压缩为 O(N·log K) 的
|
|
74
|
+
树遍历或 O(N·K·avg_deg) 的图扩散,Fiedler 谱分析表明图的代数连通度 λ₂ 直接控制
|
|
75
|
+
路由的多样性-一致性权衡。"
|
|
76
|
+
|
|
77
|
+
## 风险
|
|
78
|
+
- 图结构先验不正确时,路由被误导到次优专家
|
|
79
|
+
- 图扩散过度平滑(t 过大)导致所有 token 路由到同一专家(oversmoothing)
|
|
80
|
+
- 可学习图结构增加参数量和过拟合风险
|
|
81
|
+
- 层次树的二分类误差逐层累积,depth 过深时性能下降
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
# MoE Routing (Mixture-of-Experts Routing)
|
|
2
|
+
> **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as [v] verified / [~] retrofittable (needs validation) / [x] infeasible. Unmarked claims are theoretically possible but require engineering validation.
|
|
3
|
+
|
|
4
|
+
## Applicable Problems
|
|
5
|
+
Use in large-scale models where a small number of experts must be dynamically selected to process each token, achieving parameter scaling while keeping inference cost manageable.
|
|
6
|
+
Typical scenarios: (1) Sparse MoE layers -- each token selects top-k experts (k << K);
|
|
7
|
+
(2) Shared + Private expert mixing -- shared experts handle general features, private experts handle specialized features;
|
|
8
|
+
(3) Multi-granularity MoE -- different layers employ different granularities of expert specialization.
|
|
9
|
+
Core requirement: **sparse activation, load balancing, end-to-end trainability**.
|
|
10
|
+
|
|
11
|
+
## Mathematical Inspiration
|
|
12
|
+
- Lenses: lenses/variational.md (discrete variational relaxation, Gumbel-Softmax), lenses/probabilistic.md (information-theoretic routing)
|
|
13
|
+
- Knowledge: knowledge-base/optimization/lagrangian-duality.md (combinatorial variational, integer programming relaxation),
|
|
14
|
+
knowledge-base/probability/entropy.md (entropy regularization, information bottleneck)
|
|
15
|
+
|
|
16
|
+
## Required Mathematical Background
|
|
17
|
+
- **Mixture Model EM**: p(y|x) = sum_k pi_k(x) * p(y|x, theta_k)
|
|
18
|
+
E-step estimates responsibilities gamma_{nk} = pi_k * p(y_n|x_n, theta_k) / sum_j pi_j * p(y_n|x_n, theta_j)
|
|
19
|
+
M-step updates expert parameters theta_k and mixture weights pi_k
|
|
20
|
+
- **Top-k Sparse Gate**: G(x) = Softmax(TopK(x * W_g))
|
|
21
|
+
TopK is non-differentiable; during training, use noisy top-k or straight-through estimator
|
|
22
|
+
- **Load-Balancing Auxiliary Loss**: L_aux = alpha * K * sum_k f_k * P_k
|
|
23
|
+
f_k = fraction of tokens assigned to expert k, P_k = average gating probability for expert k
|
|
24
|
+
- **Expert Choice Routing**: Experts actively select tokens, rather than tokens selecting experts
|
|
25
|
+
score_{ki} = sim(e_k, x_i), each expert selects top-C tokens
|
|
26
|
+
|
|
27
|
+
## AI Module Form
|
|
28
|
+
```
|
|
29
|
+
Module: MoERouter
|
|
30
|
+
Input: X in R^{N x d}, K experts {E_k}_{k=1}^K, k experts activated per token
|
|
31
|
+
|
|
32
|
+
Method 1 - Standard Top-K Gate (Switch/ST-MoE):
|
|
33
|
+
logits = X @ W_gate // N x K, standard GEMM
|
|
34
|
+
noise = randn(N, K) * softplus(X @ W_noise) // learnable noise to encourage exploration
|
|
35
|
+
logits_noisy = logits + noise
|
|
36
|
+
topk_vals, topk_idx = topk(logits_noisy, k, dim=-1) // select top-k
|
|
37
|
+
gate_weights = softmax(topk_vals, dim=-1) // weights for k experts
|
|
38
|
+
// output = sum_{j in top-k} gate_weights_j * E_j(X)
|
|
39
|
+
|
|
40
|
+
Method 2 - Shared + Private Dual-Path Routing:
|
|
41
|
+
// Shared expert always activated; Private experts selected via top-k
|
|
42
|
+
shared_out = E_shared(X) // all tokens pass through shared expert
|
|
43
|
+
private_logits = X @ W_private_gate // N x K_private
|
|
44
|
+
private_topk = topk(private_logits, k_p)
|
|
45
|
+
private_out = sum_j gate_j * E_private_j(X)
|
|
46
|
+
output = shared_out + private_out // or concat + linear
|
|
47
|
+
|
|
48
|
+
Method 3 - Expert Choice (Google 2022):
|
|
49
|
+
// Reverse perspective: each expert selects top-C tokens
|
|
50
|
+
affinity = E_embeddings @ X^T // K x N, expert-token affinity
|
|
51
|
+
for k in range(K):
|
|
52
|
+
chosen_tokens = topk(affinity[k], C) // each expert selects C = N/K tokens
|
|
53
|
+
expert_k.process(chosen_tokens)
|
|
54
|
+
// Natural load balancing: each expert processes exactly C tokens
|
|
55
|
+
|
|
56
|
+
Auxiliary loss:
|
|
57
|
+
f = onehot(topk_idx).float().mean(dim=0) // K-dim, per-expert load
|
|
58
|
+
P = softmax(logits, dim=-1).mean(dim=0) // K-dim, per-expert average probability
|
|
59
|
+
L_aux = K * dot(f, P) // equals 1 when uniform, > 1 when non-uniform
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
## Implementable Structures
|
|
63
|
+
- **Gate network**: Single-layer Linear(d, K) + optional noise network
|
|
64
|
+
- **Expert parallelism**: Each expert on a separate GPU, all-to-all communication to exchange tokens
|
|
65
|
+
- **Capacity factor**: cap = C_factor * N/K; tokens exceeding capacity go through residual (not dropped)
|
|
66
|
+
- **Router Z-loss**: L_z = alpha * mean(logsumexp(logits)^2) stabilizes logit magnitudes
|
|
67
|
+
|
|
68
|
+
## GPU Feasibility
|
|
69
|
+
- **Tensorization**: Gate logits = X @ W_gate is standard GEMM (N x d) @ (d x K); expert computation is batched GEMM
|
|
70
|
+
- **GEMM-mappable**: Gate requires 1 GEMM; each expert internally is a standard FFN (2 GEMMs + activation)
|
|
71
|
+
- **Complexity**: Gate O(N * d * K); per expert O(N * d * d_ff / k); total FLOPs approximately standard FFN x k
|
|
72
|
+
- **Memory & KV-Cache**: All K expert parameters stored but only k activated; activation memory approximately standard FFN x k
|
|
73
|
+
- **Low-precision stability**: Gate softmax + top-k is safe in fp16; Router Z-loss requires fp32 logsumexp
|
|
74
|
+
- **Parallelism & Communication**: Expert parallelism requires all-to-all communication (each GPU sends/receives N/K tokens); bandwidth-sensitive
|
|
75
|
+
- **Sparse structure**: Top-k routing is naturally sparse (activates k/K fraction of parameters); sparsity = 1 - k/K
|
|
76
|
+
- **Operator fusion**: Gate -> top-k -> softmax -> weighted-sum can be fused; intra-expert FFN can be fused
|
|
77
|
+
|
|
78
|
+
## Paper-Worthy Formulation
|
|
79
|
+
"We employ noisy top-k gating for sparse mixture-of-experts routing, activating a k/K fraction of parameters to achieve O(K) parameter capacity at O(k) inference cost. Combined with the load-balancing auxiliary loss L_aux = K * <f, P> and Router Z-loss, expert utilization exceeds 95%, yielding X% performance improvement over dense models under the same FLOPs budget."
|
|
80
|
+
|
|
81
|
+
## Risks
|
|
82
|
+
- Load imbalance: A few experts are over-selected (Matthew effect), while remaining experts receive insufficient training
|
|
83
|
+
- Top-k operation is non-differentiable; straight-through estimation introduces gradient bias
|
|
84
|
+
- All-to-all communication becomes a bottleneck in multi-GPU settings, especially when k > 1 as communication volume doubles
|
|
85
|
+
- Noise injection promotes exploration but increases training variance, requiring careful annealing
|
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
# MoE Routing(混合专家路由)
|
|
2
|
+
> **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「[v] 已验证 / [~] 可改造需验证 / [x] 不可行」。未标注的视为理论可行,需工程验证。
|
|
3
|
+
|
|
4
|
+
## 适用问题
|
|
5
|
+
大规模模型中需要动态选择少量专家处理每个 token,以实现参数扩展而推理代价可控。
|
|
6
|
+
典型场景:(1) Sparse MoE 层——每个 token 选择 top-k 专家(k<<K);
|
|
7
|
+
(2) Shared+Private 专家混合——共享专家处理通用特征,私有专家处理特异特征;
|
|
8
|
+
(3) 多粒度 MoE——不同层使用不同粒度的专家分工。
|
|
9
|
+
核心诉求:**稀疏激活、负载均衡、端到端可训练**。
|
|
10
|
+
|
|
11
|
+
## 数学思想来源
|
|
12
|
+
- 透镜:lenses/variational.md(离散优化松弛、Gumbel-Softmax)、lenses/probabilistic.md(信息论路由)
|
|
13
|
+
- 知识:knowledge-base/optimization/lagrangian-duality.md(对偶分解、拉格朗日松弛)、
|
|
14
|
+
knowledge-base/probability/entropy.md(熵正则化、信息瓶颈)
|
|
15
|
+
|
|
16
|
+
## 需要的数学知识
|
|
17
|
+
- **混合模型 EM**:p(y|x) = Σ_k π_k(x) · p(y|x,θ_k)
|
|
18
|
+
E 步估计责任 γ_{nk} = π_k·p(y_n|x_n,θ_k) / Σ_j π_j·p(y_n|x_n,θ_j)
|
|
19
|
+
M 步更新专家参数 θ_k 和混合权重 π_k
|
|
20
|
+
- **Top-k 稀疏 Gate**:G(x) = Softmax(TopK(x·W_g))
|
|
21
|
+
TopK 操作不可微,训练时用 noisy top-k 或 straight-through estimator
|
|
22
|
+
- **负载均衡辅助损失**:L_aux = α · K · Σ_k f_k · P_k
|
|
23
|
+
f_k = 分配到专家 k 的 token 比例,P_k = 专家 k 的平均门控概率
|
|
24
|
+
- **Expert Choice Routing**:专家主动选择 token,而非 token 选择专家
|
|
25
|
+
score_{ki} = sim(e_k, x_i),每个专家选 top-C 个 token
|
|
26
|
+
|
|
27
|
+
## AI 模块形式
|
|
28
|
+
```
|
|
29
|
+
模块:MoERouter
|
|
30
|
+
输入:X ∈ R^{N×d},K 个专家 {E_k}_{k=1}^K,每 token 激活 k 个专家
|
|
31
|
+
|
|
32
|
+
方法1 - Standard Top-K Gate (Switch/ST-MoE):
|
|
33
|
+
logits = X @ W_gate // N×K,标准 GEMM
|
|
34
|
+
noise = randn(N, K) * softplus(X @ W_noise) // 可学习噪声,促进探索
|
|
35
|
+
logits_noisy = logits + noise
|
|
36
|
+
topk_vals, topk_idx = topk(logits_noisy, k, dim=-1) // 选 top-k
|
|
37
|
+
gate_weights = softmax(topk_vals, dim=-1) // k 个专家的权重
|
|
38
|
+
// 输出 = Σ_{j∈top-k} gate_weights_j · E_j(X)
|
|
39
|
+
|
|
40
|
+
方法2 - Shared + Private 双路路由:
|
|
41
|
+
// Shared 专家始终激活,Private 专家 top-k 选择
|
|
42
|
+
shared_out = E_shared(X) // 所有 token 都过 shared 专家
|
|
43
|
+
private_logits = X @ W_private_gate // N×K_private
|
|
44
|
+
private_topk = topk(private_logits, k_p)
|
|
45
|
+
private_out = Σ_j gate_j · E_private_j(X)
|
|
46
|
+
output = shared_out + private_out // 或 concat + linear
|
|
47
|
+
|
|
48
|
+
方法3 - Expert Choice (Google 2022):
|
|
49
|
+
// 反转视角:每个专家选择 top-C 个 token
|
|
50
|
+
affinity = E_embeddings @ X^T // K×N,专家与 token 的亲和力
|
|
51
|
+
for k in range(K):
|
|
52
|
+
chosen_tokens = topk(affinity[k], C) // 每个专家选 C=N/K 个 token
|
|
53
|
+
expert_k.process(chosen_tokens)
|
|
54
|
+
// 天然负载均衡:每个专家处理恰好 C 个 token
|
|
55
|
+
|
|
56
|
+
辅助损失:
|
|
57
|
+
f = onehot(topk_idx).float().mean(dim=0) // K 维,各专家负载
|
|
58
|
+
P = softmax(logits, dim=-1).mean(dim=0) // K 维,各专家平均概率
|
|
59
|
+
L_aux = K * dot(f, P) // 均匀时为 1,不均匀时 >1
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
## 可实现结构
|
|
63
|
+
- **Gate 网络**:单层 Linear(d, K) + optional noise network
|
|
64
|
+
- **专家并行**:each expert on separate GPU,all-to-all 通信交换 token
|
|
65
|
+
- **容量因子**:cap = C_factor · N/K,超出容量的 token 走 residual(不被丢弃)
|
|
66
|
+
- **Router Z-loss**:L_z = α·mean(logsumexp(logits)²) 稳定 logits 幅度
|
|
67
|
+
|
|
68
|
+
## GPU 可行性
|
|
69
|
+
- **张量化**:gate logits = X@W_gate 为标准 GEMM (N×d)@(d×K);专家计算为 batched GEMM
|
|
70
|
+
- **GEMM 可映射**:gate 1 次 GEMM;每个专家内部为标准 FFN(2 次 GEMM + activation)
|
|
71
|
+
- **复杂度**:gate O(N·d·K);每专家 O(N·d·d_ff/k);总 FLOPs ≈ 标准 FFN × k
|
|
72
|
+
- **显存与 KV-Cache**:K 个专家参数全存储但仅 k 个激活,激活显存 ≈ 标准 FFN × k
|
|
73
|
+
- **低精度稳定**:gate 的 softmax + top-k 在 fp16 安全;Router Z-loss 需 fp32 logsumexp
|
|
74
|
+
- **并行与通信**:专家并行需 all-to-all 通信(每个 GPU 收发 N/K 个 token),带宽敏感
|
|
75
|
+
- **稀疏结构**:top-k 路由天然稀疏(激活 k/K 的参数),稀疏度 = 1 - k/K
|
|
76
|
+
- **算子融合**:gate → top-k → softmax → weighted-sum 可融合;专家内 FFN 可 fuse
|
|
77
|
+
|
|
78
|
+
## 论文表述方式
|
|
79
|
+
"采用 noisy top-k 门控实现稀疏混合专家路由,激活 k/K 比例的参数以 O(k) 推理代价
|
|
80
|
+
获得 O(K) 参数容量,配合负载均衡辅助损失 L_aux = K·⟨f,P⟩ 和 Router Z-loss 保证
|
|
81
|
+
专家利用率 >95%,在同等 FLOPs 预算下相比 dense 模型提升 X% 性能。"
|
|
82
|
+
|
|
83
|
+
## 风险
|
|
84
|
+
- 负载不均衡:少数专家被过度选择(马太效应),其余专家得不到训练
|
|
85
|
+
- Top-k 操作不可微,straight-through 估计引入梯度偏差
|
|
86
|
+
- All-to-all 通信在多 GPU 下成为瓶颈,尤其 k>1 时通信量翻倍
|
|
87
|
+
- 噪声注入虽促进探索但增加训练方差,需 careful annealing
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
# Optimal Transport Routing
|
|
2
|
+
> **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as [v] verified / [~] retrofittable (needs validation) / [x] infeasible. Unmarked claims are theoretically possible but require engineering validation.
|
|
3
|
+
|
|
4
|
+
## Applicable Problems
|
|
5
|
+
Use when a set of input tokens/samples must be assigned to a set of experts/sub-modules while pursuing globally optimal matching cost.
|
|
6
|
+
Typical scenarios: (1) Load-balanced MoE routing -- assigning N tokens to K experts with a cost matrix defined as negative similarity; (2) Cross-layer feature alignment -- transporting layer-l features to the optimal subset of layer-(l+1); (3) Multi-task sample assignment -- assigning samples in a batch to the most appropriate task head.
|
|
7
|
+
Core requirement: **globally optimal assignment, rather than greedy per-point decisions**.
|
|
8
|
+
|
|
9
|
+
## Mathematical Inspiration
|
|
10
|
+
- Lenses: lenses/variational.md (convex variational, duality theory), lenses/geometric.md (Wasserstein distance)
|
|
11
|
+
- Knowledge: knowledge-base/optimization/lagrangian-duality.md (duality theory, constrained optimization),
|
|
12
|
+
knowledge-base/probability/entropy.md (entropy regularization, marginal constraints)
|
|
13
|
+
|
|
14
|
+
## Required Mathematical Background
|
|
15
|
+
- **Discrete Optimal Transport**: min_{P in Pi(mu,nu)} <C, P> = sum_{ij} C_{ij} P_{ij}
|
|
16
|
+
where Pi(mu,nu) = {P >= 0 : P * 1 = mu, P^T * 1 = nu} is the set of couplings with marginal constraints
|
|
17
|
+
- **Entropy-Regularized Sinkhorn**: min <C,P> - eps * H(P) => P* = diag(u) * exp(-C/eps) * diag(v)
|
|
18
|
+
Solved by alternating row/column scaling (Sinkhorn-Knopp), convergence rate O(1/eps^2)
|
|
19
|
+
- **Wasserstein-1 Distance**: W_1(mu,nu) = min_{pi in Pi} E_pi[||x-y||] = sup_{||f||_L <= 1} E_mu[f] - E_nu[f]
|
|
20
|
+
Kantorovich-Rubinstein duality, used for continuous distribution matching
|
|
21
|
+
- **Gromov-Wasserstein**: When source/target spaces have different dimensions, minimizes the structure-preserving transport cost
|
|
22
|
+
|
|
23
|
+
## AI Module Form
|
|
24
|
+
```
|
|
25
|
+
Module: OptimalTransportRouter
|
|
26
|
+
Input: token representations X in R^{N x d}, expert embeddings E in R^{K x d}, capacity constraint cap in R^K
|
|
27
|
+
|
|
28
|
+
Cost matrix: C_{ik} = -sim(X_i, E_k) or ||X_i - E_k||^2 (N x K)
|
|
29
|
+
|
|
30
|
+
Sinkhorn routing (entropy-regularized):
|
|
31
|
+
K_mat = exp(-C / eps) // Gibbs kernel, eps = 0.05 ~ 0.1
|
|
32
|
+
for t = 1..T: // T = 5 ~ 20 iterations
|
|
33
|
+
u = a / (K_mat @ v) // row scaling, a = 1/N
|
|
34
|
+
v = b / (K_mat^T @ u) // column scaling, b = cap / sum(cap)
|
|
35
|
+
P = diag(u) @ K_mat @ diag(v) // optimal transport plan, doubly stochastic matrix
|
|
36
|
+
assignment = argmax(P, dim=1) // hard assignment (at inference)
|
|
37
|
+
// At training: weighted_features = P @ E (soft assignment, differentiable)
|
|
38
|
+
|
|
39
|
+
Capacity constraint (b vector):
|
|
40
|
+
b_k = total_tokens / K // uniform allocation
|
|
41
|
+
b_k = alpha * uniform + (1 - alpha) * learned // learned non-uniform allocation
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
## Implementable Structures
|
|
45
|
+
- **Sinkhorn layer**: Custom autograd Function; forward pass performs Sinkhorn iterations, backward pass uses the implicit function theorem for gradients
|
|
46
|
+
- **Fixed iteration count**: T = 10 fixed iterations => can be unrolled into a computation graph (unrolled variational)
|
|
47
|
+
- **Log-domain stabilization**: Convert Sinkhorn to log domain to avoid exp overflow:
|
|
48
|
+
log_u = log_a - logsumexp(log_K + log_v, dim=1)
|
|
49
|
+
- **Batch OT**: Solve independently per micro-batch, parallelize Sinkhorn iterations
|
|
50
|
+
|
|
51
|
+
## GPU Feasibility
|
|
52
|
+
- **Tensorization**: Sinkhorn core is matrix-vector multiplication K @ v of shape (N x K) * (K x 1), standard GEMV
|
|
53
|
+
- **GEMM-mappable**: Computation of C via X @ E^T is GEMM (N x d) @ (d x K); Sinkhorn iterations are GEMV
|
|
54
|
+
- **Complexity**: O(N * K * T) where T = 10 ~ 20; for N = 2048, K = 64 approximately 2.6M FLOPs, negligible
|
|
55
|
+
- **Memory & KV-Cache**: Storing C (N x K) and P (N x K); for N = 2048, K = 64 approximately 1 MB
|
|
56
|
+
- **Low-precision stability**: Sinkhorn in fp16 may cause exp(-C/eps) overflow; log-domain + fp32 recommended
|
|
57
|
+
- **Parallelism & Communication**: Batch dimension is independent; Sinkhorn iterations have sequential dependencies, but each iteration's matvec is highly parallel
|
|
58
|
+
- **Sparse structure**: As eps -> 0, P approaches sparsity (permutation matrix); top-k approximation can be used for acceleration
|
|
59
|
+
- **Operator fusion**: The exp -> matvec -> division pipeline in a single Sinkhorn step can be fused into a CUDA kernel
|
|
60
|
+
|
|
61
|
+
## Paper-Worthy Formulation
|
|
62
|
+
"We formulate token-to-expert routing as an entropy-regularized optimal transport problem, obtaining an epsilon-approximate doubly stochastic assignment matrix via the Sinkhorn-Knopp algorithm within T = 10 iterations. The transport cost is guaranteed to converge to the global optimum within O(eps * log N), while marginal constraints b precisely control the load upper bound per expert."
|
|
63
|
+
|
|
64
|
+
## Risks
|
|
65
|
+
- Excessively small eps causes numerical instability in Sinkhorn (exp overflow); requires log-domain or increasing eps
|
|
66
|
+
- Fixed unrolling of T iterations limits solution precision; too many iterations increase latency
|
|
67
|
+
- The N x K cost matrix creates significant memory pressure when both N and K are large (N = 32K, K = 256 => 32 MB)
|
|
68
|
+
- Training-inference discrepancy: performance gap between soft assignment and hard assignment
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
# Optimal Transport Routing(最优传输路由)
|
|
2
|
+
> **严谨性声明**:本文件中涉及复杂度、显存、FlashAttention 融合、Tensor Core、KV-Cache 压缩的结论均标注为「[v] 已验证 / [~] 可改造需验证 / [x] 不可行」。未标注的视为理论可行,需工程验证。
|
|
3
|
+
|
|
4
|
+
## 适用问题
|
|
5
|
+
当需要将一组输入 token/样本 分配到一组专家/子模块,且追求全局最优的匹配代价时使用。
|
|
6
|
+
典型场景:(1) 负载均衡的 MoE 路由——将 N 个 token 分配到 K 个专家,代价矩阵为负相似度;
|
|
7
|
+
(2) 跨层特征对齐——将第 l 层特征传输到第 l+1 层的最优子集;
|
|
8
|
+
(3) 多任务样本分配——将 batch 中的样本分配到最合适的任务头。
|
|
9
|
+
核心诉求:**全局最优分配,而非贪心逐点决策**。
|
|
10
|
+
|
|
11
|
+
## 数学思想来源
|
|
12
|
+
- 透镜:lenses/variational.md(凸优化、对偶理论)、lenses/geometric.md(Wasserstein 距离)
|
|
13
|
+
- 知识:knowledge-base/optimization/lagrangian-duality.md(对偶理论、约束优化)、
|
|
14
|
+
knowledge-base/probability/entropy.md(熵正则化、边际约束)
|
|
15
|
+
|
|
16
|
+
## 需要的数学知识
|
|
17
|
+
- **离散最优传输**:min_{P∈Π(μ,ν)} ⟨C, P⟩ = Σ_{ij} C_{ij} P_{ij}
|
|
18
|
+
其中 Π(μ,ν) = {P ≥ 0 : P·1 = μ, P^T·1 = ν} 为边际约束耦合集
|
|
19
|
+
- **熵正则化 Sinkhorn**:min ⟨C,P⟩ - ε·H(P) → P* = diag(u)·exp(-C/ε)·diag(v)
|
|
20
|
+
通过交替缩放行/列(Sinkhorn-Knopp)求解,收敛速度 O(1/ε²)
|
|
21
|
+
- **Wasserstein-1 距离**:W_1(μ,ν) = min_{π∈Π} E_π[‖x-y‖] = sup_{‖f‖_L≤1} E_μ[f] - E_ν[f]
|
|
22
|
+
Kantorovich-Rubinstein 对偶,用于连续分布匹配
|
|
23
|
+
- **Gromov-Wasserstein**:当源/目标空间维度不同时,min 结构保持的传输代价
|
|
24
|
+
|
|
25
|
+
## AI 模块形式
|
|
26
|
+
```
|
|
27
|
+
模块:OptimalTransportRouter
|
|
28
|
+
输入:token 表示 X ∈ R^{N×d},专家嵌入 E ∈ R^{K×d},容量约束 cap ∈ R^K
|
|
29
|
+
|
|
30
|
+
代价矩阵:C_{ik} = -sim(X_i, E_k) 或 ‖X_i - E_k‖² (N×K)
|
|
31
|
+
|
|
32
|
+
Sinkhorn 路由(熵正则化):
|
|
33
|
+
K_mat = exp(-C / ε) // Gibbs kernel, ε=0.05~0.1
|
|
34
|
+
for t = 1..T: // T=5~20 次迭代
|
|
35
|
+
u = a / (K_mat @ v) // 行缩放,a = 1/N
|
|
36
|
+
v = b / (K_mat^T @ u) // 列缩放,b = cap/sum(cap)
|
|
37
|
+
P = diag(u) @ K_mat @ diag(v) // 最优传输计划,双随机矩阵
|
|
38
|
+
assignment = argmax(P, dim=1) // 硬分配(推理时)
|
|
39
|
+
// 训练时:weighted_features = P @ E (软分配,可微)
|
|
40
|
+
|
|
41
|
+
容量约束(b 向量):
|
|
42
|
+
b_k = total_tokens / K // 均匀分配
|
|
43
|
+
b_k = α·uniform + (1-α)·learned // 学习非均匀分配
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
## 可实现结构
|
|
47
|
+
- **Sinkhorn 层**:自定义 autograd Function,前向做 Sinkhorn 迭代,反向用隐函数定理求梯度
|
|
48
|
+
- **迭代次数固定**:T=10 次固定迭代 → 可展开为计算图(unrolled variational)
|
|
49
|
+
- **log-domain 稳定化**:将 Sinkhorn 转换到 log 域避免 exp 溢出:
|
|
50
|
+
log_u = log_a - logsumexp(log_K + log_v, dim=1)
|
|
51
|
+
- **Batch OT**:每个 micro-batch 独立求解,并行化 Sinkhorn 迭代
|
|
52
|
+
|
|
53
|
+
## GPU 可行性
|
|
54
|
+
- **张量化**:Sinkhorn 核心为矩阵向量乘法 K@v (N×K)·(K×1),标准 GEMV
|
|
55
|
+
- **GEMM 可映射**:C 的计算 X@E^T 为 GEMM (N×d)@(d×K);Sinkhorn 迭代为 GEMV
|
|
56
|
+
- **复杂度**:O(N·K·T) 其中 T=10~20,当 N=2048, K=64 时约 2.6M FLOPs,极小
|
|
57
|
+
- **显存与 KV-Cache**:存储 C(N×K) 和 P(N×K),N=2048,K=64 时约 1MB
|
|
58
|
+
- **低精度稳定**:Sinkhorn 在 fp16 下 exp(-C/ε) 可能溢出,建议 log-domain + fp32
|
|
59
|
+
- **并行与通信**:batch 维度独立;Sinkhorn 迭代为顺序依赖,但每次迭代的 matvec 高度并行
|
|
60
|
+
- **稀疏结构**:ε→0 时 P 趋于稀疏(permutation matrix),可用 top-k 近似加速
|
|
61
|
+
- **算子融合**:exp → matvec → division 的 Sinkhorn 单步可融合为 CUDA kernel
|
|
62
|
+
|
|
63
|
+
## 论文表述方式
|
|
64
|
+
"将 token-to-expert 路由建模为熵正则化最优传输问题,通过 Sinkhorn-Knopp 算法在 T=10 次
|
|
65
|
+
迭代内求得 ε-近似最优的双随机分配矩阵,理论保证传输代价在 O(ε·log N) 内收敛到全局最优,
|
|
66
|
+
同时通过边际约束 b 精确控制各专家的负载上限。"
|
|
67
|
+
|
|
68
|
+
## 风险
|
|
69
|
+
- ε 过小导致 Sinkhorn 数值不稳定(exp 溢出),需用 log-domain 或增大 ε
|
|
70
|
+
- T 次迭代的固定展开限制了解精度,过多 T 增加延迟
|
|
71
|
+
- N×K 代价矩阵在 N, K 都很大时内存压力显著(N=32K, K=256 → 32MB)
|
|
72
|
+
- 训练-推理不一致:软分配 vs 硬分配的性能 gap
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
# Spectral Clustering Routing
|
|
2
|
+
> **Rigor disclaimer**: Claims about complexity, memory, FlashAttention fusion, Tensor Core, and KV-Cache compression are marked as [v] verified / [~] retrofittable (needs validation) / [x] infeasible. Unmarked claims are theoretically possible but require engineering validation.
|
|
3
|
+
|
|
4
|
+
## Applicable Problems
|
|
5
|
+
Use when routing needs to be based on the intrinsic similarity structure of tokens/samples. Typical scenarios:
|
|
6
|
+
(1) Unsupervised expert assignment -- when no routing labels are available, spectral clustering automatically discovers natural token clusters;
|
|
7
|
+
(2) Adaptive expert initialization -- use spectral clustering results to initialize expert parameters at the beginning of training;
|
|
8
|
+
(3) Input-aware dynamic clustering -- different batches have different token distributions, requiring adaptive routing;
|
|
9
|
+
(4) Multi-granularity clustering -- different layers use spectral clustering at different granularities (coarse to fine).
|
|
10
|
+
Core requirement: **discover the intrinsic cluster structure of data for routing or expert initialization**.
|
|
11
|
+
|
|
12
|
+
## Mathematical Inspiration
|
|
13
|
+
- Lenses: lenses/geometric.md (spectral graph theory, Laplacian eigenmaps), lenses/variational.md (relaxation and approximation)
|
|
14
|
+
- Knowledge: knowledge-base/matrix-analysis/projection.md (eigendecomposition, Rayleigh quotient),
|
|
15
|
+
knowledge-base/differential-geometry/manifold.md (manifold learning, graph cuts)
|
|
16
|
+
|
|
17
|
+
## Required Mathematical Background
|
|
18
|
+
- **Spectral Clustering (Ng-Jordan-Weiss)**:
|
|
19
|
+
1. Construct similarity graph W_{ij} = exp(-||x_i - x_j||^2 / (2 sigma^2))
|
|
20
|
+
2. Compute normalized Laplacian L_sym = I - D^{-1/2} W D^{-1/2}
|
|
21
|
+
3. Extract the k smallest eigenvectors U_k in R^{N x k}
|
|
22
|
+
4. Apply k-means to the rows of U_k to obtain k clusters
|
|
23
|
+
- **Nystrom Approximation**: When N is too large for the full W matrix, sample m << N points
|
|
24
|
+
W approximately C * W_m^{-1} * C^T, reducing eigendecomposition to m x m
|
|
25
|
+
- **Spectral Relaxation Continuation**: Discrete cluster assignment => continuous eigenvectors => differentiable routing
|
|
26
|
+
Use softmax(U_k * W_proj) instead of hard k-means assignment
|
|
27
|
+
- **Power Iteration Acceleration**: Full eigendecomposition is unnecessary; only the top k eigenvectors are needed
|
|
28
|
+
Use Lanczos/Arnoldi iteration O(N^2 * k * iter) or randomized SVD O(N * k * log k)
|
|
29
|
+
|
|
30
|
+
## AI Module Form
|
|
31
|
+
```
|
|
32
|
+
Module: SpectralClusterRouter
|
|
33
|
+
Input: X in R^{N x d}, number of clusters K
|
|
34
|
+
|
|
35
|
+
Method 1 - Online Spectral Clustering Routing (periodic updates during training):
|
|
36
|
+
// Update cluster centers every M steps; use nearest neighbor at inference
|
|
37
|
+
W = exp(-cdist(X_sample, X_sample) / (2 sigma^2)) // m x m sampled similarity
|
|
38
|
+
L = I - D^{-1/2} W D^{-1/2} // normalized Laplacian
|
|
39
|
+
U_k = eigsh(L, k=K, which='SM') // K smallest eigenvectors
|
|
40
|
+
centers = kmeans(U_k, K) // K cluster centers
|
|
41
|
+
// Routing: project new tokens into spectral space and assign
|
|
42
|
+
proj = X @ W_proj // N -> K dimensional projection (learnable)
|
|
43
|
+
assignment = argmin(cdist(proj, centers)) // nearest center assignment
|
|
44
|
+
|
|
45
|
+
Method 2 - Differentiable Spectral Routing (end-to-end):
|
|
46
|
+
// Use softmax relaxation instead of hard assignment
|
|
47
|
+
sim_matrix = X @ X^T // N x N (or sampled m x m)
|
|
48
|
+
A = exp(sim_matrix / tau) // similarity graph (learnable tau)
|
|
49
|
+
D_inv_sqrt = diag(1 / sqrt(sum(A, dim=1) + eps))
|
|
50
|
+
L_norm = I - D_inv_sqrt @ A @ D_inv_sqrt // normalized Laplacian
|
|
51
|
+
// Approximate top K eigenvectors (power iteration + orthogonalization)
|
|
52
|
+
U_k = power_iteration_approx(L_norm, K, steps=5) // N x K
|
|
53
|
+
// Soft assignment
|
|
54
|
+
cluster_logits = U_k @ W_cluster // N x K (learnable projection)
|
|
55
|
+
route_probs = softmax(cluster_logits / tau_route) // soft routing probabilities
|
|
56
|
+
|
|
57
|
+
Method 3 - Anchor Spectral Clustering (large-scale):
|
|
58
|
+
anchors = kmeans_pp(X, m) // m anchor points, m << N
|
|
59
|
+
Z = exp(-cdist(X, anchors) / (2 sigma^2)) // N x m affinity matrix
|
|
60
|
+
L_anchor = I - D_z^{-1/2} Z^T Z D_z^{-1/2} // m x m Laplacian
|
|
61
|
+
U_k = eigsh(L_anchor, K) // m x K eigenvectors
|
|
62
|
+
route = Z @ U_k @ W_proj // N x K routing scores
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
## Implementable Structures
|
|
66
|
+
- **Periodic offline clustering**: Every N_step steps, collect token representations => offline spectral clustering => update routing table
|
|
67
|
+
- **Nystrom sampling**: Randomly sample m = 1024 representative points, reducing the N x N problem to m x m
|
|
68
|
+
- **Power iteration implementation**: 5-10 power iteration steps + Gram-Schmidt orthogonalization, GPU-friendly
|
|
69
|
+
- **Progressive training**: Early stage uses k-means coarse routing => mid-stage spectral clustering refinement => late-stage fine-tuning of routing network
|
|
70
|
+
|
|
71
|
+
## GPU Feasibility
|
|
72
|
+
- **Tensorization**: Similarity matrix X @ X^T is GEMM; Laplacian construction is element-wise + diagonal matrix operations
|
|
73
|
+
- **GEMM-mappable**: Z^T @ Z in Method 3 is GEMM (m x N) @ (N x m); Z @ U_k is GEMM (N x m) @ (m x K)
|
|
74
|
+
- **Complexity**: Full spectral clustering O(N^2 * K) does not scale; Nystrom O(N * m * K + m^3); power iteration O(N^2 * K * T)
|
|
75
|
+
- **Memory & KV-Cache**: N x N similarity matrix exceeds 64 MB when N > 4096; sampling-based dimensionality reduction is essential
|
|
76
|
+
- **Low-precision stability**: Eigendecomposition recommended in fp32; exp(-dist/sigma^2) in fp16 requires distance clipping
|
|
77
|
+
- **Parallelism & Communication**: Power iteration matvec is highly parallel; k-means assign + update steps can be batch-parallelized
|
|
78
|
+
- **Sparse structure**: k-NN graph replaces fully connected graph; W sparsity > 95%, enabling SpMM acceleration
|
|
79
|
+
- **Operator fusion**: Diagonal scaling in D^{-1/2} @ A @ D^{-1/2} can be fused; cdist + exp + normalize can be fused
|
|
80
|
+
|
|
81
|
+
## Paper-Worthy Formulation
|
|
82
|
+
"We achieve differentiable routing via continuous relaxation of spectral clustering: constructing the normalized Laplacian of the token similarity graph, reducing the O(N^2) eigendecomposition to O(Nm + K^3) through Nystrom approximation, and enabling GPU-friendly online spectral clustering via power iteration. Clustering quality, measured by Normalized Cut, guarantees an O(sqrt(log(N/K))) approximation ratio."
|
|
83
|
+
|
|
84
|
+
## Risks
|
|
85
|
+
- The memory and computation cost of the N x N similarity matrix does not scale for long sequences; sampling or k-NN sparsification is mandatory
|
|
86
|
+
- Eigendecomposition is non-differentiable (gradients undefined when eigenvalues coincide); end-to-end training requires relaxation or stop-gradient
|
|
87
|
+
- The number of clusters K must be specified a priori, and re-clustering is needed when K changes
|
|
88
|
+
- The bandwidth parameter sigma is sensitive to clustering quality: too small causes isolated points, too large causes cluster merging
|