Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study
Authors: Shawn Tan
Songlin Yang
Aaron Courville
Rameswar Panda
Yikang Shen
Venue: ICLR 2025 Poster (2025)
Keywords: transformer
attention
stick-breaking
softmax
length extrapolation
Primary area: foundation or frontier models, including LLMs
TL;DR: Using the stick-breaking process formulation as a replacement for softmax attention.

Abstract: The self-attention mechanism traditionally relies on the softmax operator, necessitating 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 implications 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 $2^{11}$ context window to perform well at $2^{14}$ with perplexity improvements.

Published as a conference paper at ICLR 2025
SCALING STICK-BREAKING ATTENTION: AN EFFICIENT IMPLEMENTATION
AND IN-DEPTH STUDY
Shawn Tan
Songlin Yang MIT
MIT-IBM Watson AI Lab shawnt and ibm. com
yangsl (etisss56mmmtt.dee
Aaron Courville
Rameswar Panda
Yikang Shen
Mila, Université de Montréal courvilafmile la. quebec
MIT-IBM Watson Al Lab rpanda s.bm. com
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-brekking process in larger scale settings. The method works 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 the current token. We repeat this on the remaining stick, until all tokens are allocated a weight, resulting in a sequence 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- 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 drop-in replacement for current softmax&RoPl attention systems, we find that stick- breckknn attention performs competitively with current methods length generalisation and downstream tasks. Stick- breckibe also performs well at length generalisation, allowing a model trained with 211 context window to perform well at 2 with perplexity improvements.
https://github. www.om//wwttt//tiiibboioooo
attention
INTRODUCTION
The Transformer architecture (Vaswani et al., 2017) uses self-attention mechanism based on the softmax operator that enables the model to weigh the importance of different tokens in the input data. However, the reliance softmax requires using positional embeddings introduce information about the order of tokens, as the attention mechanism itself is permutation-invariann, The sinusoidal position embedding as proposed in Vaswani et al. (2017) has since evolved relative positional embeddings (Shaw et al., 2018). Learned relative positional biases were used in the T5 model (Raffel al., 2020), and later fixed relative positional biases in Press et al. (2021). At the time of writing, a commonly used form of position embedding is RoPE (Su et al., 2021). Allen-Zhu & Li (2023) observe that, in context-free grammar parsing setting, attention mechanisms attend to the "most adjacent" non-terminal. This suggests an inclination to attend to the most recent entry that matches a given criteria. However, even with relative position information, it is possible to overfit on specific relative positions, resulting in failure to generalise. Kazemnejad et al. (2024) show that decoder-only Transformers with No Positional Embeddings (NoPE) can implicitly recover positional information, experimental results suggest that NoPE Transformers generalise better on length. While this is promising, a higher attention score from an irrelevant token in the sequence can function as a distractor (Kazemnejad al., 2024; Xiao al., 2023).

---

Published as a conference paper at ICLR 2025
Logits
q: Kb
Logits
-5-
V dhead
1.0
Stick-breakin Softmax
exp(z(e)
Softmax
A.j E1=1 1= (xxp(zx.))
0.0
1.0
Stick-breaking
A.,j = o(2a))
II ickej
(1 - o(x.,)))
0.5
0.0
Relative position
Figure Differences in formulation between breckibe and softmax. Stick-breaking assigns high weights to the most recent high logit, while softmax will assign equal weightage equal logits. (+) can any function R - (0, 1). In this paper we use ((x) I+empl-z)
The stick- breaking process may have properties that can alleviate the previously mentioned issues, and possess the 'most recent" bias from Allen Ailen-Zhu& & Li (2023) that we want. For token at position j attending to position i, suppose the attention weight is given by:
Ar, = 3rg
II (1- (1- = oz(,) (<K{)
II icksj
(1 (1-o (Zkg)),
where =kj are the attention logits. To illustrate via intuition, for Ai,j to be high, all 3x.j for < k have to be low. Conversely, as long as any fR.j for < k high, A..j will be low, as a token between i and has already been attended Shen al. (2017) makes similar observation Li (2023), and explicitly uses a stick- brickbbeaking process to model local structure. Csordás et al. (2021) has similar construction which they call Geometric attention after the Geometric distribution. Geometric distribution only has one parameter p. which gives the probability success per trial. The geometric distribution then gives the probability for which i trials needed for the first success: (1 (1-p)<-1p pp. But in both stick-breaking and in Geometric attention (Csordás et al., 2021), each is assigned a different value that corresponds the attention score between two tokens.
in Allen-Zht
In this paper, we expand upon prior work on this attention mechanism (Csordás et al., 2021; Shen al., 2023). We focus on the implications of replacing the softmax-base attention mechanism with the stick- process:
We compare the different properties of stick-breaking breaking attention against softmax attention,
2. We discuss numerically stable implementations of the bacckbreaking attention, and make stick-breaking amenable for large-scale training by implementing a kernel for stick-breaking in Triton,
3. We show the performance of stick-breaking attention on modelling, and evaluate 1B and 3B parameter models various NI.P tasks.
language
2 STICK-BREAKING ATTENTION
For sequence of L tokens, we have query q, E Rheadd key k; E Rhene and value v; E Rdhessd vectors for 1 iis L. Then the attention weight for token at attending to position computed by:
9 K:
1 Σ A
Arj = Big II (1 (k-)))), Big a (21,), Z1. / /dhead (<k<)
Equation the main difference between our proposal and softmax attention. As discussed earlier and in Csordás et al. (2021), this parameterisation biases towards recency. Specifically, for any pair of i and such that [j - él < [j - and ži,j Z then A., 2 A. Consequently, this imposes an ordering the attention, and we do not use position embeddings with the query and key embeddings.
2

---

Published as a conference paper at ICLR 2025
We consider two sets of logits, 2., and <i and their respective attention weights A., and A If A 2/,4 for i <. k < j. then A., = A This means that unlike softmax attention, high attention score further back in the sequence does not 'distract' from a more recent high 2,,j score. Further, if > As.j = 1, then the output is invariant to appending additional context earlier than i. We note that Σ A. I 1, which allows this attention mechanism to attend to nothing when all Bi, O. We discuss a strategy to deal with the remaining attention weight in Appendix ??.
RELATED WORK
Stick-breaking Process The stick-breaking process formulation of the Dirichlet process (Sethura- man, 1994) is also known as the GEM distribution, first coined in Ewens (1990) after Griffiths (1989), Engen (1975), and McCloskey (1965). The GEM is a specific case of what was known as a Residual Allocation Model (RAM; Allen & Lambie 1976). There are instances of the distribution being used as differentiable attention-like mechanism in neural models. Shen al. (2017) used breaking process for modelling language, and showed that the model can induce grammatical structure to some extent. Csordás et al. (2021) used boek-brea attention, which they refer to as Geometric attention, in bidirectional encoder set up. Shen al. (2023) used stick- breaking attention in decoding-onle setup, but does not explicitly study the properties of stick- breckbb
Softmax attention, Positional embeddings, and Length Generalisation Bondarenko et al. (2024) observe that softmax attention tends to attend to low-information tokens with high scores in order to 'do nothing'. Xiao et al. (2023) introduces attention sinks, a learnable token that the attention can assign attention weights to. Irie al. (2019) and Haviv et al. (2022) find that in a decoder-only setting, a Transformer with no positional emhedding can work fairly well. Kazemnejad et al. (2024) also found similar results, while also showing that NoPE has tendency to attend to the start of the sequence, while ALiBi (Press et al., 2021) has a tendency to only attend to the most recent tokens. However, Zhou al. (2024) later found that Transformers without position embeddings do not generalise to out-of-distribution sequence lengths for an addition task. Al present, Rotary Positional Embeddings (RoPE; Su et al. 2021) are the most commonly used position embedding. It encodes relative positions via multiplicative interactions with the key and query. RoPE has been found to generalise to out-o--disti lengths poorly (Press et al., 2021; Zhou et al., 2024; Kazemnejad et al., 2024), but a common trick to extend the context window RoPE-based Transfomers, is to use NWK-aware RoPE scaling (bloc97, 2023).
Conditional Computation The use of brickin for conditional computation has also been explored. Tan & Sim (2016) uses a the stick-breaking distribution as mixture over outputs for each layer in an MI.P for an acoustic model. Graves (2016) also suggested a similar formulation for language modelling. Later, Banino et al. (2021) and Tan et al. (2023) also use a stick-breaking formulation for dynamic depth modelling in a Transformer model. These prior works use conditional computation the depth of the model, while in our case, we use stick-breaking as method of restricting the computation wineth-wise.
Connection to Selective State-space Models Each stick- biikkbreakinn attention head at every time-step can be viewed as the hidden state of the final step of selective State-space Model (SSM;Gu & Dao 2023). For a given time-step j. consider the following SSM in its recurrent form and its corresponding convolutional form (as described in Merrill et al. 2024):
Let
وū.j = (1 -311) Ô Big Uj
(2)
then
0j 0
Σ ISiSi
Bag II (1 - Bk.s)) ickej
which is equivalent to the first term in Equation 1. Intuitively, this means that the output head of stick- breaking attention is equivalent to the end state of a single-gate selective SSM. Typically, an attention layer for Transformer with hidden dimension dhidden has h heads such that dhidden h dhead- For equivalence with an SSM, we need a constant query vector and each dimension as a separate head, egg. q 1 for all i, and h duidden, dhead
3

---

Published as a conference paper at ICLR 2025
Connection to Additive Relative Position Encoding (Additive RPE; Kazemnejad et al. 2024) Generally, Additive RPEs incorporate an added bias function g of the distance of the tokens - and the maximum length of the sequence L:
As exp (q) K; + b)
In the case of ALiBi (Press et al., 2021), this is a linear function -m - -). This implies that the attention weights will drop off exponentially the further j and are apart, regardless of the attention scores. In stick boock-breakigg, Equation has form that accounts for the scores from j to i with the bias b = Ek=it1 log (1 + exp(=x.)))- Specifically, log (1 + exp(=k,)) 2 m, then b - - (j 1). This suggests a learnable relative position bias that is dependent on the intermediate scores between, and
IMPLEMENTATION
Implementing stick-breaking attention naively in PyTorch results in realising the L matrix for the attention logits (where L is the length of the input). FlashAutention (Dao et al., 2022) reduces the memory footprint of attention by side-stepping the ( memory complexity of realising the attention matrix. In order to achieve this, it only realises blocks of the attention weights during computation, and accumulates the resulting weighted sum V,. We take similar approach to speeding up stick-breakin attention, allowing it to be used for longer sequences.
Forward Computing Equation directly will result in underflow issues, especially with lower precision training. We perform the operations in log- oeg-space which results in cumulative instead:
Ar, exp
log Ba. +
Σ
log (1 -Bxxs)) exp
Zt,j Σ log (1 exp(2x.5))
log
(4)
Where log (1 + exp(-)) :) is commonly known as the softplus operation. See Appendix A for the full derivation. We further numerically stabilise softplus with the following computation:
softplus(x) log (1 + exp(x)), if < 15
(5)
otherwise
to prevent overllowing of exp(x). To further speed up the computation of softplus, also write the operation Parallel Thread Execution (PTX).
Backward Let A.5 = =log Ai,j. then:
OC
0C
Oc
ac
ac
Ai.i
(6)
0A.j
dA.j
0zi
0
dA.s
Contribation from i.
Cantribution from hefore i.j
The above equations dictate the direction of the order of computation for implementation. For the forward pass (Eqn. 4), we compute from j to 1, backwards through time and accumulate
El ++1 log (1+ exp(2x.)) For backward pass (Eqn. 6), we compute from to j. accumulating 0c
4.1 TRITON SPECIFICS
We modify the Triton implementation Flash Attention for accelerating the stick-breaking attention. We used similar approach in Flash Attention 2 (Dao 2023; FA2), loading blocks of Q from High Bandwidth Memory (HBM) into Static random-accest memory (SRAM). In the forward pass, accumulate towards the start of the sequence, loading blocks of K and V the values finally writing a and as the output into HBM. the backward pass, we differ from the FA2 approach. FA2 accumulates V K, VV, and makes atomic adds toward VQ. which allows for fewer
4

