Pruned CTC:18万词表训练步显存降5.1倍

Pruned CTC for Memory-Efficient Large-Vocabulary ASR Training

Yifan Yang, Xiaoyu Yang, Zengrui Jin, Xian Shi, Yuxuan Wang, Yu Xi, Ziyang Ma, Qi Chen, Ruiyang Xu, Hui Wang, Dongchao Yang, Jin Xu, Xie Chen

eess.AS

2026-09-27

Pruned CTC只对batch内目标词加blank做对齐,全词表归一化保证损失梯度等价。18万词表全步显存降5.1倍,Qwen3上相对CE慢7% WER、识别快7到10倍。

这篇在解决什么

把预训练 LLM 接到语音识别,主流做法是 token 级交叉熵加自回归解码,论文称作 LLM-CE。训练逐步吐字,流式还得给每个音频块配文本,通常靠外部对齐。CTC(Connectionist Temporal Classification,用 blank 把变长语音帧对齐到整句文本)能用整句转录当监督,离线和流式都能跑,但标准实现会把每一帧对整个词表的激活全部写进显存。

LLM 原生词表动辄十几万到十八万。帧数乘词表这张 logit 矩阵,在编码器激活之外再占一块,大 batch 直接 OOM。合法 CTC 对齐路径只用 blank(对齐占位、坍缩时删掉的空类)和这句出现过的 token,一个 batch 里这些类通常只是全词表的一小截,其余列几乎只为 softmax 归一化而存在。

想用 LLM 原生词表做 CTC,卡的是实现把不该常驻的东西常驻了。

方法

Pruned CTC 把对齐和归一化拆开。动态规划只在 blank 加本 batch 出现过的目标 token 上跑;softmax 仍对全词表做。没选中的类通过归一化项继续进损失,对词表头的梯度仍然稠密。

路径搜索只走这句用得上的站,每一步的概率分母仍按全线路算。精确算术下,这个词表收缩与全词表 CTC 的损失和一阶梯度完全等价。没选中的类占有率为零,梯度就是它自己的 softmax 概率。

实现上还有两处设计:

LLM-CTC 把预训练因果 LLM 接到这段损失上,词表和因果注意力都留着。语音编码器 SPEAR-XLarge v2 和 LLM head 冻结,只训 projector 和 LoRA。speech-readout 在语音位置做 CTC,每个位置只能看到已到的语音前缀。query-readout 在语音后追加同一可学习向量的 Tn 份拷贝,每个 query 能看到整句。最终用后者:因果注意力下它仍能看全句,流式时也能看当前块。blank 占用 tokenizer 没用到的一行,再学一个标量偏置 η。

流式把句子切成 2 秒块,编码器与 LLM 共用一段有界左历史。CTC 损失仍对整句转录一次求和,训练不需要块级对齐。推理时 LLM 复用 KV cache,CTC 前缀 beam 跨块续搜。

结果

编码器侧用 Zipformer-M,一张 H100 80G 上扫词表。head 加 loss:Pruned CTC 在 1 万到 18 万类之间几乎不随词表涨;M=16,000、V=180,000 时从 43.0 GiB 降到 1.29 GiB,33.3 倍。全训练步在标准 CTC 还能跑完的最大词表上:

设置标准 CTCPruned CTC显存步时
M=16k, V=180k51.7 GiB10.1 GiB5.1×+17%
M=32k, V=130k77.3 GiB17.9 GiB4.3×+17%
M=64k, V=50k76.2 GiB33.8 GiB2.3×+14%

V=180k 且 M=32k 或 64k 时,标准 CTC 直接 OOM,Pruned CTC 用 18.2 和 34.5 GiB 跑完一步。

LibriSpeech、GigaSpeech、AISHELL-1 上 WER/CER 与标准 CTC 对齐到小数点后两位:test-clean 都是 2.52,test-other 都是 6.02;GigaSpeech test 11.27 对 11.24;AISHELL-1 test 4.80 对 4.79。

SPEAR-Qwen3 在 GigaSpeech 上,0.6B 到 32B 六档,LLM-CTC 相对 LLM-CE 的 test WER 相对差距不超过 7%,RTF(识别耗时除以音频时长)快 7.2 到 10.3 倍。32B 上 CE 为 9.41、CTC 为 9.94,RTF 从 0.1546 降到 0.0150。4B 及以上的 LLM-CTC 也低于同编码器的 SPEAR-XLarge CTC(10.41)和 RNN-T(10.26)。A100 上 8B 全步显存从 56.06 降到 34.67 GiB,少 38.2%,步时多 7.5%;head 加 loss 从 29.21 降到 1.14 GiB,25.7 倍。不剪枝的 reduced CTC 全步数字几乎一样,省显存主要来自词表收缩和分块重算。

4B 读出消融:因果注意力下 query-readout 的 test WER 为 10.21,speech-readout 为 10.46;speech 加双向全注意力到 10.16,但双向过不了带 KV cache 的流式。Qwen3-ASR 0.6B 和 1.7B 微调流式后,test 上相对配套离线模型 WER 相对升高不到 3%(0.6B 离线 10.73、2 秒历史 10.98;1.7B 离线 10.16、2 秒历史 10.35)。左历史从 2 秒加到 8 秒几乎不动。

为什么重要

Cut Cross-Entropy 已经证明 token 级 CE 不必物化全词表 logit。CTC 没有帧级标签,后验靠前向后向从整句推,所以必须把对齐支撑集和全词表归一化拆开。对已经在用 Zipformer 加 k2 配方的人,换 loss 就能在大词表下继续训,三套语料上准确率不用让。

对 LLM ASR,这是一条可跑的非自回归适配:因果注意力留下、原生词表留下、流式不用块级对齐。用 7% 相对 WER 换 7 到 10 倍识别速度,延迟敏感的离线识别可以认真看。新声学模型谈不上,CTC 这个目标在原生 LLM 词表上变得训得动。

LLM 全步显存只少了两三成,权重和中间激活还在。25 倍发生在 head 加 loss 那一层。

局限与存疑

论文没有单独的 Limitations 节。有限 beam 让损失高于全路径求和,剪枝边界可能让损失不连续;β=100 把丢掉的后验压到 10^{-11},但那是随机初始化上的估计,不是训练全程的证明。

选中类上两项梯度接近抵消时,舍入会占主导。附录里权重梯度相对误差最大到 1.21×10^{-4},主因是 FP32 累加。流式训练并行算所有块,推理用 KV cache 逐块算,精确算术下等价;BF16 加 FlashAttention 2 时,0.6B 的相对 logit 误差约 0.054 到 0.057,greedy 帧标签会对不上。只有 FP32 非融合注意力才基本对齐。训练和部署精度栈不同,流式一致性没有被钉死。

LLM-CTC 冻编码器、冻 head,只动 LoRA 和 projector。相对 CE 的 7% 有多少来自 CTC 的帧独立假设、有多少来自冻结,没有拆开。LLM 实验几乎全在 GigaSpeech 英文。左历史加长几乎不降 WER,有界上下文有没有被用上,没有进一步诊断。

术语

原文与代码

相关论文

全部论文解读