Block Sparse Attention With Block Retrieval
Block Sparse Attention WIth Block Retrieval!
docs: jacobfv.github.io/bsbr repo: github.com/JacobFV/bsbr
Why block sparse? The standard transformer attention mechanism computes attention scores between all pairs of tokens in a sequence. which leads to O(n²) complexity in both computation and memory
BSBR addresses this scalability issue by breaking the sequence into chucks of block size B and then combining standard attention within fixed-size chunks and a meta-attention mechanism to efficiently retrieve information between chunks. It works by compressing each chunk’s…
The chunk size (B) is a crucial hyperparameter that affects: - Memory Usage: Larger chunks use more memory but provide better local context - Computation Time: Smaller chunks are faster but may miss important long-range dependencies - Model Expressivity: Chunk size affects how…
You can control the compression factor which determines how much information is preserved in chunk states. Higher compression factors reduce memory usage and speed up computation but may lose fine-grained information:
BSBR also supports overlap between chunks which may help mitigate attention discontinuities at the block boundaries:
And besides sparse attention, the implementation also supports state reuse across layers, optimized attention masks for autoregression, and I plan on implementing more streaming-first optimizations when I get time. Imagine building agentic software that weaves threads of tokens…
Finally, how useful is a new architecture without pretrained models? Well sorry I don’t have any yet lol but the bsbr_transformers provides tools to convert pretrained huggingface transformers into bsbr ones. Just make sure to pip install bsbr[transformers] and then from…
Finishing off with an end to end example. I hope you find bsbr useful! Please also give OG’s poast a read:
Dense attention is a beautiful default and a brutal scaling law. Standard transformer attention compares every token with every other token, which gives the familiar memory and compute shape. That is tolerable until context stops being a prompt and starts becoming a working memory.
Block Sparse Attention With Block Retrieval, or BSBR, is one attempt to make that memory more structured. The sequence is divided into blocks. Attention remains dense inside a local block, where nearby tokens usually need fine-grained access. Across blocks, the model retrieves compressed block states rather than attending naively to every token.
The design pressure is simple:
- Local attention should preserve short-range precision.
- Block retrieval should preserve long-range access.
- Compression should make old context cheap enough to keep around.
- State reuse should make repeated long-context computation less wasteful.
That gives several knobs that are more operational than theoretical:
- Block size controls the local-context versus memory tradeoff.
- Compression factor controls how much information survives into block state.
- Overlap reduces discontinuities at block boundaries.
- State reuse lets layers and decoding steps avoid recomputing context that has already been summarized.
The interesting product direction is not just "longer context." It is streaming-first agent software. Agents do not merely answer once. They accumulate traces: files, tool calls, decisions, failures, user corrections, world state, and partial plans. A useful memory substrate has to weave those traces into something reusable without making every future step pay the full quadratic cost of the past.
This connects to my older interest in structured sparsity in the brain model and the more speculative architectural notes in Design Patterns for AI. The recurring intuition is that topology should be part of the interface. A model should not only learn weights; it should expose useful ways to route, compress, reuse, and inspect information.
The missing work is the part that always matters: pretrained checkpoints, benchmarks, ablations, and uncomfortable comparisons against simpler baselines. A sparse attention pattern becomes real only when it earns its complexity.