Wiola 13M，一种用于参数高效小型语言模型的门控螺旋注意力架构

# Wiola 13M：一种用于参数高效小型语言模型的门控螺旋注意力架构
来源：https://arxiv.org/html/2608.14604
###### 摘要

一千万到一亿参数范围的小型语言模型因其适用于设备端推理、快速实验和可控科学验证而备受关注，但大多数此类模型只是简单复用标准Transformer模块，未针对小规模场景进行调整。我们提出Wiola，一种仅解码器语言模型，其创新点集中于每一层的三个即插即用组件：首先，螺旋旋转位置编码通过逐维度缓慢增长因子扰动标准旋转频率，使相位轨迹随深度呈扇形展开，在不增加参数的情况下提升长程辨别能力；其次，门控螺旋注意力引入基于查询流因果累积统计的逐内容自适应标量门控，以可忽略的成本实现隐式的可微软头选择机制；最后，蝶形前馈模块用乘法交互和块内旁路路径替代传统扩展层，在匹配四倍门控线性单元参数量的同时改善浅层网络的梯度流动。我们形式化了每个组件，推导了精确的参数与计算预算，并证明门控注意力在全序列训练与缓存自回归解码间存在精确等价关系（经数值验证），确保推理时无近似误差。同时，我们在标准微型故事语料库上提供了完全可复现的训练与评估方案。参考实现已作为开源包发布，支持权重直接发布。

## I引言

大型语言模型推动了自然语言处理的快速进步，但其计算和内存需求使得设备端部署及计算预算有限的实验室难以实现。这重新激发了对小型语言模型（约一千万到一亿参数）的关注，此类模型可在单个消费级加速器上数小时内完成训练，并以较小内存占用在本地运行。在微型语料库上的研究表明，即使参数少于一千万的模型，在受控数据分布下也能生成流畅连贯的文本[8]，这使得小规模环境成为独立于数据规模研究架构的高效场景。

尽管如此，主流做法仍是直接缩放标准Transformer模块，沿袭了为大两到三个数量级模型优化的设计选择。在小规模场景中，三个特性尤为关键：位置信息必须高效编码（因序列短且每个表征维度都稀缺）；注意力头数量少，即使一个头无法专化也会浪费大量容量；前馈模块占非嵌入参数预算的主导，其参数效率基本决定了网络的计算最优形态。

我们引入Wiola模型，通过在Transformer层中引入三项针对性修改来应对这些挑战，同时保持外围训练流程标准化。我们的贡献包括：
- •**螺旋旋转位置编码**：一种无参数的旋转位置编码[2]改进方案，通过跨维度对缓慢增长的缩放因子调整旋转频率，使相位轨迹在频谱深度上呈扇形展开。该修改在极限情况下可精确退化为标准旋转编码。
- •**门控螺旋注意力**：基于查询投影因果累积均值计算的逐头标量门控，乘法应用于softmax前的分数。门控使无效头自抑制（隐式软头选择），每层仅增加数百个参数。
- •**蝶形前馈模块**：采用乘法扩展与块内旁路路径，在匹配四倍门控线性单元参数量的同时改善浅层网络的梯度传播。

除架构外，我们提供了精确的参数与计算预算，证明并通过数值验证门控注意力在并行训练路径与缓存解码路径间保持精确等价。这一特性在实践中至关重要，因为朴素定义的注意力门控可能悄然破坏训练与生成的同一性，导致样本质量下降。

## II相关工作

### II-A位置编码

原始Transformer通过固定正弦嵌入注入顺序信息[1]。旋转位置编码[2]则按绝对位置成比例旋转查询和键，使注意力分数仅依赖相对偏移。因其优雅的外推性和无需学习参数，旋转编码已成为现代仅解码器模型的默认选择[5]。其他方案如加法线性偏置[7]也针对长度外推设计。我们的螺旋编码属于旋转家族，仅修改频率调度，在保留相对偏移特性的同时改变相位在频谱上的分离方式。

### II-B前馈与门控线性单元

门控线性单元及其变体通过引入两个投影的乘法交互改进Transformer前馈模块[3,12]。这些变体通常使用三个投影矩阵并缩小内部宽度以保持参数恒定。蝶形模块延续这一传统，但增加了显式块内旁路，我们发现这在层数较少时有利于梯度流动。

### II-C注意力效率与门控

