Gated Bidirectional Linear Attention for Generative Retrieval

2026-06-05

AI Summary

This work addresses the high computational complexity and inference latency of standard self-attention in recommender systems when processing long user sequences. To reconcile recommendation quality with inference efficiency, the authors propose a hybrid encoder architecture that alternately stacks self-attention layers with a novel Gated Bidirectional Linear Attention (GBLA) mechanism. GBLA achieves linear-complexity bidirectional attention by integrating lightweight local causal convolutions, a sequence-level key-gating soft-forgetting mechanism, and gated RMSNorm. Experiments demonstrate that the proposed method matches the recommendation performance of full self-attention on the Yandex Music dataset while accelerating single-layer inference by up to 8.2×. Furthermore, strong generalization capability is validated on the Amazon dataset.

Abstract

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.

Problem

• generative retrieval
• bidirectional attention
• long sequence modeling
• recommendation systems
• attention efficiency

Innovation

• Gated Bidirectional Linear Attention
• linear-time attention
• generative retrieval
• long-sequence modeling
• hybrid encoder