模型架构设计空间：Depth、Width、Vocab Size 与架构选择

更新日期：2026-04-17

本文目标：从零设计一个 SOTA LLM 时，如何在 depth、width、head 配置、FFN ratio、vocab size 之间做最优选择。读完后能根据目标参数量直接输出一套架构配置，并理解每个选择背后的原因。

---

一、架构设计空间概述

1.1 关键维度

一个 Decoder-Only Transformer 由以下超参完全定义：

1.2 参数量公式

对标准 Decoder-Only Transformer（Pre-Norm, SwiGLU, GQA），每层参数：

\(P_{\text{attn}} = D \cdot (H \cdot d_{\text{head}}) + D \cdot (H_{kv} \cdot d_{\text{head}}) \cdot 2 + D \cdot D\)

当 \(H_{kv} = H\)（MHA）时简化为 \(4D^2\)。

\(P_{\text{ffn}} = 3 \cdot D \cdot d_{ff} \quad \text{(SwiGLU: gate, up, down 三个矩阵)}\)

\(P_{\text{norm}} = 2D \quad \text{(每层两个 RMSNorm, 可忽略)}\)

总参数量：

\(P_{\text{total}} = L \cdot (P_{\text{attn}} + P_{\text{ffn}}) + V \cdot D \cdot k\)

其中 \(k = 1\)（tied embeddings）或 \(k = 2\)（untied，input + output 各一份）。

近似公式（MHA + standard FFN）：

\(P_{\text{total}} \approx 12LD^2 + VD\)

这里 \(12D^2 = 4D^2(\text{attn}) + 8D^2(\text{FFN, 标准 4×})\)。对 SwiGLU (\(d_{ff} = \frac{8}{3}D\)) 则变为 \(12D^2\) 依然成立（\(4D^2 + 3 \times \frac{8}{3}D^2 = 4D^2 + 8D^2\)）。

1.3 "Shape" 和 "Size" 一样重要

两个 7B 模型——一个 \(L=32, D=4096\)，另一个 \(L=64, D=2896\)——参数量几乎相同，但行为截然不同。前者在每一步拥有更大的信息带宽（宽），后者拥有更多的计算步数（深）。Scaling law 研究 (Kaplan et al., 2020) 表明，在固定参数量下，不同的 depth/width 比率可以导致 5-10% 的 loss 差异——这等效于将训练数据翻倍。

---

二、深度 vs 宽度

2.1 核心权衡

2.2 理论视角

Kaplan et al. (2020) (arxiv:2001.08361)：

• 发现 loss 对深度和宽度的 power-law exponent 不同：深度的系数略大

• 结论：在相同参数量下，适度增加深度比增加宽度更 cost-effective

• 但他们同时指出，这一优势在极深模型上会饱和

Levine et al. (2020) (arxiv:2006.12467)：

• 用张量分解理论证明：深度的表达力增长是指数级的，但有一个 phase transition

• 最优深度 \(L^* \propto \log(P)\)，而最优宽度 \(D^* \propto \sqrt{P / \log P}\)

• 即：参数量翻倍时，深度只需增加常数，宽度需要 \(\sqrt{2}\times\)

• 实际含义：超过某个深度后，继续加深不如加宽

Tay et al. (2022) (arxiv:2207.10551)：

• 系统对比了不同 depth/width 配置在固定 FLOP 下的表现

• 发现 depth-to-width ratio 有一个甜蜜区间：\(L/D \approx 1:100\) 到 \(1:200\)

2.3 实际模型的 Depth/Width 选择

观察：几乎所有主流模型的 \(L/D\) 比都在 \(1:100\) 到 \(1:160\) 之间。这不是巧合——它反映了一个经验性的 Pareto 最优区间。

2.4 为什么更深通常更好

深度的优势来自组合表达力。直觉上，一个 \(L\) 层网络可以表达深度为 \(L\) 的计算图。某些函数（如多步推理、递归结构识别）需要 \(O(L)\) 步才能计算，无论 \(D\) 多大都无法在常数步内完成。

形式化地，Transformer 的每一层执行：

\(x_{l+1} = x_l + \text{FFN}(\text{Attn}(x_l))\)

