Flash RNN-T：内存高效的联合网络与 Transducer 损失计算

arXiv:2610.23409v1 公告类型：新 摘要：循环神经网络 Transducer（RNN-T）仍是流式语音识别的一种实用选择，但训练需要大型中间张量。我们提出了 Flash RNN-T，这是一种联合网络与损失实现，通过专用内核结合了可选的分块重计算，覆盖完整的对齐空间和词汇表。在中等规模和长序列工作负载上，分块和非分块 Flash 共同构成了完整的实测吞吐量—内存前沿。在 H100 上，只要能够运行，非分块 Flash 在这些工作负载上的吞吐量比最快基线高出 24%–39%，内存减少 26%–56%。分块通过重计算进一步降低内存占用。在中等规模工作负载上，当两者每个批次都使用一次反向传播调用时，分块 Flash 的内存占用比该基线低 9 倍，而吞吐量仅低 16%。在长序列工作负载上，使用 80 GB 内存时，所有采用该反向传播调度方式的测试替代方案都会耗尽内存，而分块 Flash 仍可运行。这些比较均采用各实现中内存占用最低、且吞吐量至少达到峰值吞吐量 95% 的配置。