Attention Sinks › Beyond StreamingLLM: Learnable Sink Tokens
Language AI HandbookAll chapters Attention Sinks › Beyond StreamingLLM: Learnable Sink TokensOne section · Sections in this chapter11 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The original StreamingLLM approach preserves the first tokens from the input sequence. A more principled approach is to train models with explicit, learnable sink tokens. These tokens are:
- Added to the vocabulary as special tokens
- Always prepended to the input sequence during training
- Optimized end-to-end to serve as effective attention sinks
Models trained this way develop more efficient sink representations. The sink tokens learn key vectors that are optimally positioned to capture excess attention, rather than relying on whatever tokens happen to appear at the start of the sequence.
Learnable sink tokens are special tokens added to a model's vocabulary and trained end-to-end to serve as attention sinks. Unlike implicit sinks that emerge from training on natural text, learnable sinks are explicitly optimized to absorb excess attention weight.
The training procedure is straightforward:
- Add sink tokens to the tokenizer vocabulary
- Prepend these tokens to every training sequence
- Train the model normally, allowing the sink token embeddings to learn
During inference, you prepend the same sink tokens to the input, and they immediately serve their purpose without needing to preserve specific input tokens.
The advantage of learnable sink tokens over naturally occurring ones is consistency and efficiency. When a model relies on its BOS token as a sink, it is repurposing a token that was also trained to carry other information, such as signaling the start of a sequence to the model's positional encoding. The BOS token's key vector has to serve two masters: representing "start of sequence" for the generation process and serving as an attention attractor for all subsequent positions. These goals are not always aligned. A dedicated learnable sink token has a single purpose, so its key vector can be optimized entirely for the attractor role. The result is a more compact, more reliable sink that often requires fewer sink tokens to achieve the same generation quality.
This design choice is increasingly common in recently developed models. If you are building a model from scratch, or fine-tuning a base model for production streaming, adding one or two learnable sink tokens to the vocabulary is a low-cost, high-value investment. The extra tokens add negligible computation per step while providing a principled foundation for unlimited streaming generation.
def prepare_input_with_sink_tokens(tokenizer, text, num_sink_tokens=4):
"""
Prepare input with explicit sink tokens prepended.
Args:
tokenizer: Tokenizer with sink tokens in vocabulary
text: Input text to tokenize
num_sink_tokens: Number of sink tokens to prepend
Returns:
Token IDs with sink tokens at the beginning
"""
# Assume sink tokens are named [SINK_0], [SINK_1], etc.
sink_token_ids = [
tokenizer.convert_tokens_to_ids(f"[SINK_{i}]")
for i in range(num_sink_tokens)
]
# Tokenize the actual text
text_tokens = tokenizer.encode(text, add_special_tokens=False)
# Combine: sinks first, then text
return sink_token_ids + text_tokens