{"id":"https://openalex.org/W4393967974","doi":"https://doi.org/10.48550/arxiv.2404.02258","title":"Mixture-of-Depths: Dynamically allocating compute in transformer-based language models","display_name":"Mixture-of-Depths: Dynamically allocating compute in transformer-based language models","publication_year":2024,"publication_date":"2024-04-02","ids":{"openalex":"https://openalex.org/W4393967974","doi":"https://doi.org/10.48550/arxiv.2404.02258"},"language":"en","primary_location":{"id":"pmh:oai:arXiv.org:2404.02258","doi":null,"is_oa":true,"landing_page_url":"http://arxiv.org/abs/2404.02258","pdf_url":"https://arxiv.org/pdf/2404.02258","source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":"2331-8422","issn":["2331-8422"],"is_oa":true,"is_in_doaj":false,"is_core":false,"listed_in":[],"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":"submittedVersion","is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"text"},"volume":null,"type":"preprint","indexed_in":["arxiv","datacite"],"open_access":{"is_oa":true,"oa_status":"green","oa_url":"https://arxiv.org/pdf/2404.02258","any_repository_has_fulltext":true},"authorships":[{"author_position":"first","author":{"id":"https://openalex.org/A5103096285","display_name":"David Chillo\u0301n Raposo","orcid":"https://orcid.org/0000-0001-8202-5532","observed_orcids":["https://orcid.org/0000-0001-8202-5532"]},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Raposo, David","raw_affiliation_strings":["Equal"],"raw_orcid":null,"affiliations":[{"raw_affiliation_string":"Equal","institution_ids":[]}]},{"author_position":"middle","author":{"id":"https://openalex.org/A5036044173","display_name":"Sam Ritter","orcid":null,"observed_orcids":[]},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Ritter, Sam","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5004133705","display_name":"Blake A. Richards","orcid":"https://orcid.org/0000-0001-9662-2151","observed_orcids":["https://orcid.org/0000-0001-9662-2151"]},"institutions":[{"id":"https://openalex.org/I5023651","display_name":"McGill University","ror":"https://ror.org/01pxwe438","country_code":"CA","type":"education","lineage":["https://openalex.org/I5023651"]}],"countries":["CA"],"is_corresponding":false,"raw_author_name":"Richards, Blake","raw_affiliation_strings":["McGill University & Mila ,"],"raw_orcid":null,"affiliations":[{"raw_affiliation_string":"McGill University & Mila ,","institution_ids":["https://openalex.org/I5023651"]}]},{"author_position":"middle","author":{"id":"https://openalex.org/A5066294254","display_name":"Timothy Lillicrap","orcid":"https://orcid.org/0000-0001-8918-486X","observed_orcids":["https://orcid.org/0000-0001-8918-486X"]},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Lillicrap, Timothy","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":null,"display_name":"Humphreys, Peter Conway","orcid":null,"observed_orcids":[]},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Humphreys, Peter Conway","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"last","author":{"id":"https://openalex.org/A5112705664","display_name":"Adam Santoro","orcid":null,"observed_orcids":[]},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Santoro, Adam","raw_affiliation_strings":["Equal"],"raw_orcid":null,"affiliations":[{"raw_affiliation_string":"Equal","institution_ids":[]}]}],"institutions":[],"countries_distinct_count":1,"institutions_distinct_count":1,"corresponding_author_ids":[],"corresponding_institution_ids":[],"apc_list":null,"apc_paid":null,"fwci":null,"has_fulltext":true,"cited_by_count":9,"citation_normalized_percentile":null,"cited_by_percentile_year":null,"biblio":{"volume":null,"issue":null,"first_page":null,"last_page":null},"is_retracted":false,"is_paratext":false,"is_xpac":false,"primary_topic":{"id":"https://openalex.org/T10036","display_name":"Advanced Neural Network Applications","score":0.6640999913215637,"subfield":{"id":"https://openalex.org/subfields/1707","display_name":"Computer Vision and Pattern Recognition"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},"topics":[{"id":"https://openalex.org/T10036","display_name":"Advanced Neural Network Applications","score":0.6640999913215637,"subfield":{"id":"https://openalex.org/subfields/1707","display_name":"Computer Vision and Pattern Recognition"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},{"id":"https://openalex.org/T10028","display_name":"Topic Modeling","score":0.25609999895095825,"subfield":{"id":"https://openalex.org/subfields/1702","display_name":"Artificial Intelligence"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},{"id":"https://openalex.org/T10181","display_name":"Natural Language Processing Techniques","score":0.04969999939203262,"subfield":{"id":"https://openalex.org/subfields/1702","display_name":"Artificial Intelligence"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}}],"keywords":[{"id":"https://openalex.org/keywords/transformer-based-language-models","display_name":"transformer language models","score":0.8519999980926514},{"id":"https://openalex.org/keywords/top-k-routing","display_name":"top-k routing","score":0.5820000171661377},{"id":"https://openalex.org/keywords/self-attention","display_name":"self-attention","score":0.6290000081062317},{"id":"https://openalex.org/keywords/compute-budget","display_name":"compute budget","score":0.5450000166893005},{"id":"https://openalex.org/keywords/model-compression","display_name":"model compression","score":0.4699999988079071},{"id":"https://openalex.org/keywords/conditional-computation","display_name":"conditional computation","score":0.550000011920929},{"id":"https://openalex.org/keywords/inference-acceleration","display_name":"inference acceleration","score":0.4390000104904175}],"concepts":[{"id":"https://openalex.org/C66322947","wikidata":"https://www.wikidata.org/wiki/Q11658","display_name":"Transformer","level":3,"score":0.6879242658615112},{"id":"https://openalex.org/C41008148","wikidata":"https://www.wikidata.org/wiki/Q21198","display_name":"Computer science","level":0,"score":0.5352047085762024},{"id":"https://openalex.org/C119599485","wikidata":"https://www.wikidata.org/wiki/Q43035","display_name":"Electrical engineering","level":1,"score":0.1481896936893463},{"id":"https://openalex.org/C127413603","wikidata":"https://www.wikidata.org/wiki/Q11023","display_name":"Engineering","level":0,"score":0.14802348613739014},{"id":"https://openalex.org/C165801399","wikidata":"https://www.wikidata.org/wiki/Q25428","display_name":"Voltage","level":2,"score":0.10071158409118652}],"mesh":[],"locations_count":2,"locations":[{"id":"pmh:oai:arXiv.org:2404.02258","doi":null,"is_oa":true,"landing_page_url":"http://arxiv.org/abs/2404.02258","pdf_url":"https://arxiv.org/pdf/2404.02258","source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":"2331-8422","issn":["2331-8422"],"is_oa":true,"is_in_doaj":false,"is_core":false,"listed_in":[],"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":"submittedVersion","is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"text"},{"id":"doi:10.48550/arxiv.2404.02258","doi":"https://doi.org/10.48550/arxiv.2404.02258","is_oa":true,"landing_page_url":"https://doi.org/10.48550/arxiv.2404.02258","pdf_url":null,"source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":"2331-8422","issn":["2331-8422"],"is_oa":true,"is_in_doaj":false,"is_core":false,"listed_in":[],"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":null,"is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"Preprint"}],"best_oa_location":{"id":"pmh:oai:arXiv.org:2404.02258","doi":null,"is_oa":true,"landing_page_url":"http://arxiv.org/abs/2404.02258","pdf_url":"https://arxiv.org/pdf/2404.02258","source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":"2331-8422","issn":["2331-8422"],"is_oa":true,"is_in_doaj":false,"is_core":false,"listed_in":[],"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":"submittedVersion","is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"text"},"sustainable_development_goals":[],"study_designs":[],"awards":[],"funders":[],"has_content":{"pdf":true,"grobid_xml":true},"content_urls":{"pdf":"https://content.openalex.org/works/W4393967974.pdf","grobid_xml":"https://content.openalex.org/works/W4393967974.grobid-xml"},"referenced_works_count":0,"referenced_works":[],"related_works":["https://openalex.org/W2995983533","https://openalex.org/W4416789520","https://openalex.org/W7172381945","https://openalex.org/W7134194401","https://openalex.org/W7128498855","https://openalex.org/W4404555061","https://openalex.org/W4404782163","https://openalex.org/W7171835909","https://openalex.org/W4387806199","https://openalex.org/W4406939014"],"abstract_inverted_index":{"Transformer-based":[0],"language":[1],"models":[2,151,166],"spread":[3],"FLOPs":[4,22,123,184],"uniformly":[5],"across":[6,40,125],"input":[7],"sequences.":[8],"In":[9],"this":[10,92,119,154],"work":[11],"we":[12],"demonstrate":[13],"that":[14,58],"transformers":[15],"can":[16,59,121,189],"instead":[17],"learn":[18,156],"to":[19,25,73,157,176,195],"dynamically":[20,158],"allocate":[21,159],"(or":[23],"compute)":[24],"specific":[26],"positions":[27],"in":[28,61,138,153],"a":[29,47,68,82,90,96,180],"sequence,":[30],"optimising":[31],"the":[32,35,41,53,62,79,111,114,126,146,183],"allocation":[33],"along":[34],"sequence":[36],"for":[37,170],"different":[38],"layers":[39],"model":[42,129],"depth.":[43],"Our":[44],"method":[45,120],"enforces":[46],"total":[48],"compute":[49,133],"budget":[50],"by":[51,78],"capping":[52],"number":[54],"of":[55,113,182,192],"tokens":[56,72,116],"($k$)":[57],"participate":[60],"self-attention":[63],"and":[64,128,143,173,188],"MLP":[65],"computations":[66],"at":[67,145],"given":[69],"layer.":[70],"The":[71],"be":[74,190],"processed":[75],"are":[76,117],"determined":[77],"network":[80],"using":[81],"top-$k$":[83],"routing":[84],"mechanism.":[85],"Since":[86],"$k$":[87,115],"is":[88,135],"defined":[89],"priori,":[91],"simple":[93],"procedure":[94],"uses":[95],"static":[97],"computation":[98,107],"graph":[99],"with":[100],"known":[101],"tensor":[102],"sizes,":[103],"unlike":[104],"other":[105],"conditional":[106],"techniques.":[108],"Nevertheless,":[109],"since":[110],"identities":[112],"fluid,":[118],"expend":[122],"non-uniformly":[124],"time":[127],"depth":[130],"dimensions.":[131],"Thus,":[132],"expenditure":[134],"entirely":[136],"predictable":[137],"sum":[139],"total,":[140],"but":[141,178],"dynamic":[142],"context-sensitive":[144],"token-level.":[147],"Not":[148],"only":[149],"do":[150,162],"trained":[152],"way":[155],"compute,":[160],"they":[161],"so":[163],"efficiently.":[164],"These":[165],"match":[167],"baseline":[168],"performance":[169],"equivalent":[171],"FLOPS":[172],"wall-clock":[174],"times":[175],"train,":[177],"require":[179],"fraction":[181],"per":[185],"forward":[186],"pass,":[187],"upwards":[191],"50\\%":[193],"faster":[194],"step":[196],"during":[197],"post-training":[198],"sampling.":[199]},"counts_by_year":[{"year":2026,"cited_by_count":1},{"year":2025,"cited_by_count":5},{"year":2024,"cited_by_count":3}],"updated_date":"2026-11-02T09:45:03.842371","created_date":"2024-04-05T00:00:00"}