Contextual Position Encoding: Learning to Count What's Important

TL;DR

Contextual Position Encoding (CoPE) introduces context-dependent position encoding for Transformers, enabling positions to be conditioned on surrounding context rather than solely by token indices. By computing gates that select which prior tokens to count and interpolating fractional positions, CoPE supports multi-unit addressing (e.g., words, sentences) and improves generalization to longer contexts. Across toy tasks and language/code modeling, CoPE outperforms traditional absolute/relative PE and demonstrates robust in-distribution and out-of-distribution performance, along with insightful attention patterns. The work highlights a practical, scalable approach to richer position signaling that can benefit long-range reasoning in diverse domains.

Abstract

The attention mechanism is a critical component of Large Language Models (LLMs) that allows tokens in a sequence to interact with each other, but is order-invariant. Incorporating position encoding (PE) makes it possible to address by position, such as attending to the i-th token. However, current PE methods use token counts to derive position, and thus cannot generalize to higher levels of abstraction, such as attending to the i-th sentence. In this paper, we propose a new position encoding method, Contextual Position Encoding (CoPE), that allows positions to be conditioned on context by incrementing position only on certain tokens determined by the model. This allows more general position addressing such as attending to the $i$-th particular word, noun, or sentence. We show that CoPE can solve the selective copy, counting and Flip-Flop tasks where popular position embeddings fail, and improves perplexity on language modeling and coding tasks.

Figures (7)

• Figure 1: Contextual Position Encoding (CoPE). Standard position encoding methods such as Relative PE are based on token positions. In contrast, CoPE computes gate values conditioned on the context first, then uses that to assign positions to tokens using a cumulative sum. This allows positions to be contextualized, and represent the count of different units like words, verbs or sentences. CoPE operates on each attention head and so can attend to different position types on each. In this example, attending to the last sentence using Relative PE is challenging, and the best it can do is a decaying attention ("recency bias"). CoPE can count the sentence endings and simply attend to position "0".
• Figure 2: CoPE outperforms relative PE on the counting task, especially with less training data of the task.
• Figure 3: Counting task test error rates (%) for different number of variables.
• Figure 3: Generalization to longer context length. After training on the Wikitext-103 language modeling task with a context size of 1024 (left) and 512 (right), we evaluate the model on longer context sizes and report the validation perplexity. CoPE generalizes well, outperforming existing PE methods, especially when evaluation context size is much larger than training context size (right).
• Figure 4: CoPE can focus attention on abstract elements like current paragraph (left) and section (right). Here we show attention induced by position alone on Wikitext-103. Since CoPE is contextualized, it can attend to paragraphs and sections by their position. On the left, the segments are found to be separated by newline tokens (indicated by black plus signs), while the right is separated by section titles like "= = Description = =" (similarly marked).