大量工作通过共享键值对[6]或重构注意力计算来降低注意力成本。这些方法针对大模型场景和长上下文的内存带宽。我们的目标不同：不是降低注意力成本，而是通过内容自适应门控调制每个头（数量较少）的贡献，使其更有效，这更接近条件计算而非注意力压缩。

### II-D小型语言模型

微型故事场景[8]证明小型模型能从受控分布中学习生成连贯文本，计算最优缩放分析[9]阐明了参数与词元的平衡方式。Wiola正针对这一场景设计，旨在成为其中干净可复现的基线模型。

## III模型架构

### III-A概述与符号

Wiola是预归一化仅解码器Transformer。令$d$表示隐藏宽度，$L$为层数，$H$为注意力头数，$d_h=d/H$为每头宽度。长度为$T$的词元序列经嵌入后，通过$L$个相同的解码器层处理，归一化，最后通过权重与输入嵌入共享的头投影为词表逻辑值。每层应用两个残差子块：

$$\mathbf{h}' = \mathbf{h} + \mathrm{GSA}\!\left(\mathrm{RMSNorm}(\mathbf{h})\right) \quad (1)$$
$$\mathbf{h}'' = \mathbf{h}' + \mathrm{FFN}\!\left(\mathrm{RMSNorm}(\mathbf{h}')\right) \quad (2)$$

其中$\mathrm{GSA}$为门控螺旋注意力，$\mathrm{FFN}$为蝶形模块。均方根归一化[4]应用于所有位置：

$$\mathrm{RMSNorm}(\mathbf{x}) = \frac{\mathbf{x}}{\sqrt{\frac{1}{d}\sum_{j=1}^d x_j^2 + \varepsilon}} \odot \boldsymbol{\gamma} \quad (3)$$

包含可学习增益$\boldsymbol{\gamma} \in \mathbb{R}^d$和小常数$\varepsilon$。图1展示整体数据流，图3详述注意力块。

> **图1：Wiola解码器**。每层$L$应用预归一化门控螺旋注意力子块和预归一化蝶形前馈子块，各包裹在残差连接中。输入嵌入与输出投影共享权重。

### III-B螺旋旋转位置编码

标准旋转编码[2]将每个头向量划分为$d_h/2$个坐标对，将第$i$对在位置$m$旋转角度$m\theta_i$，采用几何频率调度：

$$\theta_i = \mathrm{base}^{-2i/d_h}, \quad i=0,\dots,\tfrac{d_h}{2}-1 \quad (4)$$

其中$\mathrm{base}=10^4$。螺旋编码将每个频率乘以缓慢增长的无参数因子：

$$\varphi_i = 1 + \alpha\,\frac{\sqrt{i+1}}{\sqrt{d_h/2}} \quad (5)$$
$$\tilde{\theta}_i = \theta_i\,\varphi_i \quad (6)$$

包含单一固定超参数$\alpha$。该因子从最低对的1单调增长至最高对的$1+\alpha$，使高频对获得比例提升，其相位随位置分离更快，而最低对几乎保持不变以保留局部连续性。令$\alpha=0$可精确恢复式(4)，故螺旋编码是旋转编码的严格推广。图2绘制了最小配置所用头宽下的因子与生成频率。

当$\Theta_{m,i} = m\tilde{\theta}_i$时，应用于查询或键向量$\mathbf{u}$坐标对$(2i,2i+1)$的旋转为：

$$\begin{bmatrix} u'_{2i} \\ u'_{2i+1} \end{bmatrix} = \begin{bmatrix} \cos\Theta_{m,i} & -\sin\Theta_{m,i} \\ \sin\Theta_{m,i} & \cos\Theta_{m,i} \end{bmatrix} \begin{bmatrix} u_{2i} \\ u_{2i+1} \end{bmatrix} \quad (7)$$

因变换是按位置旋转且对查询和键相同，位置$t$的旋转查询与位置$s$的旋转键的点积仍为相对偏移$t-s$的函数，故保留了旋转编码的相对偏移特性。

> **图2：头宽$d_h=32$、$\alpha=0.05$的螺旋旋转位置编码**。(a)展示了式(5)中的无参数乘性因子$\varphi_i$，从约1.0125单调增长至1.05。(b)在对数轴上对比标准频率与螺旋频率，扰动对高频对最大。

### III-C门控螺旋注意力

