ENTP: Encoder-only Next Token Prediction (2410.01600v3)

Abstract: Next-token prediction is conventionally done using decoder-only Transformers with causal attention, as this approach allows for efficient reuse of keys and values. What if we were not compute-limited, should we still use decoder-only Transformers? In this work, we introduce Encoder-only Next Token Prediction (ENTP). We explore the differences between ENTP and decoder-only Transformers in expressive power and complexity, highlighting potential advantages of ENTP in settings with unbounded compute. We introduce the $\operatorname{Count3}$ task and show, both theoretically and experimentally, that while ENTP can perform this task easily, a decoder-only Transformer cannot. Finally, we empirically demonstrate the superior performance of ENTP across representative tasks where next-token prediction based Transformers can be evaluated, including addition, in-context learning, and LLMing.

• The paper demonstrates that encoder-only Transformers achieve next-token prediction with superior generalization and expressivity compared to decoder-only architectures.
• The authors employ both theoretical analysis and empirical evaluations, highlighting encoder advantages through tasks like triplet-counting.
• The study underscores potential future directions, including hybrid models and computational optimizations to expand the capabilities of next-token prediction.

Encoder-only Next Token Prediction: A Comprehensive Overview

The paper "ENTP: Encoder-only Next Token Prediction" presents an exploration into using encoder-only Transformers for next-token prediction tasks, typically dominated by decoder-only architectures with causal attention. The authors challenge the established belief that causal attention is indispensable for preventing prediction leakage from future tokens. They posit that this architectural choice is driven more by efficiency than necessity, and introduce Encoder-only Next Token Prediction (ENTP) as a viable alternative with potential advantages.

Core Findings

The paper explores the expressive power and computational characteristics of encoder-only versus decoder-only Transformers, leading to several insights:

• Expressive Power: The authors demonstrate that the sets of functions expressible by encoder-only and decoder-only Transformers are not directly comparable. Intriguingly, there are functions each architecture can represent that the other cannot, as well as functions that both can express.
• Computational Complexity: Through theoretical analysis, it is established that encoders recompute attention for each token and have a time complexity of $O(n^3)$ for entire sequences, contrasting with $O(n^2)$ for decoders, which utilize cached computations. However, encoders' additional memory requirements on each token represent a different trade-off.
• The Triplet-Counting Task: A novel task, suitable for ENTP but challenging for decoders, is introduced. It highlighted the limitations of decoder architectures in representing specific functions. Small and large-scale experiments demonstrate ENTP's superior performance on this task.

Empirical Evaluation

The authors conduct multiple experiments to compare encoder-only and decoder-only models across several tasks:

• Realistic Tasks: They assess models on tasks requiring length generalization and in-context learning. Encoders consistently showed lower sample complexity and better generalization capabilities.
• Next-token Prediction with Large Models: ENTP outperformed decoder-only models in LLMing tasks, indicated by lower perplexity scores on the OpenWebText dataset.
• Fine-tuning LLMs: Attempts to finetune large decoder-only models on the Triplet-Counting task were less successful, aligning with the theoretical predictions of their expressive limitations.

Implications and Future Directions

The findings suggest that ENTP provides a promising alternative to traditional decoder-only architectures for tasks requiring increased expressivity. While the efficiency of current hardware technologies supports decoder-only models, the theoretical potential of ENTP for certain tasks cannot be overlooked.

The paper opens several avenues for future exploration:

• Optimizing Computational Efficiency: Given the higher computational cost of ENTP, developing methods to reduce this overhead will be crucial.
• Hybrid Architectures: Investigating whether combining encoder and decoder architectures could leverage the strengths of both is a promising direction.
• Expanding Task Domains: Applying ENTP to other complex modeling problems beyond language tasks may reveal additional benefits and insights.

Overall, this work provides a critical reevaluation of transformer architectures, highlighting the roles of model design choices, and paving the way for broader applications of ENTP in machine learning and artificial intelligence research.