算法参考：Layer / Loss / Optimizer

本文档详细解析 neuralnet.cpp 中每个计算组件的数学原理、实现细节和引擎原语分解。

约定

矩阵布局：行主序 (rows, cols)，data_[row * cols + col]
批处理布局：列主序 batch-major (feature_dim, batch_size)
标量类型：Scalar = float
引擎原语 + DSL：算法写在 Layer 里——多数 Layer 用 dsl::compute* 单表达式融合，少数结构算子（matmul / scan / 归约 / 数据搬运）直调 ComputeEngine 原语

Layer 篇

1. Linear — 全连接层

参数： in_features, out_features

可学习参数：

w_: (out_features, in_features) — Xavier 均匀初始化
b_: (out_features, 1) — 零初始化

Forward：

out = W × x + b

实现（compute_layer_mlp.hpp，一条 DSL 表达式 = 单 kernel）：

out = dsl::compute(engine,
    dsl::matmul(W, x) + dsl::row_broadcast(b),
    out_features, batch)                 // matmul 段 + 按行广播 bias 融合为单 kernel

Backward：

grad_x = W^T × grad_out
grad_W += grad_out × x^T
grad_b += Σ_batch grad_out

实现（compute_layer_mlp.hpp）：

grad_input = dsl::compute(engine,
    dsl::matmul(W, grad_out, /*transA=*/true), in_features, batch)   // 纯 matmul 段
grad_W    ← dsl::compute_into(engine,                                // 融合累加，单次 dispatch
    dsl::leaf(grad_W) + dsl::matmul(grad_out, x, false, true), grad_W)
grad_b   ← dsl::compute_reduce(engine,                               // 归约向量输出无法并入
    dsl::row_reduce_sum(dsl::leaf(grad_out)), grad_out.rows(), grad_out.cols())  // 同一表达式，故两步：
            → dsl::compute_into(engine, dsl::leaf(grad_b) + dsl::leaf(*gb), grad_b)

2. ReLU — ReLU 激活函数

参数： 无

Forward：

out = max(x, 0)

实现：dsl::compute(engine, dsl::max(dsl::leaf(x), Scalar{0}), rows, cols)

Backward：

grad_x = grad_out   if x > 0
         0           otherwise

实现：dsl::compute(engine, dsl::select(dsl::leaf(x) > Scalar{0}, dsl::leaf(grad_out), Scalar{0}), rows, cols)

3. GeLU — QuickGeLU 激活函数

参数： β = 1.702（固定常量）

Forward：

out = x · σ(βx) = x · 1/(1 + e^(−βx))

实现：单表达式 DSL 融合（forward/backward 各一条 dsl::compute，全折叠为单个 GPU kernel）：

// forward
out = x / (1 + exp(-β·x))
// backward，s = σ(βx) 由 input_cache_ 重算（不缓存，省显存）
grad_x = grad_out · s · (1 + β·x · (1 − s))

Backward：

factor = s · (1 + βx · (1 − s)),   其中 s = σ(βx)
grad_x = grad_out · factor

4. LayerNorm — 层归一化

参数： normalized_shape, epsilon = 1e-5

可学习参数：

gamma_: (normalized_shape, 1) — 初始化为 1
beta_: (normalized_shape, 1) — 初始化为 0

Forward：

μ = (1/F) Σ_f x_{f,b}
diff = x − μ
σ² = (1/F) Σ_f diff_f²
normalized = diff / √(σ² + ε)
out = γ · normalized + β

Backward：

gy = grad_out · γ
mean_g = (1/F) Σ_f gy_f
gy_norm = gy · normalized
mean_gn = (1/F) Σ_f gy_norm,f
grad_x = (gy − mean_g − normalized · mean_gn) · 1/√(σ² + ε)
grad_γ += Σ_b gy_norm,b
grad_β += Σ_b grad_out,b

5. Softmax — 行级 Softmax

参数： 无

Forward（数值稳定）：

row_max_r = max_c x_{r,c}
shifted_{r,c} = x_{r,c} − row_max_r
exp_shift_{r,c} = e^shifted
row_sum_r = Σ_c exp_shift_{r,c}
out_{r,c} = exp_shift_{r,c} / row_sum_r

