Autoregressive Modeling with Lookahead Attention

TL;DR

This paper introduces a novel Transformer-based autoregressive model that incorporates lookahead by extrapolating multiple future continuations, improving performance on various tasks by leveraging hypothetical future information.

Contribution

The paper proposes a new lookahead attention mechanism in autoregressive models, inspired by classical AI systems, to utilize future trajectory estimates for better predictions.

Findings

Outperforms standard Transformer models on multiple tasks.

Benefits from additional computation even without true lookahead.

Discusses potential architectures and speedups for future work.

Abstract

To predict the next token, autoregressive models ordinarily examine the past. Could they also benefit from also examining hypothetical futures? We consider a novel Transformer-based autoregressive architecture that estimates the next-token distribution by extrapolating multiple continuations of the past, according to some proposal distribution, and attending to these extended strings. This architecture draws insights from classical AI systems such as board game players: when making a local decision, a policy may benefit from exploring possible future trajectories and analyzing them. On multiple tasks including morphological inflection and Boolean satisfiability, our lookahead model is able to outperform the ordinary Transformer model of comparable size. However, on some tasks, it appears to be benefiting from the extra computation without actually using the lookahead information. We discuss possible variant architectures as well as future speedups.