这是一个 residual stream 上的迭代精化过程。更多层意味着更多次精化机会。参考 Elhage et al. (2021) 对 residual stream 的"通信总线"解释：每层的 attention 和 FFN 是对总线上信息的读-处理-写操作。

2.5 为什么太深会失败

1. 梯度问题：即使有 Pre-Norm + 残差连接，\(L > 100\) 时梯度方差仍会增长。DeepNorm (Wang et al., 2022) 提出用 \(\alpha\)-残差缩放来稳定极深网络：

\(x_{l+1} = \alpha \cdot x_l + \text{sublayer}(x_l)\)

其中 \(\alpha = (2L)^{1/4}\)。

1. 流水线并行瓶颈：PP 将模型按层切分到不同 GPU。\(L\) 越大，pipeline bubble 比例越高（\(\text{bubble} \approx \frac{p-1}{m+p-1}\)，\(p\) 为 pipeline stages，\(m\) 为 micro-batches）。

2. diminishing returns：实证表明，固定参数量下，\(L\) 超过某个阈值后 loss 反而上升。这是因为过深的网络难以被当前优化器有效训练。

2.6 最优 L/D 比的经验规则

综合 Kaplan、Levine、Hoffmann (Chinchilla) 等研究，以及主流模型的实际选择，可以总结：

\(D^* \approx 96 \cdot L^* \quad \text{(适用于 1B–500B 范围)}\)

或者更精确地：

\(L^* \approx 0.7 \cdot P^{0.2}, \quad D^* \approx 0.9 \cdot P^{0.4}\)

其中 \(P\) 为总参数量（以十亿为单位）。这意味着参数量增长 10×，深度增长约 1.6×，宽度增长约 2.5×——宽度增长更快。

---

三、Attention Head 设计

3.1 Head 维度的演进

3.2 为什么 \(d_{\text{head}}=128\) 是当前最优

三个原因：

1. RoPE 频率分辨率：RoPE 将位置编码注入每对维度。\(d_{\text{head}}=128\) 提供 64 个频率分量，覆盖频率范围 \([\theta_0, \theta_{63}]\)，其中 \(\theta_i = \text{base}^{-2i/d_{\text{head}}}\)。这足以在 128K+ 上下文长度中区分位置。\(d_{\text{head}}=64\) 仅有 32 个分量，长上下文场景下频率分辨率不足。

2. FlashAttention 硬件效率：FlashAttention-⅔ 的 SRAM tile 大小在 \(d_{\text{head}}=128\) 时达到最优利用率。A100/H100 的 shared memory 为 192KB/228KB，恰好能容纳 128 维的 Q/K/V tile。

3. 容量-效率平衡：每个 head 本质上是一个独立的"信息路由器"。\(d_{\text{head}}\) 太小（32），单个 head 的 key-query 匹配空间太窄，难以精确 select；\(d_{\text{head}}\) 太大（256+），head 数减少，并行路由多样性下降。

3.3 GQA：n_kv_heads 的选择

Grouped Query Attention (Ainslie et al., 2023) 的核心思想：多个 query head 共享同一组 KV head。 KV cache 大小计算：

\(\text{KV cache} = 2 \times L \times H_{kv} \times d_{\text{head}} \times S \times \text{bytes}\)

对 LLaMA-3 70B（\(L=80, H_{kv}=8, d_{\text{head}}=128, S=8192\)，BF16）：

\(2 \times 80 \times 8 \times 128 \times 8192 \times 2 = 2.68\text{ GB}\)

如果使用 MHA (\(H_{kv}=64\))：\(2 \times 80 \times 64 \times 128 \times 8192 \times 2 = 21.5\text{ GB}\)——单个 sequence 的 KV cache 就占满一张 A100 的 ¼。

3.4 MLA：消除 Head 设计选择

Multi-head Latent Attention (DeepSeek-V2, 2024) 彻底重构了这个问题：

