Continuous Depth-wise Batching (CDB)

• Continuous Depth-wise Batching (CDB) is a method that generalizes batching from the sequence to the depth dimension, enabling immediate reuse of computation slots when a sample exits early.
• It leverages Recursive Transformers with shared weight blocks and stage-specific mechanisms to create a depth-aware pipeline that maximizes throughput.
• The technique yields significant performance gains—up to 4x speedup in simulated setups—by efficiently scheduling model inference without impacting prediction accuracy.

Searching arXiv for the specified papers and closely related work. I’ll retrieve the arXiv metadata for the cited papers to ground the article. Continuous Depth-wise Batching (CDB) is an inference paradigm introduced alongside Relaxed Recursive Transformers to exploit the looped structure of Recursive Transformers and maximize hardware utilization, especially when combined with early-exit mechanisms (Bae et al., 2024). In this setting, a model is organized as repeated applications of a shared block, and inference is scheduled not only across sequences but also across depth. The central idea is that whenever a sample leaves a depth stage—because it has exited early or completed that iteration—its slot in the next depth can be filled immediately by another sample already waiting at that depth or by a new request entering at depth 1. This reframes batching as a depth-aware pipeline and is intended to close utilization gaps left by vanilla batching or early exit.

1. Motivation and Operating Principle

Under standard Transformer inference, even with continuous sequence-wise batching, new requests can only enter the pipeline once the entire current batch has finished all layers or decoding timesteps. The stated motivation for CDB is that this leaves the device under-utilized whenever some sequences finish earlier than others, such as under early exiting, but the batch cannot advance (Bae et al., 2024).

Recursive Transformers provide the structural precondition for CDB. In the formulation described for Relaxed Recursive Transformers, the same $K$-layer block is applied repeatedly, so depth can be interpreted as loop iteration. This makes it possible to treat each loop iteration as a stage in a pipeline, all using the same weights. CDB generalizes continuous batching from the time or timestep dimension to the depth or loop dimension. A sample that finishes one stage can immediately free capacity that is reused either by another in-flight sample at a compatible depth or by a brand-new request at the first stage.

This suggests that CDB is not merely a batching convenience, but a scheduling policy tied to a particular architectural regularity: repeated execution of a shared block. Early exiting amplifies the benefit because samples do not all consume the same number of loop iterations.

2. Integration with Recursive Transformers and Early Exit

The architectural context is a Recursive Transformer in which the model consists of $B$ loop iterations of a shared $K$-layer block. Depth $\ell$ of the full $L = B \cdot K$ layers corresponds to loop index $\lceil \ell / K \rceil$ modulo the shared block structure (Bae et al., 2024). In the relaxed variant, stage-specific LoRA adapters can be attached while preserving the shared backbone.

Early exiting is applied at loop iteration $i$, that is, at depth $i \cdot K$, by evaluating a confidence test on the partial hidden state. If the confidence is at least a threshold, inference terminates immediately and the remaining loops are skipped. CDB takes this model and interprets each of the loop iterations as a pipeline stage $S_1, S_2, \ldots, S_{B}$. A queue $Q_k$ is maintained for samples waiting for stage $B$0. New requests are admitted into $B$1; a sample that does not exit at stage $B$2 is enqueued into $B$3.

The resulting data flow is stage-centric. Worker threads or an asynchronous scheduler repeatedly pull up to batch_size samples from a stage queue whose depth is ready, execute the shared $B$4-layer block together with any stage-specific LoRA adapters, and then either emit a final token or enqueue the sample into the next depth queue. Because the weights are shared across stages, stage-local work is homogeneous enough to support repeated kernel launches over the same block. A plausible implication is that the shared-block design reduces some of the coordination complexity that would arise if every depth used a distinct parameterization.

3. Formalization and Throughput Analysis

The theoretical analysis introduces the following notation. Let $B$5 denote the maximum batch size. Let $B$6 be the number of layers in the shared block, and let $B$7 be the number of loop iterations, so that $B$8. Let $B$9 be the time to process the $K$0-th layer of the shared block on a full batch of $K$1 samples, or scaled accordingly if fewer samples are present. For simplicity, the analysis assumes $K$2, so the time per loop iteration is

$K$3

Let $K$4 be the fraction of tokens or sequence positions that survive through at least $K$5 loop iterations, meaning they have not exited before stage $K$6. By definition, $K$7 and $K$8. The average compute per token under early exit is then

$K$9

Vanilla throughput, with continuous sequence-wise batching but no depth-wise reuse, is given as

$\ell$0

For CDB, at steady state, each stage $\ell$1 is kept busy with $\ell$2 samples on average until those samples exit. The wall-clock time to complete one round of outputs for $\ell$3 samples is therefore $\ell$4, yielding

$\ell$5

The relative speedup is

$\ell$6

The paper provides an illustrative case with $\ell$7: if $\ell$8 of tokens exit after the first loop, so $\ell$9, $L = B \cdot K$0 exit after the second, so $L = B \cdot K$1, and the rest go full depth with $L = B \cdot K$2, then $L = B \cdot K$3, giving a speedup of $L = B \cdot K$4 (Bae et al., 2024). The same analysis states that, with deeper early-exit distributions measured in the paper, $L = B \cdot K$5 can be as low as approximately $L = B \cdot K$6, yielding $L = B \cdot K$7–$L = B \cdot K$8 gains.

4. Scheduling Procedure and Systems Realization

The high-level algorithm initializes queues $L = B \cdot K$9 as empty, admits new requests into depth 1 while $\lceil \ell / K \rceil$0 and requests are available, and then iterates over stages looking for a nonempty queue (Bae et al., 2024). For the first nonempty stage, it dequeues up to $\lceil \ell / K \rceil$1 samples, runs SharedBlockForward(batch, stage=i), and then applies the confidence test. Samples whose confidence reaches the threshold emit their final token; samples that do not exit and are not yet at the last loop iteration are enqueued into the next stage.