---

Published as a conference paper at ICLR 2025
K, V
VV
atomic add
0
v
VQ
Accumulate O
Accumulate VQ
(a) Forward pass
(b) Backward pass
(c) Block skipping
Figure 2: Thread tile assignments for given attention head and sequence. Tiles coloured the same are processed by the same thread. For stick-breaking. the forward pass has to be computed from right-to-left, while the backward pass is computed in the reverese order. Uncoloured tiles not computed: upper right tiles are not used in causal language modelling, and in the case of block skipping, some blocks can be skipped if all entries have summed 1.
Algorithm FORWARD thread i
a 0, O
init a ditock x 1, : dutock X dinad I/ load from HBM to SRAM, Q dutock X dhead
Q < load block i of Q
for k in i. ..... do
K load block & of K
1/ K dhtock drnad
V load block V
: dulock dhead
Z + QK
1/ Z ditock X dolack
L - -softplus(2 Z)
// L dolock X dolock
A exp (Z + L
I/ cumulative sum right left. A dolock dblock
0 0 + AV
- a Σ. L
end for
store block i of M
II store from SRAM to HBM
0 store block of
synchronisations during the computation of the backward pass. However, in the case of stick-breaking, the accumulation of the gradients have to be in the reverse direction of the cumulative sum As result, we have to do atomic adds for both V K and VV. The resulting memory complexity of our approach is still OL), which allows to train large models with stick-breaking.
In our implementation, the block size dolock = 64. Algorithms details the forward pass (illustrated in Figure 2a), and Algorithm 2 details the backward pass (illustrated in Figure 2b).
Throughput We measure throughput Dolomite Engine (Mishra, 2024) with 1B class model on a node with 8 H100 GPUs. brick-breaking attains throughput of 21 billion tokens day. This is 29% performance drop compared FA2, which attains throughput of 29.5 billion tokens day. Despite the slowdown, the Triton implementation allows the method to be used for long sequences as the naive Torch implementation causes -ueoor- issues. Further optimisation of the kernel may be possible if written in CUDA, taking advantage of synchronisation features the language exposes in the hardware.
Conditional computation We can incorporate speedups due to the specific nature the stick- breaking process. Since, we accumulate AR.; Ut from the diagonal to = 1, when Σ A., = for some k. know that all A can skip all subsequent accumulations. We implement block skipping only the forward pass, and measured the time elapsed on subset of 10,000 instances of The Pile. Evaluating 16K context
where k' < k. Therefore, once the condition met,

