Dimension of each layer in multi headed transformation attention block

Browse all previously published AI Tutorials here.

Table of Contents

• Dimension of each layer in multi headed transformation attention block

• Multi-Head Self-Attention in Transformers for Document Processing (2024-2025 Review)

• Introduction

• Anatomy of a Multi-Head Self-Attention Layer

• Theoretical Insights on Dimension Choices and Representation

• Cross-Attention and Memory-Augmented Attention

• Optimizations for Long Document Inference and Fine-Tuning

• Conclusion

• References

Multi-Head Self-Attention in Transformers for Document Processing (2024-2025 Review)

Introduction

Transformer-based large language models (LLMs) rely on multi-head self-attention as a core mechanism for sequence representation and long-range context modeling. In document digitization tasks, where long texts must be chunked for processing, the attention architecture is crucial to capturing relationships within and across chunks. Recent research in 2024 and 2025 has expanded Transformers’ context lengths from the original 8K tokens to 128K or even 1M tokens (Shifting Long-Context LLMs Research from Input to Output), enabling applications like summarizing entire books or multi-chapter analysis. This literature review examines the layer dimensions in a multi-headed self-attention block (queries, keys, values, projections) and how they affect attention computation and representation learning. We also highlight extensions (cross-attention, memory-augmented attention) and practical optimizations for efficient inference and fine-tuning on long documents, citing the latest studies.

Anatomy of a Multi-Head Self-Attention Layer

Theoretical Insights on Dimension Choices and Representation

Cross-Attention and Memory-Augmented Attention

Cross-Attention: In encoder-decoder Transformers and retrieval-augmented models, cross-attention layers have a similar structure to self-attention but operate on two sequences. The query Q comes from a “target” sequence (e.g. a decoder’s hidden state or a query embedding), while keys and values come from a different source (e.g. the encoder’s outputs or an external knowledge chunk). The projection dimensions (query/key/value size per head) are typically aligned so that Q and K vectors are in the same space, enabling dot-product attention across sequences. Cross-attention thus lets a model attend to external information, such as a document chunk relevant to a query. For example, retrieval-augmented LLMs include an extra cross-attention module in each Transformer block to integrate retrieved knowledge into the model’s hidden state (Retrieval-Augmented Generation for Natural Language Processing: A Survey). During generation, the model’s internal representation (prefix) can attend to retrieved document embeddings via this cross-attention, calibrating token predictions based on the external context . This mechanism is essential in document question-answering or open-domain QA, where the model must fuse a question with relevant text chunks. Recent surveys confirm that many RAG architectures feed retrieved features through cross-attention (at one or multiple layers) to inject outside information into the decoder . The dimensions of these cross-attention layers mirror those of self-attention, ensuring compatibility with the Transformer’s embedding size while doubling the keys/values to incorporate the external sequence.

Memory-Augmented Attention: To handle ultra-long documents or streams beyond the normal context length, researchers have developed memory-augmented attention mechanisms. These introduce an external memory store or compressive context that the model can attend to, effectively extending the available context window without exploding computational cost. Infini-Attention (2024) is one such approach that adds a compressive memory to the vanilla self-attention (HERE). It combines standard local self-attention on the current chunk with a long-term linear attention over stored key-value pairs from previous segments, all within a single attention block . This design allows a Transformer to process “infinite” input streams by retaining condensed representations of earlier tokens instead of attending over an ever-growing list of keys (which would be intractable). Munkhdalai et al. (2024) demonstrate this infinite-context transformer can handle inputs of 500K–1M tokens (e.g. book-length text) by maintaining a bounded memory footprint . Another line of work, latent-space memory, compresses past activations into a fixed-size memory pool that later tokens can attend to. Wang et al. (2025) introduce MemoryLLM+ (M+), which extends a latent memory module to drastically improve long-term retention (M+: Extending MemoryLLM with Scalable Long-Term Memory). Their model co-trains a retriever that selects relevant information from a large memory bank and presents it through an attention mechanism during generation . This lets an LLM retain and recall content from over 160K tokens in the past with minimal overhead . Similarly, MELODI (ICLR 2025) uses a hierarchical compression of memories across layers and windows: short-term context is recursively compressed at each layer, while a mid-layer captures long-term context by compressing across multiple windows (MELODI: Exploring Memory Compression for Long Contexts | OpenReview) . This two-tier compression scheme enables processing very long documents in chunks, yet aggregating crucial information from the entire sequence. Compared to a dense-attention baseline that explicitly stores 64K tokens in memory, MELODI achieved better accuracy on long-document tasks with an 8× reduction in memory usage . These memory-augmented attentions maintain the Q/K/V mechanism but alter how keys and values persist across chunks – an important strategy for document digitization pipelines that feed LLMs piecewise while preserving global context.

Optimizations for Long Document Inference and Fine-Tuning

Conclusion

Research from 2024–2025 provides a deeper understanding of multi-head self-attention layers and practical ways to extend them for long-document processing. A structured breakdown of self-attention reveals how query/key/value dimensions and head counts determine the capacity of attention to learn relationships. New theoretical results show the importance of sufficient per-head dimension (rank) for expressive power ( On the Benefits of Rank in Attention Layers), while also confirming that many heads can be redundant (MoH: Multi-Head Attention as Mixture-of-Head Attention). In applied settings, cross-attention and memory-augmented attention mechanisms have become vital for integrating external knowledge and extending context windows beyond the limit of naive self-attention (HERE). For document digitization pipelines, the community has developed hierarchical and memory-efficient attention architectures that maintain tractable computation over long texts ( HDT: Hierarchical Document Transformer). Coupled with inference optimizations (like FlashAttention-2) and fine-tuning strategies (LongLoRA, LoRA) (HERE) , these advances allow modern LLMs to effectively process and analyze very large documents. Continued research is likely to further balance the dimensional design of attention layers with clever sparsity and memory techniques, pushing the limits of context length and efficiency in Transformer-based document understanding.

References

• Ashish Vaswani et al. Attention Is All You Need. 2017. (Original Transformer introduction for reference)

• Noah Amsel et al. “On the Benefits of Rank in Attention Layers.” arXiv preprint, Jul. 2024 .

• Peng Jin et al. “MoH: Multi-Head Attention as Mixture-of-Head Attention.” arXiv preprint, Oct. 2024 .

• Tsendsuren Munkhdalai et al. “Leave No Context Behind: Efficient Infinite Context Transformers with Infini-Attention.” arXiv preprint, Aug. 2024 .

• Yu Wang et al. “M+: Extending MemoryLLM with Scalable Long-Term Memory.” arXiv preprint, Feb. 2025 .

• Yinpeng Chen et al. “MELODI: Exploring Memory Compression for Long Contexts.” ICLR 2025 .

• Haoyu He et al. “HDT: Hierarchical Document Transformer.” arXiv preprint, Jul. 2024 .

• Yukang Chen et al. “LongLoRA: Efficient Fine-Tuning of Long-Context LLMs.” ICLR 2024 .

• Tri Dao. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” ICLR 2024 .

• (Additional citations from text: OpenAI 2024a; Anthropic 2024; etc., as cited in Wu et al. 2025 (Shifting Long-Context LLMs Research from Input to Output); and LoRA fine-tuning as per Zhao et al. 2024 ( LoRA based Parameter Efficient Fine-Tuning using Optimal ... - arXiv).)