Leviathan: Decoupling Input and Output Representations in Language Models¶

会议: NeurIPS 2026
arXiv: 2601.22040
OpenReview: https://openreview.net/forum?id=vGbvwN6dmK
领域: 预训练
关键词: 权重解耦, 可学习嵌入向量化, B样条基展开, 罕见词表征, 语言模型架构

一句话总结¶

本文提出 Leviathan 架构，通过可学习嵌入向量化（LEV）将输入词表索引连续化映射为输入表征，在几乎不增加总参数量（1.2B 规模仅增 0.2%）的前提下彻底解耦输入嵌入与输出分类头，并在长尾罕见词上取得高达 81% 的困惑度下降。

研究背景与动机¶

现代自回归语言模型普遍采用输入嵌入矩阵与输出投影矩阵绑定的权重共享（weight tying）范式。这一设计初衷是为了在受限参数预算下约束模型，但它强行把两项数学性质截然不同的任务耦合在一个矩阵上：输入端需要捕捉语义连续性与几何平滑度，使语义相近的词元在表示空间内彼此靠近；而输出端本质是一个跨大词表的高维分类器，必须在整个词汇空间中维持最大的判别能力与各向异性间隔。权重绑定迫使单一参数矩阵在两者之间妥协，形成了严重的表征几何瓶颈。

面对这一矛盾，最直接的直觉是解除权重绑定（untied embeddings），分别维护输入矩阵与输出矩阵。然而实验表明，简单的全解耦在有限计算预算（Chinchilla 最优规模）下不仅无法带来提升，反而因为在 200k 庞大词表上引入近 50% 的参数开销而产生严重的优化困难，在 8.4B 词元训练后验证损失仍然全面劣于权重绑定基线。低秩矩阵分解（如 ALBERT 风格的输入线性分解）虽然能压低输入端参数，却仅停留在刚性的低秩线性投影，不仅未能提供有效的表征归纳偏置，反而破坏了输入几何结构，表现甚至更差。

本文从表征参数化形式切入：输入表示需要的不是庞大离散的查表矩阵，而是具有平滑插值能力的连续生成函数。本文的核心 idea 是：提出可学习嵌入向量化（LEV）层替代离散输入嵌入表，利用组合式进位索引与可分离 B 样条张量积构建紧凑的连续生成映射，从而以极低参数代价释放全表达能力的解耦输出投影头，并在长尾罕见词上实现显著的梯度共享与平滑归纳偏置。

方法详解¶

整体框架¶

Leviathan 架构的核心是将离散词元查表嵌入替换为可学习嵌入向量化（Learned Embedding Vectorization, LEV）模块，同时保留完全独立的输出投影头（Unembedding Head）。对于任意词元索引 $i$，LEV 并不直接索引 $V \times D$ 维的密集参数矩阵，而是首先通过 $k$ 进位分解为低维离散坐标，在共享种子码本中聚合成种子表征，经由多头连续变换与一维二次 B 样条基展开进行连续插值，最后通过低秩张量积映射到 Transformer 的隐藏维度 $D$。

整个输入到输出的前向计算流程包含：词元进位分解与种子聚合、多头 B 样条基展开与张量积生成输入向量、标准 Transformer 骨干解码计算、以及完全解耦的高维分类头预测下一词分布。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入词元索引 i"] --> B["组合式进位索引分解<br/>生成 k 维坐标并检索共享码本"]
    B --> C["多头 B 样条连续基展开<br/>有界潜空间归一化与样条函数评估"]
    C --> D["低秩张量积投影<br/>跨头求和生成紧凑输入嵌入 E_i"]
    D --> E["标准 Decoder-only Transformer 骨干<br/>L 层自注意力与 SwiGLU 前馈层"]
    E --> F["完全解耦的密集输出头<br/>独立高维分类器映射至词表空间"]
    F --> G["预测下一词概率分布"]

关键设计¶

1. 组合式进位索引分解：跨词元梯度共享与参数压缩

离散查表矩阵之所以在罕见词上表现糟糕，是因为低频词在数十亿词元中仅出现寥寥数次，对应的独立嵌入行几乎无法接收到足够的梯度更新。为了让罕见词能够从频繁出现的相似结构词元中借力，LEV 引入了基于 $b$ 进制的组合式分解机制。对于词表大小为 $V$ 的词典，选定分解阶数 $k$ 与基数 $b$（满足 $b^k \ge V$），将标量索引 $i \in \{0, \dots, V-1\}$ 严格映射为 $k$ 元组 $(i_1, \dots, i_k)$，其中 $i_r \in \{0, \dots, b-1\}$。每个位置维护独立的共享码本 $C_r \in \mathbb{R}^{b \times d_{\text{seed}}}$，初始种子表征由各个分量码本向量直接求和得到： $$z(i) = \sum_{r=1}^{k} C_r[i_r]$$ 在 $V=200,376$ 的现代大词表设置下，选取 $k=3, b=59$（总容量 $59^3 = 205,379$）即可几乎无浪费地覆盖整个词表。该设计使得参数量从 $V \times D$ 大幅削减至 $k \cdot b \cdot d_{\text{seed}}$（仅数万参数），更重要的是赋予了索引空间内在的拓扑关联性，任意罕见词与常见词共享码本项，从而在反向传播中直接共享梯度信号。

