{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/stick-breaking-attention","title":"Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study","arxiv_id":"2410.17980","date":"2024-10-23","proceeding":null,"authors":["Shawn Tan","Songlin Yang","Aaron Courville","Rameswar Panda","Yikang Shen"],"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. 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.","url_abs":"https://arxiv.org/abs/2410.17980v2","url_pdf":"https://arxiv.org/pdf/2410.17980v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"stick-breaking-attention","repo_url":"https://github.com/ibm/dolomite-engine","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"stick-breaking-attention","repo_url":"https://github.com/shawntan/stickbreaking-attention","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[],"methods":[{"method_slug":"attention","method_name":"Attention"},{"method_slug":"softmax","method_name":"Softmax"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2410.17980","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2410.17980"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/shawntan/stickbreaking-attention","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/ibm/dolomite-engine","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"unverified":15},"by_repo_kind":{"official":{"samples":15,"ran":0,"repositories":2}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"8593e6db982f0994","entry":"ceil_divide","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/math.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/math.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8593e6db982f0994"}},{"code_sha256_prefix":"1a14ef87df25882f","entry":"check_power_of_2","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/math.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/math.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"1a14ef87df25882f"}},{"code_sha256_prefix":"be3f8284e0ad756c","entry":"clip_gradients","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/torch_utils.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/torch_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"be3f8284e0ad756c"}},{"code_sha256_prefix":"34aca61c3a883435","entry":"cpp_jit","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/jit.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/jit.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"34aca61c3a883435"}},{"code_sha256_prefix":"7d46751a4793398b","entry":"ctx_needs_gradients","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/custom_op.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/custom_op.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"7d46751a4793398b"}},{"code_sha256_prefix":"284647bfc008172c","entry":"custom_op","repo":"shawntan/stickbreaking-attention","repo_kind":"official","path":"stickbreaking_attention/utils.py","file_url":"https://github.com/shawntan/stickbreaking-attention/blob/HEAD/stickbreaking_attention/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"284647bfc008172c"}},{"code_sha256_prefix":"c878e222f31d20d2","entry":"get_counter_value","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/counters.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/counters.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c878e222f31d20d2"}},{"code_sha256_prefix":"b0541e2dcf262896","entry":"get_powers_of_2","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/math.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/math.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"b0541e2dcf262896"}},{"code_sha256_prefix":"90f1af006a6709d6","entry":"is_aux_loss_zero","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/training/loss.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/training/loss.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"90f1af006a6709d6"}},{"code_sha256_prefix":"a604b508bd6fa282","entry":"pad","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/tokenizers/utils.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/tokenizers/utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a604b508bd6fa282"}},{"code_sha256_prefix":"a44644f7f170ef4a","entry":"partialize_and_update_signature","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/inductor.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/inductor.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a44644f7f170ef4a"}},{"code_sha256_prefix":"dd3979e6d750a272","entry":"sigmoid","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/torch_utils.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/torch_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"dd3979e6d750a272"}},{"code_sha256_prefix":"1019f820dfe99006","entry":"stickbreaking","repo":"shawntan/stickbreaking-attention","repo_kind":"official","path":"stickbreaking_attention/sb_ref.py","file_url":"https://github.com/shawntan/stickbreaking-attention/blob/HEAD/stickbreaking_attention/sb_ref.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"1019f820dfe99006"}},{"code_sha256_prefix":"821f34b1af2cadb2","entry":"tanh","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/torch_utils.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/torch_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"821f34b1af2cadb2"}},{"code_sha256_prefix":"c13c8953208d35dd","entry":"xma_op","repo":"ibm/dolomite-engine","repo_kind":"official","path":"lm_engine/kernels/custom_op.py","file_url":"https://github.com/ibm/dolomite-engine/blob/HEAD/lm_engine/kernels/custom_op.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c13c8953208d35dd"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}