[论文解读] Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions

Stefano Massaroli, Michael Poli+12 more|VU Research Portal|Oct 28, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出 Laughing Hyena，一种蒸馏方法，可从预训练的长卷积序列模型（LCSMs）中提取紧凑的循环状态空间模型（SSMs），在自回归生成过程中实现每个 token 的 O(1) 计算和内存开销。该方法在 1.3B 参数规模下，推理吞吐量比 Transformer 高 10 倍，比 Hyena 高 1.5 倍，且蒸馏后性能无损失。

ABSTRACT

Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressive inference workloads -- naively requiring a full pass (or caching of activations) over the input sequence for each generated token -- similarly to attention-based models. In this paper, we seek to enable $\mathcal O(1)$ compute and memory cost per token in any pre-trained long convolution architecture to reduce memory footprint and increase throughput during generation. Concretely, our methods consist in extracting low-dimensional linear state-space models from each convolution layer, building upon rational interpolation and model-order reduction techniques. We further introduce architectural improvements to convolution-based layers such as Hyena: by weight-tying the filters across channels into heads, we achieve higher pre-training quality and reduce the number of filters to be distilled. The resulting model achieves 10x higher throughput than Transformers and 1.5x higher than Hyena at 1.3B parameters, without any loss in quality after distillation.

研究动机与目标

• 在长卷积序列模型（LCSMs）中实现常数时间、常数内存的自回归生成，克服标准推理中 O(K) 内存和 O(K²) 计算的瓶颈。
• 开发一种蒸馏框架，从预训练的卷积层中提取低维、稳定的状态空间模型（SSMs），同时保持模型质量。
• 通过在通道间使用权重共享的滤波器重新设计 Hyena 架构，以提升预训练质量和蒸馏效率。
• 通过减少内存占用并支持 LCSMs 中的循环推理，实现高吞吐量、大批次生成。

提出的方法

• 应用有理插值和模型降阶方法，从预训练 LCSMs 的卷积滤波器中提取紧凑的 SSM。
• 引入一种受重心法和 Prony 类方法启发的因子化解耦参数化方法，以提升 SSM 稳定性并避免数值问题。
• 以卷积滤波器上的差异度量作为近似目标，确保在下游任务中的兼容性。
• 使用 Hankel 算子谱分析确定 SSM 的最优状态维度 d。
• 对 Hyena 模块进行重构，通过在通道间引入权重共享，提升有效滤波器维度并降低蒸馏复杂度。
• 采用平衡截断和模态截断方法进行 SSM 模型降阶，误差界基于 Gramian 特征值推导得出。

实验结果

研究问题

• RQ1 能否从预训练的长卷积模型中提取出紧凑的循环 SSM，以实现每个 token 的 O(1) 推理成本？
• RQ2 在 SSM 蒸馏中，何种状态维度 d 能在近似误差与模型效率之间实现最优权衡？
• RQ3 如何改进 SSM 参数化方法，以避免数值不稳定性并提升蒸馏过程中的收敛性？
• RQ4 对 Hyena 模块的架构修改能否提升预训练质量并降低蒸馏成本？
• RQ5 将 LCSMs 蒸馏为 SSM 是否能保持下游性能，同时支持高吞吐量生成？

主要发现

• Laughing Hyena 在 1.3B 参数规模下，峰值吞吐量比同类 Transformer 高 10 倍，比 Hyena 高 1.5 倍，且蒸馏后性能无损失。
• 在 1.3B 参数规模下，Laughing Hyena 在相同内存约束下生成 512 个 token 所需内存仅为 Transformer 的 1/3。
• 该模型在 K 个 token 生成中实现 O(d) 内存和 O(dK) 时间复杂度，而 KV 缓存的 Transformer 为 O(K) 和 O(K²) 复杂度。
• 在 Hyena 模块中引入通道间权重共享，提升了预训练困惑度，并减少了需蒸馏的滤波器数量。
• 平衡截断和模态截断方法在部分层中表现出非单调的误差减少趋势，表明对数值稳定性较为敏感。
• 将滤波器差异度量作为近似目标，确保了在多样化下游任务中的鲁棒性。