Gated Bidirectional Linear Attention for Generative Retrieval

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.

• The paper introduces GBLA, a linear-time bidirectional attention mechanism that mitigates quadratic overhead while preserving generative retrieval quality.
• It employs Conv1D local causal mixing, sequence-level key gating, and gated RMSNorm to enable efficient long-context modeling and robust performance.
• Empirical evaluations demonstrate that GBLA matches full self-attention on recall metrics while achieving up to 8.2× faster inference on extended sequences.

Gated Bidirectional Linear Attention for Efficient Generative Retrieval

Introduction

Generative retrieval in recommendation leverages encoder-decoder transformers to process sequential user interaction histories and autoregressively generate item recommendations. As user histories become increasingly long—specifically in streaming and large-scale online services—the quadratic computational complexity of standard softmax-based self-attention emerges as a significant latency bottleneck for the encoder component. Empirically, bidirectional attention masks in the encoder yield substantive improvements in retrieval metrics over causal attention. However, most scalable attention proposals focus on the causal (unidirectional) case, leaving a considerable gap for efficient, high-quality bidirectional attention mechanisms suitable for inference- and training-time acceleration in industrial-scale retrieval systems.

Methodology

To mitigate the quadratic overhead while preserving bidirectional context, the paper introduces Gated Bidirectional Linear Attention (GBLA), a linear-time bidirectional attention operator. The architectural innovations in GBLA build on kernelized linear attention, extending it via three lightweight but effective techniques:

1. Local Causal Mixing (Conv1D): A single-layer 1D convolution is applied to input embeddings prior to QKV projections, improving extraction of local short-term dependencies and pattern mixing at negligible cost.
2. Sequence-Level Key Gating: GBLA introduces an explicit, learnable scalar gate for each token's key, functioning as a soft-forgetting mechanism and allowing the model to assign adaptive importance to historical tokens, thereby modulating long-range context accumulation.
3. Gated RMSNorm Output: Following recent advances in state-space and linear attention models, a gated RMSNorm with SiLU gating is applied to the output, encouraging information selectivity and stabilizing training dynamics.

The overall architecture is shown in (Figure 1):

Figure 1: The Gated Bidirectional Linear Attention (GBLA) architecture employs Conv1D, per-token key gating, and gated RMSNorm atop a kernelized bidirectional attention layer.

This design maintains linear complexity in both training and inference by replacing the expensive quadratic attention aggregation with kernelized vector-space pooling, and by fusing the three modules for efficient execution on hardware accelerators.

Empirical Evaluation

Large-Scale Production Benchmarks

On an extensive industrial dataset from Yandex Music (400M training samples, up to 32K-length histories, 9-layer bidirectional transformer encoder), the impact of replacing bidirectional self-attention blocks with GBLA is systematically evaluated. Particularly:

• Hybrid Stacking: Replacing all encoder self-attention layers with linear attention degrades retrieval quality. A hybrid stack—one self-attention layer followed by two GBLA layers ([SA, LA, LA] repeat)—achieves parity with fully bidirectional self-attention on Recall@K metrics, as shown for all $K \in \{10,100,1000\}$.
• Bidirectionality: Switching the encoder to a purely causal attention mask causes a substantial performance drop, underscoring the necessity of bidirectional connectivity for high-quality generative retrieval.

Scaling and Ablations

GBLA demonstrates consistent quality preservation with increasing history length. At sequence lengths $L = 8192$, the marginal drop in Recall@1000 is less than $0.15\%$ absolute, while at shorter histories ($L \leq 2048$), performance stays at near-equivalence to the original model.

The contribution of each GBLA component is quantitatively supported: removal of Conv1D, key gating, or gated RMSNorm individually yields incremental degradation in recall, with their aggregate effect being notably larger. This validates the necessity of each module as a non-trivial contributor to quality.

Inference and Training Efficiency

Latency benchmarks demonstrate that GBLA becomes increasingly advantageous over CUDA-optimized FlashAttention-v3 as sequence length grows.

Figure 2: Encoder inference latency for GBLA versus FlashAttention-v3 across sequence length scales on H100 GPUs.

At $L=32768$, per-layer inference is accelerated by a factor of $8.2\times$ relative to FlashAttention-v3. For production-scale batch sizes and sequences, step-time acceleration (training and inference) is substantive in the long-context regime, which is precisely where retrieval quality is most sensitive to context truncation.

Results on Public Academic Benchmarks

With the GRID experimental protocol and several Amazon product datasets (max length 128), Tiger+GBLA (hybrid) and the vanilla Tiger model show negligible differences in Recall@K and NDCG@K. For short sequences, quadratic attention is not a bottleneck and GBLA's efficiency advantage is not observed. Thus, recommendations for academic datasets with limited sequence length remain unchanged; the primary benefit of GBLA emerges in large-scale, long-context scenarios.

Theoretical and Practical Implications

The introduction of GBLA demonstrates that bidirectional linear attention, augmented with token gating, local pattern mixing, and output gating, is feasible and effective for generative retrieval. It addresses a critical aspect of the transformer scalability problem: how to preserve retrieval quality when the model is exposed to sequences far beyond typical training or inference lengths, without resorting to lossy input compression or inference-time sampling.

Practically, the hybrid encoder design provides a simple yet robust recipe for scaling deployed transformer-based recommenders, ensuring that deep bidirectional context is accessible without incurring prohibitive compute or latency costs. The public release of code indicates ease of adoption.

Theoretically, the results establish that with appropriate architectural choices, linear-complexity attention can closely approach the empirical optimality of full self-attention even in information-redundant retrieval tasks, provided bidirectional context is preserved and attention is appropriately regularized.

Future Directions

Potential avenues for further development include integration with recent state-space and delta-rule architectures, hardware specialization for GBLA-like primitives, and adapting gating mechanisms for online continual learning in dynamic recommender environments. Investigating more expressive or data-dependent kernelizations for the attention feature map may further close the quality gap or provide robustness gains.

Conclusion

GBLA provides a scalable, high-fidelity solution for bidirectional long-sequence modeling in generative retrieval, matching standard attention quality under appropriate hybridization while yielding significant latency and throughput improvements for long histories (2606.07317). As retrieval models are increasingly adopted in production recommender systems, efficient bidirectional linear attention of this type will be essential for maintaining both user experience and infrastructure efficiency.