2. 多头 B 样条连续基展开与低秩张量积：平滑连续归纳偏置

简单的离散分解或浅层 MLP 容易陷入欠拟合或缺乏有效的几何连续性归纳偏置。LEV 构建了 $h$ 个独立的连续生成头，首先将聚合的种子向量通过线性映射、层归一化和带缩放系数的 Sigmoid 函数投射到有界潜空间内： $$\tilde{z}_{\ell} = \sigma\left(\text{LN}\left(\frac{1}{2} W_{\text{seed},\ell} \, z(i)\right)\right) \in [0, 1]^{d_{\text{seed}}}$$ 其中 $\frac{1}{2}$ 的缩放因子将输入限制在 Sigmoid 的近线性活跃区以保持优化平稳。潜坐标的每个维度均通过包含 $\kappa$ 个控制节点的单变量二次 B 样条基进行连续插值拟合，并通过秩为 $r$ 的张量积算子进行特征组合，得到各头的连续映射 $\mathcal{M}_{\ell}(\tilde{z}_{\ell})$。最终词元嵌入由各头经线性变换后累加获得： $$E_i = \sum_{\ell=1}^h W_{\ell} \mathcal{M}_{\ell}(\tilde{z}_{\ell})$$ 这一设计将词元嵌入建模为定义在连续坐标上的平滑非线性函数，强迫表示空间具备局域平滑性（smoothness inductive bias），同时整套 LEV 仅占用 $1.21\text{M} + 512D$ 参数。

3. 输出投影头解耦：免额外参数代价的分类容量释放

在权重绑定的传统架构中，输出投影矩阵必须与输入矩阵共用 $V \times D$ 参数；若使用完全解耦的传统 Transformer，双矩阵将占用 $2VD$ 参数，导致词表占总参数比重过高（例如 400M 模型中膨胀至 611M，参数暴增 50%）。由于 LEV 将输入端参数量从数亿压缩至 2M 级别，模型得以毫无顾虑地保留一个完全独立、未受约束的密集输出投影矩阵 $W_{\text{out}} \in \mathbb{R}^{D \times V}$。在 1.2B 规模下，相比权重绑定的基准模型，Leviathan 增加的净参数量仅为 0.20%（FLOPs 仅增 0.26%），却彻底消除了输入平滑与输出判别之间的冲突，使输出头能够自由学习各向异性的锐利分类决策面。

损失函数 / 训练策略¶

模型采用标准自回归因果语言模型交叉熵损失训练。优化器配置为 AdamW（$\beta_1=0.9, \beta_2=0.999, \varepsilon=10^{-8}$），采用无权重衰减策略以避免不同嵌入参数化形式受到不对称的正则化干扰。学习率采用线性预热配合余弦衰减调度（衰减至峰值的 10%），梯度裁剪阈值设为 1.0。上下文长度固定为 512 词元，全局有效批大小为 512 个序列（每步 262,144 词元）。

实验关键数据¶

主实验¶

论文在 The Pile 语料库上进行了 200M、400M 和 1.2B 三种规模的严格受控对比，骨干网络完全一致。同时在 1.2B 规模下评估了 Wikitext-103 零样本迁移困惑度以及 6 个权威下游基准任务。

模型规模	模型架构	验证集损失 (nats)	验证集困惑度 (PPL)	Wikitext-103 PPL 改善	LAMBADA PPL 改善
200M (4.2B tokens)	Tied (基线)	3.073	21.61	基准	基准
200M (4.2B tokens)	Leviathan	3.052	21.16 (-2.1%)	+2.4%	-33% (落后)
400M (8.4B tokens)	Tied (基线)	2.784	16.18	基准	基准
400M (8.4B tokens)	Leviathan	2.724	15.24 (-5.6%)	+7.0%	+11%
1.2B (24.1B tokens)	Tied (基线)	2.489	12.05	基准	基准
1.2B (24.1B tokens)	Leviathan	2.402	10.94 (-9.2%)	+12.3%	+30%

