Block-Based Double Decoders
Proposes a block-based double decoder architecture combining full supervision training with inference efficiency, reducing memory and computation by over 66%.
Key Findings
Methodology
This paper introduces a novel double decoder transformer architecture utilizing doubly-causal block attention masks. The first decoder employs standard causal masking, producing contextual latent representations. The second decoder takes these latents and, guided by a block partition, performs local self-attention within blocks and cross-block full attention via a single attention operation. This design ensures each token contributes to the loss in every forward pass, addressing the sparse supervision issue of span corruption methods. The model is pretrained on SlimPajama, comparing against decoder-only and encoder-decoder baselines across scales, with detailed FLOP analysis and inference evaluations.
Key Results
- Across multiple scales, the double decoder outperforms encoder-decoder models in loss, with the largest model (100M params, 1B tokens) showing only 0.2 nat worse than decoder-only, which is a significant improvement over encoder-decoder's 0.7 nat gap. During inference, it reduces KV-cache memory by at least two-thirds and per-token compute, matching decoder-only efficiency while maintaining full supervision.
- Training results demonstrate that full token loss contribution is achieved, overcoming the sparse supervision problem of span corruption. The model maintains high scalability, with experimental data confirming near-parity with decoder-only models in performance and efficiency.
- Inference benefits include reduced memory footprint, faster token generation, and better cache reuse, making it suitable for resource-constrained environments and edge deployment.
Significance
This architecture addresses fundamental limitations of existing models by combining the training efficiency of decoder-only models with the inference advantages of encoder-decoder architectures. It enables large-scale pretraining with dense supervision signals, significantly reducing memory and computational costs during inference. The approach opens new avenues for deploying large language models on edge devices and resource-limited settings, facilitating broader accessibility and application in real-world scenarios.
Technical Contribution
The core innovation is the doubly-causal block attention mask, enabling full supervision during training while maintaining efficient inference. The architecture separates context encoding and response generation into two stacks, with a single attention operation integrating local and global dependencies. Theoretical analysis and empirical results demonstrate its scalability and efficiency, establishing a new paradigm for Transformer training and inference. The method also provides a flexible framework for future extensions such as reasoning modules or recursive attention.
Novelty
This work is the first to introduce a block-based double decoding mechanism that combines local self-attention within blocks with cross-block full attention in a single, unified attention operation. Unlike traditional encoder-decoder models with bidirectional attention and decoder-only models with sparse supervision, this architecture ensures dense token supervision and efficient inference, representing a significant departure from existing paradigms.
Limitations
- The block partitioning relies on random sampling, which may introduce variability and affect training stability. More deterministic or adaptive partitioning strategies could improve robustness.
- Current experiments are limited to models up to 100 million parameters; scalability to larger models (e.g., billions of parameters) remains to be validated.
- The attention fusion mechanism, while effective, introduces additional computational overhead, which could be further optimized for faster inference.
Future Work
Future research will explore adaptive block partitioning strategies, more efficient attention fusion algorithms, and scaling to larger models. Combining this architecture with reasoning modules like chain-of-thought or recursive attention could further enhance model capabilities. Additionally, integrating looped transformers and test-time reasoning techniques promises to push the boundaries of inference speed and reasoning depth.
AI Executive Summary
Transformer models have revolutionized natural language processing, yet their training and inference costs remain significant barriers, especially at scale. Traditional encoder-decoder architectures excel in understanding tasks but suffer from sparse supervision and dynamic sequence lengths during pretraining. Conversely, decoder-only models like GPT are highly scalable but lack the ability to efficiently incorporate full context during inference, limiting their deployment in resource-constrained environments.
This paper introduces a novel block-based double decoder architecture that combines the best of both worlds. The core idea involves splitting input sequences into contiguous blocks and employing a dual-stack transformer: a context decoder that processes the entire input with causal masking, and a generation decoder that attends locally within blocks and globally across blocks via a single attention operation. This design ensures every token contributes to the training loss, overcoming the sparse supervision problem typical of span corruption objectives. During training, the model achieves full supervision density, while at inference, it significantly reduces memory and computation, matching decoder-only efficiency.
Extensive experiments on SlimPajama datasets across multiple scales demonstrate the architecture's effectiveness. The largest model achieves near-parity with decoder-only models in loss, with only a 0.2 nat gap, while outperforming encoder-decoder baselines by a wide margin. In inference, the model reduces KV-cache memory by over 66%, enabling faster and more resource-efficient deployment. These results highlight the architecture's potential to enable large-scale, dense-supervision pretraining with low-latency inference, making it highly suitable for edge deployment and real-time applications.
Looking forward, further optimization of block partitioning, attention fusion, and scaling to larger models could unlock even greater performance. Combining this approach with reasoning modules and recursive attention mechanisms offers promising directions for future research, aiming to build more intelligent, efficient, and accessible language models.
Deep Dive
Abstract
Encoder-decoder models offer substantial inference-time savings over decoder-only models, but their pretraining objectives suffer from sparse supervision and dynamic sequence lengths, keeping them out of practice at scale. We propose block-based double decoders, a novel transformer architecture that utilizes doubly-causal block-based attention masks to train with full loss supervision and static sequence packing, combining decoder-only training efficiency with encoder-decoder inference efficiency. In scaling law experiments, block-based double decoders strongly outperform encoder-decoders and closely track decoder-only models across scales. At inference time, they cut KV-cache memory and per-token compute by at least 2/3 without sacrificing prefill caching or other existing inference optimizations available to decoder-only models.