实现：单表达式 DSL 融合（compute_layer_softmax.hpp，中间 Tensor 全部由融合 kernel 消解）：

out = exp(x − row_max(x)) / row_sum(exp(x − row_max(x)))

Backward：

ep = out ⊙ grad_out
dot_r = Σ_c ep_{r,c}
gmd = grad_out − dot
grad_x = out ⊙ gmd

实现：同样单表达式 out * (grad_out − row_sum(out * grad_out))。

6. MultiHeadAttention — 多头注意力

参数： d_model, num_heads, seq_len

可学习参数： 4 个 Linear 层（w_q, w_k, w_v, w_o），每个 (d_model, d_model)

Forward（单 fold 路径，S 不物化）：

线性投影：Q = W_q·x, K = W_k·x, V = W_v·x → (H·d_k, batch·seq)
重排：Q/K → rearrange_3d → (batch·H·d_k, seq)；V → transpose+rearrange → V_t (batch·H·seq, d_k)
缩放折进 Q：Q *= scale（scale = 1/√d_k，随表达式融合、免独立 pass）
单 fold kernel：O_t = eval_expr(make_fold_attn_o(...), {Q, K, V_t, [掩码输入]}) —— QKᵀ、掩码（因果/ALiBi/doc 变体在 body 内 select 链表达）、online softmax、 Σw·V 在同一 kernel 内按 EXPR_FOLD_BLOCK=128 逐块完成，(batch·H·seq, seq) 的 S/W 矩阵绝不物化；tri_skip 把被屏蔽的整块钳成空转
反重排：O = transpose + rearrange_3d 回 (H·d_k, batch·seq)
输出投影：out = W_o·O

Backward： recompute_W_ 两步重算 W = softmax(masked(Q·Kᵀ))（掩码树与 forward 的 fold 变体同构同序），随后 R/X 表达式 + batched_matmul 得 grad_Q/K/V——不缓存 W/m/l。

性能关键： fold 单 kernel 一次 dispatch 处理全部样本与头，无 (batch·H·seq, seq) 中间张量往返。（backward 使用 rearrange_3d + batched_matmul 把 H 个 per-head matmul 合并为 1 次 batch dispatch。）

7. PositionalEncoding — 正弦波位置编码

参数： d_model, max_len

算法（固定，不可学习）：

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Forward： out = input + PE

Backward： 梯度直接穿透（编码不可学习）

Tiling 模式： 当 tile_size > 0 时，生成 (d_model, tile_size) 基础编码，沿列方向 tile 为 (d_model, batch·tile_size)，每个样本独立使用同一份编码。

8. FeedForward — 前馈网络

参数： d_model, d_ff

结构： FFN(x) = Linear₂(GeLU(Linear₁(x)))

fc1_: (d_model, d_ff) — 扩展
fc2_: (d_ff, d_model) — 压缩

Forward/Backward 直接委托给子层的 forward/backward。

9. TransformerEncoderLayer — Pre-Norm 编码器层

参数： d_model, num_heads, d_ff, seq_len, norm_type

算法（Pre-Norm）：

x_1 = LN_1(x)
a   = SelfAttn(x_1)
r_2 = x + a
x_2 = LN_2(r_2)
f   = FFN(x_2)
out = r_2 + f

包含子层：MultiHeadAttention + 归一化×2（norm_type：LayerNorm 默认 / RMSNorm / BatchNorm，即 make_norm_layer(d_model, norm_type)）+ FeedForward

10. TransformerEncoder — ViT 编码器

参数： d_model, num_heads, d_ff, num_layers, num_patches, norm_type（默认 LayerNorm）, final_norm（默认 true = NormPlace::Final）

算法：

添加 tiled 位置编码
通过 N 个 TransformerEncoderLayer
（默认启用；final_norm 可显式关闭）编码器末端归一化：在池化之前对每个 token 施加（原版 ViT 的 ln_f）
全局平均池化（按样本聚合 num_patches 维度）

池化：