---

Published as a conference paper at ICLR 2025
Algorithm 2 BACKWARD thread
b 2--
If init b desock X 1 I/ 40 dulock X dheud I/ dulock dhead a dolock
4O - load block of 0o
load block i of
load block M
for 1... do
Do
a - a Σ L
A + exp ZZ + cumsum, L + a)
VA < A )
Á dbsockk x dblack // V Z delock X dolock
VZ + VÁ (1 ex-e((())) @ (cumsum,, VA +
b b VA
VQ < VQ VZK"
A 04 - atomic add block VK
Il accumulate gradient for V and V
VZ Q + atomic add block lc VV
end for
VQ- store block i of VQ
lengths using the the LM evaluation harness (Gao et al., 2023) with batch size of 1, we find that early halting allows a 9.3% percent speed improvement.
5 EXPERIMENTS
In this section, we compare existing attention methods against stick-breaking. We first look at modification of a synthetic task from Arora et al. (2023) to understand the inductive biases of stick-breaking. We then compare the stick-breaking against existing length extrapolation methods on 350M model setting. We then pretrain IB parameter model, and evaluate it various NLP benchmarks, and evaluate for length extrapolation and retrieval capabilities on long context using the RULER benchmark (Hsieh et al., 2024). We also report benchmark results for 3B model have trained. For reference, the size of the models are detailed in Table 1.
5.1 MULTI-QUERY REPEATED ASSOCIATIVE RECALL TASK
To illustrate the inductive bias of stick- stick-breaking attention, first analyse its behaviour simple synthetic task. Arora et al. (2023) demonstrate that there is correlation between multi-query associative recall (MQAR) and the performance of language modelling. They show that while Transformers, which are based on attention can easily handle MQAR, the current linear state-space models cannot. While stick-breaking can also solve the MQAR task, we formulated different version of this toy task and tested it on both softmax attention and stick-breaking, As before, each instance of the task an initial assignment of values to variables. However, in the query sequence, the same variable can be queried multiple times, and each query is then followed by variable assignment, which successive queries must recall. We refer to this task as multi-query repeated associative recall (MQRAR). This setting can be analogous repeatedly updated, and it where it is useful to understand the current state the variable assignment. The following is an example sequence of the task:
some scenarios in programming where variables are
Input Output
4
E 6
X
E
2
B
B
4 6
3
6
6
1
We compare Transformer models with Softmax and RoPE against stick-breaking attention, comparing their ability to handle MQRAR with increasing key-value pairs: from 32 to 192, in increments of 32. The full sequence length is 768. Both models are 2-layer Transformers with 256 hidden dimension and one attention head. We sweep through 4 learning rates of [10-4, 10- 10 10- 10-2). and report the results of the best performing model. Softmax+RoPE is able to perform perform this task up to 128 key-value pairs, while stick-breaking is able to deal with sequences up 192 key-value pairs.

---

Published as a conference paper at ICLR 2025
Figure 4: Attention visualisation of the models trained the MQRAR task. The figure shows attention for each token for the 2-layer 100-dimension Transformer. Note that in the standard Softmax RoPE setting (above), the attention head is "distracted" at the third retrieval of 'E", attending to the first instance of rather than the mroe recent one. In the stick-breaking setting (below), cach attention head attends to the prior assignment of the variable.
Table I: Model hyperparameters and total size
Mlayer
dhidden
dinter
Thead
Total params.
350M
24
1024
2730
32
2048
367 367.526,912
40
1536
4096
24
4096
1,208,083,966
3B
40
2304
9216
36
4096
3,513,473,280
For more qualitative analysis, we trained another set of 2-layer Transformers of 100 dimensions on 16 key-value pairs, and visualised the patterns of the attention head. In Figure 4, we visuallised the above given example. Note here that the retrieval of the third 'E' is distracted by the earlier assignment
'E', while stick-breaking attention correctly retrieves the later assignment of This may be due to the limitations of RoPE with fewer layers and fewer head dimensions (here we use 32).
5.2 350M MODEL. LENGTH EXTRAPOLATION
We test the ability of stick-breaking for length generalisation, where we train on a fixed context length (L 1024) and test it on longer context lengths. We started with the LLaMa (Touvron al., 2023) architecture, and modified the attention module to use the various baselines we compare against:
3.8
9192)
FIRE
RoPE 8192, 1.0)
Negazive Log lilclthodd
3.40
Negarive Log-ikelloos
NoPE
4096, 1.0)
Stick-breaking
RoPE (f 1.01
3.35
3.5
We 512
3.4
3.3
W 2048
Stick-breaking
232
233
233
24
232
233
2
212
233
2
length
length
length
(a) Position bias embeddings
(b) Comparison with 8192
(c) Sliding window
Figure 5: Comparisons against different methods of sequence length extension. L represents the training context length,
the RoPE scaling factor used, and W is the window size in sliding window attention. We compare against different position embeddings and biases with L 8192, training with L = 8192, and various sliding window sizes with = 2048. Note that the scale on the y-axis are different in all three plots.

---

Published as a conference paper at ICLR 2025
ALiBi (Press al., 2021), FIRE (Li al., 2023), NoPE (Kazemnejad et al., 2024) and the default Llama position embedding RoPE (Su et al., 2021). 1). We trained on the first 15B tokens of SlimPajama (Soboleva et al., 2023), and evaluate it on the Wikitext benchmark in the LM evaluation harness (Gao et al., 2023), with context lengths of 2048 to 64K. As control, we also trained model with a context of I. 8192, and expected, the performance longer sequences better when RoPE scaling is used (See Figure 5b).
Intuitively, we should expect model that generalises well on longer sequences to have better likelihood as the context length increased longer contexts mean more information for predicting the next word. However, most methods do not generalise well longer contexts (See Figure 5a).
Softmax-
Stick breaking
Contrary to the results in Kazemnejad al. (2024), NoPE loss increases context lengths are extended. We also find that RoPE embeddings alone do not generalise, despite being relative position encoding Using RoPE scaling with the scaling factor f 1.0 does alleviate the issue, but still results in increase in the loss. Surrrsingly, while ALiBi is a fixed linear bias that increases with relative distance, it performs better than FIRE, which learns function of the bias for given relative distance. For position embeddings, our experiments suggest that RoPE with scaling works best for length extrapolation. ALiBi is an incremental bias for the logits, effectively down-weighting positions further away, eventually approaching O. This can be viewed as 'soft' windowed attention, and given the relatively good performance, we also trained models with windowed attention of W (512, 1024, 2048]. The results suggest that windowed attention helps with preventing the loss from spiking when the context length is extended. As expected, W 512 performs the worst. However, W 2048 increases slightly as the context length is increased, while W = 1024 decreases till I, = 213 and stays constant. The W = 1024 model trained on context lengths L = 2048, which would allow the model to learn to deal with context lengths longer than the window size, while W 2048 is trained similar to the standard non-window attention model. Figure 5c shows the generalisation curves. Note that the y- axis in the plot is smaller scale than that of Figure 5a, indicating that the sliding window method a relatively good method for length extrapolation as well. all cases, stick-breaking negative ligllike still decreases the context length increases, outperforming the other methods.
50
100
150
200
Figure 3: MQRAR performance in-
creasing key-value pairs.
5.3 LANGUAGE MODEL PRETRAINING
We pretrain the IB and 3B models in this section using two-stage training scheme in Hu et al. (2024) and the Power learning rate schedule (Shen et al., 2024). In the first stage, there warmup for the learning rate to 0.01, then we apply Power decay. Our training corpus has IT tokens and mixes large-scale open-source datasets medium quality with permissive licenses. In the second stage, we exponentially decay the learning rate to zero. The stage training corpus is mix of stage data
Table 2: Results on the various NLP benchmarks for the 1B and 3B pretrained model. 'Softmax'
benchmark the standard Softmax RoPE setting.
Task ARC-e
Hella. Accuracy (normalised)
ARC-e
OBQA
PIQA
RACE SciQ
Wino.
Wiki. Ppl.
Avg.
Accuracy
IB Parameter Models
Softmax 35,8
65.6
64.8
38,8
75.0
36.5
90.5
63.4
58.8
13.8
Stick-breaking
37.7
67.6
65.4
36.6
76.0
37.4
91.9
63.1
59.5
13.4
3B Parameter Models
Softmax 42.2
73.1
73.2
40.8
78.8
37.4
93.5
67.6
63.3
11.3
Stick breaking
44.9
74.3
74.1
40.4
79.7
37.8
93.9
68.0
64.1
10.8
Gemma2-25
80.2
72.9
41.8
37.3
95.8
68.8
65.8
13.1
Qwen1 5-4-B
39.6
61.5
71.4
40.0
77.0
38.2
90.0
68.1
60.7
12.5
8