Two aspects are central to the intended behavior. First, stage 1 is always filled up to $\lceil \ell / K \rceil$2 when possible, so the device does not idle waiting for new inputs at the start of a loop. Second, whenever a sample exits early, its slot is effectively freed at all deeper stages and is immediately reused by another sample in the appropriate queue or by a fresh request in $\lceil \ell / K \rceil$3.

The implementation discussion emphasizes host-side scheduling and asynchronous execution. A small scheduler tracks $\lceil \ell / K \rceil$4 and launches a CUDA kernel for whichever stage has available data and a free compute stream. Because all depths use the same shared weights, mixed-depth inputs can, if memory allows, be concatenated so that a single fused kernel processes them in one pass. For autoregressive decoding, each sample carries its own key-value cache per layer; when a sample moves from stage $\lceil \ell / K \rceil$5 to stage $\lceil \ell / K \rceil$6, its cache moves with it, and exited samples’ caches can be freed immediately. The maximum simultaneous number of samples in flight across all depths is $\lceil \ell / K \rceil$7, though only $\lceil \ell / K \rceil$8 are present at any one depth. Multi-GPU or multi-stream setups can pipeline stages in parallel if desired.

5. Empirical Results and Performance Claims

The reported experiments simulate CDB paired with oracle early-exiting on several models, including Gemma 2B, TinyLlama 1.1B, and Pythia 1B (Bae et al., 2024). The headline findings concern throughput rather than model quality, because the paper explicitly states that no accuracy degradation is incurred by CDB itself: the model’s predictions and exit criteria remain identical. The stated trade-off arises only when higher-rank LoRA is used in relaxed recursive variants, which slightly increases model size.

Configuration	Throughput claim	Notes
Vanilla Pythia 1B	$\lceil \ell / K \rceil$9 baseline; 1,080 tokens/sec	No CDB
Pythia 1B + CSB	$i$0; $i$1 tok/s	Continuous sequence-wise batching
2-block recursive Pythia ($i$2B params) + CDB + early exit	$i$3; $i$4 tok/s	Recursive conversion with CDB
Relaxed recursive variants, LoRA rank $i$5	$i$6–$i$7	Throughput traded for improved accuracy

Across SlimPajama, RedPajama, and PG19, the average CDB gains are reported as $i$8–$i$9 over vanilla, outperforming continuous sequence-wise batching alone. For Gemma 2B, the paper states that near $i \cdot K$0 end-to-end speedup is theoretically attainable when comparing recursive and vanilla settings. Latency to first token is also described as improving because stages are never starved waiting for full-batch data.

These results should be read with two distinctions in mind. First, the measurements are tied to recursive or relaxed recursive models rather than to unmodified dense Transformers. Second, the paper describes the CDB results as simulated with oracle early-exiting, so the gains are coupled to the assumed exit behavior.

6. Relation to Earlier Depth-Based Batching and Conceptual Boundaries

CDB is distinct from the depth-based on-the-fly batching heuristic developed for dynamic computation graphs in "On-the-fly Operation Batching in Dynamic Computation Graphs" (Neubig et al., 2017). In that earlier setting, the computation is modeled as a directed acyclic graph $i \cdot K$1, node depth is defined as the length of the longest path from any leaf node to a node, and execution proceeds depth by depth, signature by signature, batching together ready nodes at equal depth that share the same signature. Grouping is performed over graph nodes, and correctness follows from the fact that if $i \cdot K$2 then $i \cdot K$3.

CDB operates on a different object and under different assumptions. Its stages are loop iterations of a shared Transformer block rather than arbitrary nodes in a dynamic computation graph. Its queues track samples waiting at successive depths, and the mechanism is specifically designed to exploit recursive weight sharing together with early exit. A common confusion arises from the shared phrase “depth-wise batching,” but the two methods address different scheduling problems. The 2017 method batches same-signature operations at the same graph depth; CDB pipelines request flow across model depths in a Recursive Transformer.

The contrast is also informative at the level of opportunity and constraint. The 2017 paper notes that strict depth-based batching can miss cross-depth opportunities and that agenda-based scheduling can recover some of them (Neubig et al., 2017). CDB, by comparison, is explicitly designed to capitalize on movement between depths, because a sample that leaves one stage immediately creates reusable capacity elsewhere in the pipeline. This suggests that the reuse of a single shared block is not incidental but constitutive of the method’s efficiency model.

7. Limitations, Trade-offs, and Scope

The paper presents CDB as a promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting (Bae et al., 2024). That phrasing is itself a boundary condition: CDB is not introduced as a universally applicable batching strategy for arbitrary Transformer deployments, but as one that depends on the recursive weight-sharing structure and benefits particularly from heterogeneous exit depths.

The reported performance gains are theoretical or simulated under oracle early-exiting, and the throughput analysis assumes a simplified per-layer timing model with $i \cdot K$4. This suggests that realized speedups may depend on how closely an implementation matches the assumed steady-state pipeline behavior, on memory available for mixed-depth or multi-stream execution, and on the actual survival fractions $i \cdot K$5. The paper nevertheless states that the model’s predictions and exit criteria remain identical under CDB itself, so the mechanism is presented as a systems-level scheduling change rather than an approximation to model computation.

Within that scope, CDB is characterized by three defining properties: it extends batching from the sequence dimension to the depth dimension, it uses stage queues over repeated shared blocks, and it derives its throughput advantage from keeping those depth stages busy while early-exiting samples release capacity.