out[:, b] = (1/P) Σ_{p=0}^{P-1} x[:, b·P + p]

反池化（backward）：

grad_x[:, b·P + p] = (1/P) · grad_out[:, b]

11. PatchEmbedding — 图像 patch 嵌入

参数： img_size, patch_size, d_model

算法：

将 (img_size², batch) 图像提取 num_patches = (img_size/patch_size)² 个不重叠 patch
每个展平 patch (patch_size²) 经 Linear 投影到 d_model 维
重排为 (d_model, batch·num_patches) — batch-major 列布局

12. CausalSelfAttention — 因果自注意力

参数： d_model, num_heads, max_len, seq_len

forward 同为单 fold 路径；因果掩码在 fold body 内以 select 链表达（语义等价）：

mask[i][j] = 0     if j ≤ i
           = -∞    if j > i

掩码绝不物化：j > i 的屏蔽项在 kernel 内直接选择 -inf（tri_skip 进一步把被屏蔽的整块钳成空转）——没有 S += mask 矩阵，也没有掩码缓存。

13. GPTBlock — Pre-Norm 解码器块

参数： d_model, num_heads, d_ff, max_len, seq_len

算法（Pre-Norm）：

x = x + CausalSelfAttn(LN_1(x))
x = x + FFN(LN_2(x))

包含子层：CausalSelfAttention + LayerNorm×2 + FeedForward

14. GPTModel — Decoder-only Transformer

参数： vocab_size, d_model, seq_len, num_heads, d_ff, num_layers

可学习参数：

token_emb_: (vocab_size, d_model) — Token 嵌入
pos_emb_: (seq_len, d_model) — 位置嵌入
N 个 GPTBlock
ln_f_: LayerNorm — 最终归一化
lm_head_: Linear(d_model, vocab_size) — 语言模型头

Forward：

all_emb = gather_rows(token_emb_, input) — 按 token ID 查表
pos_gathered = gather_rows(pos_emb_, pos_indices) — 位置嵌入
x = transpose(all_emb) + transpose(pos_gathered) — (d_model, seq·batch)
通过 N 个 GPTBlock
最终 LayerNorm
LM Head → (vocab_size, seq·batch)

Backward 关键：

Token 嵌入梯度用 scatter_add_rows 按 token ID 累加
位置嵌入梯度用 CPU 辅助聚合（数据量小）

Loss 篇

1. MSELoss — 均方误差

算法：

diff = pred − target
loss = (1/N) Σ diff²
grad = (2/N) · diff

2. CrossEntropyLoss — 交叉熵损失（含 Softmax）

算法（数值稳定）：

col_max   = col_reduce_max(logits)
shifted   = logits − col_max
exp_shift = e^shifted
col_sum   = col_reduce_sum(exp_shift)
softmax   = exp_shift / col_sum
grad      = softmax − target
log_sm    = shifted − log(col_sum)
loss      = −(1/batch) Σ target · log_sm

Backward： grad = softmax − target_onehot

Optimizer 篇

实现方式：优化器全部用 dsl::compute / dsl::compute_into（compute_optimizer.hpp），一次表达式 = 一条融合 kernel；逐元素计算不走 eager 原语——引擎只提供 52 个 virtual 方法，逐元素/广播/条件选择一律经表达式 DSL 执行。下面每条 "DSL 表达式" 即 step() 里的真实写法（leaf = 输入张量，rparam = 标量常量）。

1. SGD — 随机梯度下降

算法：

p ← p − η·g

DSL： dsl::compute_into(engine, dsl::leaf(p) + dsl::leaf(g) * dsl::rparam(-lr), p)

2. SGDWithMomentum — 带动量的 SGD

算法：

v ← β·v + (1−β)·g
p ← p − η·v

DSL（两条 compute_into）：

v ← dsl::compute_into(engine, dsl::leaf(v)*dsl::rparam(β) + dsl::leaf(g)*dsl::rparam(1-β), v)
p ← dsl::compute_into(engine, dsl::leaf(p) + dsl::leaf(v)*dsl::rparam(-lr), p)

参数： lr（学习率）, beta = 0.9（动量系数）

