Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

Technical University of Munich (TUM) · Imperial College London · Munich Center for Machine Learning (MCML) · Hasso Plattner Institute, University of Potsdam

Published September 25, 2026 · on the feed September 28, 2026 · ▲ 11 · arXiv 2608.09444

The abstract

A main promise of looped language models is depth-adaptive inference. By looping a block of shared layers a variable number of times, the model can use less compute for "easy" tokens and more for "hard" ones. However, tokens with different numbers of loops cannot share a uniform forward pass and therefore cannot be handled by standard batching systems such as vLLM. The practical value of depth-adaptive inference thus hinges on whether batching can be made efficient. We introduce the first efficient method for depth-adaptive looped LMs via continuous depth batching (CDB), which forms new batches between loop steps. Our method dynamically schedules looped and non-looped parts of the architecture, manages looped KV-caching, and predicts which tokens will exit the loop in advance so it can prepare batches asynchronously. Experiments on Ouro 1.4B and Huginn 3.5B show that fully looped architectures are best suited to depth-adaptive inference, as large non-looped layers outside the recurrent core (e.g., token embedding, LM head, and unshared transformer blocks) slow down and complicate scheduling. Overall, CDB realizes up to 99% of the estimated maximum speedup available, leaving further gains primarily dependent on model architecture and exit behavior.

The Brief

2:11

A short podcast summary of the paper. Written and voiced by AI.

Transcript

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching. By Kristian Schwethelm and two others, from the Technical University of Munich, Imperial College London, and two other institutions. Published on the twenty-fifth of September. Its code repository is worked out: about one hundred and eighty source files, with a readme and a licence.

Looped language models reuse a block of shared layers, running it a variable number of times. One of their main promises follows from that: depth-adaptive inference, where easy tokens get fewer loops and less compute, and hard tokens get more. The catch is practical. Tokens that loop different numbers of times cannot share one uniform forward pass, so standard batching systems such as vLLM cannot handle them. Whether depth-adaptive inference pays off in practice therefore comes down to one question: can batching be made efficient when every token wants a different depth?

The authors present what they describe as the first efficient method for this, continuous depth batching, or CDB. Its central move is where batches get formed. Instead of fixing a batch at the start of a forward pass, CDB forms new batches between loop steps, so tokens can be regrouped as some leave the loop and others continue.

Making that work takes three pieces of machinery. The method dynamically schedules the looped and the non-looped parts of the architecture. It manages KV-caching for the looped part. And it predicts in advance which tokens are going to exit the loop, which lets it prepare the next batches asynchronously instead of waiting to find out.

Experiments use two models, Ouro at one point four billion parameters and Huginn at three point five billion. They show that fully looped architectures are the best suited to depth-adaptive inference. Large non-looped layers sitting outside the recurrent core, such as the token embedding, the LM head and unshared transformer blocks, slow scheduling down and complicate it. Overall, CDB realizes up to ninety-nine percent of the estimated maximum speedup available.

That last figure changes where the remaining work lies. With the batching problem close to its ceiling, the authors place further gains mainly in model architecture and in how tokens exit the loop. The result about fully looped designs points the same way: the scheduler is now asking the architecture to meet it.

Paper Quest · 7 chapters · about 40 min

Play this paper

Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.

Chapter 1 🧩 · The problem

Every quest starts with a wall. Here it is compute saved on paper that never shows up as speed.

You serve a language model that runs one shared block of layers over and over. Easy tokens could stop after a few loops, hard ones could keep going, saving compute.

The road ahead

  1. 1 · The problem ~7 min A looped model could give easy tokens fewer loops, but standard batching runs every token through the same pass, so the saved compute never becomes speed.
  2. 2 · What came before ~3 min
  3. 3 · The key idea ~8 min Continuous depth batching builds batches between loop steps, runs each stage on its own queue, and learns of exits early enough to prepare batches in the background.
  4. 4 · Choose your own research ~3 min
  5. 5 · Build the method ~8 min CDB predicts exits a loop ahead, prepares the next batch on the CPU in the background, then forms the new batch between loop steps.
  6. 6 · The experiment ~6 min You’ll guess: At best, what share of the estimated maximum speedup from depth-adaptive inference did continuous depth batching actually realize on looped language models?
  7. 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Join free to play the whole quest ›

Free, with Google or an email address. Your stars and progress are kept.

KipQuest: the day’s AI research, as podcasts and quests. All papers Android app Privacy Contact Episodes are written and voiced by AI from each paper. Summaries can contain mistakes: the paper is the source.