---

Published as a conference paper at ICLR 2025
100
Softmax RoPE
Softmaxt RoPE
Softaak RoPE Stikkbreaking
Stick breaking
80
Stick-breaking
80
Nocurecy
60-
Accuracy
60
Aucureyy
60
40
40
0
16384 32768 131072
4006
131072
Sequence length
Sequence length
Sequence length
(a) Overall
(b) Needle in a Haystack (NIAH)
(c) Variable Tracking
Figure 6: Results the RULER benchmark. Without further finetuning for long context, evaluate the IB models on the RULER tasks. Subfigures: 6a shows the average over all 13 tasks, 6b is the average across the various NIAH tasks, 6c is the result for the variable tracking task.
and small amount of high-quality open-source and synthetic corpora with permissive licenses. The training batch size is 1024 and uses padding-free sequence packing for training in Dolomite Engine (Mishra, 2024). We evaluate the pretrained models on language model tasks and multiple-choice tasks from LM evaluation Harness (Gao et al., 2023). The multiple-choice tasks include: grade-school science questions (ARC; Clark et al. 2018), common sense reasoning (Hellaswag: Zellers al. 2019), open book question answering (OpenBookQA: Mihaylov et al. 2018), physical questions (PIQA; Bisk et 2020), reading comprehension (RACE; Lai et al. 2017). and Winograd schema task (Winogrande; Sakaguchi et al. 2021). Table shows the performance. Overall, we outperform our own pretrained standard attention models that are trained on the same settings. We perform better average, and attain better perplexity on Wikitext.
We also evaluated the models on MMLU with 0- shot and 5-shot settings. For the IB models, we have included the results for Laam (Zhang et al., 2024) Due to the inductive bias of stick-breaking we believed that stick-breaking would perform better on MMLU in a few-shot setting as it would be distracted by the few-shot examples provided in the context.
Table 3: MMLU few-shot results
MMLU
O-shot
5-shot
IB Parameter Model
Softmax
25.7
25.2
Stick Btickbreaking
28.4
29.3
Stick- boeck-brakkig performs surprisingly well even in the 0-shot setting, and improves with few-shot examples provided. Note that this may not always be the case, as in our Softmax + RoPE model, the performance decreases with few-shot examples in context.
TinyLlama
25.3
26.0
3B Parameter Model
Softmax
46.1
49.1
Stick- brek-breaking
50.8
52.9
We have included the Qwen .5-4B and Gemma2-2B models for comparison, and our 3B model underper- forms Gemma2-2B, smaller model, while it outper- forms Qwen1.5-4B.
Gemma2-2B
49.3
53.1
Qwen1.5-4BB
54.2
55.2
Table 3B Model GSM8K Results
Finally, we evaluate our 3B model on the GSM8K dataset (Cobbe al., 2021). Interestingly, the 5-shot setting underperforms standard attention while CoT with 8-shot sees improvement of 5.5%. The im- provements in MMLU and GSM8K on few-shot setting may suggest stick-breaking has better induc- tive bias for in-context learning, particularly in the formats that these evaluation benchmarks use.
GSM8K 8-shot, CoT
5-shot
Softmax Stick-breaking
44.1
44.2
42.3
49.7
5.3.1 IB LENGTH EXTRAPOLATION
We test the 1B brick-breaking and softmax model with the RULER benchmark (Hsich et al., 2024). The benchmark consists of -atle------hhstackk -like tasks, and is generally used for testing retrieval capabilities of long-context models that are trained specifically for long contexts. In our setting,