令$\mathbf{X} \in \mathbb{R}^{T \times d}$为归一化层输入。查询、键和值通过线性投影形成并拆分为$H$个头：$\mathbf{Q}^{(h)}, \mathbf{K}^{(h)}, \mathbf{V}^{(h)} \in \mathbb{R}^{T \times d_h}$。螺旋编码(7)应用于查询和键。

#### 因果门控

门控通过头平均上下文向量及其因果累积均值总结截至当前位置的查询流：

$$\mathbf{m}_s = \frac{1}{H}\sum_{h=1}^H \mathbf{q}^{(h)}_s \in \mathbb{R}^{d_h} \quad (8)$$
$$\mathbf{c}_t = \frac{1}{t}\sum_{s=1}^t \mathbf{m}_s \in \mathbb{R}^{d_h} \quad (9)$$

两层门控网络随后为每个头生成一个标量：

$$\mathbf{g}_t = \sigma\!\big(\mathbf{W}_2\,\mathrm{SiLU}(\mathbf{W}_1\mathbf{c}_t) + \mathbf{b}_2\big) \in (0,1)^H \quad (10)$$

其中$\mathbf{W}_1 \in \mathbb{R}^{H \times d_h}$，$\mathbf{W}_2 \in \mathbb{R}^{H \times H}$。偏置$\mathbf{b}_2$初始化为零，故训练开始时每个门控约为0.5，注意力接近无门控基线。因$\mathbf{c}_t$仅依赖$s \leq t$的位置，门控严格遵循因果关系。

#### 门控分数

对于头$h$，门控在softmax前乘以缩放分数：

$$A^{(h)}_{t,s} = g^{(h)}_t \,\frac{\big\langle \mathbf{q}^{(h)}_t,\, \mathbf{k}^{(h)}_s\big\rangle}{\sqrt{d_h}} + M_{t,s} \quad (11)$$
$$\mathbf{z}^{(h)}_t = \sum_{s=1}^T \mathrm{softmax}_s\!\big(A^{(h)}_{t,\cdot}\big)\,\mathbf{v}^{(h)}_s \quad (12)$$

其中$M$为因果掩码，$M_{t,s}=0$当$s \leq t$，否则$-\infty$。逐头输出经拼接和投影：

$$\mathrm{GSA}(\mathbf{X})_t = \big[\mathbf{z}^{(1)}_t\,\|\cdots\|\,\mathbf{z}^{(H)}_t\big]\,\mathbf{W}_O \quad (13)$$

在softmax前应用门控意味着当$g^{(h)}_t$较小时，该行趋近均匀分布，头贡献甚微，梯度压力将始终无效头的门控驱向零。这是一种隐式完全可微的软头选择，无需辅助稀疏损失。默认情况下，门控从旋转前的查询计算，使其依赖内容而非绝对位置，该选择作为配置标志公开。

> **图3：门控螺旋注意力**。查询和键接受螺旋编码。并行分支从查询的因果累积均值导出逐头门控，并在softmax前乘入分数。解码时只需缓存运行累积和这一额外量。

### III-D蝶形前馈模块

前馈模块通过单次投影将隐藏状态扩展为两部分，应用乘法门控，并在降维投影前添加块内旁路：

$$[\mathbf{a} \| \mathbf{b}] = \mathbf{x}\,\mathbf{W}_{\mathrm{up}}^\top, \quad \mathbf{a},\mathbf{b} \in \mathbb{R}^{d_{\mathrm{ff}}} \quad (14)$$
$$\mathbf{r} = \mathrm{SiLU}(\mathbf{a}) \odot \mathbf{b} + \mathbf{W}_{\mathrm{bp}}\,\mathbf{x} \quad (15)$$
$$\mathrm{FFN}(\mathbf{x}) = \mathbf{W}_{\mathrm{down}}\,\mathbf{r} \quad (16)$$

其中$\mathbf{W}_{\mathrm{up}} \in \mathbb{R}^{2d_{\mathrm{ff}} \times d}$，$\mathbf{W}_{\mathrm{down}} \in \mathbb{R}^{d \times d_{\mathrm{ff}}}$为投影矩阵，$\mathbf{W}_{\mathrm{bp}} \in \mathbb{R}^{d \times d}$为旁路投影。该模块在参数量上等同于四倍门控线性单元（含两个$4d \times d_{\mathrm{ff}}$投影），但通过SiLU门控与旁路路径的乘法交互，可在浅层网络中改善梯度流。