3. Adam — 自适应矩估计

算法：

m ← β₁·m + (1−β₁)·g
v ← β₂·v + (1−β₂)·g²
m̂ = m / (1 − β₁ᵗ),   v̂ = v / (1 − β₂ᵗ)
p ← p − η·(m̂ / (√v̂ + ε))

DSL（compute 出新张量，最后 compute_into 落回参数）：

m_new = dsl::compute(engine, dsl::leaf(m)*dsl::rparam(β1) + dsl::leaf(g)*dsl::rparam(1-β1), rows, cols)
v_new = dsl::compute(engine, dsl::leaf(v)*dsl::rparam(β2) + dsl::leaf(g)*dsl::leaf(g)*dsl::rparam(1-β2), rows, cols)
delta = dsl::compute(engine, -dsl::rparam(lr) * ((dsl::leaf(m)*dsl::rparam(inv_bc1))
                              / (dsl::sqrt(dsl::leaf(v)*dsl::rparam(inv_bc2)) + dsl::rparam(eps))),
                     rows, cols)
p    ← dsl::compute_into(engine, dsl::leaf(p) + dsl::leaf(delta), p)

参数： lr, β1=0.9, β2=0.999, ε=1e-8

偏差修正：

bc1 = 1 − β1ᵗ（一阶矩偏差修正）
bc2 = 1 − β2ᵗ（二阶矩偏差修正）

4. AdamW — 解耦权重衰减

算法：

p ← (1 − η·λ)·p       （权重衰减，独立于梯度）
m ← β₁·m + (1−β₁)·g
v ← β₂·v + (1−β₂)·g²
p ← p − η·(m̂ / (√v̂ + ε))

与 Adam 的区别：

Adam 的 L2 正则化：g += λ·p → 权重衰减被自适应学习率稀释
AdamW 的解耦衰减：p *= (1−η·λ) → 权重衰减对所有参数等效

参数： lr, β1=0.9, β2=0.999, ε=1e-8, weight_decay=0.01

5. Muon — Momentum Orthogonalized by Newton-Schulz

算法（Keller Jordan et al., 2024）：

对于每个 2D 参数 p（权重矩阵）：

SGD-Momentum： v ← μ·v + g
方向选择： Nesterov 时 update = g + μ·v（否则 update = v）
Newton-Schulz 正交化： ortho = NS_5(update)（仅 rows>1 且 cols>1 的参数）
参数更新： p ← p − η · 0.2·√max(m,n) · ortho （形状缩放：NS 输出谱范数为 1，不缩放则等效学习率随形状偏移）

对于非 2D 参数（bias 等）：p ← p − η·update（标准 SGD 更新）。

Newton-Schulz 正交化

目标： 计算矩阵 G 的最近半正交矩阵 Ortho(G) = UV^T（G = USV^T 为 SVD）。

算法：

X = G / (‖G‖_F + ε)

重复 5 次：
  A = X · X^T
  B = b·A + c·A²
  X = a·X + B·X

调优系数： a = 3.4445, b = -4.7750, c = 2.0315

这些系数使 quintic 多项式 φ^N(x) → 1 for x ∈ [0,1]，5 步内收敛。母矩阵取短边一侧（m ≤ n 走行正交化 A = X·Xᵀ（m×m）；m > n 走列正交化 G = Xᵀ·X（n×n）），显存恒为 min(m,n)²。

参考：

Keller Jordan et al., "Muon: An optimizer for hidden layers in neural networks"
https://kellerjordan.github.io/posts/muon/

DSL / 原语（每步，compute_optimizer.hpp）：

A   = dsl::compute(engine, dsl::matmul(X, X, /*transB=*/true), m, m)   // X × X^T
A_sq= dsl::compute(engine, dsl::matmul(A, A), m, m)                    // A²
B   ← dsl::compute_into(engine,                                        // B = b·A + c·A²
        dsl::leaf(A) * dsl::rparam(b) + dsl::leaf(A_sq) * dsl::rparam(c), A)
