提出线性时间双向注意力机制,显著加速长序列推荐系统。
Gated Bidirectional Linear Attention for Generative Retrieval
- 用轻量组件扩展核化线性注意力,实现双向长序列建模。
- 在32768长度下比FlashAttention-v3快8.2倍,保持高质量。
- 适合大规模流媒体推荐系统,兼顾速度与精度。
在推荐系统中,生成式召回通常采用编码器-解码器结构:编码器处理用户行为历史,自回归解码器生成推荐项。在大规模流媒体服务中,活跃用户的历史记录随时间急剧增长,导致编码器成为主要延迟瓶颈,因为Softmax注意力随序列长度呈二次方增长。实验表明,使用双向注意力可显著提升质量,但多数亚二次注意力方法聚焦于因果注意力。本文提出门控双向线性注意力(GBLA),一种线性时间的双向注意力层,通过三个轻量组件扩展核化线性注意力:局部因果混合(Conv1D)、序列级键门控实现软遗忘,以及门控RMSNorm输出。在大型Yandex Music数据集上,混合编码器(自注意力与GBLA以1:2比例交替,即一个SA块后接两个GBLA块)达到与双向自注意力相当的质量。在H100 GPU上,当历史长度为32768时,GBLA相比FlashAttention-v3最高实现8.2倍单层加速。此外,该混合设计在公开的Amazon基准测试中也表现出色,持续保持自注意力的召回质量。
原文摘要 · Abstract (English)
In recommender systems, generative retrieval typically uses an encoder-decoder setup: an encoder processes a user interaction history, and an autoregressive decoder then generates recommended items. In large-scale streaming services, active users accumulate very long histories over time. As histories grow, the encoder becomes a major latency bottleneck because softmax attention scales quadratically with sequence length. In our experiments, using bidirectional attention in the encoder substantially improves quality. However, most sub-quadratic attention methods focus on causal attention. We propose Gated Bidirectional Linear Attention (GBLA), a linear-time bidirectional attention layer that extends kernelized linear attention with three lightweight components: local causal mixing (Conv1D), sequence-level key gating for soft forgetting, and a gated RMSNorm output. On a large-scale Yandex Music dataset, a hybrid encoder that interleaves self-attention (SA) and GBLA in a 1:2 ratio (one SA block followed by two GBLA blocks) matches bidirectional self-attention quality. On H100 GPUs, GBLA reaches up to an $8.2\times$ single-layer speedup at a history length of 32768, compared to FlashAttention-v3. Finally, we show that the same hybrid design generalizes beyond our proprietary setting, consistently preserving self-attention retrieval quality on public Amazon benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。