🤖 machine learning

Block-Based Double Decoders

The paper introduces "Block-Based Double Decoders," a novel transformer architecture that combines the training efficiency of decoder-only models with the inference efficiency of encoder-decoders by utilizing doubly-causal block-based attention masks to achieve superior scaling performance while significantly reducing memory and compute costs.

Original authors: Asher Labovich, Benjamin Bradley, Vanessa Alexander, Chaitanya Harsha

Published 2026-05-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Asher Labovich, Benjamin Bradley, Vanessa Alexander, Chaitanya Harsha

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write stories. For a long time, there have been two main ways to do this, and both have a major flaw.

The Two Old Ways

The "One-Person Show" (Decoder-Only): This is like a student who reads a book and immediately tries to write the next sentence. They are very fast at writing, but they have to remember everything they've read so far in their head while they write. If the story is long, their brain (memory) gets overloaded, and they get slow.
The "Two-Person Team" (Encoder-Decoder): This is like having a Reader and a Writer. The Reader reads the whole book first, takes notes, and then hands the notes to the Writer. This saves the Writer's brain power, but the Reader is very picky. They only take notes on about 15% of the words (the "corrupted" parts) and ignore the rest. This means the team learns slowly because they are ignoring most of the story.

The New Solution: The "Block-Based Double Decoder"

The authors of this paper propose a new team structure that tries to get the best of both worlds. They call it a Block-Based Double Decoder.

Here is how it works, using a simple analogy:

Imagine you are reading a long novel, but instead of reading it one word at a time, you break it into blocks (like chapters or scenes).

The Context Decoder (The Reader): This part reads the whole story up to a certain point. It's like a fast reader who scans the first few chapters to understand the setting and characters. Because it only reads the "past," it doesn't need to hold the whole book in its memory at once. It creates a summary.
The Generation Decoder (The Writer): This part takes the summary from the Reader and the current block of text. It writes the next part of the story.

The Magic Trick: "Doubly-Causal" Blocks

The secret sauce is how they split the story. They chop the text into chunks (blocks).

Inside a single chunk, the Writer can look at all the words in that chunk (like a group discussion).
But when looking at previous chunks, the Writer can only look at the summary provided by the Reader, not the raw words.

Why is this a big deal?

No More Wasted Learning: In the old "Two-Person Team," the Reader ignored 85% of the words. In this new system, every single word gets a chance to be learned from. The model gets a "grade" on every token, making it learn much faster and smarter.
Super Memory Efficiency: When the robot actually writes the story (inference), it doesn't need to remember the whole book. It only needs to remember the summary from the Reader and the current chunk. This cuts the memory needed by about two-thirds. It's like switching from carrying a library in your backpack to just carrying a single index card.
Speed: Because the "Reader" part runs once at the beginning and then sits idle, the "Writer" part is much lighter and faster. It's like having a heavy-duty engine for the start of a race, but a lightweight, aerodynamic body for the sprint.

What did they find?

The researchers tested this new architecture against the old ones.

Training: The new model learned almost as well as the "One-Person Show" (which is the gold standard for speed) and much better than the old "Two-Person Team."
Writing (Inference): When it came time to actually generate text, the new model was a superstar. It used significantly less computer memory and was faster than the old "Two-Person Team," while still being very smart.

The Bottom Line

The paper claims that by splitting the job into two decoders and breaking the text into blocks, they created a model that learns from every word (unlike the old efficient models) but doesn't get bogged down by memory issues when it's time to write (unlike the old fast models). It's a way to get a high-performance robot that doesn't need a massive, expensive computer to run.

Technical Summary: Block-Based Double Decoders

Problem Statement
The paper addresses a fundamental trade-off in transformer architecture design between training efficiency and inference efficiency. Decoder-only models offer superior training efficiency and scalability but suffer from high inference costs, particularly regarding KV-cache memory and per-token compute. Conversely, encoder-decoder models offer significant inference-time savings (reduced KV-cache) but suffer from sparse supervision during pretraining due to objectives like span corruption (where only ~15% of tokens generate a loss signal) and dynamic sequence lengths that hinder efficient batching. Existing alternatives like PrefixLM attempt to bridge this gap but still leave many tokens untrained and fail to achieve the full efficiency of encoder-decoders.

