ICLR 2025

Published as a conference paper at ICLR 2025
SCALING STICK-BREAKING ATTENTION:
AN EFFICIENT IMPLEMENTATION
AND IN-DEPTH STUDY
Shawn Tan
MIT-IBM Watson AI Lab
shawntan@ibm.com
Songlin Yang
MIT
yangsl66@mit.edu
Aaron Courville
Mila, Université de Montréal
courvila@mila.quebec
Rameswar Panda
MIT-IBM Watson AI Lab
rpanda@ibm.com
Yikang Shen
MIT-IBM Watson AI Lab
yikang.shen@ibm.com
ABSTRACT
The self-attention mechanism traditionally relies on the softmax operator, neces-
sitating positional embeddings like RoPE, or position biases to account for token
order. But current methods using still face length generalisation challenges. We
investigate an alternative attention mechanism based on the stick-breaking process
in larger scale settings. The method works as follows: For each token before the
current, we determine a break point, which represents the proportion of the stick,
the weight of the attention, to allocate to the current token. We repeat this on the
remaining stick, until all tokens are allocated a weight, resulting in a sequence
of attention weights. This process naturally incorporates recency bias, which has
linguistic motivations for grammar parsing (Shen et al., 2017). We study the im-
plications of replacing the conventional softmax-based attention mechanism with
stick-breaking attention. We then discuss implementation of numerically stable
stick-breaking attention and adapt Flash Attention to accommodate this mechanism.
When used as a drop-in replacement for current softmax+RoPE attention systems,
we find that stick-breaking attention performs competitively with current methods
on length generalisation and downstream tasks. Stick-breaking also performs well
at length generalisation, allowing a model trained with 211 context window to
perform well at 214 with perplexity improvements.
 https://github.com/shawntan/stickbreaking-attention
1
INTRODUCTION
The Transformer architecture (Vaswani et al., 2017) uses a self-attention mechanism based on the
softmax ope...