---

Published as a conference paper at ICLR 2025
RULER evaluate both 1B models trained on 4096 contexts. Accordingly, the general capabilities of these models on longer contexts are much worse than purpose-traines models.
On average, the performance of stick-breaking dominates Softmax RoPE with scaling (Figure 6a). In the breakdown, we find that the stickbbreking model is surprisingly good at extrapolating on NIAH tasks up to 16K context lengths, while the standard model significantly drops in performance. Our results on the variable tracking (VT) task agrees with our experiments on MQRAR. The task involves tracking the variable assignments provided in the context, and find that even in-distribution (L = 4096), the standard model does not perform well at this task.
CONCLUSION & FUTURE WORK
We propose formulation for using the stick-breaking process as replacement for softmax for attention. Stick- breaking attention allows us to do away with position embeddings, while still retaining model performance. We detail the specifics of implementing the stick-breaking kernel in Triton for large scale training. We then demonstrate that stick-breaking is good at length extrapolation, performing better than other position embedding and position bias methods in our 350M class models. We also show that our pretrained stick-breaking models peform better in controlled experiment, given the same training data and training regime. On retrieval in the RULER benchmark, stick-breaking outperforms softmax attention.
The drawbacks in efficiency can be improved in future work by making similar optimisations that Flash Attention (Dao al., 2022; Dao, 2023) made for speedups. These include making full use of features CUDA, and by cache retrieval optimisations. We believe there is plenty of room for improvement in computation efficiency that can be made in future versions of stick-breaking. We also think there can be weakk--rraiing-seeceee Stick-breaking Transformer, and we have initial results in Appendix C. Overall, we find stck-breekking attention to be promising replacement for Softmax + RoPE in Transformer models.
modifications to Transformers that can augment the
ACKNOWLEDGMENTS
We would like to thank Mayank Mishra and Gaoyuan Zhang for their help during the training of the IB and 3B models.
REFERENCES
M Allen and Lambie. An environmental residual allocation model. In Proceedings of the Conference Environmental Modeling and Simulation, April 19-22, 1976, Cincinnati, Ohio, pp. 236. US Environmental Protection Agency, 1976.
Zeyuan Allen-Zhu and Yuanzhi Li. Physies of language models: Part 1, context-free grammar. arXiv preprint ar Xrv: /30.:333,....33, 2023.
Simran Arora, Sabri Eyubogiu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv::3312 04927, 2023.
Andrea Banino. Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, pp. 7432-7439, 2020.
bloc97. NWK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) con- text size without any fine-tuning and minimal perplexity degradation., 2023. URL https://www.ne rope alpewa pope_aallosslllmma___a
/atea////ntkaare___e scaled_
have_
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. Advances Neural Information Processing Systems, 36, 2024.
10

---

Published as a conference paper at ICLR 2025
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1rXhv:1803.054433 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. arXiv preprint arXiv:2110.07732 2021.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0869) 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16353 2022.
Steinar Engen. A note on the geometric series species frequency model. Biometrika, 62(3): 697-699, 1975.
Warren John Ewens. Population genetics theory-the past and the future. In Mathematical and statistical developments of evolutionary theory, pp. 177-227. Springer, 1990.
Lco Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding. Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https://zenod:: wwtpmnddo.org//eceeo http///corr//11//// 10256663.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv: I00::06000888 2016.
RC Griffiths. Genealogical-tree probabilities in the infinitely-many model. Journal of mathemat- ical biology. 27:667-6800 680, 1989.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00775 2023.
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. arXiv preprint arXiv:2203.16634 2022.
Cheng-Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404. 06654, 2024.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395 2024.
Kazuki Irie, Albert Zeyer, Ralf Schlüter, and Hermann Ney. Language modeling with deep trans- formers. arXiv preprint arXiv:1905. 042222226....... 2019.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy, The impact of positional encoding on length generalization in transformers. Advances Neural Information Processing Systems, 36, 2024.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. arXiv preprint a7X4v:1704.04683 2017.
Shanda Li. Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang. Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. arXiv preprint arXiv:2310.044: 2023.

---