BX  = dsl::compute(engine, dsl::matmul(B, X), m, n)                    // B·X
X   ← dsl::compute_into(engine,                                        // X = BX + a·X
        dsl::leaf(BX) + dsl::leaf(X) * dsl::rparam(a), BX)

参数： lr, momentum=0.95, nesterov=true, ns_steps=5, ns_eps=1e-7

Optimizer 工厂函数

auto optimizer = nn::create_optimizer(
    name, engine, params, grads, lr, weight_decay);

`name`	创建的优化器	额外参数
`"sgd"`	`SGD`	—
`"sgd_momentum"`	`SGDWithMomentum`	—
`"adam"`	`Adam`	—
`"adamw"`	`AdamW`	`weight_decay`
`"muon"`	`Muon`	—
默认	`Adam`	—

各组件复杂度对比

组件	参数量	Forward FLOPs (per sample)	Backward FLOPs
`Linear(in, out)`	in×out + out	2×in×out	4×in×out
`ReLU`	0	N	N
`GeLU`	0	~5N	~5N
`LayerNorm(F)`	2F	~5F	~8F
`Softmax`	0	~3S	~3S
`MHA(d, h, s)`	4d²	~8d²s + 2ds²	~16d²s + 4ds²
`FeedForward(d, f)`	2df + d + f	4df	8df
`GPTBlock(d, h, f, s)`	—	MHA + FFN + 4LN	MHA + FFN + 4LN

其中：N = 元素数, F = 特征维度, S = 序列长度, d = d_model, h = num_heads, f = d_ff

Navigation Menu

Repository navigation

03-algorithm-reference.md

03-algorithm-reference.md

算法参考：Layer / Loss / Optimizer

约定

Layer 篇

1. Linear — 全连接层

2. ReLU — ReLU 激活函数

3. GeLU — QuickGeLU 激活函数

4. LayerNorm — 层归一化

5. Softmax — 行级 Softmax

6. MultiHeadAttention — 多头注意力

7. PositionalEncoding — 正弦波位置编码

8. FeedForward — 前馈网络

9. TransformerEncoderLayer — Pre-Norm 编码器层

10. TransformerEncoder — ViT 编码器

11. PatchEmbedding — 图像 patch 嵌入

12. CausalSelfAttention — 因果自注意力

13. GPTBlock — Pre-Norm 解码器块

14. GPTModel — Decoder-only Transformer

Loss 篇

1. MSELoss — 均方误差

2. CrossEntropyLoss — 交叉熵损失（含 Softmax）

Optimizer 篇

1. SGD — 随机梯度下降

2. SGDWithMomentum — 带动量的 SGD

3. Adam — 自适应矩估计

4. AdamW — 解耦权重衰减

5. Muon — Momentum Orthogonalized by Newton-Schulz

Newton-Schulz 正交化

Optimizer 工厂函数

各组件复杂度对比

FilesExpand file tree

03-algorithm-reference.md

Latest commit

History

03-algorithm-reference.md

File metadata and controls

算法参考：Layer / Loss / Optimizer

约定

Layer 篇

1. Linear — 全连接层

2. ReLU — ReLU 激活函数

3. GeLU — QuickGeLU 激活函数

4. LayerNorm — 层归一化

5. Softmax — 行级 Softmax

6. MultiHeadAttention — 多头注意力

7. PositionalEncoding — 正弦波位置编码

8. FeedForward — 前馈网络

9. TransformerEncoderLayer — Pre-Norm 编码器层

10. TransformerEncoder — ViT 编码器

11. PatchEmbedding — 图像 patch 嵌入

12. CausalSelfAttention — 因果自注意力

13. GPTBlock — Pre-Norm 解码器块

14. GPTModel — Decoder-only Transformer

Loss 篇

1. MSELoss — 均方误差

2. CrossEntropyLoss — 交叉熵损失（含 Softmax）

Optimizer 篇

1. SGD — 随机梯度下降

2. SGDWithMomentum — 带动量的 SGD

3. Adam — 自适应矩估计

4. AdamW — 解耦权重衰减

5. Muon — Momentum Orthogonalized by Newton-Schulz

Newton-Schulz 正交化

Optimizer 工厂函数

各组件复杂度对比