class MLA:
    def __init__(self, d_model, n_heads, d_head, d_compress):
        self.d_c = d_compress  # 压缩维度, 远小于 n_heads * d_head
        # 下投影: d_model → d_compress (存入 KV cache 的维度)
        self.W_dkv = Linear(d_model, d_c, bias=False)
        # 上投影: d_compress → n_heads * d_head (K 和 V 各一个)
        self.W_uk = Linear(d_c, n_heads * d_head, bias=False)
        self.W_uv = Linear(d_c, n_heads * d_head, bias=False)
        # Q 也做类似压缩 (d_model → d_compress_q → n_heads * d_head)
        self.W_dq = Linear(d_model, d_compress_q, bias=False)
        self.W_uq = Linear(d_compress_q, n_heads * d_head, bias=False)
        self.W_o = Linear(n_heads * d_head, d_model, bias=False)

    def forward(self, x):
        c_kv = self.W_dkv(x)           # [B, S, d_c] — 只缓存这个!
        k = self.W_uk(c_kv)            # [B, S, H*d_head]
        v = self.W_uv(c_kv)            # [B, S, H*d_head]
        c_q = self.W_dq(x)
        q = self.W_uq(c_q)             # [B, S, H*d_head]
        # 标准 attention 计算 ...
        return self.W_o(attn_output)

MLA 的优雅之处：KV cache 只需存储 \(d_c\) 维的压缩表示（DeepSeek-V3 中 \(d_c = 512\)），而非 \(H_{kv} \times d_{\text{head}}\) 维。这使 KV cache 压缩比达到 \(\frac{H_{kv} \times d_{\text{head}}}{d_c}\)。对 DeepSeek-V3：\(\frac{128 \times 128}{512} = 32\times\) 压缩。代价是每个 token 在 decode 时需要额外的上投影计算——但这可以通过吸收进 attention 权重矩阵来消除。

---

四、Vocab Size 深度分析

4.1 Vocab Size 的核心权衡

Vocab size \(V\) 影响四个维度：

4.2 Embedding 税

"Embedding 税"指 embedding 和 LM head 参数占总参数的比例：

\(\text{Embedding tax} = \frac{V \times D \times k}{P_{\text{total}}}\)

4.3 各模型 Vocab Size 对比

4.4 Vocab Size 的 Scaling 规则

Tao et al. (2024) 系统研究了 vocab size 的 scaling law：

\(V^* \propto \sqrt{C}\)

其中 \(C\) 为总训练计算量 (FLOPs)。直觉上：

• 计算量增大时，模型能从更细粒度的 token 中学到更多

• 但 \(V\) 过大会导致 low-frequency token 的 embedding 欠训练

• \(\sqrt{C}\) 平衡了这两个效应

实践建议：

4.5 Vocab Size 与 GPU 对齐

词表大小应当是 64 或 128 的倍数。原因：GPU tensor core 以 16×16 (FP16/BF16) 或 8×8 (TF32) 为单位运算。LM head 的矩阵乘 \([B \times S, D] \times [D, V]\) 中，\(V\) 不对齐会导致最后一个 tile 产生 padding 浪费。

LLaMA-3 的 \(V=128256 = 128 \times 1002 + 0\)？实际上 128256 并不是完美的 128 倍数（\(128256 / 128 = 1002\)，是整数），所以对齐。这类细节虽小但在大规模训练中累积效果显著。

---

五、FFN Width Ratio

5.1 从标准 4× 到 SwiGLU 的 8/3×

原始 Transformer (Vaswani et al., 2017)：

\(\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2\)

参数量：\(2 \times D \times d_{ff}\)，标准 \(d_{ff} = 4D\)，即 \(8D^2\)。

SwiGLU (Shazeer, 2020)：

\(\text{SwiGLU}(x) = (\text{SiLU}(xW_{\text{gate}}) \odot xW_{\text{up}}) \cdot W_{\text{down}}\)

参数量：\(3 \times D \times d_{ff}\)（多一个 gate 矩阵）。为匹配标准 FFN 的参数量 \(8D^2\)：

\(3 \times D \times d_{ff} = 8D^2 \implies d_{ff} = \frac{8}{3}D \approx 2.67D\)

5.2 实际 FFN Ratio 对比

5.3 FFN Ratio 的设计考量

为什么不同模型选择不同的 ratio？

def compute_ffn_params(d_model, d_ff, activation="swiglu"):
    if activation == "swiglu":
        return 3  d_model  d_ff  # gate + up + down
    else:
        return 2  d_model  d_ff  # up + down

def compute_attn_params(d_model, n_heads, n_kv_heads, d_head):
    q_params = d_model  n_heads  d_head
    kv_params = d_model  n_kv_heads  d_head * 2
    o_params = n_heads  d_head  d_model
    return q_params + kv_params + o_params