Published as a conference paper at ICLR 2025
John William McCloskey. A model for the distribution of individuals by species in environment. PhD Thesis, 1965.
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404008879 2024.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv: 1o0p.0278979 2018.
Mayank Mishra. Dolomite Engine: A URL https://gthuu..cm//bb//////////
Library for Pretraining and Finetuming, 2024.
engine.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Jayaram Sethuraman. A constructive definition of dirichlet priors. Statistica sinica, pp. 639-650, 1994.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv: 00051,,,,, 2018.
Yikang Shen, Zhouhan Lin, Chin-Wei Huang, and Aaron Courville. Neural language modeling by jointly learning syntax and lexicon. arXiv preprint arXiv:1711.000 2017.
Yikang Shen, Zheyu Zhang, Tianyou Cao, Shawn Tan, Zhenfang Chen, and Chuang Gan. Mod- uleformer: Learning modular large language models from uncurated data. arXiv preprint arXiv:2306.04640, 2023.
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D Cox, and Rameswar Panda. Power scheduler: batch size and token number agnostic learning rate scheduler. arXiv preprint arXiv:2408. 13359, 2024.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. deduplicated version of RedPajama. slimpajama-a-6277--00 ceaeo-ttooe--clenedeeee URL httphhts//hugggiggaaeeee
SlimPajama: A 627B token cleaned and https://wwe www.bras. net. net/blog/ 2023.
datasets/ /ataaetts/eeebrrra PSlia//aaaaaaai277777
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint 200.:2004400864,, 2021.
Shawn Tan and Khe Chai Sim. Towards implicit complexity control using variable-depth deep neural networks for automatic speech recognition. In 2016 IEÉE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5965-5969, IEEE, 2016.
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse universal transformer. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 169-179, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.002888 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
12

---

Published as a conference paper at ICLR 2025
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:339..14553, 2023.
2023.
Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05258, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv: /905.07805 2019.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.0238. 2024.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint 24X 2024.
13

---

Published as a conference paper at ICLR 2025
A DERIVATION OF LOG-SPACE FORMULATION
a (a) the sigmoid function,
exp(x)
(xx) + exp(=x) (-z)
1 +
1
- (x)) 1 exp(x)
Since Big o azz.g))
log Big Fijj
log (1 + exp(zij))
(9)
log (1 - -B..)) = log + exp(zi,)) = log (1 + cxp(zi.)))
(10)
We can substitute these back in,
A., = exp
log Bi.
log (1 - Bkns))
(11)
k=1+1
5-1
exp
Zi log (1 + exp(zi,))
Σ
log (1 + exp(zx.))
(12)
k=it1
exp
log (1 + exp(=x3))
(13)
k=
14

---

Published as a conference paper at ICLR 2025
B 1B MODEL RULER RESULTS
fwe
multikey
100
100
Softmax RoPE
Softmax-RoPE
Softmax-RoPP RoPE
80
Stick-breaking
80
Stick-breaking
80
breakng
Accursy
60
60-
20
20
0
4096
8192
16384 32768 65536 131072
4096
8192
16384 32768 65536 131072
4096
8192
16384 32768 65536 131072 niah multiquery
niah multikey.
niah multikey.
2.00
100
Softmax+ RoPE
Suftmax-RoPL Stck-bresking
Softmaxx+Rot RoPE Stick bieekbrge
80
Stick-breeking
80
80-
Accuracy
60
60-
60-
40-
20
0
4096
6192.
16384 32769 65536 131072 niah multivalue
4096
0192
16384 32766 65536 131072
4096
8192
16384 32768, 65536 131072 niah single
niah single
100
100
Softmax+ RoPE
RufE stickbrrekking
Softmax RoPE Stick btickbreaking
80
Stick breaking
80
80
Accuracy
60
60-
60
40
40
0-
4C95
8192
16384, 32768 65336 131072
4096
8192
16384
32768
65526 131072
4096
8192
16384
32768 65336 131072 qa_2
niah single_3
Qa_1
Softmax+ RoPE
Softmax-RRPE
Softmax RoPE brek-ng
80
brick-breaking
btick-breaking
80
Accurscy
60
60
60
40
40
4096
8192
16384
65536 131072 VC
32768
4096 8192 16384
32768
65536 131072
4096
8192 16384 32768 65536 131072
Softmax RoPE Stick breaking
Accurnry
60
40
1096
8192 16384 32768 65536 131072
Figure 7: RULER length extrapolation results
15

---

