Causal Attention with Lookahead Keys

Zhuoqing Song; Peng Sun; Huizhuo Yuan; Quanquan Gu

Home/Archive/arXiv

arXiv · 2509.07301

Causal Attention with Lookahead Keys

Zhuoqing Song·Peng Sun·Huizhuo Yuan·Quanquan Gu

View original source ↗Read source document ↗Export metadata ↓

Abstract

In standard causal attention, each token's query, key, and value (QKV) are static and encode only preceding context. We introduce CAuSal aTtention with Lookahead kEys (CASTLE), an attention mechanism that continually updates each token's keys as the context unfolds. We term these updated keys lookahead keys because they belong to earlier positions yet integrate information from tokens that appear later relative to those positions, while strictly preserving the autoregressive property. Although the mechanism appears sequential, we derive a mathematical equivalence that avoids explicitly materializing lookahead keys at each position and enables efficient parallel training. On language modeling benchmarks, CASTLE consistently outperforms standard causal attention across model scales, reducing validation perplexity and improving performance on a range of downstream tasks.

Explore related subjects

cs.CL cs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeX RIS

Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu. 2025-09-29. Causal Attention with Lookahead Keys. https://arxiv.org/abs/2509.07301

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

In-context Learning vs. Instruction Tuning: The Case of Small and Multilingual Language Models

Instruction following is a critical ability for Large Language Models to be used directly by humans. This often requires supervised fine-tuning on curated instruction datasets, sometimes complemented with an alignment step. However, in multilingual scenarios, obtaining high-quality data for these stages remains challenging, motivating the exploration of In-Context Learning (ICL) as a possible alternative. In this work, we study whether ICL can serve as a substitute for Instruction Tuning in multilingual language models, while also examining how the comparison changes with model scale. Our results indicate that a gap remains between ICL and Instruction Tuning, motivating further research to reduce it.

cs.CL↗

arXiv2026

R3: Robust Rubric-Agnostic Reward Models

Reward models are essential for aligning language model outputs with human preferences, yet existing approaches often lack both controllability and interpretability. These models are typically optimized for narrow objectives, limiting their generalizability to broader downstream tasks. Moreover, their scalar outputs are difficult to interpret without contextual reasoning. To address these limitations, we introduce R3, a novel reward modeling framework that is rubric-agnostic, generalizable across evaluation dimensions, and provides interpretable, reasoned score assignments. R3 enables more transparent and flexible evaluation of language models, supporting robust alignment with diverse human values and use cases. Our models, data, and code are available as open source at https://github.com/rubricreward/r3.

cs.CL↗

arXiv2026

Post-Training Large Language Models via Reinforcement Learning from Self-Feedback

Large Language Models (LLMs) often produce plausible but poorly-calibrated answers, limiting their reliability on reasoning-intensive tasks. Recent research suggests that Chain-of-Thought (CoT) reasoning paths are inherent in pre-trained LLMs and can be elicited by simply altering the decoding process, where the presence of a CoT path correlates with higher answer confidence. Building on these insights, we present Reinforcement Learning from Self-Feedback (RLSF), a post-training stage that utilises the model's intrinsic confidence as a self-generated reward. By generating multiple CoT decoding beams from a frozen LLM, we compute the confidence of each final answer span and rank the resulting traces accordingly to create synthetic preferences. These preferences are subsequently utilised to fine-tune the policy through standard preference optimisation, requiring no human labels, gold answers, or externally curated rewards. RLSF simultaneously (i) refines the model's probability estimates--restoring well-behaved calibration--and (ii) strengthens step-by-step reasoning, yielding improved performance on arithmetic reasoning and multiple-choice question answering. By converting a model's own uncertainty into structured self-feedback, RLSF affirms reinforcement learning on intrinsic model behaviour as a principled and data-efficient component of the LLM post-training pipeline. Our results demonstrate that leveraging these inherent reasoning capabilities provides a robust path for enhancing model reliability without manual prompt engineering or external supervision.

cs.CL↗