Methodology: Block-Based Double Decoders
The authors propose a novel architecture called Block-Based Double Decoders, which combines two decoder stacks to achieve both full loss supervision and inference efficiency.

Architecture Structure:
- Context Decoder: A standard causal decoder-only transformer that processes the full input sequence and outputs causal latents ( $h_t$ ) for each token.
- Generation Decoder: A second decoder stack that takes three inputs: the causal latents from the context decoder, the raw token sequence, and a block partition.
- Doubly-Causal Block-Based Masking: The input sequence is split into contiguous blocks. Within a block, the generation decoder uses causal self-attention. Between blocks, it uses full cross-attention to the context decoder's latents of preceding blocks. This creates a "doubly-causal" structure where attention flows forward within blocks and from previous blocks.
Training Objective:
- Unlike span corruption, which masks tokens, this method ensures every token in the input sequence contributes to the loss exactly once per forward pass.
- The sequence length remains static after packing, eliminating the need for dynamic batching and token padding, which improves throughput.
- The authors implement the attention mechanism by computing separate self-attention and cross-attention logits and combining them via log-sum-exp normalization post-hoc (using PyTorch's FlexAttention), as a unified dual-key attention implementation is not yet available.
Inference Mechanism:
- At inference time, the Context Decoder runs once over the prompt (prefill). Since it is a standard causal decoder, it does not require a KV-cache for its own output (only the Generation Decoder needs to cache).
- The Generation Decoder then autoregressively generates tokens, attending to the cached context latents and its own previous outputs.
- This separation allows for the reuse of prefix-level KV-caching (common in decoder-only models) while maintaining the memory benefits of encoder-decoder architectures.

Key Contributions

Novel Architecture: Introduction of the block-based double decoder, utilizing two decoder stacks to decouple context processing from token generation.
Doubly-Causal Masking: A new attention mask that enables full token supervision (100% loss density) while maintaining static sequence lengths for efficient batching.
Inference Optimizations: The architecture reduces KV-cache memory and per-token compute by approximately 2/3 compared to decoder-only models (assuming a 2/3 context decoder, 1/3 generation decoder split) without sacrificing prefill caching capabilities.
Training Efficiency: The method avoids the sparse supervision and dynamic batching issues inherent in span-corruption-trained encoder-decoders.

Experimental Results
The authors conducted scaling law experiments comparing Decoder-Only, Standard Encoder-Decoder (SED), and Block-Based Double Decoders across parameter counts (6.25M–100M) and token budgets (62.5M–1B).

Performance: The Double Decoder strongly outperforms the SED baseline and closely tracks the Decoder-Only model. At the largest scale (100M params, 1B tokens), the Double Decoder achieves an evaluation loss only ~0.2 nats worse than the matched Decoder-Only model, whereas the SED trails both by ~0.7 nats.
Scaling Behavior: The SED models were found to be in a "data-limited regime," where increasing parameters did not improve performance due to sparse supervision. In contrast, Double Decoders and Decoder-Only models followed standard scaling laws, flattening as they hit capacity limits.
Compute Trade-off: While Double Decoders incur a slight increase in training FLOPs (~2.4% more than Decoder-Only at specific configurations) due to extra KV projections, this is outweighed by the inference benefits.

Significance and Claims
The paper claims that Block-Based Double Decoders successfully resolve the historical dichotomy between training and inference efficiency. By restoring full loss supervision to an encoder-decoder-like structure, the authors demonstrate that it is possible to achieve the inference efficiency of encoder-decoders (specifically reduced KV-cache and per-token latency) while retaining the training efficiency and scaling properties of decoder-only models.

The authors emphasize that the modest training-time loss gap (~0.2 nats) is a necessary cost for the architectural separation that enables significant inference-time savings. They argue this makes the architecture particularly suitable for edge devices and scenarios with memory constraints, while also enabling new research directions such as latent chain-of-thought reasoning during pretraining and the application of looped transformers specifically to the context decoder.

Technical Summary: Block-Based Double Decoders

More like this