高效小型语言模型的Wiola架构

# 用于高效小语言模型的 Wiola 架构 ††感谢：本研究作为独立研究贡献进行。未获得外部资助。来源：https://arxiv.org/html/2607.01394

###### 摘要

我们提出 Wiola，一个完全原创的、从第一性原理构建的小语言模型 (SLM) 架构，与任何现有模型家族（包括 GPT、LLaMA、Mistral 和 Falcon）没有结构上的延续性。Wiola 引入了五个独立的创新组件：(i) 螺旋旋转位置编码 (SRPE)，它将 token 位置嵌入到一个三维螺旋流形上，结合了绝对、相对和分层位置信号；(ii) 门控跨层注意力 (GCLA)，它为每个解码器层提供对前两层压缩摘要的软跨层注意力访问，以实现层间一致性；(iii) 自适应 Token 合并 (ATM)，它会动态合并网络中间层中语义冗余的相邻 token，以在不损失信息的情况下降低注意力复杂度；(iv) 双流前馈网络 (DSFF)，它用由学习到的逐维度门融合的两个并行流取代传统的 MLP；(v) WiolaRMSNorm，一种改进的归一化方法，引入了一个学习到的逐维度偏移向量，以防止表示坍缩。我们提供了完整的数学推导、架构框图、复杂度分析，以及与 GPT-2、LLaMA-2 和 Mistral 的系统性比较。Wiola 以四种尺寸（1.2 亿、3.6 亿、7 亿和 15 亿参数）发布，完全兼容 HuggingFace Transformers 生态系统，所有 22 项架构单元测试均通过。

## I. 引言