Published as a conference paper at ICLR 2025
Table In these experiments, we use Player 40, uses Grouped Query Attention with 7head 12 and 4 key and value heads, and dhead = 128. For the 6.5B model, we use Membd 1536, with 64 experts and k 8, with each expert having a hidden size 512. For the 28.6B model, we use Nlemdd 4096, with 72 experts and k 12, with each expert having hidden size of 768.
RB
GN
ARC
Hella.
OBQA
PIQA
Wino. Accuracy
MMLU
GSMSK 8-shot, cot
MATH 4-shot
Avg.
Accuracy (normalised)
1.2B parumeters, 1.257 tokens
Softmax
55.8
66.7
38.4
76.0
63.9
28.3
27.0
7.6
45.5
SB
X
X
54.0
66.9
39.6
76.8
65.5
37.3
29.6
8.2
47.2
X
54.1
66.6
39.8
76.2
66.9
39.6
29.2
7.8
47.5
V
56.4
66.7
38.4
76.7
64.8
39.8
29.3
8.2
47.5
1B activation / 6,5B parameters, 2T tokens
Softmax
66.3
76.9
44.6
80.7
72.5
55.5
50.8
21.6
58.6
SB
/
68.7
77.9
46.6
80.5
74.4
56.9
55.3
22.5
60.4
6B activation / 28.6B parumeters, 3T tokens
Softmax
76.8
84.4
51.2
84.1
80.1
67.7
70.4
34.5
68.6
SB
77.2
84.7
51.4
83.7
81.5
71.2
71.9
35.6
69.6
C STICK-BREAKING SPECIFIC TRANSFORMER MODIFICATIONS
While we demonstrate the effectiveness of stick- breaking as a drop-in in replacement for softmax attention the main paper, found several minor modifications the standard Transformer architecture can help enhance the performance of the Stick-breaking Transformer various tasks.
EE A., 1, experimented with using remaining attention weight is assigned to an embedding r, by modifying the output of the attention layer as follows:
Remainder Bias (RB) As
0j ArJ U. (1 Σ A 1
(14)
Each head has its corresponding embedding, resulting in a matrix of additional parameters of Whead dhead. We refer this as remainder bias. This is similar to an attention sink token (Xiao et al., 2023). The authors showed that having an attention sink token than the standard attention, and saw better performance when learnable attention sink embedding was used. In our experiments, we saw improvements when the remainder bias was used.
resulted in poorer performance
Group Norm (GN) Due the unnormalised weight of stick-breaking attention, the resulting o, can have varying norm magnitudes. We experimented with applying a head-wise norm, which was implemented with Group Norm. This also done with Differential Transformers (Ye et al., 2024), for similar reasons. We found improvements when Group Norm was used in conjunction with RB.
In Table 5, we perform experiments with RB and GN on 1B models where we found improvements on IB dense models. We then trained a 6.5B and 28.6 parameter MoE model with both GN and RB, and show that they outperform corresponding Transformer with softmax attention.
C.1 CONTEXT WINDOW EXTENSION
All models are pre-trained with context window of 4096. To extend their context lengths, continue training for 6250 steps at context length of 16k. Our learning rate is scheduled to increase from 10 to 0.000125 in 150 steps, and then decays exponentially to O until 6250 steps. The effective batch size has 4M tokens.
For Softmax
set the dynamic RoPE scaling factor to 4. In Figure, we show that by applying rope scaling, the model can somewhat extrapolate to 16k context lengths, but the log- log-liieldoood spikes extend further.
Applying this method context window extension to both models, we find that the Softmax+RoPE model improves slightly, while still showing issues with longer context lengths, albeit mitigated
16

---

Published as a conference paper at ICLR 2025
slightly. Stick-breaking however exhibits lower NLL extend the context window further to 32k and 64k.
The behaviour of RoPE out-of-the-bow for length extrapolation is also demonstrated in Section 5.2. also show that training at longer sequences does not mitigate the out-of-distribution performance of Softmax+RoPE: models. Here, we show that while naive continued fine-tuning helps, the im- provement on long-contex: windows is not as substantial as with stick-breaking. This suggests better context window extension behaviour and further length extrapolation for the stick-breaking attention mechanism.
4.75
Stick breckbbeakiin extrapolation
Stick breaking
4.50
Softmax extrapolation (RoPE scale
Negative Log-llkelhhodd
Softmax extension (RoPE scale 4)
4.25
4.00
3.75
2
2
Context length
Figure 8: Stick-breaking and Softmax+RoPE. context window extrapolation performance on longer context lengths in the Pile 10k evaluation dataset.
17


Reviews:
[{"review_id": "Ouqb5unst2", "reviewer": "Reviewer_wgvD", "summary": "This paper does a thorough evaluation of stick-breaking attention (also known as geometric attention) for synthetic test tasks and natural language tasks, showing its effectiveness at capturing a preference for locality, as is natural and correct for natural language tasks. The paper examines a fast, numerically stable implementation of stick-breaking attention in Triton, which enables larger scale experiments.", "questions": "- The differences between many models in Table 2 are fairly small and nothing is said about the detailed validation of the results. Are these from single runs rather than averages from 3-5 runs with different random initialization? How much variance would there be here, how confident can we be that a result of 63.4 is better than 63.1 for Winogrande on Softmax vs. SB w/o remainder correction, for example?\n- I took the text around lines 101-107 as suggesting that things should work better doing stick-breaking with remainder, but the actual results go the other way. It would be good to provide more understanding of why this harms rather than helping.\n- For the paper, you should explain lines 300-301. This isn't a question. I figured it out, but the text of the paper should explain MQRAR better for people who haven't seen it.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- The paper provides thorough and useful experimental results on the performance of stick-breaking attention, providing good exploration of its effectiveness for artificial and natural language tasks, for length generalization, etc. The experimentation seem well thought through and well done.\n- The paper is generally clear and easy to read.\n- The paper is very honest about what it contributes and what it uses from prior work.\n- The paper provides useful and new empirical results on different forms of attention on different tasks.\n- The paper examines building an efficient, numerically stable Triton implementation of stick-breaking attention.\n- It's good to include comparisons to Gemma2-2B and Qwen1.5-4B so that people can easily see that the results are decent, even though the papers own results are the apples-to-apples comparisons.", "weaknesses": "- The paper lacks originality in machine learning ideas. Stick-breaking attention has been previously explored by Yikang Shen (in multiple papers) and especially by Csordas et al. (2021), the latter under the name \"Geometric attention\". \"Stick-breaking attention\" is a better name for the model used, but the model is exactly the same as in these prior works, limiting the originality of this paper. The value is mainly in the more extensive experimentation, including showing performance on larger scale, standard natural language benchmarks.\n- This paper lacks somewhat in significance because of this. It does have some significance, since it is really good to see that these ideas really do give gains on standard NL tasks like ARC, Hellaswag or RACE, but the basic correctness of the idea had already been established.\n- The differences between many models in Table 2 are fairly small and nothing is said about the detailed validation of the results. Are these from single runs rather than averages from 3-5 runs with different random initialization? How much variance would there be here, how confident can we be that a result of 63.4 is better than 63.1 for Winogrande on Softmax vs. SB w/o remainder correction, for example?\n- How to produce an efficient, numerically stable Triton kernel basically follows the methods of FlashAttention and standard good practice (using log(1 + exp(x)), etc.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ClHKfSc126", "reviewer": "Reviewer_8GF8", "summary": "The paper introduce an alternative attention that includes a form of positional embeddings, avoiding the need to add PE techniques such as RoPE. The method is based on the stick-breaking process, which means for each token in the context (keys), they assign a break point that represents the remaining stick to the current token. This means that if 2 tokens have equal logits, the one closer to the current token will receive more stick value, representing the positional order.\nConducting this process through out the entire context results in a sequence of attention weights in lieu of traditional attention weights.\nThe paper suggests this process incorporate recency bias in nature without PE.\nExperiments done show that the new method is better at length generalization in perplexity compared to attention+RoPE.", "questions": "- What is $\\sigma$ function? It is never defined. is it sigmoid ?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a novel addition to the zoo of \"attention alternatives\" by leveraging the stick breaking process, which performs both \"attention\" and positional embeddings intrinsically. The mechanism has a recency bias, meaning a token can prefer to allocate all its \"energy\" to few recent tokens, but it can also skip over and only attend to far-away tokens.\n- The paper also include details implementation in Triton for flash-attention style efficiency and speed-up optimization, which is hugely appreciated, especially when nowadays efficiency and scalability is valued more for massive training of LLMs.\n- The paper conducted diverse range of experiments, from throughput (only 20% slower than flash attention), perplexity and language modeling tasks (MMLU, ARC-c, hellaswag...), which all shows promising results.\n- The main driver of this work is to solve the length generalization challenge of LLMs, which show good results.", "weaknesses": "- Method can be explained more clearly with diagrams, formulation should be defined more thoroughly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SULyHU7fYi", "reviewer": "Reviewer_MyhJ", "summary": "This paper is about stick-breaking or geometric attention, an alternative to softmax that has a built-in bias towards more recent tokens. Given a query position j, each key/value position i < j computes a probability of \"yes\" or \"no\", and the attention weight on i is the probability that position j is \"yes\" but all intervening positions are \"no\".\n\nStick-breaking/geometric attention was introduced in previous work. It makes position embeddings unnecessary and can be computed in O(log n) parallel time. This paper presents an improved implementation, and shows that stick-breaking/geometric attention improves performance on the multi-query associative recall task and various benchmarks.", "questions": "What is the exact relationship of stick-breaking attention to geometric attention (Csordas et al 2021) and the stick-breaking attention of Shen et al (2023)?\n\nEq (1):\n- I'd suggest not using \\cdot, as it might be misinterpreted as an inner product in this context.\n- The first summation only goes up to i-1, so a position cannot attend to itself, which is different from usual future masking. I didn't see this decision discussed; what's the reason for it?\n\nLine 94: typo, should be $z_{i,j} = z_{i',j}$? \n\nEq (4): I don't know if it matters, but you could use log1p(exp(z)) instead of log(1+exp(z)).\n\nIt may be interesting to node that while softmax with a temperature factor (α = softmax (z/T)) as T -> 0 approaches average hard attention, which is used in many theoretical studies of transformers (e.g., https://arxiv.org/abs/1901.03429, https://arxiv.org/abs/2106.16213), stick-breaking attention (β = σ(z/T)) as T -> 0 approaches strictly future-masked rightmost hard attention, which is the kind of attention studied by https://arxiv.org/abs/2310.13897.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This softmax alternative is simple and induces a bias towards attending to recent positions in a very natural way.\n\nThe experimental results all look strong.", "weaknesses": "As far as I can tell, stick-breaking attention is exactly the same as geometric attention (Csordas et al 2021), and stick-breaking was previously introduced by Shen et al (2023). Both papers are cited in the introduction, and the introduction concludes with an accurate list of the novel contributions of the paper. However, \n- The paper's short title \"Stick-Breaking Attention\" may give the impression that this is the first paper about stick-breaking attention.\n- The abstract does not mention previous work; on the contrary, it says \"We propose an alternative attention mechanism.\"\n- The introduction, probably inadvertently, could be mis-read as saying that geometric attention only has one parameter (\"Geometric attention, named after the Geometric distribution, which only has one parameter\").\n\nEq. (2): Putting the remainder of the attention onto position j itself does not seem like the right choice. Probability $(1-\\beta_{i,j})$ is the probability of pushing the attention to the left of position i, so $\\prod_i (1-\\beta_{i,j})$ is the probability of pushing the attention all the way to the left. So it's not surprising that this turned out not to work well. Letting the attention weights sum to less than one (or equivalently, putting the remainder of the attention onto the zero vector) seems like the most sensible thing to do.\n\nThe Flash Attention-like implementation of stick-breaking attention is 20% slower than Flash Attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yCvlLobBsh", "reviewer": "Reviewer_q9HG", "summary": "This paper presents a new type of attention mechanism called \"stick-breaking attention\". It is meant to have a bias for attending to recent positions. Instead of using dot-products between keys and queries as logits and renormalizing them using a softmax to form the attention weights, the dot-products are each passed through the logistic function and then used as the probabilities in a stick-breaking process. In other words, the model first decides how much to attend to the current token, *then* decides how much of the remainder to allocate to attending to the previous token, and so on. It is a generalization of geometric attention proposed in prior work, where the probabilities at each timestep can be different. The authors discuss details for implementing it efficiently. They test it against standard scaled dot-product attention on a simple synthetic task that advantages recency bias, and on a variety of natural language benchmarks. Stick-breaking attention gets better scores on most tasks.", "questions": "1. The keys and queries do not have positional encodings, but the values still do, right? Does the lack of PEs on queries and keys reduce the model's expressivity (i.e., are there certain functions that it can no longer implement because of the loss of PEs)? 474: Stick breaking attention doesn't completely get rid of PEs, right?\n1. Although I intuitively understand the motivation for using stick-breaking, can you make a more rigorous case as to *why* stick-breaking has a recency bias? Is it not true that both softmax and stick-breaking attention can implement arbitrary attention patterns?\n1. Related to the above, did you verify that the model actually attends to more recent tokens? Does this happen in wikitext?\n1. 096: I'm not sure what this discussion is trying to say. Can you explain this more?\n1. 094: I think $z_{i,j} = z_{i,j}$ is a typo.\n1. Instead of using Eq 5, which is an inexact approximation, why not use the identity $\\log(1 + \\exp x) = c + \\log(\\exp(-c) + \\exp(x - c))$, where $c = \\max(0, x)$? This would solve the overflow problem while being an exact solution.\n1. 208: Doesn't this make your method much less parallelizable than standard attention? I think you can actually get the same parallel time complexity using a variant of parallel prefix sum -- have you considered using that?\n1. 304: How much of each sequence is devoted to the initial set of pairs, and repeated pairs?\n1. Table 2, Figure 6: Why do you suppose stick-breaking attention helps on these tasks? It even gets slightly lower perplexity on wikitext, which is not a task that is specifically advantageous for recency-biased models. Can you spend more time discussing what kind of capability each natural language benchmark represents? Why do we expect recency bias to help with them? Do the gains come from recency bias, or is it more about length generalization thanks to the use of relative PEs?\n1. Why does stick-breaking attention do better at length generalization on RULER if the retrieved data is not necessarily recent? I don't see why stick-breaking would be advantageous here.\n1. Why does stick-breaking attention eventually fail on RULER at longer lengths? What is the failure mode?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experiments include a variety of natural language benchmarks.\n1. The authors include a thorough comparison to similar prior work.\n1. The results as presented show that stick-breaking attention outperforms standard attention on a variety of benchmarks.\n1. Their new attention mechanism is reasonably fast, and further optimizations are possible.", "weaknesses": "1. ~~Their implementation relies on a very rough approximation of the exponential function to avoid overflow (Eq 5), so their method is implementing a function that is quite different from what is proposed. I think this can be avoided (see Questions), and I am curious to see if using an exact solution affects the results.~~\n1. To improve the readability of the paper, a more intuitive explanation of stick-breaking would be useful, for those who are not familiar with the term. For starters, why is it called stick-breaking?\n1. Although the results are encouraging in Section 5.3, it is not clear to me that there was much hyperparameter tuning on the baselines. I would like to see more discussion of this.\n1. I would like to see more discussion of *why* recency bias helps on the natural language benchmarks tested (see Questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)