{"paper":{"title":"Mixture-of-Depths: Dynamically allocating compute in transformer-based language models","license":"http://creativecommons.org/licenses/by/4.0/","headline":"Transformer language models can learn to dynamically allocate compute to select tokens at each layer.","cross_cats":["cs.CL"],"primary_cat":"cs.LG","authors_text":"Adam Santoro, Blake Richards, David Raposo, Peter Conway Humphreys, Sam Ritter, Timothy Lillicrap","submitted_at":"2024-04-02T19:28:11Z","abstract_excerpt":"Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori,"},"claims":{"count":4,"items":[{"kind":"strongest_claim","text":"Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50% faster to step during post-training sampling.","source":"verdict.strongest_claim","status":"machine_extracted","claim_id":"C1","attestation":"unclaimed"},{"kind":"weakest_assumption","text":"The assumption that a learned top-k router can reliably identify which tokens merit full processing at each layer without degrading overall model capacity or introducing training instabilities, and that this holds across model scales and tasks.","source":"verdict.weakest_assumption","status":"machine_extracted","claim_id":"C2","attestation":"unclaimed"},{"kind":"one_line_summary","text":"Mixture-of-Depths enables transformers to dynamically allocate compute by routing only the top-k tokens through each layer's full computations, matching baseline performance with a fraction of the FLOPs per forward pass and up to 50% faster sampling.","source":"verdict.one_line_summary","status":"machine_extracted","claim_id":"C3","attestation":"unclaimed"},{"kind":"headline","text":"Transformer language models can learn to dynamically allocate compute to select tokens at each layer.","source":"verdict.pith_extraction.headline","status":"machine_extracted","claim_id":"C4","attestation":"unclaimed"}],"snapshot_sha256":"2487a200c9eb271af365818e08fa621e4c53c12a8d37f2ee982a5b3b838e1b76"},"source":{"id":"2404.02258","kind":"arxiv","version":1},"verdict":{"id":"47fc0ddc-7aad-4707-9d3b-017bc1794454","model_set":{"reader":"grok-4.3"},"created_at":"2026-05-17T02:13:07.402130Z","strongest_claim":"Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50% faster to step during post-training sampling.","one_line_summary":"Mixture-of-Depths enables transformers to dynamically allocate compute by routing only the top-k tokens through each layer's full computations, matching baseline performance with a fraction of the FLOPs per forward pass and up to 50% faster sampling.","pipeline_version":"pith-pipeline@v0.9.0","weakest_assumption":"The assumption that a learned top-k router can reliably identify which tokens merit full processing at each layer without degrading overall model capacity or introducing training instabilities, and that this holds across model scales and tasks.","pith_extraction_headline":"Transformer language models can learn to dynamically allocate compute to select tokens at each layer."},"references":{"count":11,"sample":[{"doi":"","year":2002,"title":"Controlling computation versus quality for neural sequence models","work_id":"21f2a3df-f950-438d-8553-6db49dc7afe8","ref_index":2,"cited_arxiv_id":"","is_internal_anchor":false},{"doi":"","year":null,"title":"Universal Transformers","work_id":"8e5baefe-d209-411c-aefc-5acaa9275c8a","ref_index":3,"cited_arxiv_id":"1807.03819","is_internal_anchor":true},{"doi":"","year":1910,"title":"Depth-adaptive transformer","work_id":"6d994f34-ab13-423e-96ec-52b63f47a6d6","ref_index":5,"cited_arxiv_id":"","is_internal_anchor":false},{"doi":"","year":null,"title":"Adaptive Computation Time for Recurrent Neural Networks","work_id":"75565443-173e-479c-b0e7-d2464e7630be","ref_index":7,"cited_arxiv_id":"1603.08983","is_internal_anchor":true},{"doi":"","year":null,"title":"Towards a unified view of parameter-efficient transfer learning","work_id":"e758988d-4ae2-44ff-94fc-4417c242b1e4","ref_index":8,"cited_arxiv_id":"","is_internal_anchor":false}],"resolved_work":11,"snapshot_sha256":"eb3b74aaf8e7de70ad3073dbc0bfbc87e51628e687adfc5c29387dbc73af8d74","internal_anchors":5},"formal_canon":{"evidence_count":1,"snapshot_sha256":"740ed61cf7e8970d28f0d74fd67712b045d4b6754137d9ae6d7eec1875379ca6"},"author_claims":{"count":0,"strong_count":0,"snapshot_sha256":"258153158e38e3291e3d48162225fcdb2d5a3ed65a07baac614ab91432fd4f57"},"builder_version":"pith-number-builder-2026-05-17-v1"}