Transformer[1 (https://arxiv.org/html/2607.01394#bib.bib1)] 推动了自然语言处理的显著进步。然而，主导模型家族——GPT[2 (https://arxiv.org/html/2607.01394#bib.bib2)]、LLaMA[4 (https://arxiv.org/html/2607.01394#bib.bib4)]、Mistral[5 (https://arxiv.org/html/2607.01394#bib.bib5)] 及其衍生模型——共享相同的结构谱系，仅在位置编码或注意力分组上存在渐进式差异。这种保守性留下了一些基本的架构问题：不同的位置几何能否更好地捕捉多尺度语言结构？层间信息路由能否改善生成文本中的长程连贯性？能否利用 token 级别的冗余来降低二次注意力成本？Wiola 是一个全新设计的 SLM，通过五个新颖的架构组件来解决这三个问题。每个子组件都源自独立的数学原理，并经过验证在结构上与所有先前发表的公式不同。本文的主要贡献如下：

1.  1. SRPE：一种三维螺旋位置编码，将绝对、相对和分层位置结合在一个统一的流形上，无需额外参数。
2.  2. GCLA：门控跨层注意力，通过压缩的层摘要提供层间连贯性，计算开销可忽略不计。
3.  3. ATM：中间层的动态贪心 token 合并，在训练期间减少 5-9% 的注意力 FLOPs，并精确恢复原始长度。
4.  4. DSFF：一个双流并行 FFN，带有学习到的逐维度融合，将局部和全局特征提取分开。
5.  5. WiolaRMSNorm：改进的 RMS 归一化，带有逐维度偏移，可抵消深层堆栈中的表示坍缩。
6.  6. 一个生产级实现，包含 22 个通过单元测试和完整的 HuggingFace Hub 集成。

## II. 相关工作

### II-A 位置编码

绝对正弦编码[1 (https://arxiv.org/html/2607.01394#bib.bib1)]和可学习的绝对编码[3 (https://arxiv.org/html/2607.01394#bib.bib3)]无法泛化到训练长度之外。相对编码如 ALiBi[7 (https://arxiv.org/html/2607.01394#bib.bib7)]和 T5-bias[8 (https://arxiv.org/html/2607.01394#bib.bib8)] 在注意力 logits 中编码成对偏移。RoPE[6 (https://arxiv.org/html/2607.01394#bib.bib6)] 将位置编码为复值旋转，确保注意力仅取决于相对偏移 \(p - q\)。扩展方法（如 YaRN[9 (https://arxiv.org/html/2607.01394#bib.bib9)]）重新参数化相同的平面二维圆。Wiola 的 SRPE 是首个将位置放置在具有双旋转角和一个正弦径向分量的三维螺旋上的编码，通过解析方式编码多尺度结构，无需学习参数。

### II-B 注意力变体

多查询注意力 (MQA)[11 (https://arxiv.org/html/2607.01394#bib.bib11)]和分组查询注意力 (GQA)[10 (https://arxiv.org/html/2607.01394#bib.bib10)] 减少了 KV 缓存内存。滑动窗口注意力[5 (https://arxiv.org/html/2607.01394#bib.bib5)] 将二次成本限制在局部窗口内。编码器-解码器模型中存在跨层注意力，但在仅解码器的自回归语言模型中不存在。GCLA 是首次将来自压缩的*先前层摘要*的跨层注意力注入到每个解码器层中的公式。

### II-C 前馈网络

SwiGLU[12 (https://arxiv.org/html/2607.01394#bib.bib12)]和 GELU[13 (https://arxiv.org/html/2607.01394#bib.bib13)] 变体的单流 MLP 无处不在。混合专家 (MoE)[15 (https://arxiv.org/html/2607.01394#bib.bib15)] 将 token 稀疏路由到专家 FFN。DSFF 则不同：两个不同宽度和激活函数的并行*密集*流由一个学习到的逐维度门融合——不是稀疏路由，也不是单流。

### II-D Token 压缩

视觉 Transformer 的 token 合并 (ToMe[16 (https://arxiv.org/html/2607.01394#bib.bib16)]) 使用二分匹配。ATM 将相邻 token 余弦相似度合并应用于因果解码器中间三分之一层的语言模型隐藏状态——这是一种此前未被探索的迁移。

## III. 符号表示

标量：斜体 (\(x, d, T\))。向量：粗体小写 (\(\bm{x}\))。矩阵：粗体大写 (\(\mathbf{W}\))。拼接：\([\bm{a}; \bm{b}]\)。逐元素乘积：\(\odot\)。Sigmoid：\(\sigma(x) = (1 + e^{-x})^{-1}\)。\([n] \triangleq \{0, \dots, n-1\}\)。表 I (https://arxiv.org/html/2607.01394#S3.T1) 列出了 wiola-360m 配置的核心超参数符号及其默认值。

表 I：核心超参数符号（wiola-360m 默认值）

## IV. Wiola 架构

### IV-A 宏观结构

Wiola 是一个自回归、仅解码器的语言模型。Token ID 被嵌入为 \(\mathbf{X}^{(0)} \in \mathbb{R}^{T \times d}\)，经过 \(L\) 个解码器层，归一化后通过一个绑定的线性头投影为 logits。对于层 \(\ell \in [L]\)：
\[
\begin{aligned}
\tilde{\mathbf{X}}^{(\ell)} &= \operatorname{WRMSNorm}_{\ell}\!\left(\mathbf{X}^{(\ell)}\right), \tag{1}\\
\mathbf{A}^{(\ell)} &= \operatorname{GCLA}_{\ell}\!\left(\tilde{\mathbf{X}}^{(\ell)}, \mathcal{C}^{(\ell)}\right), \tag{2}\\
\mathbf{X}^{(\ell+\frac{1}{2})} &= \mathbf{X}^{(\ell)} + \mathbf{A}^{(\ell)}, \tag{3}\\
\hat{\mathbf{X}}^{(\ell)} &= \operatorname{WRMSNorm}_{\ell}^{\prime}\!\left(\mathbf{X}^{(\ell+\frac{1}{2})}\right), \tag{4}\\
\mathbf{F}^{(\ell)} &= \operatorname{DSFF}_{\ell}\!\left(\hat{\mathbf{X}}^{(\ell)}\right), \tag{5}\\
\mathbf{X}^{(\ell+1)} &= \mathbf{X}^{(\ell+\frac{1}{2})} + \mathbf{F}^{(\ell)}. \tag{6}
\end{aligned}
\]
ATM 在训练期间插入到 (1) 和 (2) 之间，适用于中间三分之一的层。输出 logits 为：
\[
\mathbf{Z} = \operatorname{WRMSNorm}_{\mathrm{final}}\!\left(\mathbf{X}^{(L)}\right) \mathbf{W}_{\mathrm{head}}, \quad \mathbf{W}_{\mathrm{head}} = \mathbf{E}^{\top} \in \mathbb{R}^{d \times V}. \tag{7}
\]

### IV-B 层块图

图 1 (https://arxiv.org/html/2607.01394#S4.F1) 展示了完整的 Wiola 解码器层。参见图注

图 1：Wiola 解码器层。橙色虚线箭头：来自先前层的跨层摘要 \(\mathcal{C}^{(\ell)}\) 注入到 GCLA。ATM 仅在训练期间在中间三分之一的层中激活。

## V. WiolaRMSNorm

标准 RMSNorm[17 (https://arxiv.org/html/2607.01394#bib.bib17)] 归一化：
\[
\mathrm{RMSNorm}(\bm{x}) = \bm{\gamma} \odot \frac{\bm{x}}{\operatorname{RMS}(\bm{x})}, \quad \operatorname{RMS}(\bm{x}) = \sqrt{\tfrac{1}{d}\textstyle\sum_{i=1}^{d} x_i^2 + \epsilon}. \tag{8}
\]
它无法移动层分布的零参考点。Dong 等人[19 (https://arxiv.org/html/2607.01394#bib.bib19)] 表明，深度注意力网络会遭遇*表示坍缩*，其中隐藏状态收敛到一个退化的低秩子空间。仅靠重缩放无法抵消这一点。WiolaRMSNorm 引入了一个学习到的逐维度偏移 \(\bm{\delta} \in \mathbb{R}^{d}\)，它在*归一化之前*移动输入：
\[
\boxed{\operatorname{WRMSNorm}(\bm{x}) = \bm{\gamma} \odot \frac{\bm{x} + \bm{\delta}}{\sqrt{\tfrac{1}{d}\sum_{i=1}^{d} (x_i + \delta_i)^2 + \epsilon}}}. \tag{9}
\]
令 \(\bm{z} = \bm{x} + \bm{\delta}\)，则 \(\operatorname{WRMSNorm}(\bm{x}) = \bm{\gamma} \odot \bm{z} / \operatorname{RMS}(\bm{z})\)。令 \(\bm{\delta} = \bm{0}\) 则完全恢复 (8)，因此 WiolaRMSNorm 严格泛化了 RMSNorm。对 \(\delta_i\) 的梯度为：
\[
\frac{\partial \mathcal{L}}{\partial \delta_i} = \frac{\gamma_i}{r}\!\left(\frac{\partial \mathcal{L}}{\partial \hat{x}_i} - \frac{z_i}{d r^2}\sum_k \gamma_k \frac{\partial \mathcal{L}}{\partial \hat{x}_k} z_k\right), \quad r = \operatorname{RMS}(\bm{z}), \tag{10}
\]
通常非零，确保 \(\bm{\delta}\) 在训练期间偏离 \(\bm{0}\)。每层相对于 RMSNorm 的额外开销为 \(d\) 个参数（\(\bm{\delta}\)）。每个模型有 \(2L\) 个归一化，对于 wiola-360m，总开销为 \(2Ld = 32,768\) 个参数（占总参数的 \(0.009\%\)）。图 2 (https://arxiv.org/html/2607.01394#S5.F2) 显示了 WiolaRMSNorm 的数据流。参见图注

图 2：WiolaRMSNorm 数据流。偏移量 \(\bm{\delta}\) 在 RMS 计算*之前*移动输入，改变了归一化目标本身，而不是添加后归一化偏置。

## VI. 螺旋旋转位置编码 (SRPE)

### VI-A 动机

RoPE[6 (https://arxiv.org/html/2607.01394#bib.bib6)] 将位置 \(p\) 映射到每对维度的二维旋转，精确编码相对偏移，但仅表示一个位置尺度。自然语言至少有三种尺度：子词 token、短语级成分（3–15 个 token）和话语单元（句子、段落）。SRPE 将位置嵌入到一个*三维螺旋流形*上，在单个解析公式中编码所有三种尺度，且无需额外的学习参数。

### VI-B 数学推导

对于位置 \(p \in [T]\) 和维度对索引 \(j \in [d_h/2]\)：

步骤 1 — 主逆频率：
\[
\omega_j = \theta_0^{-2j/d_h}. \tag{11}
\]

步骤 2 — 双旋转角：
\[
\begin{aligned}
\theta_j^{(1)}(p) &= p \omega_j, &\quad \theta_j^{(2)}(p) &= \frac{p \omega_j}{k_s}, \tag{12}\\
\Theta_j(p) &= p \omega_j \!\left(1 + \tfrac{1}{k_s}\right). \tag{13}
\end{aligned}
\]

步骤 3 — 径向调制：
\[
r_j(p) = 1 + a_s \sin\!\left(p f_s \omega_j\right). \tag{14}
\]

步骤 4 — 编码系数：
\[
\begin{aligned}
c_j(p) &= r_j(p) \cos \Theta_j(p), \tag{15}\\
s_j(p) &= r_j(p) \sin \Theta_j(p). \tag{16}
\end{aligned}
\]

步骤 5 — 应用于查询 \(\bm{q} \in \mathbb{R}^{d_h}\)：
\[
\begin{aligned}
\operatorname{SRPE}(\bm{q}, p)_j &= q_j c_j(p) - q_{j+d_h/2} s_j(p), \tag{17}\\
\operatorname{SRPE}(\bm{q}, p)_{j+d_h/2} &= q_j s_j(p) + q_{j+d_h/2} c_j(p). \tag{18}
\end{aligned}
\]
相同的旋转应用于键 \(\bm{k}\)。矩阵形式：\(\operatorname{SRPE}(\bm{q}, p) = \mathbf{R}(p) \bm{q}\)，其中 \(\mathbf{R}(p) = \bigoplus_j \bigl[\begin{smallmatrix} c_j & -s_j \\ s_j & c_j \end{smallmatrix}\bigr]\)。

相对位置性质：来自第 \(j\) 对的点积贡献为：
\[
r_j(p) \, r_j(q) \cos\!\bigl(\Theta_j(p) - \Theta_j(q)\bigr), \tag{19}
\]
其中 \(\Theta_j(p) - \Theta_j(q) = (p - q) \omega_j (1 + 1/k_s)\) 仅取决于相对偏移 \(\Delta = p - q\)。径向乘积 \(r_j(p) r_j(q)\) 引入了受控的绝对位置依赖性，编码话语结构。表 II (https://arxiv.org/html/2607.01394#S6.T2) 比较了 SRPE 和 RoPE。

表 II：SRPE 与 RoPE 对比

## VII. 门控跨层注意力 (GCLA)

### VII-A 跨层摘要缓存

在层 \(\ell\) 产生 \(\mathbf{X}^{(\ell+1)} \in \mathbb{R}^{T \times d}\) 后，通过平均池化形成一个摘要：
\[
\bm{s}^{(\ell)} = \frac{1}{T} \sum_{t=1}^{T} \mathbf{X}^{(\ell+1)}_{t,:} \in \mathbb{R}^{d}. \tag{20}
\]
下一层的上下文矩阵使用最近的 \(\Lambda = 2\) 个摘要：
\[
\mathcal{C}^{(\ell+1)} = \bigl[\bm{s}^{(\ell-1)}; \bm{s}^{(\ell)}\bigr] \in \mathbb{R}^{\Lambda \times d}. \tag{21}
\]

### VII-B 带有 SRPE 和 GQA 的自注意力

投影：\(\mathbf{Q} = \tilde{\mathbf{X}} \mathbf{W}_Q\)，\(\mathbf{K} = \tilde{\mathbf{X}} \mathbf{W}_K\)，\(\mathbf{V} = \tilde{\mathbf{X}} \mathbf{W}_V\)，其中 \(\mathbf{W}_Q \in \mathbb{R}^{d \times H d_h}\)，\(\mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{d \times H_{\mathrm{kv}} d_h}\)。每个头应用 SRPE：\(\tilde{\mathbf{Q}}_h = \operatorname{SRPE}(\mathbf{Q}_h)\)，\(\tilde{\mathbf{K}}_h = \operatorname{SRPE}(\mathbf{K}_h)\)。对于头 \(h\)，GQA 组 \(g = h \bmod H_{\mathrm{kv}}\) 的因果自注意力：
\[
\begin{aligned}
\mathbf{A}_h &= \operatorname{softmax}\!\left( \frac{\tilde{\mathbf{Q}}_h \tilde{\mathbf{K}}_g^{\top} + \mathbf{M}}{\sqrt{d_h}} \right), \tag{22}\\
\mathbf{O}_h^{\mathrm{self}} &= \mathbf{A}_h \mathbf{V}_g, \tag{23}
\end{aligned}
\]
其中 \(\mathbf{M}\) 是因果掩码（对角线上方为 \(-\infty\)）。

### VII-C 跨层上下文子注意力