Source-linked AI summary
BCMT: Blockwise Causal Memory Transformer
Rachid Arezki
TL;DR
Dense self-attention becomes costly for long contexts because its complexity grows quadratically with sequence length. BCMT separates local attention from global context propagation using adaptive block summaries and exponential causal memory, achieving comparable validation performance to Dense Transformers while improving throughput and reducing memory consumption.
Problem
Dense causal self-attention has quadratic sequence-length complexity, limiting its applicability to long-context language modeling.
Method
BCMT applies dense causal self-attention within local blocks and propagates global context through exponential causal memory built from adaptive block summaries.
Results
BCMT achieves validation performance comparable to a Dense Transformer while significantly improving training throughput and reducing GPU memory consumption on WikiText.
Takeaways & Limitations
Adaptive block-summary memory is a viable alternative to dense global attention for long-context language modeling.
Takeaways & Limitations
BCMT cannot exactly retrieve distant information, so tasks requiring precise localization or faithful reconstruction may benefit from explicit global attention or specialized memory systems.
Abstract
from arXiv · showhide
Transformer architectures rely on dense self-attention to model long-range dependencies, but this mechanism exhibits quadratic complexity with respect to sequence length. We introduce BCMT (Blockwise Causal Memory Transformer), an architecture for long-context language modeling that decouples local token interactions from global context propagation. Dense causal self-attention is applied independently within local blocks, while each block produces an adaptive summary aggregated through an exponential causal memory. This memory is subsequently injected back into the token representations, enabling efficient propagation of long-range contextual information without relying on explicit global attention. Unlike standard Transformers and recurrent memory architectures, BCMT maintains neither dense interactions between distant tokens nor learned memory states. Its memory mechanism is fully parallelizable and remains compatible with standard implementations of dense self-attention. Experiments on language modeling with context lengths of up to 1024 tokens show that BCMT achieves validation performance comparable to that of Dense Transformers while significantly improving training throughput and reducing memory consumption. An ablation study further confirms that these improvements arise from the proposed memory mechanism. These results demonstrate that an exponential causal memory constructed from block summaries provides an effective alternative to dense global attention mechanisms for long-context language modeling.
1 Introduction
BCMT separates local token modeling from global context propagation by combining block-local causal self-attention with an exponential causal memory over adaptive block summaries. On contexts up to 1024 tokens, it matches Dense Transformer validation performance while improving throughput and memory efficiency, with ablations attributing these gains to the memory mechanism.
- Causal self-attention has quadratic complexity with sequence length, limiting its applicability as context grows.
- BCMT applies dense causal self-attention independently within local blocks and propagates long-range dependencies through an exponential causal memory built from adaptive block summaries.
- The memory is injected into token representations through gating, without explicit attention between tokens in different blocks, learned hidden states, or persistent memory tokens.Its exponential causal aggregation is fully parallelizable.
- Up to 1024 tokens, BCMT achieves validation performance comparable to a Dense Transformer while significantly improving training throughput and reducing memory consumption.
- Ablation results show that the improvements arise from the proposed memory mechanism rather than sequence partitioning alone.
2 Related Work
Prior approaches improve long-context scalability through sparse or approximate attention, structured attention, local block processing, compression, or recurrent and hierarchical memory. BCMT instead confines dense self-attention to local blocks and propagates context through adaptive block summaries aggregated by an exponential causal memory.
- Attention alternatives: Dense self-attention is expressive but has quadratic sequence-length complexity, motivating sparse, approximate, and structured alternatives for long-context modeling.Longformer and BigBird use sparse connections; Linformer and Performer modify attention; LongNet uses dilated attention, while retaining explicit distant-token connections.
- BCMT’s context propagation: BCMT preserves local dense attention while propagating long-range dependencies through adaptive block summaries aggregated by an exponential causal memory.The summaries form an auxiliary memory channel rather than replacing tokens or reducing the sequence length processed by later layers.
- Local block processing: Local block architectures reduce attention cost and preserve dense intra-block processing, but lack explicit communication between blocks unless block sizes or additional mechanisms increase.BlockBERT independently processes contiguous blocks with dense self-attention, while its main limitation is the absence of cross-block communication.
- Hierarchical representations: Unlike hierarchical architectures, BCMT uses a single level of adaptive block summaries rather than multiple progressively coarser representation levels.Its compact summaries are aggregated through exponential causal memory to progressively propagate context along the sequence.
- Memory-based architectures: Unlike recurrent memory architectures, BCMT uses neither recurrent hidden states, persistent memory tokens, nor learned recurrent dynamics.It combines dense causal self-attention within independent local blocks with exponential causal aggregation of adaptive summaries.
- Central architectural principle: BCMT separates local token interactions from global context propagation, providing a fully parallelizable alternative compatible with standard dense self-attention implementations.Long-range dependencies emerge through compact contextual representations propagated across blocks rather than global attention or evolving learned memory states.
3 Method … 3.3 Local Multi-Head Causal Attention (H)
BCMT replaces global causal attention with local blockwise self-attention and an inter-block exponential causal memory built from adaptive summaries. The method partitions sequences into contiguous blocks, computes masked multi-head local representations, and uses them to construct global context.
- 3 Method: BCMT models nearby-token interactions with dense causal self-attention within local blocks while propagating long-range dependencies through exponential causal memory.These mechanisms preserve strictly local attention while providing global context across blocks.
- 3.1 Problem Formulation: For autoregressive language modeling, BCMT models p(x_t | x_<t) while replacing global attention with local attention and inter-block memory.The replacement reduces self-attention cost for long sequences while preserving autoregressive causality.
- 3.2 Block Decomposition: The input sequence is partitioned into N_b contiguous, non-overlapping blocks with maximum size L.This block decomposition defines the units over which local attention and memory operations are applied.
- 3.2 Block Decomposition: Each block X^(i) ∈ R^(L_i×D) contains L_i ≤ L tokens, with only the final block allowed to be shorter.All blocks have size L unless T is not a multiple of L.
- 3.2 Block Decomposition: Token interactions are restricted to individual blocks through dense causal self-attention, while cross-block dependencies are handled by memory.This separation is the fundamental organizational principle of BCMT.
- 3.3 Local Multi-Head Causal Attention (H): Each block is processed independently using a standard multi-head causal self-attention layer.For each head, query, key, and value projections operate in the representation space with D = h d_h.
- 3.3 Local Multi-Head Causal Attention (H): Head outputs are combined with a standard causal mask, concatenated, projected back into representation space, and concatenated across blocks to form H.H is the local representation used to construct adaptive block summaries for subsequent inter-block memory computation.
3.4 Adaptive Block Summaries (S)
BCMT summarizes each local block into a single contextual vector using learned, content-dependent token weighting rather than uniform averaging. These compact summaries are then passed to the exponential causal memory mechanism.
- Adaptive Block Summaries (S): Each block is summarized into a single contextual vector using adaptive learned weighting over its token representations.This replaces uniform averaging with content-dependent aggregation from the local representations H(i).
- Adaptive Block Summaries (S): Token scores are obtained through a linear projection and normalized with softmax to form the adaptive weighting.The normalized scores determine how token representations contribute to each block summary.
- Adaptive Block Summaries (S): The resulting vectors S(i) provide compact representations of blocks and serve as inputs to the exponential causal memory.This connects adaptive block summarization to the subsequent memory mechanism.
3.5 Exponential Causal Memory
BCMT constructs inter-block memory by applying normalized causal exponential aggregation to block summaries, emphasizing recent summaries while retaining diminishing contributions from earlier blocks. The projected memory is shifted by one block to preserve strict autoregressive causality, so each block receives information only from preceding blocks.
- Memory construction: Block summaries are aggregated causally into an inter-block memory vector that propagates context throughout the sequence.The aggregation uses summaries observed up to the current block.
- Memory construction: The exponential decay factor weights recent summaries more heavily while preserving progressively decreasing contributions from the entire history.This distinguishes the aggregation from a uniform average.
- Implementation: α = 0.8 is used throughout all experiments, while preliminary tests found similar performance for α values between 0.7 and 0.9.The decay factor is fixed rather than learned.
- Causal injection: The projected memory is shifted by one block to preserve strict autoregressive causality.This prevents the current block from accessing its own summary.
- Causal injection: Consequently, each block receives information exclusively from preceding blocks and never accesses its own summary.The memory representations are projected into the local representation space before the causal shift.
3.6 Autoregressive Causality · 3.7 Memory Broadcast to Tokens
BCMT preserves strict autoregressive causality through causal within-block attention, summary-based inter-block memory, and a one-block memory shift. Each block’s memory is broadcast uniformly to its tokens before gated fusion with local representations.
- 3.6 Autoregressive Causality: Within each block, a lower-triangular causal mask restricts every token’s attention to itself and preceding tokens.Therefore, the local attention pathway cannot access future positions.
- 3.6 Autoregressive Causality: Each adaptive summary S(i) uses only block i’s local representations, while memory aggregates summaries from blocks observed through the current block.The inter-block memory is constructed causally from block summaries.
- 3.6 Autoregressive Causality: A one-block causal shift ensures block i receives contextual information only from preceding blocks, excluding its own and future summaries.The shift occurs before memory injection into local representations.
- 3.6 Autoregressive Causality: Global causality combines strictly causal local attention, observed-block summary memory, and causal memory shifting.Together, these properties eliminate computational paths from future tokens to current representations.
- 3.6 Autoregressive Causality: BCMT satisfies autoregressive causality: no future token information can influence the current token representation.This meets the causality constraints required for autoregressive language modeling.
- 3.7 Memory Broadcast to Tokens: Each block’s memory representation is broadcast to all tokens belonging to that block.The memory vector is replicated across the L positions to form a block-level matrix.
- 3.7 Memory Broadcast to Tokens: The per-block broadcast matrices are concatenated along the sequence dimension to produce the token-level memory tensor.This constructs the memory representation used across the sequence.
- 3.7 Memory Broadcast to Tokens: Every token shares its block’s contextual representation, which is derived exclusively from preceding blocks before gated fusion with local representations.The resulting tensor is combined with local representations through the stated gated fusion mechanism.
3.8 Gated Memory Injection · 3.9 Conceptual Interpretation
BCMT injects exponential causal memory into local token representations through token- and dimension-specific gates, while separating local dependency modeling from long-range context propagation. Its architecture assigns fine-grained interactions to blockwise self-attention and progressively aggregated global information to a complementary memory pathway.
- 3.8 Gated Memory Injection: The exponential causal memory complements local self-attention representations, with its contribution dynamically modulated by a learned gating mechanism.
- 3.8 Gated Memory Injection: Gate coefficients are computed independently for each token and representation dimension, controlling how much contextual memory enters each local representation.
- 3.8 Gated Memory Injection: When gate coefficients approach zero, local representations are largely preserved; when they approach one, exponential causal memory contributes more strongly.
- 3.8 Gated Memory Injection: The gating mechanism adaptively regulates memory use according to context while preserving local information whenever it is sufficient.
- 3.9 Conceptual Interpretation: BCMT explicitly separates local dependency modeling from long-range dependency modeling.
- 3.9 Conceptual Interpretation: Local causal self-attention models neighboring-token interactions within each block, while adaptive block-summary memory propagates more distant dependencies.
- 3.9 Conceptual Interpretation: This separation gives self-attention a fine-grained role and assigns the exponential causal memory a more global, progressively aggregated representation of sequence history.
- 3.9 Conceptual Interpretation: Unlike global-attention Transformers, BCMT distributes dependency modeling across complementary mechanisms operating at different temporal resolutions.
3.10 Computational Complexity
BCMT decomposes computation into local attention, block summarization, exponential causal memory, and memory injection. With fixed block size, local attention yields O(TL) overall complexity instead of O(T^2) for global attention, while memory costs also remain linear.
- Computational decomposition: BCMT performs four operations: local intra-block attention, adaptive block summarization, exponential causal memory construction, and memory injection.These operations are applied successively.
- Local attention: For fixed block size, local attention and its attention-matrix storage grow linearly with sequence length.The attention complexity is O(TL), and the corresponding memory requirement follows the same analysis.
- Adaptive block summaries: Adaptive block summarization remains cheaper than attention when L ≫D−1.Pooling scores and weighted summaries require a single pass over local representations.
- Exponential causal memory: Exponential causal memory construction is linear in the number of blocks.It aggregates block summaries through a normalized causal exponential average, followed by projection and causal shifting.
- Overall complexity: Dense Transformer: O(T 2), BCMT: O(TL).With L ≪T, BCMT becomes linear in sequence length while global-attention Transformers remain quadratic.
- Overall complexity: BCMT reduces computational and memory costs by restricting self-attention to local blocks and propagating long-range dependencies through an exponential causal memory.The memory mechanism has linear computational cost.
4 Experiments
Across context lengths and training scales, BCMT preserves validation performance close to the Dense Transformer while improving throughput and reducing peak GPU memory. Block size controls this trade-off, and ablations show that exponential causal memory—not local partitioning alone—provides the contextual benefit.
- Context-length sensitivity: At 1024 tokens, BCMT reaches 204.2K tokens/s versus 119.9K for the Dense Transformer (+70.3%), while maintaining validation loss of 4.6130 ± 0.0069 versus 4.5752 ± 0.0063.Across evaluated lengths, modeling remains close while computational efficiency improves with longer contexts.
- Block-size trade-off: 57.4% higher throughput and 23.1% lower peak GPU memory make BCMT-256 the best evaluated trade-off, with validation loss 4.5931 ± 0.0045 versus 4.5752 ± 0.0045.Increasing block size improves predictive performance but gradually reduces computational gains.
- Block-size trade-off: 76.5% higher throughput and 29.0% lower peak GPU memory make BCMT-64 the most efficient configuration, but its validation loss rises to 4.6344 ± 0.0037.Smaller blocks limit locally available contextual information.
- Training-data scaling: 57.8% higher throughput, 36.7% shorter epoch time, and 23.1% lower peak GPU memory persist after training on approximately 38 million tokens.BCMT-256 records validation loss 4.1470 ± 0.0003 versus 4.1115±0.0033 for the Dense Transformer, with very small standard deviations.
- Memory ablation: 0.0682 higher validation loss results when memory is removed for 64-token blocks, versus +0.0386 for 128, +0.0121 for 256, and +0.0055 for 512-token blocks.The memory contribution is strongest when local attention covers the least context.
- Memory ablation: Removing inter-block memory consistently worsens validation loss, showing that BCMT’s gains require both local attention and exponential causal memory.Local attention models short-range dependencies, while memory compensates for context lost through block-wise processing.
5 Conclusion
BCMT separates local token interactions from global context propagation, using blockwise dense causal attention and exponential memory over adaptive summaries. Experiments show comparable validation performance to Dense Transformers with higher throughput and lower GPU memory use, while ablations attribute these gains to the memory mechanism.
- Architecture: BCMT separates local interactions, modeled by dense causal self-attention within blocks, from long-range propagation through exponential causal memory over adaptive block summaries.This design avoids explicit attention between all sequence tokens while preserving compatibility with standard dense self-attention implementations.
- Empirical findings: On WikiText, BCMT achieves validation performance comparable to a Dense Transformer while significantly improving training throughput and reducing GPU memory consumption.The conclusion reports these improvements without specifying numerical values.
- Empirical findings: The ablation study confirms that the improvements arise from the proposed memory mechanism rather than local block processing alone.This finding highlights context propagation across blocks as essential to the reported gains.
- Implications: A compact memory built from adaptive block summaries provides a viable alternative to dense global attention for long-context language modeling.The paper further presents separating local modeling from global context propagation as a promising principle for efficient sequence models.
- Future work: Future work includes more expressive summarization, hierarchical or multi-scale memory, and applying BCMT to larger models and tasks requiring stronger reasoning or long-range retrieval.These directions address broader model scale and capabilities beyond the reported experiments.
- Implications: BCMT exemplifies architectures that handle local modeling and global context propagation with distinct, complementary mechanisms as a simple, efficient, fully parallelizable alternative to dense global attention.The conclusion frames this approach as contributing to a new family of sequence-model architectures.
6 Limitations
BCMT’s compact exponential memory reduces computational cost but cannot faithfully retrieve all distant-token information. Its evaluation and memory design remain limited in scope, motivating broader comparisons and adaptive multi-scale extensions.
- Memory representation: BCMT reduces computational cost by propagating long-range dependencies through adaptive block summaries and exponential causal memory rather than explicit all-token interactions.This efficiency relies on a deliberately compact representation of historical context.
- Memory representation: Exponential causal aggregation emphasizes recent blocks and progressively downweights earlier ones, preventing exact retrieval or faithful reconstruction of distant information.The memory provides a global contextual signal but does not preserve every observed token in detail.
- Evaluation scope: The experiments primarily use WikiText language modeling and compare BCMT with a standard Dense Transformer, limiting evaluation across the broader long-context architecture landscape.This focused methodology isolates the proposed memory mechanism’s contribution.
- Evaluation scope: Comparisons with Transformer-XL, Recurrent Memory Transformer, RetNet, RWKV, and Mamba remain an important next step for positioning BCMT precisely.These models represent recent memory-augmented long-context architectures identified for future comparison.
- Future directions: The current BCMT uses a single memory level and fixed exponential decay factor, leaving hierarchical, multi-scale, or adaptive memory as future directions.Such extensions could better represent dependencies across temporal ranges while preserving linear computational complexity.