def attn_to_ffn_ratio(d_model, d_ff, n_heads, n_kv_heads, d_head):
    attn = compute_attn_params(d_model, n_heads, n_kv_heads, d_head)
    ffn = compute_ffn_params(d_model, d_ff)
    return ffn / attn

关键 insight：增大 \(d_{ff}/D\) 意味着将更多参数分配给 FFN（知识存储）而非 attention（信息路由）。(Geva et al., 2021) 表明 FFN 层充当 key-value memory——更宽的 FFN 能存储更多事实知识。

---

六、Tied vs Untied Embeddings

6.1 Weight Tying 机制

Weight tying (Press & Wolf, 2017) 将 input embedding 矩阵 \(W_E \in \mathbb{R}^{V \times D}\) 直接复用为 output LM head \(W_{\text{LM}} = W_E^T\)：

class TiedTransformerLM:
    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff):
        self.embed = Embedding(vocab_size, d_model)
        self.layers = [TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
        self.norm = RMSNorm(d_model)
        # 没有单独的 lm_head!

    def forward(self, token_ids):
        x = self.embed(token_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        logits = x @ self.embed.weight.T  # 直接用 embedding 权重转置
        return logits

6.2 参数节省量

\(\Delta P = V \times D\)

6.3 各模型的选择

6.4 何时 Tie，何时 Untie

决策规则：

\(\text{Tie if } \frac{V \times D}{P_{\text{total}}} > 5\% \text{ AND } P_{\text{total}} < 3B\)

理由：

• 小模型 + 大 vocab：embedding 税太高，tying 是性价比最优的参数节省方式

• 大模型 + 任意 vocab：untied 的表达力提升 > 额外参数成本

• 经验数据：Chowdhery et al. (2022, PaLM) 在 ablation 中发现，8B 以上模型 untied 一致优于 tied

---

七、Complete Architecture Recipe

7.1 推荐配置表

规模	\(L\)	\(D\)	\(H\)	\(d_{\text{head}}\)	\(H_{kv}\)	\(d_{ff}\)	\(V\)	Norm	Act	Tied?	为什么
1B	24	2048	16	128	16 (MHA)	5504	32K–64K	RMSNorm Pre	SwiGLU	Yes	1B 小模型 MHA 即可（KV cache 不大）；tying 节省 ~6% 参数；\(d_{ff} = \frac{8}{3} \times 2048 \approx 5461\)，取 5504 对齐 256
3B	28	3072	24	128	8	8192	64K	RMSNorm Pre	SwiGLU	No	3B 开始 untie；GQA 8 KV heads 为推理优化；\(d_{ff}/D = 2.67\) 标准 ratio
7B	32	4096	32	128	8	14336	128K	RMSNorm Pre	SwiGLU	No	行业标准配置 (LLaMA-3 验证)；\(d_{ff}/D = 3.5\) 比理论最优略宽，换取更强 FFN
13B	40	5120	40	128	8	13824	128K	RMSNorm Pre	SwiGLU	No	\(L/D = 1:128\)；40 heads 不是 2 的幂但 GQA 8 组仍可整除
34B	48	6656	52	128	8	17920	128K	RMSNorm Pre	SwiGLU	No	中间规模；52 heads × 128 = 6656；8 KV heads 意味着约 6.5:1 的 Q/KV ratio
70B	80	8192	64	128	8	28672	128K	RMSNorm Pre	SwiGLU	No	LLaMA-70B 验证；\(d_{ff}/D = 3.5\)；80 层需要谨慎的 init 和 lr 调度
200B	96	12288	96	128	8	32768	128K–256K	RMSNorm Pre	SwiGLU	No	GPT-3 级 depth；宽度增长为主；\(d_{ff}/D = 2.67\) 回归理论值以控制总参数
400B+	96–128	16384+	128	128	16	43008+	128K–256K	DeepNorm 或 RMSNorm Pre	SwiGLU	No	超大模型 depth 不再大幅增长（流水线瓶颈）；KV heads 增到 16 因为总 head 数增加需要保持 GQA 比例合理

7.2 配置生成伪代码

import math

def design_architecture(target_params_B: float, vocab_size: int = 128000):
    """根据目标参数量生成架构配置"""
    P = target_params_B * 1e9

    # 估算 L 和 D (经验公式)
    L = int(round(0.7  (target_params_B  0.2)  32))
    L = max(24, min(128, L))
    L = (L // 2) * 2  # 偶数，方便 PP 切分

    # 从 P ≈ 12LD^2 + 2VD 反解 D
    # 12LD^2 + 2VD - P = 0
    a = 12 * L
    b = 2 * vocab_size
    c = -P
    D = int((-b + math.sqrt(b2 - 4ac)) / (2*a))
    D = (D // 128) * 128  # 对齐到 128

    d_head = 128
    H = D // d_head
    H_kv = max(1, min(H, 8))  # 默认 8 KV heads
    if target_params_B <= 1.5:
        H_kv = H  # 小模型用 MHA

    d_ff = int(round(8/3  D / 256)  256)  # SwiGLU 理论值，对齐 256
    if target_params_B >= 7:
        d_ff = int(round(3.5  D / 256)  256)  # 大模型用更宽 FFN

    tied = target_params_B < 2.0

    actual_params = L * (
        D  H  d_head + D  H_kv  d_head  2 + H  d_head * D  # attn
        + 3  D  d_ff  # SwiGLU FFN
    ) + vocab_size  D  (1 if tied else 2)

    return {
        "n_layers": L, "d_model": D, "n_heads": H,
        "d_head": d_head, "n_kv_heads": H_kv, "d_ff": d_ff,
        "vocab_size": vocab_size, "tied": tied,
        "actual_params_B": actual_params / 1e9
    }

7.3 MoE 场景的调整

MoE 模型需要额外考虑 expert 数量和 top-k 路由：

---

八、追问延伸

8.1 Depth/Width vs 下游任务

不同任务对 depth/width 的需求不同：

8.2 Normalization 的进一步选择

除了 RMSNorm Pre-Norm（当前标准），还有新的方向：

8.3 Activation Function 还有更好的选择吗？

SwiGLU 是 2024-2025 的事实标准，但近期研究的值得关注的方向：

8.4 未来方向

1. Mixture of Depths (Raposo et al., 2024)：不同 token 使用不同数量的层，动态调整"深度"

2. Hyper-connections (Zhu et al., 2024)：泛化残差连接，让每层的输入可以来自所有前序层而非仅上一层

3. 线性 Attention + SSM 混合：用 Mamba 层替代部分 attention 层 (Lieber et al., 2024, Jamba)，改变"深度"的含义

---

参考文献

1. Vaswani, A., et al. (2017). Attention Is All You Need. arxiv:1706.03762

2. Kaplan, J., et al. (2020). Scaling Laws for Neural Language Models. arxiv:2001.08361

3. Levine, Y., et al. (2020). The Depth-to-Width Interplay in Self-Attention. arxiv:2006.12467

4. Shazeer, N. (2020). GLU Variants Improve Transformer. arxiv:2002.05202

5. Hoffmann, J., et al. (2022). Training Compute-Optimal Large Language Models (Chinchilla). arxiv:2203.15556

6. Chowdhery, A., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arxiv:2204.02311

7. Tay, Y., et al. (2022). Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers. arxiv:2207.10551

8. Wang, H., et al. (2022). DeepNet: Scaling Transformers to 1,000 Layers. arxiv:2203.00555

9. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arxiv:2305.13245

10. Press, O. & Wolf, L. (2017). Using the Output Embedding to Improve Language Models. arxiv:1608.05859

11. Geva, M., et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. arxiv:2012.14913

12. Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. arxiv:2110.14737

13. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arxiv:2405.04434

14. Tao, C., et al. (2024). Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies. arxiv:2407.13623

15. Feng, G., et al. (2023). Towards Revealing the Mystery behind Chain of Thought. arxiv:2305.18654

16. Raposo, D., et al. (2024). Mixture-of-Depths: Dynamically allocating compute in transformer-based language models. arxiv:2404.02258

17. Loshchilov, I., et al. (2024). nGPT: Normalized Transformer with Representation Learning on the Hypersphere. arxiv:2410.01131

18. Zhu, Z., et al. (2024). Hyper-Connections. arxiv:2409.19606

19. Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arxiv:2403.19887

---

↑ 上级 · A. 基础理论