在下游零样本评测任务中，1.2B 规模下的 Leviathan 全面超越 Tied 基线：HellaSwag 达到 33.2% (vs 31.7%)，ARC-Easy 达到 42.6% (vs 38.7%)，ARC-Challenge 达到 23.5% (vs 23.1%)，PIQA 达到 63.5% (vs 61.3%)，WinoGrande 达到 51.5% (vs 48.5%)，LAMBADA 达到 33.4% (vs 32.6%)，平均绝对准确率提升达 2.3 个百分点。

消融实验¶

在 400M 规模下，作者对比了输入输出参数化形式（包括解耦上限 Untied 与线性分解 Albert），并对 LEV 内部组件进行了递进式消融。

配置	验证集损失 (nats)	相对 Tied 改善率 (%)	说明
Tied (权重绑定基线)	2.784	0.00%	标准输入输出共享矩阵
Untied (完全独立矩阵)	2.801	-0.61%	参数量膨胀 50% (611M)，优化困难导致掉点
Albert (低秩线性分解)	2.807	-0.83%	参数与 LEV 对齐，但破坏了几何结构且无非线性
Flat MLP 消融	2.753	+1.13%	仅用双层 SiLU MLP 连续生成，验证连续参数化收益
Factored MLP 消融	2.742	+1.50%	引入 $b$ 进制分解但无 B 样条，验证进位索引收益
Leviathan (完整模型)	2.724	+2.13%	进位分解 + B 样条平滑基展开 + 张量积组合

关键发现¶

罕见词是性能飞跃的核心来源：按词元在预训练语料中的出现频次分桶分析表明，在 1.2B 规模下，出现频次仅 1-9 次的长尾罕见词上的困惑度大幅下降了 81%（损失降低 1.65 nats）；而对于出现频率极高（>1M 次）的常见词，两者的性能差异几乎为零。这证实连续函数参数化在梯度极度匮乏的数据稀疏区提供了决定性的插值平滑效果。
训练扩展效应递增且早期优化滞后被抹平：在 200M 小模型上，Leviathan 在训练前期（3B tokens 之前）因连续映射较难初始化而略逊于 Tied 基线；但随着模型规模提升至 1.2B，优化劣势完全消失，全程单调领先，达到 Tied 基线最终损失所需的样本数减少至原来的 $1/2.1$（计算吞吐折算后挂钟时间节约 50%）。
词表缩放鲁棒性：在 50k（GPT-2）、100k（cl100k）与 200k（o200k）不同词表规模下，Leviathan 相对 Tied 的相对损失改善均稳定在 1.95%~2.12% 之间，证明该架构收益不依赖特定的词表尺寸。

亮点与洞察¶

跳出「查表 vs 解绑」的非此即彼思维：以往工作认为解除绑定必然伴随巨大的参数膨胀与优化恶化，而本文发现只需改变输入端的表征生成方式（从离散离散查表变为连续函数生成），就能既享受解绑输出分类头的容量红利，又免除优化发散的困扰。
长尾表征退化的数学良药：语言模型词表天然符合齐普夫定律，绝大部分词汇是极端低频的长尾词。传统查表嵌入使这部分表征严重各向异性退化；而组合式连续样条参数化通过基数共享，强制低频词嵌入与高频词共享潜空间几何与梯度更新。
零推理开销工程复用：尽管 LEV 在训练阶段引入了轻微的计算前向开销，但在部署推理阶段，整个词表的连续映射可以一次性离线缓存为静态 $V \times D$ 嵌入表（单张 H200 上仅需 76.8ms），推理期的内存与前向耗时与传统模型完全一致。

局限与展望¶

模型规模尚未外推至百亿级：受算力限制，论文验证的最大模型为 1.2B、预训练词元为 24.1B。在 7B+ 或 Llama 等数十万亿词元充分训练的超大模型中，嵌入层占参数比重更小，连续生成带来的收益是否会被超大规模参数稀释仍需验证。
单一语料库与分词体系限制：目前评测全部建立在 The Pile 语料和 Byte-level BPE 分词器上，尚未在多语言混排、代码专用或形态丰富语言（如阿拉伯语、芬兰语）上展开全面对比。
未来方向：探索将 LEV 的组合式连续映射扩展至输出解嵌入端，以及结合多模态词元的连续嵌入生成。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将可分离 B 样条与组合进制编码引入语言模型输入嵌入，彻底重构了输入参数化形式]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 3 个模型尺度、3 种词表大小、全面的机制消融与细致的词频分桶长尾分析]
写作质量: ⭐⭐⭐⭐⭐ [数学推导清晰严密，论据链条扎实，实验设计严谨自洽]
价值: ⭐⭐⭐⭐⭐ [为解决词表膨胀与长尾表征退化提供了极具潜力的全新底层架构思路，推理期无缝兼容现有引擎]