<!DOCTYPE html><html lang="en">
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8">
<title>[2403.10616] DiPaCo: Distributed Path Composition</title><meta property="og:description" content="Progress in machine learning (ML) has been fueled by scaling neural network models. This scaling has been enabled by ever more heroic feats of engineering, necessary for accommodating ML approaches that require high ba…">
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="twitter:card" content="summary">
<meta name="twitter:title" content="DiPaCo: Distributed Path Composition">
<meta name="twitter:image:src" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta name="twitter:image:alt" content="ar5iv logo">
<meta property="og:title" content="DiPaCo: Distributed Path Composition">
<meta property="og:site_name" content="ar5iv">
<meta property="og:image" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta property="og:type" content="article">
<meta property="og:url" content="https://ar5iv.labs.arxiv.org/html/2403.10616">

<!--Generated on Fri Apr  5 17:24:29 2024 by LaTeXML (version 0.8.8) http://dlmf.nist.gov/LaTeXML/.-->
<meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">

<script>
  function detectColorScheme(){
    var theme="light";
    var current_theme = localStorage.getItem("ar5iv_theme");
    if(current_theme){
      if(current_theme == "dark"){
        theme = "dark";
      } }
    else if(!window.matchMedia) { return false; }
    else if(window.matchMedia("(prefers-color-scheme: dark)").matches) {
      theme = "dark"; }
    if (theme=="dark") {
      document.documentElement.setAttribute("data-theme", "dark");
    } else {
      document.documentElement.setAttribute("data-theme", "light"); } }

  detectColorScheme();

  function toggleColorScheme(){
    var current_theme = localStorage.getItem("ar5iv_theme");
    if (current_theme) {
      if (current_theme == "light") {
        localStorage.setItem("ar5iv_theme", "dark"); }
      else {
        localStorage.setItem("ar5iv_theme", "light"); } }
    else {
        localStorage.setItem("ar5iv_theme", "dark"); }
    detectColorScheme(); }
</script>
<link media="all" rel="stylesheet" href="/assets/ar5iv-fonts.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv-site.0.2.2.css">
</head>
<body>
<div class="ltx_page_main">
<div class="ltx_page_content">
<article class="ltx_document ltx_authors_1line">
<div id="p1" class="ltx_para">
<span id="p1.1" class="ltx_ERROR undefined">\pdftrailerid</span>
<p id="p1.2" class="ltx_p">redacted

 

  











<span id="p1.2.1" class="ltx_ERROR undefined">\correspondingauthor</span>douillard@google.com

<span id="p1.2.2" class="ltx_ERROR undefined">\reportnumber</span> 












</p>
</div>
<h1 class="ltx_title ltx_title_document">DiPaCo: Distributed Path Composition</h1>
<div class="ltx_authors">
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Arthur Douillard
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal core contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Qixuan Feng
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal core contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Andrei A. Rusu
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal core contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Adhiguna Kuncoro
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Yani Donchev
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Rachita Chhaparia
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Ionel Gog
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Marc’Aurelio Ranzato
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal leading contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Jiajun Shen
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal leading contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Arthur Szlam
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">Equal leading contributions
</span>
<span class="ltx_contact ltx_role_affiliation">Google DeepMind
</span></span></span>
</div>

<div class="ltx_abstract">
<h6 class="ltx_title ltx_title_abstract">Abstract</h6>
<p id="id1.id1" class="ltx_p">Progress in machine learning (ML) has been fueled by scaling neural network models. This scaling has been enabled by ever more heroic feats of engineering, necessary for accommodating ML approaches that require high bandwidth communication between devices working in parallel.
In this work, we propose a co-designed modular architecture and training approach for ML models, dubbed DIstributed PAth COmposition (DiPaCo). During training, DiPaCo distributes computation by paths through a set of shared modules. Together with a Local-SGD inspired optimization (DiLoCo) that keeps modules in sync with drastically reduced communication,
Our approach facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions. At inference time, only a single path needs to be executed for each input, without the need for any model compression. We
consider this approach as a first prototype towards a new paradigm of large-scale learning, one that is less synchronous and more modular.
</p>
<p id="id2.id2" class="ltx_p">Our experiments on the widely used C4 benchmark show that, for the same amount of training steps but less wall-clock time, DiPaCo exceeds the performance of a 1 billion-parameter dense transformer language model by choosing one of 256 possible paths, each with a size of 150 million parameters.</p>
</div>
<div class="ltx_classification">
<h6 class="ltx_title ltx_title_classification">keywords: </h6>modularity, large-scale, mixture of experts, language modeling, distributed learning
</div>
<section id="S1" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">1 </span>Introduction</h2>

<div id="S1.p1" class="ltx_para">
<p id="S1.p1.1" class="ltx_p">Progress in machine learning and AI has been driven by spending more FLOPs on larger neural network models, trained on bigger data sets. This scaling has been accomplished via data and model parallelism <cite class="ltx_cite ltx_citemacro_citep">(Dean et al., <a href="#bib.bib11" title="" class="ltx_ref">2012</a>)</cite> and pipelining <cite class="ltx_cite ltx_citemacro_citep">(Narayanan et al., <a href="#bib.bib28" title="" class="ltx_ref">2020</a>)</cite> to distribute computation, enabling the concurrent use of a large number of devices <cite class="ltx_cite ltx_citemacro_citep">(OpenAI et al., <a href="#bib.bib29" title="" class="ltx_ref">2023</a>; Touvron et al., <a href="#bib.bib41" title="" class="ltx_ref">2023</a>; Anil et al., <a href="#bib.bib2" title="" class="ltx_ref">2023</a>)</cite>. Although model architectures <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib25" title="" class="ltx_ref">2021</a>; OpenAI et al., <a href="#bib.bib29" title="" class="ltx_ref">2023</a>)</cite> have also been used to allow computational parallelism, and optimization procedures to prefer larger batches <cite class="ltx_cite ltx_citemacro_citep">(Goyal et al., <a href="#bib.bib18" title="" class="ltx_ref">2017</a>)</cite> (again allowing more data parallelism), the current training paradigm has not fundamentally changed model architecture or optimization procedure to facilitate distributed training.
State of the art models are still essentially monoliths, and their optimization requires exchanges of parameters, gradients, and activations at every step of the learning process.</p>
</div>
<div id="S1.p2" class="ltx_para">
<p id="S1.p2.1" class="ltx_p">This approach incurs engineering and infrastructure challenges associated with provisioning and managing the large number of tightly interconnected devices required for
the lengthy training process. The training process itself is often restarted for each new model release, essentially discarding much of the computation for training the last model. Moreover, training monoliths incurs human-organizational challenges, as it is difficult to localize the effects to the final model of changes to any step in the process (beyond data preparation).
In particular, it is difficult to leverage the potential of the greater ML community, rather than a single organization <cite class="ltx_cite ltx_citemacro_citep">(Raffel, <a href="#bib.bib34" title="" class="ltx_ref">2023b</a>)</cite>.
Because of these challenges, it may become more difficult to continue to scale with the current approach.</p>
</div>
<div id="S1.p3" class="ltx_para">
<p id="S1.p3.1" class="ltx_p">As in <cite class="ltx_cite ltx_citemacro_citep">(Ryabinin and Gusev, <a href="#bib.bib38" title="" class="ltx_ref">2020</a>; Barham et al., <a href="#bib.bib4" title="" class="ltx_ref">2022</a>; Borzunov et al., <a href="#bib.bib6" title="" class="ltx_ref">2022</a>; Raffel, <a href="#bib.bib34" title="" class="ltx_ref">2023b</a>)</cite> and depicted in <a href="#S1.F1" title="Figure 1 ‣ 1 Introduction ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 1</span></a>, we envision an alternative approach where models are scalable both in terms of ability to ingest and train on large amounts of data and in terms of enabling collaboration with many contributors, and can be continuously updated and expanded, by virtue of modular design. Modular ML may have many other benefits, see <cite class="ltx_cite ltx_citemacro_cite">Pfeiffer et al. (<a href="#bib.bib31" title="" class="ltx_ref">2023</a>)</cite> for further discussion and references.</p>
</div>
<div id="S1.p4" class="ltx_para">
<p id="S1.p4.1" class="ltx_p">In this work, we take a step towards this scalable modular ML paradigm,
proposing an architecture and training algorithm, DIstributed PAths COmposition (DiPaCo).
DiPaCo’s architecture and optimization have been co-designed to reduce communication and enable better scaling. The high level idea is to <span id="S1.p4.1.1" class="ltx_text ltx_font_italic">distribute computation by path</span>; here, a “path” means a sequence of modules that define an input-output function.
Paths are small relative to the entire model, thus requiring only a handful of tightly connected devices to train or evaluate.
During both training and deployment, a query is routed to a replica of a path rather than a replica of the whole model; in other words, the DiPaCo architecture is sparsely activated. We give a detailed exposition of DiPaCo in <a href="#S2" title="2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">section 2</span></a>, and the infrastructure we used to implement it in <a href="#S3" title="3 Infrastructure ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">section 3</span></a>.</p>
</div>
<div id="S1.p5" class="ltx_para">
<p id="S1.p5.1" class="ltx_p">In <a href="#S4" title="4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">section 4</span></a> we demonstrate the feasibility of DiPaCo by training a language model on the C4 dataset <cite class="ltx_cite ltx_citemacro_citep">(Raffel et al., <a href="#bib.bib35" title="" class="ltx_ref">2020</a>)</cite> with paths of size 150 million parameters, matching the performance in terms of validation perplexity of a 1.3 billion model, but with 45% less wall clock training time. While the dense 1.3B system required the use of all co-located devices, DiPaCo uses 256 islands of compute, each of which is one-eighth the number of devices used to train the baseline. Neither during train or evaluation time is it necessary to co-locate the full DiPaCo model.</p>
</div>
<figure id="S1.F1" class="ltx_figure"><img src="/html/2403.10616/assets/x1.png" id="S1.F1.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="664" height="406" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S1.F1.8.3.1" class="ltx_text" style="font-size:90%;">Figure 1</span>: </span><span id="S1.F1.4.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Long-term Goal<span id="S1.F1.4.2.2" class="ltx_text ltx_font_medium">: Ultimately, we envision a modular network where different components, <span id="S1.F1.4.2.2.1" class="ltx_text ltx_font_italic">paths</span> <math id="S1.F1.3.1.1.m1.1" class="ltx_Math" alttext="\pi_{i}" display="inline"><semantics id="S1.F1.3.1.1.m1.1b"><msub id="S1.F1.3.1.1.m1.1.1" xref="S1.F1.3.1.1.m1.1.1.cmml"><mi id="S1.F1.3.1.1.m1.1.1.2" xref="S1.F1.3.1.1.m1.1.1.2.cmml">π</mi><mi id="S1.F1.3.1.1.m1.1.1.3" xref="S1.F1.3.1.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S1.F1.3.1.1.m1.1c"><apply id="S1.F1.3.1.1.m1.1.1.cmml" xref="S1.F1.3.1.1.m1.1.1"><csymbol cd="ambiguous" id="S1.F1.3.1.1.m1.1.1.1.cmml" xref="S1.F1.3.1.1.m1.1.1">subscript</csymbol><ci id="S1.F1.3.1.1.m1.1.1.2.cmml" xref="S1.F1.3.1.1.m1.1.1.2">𝜋</ci><ci id="S1.F1.3.1.1.m1.1.1.3.cmml" xref="S1.F1.3.1.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S1.F1.3.1.1.m1.1d">\pi_{i}</annotation></semantics></math>, are optimized for different tasks, <math id="S1.F1.4.2.2.m2.1" class="ltx_Math" alttext="\mathcal{D}_{j}" display="inline"><semantics id="S1.F1.4.2.2.m2.1b"><msub id="S1.F1.4.2.2.m2.1.1" xref="S1.F1.4.2.2.m2.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S1.F1.4.2.2.m2.1.1.2" xref="S1.F1.4.2.2.m2.1.1.2.cmml">𝒟</mi><mi id="S1.F1.4.2.2.m2.1.1.3" xref="S1.F1.4.2.2.m2.1.1.3.cmml">j</mi></msub><annotation-xml encoding="MathML-Content" id="S1.F1.4.2.2.m2.1c"><apply id="S1.F1.4.2.2.m2.1.1.cmml" xref="S1.F1.4.2.2.m2.1.1"><csymbol cd="ambiguous" id="S1.F1.4.2.2.m2.1.1.1.cmml" xref="S1.F1.4.2.2.m2.1.1">subscript</csymbol><ci id="S1.F1.4.2.2.m2.1.1.2.cmml" xref="S1.F1.4.2.2.m2.1.1.2">𝒟</ci><ci id="S1.F1.4.2.2.m2.1.1.3.cmml" xref="S1.F1.4.2.2.m2.1.1.3">𝑗</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S1.F1.4.2.2.m2.1d">\mathcal{D}_{j}</annotation></semantics></math>, each designed by different researchers. The paths, trained on any available hardware type, communicate infrequently across the world, exchanging useful information and enabling new forms of composition.</span></span></figcaption>
</figure>
</section>
<section id="S2" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">2 </span>Approach</h2>

<div id="S2.p1" class="ltx_para">
<p id="S2.p1.1" class="ltx_p">In this section we give a detailed description of DiPaCo. We start by describing the assumptions and setting in which we will work.</p>
</div>
<section id="S2.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.1 </span>Setting</h3>

<div id="S2.SS1.p1" class="ltx_para">
<p id="S2.SS1.p1.1" class="ltx_p">Our goal in this work is to demonstrate scalable ML models in the setting of many smaller islands of compute, as opposed to one tightly connected compute landmass. Accordingly, we assume: 
<br class="ltx_break">1. Training compute (FLOPs) is relatively cheap
<br class="ltx_break">2. Communication is relatively expensive
<br class="ltx_break"></p>
</div>
<div id="S2.SS1.p2" class="ltx_para">
<p id="S2.SS1.p2.1" class="ltx_p">These assumptions are <span id="S2.SS1.p2.1.1" class="ltx_text ltx_font_italic">not</span> realistic in the current ML training paradigm.
As discussed above, devices are usually co-located and run by a single organization, FLOPs are expensive, constrained by the purchase and installation of many accelerators, but communication between devices is relatively cheap because of the co-location.</p>
</div>
<div id="S2.SS1.p3" class="ltx_para">
<p id="S2.SS1.p3.1" class="ltx_p">Nevertheless, in our view,
these assumptions
may be realistic in the near future, if our compute requirements (or model sizes) grow beyond what can be reasonably co-located, and so communication costs become relatively more expensive.
Conversely, progress in distributed training of ML models may allow simpler infrastructure build-outs, leading eventually to more available compute. As it is, infrastructure is designed around the standard approach to training large monoliths; and ML models are architected to take advantage of the current infrastructure and training approaches. This feedback loop may be leading the community to a spurious local minimum where compute is more constrained than it needs to be.</p>
</div>
<div id="S2.SS1.p4" class="ltx_para">
<p id="S2.SS1.p4.1" class="ltx_p">In addition (and closely related) to the two assumptions above, we also assume that both during train and evaluation, we cannot instantiate models as large as we would like to have on any single compute island.</p>
</div>
<div id="S2.SS1.p5" class="ltx_para">
<p id="S2.SS1.p5.1" class="ltx_p">We will work in the context of language modeling (LM).
Based on the above assumptions, we will evaluate models by evaluation perplexity (PPL) against wall-clock time, keeping evaluation complexity in mind . We choose this metric because it encourages exploring how to distribute model training, and more generally, points to in our opinion plausible new paradigms for ML.</p>
</div>
</section>
<section id="S2.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.2 </span>Overview of System</h3>

<div id="S2.SS2.p1" class="ltx_para">
<p id="S2.SS2.p1.1" class="ltx_p">The core idea of DiPaCo is to train a sparsely-activated modular system where data and computation are distributed by the choice of path through the modules. There are two key ideas to making this work:</p>
</div>
<section id="S2.SS2.SSS0.Px1" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Coarse Routing:</h5>

<div id="S2.SS2.SSS0.Px1.p1" class="ltx_para">
<p id="S2.SS2.SSS0.Px1.p1.1" class="ltx_p">Sparsely routed Mixture of Experts (MoE) have shown great results in language modeling <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib25" title="" class="ltx_ref">2021</a>)</cite>. The standard approach in a transformer MoE LM is to make a routing decision at each token, based on the feature at each routed layer. In contrast, in this work, during training we will route once per document and offline, as in  <cite class="ltx_cite ltx_citemacro_citep">(Gross et al., <a href="#bib.bib19" title="" class="ltx_ref">2017</a>; Gururangan et al., <a href="#bib.bib20" title="" class="ltx_ref">2023</a>)</cite>.</p>
</div>
<div id="S2.SS2.SSS0.Px1.p2" class="ltx_para">
<p id="S2.SS2.SSS0.Px1.p2.1" class="ltx_p">Routing once per document allows batching computation across all tokens of a sequence, without the need to swap modules in and out as a sequence is processed. This in turn allows parameters to be distributed across distant workers.
In addition, instead of learning the router along with the model, we will compute routing decisions
<span id="S2.SS2.SSS0.Px1.p2.1.1" class="ltx_text ltx_font_italic">offline</span>. This enables pre-sharding data by path, before the start of training. This is critical to distribute training across paths, as each worker can now train a path by processing its own shard of data, using its own set of parameters. In <a href="#S2.SS4" title="2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.4</span></a>, we will describe in detail how we approach coarse routing.</p>
</div>
</section>
<section id="S2.SS2.SSS0.Px2" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">DiLoCo to keep modules in sync</h5>

<div id="S2.SS2.SSS0.Px2.p1" class="ltx_para">
<p id="S2.SS2.SSS0.Px2.p1.4" class="ltx_p">Paths cannot be trained completely independently, because some modules might be shared across multiple paths. To support module sharing across paths, we use DiLoCo <cite class="ltx_cite ltx_citemacro_citep">(Douillard et al., <a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite> for low communication data parallelism, see <a href="#S2.SS5" title="2.5 DiLoCo: Review ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.5</span></a>. If there are <math id="S2.SS2.SSS0.Px2.p1.1.m1.1" class="ltx_Math" alttext="P" display="inline"><semantics id="S2.SS2.SSS0.Px2.p1.1.m1.1a"><mi id="S2.SS2.SSS0.Px2.p1.1.m1.1.1" xref="S2.SS2.SSS0.Px2.p1.1.m1.1.1.cmml">P</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px2.p1.1.m1.1b"><ci id="S2.SS2.SSS0.Px2.p1.1.m1.1.1.cmml" xref="S2.SS2.SSS0.Px2.p1.1.m1.1.1">𝑃</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px2.p1.1.m1.1c">P</annotation></semantics></math> paths (assigned to <math id="S2.SS2.SSS0.Px2.p1.2.m2.1" class="ltx_Math" alttext="P" display="inline"><semantics id="S2.SS2.SSS0.Px2.p1.2.m2.1a"><mi id="S2.SS2.SSS0.Px2.p1.2.m2.1.1" xref="S2.SS2.SSS0.Px2.p1.2.m2.1.1.cmml">P</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px2.p1.2.m2.1b"><ci id="S2.SS2.SSS0.Px2.p1.2.m2.1.1.cmml" xref="S2.SS2.SSS0.Px2.p1.2.m2.1.1">𝑃</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px2.p1.2.m2.1c">P</annotation></semantics></math> workers) that share module <math id="S2.SS2.SSS0.Px2.p1.3.m3.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS2.SSS0.Px2.p1.3.m3.1a"><mi id="S2.SS2.SSS0.Px2.p1.3.m3.1.1" xref="S2.SS2.SSS0.Px2.p1.3.m3.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px2.p1.3.m3.1b"><ci id="S2.SS2.SSS0.Px2.p1.3.m3.1.1.cmml" xref="S2.SS2.SSS0.Px2.p1.3.m3.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px2.p1.3.m3.1c">i</annotation></semantics></math>, each corresponding worker performs SGD on its own shard of data, and every <span id="S2.SS2.SSS0.Px2.p1.4.1" class="ltx_text ltx_font_italic">few hundred</span> steps, workers average the difference of the parameters of module <math id="S2.SS2.SSS0.Px2.p1.4.m4.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS2.SSS0.Px2.p1.4.m4.1a"><mi id="S2.SS2.SSS0.Px2.p1.4.m4.1.1" xref="S2.SS2.SSS0.Px2.p1.4.m4.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px2.p1.4.m4.1b"><ci id="S2.SS2.SSS0.Px2.p1.4.m4.1.1.cmml" xref="S2.SS2.SSS0.Px2.p1.4.m4.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px2.p1.4.m4.1c">i</annotation></semantics></math> before and after the local SGD phase. These averages are then used to update a global parameter vector which is then re-distributed across the workers to sync them, a procedure described in <a href="#S2.SS6" title="2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.6</span></a>.</p>
</div>
<div id="S2.SS2.SSS0.Px2.p2" class="ltx_para">
<p id="S2.SS2.SSS0.Px2.p2.1" class="ltx_p">With these two choices,
neither at training nor at test time does the entire network (collection of paths) need to be materialized together, as shown in <a href="#S1.F1" title="Figure 1 ‣ 1 Introduction ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 1</span></a>.</p>
</div>
</section>
</section>
<section id="S2.SS3" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.3 </span>Notation</h3>

<div id="S2.SS3.p1" class="ltx_para">
<p id="S2.SS3.p1.9" class="ltx_p">In this section we introduce the basic notation used throughout this work.
We assume that we have a base model architecture with parameters <math id="S2.SS3.p1.1.m1.1" class="ltx_Math" alttext="\theta" display="inline"><semantics id="S2.SS3.p1.1.m1.1a"><mi id="S2.SS3.p1.1.m1.1.1" xref="S2.SS3.p1.1.m1.1.1.cmml">θ</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.1.m1.1b"><ci id="S2.SS3.p1.1.m1.1.1.cmml" xref="S2.SS3.p1.1.m1.1.1">𝜃</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.1.m1.1c">\theta</annotation></semantics></math>. We take a partition of the parameter indices into <math id="S2.SS3.p1.2.m2.1" class="ltx_Math" alttext="L" display="inline"><semantics id="S2.SS3.p1.2.m2.1a"><mi id="S2.SS3.p1.2.m2.1.1" xref="S2.SS3.p1.2.m2.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.2.m2.1b"><ci id="S2.SS3.p1.2.m2.1.1.cmml" xref="S2.SS3.p1.2.m2.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.2.m2.1c">L</annotation></semantics></math> subsets <math id="S2.SS3.p1.3.m3.5" class="ltx_Math" alttext="B_{l},l\in[1,\ldots,L]" display="inline"><semantics id="S2.SS3.p1.3.m3.5a"><mrow id="S2.SS3.p1.3.m3.5.5" xref="S2.SS3.p1.3.m3.5.5.cmml"><mrow id="S2.SS3.p1.3.m3.5.5.1.1" xref="S2.SS3.p1.3.m3.5.5.1.2.cmml"><msub id="S2.SS3.p1.3.m3.5.5.1.1.1" xref="S2.SS3.p1.3.m3.5.5.1.1.1.cmml"><mi id="S2.SS3.p1.3.m3.5.5.1.1.1.2" xref="S2.SS3.p1.3.m3.5.5.1.1.1.2.cmml">B</mi><mi id="S2.SS3.p1.3.m3.5.5.1.1.1.3" xref="S2.SS3.p1.3.m3.5.5.1.1.1.3.cmml">l</mi></msub><mo id="S2.SS3.p1.3.m3.5.5.1.1.2" xref="S2.SS3.p1.3.m3.5.5.1.2.cmml">,</mo><mi id="S2.SS3.p1.3.m3.4.4" xref="S2.SS3.p1.3.m3.4.4.cmml">l</mi></mrow><mo id="S2.SS3.p1.3.m3.5.5.2" xref="S2.SS3.p1.3.m3.5.5.2.cmml">∈</mo><mrow id="S2.SS3.p1.3.m3.5.5.3.2" xref="S2.SS3.p1.3.m3.5.5.3.1.cmml"><mo stretchy="false" id="S2.SS3.p1.3.m3.5.5.3.2.1" xref="S2.SS3.p1.3.m3.5.5.3.1.cmml">[</mo><mn id="S2.SS3.p1.3.m3.1.1" xref="S2.SS3.p1.3.m3.1.1.cmml">1</mn><mo id="S2.SS3.p1.3.m3.5.5.3.2.2" xref="S2.SS3.p1.3.m3.5.5.3.1.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p1.3.m3.2.2" xref="S2.SS3.p1.3.m3.2.2.cmml">…</mi><mo id="S2.SS3.p1.3.m3.5.5.3.2.3" xref="S2.SS3.p1.3.m3.5.5.3.1.cmml">,</mo><mi id="S2.SS3.p1.3.m3.3.3" xref="S2.SS3.p1.3.m3.3.3.cmml">L</mi><mo stretchy="false" id="S2.SS3.p1.3.m3.5.5.3.2.4" xref="S2.SS3.p1.3.m3.5.5.3.1.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.3.m3.5b"><apply id="S2.SS3.p1.3.m3.5.5.cmml" xref="S2.SS3.p1.3.m3.5.5"><in id="S2.SS3.p1.3.m3.5.5.2.cmml" xref="S2.SS3.p1.3.m3.5.5.2"></in><list id="S2.SS3.p1.3.m3.5.5.1.2.cmml" xref="S2.SS3.p1.3.m3.5.5.1.1"><apply id="S2.SS3.p1.3.m3.5.5.1.1.1.cmml" xref="S2.SS3.p1.3.m3.5.5.1.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.3.m3.5.5.1.1.1.1.cmml" xref="S2.SS3.p1.3.m3.5.5.1.1.1">subscript</csymbol><ci id="S2.SS3.p1.3.m3.5.5.1.1.1.2.cmml" xref="S2.SS3.p1.3.m3.5.5.1.1.1.2">𝐵</ci><ci id="S2.SS3.p1.3.m3.5.5.1.1.1.3.cmml" xref="S2.SS3.p1.3.m3.5.5.1.1.1.3">𝑙</ci></apply><ci id="S2.SS3.p1.3.m3.4.4.cmml" xref="S2.SS3.p1.3.m3.4.4">𝑙</ci></list><list id="S2.SS3.p1.3.m3.5.5.3.1.cmml" xref="S2.SS3.p1.3.m3.5.5.3.2"><cn type="integer" id="S2.SS3.p1.3.m3.1.1.cmml" xref="S2.SS3.p1.3.m3.1.1">1</cn><ci id="S2.SS3.p1.3.m3.2.2.cmml" xref="S2.SS3.p1.3.m3.2.2">…</ci><ci id="S2.SS3.p1.3.m3.3.3.cmml" xref="S2.SS3.p1.3.m3.3.3">𝐿</ci></list></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.3.m3.5c">B_{l},l\in[1,\ldots,L]</annotation></semantics></math>.
For each <math id="S2.SS3.p1.4.m4.1" class="ltx_Math" alttext="l" display="inline"><semantics id="S2.SS3.p1.4.m4.1a"><mi id="S2.SS3.p1.4.m4.1.1" xref="S2.SS3.p1.4.m4.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.4.m4.1b"><ci id="S2.SS3.p1.4.m4.1.1.cmml" xref="S2.SS3.p1.4.m4.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.4.m4.1c">l</annotation></semantics></math>, we choose a number <math id="S2.SS3.p1.5.m5.1" class="ltx_Math" alttext="K_{l}" display="inline"><semantics id="S2.SS3.p1.5.m5.1a"><msub id="S2.SS3.p1.5.m5.1.1" xref="S2.SS3.p1.5.m5.1.1.cmml"><mi id="S2.SS3.p1.5.m5.1.1.2" xref="S2.SS3.p1.5.m5.1.1.2.cmml">K</mi><mi id="S2.SS3.p1.5.m5.1.1.3" xref="S2.SS3.p1.5.m5.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.5.m5.1b"><apply id="S2.SS3.p1.5.m5.1.1.cmml" xref="S2.SS3.p1.5.m5.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.5.m5.1.1.1.cmml" xref="S2.SS3.p1.5.m5.1.1">subscript</csymbol><ci id="S2.SS3.p1.5.m5.1.1.2.cmml" xref="S2.SS3.p1.5.m5.1.1.2">𝐾</ci><ci id="S2.SS3.p1.5.m5.1.1.3.cmml" xref="S2.SS3.p1.5.m5.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.5.m5.1c">K_{l}</annotation></semantics></math> that represents the number of possible choices for the parameters in <math id="S2.SS3.p1.6.m6.1" class="ltx_Math" alttext="B_{l}" display="inline"><semantics id="S2.SS3.p1.6.m6.1a"><msub id="S2.SS3.p1.6.m6.1.1" xref="S2.SS3.p1.6.m6.1.1.cmml"><mi id="S2.SS3.p1.6.m6.1.1.2" xref="S2.SS3.p1.6.m6.1.1.2.cmml">B</mi><mi id="S2.SS3.p1.6.m6.1.1.3" xref="S2.SS3.p1.6.m6.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.6.m6.1b"><apply id="S2.SS3.p1.6.m6.1.1.cmml" xref="S2.SS3.p1.6.m6.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.6.m6.1.1.1.cmml" xref="S2.SS3.p1.6.m6.1.1">subscript</csymbol><ci id="S2.SS3.p1.6.m6.1.1.2.cmml" xref="S2.SS3.p1.6.m6.1.1.2">𝐵</ci><ci id="S2.SS3.p1.6.m6.1.1.3.cmml" xref="S2.SS3.p1.6.m6.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.6.m6.1c">B_{l}</annotation></semantics></math>. We will call a set of parameters associated to a <math id="S2.SS3.p1.7.m7.1" class="ltx_Math" alttext="B_{l}" display="inline"><semantics id="S2.SS3.p1.7.m7.1a"><msub id="S2.SS3.p1.7.m7.1.1" xref="S2.SS3.p1.7.m7.1.1.cmml"><mi id="S2.SS3.p1.7.m7.1.1.2" xref="S2.SS3.p1.7.m7.1.1.2.cmml">B</mi><mi id="S2.SS3.p1.7.m7.1.1.3" xref="S2.SS3.p1.7.m7.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.7.m7.1b"><apply id="S2.SS3.p1.7.m7.1.1.cmml" xref="S2.SS3.p1.7.m7.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.7.m7.1.1.1.cmml" xref="S2.SS3.p1.7.m7.1.1">subscript</csymbol><ci id="S2.SS3.p1.7.m7.1.1.2.cmml" xref="S2.SS3.p1.7.m7.1.1.2">𝐵</ci><ci id="S2.SS3.p1.7.m7.1.1.3.cmml" xref="S2.SS3.p1.7.m7.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.7.m7.1c">B_{l}</annotation></semantics></math> a “module” or, as in  <cite class="ltx_cite ltx_citemacro_citep">(Jacobs et al., <a href="#bib.bib21" title="" class="ltx_ref">1991</a>; Jordan and Jacobs, <a href="#bib.bib22" title="" class="ltx_ref">1994</a>)</cite>, an “expert”. In this terminology, <math id="S2.SS3.p1.8.m8.1" class="ltx_Math" alttext="K_{l}" display="inline"><semantics id="S2.SS3.p1.8.m8.1a"><msub id="S2.SS3.p1.8.m8.1.1" xref="S2.SS3.p1.8.m8.1.1.cmml"><mi id="S2.SS3.p1.8.m8.1.1.2" xref="S2.SS3.p1.8.m8.1.1.2.cmml">K</mi><mi id="S2.SS3.p1.8.m8.1.1.3" xref="S2.SS3.p1.8.m8.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.8.m8.1b"><apply id="S2.SS3.p1.8.m8.1.1.cmml" xref="S2.SS3.p1.8.m8.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.8.m8.1.1.1.cmml" xref="S2.SS3.p1.8.m8.1.1">subscript</csymbol><ci id="S2.SS3.p1.8.m8.1.1.2.cmml" xref="S2.SS3.p1.8.m8.1.1.2">𝐾</ci><ci id="S2.SS3.p1.8.m8.1.1.3.cmml" xref="S2.SS3.p1.8.m8.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.8.m8.1c">K_{l}</annotation></semantics></math> is the number of distinct modules associated to <math id="S2.SS3.p1.9.m9.1" class="ltx_Math" alttext="B_{l}" display="inline"><semantics id="S2.SS3.p1.9.m9.1a"><msub id="S2.SS3.p1.9.m9.1.1" xref="S2.SS3.p1.9.m9.1.1.cmml"><mi id="S2.SS3.p1.9.m9.1.1.2" xref="S2.SS3.p1.9.m9.1.1.2.cmml">B</mi><mi id="S2.SS3.p1.9.m9.1.1.3" xref="S2.SS3.p1.9.m9.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p1.9.m9.1b"><apply id="S2.SS3.p1.9.m9.1.1.cmml" xref="S2.SS3.p1.9.m9.1.1"><csymbol cd="ambiguous" id="S2.SS3.p1.9.m9.1.1.1.cmml" xref="S2.SS3.p1.9.m9.1.1">subscript</csymbol><ci id="S2.SS3.p1.9.m9.1.1.2.cmml" xref="S2.SS3.p1.9.m9.1.1.2">𝐵</ci><ci id="S2.SS3.p1.9.m9.1.1.3.cmml" xref="S2.SS3.p1.9.m9.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p1.9.m9.1c">B_{l}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS3.p2" class="ltx_para">
<p id="S2.SS3.p2.14" class="ltx_p">For a simple example, consider a four-layer fully connected network <math id="S2.SS3.p2.1.m1.1" class="ltx_Math" alttext="A" display="inline"><semantics id="S2.SS3.p2.1.m1.1a"><mi id="S2.SS3.p2.1.m1.1.1" xref="S2.SS3.p2.1.m1.1.1.cmml">A</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.1.m1.1b"><ci id="S2.SS3.p2.1.m1.1.1.cmml" xref="S2.SS3.p2.1.m1.1.1">𝐴</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.1.m1.1c">A</annotation></semantics></math>. We might take the first two layers to be <math id="S2.SS3.p2.2.m2.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS3.p2.2.m2.1a"><msub id="S2.SS3.p2.2.m2.1.1" xref="S2.SS3.p2.2.m2.1.1.cmml"><mi id="S2.SS3.p2.2.m2.1.1.2" xref="S2.SS3.p2.2.m2.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.2.m2.1.1.3" xref="S2.SS3.p2.2.m2.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.2.m2.1b"><apply id="S2.SS3.p2.2.m2.1.1.cmml" xref="S2.SS3.p2.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.2.m2.1.1.1.cmml" xref="S2.SS3.p2.2.m2.1.1">subscript</csymbol><ci id="S2.SS3.p2.2.m2.1.1.2.cmml" xref="S2.SS3.p2.2.m2.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.2.m2.1.1.3.cmml" xref="S2.SS3.p2.2.m2.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.2.m2.1c">B_{1}</annotation></semantics></math> and the third and fourth layers to be <math id="S2.SS3.p2.3.m3.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS3.p2.3.m3.1a"><msub id="S2.SS3.p2.3.m3.1.1" xref="S2.SS3.p2.3.m3.1.1.cmml"><mi id="S2.SS3.p2.3.m3.1.1.2" xref="S2.SS3.p2.3.m3.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.3.m3.1.1.3" xref="S2.SS3.p2.3.m3.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.3.m3.1b"><apply id="S2.SS3.p2.3.m3.1.1.cmml" xref="S2.SS3.p2.3.m3.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.3.m3.1.1.1.cmml" xref="S2.SS3.p2.3.m3.1.1">subscript</csymbol><ci id="S2.SS3.p2.3.m3.1.1.2.cmml" xref="S2.SS3.p2.3.m3.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.3.m3.1.1.3.cmml" xref="S2.SS3.p2.3.m3.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.3.m3.1c">B_{2}</annotation></semantics></math>, so <math id="S2.SS3.p2.4.m4.1" class="ltx_Math" alttext="L=2" display="inline"><semantics id="S2.SS3.p2.4.m4.1a"><mrow id="S2.SS3.p2.4.m4.1.1" xref="S2.SS3.p2.4.m4.1.1.cmml"><mi id="S2.SS3.p2.4.m4.1.1.2" xref="S2.SS3.p2.4.m4.1.1.2.cmml">L</mi><mo id="S2.SS3.p2.4.m4.1.1.1" xref="S2.SS3.p2.4.m4.1.1.1.cmml">=</mo><mn id="S2.SS3.p2.4.m4.1.1.3" xref="S2.SS3.p2.4.m4.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.4.m4.1b"><apply id="S2.SS3.p2.4.m4.1.1.cmml" xref="S2.SS3.p2.4.m4.1.1"><eq id="S2.SS3.p2.4.m4.1.1.1.cmml" xref="S2.SS3.p2.4.m4.1.1.1"></eq><ci id="S2.SS3.p2.4.m4.1.1.2.cmml" xref="S2.SS3.p2.4.m4.1.1.2">𝐿</ci><cn type="integer" id="S2.SS3.p2.4.m4.1.1.3.cmml" xref="S2.SS3.p2.4.m4.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.4.m4.1c">L=2</annotation></semantics></math>. If we choose <math id="S2.SS3.p2.5.m5.1" class="ltx_Math" alttext="K_{i}=3" display="inline"><semantics id="S2.SS3.p2.5.m5.1a"><mrow id="S2.SS3.p2.5.m5.1.1" xref="S2.SS3.p2.5.m5.1.1.cmml"><msub id="S2.SS3.p2.5.m5.1.1.2" xref="S2.SS3.p2.5.m5.1.1.2.cmml"><mi id="S2.SS3.p2.5.m5.1.1.2.2" xref="S2.SS3.p2.5.m5.1.1.2.2.cmml">K</mi><mi id="S2.SS3.p2.5.m5.1.1.2.3" xref="S2.SS3.p2.5.m5.1.1.2.3.cmml">i</mi></msub><mo id="S2.SS3.p2.5.m5.1.1.1" xref="S2.SS3.p2.5.m5.1.1.1.cmml">=</mo><mn id="S2.SS3.p2.5.m5.1.1.3" xref="S2.SS3.p2.5.m5.1.1.3.cmml">3</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.5.m5.1b"><apply id="S2.SS3.p2.5.m5.1.1.cmml" xref="S2.SS3.p2.5.m5.1.1"><eq id="S2.SS3.p2.5.m5.1.1.1.cmml" xref="S2.SS3.p2.5.m5.1.1.1"></eq><apply id="S2.SS3.p2.5.m5.1.1.2.cmml" xref="S2.SS3.p2.5.m5.1.1.2"><csymbol cd="ambiguous" id="S2.SS3.p2.5.m5.1.1.2.1.cmml" xref="S2.SS3.p2.5.m5.1.1.2">subscript</csymbol><ci id="S2.SS3.p2.5.m5.1.1.2.2.cmml" xref="S2.SS3.p2.5.m5.1.1.2.2">𝐾</ci><ci id="S2.SS3.p2.5.m5.1.1.2.3.cmml" xref="S2.SS3.p2.5.m5.1.1.2.3">𝑖</ci></apply><cn type="integer" id="S2.SS3.p2.5.m5.1.1.3.cmml" xref="S2.SS3.p2.5.m5.1.1.3">3</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.5.m5.1c">K_{i}=3</annotation></semantics></math> for <math id="S2.SS3.p2.6.m6.2" class="ltx_Math" alttext="i=1,2" display="inline"><semantics id="S2.SS3.p2.6.m6.2a"><mrow id="S2.SS3.p2.6.m6.2.3" xref="S2.SS3.p2.6.m6.2.3.cmml"><mi id="S2.SS3.p2.6.m6.2.3.2" xref="S2.SS3.p2.6.m6.2.3.2.cmml">i</mi><mo id="S2.SS3.p2.6.m6.2.3.1" xref="S2.SS3.p2.6.m6.2.3.1.cmml">=</mo><mrow id="S2.SS3.p2.6.m6.2.3.3.2" xref="S2.SS3.p2.6.m6.2.3.3.1.cmml"><mn id="S2.SS3.p2.6.m6.1.1" xref="S2.SS3.p2.6.m6.1.1.cmml">1</mn><mo id="S2.SS3.p2.6.m6.2.3.3.2.1" xref="S2.SS3.p2.6.m6.2.3.3.1.cmml">,</mo><mn id="S2.SS3.p2.6.m6.2.2" xref="S2.SS3.p2.6.m6.2.2.cmml">2</mn></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.6.m6.2b"><apply id="S2.SS3.p2.6.m6.2.3.cmml" xref="S2.SS3.p2.6.m6.2.3"><eq id="S2.SS3.p2.6.m6.2.3.1.cmml" xref="S2.SS3.p2.6.m6.2.3.1"></eq><ci id="S2.SS3.p2.6.m6.2.3.2.cmml" xref="S2.SS3.p2.6.m6.2.3.2">𝑖</ci><list id="S2.SS3.p2.6.m6.2.3.3.1.cmml" xref="S2.SS3.p2.6.m6.2.3.3.2"><cn type="integer" id="S2.SS3.p2.6.m6.1.1.cmml" xref="S2.SS3.p2.6.m6.1.1">1</cn><cn type="integer" id="S2.SS3.p2.6.m6.2.2.cmml" xref="S2.SS3.p2.6.m6.2.2">2</cn></list></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.6.m6.2c">i=1,2</annotation></semantics></math>, then we have a <math id="S2.SS3.p2.7.m7.1" class="ltx_Math" alttext="3" display="inline"><semantics id="S2.SS3.p2.7.m7.1a"><mn id="S2.SS3.p2.7.m7.1.1" xref="S2.SS3.p2.7.m7.1.1.cmml">3</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.7.m7.1b"><cn type="integer" id="S2.SS3.p2.7.m7.1.1.cmml" xref="S2.SS3.p2.7.m7.1.1">3</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.7.m7.1c">3</annotation></semantics></math> modules in <math id="S2.SS3.p2.8.m8.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS3.p2.8.m8.1a"><msub id="S2.SS3.p2.8.m8.1.1" xref="S2.SS3.p2.8.m8.1.1.cmml"><mi id="S2.SS3.p2.8.m8.1.1.2" xref="S2.SS3.p2.8.m8.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.8.m8.1.1.3" xref="S2.SS3.p2.8.m8.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.8.m8.1b"><apply id="S2.SS3.p2.8.m8.1.1.cmml" xref="S2.SS3.p2.8.m8.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.8.m8.1.1.1.cmml" xref="S2.SS3.p2.8.m8.1.1">subscript</csymbol><ci id="S2.SS3.p2.8.m8.1.1.2.cmml" xref="S2.SS3.p2.8.m8.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.8.m8.1.1.3.cmml" xref="S2.SS3.p2.8.m8.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.8.m8.1c">B_{1}</annotation></semantics></math>, and <math id="S2.SS3.p2.9.m9.1" class="ltx_Math" alttext="3" display="inline"><semantics id="S2.SS3.p2.9.m9.1a"><mn id="S2.SS3.p2.9.m9.1.1" xref="S2.SS3.p2.9.m9.1.1.cmml">3</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.9.m9.1b"><cn type="integer" id="S2.SS3.p2.9.m9.1.1.cmml" xref="S2.SS3.p2.9.m9.1.1">3</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.9.m9.1c">3</annotation></semantics></math> modules in <math id="S2.SS3.p2.10.m10.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS3.p2.10.m10.1a"><msub id="S2.SS3.p2.10.m10.1.1" xref="S2.SS3.p2.10.m10.1.1.cmml"><mi id="S2.SS3.p2.10.m10.1.1.2" xref="S2.SS3.p2.10.m10.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.10.m10.1.1.3" xref="S2.SS3.p2.10.m10.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.10.m10.1b"><apply id="S2.SS3.p2.10.m10.1.1.cmml" xref="S2.SS3.p2.10.m10.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.10.m10.1.1.1.cmml" xref="S2.SS3.p2.10.m10.1.1">subscript</csymbol><ci id="S2.SS3.p2.10.m10.1.1.2.cmml" xref="S2.SS3.p2.10.m10.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.10.m10.1.1.3.cmml" xref="S2.SS3.p2.10.m10.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.10.m10.1c">B_{2}</annotation></semantics></math>, for a total of <math id="S2.SS3.p2.11.m11.1" class="ltx_Math" alttext="6" display="inline"><semantics id="S2.SS3.p2.11.m11.1a"><mn id="S2.SS3.p2.11.m11.1.1" xref="S2.SS3.p2.11.m11.1.1.cmml">6</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.11.m11.1b"><cn type="integer" id="S2.SS3.p2.11.m11.1.1.cmml" xref="S2.SS3.p2.11.m11.1.1">6</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.11.m11.1c">6</annotation></semantics></math> modules. Any choice of module from <math id="S2.SS3.p2.12.m12.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS3.p2.12.m12.1a"><msub id="S2.SS3.p2.12.m12.1.1" xref="S2.SS3.p2.12.m12.1.1.cmml"><mi id="S2.SS3.p2.12.m12.1.1.2" xref="S2.SS3.p2.12.m12.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.12.m12.1.1.3" xref="S2.SS3.p2.12.m12.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.12.m12.1b"><apply id="S2.SS3.p2.12.m12.1.1.cmml" xref="S2.SS3.p2.12.m12.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.12.m12.1.1.1.cmml" xref="S2.SS3.p2.12.m12.1.1">subscript</csymbol><ci id="S2.SS3.p2.12.m12.1.1.2.cmml" xref="S2.SS3.p2.12.m12.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.12.m12.1.1.3.cmml" xref="S2.SS3.p2.12.m12.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.12.m12.1c">B_{1}</annotation></semantics></math> and module form <math id="S2.SS3.p2.13.m13.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS3.p2.13.m13.1a"><msub id="S2.SS3.p2.13.m13.1.1" xref="S2.SS3.p2.13.m13.1.1.cmml"><mi id="S2.SS3.p2.13.m13.1.1.2" xref="S2.SS3.p2.13.m13.1.1.2.cmml">B</mi><mn id="S2.SS3.p2.13.m13.1.1.3" xref="S2.SS3.p2.13.m13.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.13.m13.1b"><apply id="S2.SS3.p2.13.m13.1.1.cmml" xref="S2.SS3.p2.13.m13.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.13.m13.1.1.1.cmml" xref="S2.SS3.p2.13.m13.1.1">subscript</csymbol><ci id="S2.SS3.p2.13.m13.1.1.2.cmml" xref="S2.SS3.p2.13.m13.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p2.13.m13.1.1.3.cmml" xref="S2.SS3.p2.13.m13.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.13.m13.1c">B_{2}</annotation></semantics></math> defines a neural network; as in <cite class="ltx_cite ltx_citemacro_cite">Dean (<a href="#bib.bib10" title="" class="ltx_ref">2021</a>)</cite>, we call each of these <math id="S2.SS3.p2.14.m14.1" class="ltx_Math" alttext="9" display="inline"><semantics id="S2.SS3.p2.14.m14.1a"><mn id="S2.SS3.p2.14.m14.1.1" xref="S2.SS3.p2.14.m14.1.1.cmml">9</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.14.m14.1b"><cn type="integer" id="S2.SS3.p2.14.m14.1.1.cmml" xref="S2.SS3.p2.14.m14.1.1">9</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.14.m14.1c">9</annotation></semantics></math> possible networks a “path”, see <a href="#S2.F2" title="Figure 2 ‣ 2.3 Notation ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 2</span></a>.</p>
</div>
<figure id="S2.F2" class="ltx_figure">
<div class="ltx_flex_figure">
<div class="ltx_flex_cell ltx_flex_size_2">
<div id="S2.F2.1" class="ltx_block ltx_figure_panel ltx_minipage ltx_align_middle" style="width:212.5pt;">
<img src="/html/2403.10616/assets/x2.png" id="S2.F2.1.g1" class="ltx_graphics ltx_img_landscape" width="298" height="117" alt="Refer to caption">
</div>
</div>
<div class="ltx_flex_cell ltx_flex_size_2">
<div id="S2.F2.2" class="ltx_block ltx_figure_panel ltx_minipage ltx_align_middle" style="width:212.5pt;">
<img src="/html/2403.10616/assets/x3.png" id="S2.F2.2.g1" class="ltx_graphics ltx_img_landscape" width="298" height="117" alt="Refer to caption">
</div>
</div>
</div>
<figcaption class="ltx_caption"><span class="ltx_tag ltx_tag_figure"><span id="S2.F2.18.8.1" class="ltx_text" style="font-size:90%;">Figure 2</span>: </span><span id="S2.F2.16.7" class="ltx_text" style="font-size:90%;">An illustration of the first example from Section <a href="#S2.SS3" title="2.3 Notation ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.3</span></a>. A <math id="S2.F2.10.1.m1.1" class="ltx_Math" alttext="4" display="inline"><semantics id="S2.F2.10.1.m1.1b"><mn id="S2.F2.10.1.m1.1.1" xref="S2.F2.10.1.m1.1.1.cmml">4</mn><annotation-xml encoding="MathML-Content" id="S2.F2.10.1.m1.1c"><cn type="integer" id="S2.F2.10.1.m1.1.1.cmml" xref="S2.F2.10.1.m1.1.1">4</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.10.1.m1.1d">4</annotation></semantics></math> layer neural network, with block <math id="S2.F2.11.2.m2.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.F2.11.2.m2.1b"><msub id="S2.F2.11.2.m2.1.1" xref="S2.F2.11.2.m2.1.1.cmml"><mi id="S2.F2.11.2.m2.1.1.2" xref="S2.F2.11.2.m2.1.1.2.cmml">B</mi><mn id="S2.F2.11.2.m2.1.1.3" xref="S2.F2.11.2.m2.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.F2.11.2.m2.1c"><apply id="S2.F2.11.2.m2.1.1.cmml" xref="S2.F2.11.2.m2.1.1"><csymbol cd="ambiguous" id="S2.F2.11.2.m2.1.1.1.cmml" xref="S2.F2.11.2.m2.1.1">subscript</csymbol><ci id="S2.F2.11.2.m2.1.1.2.cmml" xref="S2.F2.11.2.m2.1.1.2">𝐵</ci><cn type="integer" id="S2.F2.11.2.m2.1.1.3.cmml" xref="S2.F2.11.2.m2.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.11.2.m2.1d">B_{1}</annotation></semantics></math> consisting of the first <math id="S2.F2.12.3.m3.1" class="ltx_Math" alttext="2" display="inline"><semantics id="S2.F2.12.3.m3.1b"><mn id="S2.F2.12.3.m3.1.1" xref="S2.F2.12.3.m3.1.1.cmml">2</mn><annotation-xml encoding="MathML-Content" id="S2.F2.12.3.m3.1c"><cn type="integer" id="S2.F2.12.3.m3.1.1.cmml" xref="S2.F2.12.3.m3.1.1">2</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.12.3.m3.1d">2</annotation></semantics></math> layers and <math id="S2.F2.13.4.m4.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.F2.13.4.m4.1b"><msub id="S2.F2.13.4.m4.1.1" xref="S2.F2.13.4.m4.1.1.cmml"><mi id="S2.F2.13.4.m4.1.1.2" xref="S2.F2.13.4.m4.1.1.2.cmml">B</mi><mn id="S2.F2.13.4.m4.1.1.3" xref="S2.F2.13.4.m4.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.F2.13.4.m4.1c"><apply id="S2.F2.13.4.m4.1.1.cmml" xref="S2.F2.13.4.m4.1.1"><csymbol cd="ambiguous" id="S2.F2.13.4.m4.1.1.1.cmml" xref="S2.F2.13.4.m4.1.1">subscript</csymbol><ci id="S2.F2.13.4.m4.1.1.2.cmml" xref="S2.F2.13.4.m4.1.1.2">𝐵</ci><cn type="integer" id="S2.F2.13.4.m4.1.1.3.cmml" xref="S2.F2.13.4.m4.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.13.4.m4.1d">B_{2}</annotation></semantics></math> consisting of the next <math id="S2.F2.14.5.m5.1" class="ltx_Math" alttext="2" display="inline"><semantics id="S2.F2.14.5.m5.1b"><mn id="S2.F2.14.5.m5.1.1" xref="S2.F2.14.5.m5.1.1.cmml">2</mn><annotation-xml encoding="MathML-Content" id="S2.F2.14.5.m5.1c"><cn type="integer" id="S2.F2.14.5.m5.1.1.cmml" xref="S2.F2.14.5.m5.1.1">2</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.14.5.m5.1d">2</annotation></semantics></math> layers. Each block has <math id="S2.F2.15.6.m6.1" class="ltx_Math" alttext="3" display="inline"><semantics id="S2.F2.15.6.m6.1b"><mn id="S2.F2.15.6.m6.1.1" xref="S2.F2.15.6.m6.1.1.cmml">3</mn><annotation-xml encoding="MathML-Content" id="S2.F2.15.6.m6.1c"><cn type="integer" id="S2.F2.15.6.m6.1.1.cmml" xref="S2.F2.15.6.m6.1.1">3</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.15.6.m6.1d">3</annotation></semantics></math> choices of module (each with its own parameters), represented by different colors. On the left, we show all of the <math id="S2.F2.16.7.m7.1" class="ltx_Math" alttext="9" display="inline"><semantics id="S2.F2.16.7.m7.1b"><mn id="S2.F2.16.7.m7.1.1" xref="S2.F2.16.7.m7.1.1.cmml">9</mn><annotation-xml encoding="MathML-Content" id="S2.F2.16.7.m7.1c"><cn type="integer" id="S2.F2.16.7.m7.1.1.cmml" xref="S2.F2.16.7.m7.1.1">9</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.F2.16.7.m7.1d">9</annotation></semantics></math> possible paths. On the right, we show a single path. </span></figcaption>
</figure>
<div id="S2.SS3.p3" class="ltx_para">
<p id="S2.SS3.p3.11" class="ltx_p">In this example, and most of the cases we will consider in the rest of this paper, the “blocks” <math id="S2.SS3.p3.1.m1.1" class="ltx_Math" alttext="B_{i}" display="inline"><semantics id="S2.SS3.p3.1.m1.1a"><msub id="S2.SS3.p3.1.m1.1.1" xref="S2.SS3.p3.1.m1.1.1.cmml"><mi id="S2.SS3.p3.1.m1.1.1.2" xref="S2.SS3.p3.1.m1.1.1.2.cmml">B</mi><mi id="S2.SS3.p3.1.m1.1.1.3" xref="S2.SS3.p3.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.1.m1.1b"><apply id="S2.SS3.p3.1.m1.1.1.cmml" xref="S2.SS3.p3.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS3.p3.1.m1.1.1.1.cmml" xref="S2.SS3.p3.1.m1.1.1">subscript</csymbol><ci id="S2.SS3.p3.1.m1.1.1.2.cmml" xref="S2.SS3.p3.1.m1.1.1.2">𝐵</ci><ci id="S2.SS3.p3.1.m1.1.1.3.cmml" xref="S2.SS3.p3.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.1.m1.1c">B_{i}</annotation></semantics></math> determine contiguous sub-networks of <math id="S2.SS3.p3.2.m2.1" class="ltx_Math" alttext="A" display="inline"><semantics id="S2.SS3.p3.2.m2.1a"><mi id="S2.SS3.p3.2.m2.1.1" xref="S2.SS3.p3.2.m2.1.1.cmml">A</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.2.m2.1b"><ci id="S2.SS3.p3.2.m2.1.1.cmml" xref="S2.SS3.p3.2.m2.1.1">𝐴</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.2.m2.1c">A</annotation></semantics></math> that can act as an input-output mapping, but this is not necessary. We could just as easily let <math id="S2.SS3.p3.3.m3.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS3.p3.3.m3.1a"><msub id="S2.SS3.p3.3.m3.1.1" xref="S2.SS3.p3.3.m3.1.1.cmml"><mi id="S2.SS3.p3.3.m3.1.1.2" xref="S2.SS3.p3.3.m3.1.1.2.cmml">B</mi><mn id="S2.SS3.p3.3.m3.1.1.3" xref="S2.SS3.p3.3.m3.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.3.m3.1b"><apply id="S2.SS3.p3.3.m3.1.1.cmml" xref="S2.SS3.p3.3.m3.1.1"><csymbol cd="ambiguous" id="S2.SS3.p3.3.m3.1.1.1.cmml" xref="S2.SS3.p3.3.m3.1.1">subscript</csymbol><ci id="S2.SS3.p3.3.m3.1.1.2.cmml" xref="S2.SS3.p3.3.m3.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p3.3.m3.1.1.3.cmml" xref="S2.SS3.p3.3.m3.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.3.m3.1c">B_{1}</annotation></semantics></math> be layers <math id="S2.SS3.p3.4.m4.1" class="ltx_Math" alttext="1" display="inline"><semantics id="S2.SS3.p3.4.m4.1a"><mn id="S2.SS3.p3.4.m4.1.1" xref="S2.SS3.p3.4.m4.1.1.cmml">1</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.4.m4.1b"><cn type="integer" id="S2.SS3.p3.4.m4.1.1.cmml" xref="S2.SS3.p3.4.m4.1.1">1</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.4.m4.1c">1</annotation></semantics></math> and <math id="S2.SS3.p3.5.m5.1" class="ltx_Math" alttext="3" display="inline"><semantics id="S2.SS3.p3.5.m5.1a"><mn id="S2.SS3.p3.5.m5.1.1" xref="S2.SS3.p3.5.m5.1.1.cmml">3</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.5.m5.1b"><cn type="integer" id="S2.SS3.p3.5.m5.1.1.cmml" xref="S2.SS3.p3.5.m5.1.1">3</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.5.m5.1c">3</annotation></semantics></math> and <math id="S2.SS3.p3.6.m6.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS3.p3.6.m6.1a"><msub id="S2.SS3.p3.6.m6.1.1" xref="S2.SS3.p3.6.m6.1.1.cmml"><mi id="S2.SS3.p3.6.m6.1.1.2" xref="S2.SS3.p3.6.m6.1.1.2.cmml">B</mi><mn id="S2.SS3.p3.6.m6.1.1.3" xref="S2.SS3.p3.6.m6.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.6.m6.1b"><apply id="S2.SS3.p3.6.m6.1.1.cmml" xref="S2.SS3.p3.6.m6.1.1"><csymbol cd="ambiguous" id="S2.SS3.p3.6.m6.1.1.1.cmml" xref="S2.SS3.p3.6.m6.1.1">subscript</csymbol><ci id="S2.SS3.p3.6.m6.1.1.2.cmml" xref="S2.SS3.p3.6.m6.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p3.6.m6.1.1.3.cmml" xref="S2.SS3.p3.6.m6.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.6.m6.1c">B_{2}</annotation></semantics></math> be layers <math id="S2.SS3.p3.7.m7.1" class="ltx_Math" alttext="2" display="inline"><semantics id="S2.SS3.p3.7.m7.1a"><mn id="S2.SS3.p3.7.m7.1.1" xref="S2.SS3.p3.7.m7.1.1.cmml">2</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.7.m7.1b"><cn type="integer" id="S2.SS3.p3.7.m7.1.1.cmml" xref="S2.SS3.p3.7.m7.1.1">2</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.7.m7.1c">2</annotation></semantics></math> and <math id="S2.SS3.p3.8.m8.1" class="ltx_Math" alttext="4" display="inline"><semantics id="S2.SS3.p3.8.m8.1a"><mn id="S2.SS3.p3.8.m8.1.1" xref="S2.SS3.p3.8.m8.1.1.cmml">4</mn><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.8.m8.1b"><cn type="integer" id="S2.SS3.p3.8.m8.1.1.cmml" xref="S2.SS3.p3.8.m8.1.1">4</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.8.m8.1c">4</annotation></semantics></math>; or <math id="S2.SS3.p3.9.m9.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS3.p3.9.m9.1a"><msub id="S2.SS3.p3.9.m9.1.1" xref="S2.SS3.p3.9.m9.1.1.cmml"><mi id="S2.SS3.p3.9.m9.1.1.2" xref="S2.SS3.p3.9.m9.1.1.2.cmml">B</mi><mn id="S2.SS3.p3.9.m9.1.1.3" xref="S2.SS3.p3.9.m9.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.9.m9.1b"><apply id="S2.SS3.p3.9.m9.1.1.cmml" xref="S2.SS3.p3.9.m9.1.1"><csymbol cd="ambiguous" id="S2.SS3.p3.9.m9.1.1.1.cmml" xref="S2.SS3.p3.9.m9.1.1">subscript</csymbol><ci id="S2.SS3.p3.9.m9.1.1.2.cmml" xref="S2.SS3.p3.9.m9.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p3.9.m9.1.1.3.cmml" xref="S2.SS3.p3.9.m9.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.9.m9.1c">B_{1}</annotation></semantics></math> be all the biases in the network, and <math id="S2.SS3.p3.10.m10.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS3.p3.10.m10.1a"><msub id="S2.SS3.p3.10.m10.1.1" xref="S2.SS3.p3.10.m10.1.1.cmml"><mi id="S2.SS3.p3.10.m10.1.1.2" xref="S2.SS3.p3.10.m10.1.1.2.cmml">B</mi><mn id="S2.SS3.p3.10.m10.1.1.3" xref="S2.SS3.p3.10.m10.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.10.m10.1b"><apply id="S2.SS3.p3.10.m10.1.1.cmml" xref="S2.SS3.p3.10.m10.1.1"><csymbol cd="ambiguous" id="S2.SS3.p3.10.m10.1.1.1.cmml" xref="S2.SS3.p3.10.m10.1.1">subscript</csymbol><ci id="S2.SS3.p3.10.m10.1.1.2.cmml" xref="S2.SS3.p3.10.m10.1.1.2">𝐵</ci><cn type="integer" id="S2.SS3.p3.10.m10.1.1.3.cmml" xref="S2.SS3.p3.10.m10.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.10.m10.1c">B_{2}</annotation></semantics></math> be all the linear parameters. Nevertheless, we will call <math id="S2.SS3.p3.11.m11.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS3.p3.11.m11.1a"><mi id="S2.SS3.p3.11.m11.1.1" xref="S2.SS3.p3.11.m11.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p3.11.m11.1b"><ci id="S2.SS3.p3.11.m11.1.1.cmml" xref="S2.SS3.p3.11.m11.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p3.11.m11.1c">i</annotation></semantics></math> the “level” of the module.</p>
</div>
<div id="S2.SS3.p4" class="ltx_para">
<p id="S2.SS3.p4.19" class="ltx_p">In order to transform an input <math id="S2.SS3.p4.1.m1.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S2.SS3.p4.1.m1.1a"><mi id="S2.SS3.p4.1.m1.1.1" xref="S2.SS3.p4.1.m1.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.1.m1.1b"><ci id="S2.SS3.p4.1.m1.1.1.cmml" xref="S2.SS3.p4.1.m1.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.1.m1.1c">x</annotation></semantics></math> to an output <math id="S2.SS3.p4.2.m2.1" class="ltx_Math" alttext="y" display="inline"><semantics id="S2.SS3.p4.2.m2.1a"><mi id="S2.SS3.p4.2.m2.1.1" xref="S2.SS3.p4.2.m2.1.1.cmml">y</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.2.m2.1b"><ci id="S2.SS3.p4.2.m2.1.1.cmml" xref="S2.SS3.p4.2.m2.1.1">𝑦</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.2.m2.1c">y</annotation></semantics></math>, we need to figure out which parameters will be used to operate on <math id="S2.SS3.p4.3.m3.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S2.SS3.p4.3.m3.1a"><mi id="S2.SS3.p4.3.m3.1.1" xref="S2.SS3.p4.3.m3.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.3.m3.1b"><ci id="S2.SS3.p4.3.m3.1.1.cmml" xref="S2.SS3.p4.3.m3.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.3.m3.1c">x</annotation></semantics></math>. The function that takes <math id="S2.SS3.p4.4.m4.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S2.SS3.p4.4.m4.1a"><mi id="S2.SS3.p4.4.m4.1.1" xref="S2.SS3.p4.4.m4.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.4.m4.1b"><ci id="S2.SS3.p4.4.m4.1.1.cmml" xref="S2.SS3.p4.4.m4.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.4.m4.1c">x</annotation></semantics></math> to a choice of parameters will be called the “router”, and is denoted by <math id="S2.SS3.p4.5.m5.1" class="ltx_Math" alttext="r" display="inline"><semantics id="S2.SS3.p4.5.m5.1a"><mi id="S2.SS3.p4.5.m5.1.1" xref="S2.SS3.p4.5.m5.1.1.cmml">r</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.5.m5.1b"><ci id="S2.SS3.p4.5.m5.1.1.cmml" xref="S2.SS3.p4.5.m5.1.1">𝑟</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.5.m5.1c">r</annotation></semantics></math>; thus <math id="S2.SS3.p4.6.m6.1" class="ltx_Math" alttext="r(x)" display="inline"><semantics id="S2.SS3.p4.6.m6.1a"><mrow id="S2.SS3.p4.6.m6.1.2" xref="S2.SS3.p4.6.m6.1.2.cmml"><mi id="S2.SS3.p4.6.m6.1.2.2" xref="S2.SS3.p4.6.m6.1.2.2.cmml">r</mi><mo lspace="0em" rspace="0em" id="S2.SS3.p4.6.m6.1.2.1" xref="S2.SS3.p4.6.m6.1.2.1.cmml">​</mo><mrow id="S2.SS3.p4.6.m6.1.2.3.2" xref="S2.SS3.p4.6.m6.1.2.cmml"><mo stretchy="false" id="S2.SS3.p4.6.m6.1.2.3.2.1" xref="S2.SS3.p4.6.m6.1.2.cmml">(</mo><mi id="S2.SS3.p4.6.m6.1.1" xref="S2.SS3.p4.6.m6.1.1.cmml">x</mi><mo stretchy="false" id="S2.SS3.p4.6.m6.1.2.3.2.2" xref="S2.SS3.p4.6.m6.1.2.cmml">)</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.6.m6.1b"><apply id="S2.SS3.p4.6.m6.1.2.cmml" xref="S2.SS3.p4.6.m6.1.2"><times id="S2.SS3.p4.6.m6.1.2.1.cmml" xref="S2.SS3.p4.6.m6.1.2.1"></times><ci id="S2.SS3.p4.6.m6.1.2.2.cmml" xref="S2.SS3.p4.6.m6.1.2.2">𝑟</ci><ci id="S2.SS3.p4.6.m6.1.1.cmml" xref="S2.SS3.p4.6.m6.1.1">𝑥</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.6.m6.1c">r(x)</annotation></semantics></math> has the form <math id="S2.SS3.p4.7.m7.3" class="ltx_Math" alttext="[j_{1},...,j_{L}]" display="inline"><semantics id="S2.SS3.p4.7.m7.3a"><mrow id="S2.SS3.p4.7.m7.3.3.2" xref="S2.SS3.p4.7.m7.3.3.3.cmml"><mo stretchy="false" id="S2.SS3.p4.7.m7.3.3.2.3" xref="S2.SS3.p4.7.m7.3.3.3.cmml">[</mo><msub id="S2.SS3.p4.7.m7.2.2.1.1" xref="S2.SS3.p4.7.m7.2.2.1.1.cmml"><mi id="S2.SS3.p4.7.m7.2.2.1.1.2" xref="S2.SS3.p4.7.m7.2.2.1.1.2.cmml">j</mi><mn id="S2.SS3.p4.7.m7.2.2.1.1.3" xref="S2.SS3.p4.7.m7.2.2.1.1.3.cmml">1</mn></msub><mo id="S2.SS3.p4.7.m7.3.3.2.4" xref="S2.SS3.p4.7.m7.3.3.3.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p4.7.m7.1.1" xref="S2.SS3.p4.7.m7.1.1.cmml">…</mi><mo id="S2.SS3.p4.7.m7.3.3.2.5" xref="S2.SS3.p4.7.m7.3.3.3.cmml">,</mo><msub id="S2.SS3.p4.7.m7.3.3.2.2" xref="S2.SS3.p4.7.m7.3.3.2.2.cmml"><mi id="S2.SS3.p4.7.m7.3.3.2.2.2" xref="S2.SS3.p4.7.m7.3.3.2.2.2.cmml">j</mi><mi id="S2.SS3.p4.7.m7.3.3.2.2.3" xref="S2.SS3.p4.7.m7.3.3.2.2.3.cmml">L</mi></msub><mo stretchy="false" id="S2.SS3.p4.7.m7.3.3.2.6" xref="S2.SS3.p4.7.m7.3.3.3.cmml">]</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.7.m7.3b"><list id="S2.SS3.p4.7.m7.3.3.3.cmml" xref="S2.SS3.p4.7.m7.3.3.2"><apply id="S2.SS3.p4.7.m7.2.2.1.1.cmml" xref="S2.SS3.p4.7.m7.2.2.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.7.m7.2.2.1.1.1.cmml" xref="S2.SS3.p4.7.m7.2.2.1.1">subscript</csymbol><ci id="S2.SS3.p4.7.m7.2.2.1.1.2.cmml" xref="S2.SS3.p4.7.m7.2.2.1.1.2">𝑗</ci><cn type="integer" id="S2.SS3.p4.7.m7.2.2.1.1.3.cmml" xref="S2.SS3.p4.7.m7.2.2.1.1.3">1</cn></apply><ci id="S2.SS3.p4.7.m7.1.1.cmml" xref="S2.SS3.p4.7.m7.1.1">…</ci><apply id="S2.SS3.p4.7.m7.3.3.2.2.cmml" xref="S2.SS3.p4.7.m7.3.3.2.2"><csymbol cd="ambiguous" id="S2.SS3.p4.7.m7.3.3.2.2.1.cmml" xref="S2.SS3.p4.7.m7.3.3.2.2">subscript</csymbol><ci id="S2.SS3.p4.7.m7.3.3.2.2.2.cmml" xref="S2.SS3.p4.7.m7.3.3.2.2.2">𝑗</ci><ci id="S2.SS3.p4.7.m7.3.3.2.2.3.cmml" xref="S2.SS3.p4.7.m7.3.3.2.2.3">𝐿</ci></apply></list></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.7.m7.3c">[j_{1},...,j_{L}]</annotation></semantics></math>, where each <math id="S2.SS3.p4.8.m8.1" class="ltx_Math" alttext="j_{l}" display="inline"><semantics id="S2.SS3.p4.8.m8.1a"><msub id="S2.SS3.p4.8.m8.1.1" xref="S2.SS3.p4.8.m8.1.1.cmml"><mi id="S2.SS3.p4.8.m8.1.1.2" xref="S2.SS3.p4.8.m8.1.1.2.cmml">j</mi><mi id="S2.SS3.p4.8.m8.1.1.3" xref="S2.SS3.p4.8.m8.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.8.m8.1b"><apply id="S2.SS3.p4.8.m8.1.1.cmml" xref="S2.SS3.p4.8.m8.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.8.m8.1.1.1.cmml" xref="S2.SS3.p4.8.m8.1.1">subscript</csymbol><ci id="S2.SS3.p4.8.m8.1.1.2.cmml" xref="S2.SS3.p4.8.m8.1.1.2">𝑗</ci><ci id="S2.SS3.p4.8.m8.1.1.3.cmml" xref="S2.SS3.p4.8.m8.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.8.m8.1c">j_{l}</annotation></semantics></math> indexes one of the <math id="S2.SS3.p4.9.m9.1" class="ltx_Math" alttext="K_{l}" display="inline"><semantics id="S2.SS3.p4.9.m9.1a"><msub id="S2.SS3.p4.9.m9.1.1" xref="S2.SS3.p4.9.m9.1.1.cmml"><mi id="S2.SS3.p4.9.m9.1.1.2" xref="S2.SS3.p4.9.m9.1.1.2.cmml">K</mi><mi id="S2.SS3.p4.9.m9.1.1.3" xref="S2.SS3.p4.9.m9.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.9.m9.1b"><apply id="S2.SS3.p4.9.m9.1.1.cmml" xref="S2.SS3.p4.9.m9.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.9.m9.1.1.1.cmml" xref="S2.SS3.p4.9.m9.1.1">subscript</csymbol><ci id="S2.SS3.p4.9.m9.1.1.2.cmml" xref="S2.SS3.p4.9.m9.1.1.2">𝐾</ci><ci id="S2.SS3.p4.9.m9.1.1.3.cmml" xref="S2.SS3.p4.9.m9.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.9.m9.1c">K_{l}</annotation></semantics></math> choices of parameters for <math id="S2.SS3.p4.10.m10.1" class="ltx_Math" alttext="B_{l}" display="inline"><semantics id="S2.SS3.p4.10.m10.1a"><msub id="S2.SS3.p4.10.m10.1.1" xref="S2.SS3.p4.10.m10.1.1.cmml"><mi id="S2.SS3.p4.10.m10.1.1.2" xref="S2.SS3.p4.10.m10.1.1.2.cmml">B</mi><mi id="S2.SS3.p4.10.m10.1.1.3" xref="S2.SS3.p4.10.m10.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.10.m10.1b"><apply id="S2.SS3.p4.10.m10.1.1.cmml" xref="S2.SS3.p4.10.m10.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.10.m10.1.1.1.cmml" xref="S2.SS3.p4.10.m10.1.1">subscript</csymbol><ci id="S2.SS3.p4.10.m10.1.1.2.cmml" xref="S2.SS3.p4.10.m10.1.1.2">𝐵</ci><ci id="S2.SS3.p4.10.m10.1.1.3.cmml" xref="S2.SS3.p4.10.m10.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.10.m10.1c">B_{l}</annotation></semantics></math>, and <math id="S2.SS3.p4.11.m11.1" class="ltx_Math" alttext="r" display="inline"><semantics id="S2.SS3.p4.11.m11.1a"><mi id="S2.SS3.p4.11.m11.1.1" xref="S2.SS3.p4.11.m11.1.1.cmml">r</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.11.m11.1b"><ci id="S2.SS3.p4.11.m11.1.1.cmml" xref="S2.SS3.p4.11.m11.1.1">𝑟</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.11.m11.1c">r</annotation></semantics></math> maps an input to a path <math id="S2.SS3.p4.12.m12.3" class="ltx_Math" alttext="\pi=\pi_{j_{1},\ldots,j_{L}}" display="inline"><semantics id="S2.SS3.p4.12.m12.3a"><mrow id="S2.SS3.p4.12.m12.3.4" xref="S2.SS3.p4.12.m12.3.4.cmml"><mi id="S2.SS3.p4.12.m12.3.4.2" xref="S2.SS3.p4.12.m12.3.4.2.cmml">π</mi><mo id="S2.SS3.p4.12.m12.3.4.1" xref="S2.SS3.p4.12.m12.3.4.1.cmml">=</mo><msub id="S2.SS3.p4.12.m12.3.4.3" xref="S2.SS3.p4.12.m12.3.4.3.cmml"><mi id="S2.SS3.p4.12.m12.3.4.3.2" xref="S2.SS3.p4.12.m12.3.4.3.2.cmml">π</mi><mrow id="S2.SS3.p4.12.m12.3.3.3.3" xref="S2.SS3.p4.12.m12.3.3.3.4.cmml"><msub id="S2.SS3.p4.12.m12.2.2.2.2.1" xref="S2.SS3.p4.12.m12.2.2.2.2.1.cmml"><mi id="S2.SS3.p4.12.m12.2.2.2.2.1.2" xref="S2.SS3.p4.12.m12.2.2.2.2.1.2.cmml">j</mi><mn id="S2.SS3.p4.12.m12.2.2.2.2.1.3" xref="S2.SS3.p4.12.m12.2.2.2.2.1.3.cmml">1</mn></msub><mo id="S2.SS3.p4.12.m12.3.3.3.3.3" xref="S2.SS3.p4.12.m12.3.3.3.4.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p4.12.m12.1.1.1.1" xref="S2.SS3.p4.12.m12.1.1.1.1.cmml">…</mi><mo id="S2.SS3.p4.12.m12.3.3.3.3.4" xref="S2.SS3.p4.12.m12.3.3.3.4.cmml">,</mo><msub id="S2.SS3.p4.12.m12.3.3.3.3.2" xref="S2.SS3.p4.12.m12.3.3.3.3.2.cmml"><mi id="S2.SS3.p4.12.m12.3.3.3.3.2.2" xref="S2.SS3.p4.12.m12.3.3.3.3.2.2.cmml">j</mi><mi id="S2.SS3.p4.12.m12.3.3.3.3.2.3" xref="S2.SS3.p4.12.m12.3.3.3.3.2.3.cmml">L</mi></msub></mrow></msub></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.12.m12.3b"><apply id="S2.SS3.p4.12.m12.3.4.cmml" xref="S2.SS3.p4.12.m12.3.4"><eq id="S2.SS3.p4.12.m12.3.4.1.cmml" xref="S2.SS3.p4.12.m12.3.4.1"></eq><ci id="S2.SS3.p4.12.m12.3.4.2.cmml" xref="S2.SS3.p4.12.m12.3.4.2">𝜋</ci><apply id="S2.SS3.p4.12.m12.3.4.3.cmml" xref="S2.SS3.p4.12.m12.3.4.3"><csymbol cd="ambiguous" id="S2.SS3.p4.12.m12.3.4.3.1.cmml" xref="S2.SS3.p4.12.m12.3.4.3">subscript</csymbol><ci id="S2.SS3.p4.12.m12.3.4.3.2.cmml" xref="S2.SS3.p4.12.m12.3.4.3.2">𝜋</ci><list id="S2.SS3.p4.12.m12.3.3.3.4.cmml" xref="S2.SS3.p4.12.m12.3.3.3.3"><apply id="S2.SS3.p4.12.m12.2.2.2.2.1.cmml" xref="S2.SS3.p4.12.m12.2.2.2.2.1"><csymbol cd="ambiguous" id="S2.SS3.p4.12.m12.2.2.2.2.1.1.cmml" xref="S2.SS3.p4.12.m12.2.2.2.2.1">subscript</csymbol><ci id="S2.SS3.p4.12.m12.2.2.2.2.1.2.cmml" xref="S2.SS3.p4.12.m12.2.2.2.2.1.2">𝑗</ci><cn type="integer" id="S2.SS3.p4.12.m12.2.2.2.2.1.3.cmml" xref="S2.SS3.p4.12.m12.2.2.2.2.1.3">1</cn></apply><ci id="S2.SS3.p4.12.m12.1.1.1.1.cmml" xref="S2.SS3.p4.12.m12.1.1.1.1">…</ci><apply id="S2.SS3.p4.12.m12.3.3.3.3.2.cmml" xref="S2.SS3.p4.12.m12.3.3.3.3.2"><csymbol cd="ambiguous" id="S2.SS3.p4.12.m12.3.3.3.3.2.1.cmml" xref="S2.SS3.p4.12.m12.3.3.3.3.2">subscript</csymbol><ci id="S2.SS3.p4.12.m12.3.3.3.3.2.2.cmml" xref="S2.SS3.p4.12.m12.3.3.3.3.2.2">𝑗</ci><ci id="S2.SS3.p4.12.m12.3.3.3.3.2.3.cmml" xref="S2.SS3.p4.12.m12.3.3.3.3.2.3">𝐿</ci></apply></list></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.12.m12.3c">\pi=\pi_{j_{1},\ldots,j_{L}}</annotation></semantics></math>. We will also often collapse the sequence <math id="S2.SS3.p4.13.m13.3" class="ltx_Math" alttext="j_{1},...,j_{L}" display="inline"><semantics id="S2.SS3.p4.13.m13.3a"><mrow id="S2.SS3.p4.13.m13.3.3.2" xref="S2.SS3.p4.13.m13.3.3.3.cmml"><msub id="S2.SS3.p4.13.m13.2.2.1.1" xref="S2.SS3.p4.13.m13.2.2.1.1.cmml"><mi id="S2.SS3.p4.13.m13.2.2.1.1.2" xref="S2.SS3.p4.13.m13.2.2.1.1.2.cmml">j</mi><mn id="S2.SS3.p4.13.m13.2.2.1.1.3" xref="S2.SS3.p4.13.m13.2.2.1.1.3.cmml">1</mn></msub><mo id="S2.SS3.p4.13.m13.3.3.2.3" xref="S2.SS3.p4.13.m13.3.3.3.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p4.13.m13.1.1" xref="S2.SS3.p4.13.m13.1.1.cmml">…</mi><mo id="S2.SS3.p4.13.m13.3.3.2.4" xref="S2.SS3.p4.13.m13.3.3.3.cmml">,</mo><msub id="S2.SS3.p4.13.m13.3.3.2.2" xref="S2.SS3.p4.13.m13.3.3.2.2.cmml"><mi id="S2.SS3.p4.13.m13.3.3.2.2.2" xref="S2.SS3.p4.13.m13.3.3.2.2.2.cmml">j</mi><mi id="S2.SS3.p4.13.m13.3.3.2.2.3" xref="S2.SS3.p4.13.m13.3.3.2.2.3.cmml">L</mi></msub></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.13.m13.3b"><list id="S2.SS3.p4.13.m13.3.3.3.cmml" xref="S2.SS3.p4.13.m13.3.3.2"><apply id="S2.SS3.p4.13.m13.2.2.1.1.cmml" xref="S2.SS3.p4.13.m13.2.2.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.13.m13.2.2.1.1.1.cmml" xref="S2.SS3.p4.13.m13.2.2.1.1">subscript</csymbol><ci id="S2.SS3.p4.13.m13.2.2.1.1.2.cmml" xref="S2.SS3.p4.13.m13.2.2.1.1.2">𝑗</ci><cn type="integer" id="S2.SS3.p4.13.m13.2.2.1.1.3.cmml" xref="S2.SS3.p4.13.m13.2.2.1.1.3">1</cn></apply><ci id="S2.SS3.p4.13.m13.1.1.cmml" xref="S2.SS3.p4.13.m13.1.1">…</ci><apply id="S2.SS3.p4.13.m13.3.3.2.2.cmml" xref="S2.SS3.p4.13.m13.3.3.2.2"><csymbol cd="ambiguous" id="S2.SS3.p4.13.m13.3.3.2.2.1.cmml" xref="S2.SS3.p4.13.m13.3.3.2.2">subscript</csymbol><ci id="S2.SS3.p4.13.m13.3.3.2.2.2.cmml" xref="S2.SS3.p4.13.m13.3.3.2.2.2">𝑗</ci><ci id="S2.SS3.p4.13.m13.3.3.2.2.3.cmml" xref="S2.SS3.p4.13.m13.3.3.2.2.3">𝐿</ci></apply></list></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.13.m13.3c">j_{1},...,j_{L}</annotation></semantics></math> into a single number <math id="S2.SS3.p4.14.m14.1" class="ltx_Math" alttext="j" display="inline"><semantics id="S2.SS3.p4.14.m14.1a"><mi id="S2.SS3.p4.14.m14.1.1" xref="S2.SS3.p4.14.m14.1.1.cmml">j</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.14.m14.1b"><ci id="S2.SS3.p4.14.m14.1.1.cmml" xref="S2.SS3.p4.14.m14.1.1">𝑗</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.14.m14.1c">j</annotation></semantics></math>, where <math id="S2.SS3.p4.15.m15.4" class="ltx_Math" alttext="j\in\{1,\ldots,P=\Pi_{i=1}^{L}K_{i}\}" display="inline"><semantics id="S2.SS3.p4.15.m15.4a"><mrow id="S2.SS3.p4.15.m15.4.4" xref="S2.SS3.p4.15.m15.4.4.cmml"><mi id="S2.SS3.p4.15.m15.4.4.3" xref="S2.SS3.p4.15.m15.4.4.3.cmml">j</mi><mo id="S2.SS3.p4.15.m15.4.4.2" xref="S2.SS3.p4.15.m15.4.4.2.cmml">∈</mo><mrow id="S2.SS3.p4.15.m15.4.4.1.1" xref="S2.SS3.p4.15.m15.4.4.1.2.cmml"><mo stretchy="false" id="S2.SS3.p4.15.m15.4.4.1.1.2" xref="S2.SS3.p4.15.m15.4.4.1.2.cmml">{</mo><mrow id="S2.SS3.p4.15.m15.4.4.1.1.1" xref="S2.SS3.p4.15.m15.4.4.1.1.1.cmml"><mrow id="S2.SS3.p4.15.m15.4.4.1.1.1.2.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.2.1.cmml"><mn id="S2.SS3.p4.15.m15.1.1" xref="S2.SS3.p4.15.m15.1.1.cmml">1</mn><mo id="S2.SS3.p4.15.m15.4.4.1.1.1.2.2.1" xref="S2.SS3.p4.15.m15.4.4.1.1.1.2.1.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p4.15.m15.2.2" xref="S2.SS3.p4.15.m15.2.2.cmml">…</mi><mo id="S2.SS3.p4.15.m15.4.4.1.1.1.2.2.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.2.1.cmml">,</mo><mi id="S2.SS3.p4.15.m15.3.3" xref="S2.SS3.p4.15.m15.3.3.cmml">P</mi></mrow><mo id="S2.SS3.p4.15.m15.4.4.1.1.1.1" xref="S2.SS3.p4.15.m15.4.4.1.1.1.1.cmml">=</mo><mrow id="S2.SS3.p4.15.m15.4.4.1.1.1.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.cmml"><msubsup id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.cmml"><mi mathvariant="normal" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.2.cmml">Π</mi><mrow id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.cmml"><mi id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.2.cmml">i</mi><mo id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.1" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.1.cmml">=</mo><mn id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.3.cmml">1</mn></mrow><mi id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.3.cmml">L</mi></msubsup><mo lspace="0em" rspace="0em" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.1" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.1.cmml">​</mo><msub id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.cmml"><mi id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.2" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.2.cmml">K</mi><mi id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.3" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.3.cmml">i</mi></msub></mrow></mrow><mo stretchy="false" id="S2.SS3.p4.15.m15.4.4.1.1.3" xref="S2.SS3.p4.15.m15.4.4.1.2.cmml">}</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.15.m15.4b"><apply id="S2.SS3.p4.15.m15.4.4.cmml" xref="S2.SS3.p4.15.m15.4.4"><in id="S2.SS3.p4.15.m15.4.4.2.cmml" xref="S2.SS3.p4.15.m15.4.4.2"></in><ci id="S2.SS3.p4.15.m15.4.4.3.cmml" xref="S2.SS3.p4.15.m15.4.4.3">𝑗</ci><set id="S2.SS3.p4.15.m15.4.4.1.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1"><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1"><eq id="S2.SS3.p4.15.m15.4.4.1.1.1.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.1"></eq><list id="S2.SS3.p4.15.m15.4.4.1.1.1.2.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.2.2"><cn type="integer" id="S2.SS3.p4.15.m15.1.1.cmml" xref="S2.SS3.p4.15.m15.1.1">1</cn><ci id="S2.SS3.p4.15.m15.2.2.cmml" xref="S2.SS3.p4.15.m15.2.2">…</ci><ci id="S2.SS3.p4.15.m15.3.3.cmml" xref="S2.SS3.p4.15.m15.3.3">𝑃</ci></list><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3"><times id="S2.SS3.p4.15.m15.4.4.1.1.1.3.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.1"></times><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2"><csymbol cd="ambiguous" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2">superscript</csymbol><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2"><csymbol cd="ambiguous" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2">subscript</csymbol><ci id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.2">Π</ci><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3"><eq id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.1"></eq><ci id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.2">𝑖</ci><cn type="integer" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.2.3.3">1</cn></apply></apply><ci id="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.2.3">𝐿</ci></apply><apply id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3"><csymbol cd="ambiguous" id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.1.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3">subscript</csymbol><ci id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.2.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.2">𝐾</ci><ci id="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.3.cmml" xref="S2.SS3.p4.15.m15.4.4.1.1.1.3.3.3">𝑖</ci></apply></apply></apply></set></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.15.m15.4c">j\in\{1,\ldots,P=\Pi_{i=1}^{L}K_{i}\}</annotation></semantics></math>. If the training dataset is denoted by <math id="S2.SS3.p4.16.m16.1" class="ltx_Math" alttext="X" display="inline"><semantics id="S2.SS3.p4.16.m16.1a"><mi id="S2.SS3.p4.16.m16.1.1" xref="S2.SS3.p4.16.m16.1.1.cmml">X</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.16.m16.1b"><ci id="S2.SS3.p4.16.m16.1.1.cmml" xref="S2.SS3.p4.16.m16.1.1">𝑋</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.16.m16.1c">X</annotation></semantics></math>, the subset of data that is routed to path <math id="S2.SS3.p4.17.m17.1" class="ltx_Math" alttext="j" display="inline"><semantics id="S2.SS3.p4.17.m17.1a"><mi id="S2.SS3.p4.17.m17.1.1" xref="S2.SS3.p4.17.m17.1.1.cmml">j</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.17.m17.1b"><ci id="S2.SS3.p4.17.m17.1.1.cmml" xref="S2.SS3.p4.17.m17.1.1">𝑗</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.17.m17.1c">j</annotation></semantics></math> will be called the <math id="S2.SS3.p4.18.m18.1" class="ltx_Math" alttext="j" display="inline"><semantics id="S2.SS3.p4.18.m18.1a"><mi id="S2.SS3.p4.18.m18.1.1" xref="S2.SS3.p4.18.m18.1.1.cmml">j</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.18.m18.1b"><ci id="S2.SS3.p4.18.m18.1.1.cmml" xref="S2.SS3.p4.18.m18.1.1">𝑗</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.18.m18.1c">j</annotation></semantics></math>-th “shard” <math id="S2.SS3.p4.19.m19.1" class="ltx_Math" alttext="\mathcal{D}_{j}" display="inline"><semantics id="S2.SS3.p4.19.m19.1a"><msub id="S2.SS3.p4.19.m19.1.1" xref="S2.SS3.p4.19.m19.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS3.p4.19.m19.1.1.2" xref="S2.SS3.p4.19.m19.1.1.2.cmml">𝒟</mi><mi id="S2.SS3.p4.19.m19.1.1.3" xref="S2.SS3.p4.19.m19.1.1.3.cmml">j</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p4.19.m19.1b"><apply id="S2.SS3.p4.19.m19.1.1.cmml" xref="S2.SS3.p4.19.m19.1.1"><csymbol cd="ambiguous" id="S2.SS3.p4.19.m19.1.1.1.cmml" xref="S2.SS3.p4.19.m19.1.1">subscript</csymbol><ci id="S2.SS3.p4.19.m19.1.1.2.cmml" xref="S2.SS3.p4.19.m19.1.1.2">𝒟</ci><ci id="S2.SS3.p4.19.m19.1.1.3.cmml" xref="S2.SS3.p4.19.m19.1.1.3">𝑗</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p4.19.m19.1c">\mathcal{D}_{j}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS3.p5" class="ltx_para">
<p id="S2.SS3.p5.2" class="ltx_p">When the router maps an <math id="S2.SS3.p5.1.m1.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S2.SS3.p5.1.m1.1a"><mi id="S2.SS3.p5.1.m1.1.1" xref="S2.SS3.p5.1.m1.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p5.1.m1.1b"><ci id="S2.SS3.p5.1.m1.1.1.cmml" xref="S2.SS3.p5.1.m1.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p5.1.m1.1c">x</annotation></semantics></math> to a distribution over <math id="S2.SS3.p5.2.m2.3" class="ltx_Math" alttext="\{1,\ldots,P\}" display="inline"><semantics id="S2.SS3.p5.2.m2.3a"><mrow id="S2.SS3.p5.2.m2.3.4.2" xref="S2.SS3.p5.2.m2.3.4.1.cmml"><mo stretchy="false" id="S2.SS3.p5.2.m2.3.4.2.1" xref="S2.SS3.p5.2.m2.3.4.1.cmml">{</mo><mn id="S2.SS3.p5.2.m2.1.1" xref="S2.SS3.p5.2.m2.1.1.cmml">1</mn><mo id="S2.SS3.p5.2.m2.3.4.2.2" xref="S2.SS3.p5.2.m2.3.4.1.cmml">,</mo><mi mathvariant="normal" id="S2.SS3.p5.2.m2.2.2" xref="S2.SS3.p5.2.m2.2.2.cmml">…</mi><mo id="S2.SS3.p5.2.m2.3.4.2.3" xref="S2.SS3.p5.2.m2.3.4.1.cmml">,</mo><mi id="S2.SS3.p5.2.m2.3.3" xref="S2.SS3.p5.2.m2.3.3.cmml">P</mi><mo stretchy="false" id="S2.SS3.p5.2.m2.3.4.2.4" xref="S2.SS3.p5.2.m2.3.4.1.cmml">}</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p5.2.m2.3b"><set id="S2.SS3.p5.2.m2.3.4.1.cmml" xref="S2.SS3.p5.2.m2.3.4.2"><cn type="integer" id="S2.SS3.p5.2.m2.1.1.cmml" xref="S2.SS3.p5.2.m2.1.1">1</cn><ci id="S2.SS3.p5.2.m2.2.2.cmml" xref="S2.SS3.p5.2.m2.2.2">…</ci><ci id="S2.SS3.p5.2.m2.3.3.cmml" xref="S2.SS3.p5.2.m2.3.3">𝑃</ci></set></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p5.2.m2.3c">\{1,\ldots,P\}</annotation></semantics></math> rather to a deterministic choice, we say that the router is “soft”. In this work, we will consider “hard” routers instead, where the choice is single valued and deterministic.</p>
</div>
<div id="S2.SS3.p6" class="ltx_para">
<p id="S2.SS3.p6.18" class="ltx_p">Finally, it is useful to refer to module parameters. Assume that the <math id="S2.SS3.p6.1.m1.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS3.p6.1.m1.1a"><mi id="S2.SS3.p6.1.m1.1.1" xref="S2.SS3.p6.1.m1.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.1.m1.1b"><ci id="S2.SS3.p6.1.m1.1.1.cmml" xref="S2.SS3.p6.1.m1.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.1.m1.1c">i</annotation></semantics></math>-th path goes through module <math id="S2.SS3.p6.2.m2.1" class="ltx_Math" alttext="e" display="inline"><semantics id="S2.SS3.p6.2.m2.1a"><mi id="S2.SS3.p6.2.m2.1.1" xref="S2.SS3.p6.2.m2.1.1.cmml">e</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.2.m2.1b"><ci id="S2.SS3.p6.2.m2.1.1.cmml" xref="S2.SS3.p6.2.m2.1.1">𝑒</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.2.m2.1c">e</annotation></semantics></math> at level <math id="S2.SS3.p6.3.m3.1" class="ltx_Math" alttext="l" display="inline"><semantics id="S2.SS3.p6.3.m3.1a"><mi id="S2.SS3.p6.3.m3.1.1" xref="S2.SS3.p6.3.m3.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.3.m3.1b"><ci id="S2.SS3.p6.3.m3.1.1.cmml" xref="S2.SS3.p6.3.m3.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.3.m3.1c">l</annotation></semantics></math>, with <math id="S2.SS3.p6.4.m4.2" class="ltx_Math" alttext="e\in[1,K_{l}]" display="inline"><semantics id="S2.SS3.p6.4.m4.2a"><mrow id="S2.SS3.p6.4.m4.2.2" xref="S2.SS3.p6.4.m4.2.2.cmml"><mi id="S2.SS3.p6.4.m4.2.2.3" xref="S2.SS3.p6.4.m4.2.2.3.cmml">e</mi><mo id="S2.SS3.p6.4.m4.2.2.2" xref="S2.SS3.p6.4.m4.2.2.2.cmml">∈</mo><mrow id="S2.SS3.p6.4.m4.2.2.1.1" xref="S2.SS3.p6.4.m4.2.2.1.2.cmml"><mo stretchy="false" id="S2.SS3.p6.4.m4.2.2.1.1.2" xref="S2.SS3.p6.4.m4.2.2.1.2.cmml">[</mo><mn id="S2.SS3.p6.4.m4.1.1" xref="S2.SS3.p6.4.m4.1.1.cmml">1</mn><mo id="S2.SS3.p6.4.m4.2.2.1.1.3" xref="S2.SS3.p6.4.m4.2.2.1.2.cmml">,</mo><msub id="S2.SS3.p6.4.m4.2.2.1.1.1" xref="S2.SS3.p6.4.m4.2.2.1.1.1.cmml"><mi id="S2.SS3.p6.4.m4.2.2.1.1.1.2" xref="S2.SS3.p6.4.m4.2.2.1.1.1.2.cmml">K</mi><mi id="S2.SS3.p6.4.m4.2.2.1.1.1.3" xref="S2.SS3.p6.4.m4.2.2.1.1.1.3.cmml">l</mi></msub><mo stretchy="false" id="S2.SS3.p6.4.m4.2.2.1.1.4" xref="S2.SS3.p6.4.m4.2.2.1.2.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.4.m4.2b"><apply id="S2.SS3.p6.4.m4.2.2.cmml" xref="S2.SS3.p6.4.m4.2.2"><in id="S2.SS3.p6.4.m4.2.2.2.cmml" xref="S2.SS3.p6.4.m4.2.2.2"></in><ci id="S2.SS3.p6.4.m4.2.2.3.cmml" xref="S2.SS3.p6.4.m4.2.2.3">𝑒</ci><interval closure="closed" id="S2.SS3.p6.4.m4.2.2.1.2.cmml" xref="S2.SS3.p6.4.m4.2.2.1.1"><cn type="integer" id="S2.SS3.p6.4.m4.1.1.cmml" xref="S2.SS3.p6.4.m4.1.1">1</cn><apply id="S2.SS3.p6.4.m4.2.2.1.1.1.cmml" xref="S2.SS3.p6.4.m4.2.2.1.1.1"><csymbol cd="ambiguous" id="S2.SS3.p6.4.m4.2.2.1.1.1.1.cmml" xref="S2.SS3.p6.4.m4.2.2.1.1.1">subscript</csymbol><ci id="S2.SS3.p6.4.m4.2.2.1.1.1.2.cmml" xref="S2.SS3.p6.4.m4.2.2.1.1.1.2">𝐾</ci><ci id="S2.SS3.p6.4.m4.2.2.1.1.1.3.cmml" xref="S2.SS3.p6.4.m4.2.2.1.1.1.3">𝑙</ci></apply></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.4.m4.2c">e\in[1,K_{l}]</annotation></semantics></math>, <math id="S2.SS3.p6.5.m5.2" class="ltx_Math" alttext="l\in[1,L]" display="inline"><semantics id="S2.SS3.p6.5.m5.2a"><mrow id="S2.SS3.p6.5.m5.2.3" xref="S2.SS3.p6.5.m5.2.3.cmml"><mi id="S2.SS3.p6.5.m5.2.3.2" xref="S2.SS3.p6.5.m5.2.3.2.cmml">l</mi><mo id="S2.SS3.p6.5.m5.2.3.1" xref="S2.SS3.p6.5.m5.2.3.1.cmml">∈</mo><mrow id="S2.SS3.p6.5.m5.2.3.3.2" xref="S2.SS3.p6.5.m5.2.3.3.1.cmml"><mo stretchy="false" id="S2.SS3.p6.5.m5.2.3.3.2.1" xref="S2.SS3.p6.5.m5.2.3.3.1.cmml">[</mo><mn id="S2.SS3.p6.5.m5.1.1" xref="S2.SS3.p6.5.m5.1.1.cmml">1</mn><mo id="S2.SS3.p6.5.m5.2.3.3.2.2" xref="S2.SS3.p6.5.m5.2.3.3.1.cmml">,</mo><mi id="S2.SS3.p6.5.m5.2.2" xref="S2.SS3.p6.5.m5.2.2.cmml">L</mi><mo stretchy="false" id="S2.SS3.p6.5.m5.2.3.3.2.3" xref="S2.SS3.p6.5.m5.2.3.3.1.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.5.m5.2b"><apply id="S2.SS3.p6.5.m5.2.3.cmml" xref="S2.SS3.p6.5.m5.2.3"><in id="S2.SS3.p6.5.m5.2.3.1.cmml" xref="S2.SS3.p6.5.m5.2.3.1"></in><ci id="S2.SS3.p6.5.m5.2.3.2.cmml" xref="S2.SS3.p6.5.m5.2.3.2">𝑙</ci><interval closure="closed" id="S2.SS3.p6.5.m5.2.3.3.1.cmml" xref="S2.SS3.p6.5.m5.2.3.3.2"><cn type="integer" id="S2.SS3.p6.5.m5.1.1.cmml" xref="S2.SS3.p6.5.m5.1.1">1</cn><ci id="S2.SS3.p6.5.m5.2.2.cmml" xref="S2.SS3.p6.5.m5.2.2">𝐿</ci></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.5.m5.2c">l\in[1,L]</annotation></semantics></math>.
Since during training paths are not fully synchronized, we denote the local copy of the parameters of module <math id="S2.SS3.p6.6.m6.2" class="ltx_Math" alttext="(l,e)" display="inline"><semantics id="S2.SS3.p6.6.m6.2a"><mrow id="S2.SS3.p6.6.m6.2.3.2" xref="S2.SS3.p6.6.m6.2.3.1.cmml"><mo stretchy="false" id="S2.SS3.p6.6.m6.2.3.2.1" xref="S2.SS3.p6.6.m6.2.3.1.cmml">(</mo><mi id="S2.SS3.p6.6.m6.1.1" xref="S2.SS3.p6.6.m6.1.1.cmml">l</mi><mo id="S2.SS3.p6.6.m6.2.3.2.2" xref="S2.SS3.p6.6.m6.2.3.1.cmml">,</mo><mi id="S2.SS3.p6.6.m6.2.2" xref="S2.SS3.p6.6.m6.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.6.m6.2.3.2.3" xref="S2.SS3.p6.6.m6.2.3.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.6.m6.2b"><interval closure="open" id="S2.SS3.p6.6.m6.2.3.1.cmml" xref="S2.SS3.p6.6.m6.2.3.2"><ci id="S2.SS3.p6.6.m6.1.1.cmml" xref="S2.SS3.p6.6.m6.1.1">𝑙</ci><ci id="S2.SS3.p6.6.m6.2.2.cmml" xref="S2.SS3.p6.6.m6.2.2">𝑒</ci></interval></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.6.m6.2c">(l,e)</annotation></semantics></math> in path <math id="S2.SS3.p6.7.m7.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS3.p6.7.m7.1a"><mi id="S2.SS3.p6.7.m7.1.1" xref="S2.SS3.p6.7.m7.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.7.m7.1b"><ci id="S2.SS3.p6.7.m7.1.1.cmml" xref="S2.SS3.p6.7.m7.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.7.m7.1c">i</annotation></semantics></math> at iteration <math id="S2.SS3.p6.8.m8.1" class="ltx_Math" alttext="t" display="inline"><semantics id="S2.SS3.p6.8.m8.1a"><mi id="S2.SS3.p6.8.m8.1.1" xref="S2.SS3.p6.8.m8.1.1.cmml">t</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.8.m8.1b"><ci id="S2.SS3.p6.8.m8.1.1.cmml" xref="S2.SS3.p6.8.m8.1.1">𝑡</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.8.m8.1c">t</annotation></semantics></math> with <math id="S2.SS3.p6.9.m9.2" class="ltx_Math" alttext="\theta(l,e)^{t}_{i}" display="inline"><semantics id="S2.SS3.p6.9.m9.2a"><mrow id="S2.SS3.p6.9.m9.2.3" xref="S2.SS3.p6.9.m9.2.3.cmml"><mi id="S2.SS3.p6.9.m9.2.3.2" xref="S2.SS3.p6.9.m9.2.3.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="S2.SS3.p6.9.m9.2.3.1" xref="S2.SS3.p6.9.m9.2.3.1.cmml">​</mo><msubsup id="S2.SS3.p6.9.m9.2.3.3" xref="S2.SS3.p6.9.m9.2.3.3.cmml"><mrow id="S2.SS3.p6.9.m9.2.3.3.2.2.2" xref="S2.SS3.p6.9.m9.2.3.3.2.2.1.cmml"><mo stretchy="false" id="S2.SS3.p6.9.m9.2.3.3.2.2.2.1" xref="S2.SS3.p6.9.m9.2.3.3.2.2.1.cmml">(</mo><mi id="S2.SS3.p6.9.m9.1.1" xref="S2.SS3.p6.9.m9.1.1.cmml">l</mi><mo id="S2.SS3.p6.9.m9.2.3.3.2.2.2.2" xref="S2.SS3.p6.9.m9.2.3.3.2.2.1.cmml">,</mo><mi id="S2.SS3.p6.9.m9.2.2" xref="S2.SS3.p6.9.m9.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.9.m9.2.3.3.2.2.2.3" xref="S2.SS3.p6.9.m9.2.3.3.2.2.1.cmml">)</mo></mrow><mi id="S2.SS3.p6.9.m9.2.3.3.3" xref="S2.SS3.p6.9.m9.2.3.3.3.cmml">i</mi><mi id="S2.SS3.p6.9.m9.2.3.3.2.3" xref="S2.SS3.p6.9.m9.2.3.3.2.3.cmml">t</mi></msubsup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.9.m9.2b"><apply id="S2.SS3.p6.9.m9.2.3.cmml" xref="S2.SS3.p6.9.m9.2.3"><times id="S2.SS3.p6.9.m9.2.3.1.cmml" xref="S2.SS3.p6.9.m9.2.3.1"></times><ci id="S2.SS3.p6.9.m9.2.3.2.cmml" xref="S2.SS3.p6.9.m9.2.3.2">𝜃</ci><apply id="S2.SS3.p6.9.m9.2.3.3.cmml" xref="S2.SS3.p6.9.m9.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.9.m9.2.3.3.1.cmml" xref="S2.SS3.p6.9.m9.2.3.3">subscript</csymbol><apply id="S2.SS3.p6.9.m9.2.3.3.2.cmml" xref="S2.SS3.p6.9.m9.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.9.m9.2.3.3.2.1.cmml" xref="S2.SS3.p6.9.m9.2.3.3">superscript</csymbol><interval closure="open" id="S2.SS3.p6.9.m9.2.3.3.2.2.1.cmml" xref="S2.SS3.p6.9.m9.2.3.3.2.2.2"><ci id="S2.SS3.p6.9.m9.1.1.cmml" xref="S2.SS3.p6.9.m9.1.1">𝑙</ci><ci id="S2.SS3.p6.9.m9.2.2.cmml" xref="S2.SS3.p6.9.m9.2.2">𝑒</ci></interval><ci id="S2.SS3.p6.9.m9.2.3.3.2.3.cmml" xref="S2.SS3.p6.9.m9.2.3.3.2.3">𝑡</ci></apply><ci id="S2.SS3.p6.9.m9.2.3.3.3.cmml" xref="S2.SS3.p6.9.m9.2.3.3.3">𝑖</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.9.m9.2c">\theta(l,e)^{t}_{i}</annotation></semantics></math>. The collection of parameters used by the <math id="S2.SS3.p6.10.m10.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS3.p6.10.m10.1a"><mi id="S2.SS3.p6.10.m10.1.1" xref="S2.SS3.p6.10.m10.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.10.m10.1b"><ci id="S2.SS3.p6.10.m10.1.1.cmml" xref="S2.SS3.p6.10.m10.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.10.m10.1c">i</annotation></semantics></math>-th path (across all levels) is denoted by <math id="S2.SS3.p6.11.m11.1" class="ltx_Math" alttext="\theta^{t}_{i}" display="inline"><semantics id="S2.SS3.p6.11.m11.1a"><msubsup id="S2.SS3.p6.11.m11.1.1" xref="S2.SS3.p6.11.m11.1.1.cmml"><mi id="S2.SS3.p6.11.m11.1.1.2.2" xref="S2.SS3.p6.11.m11.1.1.2.2.cmml">θ</mi><mi id="S2.SS3.p6.11.m11.1.1.3" xref="S2.SS3.p6.11.m11.1.1.3.cmml">i</mi><mi id="S2.SS3.p6.11.m11.1.1.2.3" xref="S2.SS3.p6.11.m11.1.1.2.3.cmml">t</mi></msubsup><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.11.m11.1b"><apply id="S2.SS3.p6.11.m11.1.1.cmml" xref="S2.SS3.p6.11.m11.1.1"><csymbol cd="ambiguous" id="S2.SS3.p6.11.m11.1.1.1.cmml" xref="S2.SS3.p6.11.m11.1.1">subscript</csymbol><apply id="S2.SS3.p6.11.m11.1.1.2.cmml" xref="S2.SS3.p6.11.m11.1.1"><csymbol cd="ambiguous" id="S2.SS3.p6.11.m11.1.1.2.1.cmml" xref="S2.SS3.p6.11.m11.1.1">superscript</csymbol><ci id="S2.SS3.p6.11.m11.1.1.2.2.cmml" xref="S2.SS3.p6.11.m11.1.1.2.2">𝜃</ci><ci id="S2.SS3.p6.11.m11.1.1.2.3.cmml" xref="S2.SS3.p6.11.m11.1.1.2.3">𝑡</ci></apply><ci id="S2.SS3.p6.11.m11.1.1.3.cmml" xref="S2.SS3.p6.11.m11.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.11.m11.1c">\theta^{t}_{i}</annotation></semantics></math>. When modules across paths are synchronized, the parameter values of module <math id="S2.SS3.p6.12.m12.2" class="ltx_Math" alttext="(l,e)" display="inline"><semantics id="S2.SS3.p6.12.m12.2a"><mrow id="S2.SS3.p6.12.m12.2.3.2" xref="S2.SS3.p6.12.m12.2.3.1.cmml"><mo stretchy="false" id="S2.SS3.p6.12.m12.2.3.2.1" xref="S2.SS3.p6.12.m12.2.3.1.cmml">(</mo><mi id="S2.SS3.p6.12.m12.1.1" xref="S2.SS3.p6.12.m12.1.1.cmml">l</mi><mo id="S2.SS3.p6.12.m12.2.3.2.2" xref="S2.SS3.p6.12.m12.2.3.1.cmml">,</mo><mi id="S2.SS3.p6.12.m12.2.2" xref="S2.SS3.p6.12.m12.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.12.m12.2.3.2.3" xref="S2.SS3.p6.12.m12.2.3.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.12.m12.2b"><interval closure="open" id="S2.SS3.p6.12.m12.2.3.1.cmml" xref="S2.SS3.p6.12.m12.2.3.2"><ci id="S2.SS3.p6.12.m12.1.1.cmml" xref="S2.SS3.p6.12.m12.1.1">𝑙</ci><ci id="S2.SS3.p6.12.m12.2.2.cmml" xref="S2.SS3.p6.12.m12.2.2">𝑒</ci></interval></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.12.m12.2c">(l,e)</annotation></semantics></math> is the same across all paths, and therefore we denote the global copy of the parameters of module <math id="S2.SS3.p6.13.m13.2" class="ltx_Math" alttext="(l,e)" display="inline"><semantics id="S2.SS3.p6.13.m13.2a"><mrow id="S2.SS3.p6.13.m13.2.3.2" xref="S2.SS3.p6.13.m13.2.3.1.cmml"><mo stretchy="false" id="S2.SS3.p6.13.m13.2.3.2.1" xref="S2.SS3.p6.13.m13.2.3.1.cmml">(</mo><mi id="S2.SS3.p6.13.m13.1.1" xref="S2.SS3.p6.13.m13.1.1.cmml">l</mi><mo id="S2.SS3.p6.13.m13.2.3.2.2" xref="S2.SS3.p6.13.m13.2.3.1.cmml">,</mo><mi id="S2.SS3.p6.13.m13.2.2" xref="S2.SS3.p6.13.m13.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.13.m13.2.3.2.3" xref="S2.SS3.p6.13.m13.2.3.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.13.m13.2b"><interval closure="open" id="S2.SS3.p6.13.m13.2.3.1.cmml" xref="S2.SS3.p6.13.m13.2.3.2"><ci id="S2.SS3.p6.13.m13.1.1.cmml" xref="S2.SS3.p6.13.m13.1.1">𝑙</ci><ci id="S2.SS3.p6.13.m13.2.2.cmml" xref="S2.SS3.p6.13.m13.2.2">𝑒</ci></interval></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.13.m13.2c">(l,e)</annotation></semantics></math> at iteration <math id="S2.SS3.p6.14.m14.1" class="ltx_Math" alttext="t" display="inline"><semantics id="S2.SS3.p6.14.m14.1a"><mi id="S2.SS3.p6.14.m14.1.1" xref="S2.SS3.p6.14.m14.1.1.cmml">t</mi><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.14.m14.1b"><ci id="S2.SS3.p6.14.m14.1.1.cmml" xref="S2.SS3.p6.14.m14.1.1">𝑡</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.14.m14.1c">t</annotation></semantics></math> by omitting the path index, as in <math id="S2.SS3.p6.15.m15.2" class="ltx_Math" alttext="\theta(l,e)^{t}" display="inline"><semantics id="S2.SS3.p6.15.m15.2a"><mrow id="S2.SS3.p6.15.m15.2.3" xref="S2.SS3.p6.15.m15.2.3.cmml"><mi id="S2.SS3.p6.15.m15.2.3.2" xref="S2.SS3.p6.15.m15.2.3.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="S2.SS3.p6.15.m15.2.3.1" xref="S2.SS3.p6.15.m15.2.3.1.cmml">​</mo><msup id="S2.SS3.p6.15.m15.2.3.3" xref="S2.SS3.p6.15.m15.2.3.3.cmml"><mrow id="S2.SS3.p6.15.m15.2.3.3.2.2" xref="S2.SS3.p6.15.m15.2.3.3.2.1.cmml"><mo stretchy="false" id="S2.SS3.p6.15.m15.2.3.3.2.2.1" xref="S2.SS3.p6.15.m15.2.3.3.2.1.cmml">(</mo><mi id="S2.SS3.p6.15.m15.1.1" xref="S2.SS3.p6.15.m15.1.1.cmml">l</mi><mo id="S2.SS3.p6.15.m15.2.3.3.2.2.2" xref="S2.SS3.p6.15.m15.2.3.3.2.1.cmml">,</mo><mi id="S2.SS3.p6.15.m15.2.2" xref="S2.SS3.p6.15.m15.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.15.m15.2.3.3.2.2.3" xref="S2.SS3.p6.15.m15.2.3.3.2.1.cmml">)</mo></mrow><mi id="S2.SS3.p6.15.m15.2.3.3.3" xref="S2.SS3.p6.15.m15.2.3.3.3.cmml">t</mi></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.15.m15.2b"><apply id="S2.SS3.p6.15.m15.2.3.cmml" xref="S2.SS3.p6.15.m15.2.3"><times id="S2.SS3.p6.15.m15.2.3.1.cmml" xref="S2.SS3.p6.15.m15.2.3.1"></times><ci id="S2.SS3.p6.15.m15.2.3.2.cmml" xref="S2.SS3.p6.15.m15.2.3.2">𝜃</ci><apply id="S2.SS3.p6.15.m15.2.3.3.cmml" xref="S2.SS3.p6.15.m15.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.15.m15.2.3.3.1.cmml" xref="S2.SS3.p6.15.m15.2.3.3">superscript</csymbol><interval closure="open" id="S2.SS3.p6.15.m15.2.3.3.2.1.cmml" xref="S2.SS3.p6.15.m15.2.3.3.2.2"><ci id="S2.SS3.p6.15.m15.1.1.cmml" xref="S2.SS3.p6.15.m15.1.1">𝑙</ci><ci id="S2.SS3.p6.15.m15.2.2.cmml" xref="S2.SS3.p6.15.m15.2.2">𝑒</ci></interval><ci id="S2.SS3.p6.15.m15.2.3.3.3.cmml" xref="S2.SS3.p6.15.m15.2.3.3.3">𝑡</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.15.m15.2c">\theta(l,e)^{t}</annotation></semantics></math>. Similarly, we denote local and global gradients of module <math id="S2.SS3.p6.16.m16.2" class="ltx_Math" alttext="(l,e)" display="inline"><semantics id="S2.SS3.p6.16.m16.2a"><mrow id="S2.SS3.p6.16.m16.2.3.2" xref="S2.SS3.p6.16.m16.2.3.1.cmml"><mo stretchy="false" id="S2.SS3.p6.16.m16.2.3.2.1" xref="S2.SS3.p6.16.m16.2.3.1.cmml">(</mo><mi id="S2.SS3.p6.16.m16.1.1" xref="S2.SS3.p6.16.m16.1.1.cmml">l</mi><mo id="S2.SS3.p6.16.m16.2.3.2.2" xref="S2.SS3.p6.16.m16.2.3.1.cmml">,</mo><mi id="S2.SS3.p6.16.m16.2.2" xref="S2.SS3.p6.16.m16.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.16.m16.2.3.2.3" xref="S2.SS3.p6.16.m16.2.3.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.16.m16.2b"><interval closure="open" id="S2.SS3.p6.16.m16.2.3.1.cmml" xref="S2.SS3.p6.16.m16.2.3.2"><ci id="S2.SS3.p6.16.m16.1.1.cmml" xref="S2.SS3.p6.16.m16.1.1">𝑙</ci><ci id="S2.SS3.p6.16.m16.2.2.cmml" xref="S2.SS3.p6.16.m16.2.2">𝑒</ci></interval></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.16.m16.2c">(l,e)</annotation></semantics></math> with <math id="S2.SS3.p6.17.m17.2" class="ltx_Math" alttext="\Delta(l,e)^{t}_{i}" display="inline"><semantics id="S2.SS3.p6.17.m17.2a"><mrow id="S2.SS3.p6.17.m17.2.3" xref="S2.SS3.p6.17.m17.2.3.cmml"><mi mathvariant="normal" id="S2.SS3.p6.17.m17.2.3.2" xref="S2.SS3.p6.17.m17.2.3.2.cmml">Δ</mi><mo lspace="0em" rspace="0em" id="S2.SS3.p6.17.m17.2.3.1" xref="S2.SS3.p6.17.m17.2.3.1.cmml">​</mo><msubsup id="S2.SS3.p6.17.m17.2.3.3" xref="S2.SS3.p6.17.m17.2.3.3.cmml"><mrow id="S2.SS3.p6.17.m17.2.3.3.2.2.2" xref="S2.SS3.p6.17.m17.2.3.3.2.2.1.cmml"><mo stretchy="false" id="S2.SS3.p6.17.m17.2.3.3.2.2.2.1" xref="S2.SS3.p6.17.m17.2.3.3.2.2.1.cmml">(</mo><mi id="S2.SS3.p6.17.m17.1.1" xref="S2.SS3.p6.17.m17.1.1.cmml">l</mi><mo id="S2.SS3.p6.17.m17.2.3.3.2.2.2.2" xref="S2.SS3.p6.17.m17.2.3.3.2.2.1.cmml">,</mo><mi id="S2.SS3.p6.17.m17.2.2" xref="S2.SS3.p6.17.m17.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.17.m17.2.3.3.2.2.2.3" xref="S2.SS3.p6.17.m17.2.3.3.2.2.1.cmml">)</mo></mrow><mi id="S2.SS3.p6.17.m17.2.3.3.3" xref="S2.SS3.p6.17.m17.2.3.3.3.cmml">i</mi><mi id="S2.SS3.p6.17.m17.2.3.3.2.3" xref="S2.SS3.p6.17.m17.2.3.3.2.3.cmml">t</mi></msubsup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.17.m17.2b"><apply id="S2.SS3.p6.17.m17.2.3.cmml" xref="S2.SS3.p6.17.m17.2.3"><times id="S2.SS3.p6.17.m17.2.3.1.cmml" xref="S2.SS3.p6.17.m17.2.3.1"></times><ci id="S2.SS3.p6.17.m17.2.3.2.cmml" xref="S2.SS3.p6.17.m17.2.3.2">Δ</ci><apply id="S2.SS3.p6.17.m17.2.3.3.cmml" xref="S2.SS3.p6.17.m17.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.17.m17.2.3.3.1.cmml" xref="S2.SS3.p6.17.m17.2.3.3">subscript</csymbol><apply id="S2.SS3.p6.17.m17.2.3.3.2.cmml" xref="S2.SS3.p6.17.m17.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.17.m17.2.3.3.2.1.cmml" xref="S2.SS3.p6.17.m17.2.3.3">superscript</csymbol><interval closure="open" id="S2.SS3.p6.17.m17.2.3.3.2.2.1.cmml" xref="S2.SS3.p6.17.m17.2.3.3.2.2.2"><ci id="S2.SS3.p6.17.m17.1.1.cmml" xref="S2.SS3.p6.17.m17.1.1">𝑙</ci><ci id="S2.SS3.p6.17.m17.2.2.cmml" xref="S2.SS3.p6.17.m17.2.2">𝑒</ci></interval><ci id="S2.SS3.p6.17.m17.2.3.3.2.3.cmml" xref="S2.SS3.p6.17.m17.2.3.3.2.3">𝑡</ci></apply><ci id="S2.SS3.p6.17.m17.2.3.3.3.cmml" xref="S2.SS3.p6.17.m17.2.3.3.3">𝑖</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.17.m17.2c">\Delta(l,e)^{t}_{i}</annotation></semantics></math> and <math id="S2.SS3.p6.18.m18.2" class="ltx_Math" alttext="\Delta(l,e)^{t}" display="inline"><semantics id="S2.SS3.p6.18.m18.2a"><mrow id="S2.SS3.p6.18.m18.2.3" xref="S2.SS3.p6.18.m18.2.3.cmml"><mi mathvariant="normal" id="S2.SS3.p6.18.m18.2.3.2" xref="S2.SS3.p6.18.m18.2.3.2.cmml">Δ</mi><mo lspace="0em" rspace="0em" id="S2.SS3.p6.18.m18.2.3.1" xref="S2.SS3.p6.18.m18.2.3.1.cmml">​</mo><msup id="S2.SS3.p6.18.m18.2.3.3" xref="S2.SS3.p6.18.m18.2.3.3.cmml"><mrow id="S2.SS3.p6.18.m18.2.3.3.2.2" xref="S2.SS3.p6.18.m18.2.3.3.2.1.cmml"><mo stretchy="false" id="S2.SS3.p6.18.m18.2.3.3.2.2.1" xref="S2.SS3.p6.18.m18.2.3.3.2.1.cmml">(</mo><mi id="S2.SS3.p6.18.m18.1.1" xref="S2.SS3.p6.18.m18.1.1.cmml">l</mi><mo id="S2.SS3.p6.18.m18.2.3.3.2.2.2" xref="S2.SS3.p6.18.m18.2.3.3.2.1.cmml">,</mo><mi id="S2.SS3.p6.18.m18.2.2" xref="S2.SS3.p6.18.m18.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS3.p6.18.m18.2.3.3.2.2.3" xref="S2.SS3.p6.18.m18.2.3.3.2.1.cmml">)</mo></mrow><mi id="S2.SS3.p6.18.m18.2.3.3.3" xref="S2.SS3.p6.18.m18.2.3.3.3.cmml">t</mi></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p6.18.m18.2b"><apply id="S2.SS3.p6.18.m18.2.3.cmml" xref="S2.SS3.p6.18.m18.2.3"><times id="S2.SS3.p6.18.m18.2.3.1.cmml" xref="S2.SS3.p6.18.m18.2.3.1"></times><ci id="S2.SS3.p6.18.m18.2.3.2.cmml" xref="S2.SS3.p6.18.m18.2.3.2">Δ</ci><apply id="S2.SS3.p6.18.m18.2.3.3.cmml" xref="S2.SS3.p6.18.m18.2.3.3"><csymbol cd="ambiguous" id="S2.SS3.p6.18.m18.2.3.3.1.cmml" xref="S2.SS3.p6.18.m18.2.3.3">superscript</csymbol><interval closure="open" id="S2.SS3.p6.18.m18.2.3.3.2.1.cmml" xref="S2.SS3.p6.18.m18.2.3.3.2.2"><ci id="S2.SS3.p6.18.m18.1.1.cmml" xref="S2.SS3.p6.18.m18.1.1">𝑙</ci><ci id="S2.SS3.p6.18.m18.2.2.cmml" xref="S2.SS3.p6.18.m18.2.2">𝑒</ci></interval><ci id="S2.SS3.p6.18.m18.2.3.3.3.cmml" xref="S2.SS3.p6.18.m18.2.3.3.3">𝑡</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p6.18.m18.2c">\Delta(l,e)^{t}</annotation></semantics></math>, respectively.</p>
</div>
</section>
<section id="S2.SS4" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.4 </span>Coarse Routing</h3>

<div id="S2.SS4.p1" class="ltx_para">
<p id="S2.SS4.p1.2" class="ltx_p">In this section we describe how we route sequences.
The basic intuition is to use some context to decide which path is most suitable to process a given sequence. Since we focus on language modeling, we use as context the first <math id="S2.SS4.p1.1.m1.1" class="ltx_Math" alttext="32" display="inline"><semantics id="S2.SS4.p1.1.m1.1a"><mn id="S2.SS4.p1.1.m1.1.1" xref="S2.SS4.p1.1.m1.1.1.cmml">32</mn><annotation-xml encoding="MathML-Content" id="S2.SS4.p1.1.m1.1b"><cn type="integer" id="S2.SS4.p1.1.m1.1.1.cmml" xref="S2.SS4.p1.1.m1.1.1">32</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p1.1.m1.1c">32</annotation></semantics></math> tokens of each sequence. At training time, we then use the remaining tokens for learning the model parameters. At test time, we use the remaining tokens to compute perplexity on a held-out validation set. All methods, including dense baselines, will be evaluated in the same way, by calculating perplexity using all but the first <math id="S2.SS4.p1.2.m2.1" class="ltx_Math" alttext="32" display="inline"><semantics id="S2.SS4.p1.2.m2.1a"><mn id="S2.SS4.p1.2.m2.1.1" xref="S2.SS4.p1.2.m2.1.1.cmml">32</mn><annotation-xml encoding="MathML-Content" id="S2.SS4.p1.2.m2.1b"><cn type="integer" id="S2.SS4.p1.2.m2.1.1.cmml" xref="S2.SS4.p1.2.m2.1.1">32</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p1.2.m2.1c">32</annotation></semantics></math> tokens of each sequence, which was used to determine the routing decision of our models.</p>
</div>
<div id="S2.SS4.p2" class="ltx_para">
<p id="S2.SS4.p2.2" class="ltx_p">Notice that the outcome of routing is assigning each sequence to a certain path in the network. Collectively, all sequences that are assigned to a particular path form a shard of the original dataset. In this work, there is a one-to-one association between paths and corresponding shards; path <math id="S2.SS4.p2.1.m1.1" class="ltx_Math" alttext="\pi_{i}" display="inline"><semantics id="S2.SS4.p2.1.m1.1a"><msub id="S2.SS4.p2.1.m1.1.1" xref="S2.SS4.p2.1.m1.1.1.cmml"><mi id="S2.SS4.p2.1.m1.1.1.2" xref="S2.SS4.p2.1.m1.1.1.2.cmml">π</mi><mi id="S2.SS4.p2.1.m1.1.1.3" xref="S2.SS4.p2.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS4.p2.1.m1.1b"><apply id="S2.SS4.p2.1.m1.1.1.cmml" xref="S2.SS4.p2.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS4.p2.1.m1.1.1.1.cmml" xref="S2.SS4.p2.1.m1.1.1">subscript</csymbol><ci id="S2.SS4.p2.1.m1.1.1.2.cmml" xref="S2.SS4.p2.1.m1.1.1.2">𝜋</ci><ci id="S2.SS4.p2.1.m1.1.1.3.cmml" xref="S2.SS4.p2.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p2.1.m1.1c">\pi_{i}</annotation></semantics></math> is associated to shard <math id="S2.SS4.p2.2.m2.1" class="ltx_Math" alttext="\mathcal{D}_{i}" display="inline"><semantics id="S2.SS4.p2.2.m2.1a"><msub id="S2.SS4.p2.2.m2.1.1" xref="S2.SS4.p2.2.m2.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS4.p2.2.m2.1.1.2" xref="S2.SS4.p2.2.m2.1.1.2.cmml">𝒟</mi><mi id="S2.SS4.p2.2.m2.1.1.3" xref="S2.SS4.p2.2.m2.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS4.p2.2.m2.1b"><apply id="S2.SS4.p2.2.m2.1.1.cmml" xref="S2.SS4.p2.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS4.p2.2.m2.1.1.1.cmml" xref="S2.SS4.p2.2.m2.1.1">subscript</csymbol><ci id="S2.SS4.p2.2.m2.1.1.2.cmml" xref="S2.SS4.p2.2.m2.1.1.2">𝒟</ci><ci id="S2.SS4.p2.2.m2.1.1.3.cmml" xref="S2.SS4.p2.2.m2.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p2.2.m2.1c">\mathcal{D}_{i}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS4.p3" class="ltx_para">
<p id="S2.SS4.p3.1" class="ltx_p">Next, we describe how we use this prefix of <math id="S2.SS4.p3.1.m1.1" class="ltx_Math" alttext="32" display="inline"><semantics id="S2.SS4.p3.1.m1.1a"><mn id="S2.SS4.p3.1.m1.1.1" xref="S2.SS4.p3.1.m1.1.1.cmml">32</mn><annotation-xml encoding="MathML-Content" id="S2.SS4.p3.1.m1.1b"><cn type="integer" id="S2.SS4.p3.1.m1.1.1.cmml" xref="S2.SS4.p3.1.m1.1.1">32</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p3.1.m1.1c">32</annotation></semantics></math> tokens to route, as well as other variants of routing.</p>
</div>
<figure id="S2.F3" class="ltx_figure"><img src="/html/2403.10616/assets/x4.png" id="S2.F3.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="332" height="98" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S2.F3.7.3.1" class="ltx_text" style="font-size:90%;">Figure 3</span>: </span><span id="S2.F3.4.2" class="ltx_text ltx_font_bold" style="font-size:90%;">routing More Frequently at Test-Time<span id="S2.F3.4.2.2" class="ltx_text ltx_font_medium">: At training time (left panel), the router selects the path <math id="S2.F3.3.1.1.m1.1" class="ltx_Math" alttext="\pi_{i}" display="inline"><semantics id="S2.F3.3.1.1.m1.1b"><msub id="S2.F3.3.1.1.m1.1.1" xref="S2.F3.3.1.1.m1.1.1.cmml"><mi id="S2.F3.3.1.1.m1.1.1.2" xref="S2.F3.3.1.1.m1.1.1.2.cmml">π</mi><mi id="S2.F3.3.1.1.m1.1.1.3" xref="S2.F3.3.1.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.F3.3.1.1.m1.1c"><apply id="S2.F3.3.1.1.m1.1.1.cmml" xref="S2.F3.3.1.1.m1.1.1"><csymbol cd="ambiguous" id="S2.F3.3.1.1.m1.1.1.1.cmml" xref="S2.F3.3.1.1.m1.1.1">subscript</csymbol><ci id="S2.F3.3.1.1.m1.1.1.2.cmml" xref="S2.F3.3.1.1.m1.1.1.2">𝜋</ci><ci id="S2.F3.3.1.1.m1.1.1.3.cmml" xref="S2.F3.3.1.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F3.3.1.1.m1.1d">\pi_{i}</annotation></semantics></math> using the prefix <math id="S2.F3.4.2.2.m2.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.F3.4.2.2.m2.1b"><mi id="S2.F3.4.2.2.m2.1.1" xref="S2.F3.4.2.2.m2.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.F3.4.2.2.m2.1c"><ci id="S2.F3.4.2.2.m2.1.1.cmml" xref="S2.F3.4.2.2.m2.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.F3.4.2.2.m2.1d">z</annotation></semantics></math>. We train the chosen path on the whole sequence using the usual language modeling loss. At test time (right panel), the path selected by the router given the prefix is used to score the next chunk of tokens. Then, we re-use the router to choose the most likely path given the new chunk. This process repeats until the whole sequence has been scored.</span></span></figcaption>
</figure>
<section id="S2.SS4.SSS1" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.4.1 </span>Generative routing</h4>

<div id="S2.SS4.SSS1.p1" class="ltx_para">
<p id="S2.SS4.SSS1.p1.11" class="ltx_p">In generative routing, the decision about which shard to assign a sequence to is not informed by the task at hand, namely language modeling. Instead, the choice is based on minimizing feature reconstruction error. Given a representation <math id="S2.SS4.SSS1.p1.1.m1.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS1.p1.1.m1.1a"><mi id="S2.SS4.SSS1.p1.1.m1.1.1" xref="S2.SS4.SSS1.p1.1.m1.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.1.m1.1b"><ci id="S2.SS4.SSS1.p1.1.m1.1.1.cmml" xref="S2.SS4.SSS1.p1.1.m1.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.1.m1.1c">z</annotation></semantics></math> of the first <math id="S2.SS4.SSS1.p1.2.m2.1" class="ltx_Math" alttext="32" display="inline"><semantics id="S2.SS4.SSS1.p1.2.m2.1a"><mn id="S2.SS4.SSS1.p1.2.m2.1.1" xref="S2.SS4.SSS1.p1.2.m2.1.1.cmml">32</mn><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.2.m2.1b"><cn type="integer" id="S2.SS4.SSS1.p1.2.m2.1.1.cmml" xref="S2.SS4.SSS1.p1.2.m2.1.1">32</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.2.m2.1c">32</annotation></semantics></math> tokens of a sequence (context), we perform <math id="S2.SS4.SSS1.p1.3.m3.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS1.p1.3.m3.1a"><mi id="S2.SS4.SSS1.p1.3.m3.1.1" xref="S2.SS4.SSS1.p1.3.m3.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.3.m3.1b"><ci id="S2.SS4.SSS1.p1.3.m3.1.1.cmml" xref="S2.SS4.SSS1.p1.3.m3.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.3.m3.1c">k</annotation></semantics></math>-Means on the features <math id="S2.SS4.SSS1.p1.4.m4.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS1.p1.4.m4.1a"><mi id="S2.SS4.SSS1.p1.4.m4.1.1" xref="S2.SS4.SSS1.p1.4.m4.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.4.m4.1b"><ci id="S2.SS4.SSS1.p1.4.m4.1.1.cmml" xref="S2.SS4.SSS1.p1.4.m4.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.4.m4.1c">z</annotation></semantics></math> of each sequence, and then we use the <math id="S2.SS4.SSS1.p1.5.m5.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS1.p1.5.m5.1a"><mi id="S2.SS4.SSS1.p1.5.m5.1.1" xref="S2.SS4.SSS1.p1.5.m5.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.5.m5.1b"><ci id="S2.SS4.SSS1.p1.5.m5.1.1.cmml" xref="S2.SS4.SSS1.p1.5.m5.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.5.m5.1c">k</annotation></semantics></math>-means assignment algorithm to shard the data into <math id="S2.SS4.SSS1.p1.6.m6.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS1.p1.6.m6.1a"><mi id="S2.SS4.SSS1.p1.6.m6.1.1" xref="S2.SS4.SSS1.p1.6.m6.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.6.m6.1b"><ci id="S2.SS4.SSS1.p1.6.m6.1.1.cmml" xref="S2.SS4.SSS1.p1.6.m6.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.6.m6.1c">k</annotation></semantics></math> shards. If <math id="S2.SS4.SSS1.p1.7.m7.3" class="ltx_Math" alttext="\{c_{1},\dots,c_{k}\}" display="inline"><semantics id="S2.SS4.SSS1.p1.7.m7.3a"><mrow id="S2.SS4.SSS1.p1.7.m7.3.3.2" xref="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml"><mo stretchy="false" id="S2.SS4.SSS1.p1.7.m7.3.3.2.3" xref="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml">{</mo><msub id="S2.SS4.SSS1.p1.7.m7.2.2.1.1" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1.cmml"><mi id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.2" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1.2.cmml">c</mi><mn id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.3" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1.3.cmml">1</mn></msub><mo id="S2.SS4.SSS1.p1.7.m7.3.3.2.4" xref="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml">,</mo><mi mathvariant="normal" id="S2.SS4.SSS1.p1.7.m7.1.1" xref="S2.SS4.SSS1.p1.7.m7.1.1.cmml">…</mi><mo id="S2.SS4.SSS1.p1.7.m7.3.3.2.5" xref="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml">,</mo><msub id="S2.SS4.SSS1.p1.7.m7.3.3.2.2" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2.cmml"><mi id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.2" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2.2.cmml">c</mi><mi id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.3" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2.3.cmml">k</mi></msub><mo stretchy="false" id="S2.SS4.SSS1.p1.7.m7.3.3.2.6" xref="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml">}</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.7.m7.3b"><set id="S2.SS4.SSS1.p1.7.m7.3.3.3.cmml" xref="S2.SS4.SSS1.p1.7.m7.3.3.2"><apply id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.cmml" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1"><csymbol cd="ambiguous" id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.1.cmml" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1">subscript</csymbol><ci id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.2.cmml" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1.2">𝑐</ci><cn type="integer" id="S2.SS4.SSS1.p1.7.m7.2.2.1.1.3.cmml" xref="S2.SS4.SSS1.p1.7.m7.2.2.1.1.3">1</cn></apply><ci id="S2.SS4.SSS1.p1.7.m7.1.1.cmml" xref="S2.SS4.SSS1.p1.7.m7.1.1">…</ci><apply id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.cmml" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2"><csymbol cd="ambiguous" id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.1.cmml" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2">subscript</csymbol><ci id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.2.cmml" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2.2">𝑐</ci><ci id="S2.SS4.SSS1.p1.7.m7.3.3.2.2.3.cmml" xref="S2.SS4.SSS1.p1.7.m7.3.3.2.2.3">𝑘</ci></apply></set></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.7.m7.3c">\{c_{1},\dots,c_{k}\}</annotation></semantics></math> are the <math id="S2.SS4.SSS1.p1.8.m8.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS1.p1.8.m8.1a"><mi id="S2.SS4.SSS1.p1.8.m8.1.1" xref="S2.SS4.SSS1.p1.8.m8.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.8.m8.1b"><ci id="S2.SS4.SSS1.p1.8.m8.1.1.cmml" xref="S2.SS4.SSS1.p1.8.m8.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.8.m8.1c">k</annotation></semantics></math> prototypes learned by <math id="S2.SS4.SSS1.p1.9.m9.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS1.p1.9.m9.1a"><mi id="S2.SS4.SSS1.p1.9.m9.1.1" xref="S2.SS4.SSS1.p1.9.m9.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.9.m9.1b"><ci id="S2.SS4.SSS1.p1.9.m9.1.1.cmml" xref="S2.SS4.SSS1.p1.9.m9.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.9.m9.1c">k</annotation></semantics></math>-means, the sequence with prefix <math id="S2.SS4.SSS1.p1.10.m10.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS1.p1.10.m10.1a"><mi id="S2.SS4.SSS1.p1.10.m10.1.1" xref="S2.SS4.SSS1.p1.10.m10.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.10.m10.1b"><ci id="S2.SS4.SSS1.p1.10.m10.1.1.cmml" xref="S2.SS4.SSS1.p1.10.m10.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.10.m10.1c">z</annotation></semantics></math> is assigned to shard <math id="S2.SS4.SSS1.p1.11.m11.1" class="ltx_Math" alttext="\mathcal{D}_{r(z)}" display="inline"><semantics id="S2.SS4.SSS1.p1.11.m11.1a"><msub id="S2.SS4.SSS1.p1.11.m11.1.2" xref="S2.SS4.SSS1.p1.11.m11.1.2.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS4.SSS1.p1.11.m11.1.2.2" xref="S2.SS4.SSS1.p1.11.m11.1.2.2.cmml">𝒟</mi><mrow id="S2.SS4.SSS1.p1.11.m11.1.1.1" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.cmml"><mi id="S2.SS4.SSS1.p1.11.m11.1.1.1.3" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.3.cmml">r</mi><mo lspace="0em" rspace="0em" id="S2.SS4.SSS1.p1.11.m11.1.1.1.2" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.2.cmml">​</mo><mrow id="S2.SS4.SSS1.p1.11.m11.1.1.1.4.2" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.SSS1.p1.11.m11.1.1.1.4.2.1" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.cmml">(</mo><mi id="S2.SS4.SSS1.p1.11.m11.1.1.1.1" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.1.cmml">z</mi><mo stretchy="false" id="S2.SS4.SSS1.p1.11.m11.1.1.1.4.2.2" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.cmml">)</mo></mrow></mrow></msub><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS1.p1.11.m11.1b"><apply id="S2.SS4.SSS1.p1.11.m11.1.2.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.2"><csymbol cd="ambiguous" id="S2.SS4.SSS1.p1.11.m11.1.2.1.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.2">subscript</csymbol><ci id="S2.SS4.SSS1.p1.11.m11.1.2.2.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.2.2">𝒟</ci><apply id="S2.SS4.SSS1.p1.11.m11.1.1.1.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.1.1"><times id="S2.SS4.SSS1.p1.11.m11.1.1.1.2.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.2"></times><ci id="S2.SS4.SSS1.p1.11.m11.1.1.1.3.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.3">𝑟</ci><ci id="S2.SS4.SSS1.p1.11.m11.1.1.1.1.cmml" xref="S2.SS4.SSS1.p1.11.m11.1.1.1.1">𝑧</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS1.p1.11.m11.1c">\mathcal{D}_{r(z)}</annotation></semantics></math> via:</p>
<table id="S2.E1" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="S2.E1.m1.4" class="ltx_Math" alttext="r(z)=\arg\min_{i\in[1,k]}||z-c_{i}||^{2}." display="block"><semantics id="S2.E1.m1.4a"><mrow id="S2.E1.m1.4.4.1" xref="S2.E1.m1.4.4.1.1.cmml"><mrow id="S2.E1.m1.4.4.1.1" xref="S2.E1.m1.4.4.1.1.cmml"><mrow id="S2.E1.m1.4.4.1.1.3" xref="S2.E1.m1.4.4.1.1.3.cmml"><mi id="S2.E1.m1.4.4.1.1.3.2" xref="S2.E1.m1.4.4.1.1.3.2.cmml">r</mi><mo lspace="0em" rspace="0em" id="S2.E1.m1.4.4.1.1.3.1" xref="S2.E1.m1.4.4.1.1.3.1.cmml">​</mo><mrow id="S2.E1.m1.4.4.1.1.3.3.2" xref="S2.E1.m1.4.4.1.1.3.cmml"><mo stretchy="false" id="S2.E1.m1.4.4.1.1.3.3.2.1" xref="S2.E1.m1.4.4.1.1.3.cmml">(</mo><mi id="S2.E1.m1.3.3" xref="S2.E1.m1.3.3.cmml">z</mi><mo stretchy="false" id="S2.E1.m1.4.4.1.1.3.3.2.2" xref="S2.E1.m1.4.4.1.1.3.cmml">)</mo></mrow></mrow><mo id="S2.E1.m1.4.4.1.1.2" xref="S2.E1.m1.4.4.1.1.2.cmml">=</mo><mrow id="S2.E1.m1.4.4.1.1.1" xref="S2.E1.m1.4.4.1.1.1.cmml"><mrow id="S2.E1.m1.4.4.1.1.1.3" xref="S2.E1.m1.4.4.1.1.1.3.cmml"><mi id="S2.E1.m1.4.4.1.1.1.3.1" xref="S2.E1.m1.4.4.1.1.1.3.1.cmml">arg</mi><mo lspace="0.167em" id="S2.E1.m1.4.4.1.1.1.3a" xref="S2.E1.m1.4.4.1.1.1.3.cmml">⁡</mo><munder id="S2.E1.m1.4.4.1.1.1.3.2" xref="S2.E1.m1.4.4.1.1.1.3.2.cmml"><mi id="S2.E1.m1.4.4.1.1.1.3.2.2" xref="S2.E1.m1.4.4.1.1.1.3.2.2.cmml">min</mi><mrow id="S2.E1.m1.2.2.2" xref="S2.E1.m1.2.2.2.cmml"><mi id="S2.E1.m1.2.2.2.4" xref="S2.E1.m1.2.2.2.4.cmml">i</mi><mo id="S2.E1.m1.2.2.2.3" xref="S2.E1.m1.2.2.2.3.cmml">∈</mo><mrow id="S2.E1.m1.2.2.2.5.2" xref="S2.E1.m1.2.2.2.5.1.cmml"><mo stretchy="false" id="S2.E1.m1.2.2.2.5.2.1" xref="S2.E1.m1.2.2.2.5.1.cmml">[</mo><mn id="S2.E1.m1.1.1.1.1" xref="S2.E1.m1.1.1.1.1.cmml">1</mn><mo id="S2.E1.m1.2.2.2.5.2.2" xref="S2.E1.m1.2.2.2.5.1.cmml">,</mo><mi id="S2.E1.m1.2.2.2.2" xref="S2.E1.m1.2.2.2.2.cmml">k</mi><mo stretchy="false" id="S2.E1.m1.2.2.2.5.2.3" xref="S2.E1.m1.2.2.2.5.1.cmml">]</mo></mrow></mrow></munder></mrow><mo lspace="0em" rspace="0em" id="S2.E1.m1.4.4.1.1.1.2" xref="S2.E1.m1.4.4.1.1.1.2.cmml">​</mo><msup id="S2.E1.m1.4.4.1.1.1.1" xref="S2.E1.m1.4.4.1.1.1.1.cmml"><mrow id="S2.E1.m1.4.4.1.1.1.1.1.1" xref="S2.E1.m1.4.4.1.1.1.1.1.2.cmml"><mo stretchy="false" id="S2.E1.m1.4.4.1.1.1.1.1.1.2" xref="S2.E1.m1.4.4.1.1.1.1.1.2.1.cmml">‖</mo><mrow id="S2.E1.m1.4.4.1.1.1.1.1.1.1" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.cmml"><mi id="S2.E1.m1.4.4.1.1.1.1.1.1.1.2" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.2.cmml">z</mi><mo id="S2.E1.m1.4.4.1.1.1.1.1.1.1.1" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.1.cmml">−</mo><msub id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.cmml"><mi id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.2" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.2.cmml">c</mi><mi id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.3" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.3.cmml">i</mi></msub></mrow><mo stretchy="false" id="S2.E1.m1.4.4.1.1.1.1.1.1.3" xref="S2.E1.m1.4.4.1.1.1.1.1.2.1.cmml">‖</mo></mrow><mn id="S2.E1.m1.4.4.1.1.1.1.3" xref="S2.E1.m1.4.4.1.1.1.1.3.cmml">2</mn></msup></mrow></mrow><mo lspace="0em" id="S2.E1.m1.4.4.1.2" xref="S2.E1.m1.4.4.1.1.cmml">.</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E1.m1.4b"><apply id="S2.E1.m1.4.4.1.1.cmml" xref="S2.E1.m1.4.4.1"><eq id="S2.E1.m1.4.4.1.1.2.cmml" xref="S2.E1.m1.4.4.1.1.2"></eq><apply id="S2.E1.m1.4.4.1.1.3.cmml" xref="S2.E1.m1.4.4.1.1.3"><times id="S2.E1.m1.4.4.1.1.3.1.cmml" xref="S2.E1.m1.4.4.1.1.3.1"></times><ci id="S2.E1.m1.4.4.1.1.3.2.cmml" xref="S2.E1.m1.4.4.1.1.3.2">𝑟</ci><ci id="S2.E1.m1.3.3.cmml" xref="S2.E1.m1.3.3">𝑧</ci></apply><apply id="S2.E1.m1.4.4.1.1.1.cmml" xref="S2.E1.m1.4.4.1.1.1"><times id="S2.E1.m1.4.4.1.1.1.2.cmml" xref="S2.E1.m1.4.4.1.1.1.2"></times><apply id="S2.E1.m1.4.4.1.1.1.3.cmml" xref="S2.E1.m1.4.4.1.1.1.3"><arg id="S2.E1.m1.4.4.1.1.1.3.1.cmml" xref="S2.E1.m1.4.4.1.1.1.3.1"></arg><apply id="S2.E1.m1.4.4.1.1.1.3.2.cmml" xref="S2.E1.m1.4.4.1.1.1.3.2"><csymbol cd="ambiguous" id="S2.E1.m1.4.4.1.1.1.3.2.1.cmml" xref="S2.E1.m1.4.4.1.1.1.3.2">subscript</csymbol><min id="S2.E1.m1.4.4.1.1.1.3.2.2.cmml" xref="S2.E1.m1.4.4.1.1.1.3.2.2"></min><apply id="S2.E1.m1.2.2.2.cmml" xref="S2.E1.m1.2.2.2"><in id="S2.E1.m1.2.2.2.3.cmml" xref="S2.E1.m1.2.2.2.3"></in><ci id="S2.E1.m1.2.2.2.4.cmml" xref="S2.E1.m1.2.2.2.4">𝑖</ci><interval closure="closed" id="S2.E1.m1.2.2.2.5.1.cmml" xref="S2.E1.m1.2.2.2.5.2"><cn type="integer" id="S2.E1.m1.1.1.1.1.cmml" xref="S2.E1.m1.1.1.1.1">1</cn><ci id="S2.E1.m1.2.2.2.2.cmml" xref="S2.E1.m1.2.2.2.2">𝑘</ci></interval></apply></apply></apply><apply id="S2.E1.m1.4.4.1.1.1.1.cmml" xref="S2.E1.m1.4.4.1.1.1.1"><csymbol cd="ambiguous" id="S2.E1.m1.4.4.1.1.1.1.2.cmml" xref="S2.E1.m1.4.4.1.1.1.1">superscript</csymbol><apply id="S2.E1.m1.4.4.1.1.1.1.1.2.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1"><csymbol cd="latexml" id="S2.E1.m1.4.4.1.1.1.1.1.2.1.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.2">norm</csymbol><apply id="S2.E1.m1.4.4.1.1.1.1.1.1.1.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1"><minus id="S2.E1.m1.4.4.1.1.1.1.1.1.1.1.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.1"></minus><ci id="S2.E1.m1.4.4.1.1.1.1.1.1.1.2.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.2">𝑧</ci><apply id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.1.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3">subscript</csymbol><ci id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.2.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.2">𝑐</ci><ci id="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.3.cmml" xref="S2.E1.m1.4.4.1.1.1.1.1.1.1.3.3">𝑖</ci></apply></apply></apply><cn type="integer" id="S2.E1.m1.4.4.1.1.1.1.3.cmml" xref="S2.E1.m1.4.4.1.1.1.1.3">2</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E1.m1.4c">r(z)=\arg\min_{i\in[1,k]}||z-c_{i}||^{2}.</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
<td rowspan="1" class="ltx_eqn_cell ltx_eqn_eqno ltx_align_middle ltx_align_right"><span class="ltx_tag ltx_tag_equation ltx_align_right">(1)</span></td>
</tr></tbody>
</table>
</div>
</section>
<section id="S2.SS4.SSS2" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.4.2 </span>Discriminative routing</h4>

<div id="S2.SS4.SSS2.p1" class="ltx_para">
<p id="S2.SS4.SSS2.p1.8" class="ltx_p">Generative routing is agnostic of the language modeling task. In discriminative routing the sharding takes into account how well experts perform on each sequence. The training set is split into two parts. The first and largest part is used to train the paths on data sharded using a generative sharding appproach. The second and much smaller part is used to evaluate paths on each sequence. Let <math id="S2.SS4.SSS2.p1.1.m1.1" class="ltx_Math" alttext="s_{i}" display="inline"><semantics id="S2.SS4.SSS2.p1.1.m1.1a"><msub id="S2.SS4.SSS2.p1.1.m1.1.1" xref="S2.SS4.SSS2.p1.1.m1.1.1.cmml"><mi id="S2.SS4.SSS2.p1.1.m1.1.1.2" xref="S2.SS4.SSS2.p1.1.m1.1.1.2.cmml">s</mi><mi id="S2.SS4.SSS2.p1.1.m1.1.1.3" xref="S2.SS4.SSS2.p1.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.1.m1.1b"><apply id="S2.SS4.SSS2.p1.1.m1.1.1.cmml" xref="S2.SS4.SSS2.p1.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.1.m1.1.1.1.cmml" xref="S2.SS4.SSS2.p1.1.m1.1.1">subscript</csymbol><ci id="S2.SS4.SSS2.p1.1.m1.1.1.2.cmml" xref="S2.SS4.SSS2.p1.1.m1.1.1.2">𝑠</ci><ci id="S2.SS4.SSS2.p1.1.m1.1.1.3.cmml" xref="S2.SS4.SSS2.p1.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.1.m1.1c">s_{i}</annotation></semantics></math> be the log-likelihood score of the path <math id="S2.SS4.SSS2.p1.2.m2.1" class="ltx_Math" alttext="\pi_{i}" display="inline"><semantics id="S2.SS4.SSS2.p1.2.m2.1a"><msub id="S2.SS4.SSS2.p1.2.m2.1.1" xref="S2.SS4.SSS2.p1.2.m2.1.1.cmml"><mi id="S2.SS4.SSS2.p1.2.m2.1.1.2" xref="S2.SS4.SSS2.p1.2.m2.1.1.2.cmml">π</mi><mi id="S2.SS4.SSS2.p1.2.m2.1.1.3" xref="S2.SS4.SSS2.p1.2.m2.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.2.m2.1b"><apply id="S2.SS4.SSS2.p1.2.m2.1.1.cmml" xref="S2.SS4.SSS2.p1.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.2.m2.1.1.1.cmml" xref="S2.SS4.SSS2.p1.2.m2.1.1">subscript</csymbol><ci id="S2.SS4.SSS2.p1.2.m2.1.1.2.cmml" xref="S2.SS4.SSS2.p1.2.m2.1.1.2">𝜋</ci><ci id="S2.SS4.SSS2.p1.2.m2.1.1.3.cmml" xref="S2.SS4.SSS2.p1.2.m2.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.2.m2.1c">\pi_{i}</annotation></semantics></math> on the sequence associated with prefix <math id="S2.SS4.SSS2.p1.3.m3.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS2.p1.3.m3.1a"><mi id="S2.SS4.SSS2.p1.3.m3.1.1" xref="S2.SS4.SSS2.p1.3.m3.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.3.m3.1b"><ci id="S2.SS4.SSS2.p1.3.m3.1.1.cmml" xref="S2.SS4.SSS2.p1.3.m3.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.3.m3.1c">z</annotation></semantics></math>. The index of the path attaining the largest likelihood score, <math id="S2.SS4.SSS2.p1.4.m4.1" class="ltx_Math" alttext="i^{*}=\arg\max_{j}s_{j}" display="inline"><semantics id="S2.SS4.SSS2.p1.4.m4.1a"><mrow id="S2.SS4.SSS2.p1.4.m4.1.1" xref="S2.SS4.SSS2.p1.4.m4.1.1.cmml"><msup id="S2.SS4.SSS2.p1.4.m4.1.1.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.2.cmml"><mi id="S2.SS4.SSS2.p1.4.m4.1.1.2.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.2.2.cmml">i</mi><mo id="S2.SS4.SSS2.p1.4.m4.1.1.2.3" xref="S2.SS4.SSS2.p1.4.m4.1.1.2.3.cmml">∗</mo></msup><mo id="S2.SS4.SSS2.p1.4.m4.1.1.1" xref="S2.SS4.SSS2.p1.4.m4.1.1.1.cmml">=</mo><mrow id="S2.SS4.SSS2.p1.4.m4.1.1.3" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.cmml"><mi id="S2.SS4.SSS2.p1.4.m4.1.1.3.1" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.1.cmml">arg</mi><mo lspace="0.167em" id="S2.SS4.SSS2.p1.4.m4.1.1.3a" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.cmml">⁡</mo><mrow id="S2.SS4.SSS2.p1.4.m4.1.1.3.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.cmml"><msub id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.cmml"><mi id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.2.cmml">max</mi><mi id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.3" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.3.cmml">j</mi></msub><mo lspace="0.167em" id="S2.SS4.SSS2.p1.4.m4.1.1.3.2a" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.cmml">⁡</mo><msub id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.cmml"><mi id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.2" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.2.cmml">s</mi><mi id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.3" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.3.cmml">j</mi></msub></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.4.m4.1b"><apply id="S2.SS4.SSS2.p1.4.m4.1.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1"><eq id="S2.SS4.SSS2.p1.4.m4.1.1.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.1"></eq><apply id="S2.SS4.SSS2.p1.4.m4.1.1.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.2"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.4.m4.1.1.2.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.2">superscript</csymbol><ci id="S2.SS4.SSS2.p1.4.m4.1.1.2.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.2.2">𝑖</ci><times id="S2.SS4.SSS2.p1.4.m4.1.1.2.3.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.2.3"></times></apply><apply id="S2.SS4.SSS2.p1.4.m4.1.1.3.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3"><arg id="S2.SS4.SSS2.p1.4.m4.1.1.3.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.1"></arg><apply id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2"><apply id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1">subscript</csymbol><max id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.2"></max><ci id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.3.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.1.3">𝑗</ci></apply><apply id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.1.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2">subscript</csymbol><ci id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.2.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.2">𝑠</ci><ci id="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.3.cmml" xref="S2.SS4.SSS2.p1.4.m4.1.1.3.2.2.3">𝑗</ci></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.4.m4.1c">i^{*}=\arg\max_{j}s_{j}</annotation></semantics></math>, is used as label for <math id="S2.SS4.SSS2.p1.5.m5.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS2.p1.5.m5.1a"><mi id="S2.SS4.SSS2.p1.5.m5.1.1" xref="S2.SS4.SSS2.p1.5.m5.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.5.m5.1b"><ci id="S2.SS4.SSS2.p1.5.m5.1.1.cmml" xref="S2.SS4.SSS2.p1.5.m5.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.5.m5.1c">z</annotation></semantics></math>. We then train a logistic regression classifier to directly predict <math id="S2.SS4.SSS2.p1.6.m6.1" class="ltx_Math" alttext="i^{*}" display="inline"><semantics id="S2.SS4.SSS2.p1.6.m6.1a"><msup id="S2.SS4.SSS2.p1.6.m6.1.1" xref="S2.SS4.SSS2.p1.6.m6.1.1.cmml"><mi id="S2.SS4.SSS2.p1.6.m6.1.1.2" xref="S2.SS4.SSS2.p1.6.m6.1.1.2.cmml">i</mi><mo id="S2.SS4.SSS2.p1.6.m6.1.1.3" xref="S2.SS4.SSS2.p1.6.m6.1.1.3.cmml">∗</mo></msup><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.6.m6.1b"><apply id="S2.SS4.SSS2.p1.6.m6.1.1.cmml" xref="S2.SS4.SSS2.p1.6.m6.1.1"><csymbol cd="ambiguous" id="S2.SS4.SSS2.p1.6.m6.1.1.1.cmml" xref="S2.SS4.SSS2.p1.6.m6.1.1">superscript</csymbol><ci id="S2.SS4.SSS2.p1.6.m6.1.1.2.cmml" xref="S2.SS4.SSS2.p1.6.m6.1.1.2">𝑖</ci><times id="S2.SS4.SSS2.p1.6.m6.1.1.3.cmml" xref="S2.SS4.SSS2.p1.6.m6.1.1.3"></times></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.6.m6.1c">i^{*}</annotation></semantics></math> from the features <math id="S2.SS4.SSS2.p1.7.m7.1" class="ltx_Math" alttext="z" display="inline"><semantics id="S2.SS4.SSS2.p1.7.m7.1a"><mi id="S2.SS4.SSS2.p1.7.m7.1.1" xref="S2.SS4.SSS2.p1.7.m7.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.7.m7.1b"><ci id="S2.SS4.SSS2.p1.7.m7.1.1.cmml" xref="S2.SS4.SSS2.p1.7.m7.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.7.m7.1c">z</annotation></semantics></math>; this trained classifier is used as <math id="S2.SS4.SSS2.p1.8.m8.1" class="ltx_Math" alttext="r" display="inline"><semantics id="S2.SS4.SSS2.p1.8.m8.1a"><mi id="S2.SS4.SSS2.p1.8.m8.1.1" xref="S2.SS4.SSS2.p1.8.m8.1.1.cmml">r</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS2.p1.8.m8.1b"><ci id="S2.SS4.SSS2.p1.8.m8.1.1.cmml" xref="S2.SS4.SSS2.p1.8.m8.1.1">𝑟</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS2.p1.8.m8.1c">r</annotation></semantics></math>.
Finally, we use such classifier to re-shard the entire training and validation datasets.</p>
</div>
<div id="S2.SS4.SSS2.p2" class="ltx_para">
<p id="S2.SS4.SSS2.p2.1" class="ltx_p">This process is an approximation to Expectation Maximization <cite class="ltx_cite ltx_citemacro_citep">(Dempster et al., <a href="#bib.bib12" title="" class="ltx_ref">1977</a>)</cite>, a coordinate descent method alternating between updating the parameters of the model (paths), and latent variables (path assignments). It is an approximation because, for simplicity, we use as label the top scoring path as opposed to the full posterior distribution over path scores. Of course, this re-assignment process can be repeated for as many times as desired.</p>
</div>
<div id="S2.SS4.SSS2.p3" class="ltx_para">
<p id="S2.SS4.SSS2.p3.1" class="ltx_p">We invite those readers interested in diving deeper in the routing mechanism to read <a href="#Sx2.SS2" title="7.2 Routing Details ‣ Supplementary Materials ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 7.2</span></a> of the Appendix for further details.</p>
</div>
</section>
<section id="S2.SS4.SSS3" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.4.3 </span>Routing More Frequently at Test Time</h4>

<div id="S2.SS4.SSS3.p1" class="ltx_para">
<p id="S2.SS4.SSS3.p1.4" class="ltx_p">Routing only once per sequence is critical to distribute training across paths. However, at test time we can afford routing more frequently. Like <a href="#S2.F3" title="Figure 3 ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 3</span></a> shows, we can divide a sequence at test time into chunks of <math id="S2.SS4.SSS3.p1.1.m1.1" class="ltx_Math" alttext="W&gt;1" display="inline"><semantics id="S2.SS4.SSS3.p1.1.m1.1a"><mrow id="S2.SS4.SSS3.p1.1.m1.1.1" xref="S2.SS4.SSS3.p1.1.m1.1.1.cmml"><mi id="S2.SS4.SSS3.p1.1.m1.1.1.2" xref="S2.SS4.SSS3.p1.1.m1.1.1.2.cmml">W</mi><mo id="S2.SS4.SSS3.p1.1.m1.1.1.1" xref="S2.SS4.SSS3.p1.1.m1.1.1.1.cmml">&gt;</mo><mn id="S2.SS4.SSS3.p1.1.m1.1.1.3" xref="S2.SS4.SSS3.p1.1.m1.1.1.3.cmml">1</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS3.p1.1.m1.1b"><apply id="S2.SS4.SSS3.p1.1.m1.1.1.cmml" xref="S2.SS4.SSS3.p1.1.m1.1.1"><gt id="S2.SS4.SSS3.p1.1.m1.1.1.1.cmml" xref="S2.SS4.SSS3.p1.1.m1.1.1.1"></gt><ci id="S2.SS4.SSS3.p1.1.m1.1.1.2.cmml" xref="S2.SS4.SSS3.p1.1.m1.1.1.2">𝑊</ci><cn type="integer" id="S2.SS4.SSS3.p1.1.m1.1.1.3.cmml" xref="S2.SS4.SSS3.p1.1.m1.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS3.p1.1.m1.1c">W&gt;1</annotation></semantics></math> consecutive tokens, e.g., <math id="S2.SS4.SSS3.p1.2.m2.1" class="ltx_Math" alttext="W=128" display="inline"><semantics id="S2.SS4.SSS3.p1.2.m2.1a"><mrow id="S2.SS4.SSS3.p1.2.m2.1.1" xref="S2.SS4.SSS3.p1.2.m2.1.1.cmml"><mi id="S2.SS4.SSS3.p1.2.m2.1.1.2" xref="S2.SS4.SSS3.p1.2.m2.1.1.2.cmml">W</mi><mo id="S2.SS4.SSS3.p1.2.m2.1.1.1" xref="S2.SS4.SSS3.p1.2.m2.1.1.1.cmml">=</mo><mn id="S2.SS4.SSS3.p1.2.m2.1.1.3" xref="S2.SS4.SSS3.p1.2.m2.1.1.3.cmml">128</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS3.p1.2.m2.1b"><apply id="S2.SS4.SSS3.p1.2.m2.1.1.cmml" xref="S2.SS4.SSS3.p1.2.m2.1.1"><eq id="S2.SS4.SSS3.p1.2.m2.1.1.1.cmml" xref="S2.SS4.SSS3.p1.2.m2.1.1.1"></eq><ci id="S2.SS4.SSS3.p1.2.m2.1.1.2.cmml" xref="S2.SS4.SSS3.p1.2.m2.1.1.2">𝑊</ci><cn type="integer" id="S2.SS4.SSS3.p1.2.m2.1.1.3.cmml" xref="S2.SS4.SSS3.p1.2.m2.1.1.3">128</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS3.p1.2.m2.1c">W=128</annotation></semantics></math>. We can then feed a router with the <math id="S2.SS4.SSS3.p1.3.m3.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S2.SS4.SSS3.p1.3.m3.1a"><mi id="S2.SS4.SSS3.p1.3.m3.1.1" xref="S2.SS4.SSS3.p1.3.m3.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS3.p1.3.m3.1b"><ci id="S2.SS4.SSS3.p1.3.m3.1.1.cmml" xref="S2.SS4.SSS3.p1.3.m3.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS3.p1.3.m3.1c">i</annotation></semantics></math>-th chunk to predict the best path to use at the <math id="S2.SS4.SSS3.p1.4.m4.1" class="ltx_Math" alttext="(i+1)" display="inline"><semantics id="S2.SS4.SSS3.p1.4.m4.1a"><mrow id="S2.SS4.SSS3.p1.4.m4.1.1.1" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.SSS3.p1.4.m4.1.1.1.2" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.cmml">(</mo><mrow id="S2.SS4.SSS3.p1.4.m4.1.1.1.1" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.cmml"><mi id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.2" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.2.cmml">i</mi><mo id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.1" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.1.cmml">+</mo><mn id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.3" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.3.cmml">1</mn></mrow><mo stretchy="false" id="S2.SS4.SSS3.p1.4.m4.1.1.1.3" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS3.p1.4.m4.1b"><apply id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.cmml" xref="S2.SS4.SSS3.p1.4.m4.1.1.1"><plus id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.1.cmml" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.1"></plus><ci id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.2.cmml" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.2">𝑖</ci><cn type="integer" id="S2.SS4.SSS3.p1.4.m4.1.1.1.1.3.cmml" xref="S2.SS4.SSS3.p1.4.m4.1.1.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS3.p1.4.m4.1c">(i+1)</annotation></semantics></math>-th chunk, using a process similar to the one described in <a href="#S2.SS4.SSS2" title="2.4.2 Discriminative routing ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.4.2</span></a>. The cost of switching paths is small, since the router is run infrequently and in scoring (as opposed to auto-regressive generation) mode. Moreover, only text needs to be communicated to the router and next selected path.</p>
</div>
</section>
<section id="S2.SS4.SSS4" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.4.4 </span>Overlapping Shards</h4>

<div id="S2.SS4.SSS4.p1" class="ltx_para">
<p id="S2.SS4.SSS4.p1.1" class="ltx_p">Above, each sequence was associated to one and only one shard; in other words, we pre-sharded the dataset into a set of <math id="S2.SS4.SSS4.p1.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS4.p1.1.m1.1a"><mi id="S2.SS4.SSS4.p1.1.m1.1.1" xref="S2.SS4.SSS4.p1.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS4.p1.1.m1.1b"><ci id="S2.SS4.SSS4.p1.1.m1.1.1.cmml" xref="S2.SS4.SSS4.p1.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS4.p1.1.m1.1c">k</annotation></semantics></math> disjoint partitions.</p>
</div>
<div id="S2.SS4.SSS4.p2" class="ltx_para">
<p id="S2.SS4.SSS4.p2.2" class="ltx_p">When the number of shards <math id="S2.SS4.SSS4.p2.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.SSS4.p2.1.m1.1a"><mi id="S2.SS4.SSS4.p2.1.m1.1.1" xref="S2.SS4.SSS4.p2.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS4.p2.1.m1.1b"><ci id="S2.SS4.SSS4.p2.1.m1.1.1.cmml" xref="S2.SS4.SSS4.p2.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS4.p2.1.m1.1c">k</annotation></semantics></math> is large, the mixture might overfit, particularly when paths share few parameters.
To combat this issue we might assign each sequence not just to one shard but to its <math id="S2.SS4.SSS4.p2.2.m2.1" class="ltx_Math" alttext="n" display="inline"><semantics id="S2.SS4.SSS4.p2.2.m2.1a"><mi id="S2.SS4.SSS4.p2.2.m2.1.1" xref="S2.SS4.SSS4.p2.2.m2.1.1.cmml">n</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.SSS4.p2.2.m2.1b"><ci id="S2.SS4.SSS4.p2.2.m2.1.1.cmml" xref="S2.SS4.SSS4.p2.2.m2.1.1">𝑛</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.SSS4.p2.2.m2.1c">n</annotation></semantics></math> “closest” shards.
This makes shards overlap with each other, smoothing out the boundaries among themselves, which may improve generalization when the number of shards is relatively large. This can be done <span id="S2.SS4.SSS4.p2.2.1" class="ltx_text ltx_font_italic">independently</span> at train time and at evaluation time; that is, one might overlap shards only during train, at eval, or both, or neither. At train time, overlapping the shards limits the specialization of each path, and increases the storage necessary for each shard, but otherwise does not have a computational cost, as each path is still trained independently on its shard. However, overlapping the shards at evaluation means that the model has to be forwarded through multiple paths, which increases the evaluation computational cost.</p>
</div>
<div id="S2.SS4.SSS4.p3" class="ltx_para">
<p id="S2.SS4.SSS4.p3.1" class="ltx_p">In this work, when we overlap the shards in training, we use the top-2 choices, while we never overlap the shards at evaluation. We leave to future work other approaches to overlap shards.</p>
</div>
</section>
</section>
<section id="S2.SS5" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.5 </span>DiLoCo: Review</h3>

<div id="S2.SS5.p1" class="ltx_para">
<p id="S2.SS5.p1.1" class="ltx_p">In this section, we summarize DiLoCo <cite class="ltx_cite ltx_citemacro_citep">(Douillard et al., <a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite> as it lays the foundation of our distributed optimization algorithm, which will be described in the next section. Together with coarse routing of <a href="#S2.SS4" title="2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.4</span></a>, DiLoCo is a critical component of DiPaCo.</p>
</div>
<div id="S2.SS5.p2" class="ltx_para">
<p id="S2.SS5.p2.3" class="ltx_p">DiLoCo optimizes a dense model across <math id="S2.SS5.p2.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS5.p2.1.m1.1a"><mi id="S2.SS5.p2.1.m1.1.1" xref="S2.SS5.p2.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS5.p2.1.m1.1b"><ci id="S2.SS5.p2.1.m1.1.1.cmml" xref="S2.SS5.p2.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p2.1.m1.1c">k</annotation></semantics></math> workers. First, the original dataset is sharded into <math id="S2.SS5.p2.2.m2.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS5.p2.2.m2.1a"><mi id="S2.SS5.p2.2.m2.1.1" xref="S2.SS5.p2.2.m2.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS5.p2.2.m2.1b"><ci id="S2.SS5.p2.2.m2.1.1.cmml" xref="S2.SS5.p2.2.m2.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p2.2.m2.1c">k</annotation></semantics></math> shards, and each shard is associated to a worker. Second, all workers start from the same copy of the model parameters (typically a pretrained model), and perform <math id="S2.SS5.p2.3.m3.1" class="ltx_Math" alttext="H" display="inline"><semantics id="S2.SS5.p2.3.m3.1a"><mi id="S2.SS5.p2.3.m3.1.1" xref="S2.SS5.p2.3.m3.1.1.cmml">H</mi><annotation-xml encoding="MathML-Content" id="S2.SS5.p2.3.m3.1b"><ci id="S2.SS5.p2.3.m3.1.1.cmml" xref="S2.SS5.p2.3.m3.1.1">𝐻</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p2.3.m3.1c">H</annotation></semantics></math> (inner) optimization steps over their own shard independently. Third, each worker sends to a central CPU server the difference in parameter space between their new parameters and their initial ones. We refer to these differences as <span id="S2.SS5.p2.3.1" class="ltx_text ltx_font_italic">outer gradients</span>. The central server averages these outer gradients into a single update vector. Finally, the global parameter vector is updated using an (outer) optimizer, and the new global parameter vector is re-dispatched to all workers for another phase of local training. The process repeats for as many rounds as desired.</p>
</div>
<div id="S2.SS5.p3" class="ltx_para">
<p id="S2.SS5.p3.2" class="ltx_p">In language modeling applications using transformers, the inner and outer optimizers that have been shown to be most effective are respectively AdamW <cite class="ltx_cite ltx_citemacro_citep">(Kingma and Ba, <a href="#bib.bib23" title="" class="ltx_ref">2014</a>)</cite> and Nesterov momentum <cite class="ltx_cite ltx_citemacro_citep">(Sutskever et al., <a href="#bib.bib40" title="" class="ltx_ref">2013</a>)</cite>. Empirically in prior work, <math id="S2.SS5.p3.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS5.p3.1.m1.1a"><mi id="S2.SS5.p3.1.m1.1.1" xref="S2.SS5.p3.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS5.p3.1.m1.1b"><ci id="S2.SS5.p3.1.m1.1.1.cmml" xref="S2.SS5.p3.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p3.1.m1.1c">k</annotation></semantics></math> has been set up to 64 and <math id="S2.SS5.p3.2.m2.1" class="ltx_Math" alttext="H" display="inline"><semantics id="S2.SS5.p3.2.m2.1a"><mi id="S2.SS5.p3.2.m2.1.1" xref="S2.SS5.p3.2.m2.1.1.cmml">H</mi><annotation-xml encoding="MathML-Content" id="S2.SS5.p3.2.m2.1b"><ci id="S2.SS5.p3.2.m2.1.1.cmml" xref="S2.SS5.p3.2.m2.1.1">𝐻</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p3.2.m2.1c">H</annotation></semantics></math> up to thousand. Note that other alternatives, such as FedOpt <cite class="ltx_cite ltx_citemacro_citep">(Reddi et al., <a href="#bib.bib36" title="" class="ltx_ref">2021</a>)</cite>, are compatible with this framework.</p>
</div>
</section>
<section id="S2.SS6" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.6 </span>DiPaCo</h3>

<figure id="S2.F4" class="ltx_figure"><img src="/html/2403.10616/assets/x5.png" id="S2.F4.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="298" height="106" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S2.F4.21.9.1" class="ltx_text" style="font-size:90%;">Figure 4</span>: </span><span id="S2.F4.16.8" class="ltx_text ltx_font_bold" style="font-size:90%;">DiPaCo<span id="S2.F4.16.8.8" class="ltx_text ltx_font_medium">: (left) The dataset is pre-sharded into <math id="S2.F4.9.1.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.F4.9.1.1.m1.1b"><mi id="S2.F4.9.1.1.m1.1.1" xref="S2.F4.9.1.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.F4.9.1.1.m1.1c"><ci id="S2.F4.9.1.1.m1.1.1.cmml" xref="S2.F4.9.1.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.9.1.1.m1.1d">k</annotation></semantics></math> shards, <math id="S2.F4.10.2.2.m2.1" class="ltx_Math" alttext="\mathcal{D}_{i}" display="inline"><semantics id="S2.F4.10.2.2.m2.1b"><msub id="S2.F4.10.2.2.m2.1.1" xref="S2.F4.10.2.2.m2.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.F4.10.2.2.m2.1.1.2" xref="S2.F4.10.2.2.m2.1.1.2.cmml">𝒟</mi><mi id="S2.F4.10.2.2.m2.1.1.3" xref="S2.F4.10.2.2.m2.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.F4.10.2.2.m2.1c"><apply id="S2.F4.10.2.2.m2.1.1.cmml" xref="S2.F4.10.2.2.m2.1.1"><csymbol cd="ambiguous" id="S2.F4.10.2.2.m2.1.1.1.cmml" xref="S2.F4.10.2.2.m2.1.1">subscript</csymbol><ci id="S2.F4.10.2.2.m2.1.1.2.cmml" xref="S2.F4.10.2.2.m2.1.1.2">𝒟</ci><ci id="S2.F4.10.2.2.m2.1.1.3.cmml" xref="S2.F4.10.2.2.m2.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.10.2.2.m2.1d">\mathcal{D}_{i}</annotation></semantics></math> (here <math id="S2.F4.11.3.3.m3.1" class="ltx_Math" alttext="k=4" display="inline"><semantics id="S2.F4.11.3.3.m3.1b"><mrow id="S2.F4.11.3.3.m3.1.1" xref="S2.F4.11.3.3.m3.1.1.cmml"><mi id="S2.F4.11.3.3.m3.1.1.2" xref="S2.F4.11.3.3.m3.1.1.2.cmml">k</mi><mo id="S2.F4.11.3.3.m3.1.1.1" xref="S2.F4.11.3.3.m3.1.1.1.cmml">=</mo><mn id="S2.F4.11.3.3.m3.1.1.3" xref="S2.F4.11.3.3.m3.1.1.3.cmml">4</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.F4.11.3.3.m3.1c"><apply id="S2.F4.11.3.3.m3.1.1.cmml" xref="S2.F4.11.3.3.m3.1.1"><eq id="S2.F4.11.3.3.m3.1.1.1.cmml" xref="S2.F4.11.3.3.m3.1.1.1"></eq><ci id="S2.F4.11.3.3.m3.1.1.2.cmml" xref="S2.F4.11.3.3.m3.1.1.2">𝑘</ci><cn type="integer" id="S2.F4.11.3.3.m3.1.1.3.cmml" xref="S2.F4.11.3.3.m3.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.11.3.3.m3.1d">k=4</annotation></semantics></math>). (middle) Compact view of a <math id="S2.F4.12.4.4.m4.1" class="ltx_Math" alttext="2\times 2" display="inline"><semantics id="S2.F4.12.4.4.m4.1b"><mrow id="S2.F4.12.4.4.m4.1.1" xref="S2.F4.12.4.4.m4.1.1.cmml"><mn id="S2.F4.12.4.4.m4.1.1.2" xref="S2.F4.12.4.4.m4.1.1.2.cmml">2</mn><mo lspace="0.222em" rspace="0.222em" id="S2.F4.12.4.4.m4.1.1.1" xref="S2.F4.12.4.4.m4.1.1.1.cmml">×</mo><mn id="S2.F4.12.4.4.m4.1.1.3" xref="S2.F4.12.4.4.m4.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.F4.12.4.4.m4.1c"><apply id="S2.F4.12.4.4.m4.1.1.cmml" xref="S2.F4.12.4.4.m4.1.1"><times id="S2.F4.12.4.4.m4.1.1.1.cmml" xref="S2.F4.12.4.4.m4.1.1.1"></times><cn type="integer" id="S2.F4.12.4.4.m4.1.1.2.cmml" xref="S2.F4.12.4.4.m4.1.1.2">2</cn><cn type="integer" id="S2.F4.12.4.4.m4.1.1.3.cmml" xref="S2.F4.12.4.4.m4.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.12.4.4.m4.1d">2\times 2</annotation></semantics></math> DiPaCo, which is never instantiated. In this toy illustration, there are three levels. Level 2 and 3 have a mixture with two modules each. Level 1 has a single module shared by all paths. (right)
We associate each shard <math id="S2.F4.13.5.5.m5.1" class="ltx_Math" alttext="\mathcal{D}_{i}" display="inline"><semantics id="S2.F4.13.5.5.m5.1b"><msub id="S2.F4.13.5.5.m5.1.1" xref="S2.F4.13.5.5.m5.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.F4.13.5.5.m5.1.1.2" xref="S2.F4.13.5.5.m5.1.1.2.cmml">𝒟</mi><mi id="S2.F4.13.5.5.m5.1.1.3" xref="S2.F4.13.5.5.m5.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S2.F4.13.5.5.m5.1c"><apply id="S2.F4.13.5.5.m5.1.1.cmml" xref="S2.F4.13.5.5.m5.1.1"><csymbol cd="ambiguous" id="S2.F4.13.5.5.m5.1.1.1.cmml" xref="S2.F4.13.5.5.m5.1.1">subscript</csymbol><ci id="S2.F4.13.5.5.m5.1.1.2.cmml" xref="S2.F4.13.5.5.m5.1.1.2">𝒟</ci><ci id="S2.F4.13.5.5.m5.1.1.3.cmml" xref="S2.F4.13.5.5.m5.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.13.5.5.m5.1d">\mathcal{D}_{i}</annotation></semantics></math> to a <span id="S2.F4.16.8.8.1" class="ltx_text ltx_font_italic">path</span> <math id="S2.F4.14.6.6.m6.4" class="ltx_Math" alttext="\pi_{i},\,\forall i\in[1,4]" display="inline"><semantics id="S2.F4.14.6.6.m6.4b"><mrow id="S2.F4.14.6.6.m6.4.4" xref="S2.F4.14.6.6.m6.4.4.cmml"><mrow id="S2.F4.14.6.6.m6.4.4.2.2" xref="S2.F4.14.6.6.m6.4.4.2.3.cmml"><msub id="S2.F4.14.6.6.m6.3.3.1.1.1" xref="S2.F4.14.6.6.m6.3.3.1.1.1.cmml"><mi id="S2.F4.14.6.6.m6.3.3.1.1.1.2" xref="S2.F4.14.6.6.m6.3.3.1.1.1.2.cmml">π</mi><mi id="S2.F4.14.6.6.m6.3.3.1.1.1.3" xref="S2.F4.14.6.6.m6.3.3.1.1.1.3.cmml">i</mi></msub><mo rspace="0.337em" id="S2.F4.14.6.6.m6.4.4.2.2.3" xref="S2.F4.14.6.6.m6.4.4.2.3.cmml">,</mo><mrow id="S2.F4.14.6.6.m6.4.4.2.2.2" xref="S2.F4.14.6.6.m6.4.4.2.2.2.cmml"><mo rspace="0.167em" id="S2.F4.14.6.6.m6.4.4.2.2.2.1" xref="S2.F4.14.6.6.m6.4.4.2.2.2.1.cmml">∀</mo><mi id="S2.F4.14.6.6.m6.4.4.2.2.2.2" xref="S2.F4.14.6.6.m6.4.4.2.2.2.2.cmml">i</mi></mrow></mrow><mo id="S2.F4.14.6.6.m6.4.4.3" xref="S2.F4.14.6.6.m6.4.4.3.cmml">∈</mo><mrow id="S2.F4.14.6.6.m6.4.4.4.2" xref="S2.F4.14.6.6.m6.4.4.4.1.cmml"><mo stretchy="false" id="S2.F4.14.6.6.m6.4.4.4.2.1" xref="S2.F4.14.6.6.m6.4.4.4.1.cmml">[</mo><mn id="S2.F4.14.6.6.m6.1.1" xref="S2.F4.14.6.6.m6.1.1.cmml">1</mn><mo id="S2.F4.14.6.6.m6.4.4.4.2.2" xref="S2.F4.14.6.6.m6.4.4.4.1.cmml">,</mo><mn id="S2.F4.14.6.6.m6.2.2" xref="S2.F4.14.6.6.m6.2.2.cmml">4</mn><mo stretchy="false" id="S2.F4.14.6.6.m6.4.4.4.2.3" xref="S2.F4.14.6.6.m6.4.4.4.1.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.F4.14.6.6.m6.4c"><apply id="S2.F4.14.6.6.m6.4.4.cmml" xref="S2.F4.14.6.6.m6.4.4"><in id="S2.F4.14.6.6.m6.4.4.3.cmml" xref="S2.F4.14.6.6.m6.4.4.3"></in><list id="S2.F4.14.6.6.m6.4.4.2.3.cmml" xref="S2.F4.14.6.6.m6.4.4.2.2"><apply id="S2.F4.14.6.6.m6.3.3.1.1.1.cmml" xref="S2.F4.14.6.6.m6.3.3.1.1.1"><csymbol cd="ambiguous" id="S2.F4.14.6.6.m6.3.3.1.1.1.1.cmml" xref="S2.F4.14.6.6.m6.3.3.1.1.1">subscript</csymbol><ci id="S2.F4.14.6.6.m6.3.3.1.1.1.2.cmml" xref="S2.F4.14.6.6.m6.3.3.1.1.1.2">𝜋</ci><ci id="S2.F4.14.6.6.m6.3.3.1.1.1.3.cmml" xref="S2.F4.14.6.6.m6.3.3.1.1.1.3">𝑖</ci></apply><apply id="S2.F4.14.6.6.m6.4.4.2.2.2.cmml" xref="S2.F4.14.6.6.m6.4.4.2.2.2"><csymbol cd="latexml" id="S2.F4.14.6.6.m6.4.4.2.2.2.1.cmml" xref="S2.F4.14.6.6.m6.4.4.2.2.2.1">for-all</csymbol><ci id="S2.F4.14.6.6.m6.4.4.2.2.2.2.cmml" xref="S2.F4.14.6.6.m6.4.4.2.2.2.2">𝑖</ci></apply></list><interval closure="closed" id="S2.F4.14.6.6.m6.4.4.4.1.cmml" xref="S2.F4.14.6.6.m6.4.4.4.2"><cn type="integer" id="S2.F4.14.6.6.m6.1.1.cmml" xref="S2.F4.14.6.6.m6.1.1">1</cn><cn type="integer" id="S2.F4.14.6.6.m6.2.2.cmml" xref="S2.F4.14.6.6.m6.2.2">4</cn></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.14.6.6.m6.4d">\pi_{i},\,\forall i\in[1,4]</annotation></semantics></math>. In this toy illustration, a path is the composition of three neural network blocks. The color refers to the id of a module. The figure shows the modular network unrolled across the four paths. These are trained by using DiLoCo which requires communicating only every few hundred steps. In this example, module 2a (in <span id="S2.F4.16.8.8.2" class="ltx_text" style="color:#E1144B;">red</span>) is shared by paths <math id="S2.F4.15.7.7.m7.1" class="ltx_Math" alttext="\pi_{1}" display="inline"><semantics id="S2.F4.15.7.7.m7.1b"><msub id="S2.F4.15.7.7.m7.1.1" xref="S2.F4.15.7.7.m7.1.1.cmml"><mi id="S2.F4.15.7.7.m7.1.1.2" xref="S2.F4.15.7.7.m7.1.1.2.cmml">π</mi><mn id="S2.F4.15.7.7.m7.1.1.3" xref="S2.F4.15.7.7.m7.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.F4.15.7.7.m7.1c"><apply id="S2.F4.15.7.7.m7.1.1.cmml" xref="S2.F4.15.7.7.m7.1.1"><csymbol cd="ambiguous" id="S2.F4.15.7.7.m7.1.1.1.cmml" xref="S2.F4.15.7.7.m7.1.1">subscript</csymbol><ci id="S2.F4.15.7.7.m7.1.1.2.cmml" xref="S2.F4.15.7.7.m7.1.1.2">𝜋</ci><cn type="integer" id="S2.F4.15.7.7.m7.1.1.3.cmml" xref="S2.F4.15.7.7.m7.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.15.7.7.m7.1d">\pi_{1}</annotation></semantics></math> and <math id="S2.F4.16.8.8.m8.1" class="ltx_Math" alttext="\pi_{2}" display="inline"><semantics id="S2.F4.16.8.8.m8.1b"><msub id="S2.F4.16.8.8.m8.1.1" xref="S2.F4.16.8.8.m8.1.1.cmml"><mi id="S2.F4.16.8.8.m8.1.1.2" xref="S2.F4.16.8.8.m8.1.1.2.cmml">π</mi><mn id="S2.F4.16.8.8.m8.1.1.3" xref="S2.F4.16.8.8.m8.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.F4.16.8.8.m8.1c"><apply id="S2.F4.16.8.8.m8.1.1.cmml" xref="S2.F4.16.8.8.m8.1.1"><csymbol cd="ambiguous" id="S2.F4.16.8.8.m8.1.1.1.cmml" xref="S2.F4.16.8.8.m8.1.1">subscript</csymbol><ci id="S2.F4.16.8.8.m8.1.1.2.cmml" xref="S2.F4.16.8.8.m8.1.1.2">𝜋</ci><cn type="integer" id="S2.F4.16.8.8.m8.1.1.3.cmml" xref="S2.F4.16.8.8.m8.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.F4.16.8.8.m8.1d">\pi_{2}</annotation></semantics></math>. Workers associated to paths might use different hardware types (different kind of GPUs or TPUs) and might be placed in far away geographic areas.</span></span></figcaption>
</figure>
<div id="S2.SS6.p1" class="ltx_para">
<p id="S2.SS6.p1.1" class="ltx_p">In <a href="#S2.SS4" title="2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.4</span></a>, we described methods to pre-shard data, and in <a href="#S2.SS5" title="2.5 DiLoCo: Review ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.5</span></a> we described how a model can be trained across several workers that communicated infrequently.
This section combines these two approaches towards the efficient training of a composable mixture, a network whose levels are replaced by a mixture of experts. We dub such architecture and training algorithm DiPaCo. In our experiments below, a level is composed of several consecutive transformer blocks.</p>
</div>
<div id="S2.SS6.p2" class="ltx_para">
<p id="S2.SS6.p2.13" class="ltx_p">In the toy illustration of <a href="#S2.F4" title="Figure 4 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 4</span></a> there are three levels, <math id="S2.SS6.p2.1.m1.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS6.p2.1.m1.1a"><msub id="S2.SS6.p2.1.m1.1.1" xref="S2.SS6.p2.1.m1.1.1.cmml"><mi id="S2.SS6.p2.1.m1.1.1.2" xref="S2.SS6.p2.1.m1.1.1.2.cmml">B</mi><mn id="S2.SS6.p2.1.m1.1.1.3" xref="S2.SS6.p2.1.m1.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.1.m1.1b"><apply id="S2.SS6.p2.1.m1.1.1.cmml" xref="S2.SS6.p2.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.1.m1.1.1.1.cmml" xref="S2.SS6.p2.1.m1.1.1">subscript</csymbol><ci id="S2.SS6.p2.1.m1.1.1.2.cmml" xref="S2.SS6.p2.1.m1.1.1.2">𝐵</ci><cn type="integer" id="S2.SS6.p2.1.m1.1.1.3.cmml" xref="S2.SS6.p2.1.m1.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.1.m1.1c">B_{1}</annotation></semantics></math>, <math id="S2.SS6.p2.2.m2.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS6.p2.2.m2.1a"><msub id="S2.SS6.p2.2.m2.1.1" xref="S2.SS6.p2.2.m2.1.1.cmml"><mi id="S2.SS6.p2.2.m2.1.1.2" xref="S2.SS6.p2.2.m2.1.1.2.cmml">B</mi><mn id="S2.SS6.p2.2.m2.1.1.3" xref="S2.SS6.p2.2.m2.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.2.m2.1b"><apply id="S2.SS6.p2.2.m2.1.1.cmml" xref="S2.SS6.p2.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.2.m2.1.1.1.cmml" xref="S2.SS6.p2.2.m2.1.1">subscript</csymbol><ci id="S2.SS6.p2.2.m2.1.1.2.cmml" xref="S2.SS6.p2.2.m2.1.1.2">𝐵</ci><cn type="integer" id="S2.SS6.p2.2.m2.1.1.3.cmml" xref="S2.SS6.p2.2.m2.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.2.m2.1c">B_{2}</annotation></semantics></math>, and <math id="S2.SS6.p2.3.m3.1" class="ltx_Math" alttext="B_{3}" display="inline"><semantics id="S2.SS6.p2.3.m3.1a"><msub id="S2.SS6.p2.3.m3.1.1" xref="S2.SS6.p2.3.m3.1.1.cmml"><mi id="S2.SS6.p2.3.m3.1.1.2" xref="S2.SS6.p2.3.m3.1.1.2.cmml">B</mi><mn id="S2.SS6.p2.3.m3.1.1.3" xref="S2.SS6.p2.3.m3.1.1.3.cmml">3</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.3.m3.1b"><apply id="S2.SS6.p2.3.m3.1.1.cmml" xref="S2.SS6.p2.3.m3.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.3.m3.1.1.1.cmml" xref="S2.SS6.p2.3.m3.1.1">subscript</csymbol><ci id="S2.SS6.p2.3.m3.1.1.2.cmml" xref="S2.SS6.p2.3.m3.1.1.2">𝐵</ci><cn type="integer" id="S2.SS6.p2.3.m3.1.1.3.cmml" xref="S2.SS6.p2.3.m3.1.1.3">3</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.3.m3.1c">B_{3}</annotation></semantics></math>. There is only one module (equivalently, one set of parameters) in <math id="S2.SS6.p2.4.m4.1" class="ltx_Math" alttext="B_{1}" display="inline"><semantics id="S2.SS6.p2.4.m4.1a"><msub id="S2.SS6.p2.4.m4.1.1" xref="S2.SS6.p2.4.m4.1.1.cmml"><mi id="S2.SS6.p2.4.m4.1.1.2" xref="S2.SS6.p2.4.m4.1.1.2.cmml">B</mi><mn id="S2.SS6.p2.4.m4.1.1.3" xref="S2.SS6.p2.4.m4.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.4.m4.1b"><apply id="S2.SS6.p2.4.m4.1.1.cmml" xref="S2.SS6.p2.4.m4.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.4.m4.1.1.1.cmml" xref="S2.SS6.p2.4.m4.1.1">subscript</csymbol><ci id="S2.SS6.p2.4.m4.1.1.2.cmml" xref="S2.SS6.p2.4.m4.1.1.2">𝐵</ci><cn type="integer" id="S2.SS6.p2.4.m4.1.1.3.cmml" xref="S2.SS6.p2.4.m4.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.4.m4.1c">B_{1}</annotation></semantics></math>, and it is shared across all paths. There are two modules in <math id="S2.SS6.p2.5.m5.1" class="ltx_Math" alttext="B_{2}" display="inline"><semantics id="S2.SS6.p2.5.m5.1a"><msub id="S2.SS6.p2.5.m5.1.1" xref="S2.SS6.p2.5.m5.1.1.cmml"><mi id="S2.SS6.p2.5.m5.1.1.2" xref="S2.SS6.p2.5.m5.1.1.2.cmml">B</mi><mn id="S2.SS6.p2.5.m5.1.1.3" xref="S2.SS6.p2.5.m5.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.5.m5.1b"><apply id="S2.SS6.p2.5.m5.1.1.cmml" xref="S2.SS6.p2.5.m5.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.5.m5.1.1.1.cmml" xref="S2.SS6.p2.5.m5.1.1">subscript</csymbol><ci id="S2.SS6.p2.5.m5.1.1.2.cmml" xref="S2.SS6.p2.5.m5.1.1.2">𝐵</ci><cn type="integer" id="S2.SS6.p2.5.m5.1.1.3.cmml" xref="S2.SS6.p2.5.m5.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.5.m5.1c">B_{2}</annotation></semantics></math> (meaning that there will be two distinct parameter vectors for the module at that level); the first is shared by paths <math id="S2.SS6.p2.6.m6.1" class="ltx_Math" alttext="\pi_{1}" display="inline"><semantics id="S2.SS6.p2.6.m6.1a"><msub id="S2.SS6.p2.6.m6.1.1" xref="S2.SS6.p2.6.m6.1.1.cmml"><mi id="S2.SS6.p2.6.m6.1.1.2" xref="S2.SS6.p2.6.m6.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.6.m6.1.1.3" xref="S2.SS6.p2.6.m6.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.6.m6.1b"><apply id="S2.SS6.p2.6.m6.1.1.cmml" xref="S2.SS6.p2.6.m6.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.6.m6.1.1.1.cmml" xref="S2.SS6.p2.6.m6.1.1">subscript</csymbol><ci id="S2.SS6.p2.6.m6.1.1.2.cmml" xref="S2.SS6.p2.6.m6.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.6.m6.1.1.3.cmml" xref="S2.SS6.p2.6.m6.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.6.m6.1c">\pi_{1}</annotation></semantics></math> and <math id="S2.SS6.p2.7.m7.1" class="ltx_Math" alttext="\pi_{2}" display="inline"><semantics id="S2.SS6.p2.7.m7.1a"><msub id="S2.SS6.p2.7.m7.1.1" xref="S2.SS6.p2.7.m7.1.1.cmml"><mi id="S2.SS6.p2.7.m7.1.1.2" xref="S2.SS6.p2.7.m7.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.7.m7.1.1.3" xref="S2.SS6.p2.7.m7.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.7.m7.1b"><apply id="S2.SS6.p2.7.m7.1.1.cmml" xref="S2.SS6.p2.7.m7.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.7.m7.1.1.1.cmml" xref="S2.SS6.p2.7.m7.1.1">subscript</csymbol><ci id="S2.SS6.p2.7.m7.1.1.2.cmml" xref="S2.SS6.p2.7.m7.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.7.m7.1.1.3.cmml" xref="S2.SS6.p2.7.m7.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.7.m7.1c">\pi_{2}</annotation></semantics></math>, and the second by paths <math id="S2.SS6.p2.8.m8.1" class="ltx_Math" alttext="\pi_{3}" display="inline"><semantics id="S2.SS6.p2.8.m8.1a"><msub id="S2.SS6.p2.8.m8.1.1" xref="S2.SS6.p2.8.m8.1.1.cmml"><mi id="S2.SS6.p2.8.m8.1.1.2" xref="S2.SS6.p2.8.m8.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.8.m8.1.1.3" xref="S2.SS6.p2.8.m8.1.1.3.cmml">3</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.8.m8.1b"><apply id="S2.SS6.p2.8.m8.1.1.cmml" xref="S2.SS6.p2.8.m8.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.8.m8.1.1.1.cmml" xref="S2.SS6.p2.8.m8.1.1">subscript</csymbol><ci id="S2.SS6.p2.8.m8.1.1.2.cmml" xref="S2.SS6.p2.8.m8.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.8.m8.1.1.3.cmml" xref="S2.SS6.p2.8.m8.1.1.3">3</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.8.m8.1c">\pi_{3}</annotation></semantics></math> and <math id="S2.SS6.p2.9.m9.1" class="ltx_Math" alttext="\pi_{4}" display="inline"><semantics id="S2.SS6.p2.9.m9.1a"><msub id="S2.SS6.p2.9.m9.1.1" xref="S2.SS6.p2.9.m9.1.1.cmml"><mi id="S2.SS6.p2.9.m9.1.1.2" xref="S2.SS6.p2.9.m9.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.9.m9.1.1.3" xref="S2.SS6.p2.9.m9.1.1.3.cmml">4</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.9.m9.1b"><apply id="S2.SS6.p2.9.m9.1.1.cmml" xref="S2.SS6.p2.9.m9.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.9.m9.1.1.1.cmml" xref="S2.SS6.p2.9.m9.1.1">subscript</csymbol><ci id="S2.SS6.p2.9.m9.1.1.2.cmml" xref="S2.SS6.p2.9.m9.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.9.m9.1.1.3.cmml" xref="S2.SS6.p2.9.m9.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.9.m9.1c">\pi_{4}</annotation></semantics></math>. Similarly, at the third level <math id="S2.SS6.p2.10.m10.1" class="ltx_Math" alttext="\pi_{1}" display="inline"><semantics id="S2.SS6.p2.10.m10.1a"><msub id="S2.SS6.p2.10.m10.1.1" xref="S2.SS6.p2.10.m10.1.1.cmml"><mi id="S2.SS6.p2.10.m10.1.1.2" xref="S2.SS6.p2.10.m10.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.10.m10.1.1.3" xref="S2.SS6.p2.10.m10.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.10.m10.1b"><apply id="S2.SS6.p2.10.m10.1.1.cmml" xref="S2.SS6.p2.10.m10.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.10.m10.1.1.1.cmml" xref="S2.SS6.p2.10.m10.1.1">subscript</csymbol><ci id="S2.SS6.p2.10.m10.1.1.2.cmml" xref="S2.SS6.p2.10.m10.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.10.m10.1.1.3.cmml" xref="S2.SS6.p2.10.m10.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.10.m10.1c">\pi_{1}</annotation></semantics></math> and <math id="S2.SS6.p2.11.m11.1" class="ltx_Math" alttext="\pi_{3}" display="inline"><semantics id="S2.SS6.p2.11.m11.1a"><msub id="S2.SS6.p2.11.m11.1.1" xref="S2.SS6.p2.11.m11.1.1.cmml"><mi id="S2.SS6.p2.11.m11.1.1.2" xref="S2.SS6.p2.11.m11.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.11.m11.1.1.3" xref="S2.SS6.p2.11.m11.1.1.3.cmml">3</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.11.m11.1b"><apply id="S2.SS6.p2.11.m11.1.1.cmml" xref="S2.SS6.p2.11.m11.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.11.m11.1.1.1.cmml" xref="S2.SS6.p2.11.m11.1.1">subscript</csymbol><ci id="S2.SS6.p2.11.m11.1.1.2.cmml" xref="S2.SS6.p2.11.m11.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.11.m11.1.1.3.cmml" xref="S2.SS6.p2.11.m11.1.1.3">3</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.11.m11.1c">\pi_{3}</annotation></semantics></math> share parameters, and <math id="S2.SS6.p2.12.m12.1" class="ltx_Math" alttext="\pi_{2}" display="inline"><semantics id="S2.SS6.p2.12.m12.1a"><msub id="S2.SS6.p2.12.m12.1.1" xref="S2.SS6.p2.12.m12.1.1.cmml"><mi id="S2.SS6.p2.12.m12.1.1.2" xref="S2.SS6.p2.12.m12.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.12.m12.1.1.3" xref="S2.SS6.p2.12.m12.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.12.m12.1b"><apply id="S2.SS6.p2.12.m12.1.1.cmml" xref="S2.SS6.p2.12.m12.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.12.m12.1.1.1.cmml" xref="S2.SS6.p2.12.m12.1.1">subscript</csymbol><ci id="S2.SS6.p2.12.m12.1.1.2.cmml" xref="S2.SS6.p2.12.m12.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.12.m12.1.1.3.cmml" xref="S2.SS6.p2.12.m12.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.12.m12.1c">\pi_{2}</annotation></semantics></math> and <math id="S2.SS6.p2.13.m13.1" class="ltx_Math" alttext="\pi_{4}" display="inline"><semantics id="S2.SS6.p2.13.m13.1a"><msub id="S2.SS6.p2.13.m13.1.1" xref="S2.SS6.p2.13.m13.1.1.cmml"><mi id="S2.SS6.p2.13.m13.1.1.2" xref="S2.SS6.p2.13.m13.1.1.2.cmml">π</mi><mn id="S2.SS6.p2.13.m13.1.1.3" xref="S2.SS6.p2.13.m13.1.1.3.cmml">4</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p2.13.m13.1b"><apply id="S2.SS6.p2.13.m13.1.1.cmml" xref="S2.SS6.p2.13.m13.1.1"><csymbol cd="ambiguous" id="S2.SS6.p2.13.m13.1.1.1.cmml" xref="S2.SS6.p2.13.m13.1.1">subscript</csymbol><ci id="S2.SS6.p2.13.m13.1.1.2.cmml" xref="S2.SS6.p2.13.m13.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p2.13.m13.1.1.3.cmml" xref="S2.SS6.p2.13.m13.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p2.13.m13.1c">\pi_{4}</annotation></semantics></math> share parameters.</p>
</div>
<div id="S2.SS6.p3" class="ltx_para">
<p id="S2.SS6.p3.2" class="ltx_p">The resulting <math id="S2.SS6.p3.1.m1.1" class="ltx_Math" alttext="2\times 2" display="inline"><semantics id="S2.SS6.p3.1.m1.1a"><mrow id="S2.SS6.p3.1.m1.1.1" xref="S2.SS6.p3.1.m1.1.1.cmml"><mn id="S2.SS6.p3.1.m1.1.1.2" xref="S2.SS6.p3.1.m1.1.1.2.cmml">2</mn><mo lspace="0.222em" rspace="0.222em" id="S2.SS6.p3.1.m1.1.1.1" xref="S2.SS6.p3.1.m1.1.1.1.cmml">×</mo><mn id="S2.SS6.p3.1.m1.1.1.3" xref="S2.SS6.p3.1.m1.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.p3.1.m1.1b"><apply id="S2.SS6.p3.1.m1.1.1.cmml" xref="S2.SS6.p3.1.m1.1.1"><times id="S2.SS6.p3.1.m1.1.1.1.cmml" xref="S2.SS6.p3.1.m1.1.1.1"></times><cn type="integer" id="S2.SS6.p3.1.m1.1.1.2.cmml" xref="S2.SS6.p3.1.m1.1.1.2">2</cn><cn type="integer" id="S2.SS6.p3.1.m1.1.1.3.cmml" xref="S2.SS6.p3.1.m1.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p3.1.m1.1c">2\times 2</annotation></semantics></math> DiPaCo has <math id="S2.SS6.p3.2.m2.1" class="ltx_Math" alttext="4" display="inline"><semantics id="S2.SS6.p3.2.m2.1a"><mn id="S2.SS6.p3.2.m2.1.1" xref="S2.SS6.p3.2.m2.1.1.cmml">4</mn><annotation-xml encoding="MathML-Content" id="S2.SS6.p3.2.m2.1b"><cn type="integer" id="S2.SS6.p3.2.m2.1.1.cmml" xref="S2.SS6.p3.2.m2.1.1">4</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p3.2.m2.1c">4</annotation></semantics></math> paths in total (as shown on the middle panel of <a href="#S2.F4" title="Figure 4 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 4</span></a>). However, the full model need never be fully instantiated, neither during training nor testing. Only paths are realized (as shown on the right panel of <a href="#S2.F4" title="Figure 4 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 4</span></a>), and these are trained in parallel with infrequent communication.</p>
</div>
<div id="S2.SS6.p4" class="ltx_para">
<p id="S2.SS6.p4.2" class="ltx_p">At training time DiLoCo is applied at the level of modules, as opposed to the entire network; see lines 11-16 of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>. The outer optimizer updates independently the parameters of each module after receiving the corresponding outer gradients from workers processing paths passing through that particular module. For instance and with reference to <a href="#S2.F4" title="Figure 4 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 4</span></a>, the outer gradient of module 3b is calculated by averaging the gradients of module 3b of the worker processing <math id="S2.SS6.p4.1.m1.1" class="ltx_Math" alttext="\pi_{2}" display="inline"><semantics id="S2.SS6.p4.1.m1.1a"><msub id="S2.SS6.p4.1.m1.1.1" xref="S2.SS6.p4.1.m1.1.1.cmml"><mi id="S2.SS6.p4.1.m1.1.1.2" xref="S2.SS6.p4.1.m1.1.1.2.cmml">π</mi><mn id="S2.SS6.p4.1.m1.1.1.3" xref="S2.SS6.p4.1.m1.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p4.1.m1.1b"><apply id="S2.SS6.p4.1.m1.1.1.cmml" xref="S2.SS6.p4.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS6.p4.1.m1.1.1.1.cmml" xref="S2.SS6.p4.1.m1.1.1">subscript</csymbol><ci id="S2.SS6.p4.1.m1.1.1.2.cmml" xref="S2.SS6.p4.1.m1.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p4.1.m1.1.1.3.cmml" xref="S2.SS6.p4.1.m1.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p4.1.m1.1c">\pi_{2}</annotation></semantics></math> and <math id="S2.SS6.p4.2.m2.1" class="ltx_Math" alttext="\pi_{4}" display="inline"><semantics id="S2.SS6.p4.2.m2.1a"><msub id="S2.SS6.p4.2.m2.1.1" xref="S2.SS6.p4.2.m2.1.1.cmml"><mi id="S2.SS6.p4.2.m2.1.1.2" xref="S2.SS6.p4.2.m2.1.1.2.cmml">π</mi><mn id="S2.SS6.p4.2.m2.1.1.3" xref="S2.SS6.p4.2.m2.1.1.3.cmml">4</mn></msub><annotation-xml encoding="MathML-Content" id="S2.SS6.p4.2.m2.1b"><apply id="S2.SS6.p4.2.m2.1.1.cmml" xref="S2.SS6.p4.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS6.p4.2.m2.1.1.1.cmml" xref="S2.SS6.p4.2.m2.1.1">subscript</csymbol><ci id="S2.SS6.p4.2.m2.1.1.2.cmml" xref="S2.SS6.p4.2.m2.1.1.2">𝜋</ci><cn type="integer" id="S2.SS6.p4.2.m2.1.1.3.cmml" xref="S2.SS6.p4.2.m2.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.p4.2.m2.1c">\pi_{4}</annotation></semantics></math>. Such outer gradient is then used to update the parameters of module 3b.
A worker does not need to be powerful enough (in terms of memory and compute) to host the entire model, but just a single path. Like in DiLoCo work <cite class="ltx_cite ltx_citemacro_citep">(Douillard et al., <a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite>, the inner optimizer is AdamW and the outer optimizer is Nesterov momentum.
Similarly, at test time, the paths are instantiated, and served independently, with text routed to each path via a router.</p>
</div>
<figure id="alg1" class="ltx_float ltx_float_algorithm ltx_framed ltx_framed_top">
<figcaption class="ltx_caption"><span class="ltx_tag ltx_tag_float"><span id="alg1.2.1.1" class="ltx_text ltx_font_bold">Algorithm 1</span> </span> DiPaCo (see notation in <a href="#S2.SS3" title="2.3 Notation ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">section</span> <span class="ltx_text ltx_ref_tag">2.3</span></a>)</figcaption>
<div id="alg1.3" class="ltx_listing ltx_listing">
<div id="algx1.l1" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l1.1.1.1" class="ltx_text" style="font-size:80%;">1:</span></span>Num. levels <math id="algx1.l1.m1.1" class="ltx_Math" alttext="L" display="inline"><semantics id="algx1.l1.m1.1a"><mi id="algx1.l1.m1.1.1" xref="algx1.l1.m1.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="algx1.l1.m1.1b"><ci id="algx1.l1.m1.1.1.cmml" xref="algx1.l1.m1.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m1.1c">L</annotation></semantics></math>, num. experts per level <math id="algx1.l1.m2.1" class="ltx_Math" alttext="K_{l}" display="inline"><semantics id="algx1.l1.m2.1a"><msub id="algx1.l1.m2.1.1" xref="algx1.l1.m2.1.1.cmml"><mi id="algx1.l1.m2.1.1.2" xref="algx1.l1.m2.1.1.2.cmml">K</mi><mi id="algx1.l1.m2.1.1.3" xref="algx1.l1.m2.1.1.3.cmml">l</mi></msub><annotation-xml encoding="MathML-Content" id="algx1.l1.m2.1b"><apply id="algx1.l1.m2.1.1.cmml" xref="algx1.l1.m2.1.1"><csymbol cd="ambiguous" id="algx1.l1.m2.1.1.1.cmml" xref="algx1.l1.m2.1.1">subscript</csymbol><ci id="algx1.l1.m2.1.1.2.cmml" xref="algx1.l1.m2.1.1.2">𝐾</ci><ci id="algx1.l1.m2.1.1.3.cmml" xref="algx1.l1.m2.1.1.3">𝑙</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m2.1c">K_{l}</annotation></semantics></math>, paths <math id="algx1.l1.m3.1" class="ltx_Math" alttext="\pi_{i}" display="inline"><semantics id="algx1.l1.m3.1a"><msub id="algx1.l1.m3.1.1" xref="algx1.l1.m3.1.1.cmml"><mi id="algx1.l1.m3.1.1.2" xref="algx1.l1.m3.1.1.2.cmml">π</mi><mi id="algx1.l1.m3.1.1.3" xref="algx1.l1.m3.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="algx1.l1.m3.1b"><apply id="algx1.l1.m3.1.1.cmml" xref="algx1.l1.m3.1.1"><csymbol cd="ambiguous" id="algx1.l1.m3.1.1.1.cmml" xref="algx1.l1.m3.1.1">subscript</csymbol><ci id="algx1.l1.m3.1.1.2.cmml" xref="algx1.l1.m3.1.1.2">𝜋</ci><ci id="algx1.l1.m3.1.1.3.cmml" xref="algx1.l1.m3.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m3.1c">\pi_{i}</annotation></semantics></math> with <math id="algx1.l1.m4.3" class="ltx_Math" alttext="i\in{1,\ldots,P}" display="inline"><semantics id="algx1.l1.m4.3a"><mrow id="algx1.l1.m4.3.4" xref="algx1.l1.m4.3.4.cmml"><mi id="algx1.l1.m4.3.4.2" xref="algx1.l1.m4.3.4.2.cmml">i</mi><mo id="algx1.l1.m4.3.4.1" xref="algx1.l1.m4.3.4.1.cmml">∈</mo><mrow id="algx1.l1.m4.3.4.3.2" xref="algx1.l1.m4.3.4.3.1.cmml"><mn id="algx1.l1.m4.1.1" xref="algx1.l1.m4.1.1.cmml">1</mn><mo id="algx1.l1.m4.3.4.3.2.1" xref="algx1.l1.m4.3.4.3.1.cmml">,</mo><mi mathvariant="normal" id="algx1.l1.m4.2.2" xref="algx1.l1.m4.2.2.cmml">…</mi><mo id="algx1.l1.m4.3.4.3.2.2" xref="algx1.l1.m4.3.4.3.1.cmml">,</mo><mi id="algx1.l1.m4.3.3" xref="algx1.l1.m4.3.3.cmml">P</mi></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l1.m4.3b"><apply id="algx1.l1.m4.3.4.cmml" xref="algx1.l1.m4.3.4"><in id="algx1.l1.m4.3.4.1.cmml" xref="algx1.l1.m4.3.4.1"></in><ci id="algx1.l1.m4.3.4.2.cmml" xref="algx1.l1.m4.3.4.2">𝑖</ci><list id="algx1.l1.m4.3.4.3.1.cmml" xref="algx1.l1.m4.3.4.3.2"><cn type="integer" id="algx1.l1.m4.1.1.cmml" xref="algx1.l1.m4.1.1">1</cn><ci id="algx1.l1.m4.2.2.cmml" xref="algx1.l1.m4.2.2">…</ci><ci id="algx1.l1.m4.3.3.cmml" xref="algx1.l1.m4.3.3">𝑃</ci></list></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m4.3c">i\in{1,\ldots,P}</annotation></semantics></math> and <math id="algx1.l1.m5.1" class="ltx_Math" alttext="P=\prod_{i=1}^{L}K_{i}" display="inline"><semantics id="algx1.l1.m5.1a"><mrow id="algx1.l1.m5.1.1" xref="algx1.l1.m5.1.1.cmml"><mi id="algx1.l1.m5.1.1.2" xref="algx1.l1.m5.1.1.2.cmml">P</mi><mo rspace="0.111em" id="algx1.l1.m5.1.1.1" xref="algx1.l1.m5.1.1.1.cmml">=</mo><mrow id="algx1.l1.m5.1.1.3" xref="algx1.l1.m5.1.1.3.cmml"><msubsup id="algx1.l1.m5.1.1.3.1" xref="algx1.l1.m5.1.1.3.1.cmml"><mo id="algx1.l1.m5.1.1.3.1.2.2" xref="algx1.l1.m5.1.1.3.1.2.2.cmml">∏</mo><mrow id="algx1.l1.m5.1.1.3.1.2.3" xref="algx1.l1.m5.1.1.3.1.2.3.cmml"><mi id="algx1.l1.m5.1.1.3.1.2.3.2" xref="algx1.l1.m5.1.1.3.1.2.3.2.cmml">i</mi><mo id="algx1.l1.m5.1.1.3.1.2.3.1" xref="algx1.l1.m5.1.1.3.1.2.3.1.cmml">=</mo><mn id="algx1.l1.m5.1.1.3.1.2.3.3" xref="algx1.l1.m5.1.1.3.1.2.3.3.cmml">1</mn></mrow><mi id="algx1.l1.m5.1.1.3.1.3" xref="algx1.l1.m5.1.1.3.1.3.cmml">L</mi></msubsup><msub id="algx1.l1.m5.1.1.3.2" xref="algx1.l1.m5.1.1.3.2.cmml"><mi id="algx1.l1.m5.1.1.3.2.2" xref="algx1.l1.m5.1.1.3.2.2.cmml">K</mi><mi id="algx1.l1.m5.1.1.3.2.3" xref="algx1.l1.m5.1.1.3.2.3.cmml">i</mi></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l1.m5.1b"><apply id="algx1.l1.m5.1.1.cmml" xref="algx1.l1.m5.1.1"><eq id="algx1.l1.m5.1.1.1.cmml" xref="algx1.l1.m5.1.1.1"></eq><ci id="algx1.l1.m5.1.1.2.cmml" xref="algx1.l1.m5.1.1.2">𝑃</ci><apply id="algx1.l1.m5.1.1.3.cmml" xref="algx1.l1.m5.1.1.3"><apply id="algx1.l1.m5.1.1.3.1.cmml" xref="algx1.l1.m5.1.1.3.1"><csymbol cd="ambiguous" id="algx1.l1.m5.1.1.3.1.1.cmml" xref="algx1.l1.m5.1.1.3.1">superscript</csymbol><apply id="algx1.l1.m5.1.1.3.1.2.cmml" xref="algx1.l1.m5.1.1.3.1"><csymbol cd="ambiguous" id="algx1.l1.m5.1.1.3.1.2.1.cmml" xref="algx1.l1.m5.1.1.3.1">subscript</csymbol><csymbol cd="latexml" id="algx1.l1.m5.1.1.3.1.2.2.cmml" xref="algx1.l1.m5.1.1.3.1.2.2">product</csymbol><apply id="algx1.l1.m5.1.1.3.1.2.3.cmml" xref="algx1.l1.m5.1.1.3.1.2.3"><eq id="algx1.l1.m5.1.1.3.1.2.3.1.cmml" xref="algx1.l1.m5.1.1.3.1.2.3.1"></eq><ci id="algx1.l1.m5.1.1.3.1.2.3.2.cmml" xref="algx1.l1.m5.1.1.3.1.2.3.2">𝑖</ci><cn type="integer" id="algx1.l1.m5.1.1.3.1.2.3.3.cmml" xref="algx1.l1.m5.1.1.3.1.2.3.3">1</cn></apply></apply><ci id="algx1.l1.m5.1.1.3.1.3.cmml" xref="algx1.l1.m5.1.1.3.1.3">𝐿</ci></apply><apply id="algx1.l1.m5.1.1.3.2.cmml" xref="algx1.l1.m5.1.1.3.2"><csymbol cd="ambiguous" id="algx1.l1.m5.1.1.3.2.1.cmml" xref="algx1.l1.m5.1.1.3.2">subscript</csymbol><ci id="algx1.l1.m5.1.1.3.2.2.cmml" xref="algx1.l1.m5.1.1.3.2.2">𝐾</ci><ci id="algx1.l1.m5.1.1.3.2.3.cmml" xref="algx1.l1.m5.1.1.3.2.3">𝑖</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m5.1c">P=\prod_{i=1}^{L}K_{i}</annotation></semantics></math>. Let <math id="algx1.l1.m6.2" class="ltx_Math" alttext="P_{l,e}" display="inline"><semantics id="algx1.l1.m6.2a"><msub id="algx1.l1.m6.2.3" xref="algx1.l1.m6.2.3.cmml"><mi id="algx1.l1.m6.2.3.2" xref="algx1.l1.m6.2.3.2.cmml">P</mi><mrow id="algx1.l1.m6.2.2.2.4" xref="algx1.l1.m6.2.2.2.3.cmml"><mi id="algx1.l1.m6.1.1.1.1" xref="algx1.l1.m6.1.1.1.1.cmml">l</mi><mo id="algx1.l1.m6.2.2.2.4.1" xref="algx1.l1.m6.2.2.2.3.cmml">,</mo><mi id="algx1.l1.m6.2.2.2.2" xref="algx1.l1.m6.2.2.2.2.cmml">e</mi></mrow></msub><annotation-xml encoding="MathML-Content" id="algx1.l1.m6.2b"><apply id="algx1.l1.m6.2.3.cmml" xref="algx1.l1.m6.2.3"><csymbol cd="ambiguous" id="algx1.l1.m6.2.3.1.cmml" xref="algx1.l1.m6.2.3">subscript</csymbol><ci id="algx1.l1.m6.2.3.2.cmml" xref="algx1.l1.m6.2.3.2">𝑃</ci><list id="algx1.l1.m6.2.2.2.3.cmml" xref="algx1.l1.m6.2.2.2.4"><ci id="algx1.l1.m6.1.1.1.1.cmml" xref="algx1.l1.m6.1.1.1.1">𝑙</ci><ci id="algx1.l1.m6.2.2.2.2.cmml" xref="algx1.l1.m6.2.2.2.2">𝑒</ci></list></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m6.2c">P_{l,e}</annotation></semantics></math> be the number of paths going through module <math id="algx1.l1.m7.1" class="ltx_Math" alttext="e" display="inline"><semantics id="algx1.l1.m7.1a"><mi id="algx1.l1.m7.1.1" xref="algx1.l1.m7.1.1.cmml">e</mi><annotation-xml encoding="MathML-Content" id="algx1.l1.m7.1b"><ci id="algx1.l1.m7.1.1.cmml" xref="algx1.l1.m7.1.1">𝑒</ci></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m7.1c">e</annotation></semantics></math> at level <math id="algx1.l1.m8.1" class="ltx_Math" alttext="l" display="inline"><semantics id="algx1.l1.m8.1a"><mi id="algx1.l1.m8.1.1" xref="algx1.l1.m8.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="algx1.l1.m8.1b"><ci id="algx1.l1.m8.1.1.cmml" xref="algx1.l1.m8.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="algx1.l1.m8.1c">l</annotation></semantics></math>.

</div>
<div id="algx1.l2" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l2.1.1.1" class="ltx_text" style="font-size:80%;">2:</span></span>Pre-sharded training set into <math id="algx1.l2.m1.3" class="ltx_Math" alttext="\{\mathcal{D}_{1},\ldots,\mathcal{D}_{P}\}" display="inline"><semantics id="algx1.l2.m1.3a"><mrow id="algx1.l2.m1.3.3.2" xref="algx1.l2.m1.3.3.3.cmml"><mo stretchy="false" id="algx1.l2.m1.3.3.2.3" xref="algx1.l2.m1.3.3.3.cmml">{</mo><msub id="algx1.l2.m1.2.2.1.1" xref="algx1.l2.m1.2.2.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="algx1.l2.m1.2.2.1.1.2" xref="algx1.l2.m1.2.2.1.1.2.cmml">𝒟</mi><mn id="algx1.l2.m1.2.2.1.1.3" xref="algx1.l2.m1.2.2.1.1.3.cmml">1</mn></msub><mo id="algx1.l2.m1.3.3.2.4" xref="algx1.l2.m1.3.3.3.cmml">,</mo><mi mathvariant="normal" id="algx1.l2.m1.1.1" xref="algx1.l2.m1.1.1.cmml">…</mi><mo id="algx1.l2.m1.3.3.2.5" xref="algx1.l2.m1.3.3.3.cmml">,</mo><msub id="algx1.l2.m1.3.3.2.2" xref="algx1.l2.m1.3.3.2.2.cmml"><mi class="ltx_font_mathcaligraphic" id="algx1.l2.m1.3.3.2.2.2" xref="algx1.l2.m1.3.3.2.2.2.cmml">𝒟</mi><mi id="algx1.l2.m1.3.3.2.2.3" xref="algx1.l2.m1.3.3.2.2.3.cmml">P</mi></msub><mo stretchy="false" id="algx1.l2.m1.3.3.2.6" xref="algx1.l2.m1.3.3.3.cmml">}</mo></mrow><annotation-xml encoding="MathML-Content" id="algx1.l2.m1.3b"><set id="algx1.l2.m1.3.3.3.cmml" xref="algx1.l2.m1.3.3.2"><apply id="algx1.l2.m1.2.2.1.1.cmml" xref="algx1.l2.m1.2.2.1.1"><csymbol cd="ambiguous" id="algx1.l2.m1.2.2.1.1.1.cmml" xref="algx1.l2.m1.2.2.1.1">subscript</csymbol><ci id="algx1.l2.m1.2.2.1.1.2.cmml" xref="algx1.l2.m1.2.2.1.1.2">𝒟</ci><cn type="integer" id="algx1.l2.m1.2.2.1.1.3.cmml" xref="algx1.l2.m1.2.2.1.1.3">1</cn></apply><ci id="algx1.l2.m1.1.1.cmml" xref="algx1.l2.m1.1.1">…</ci><apply id="algx1.l2.m1.3.3.2.2.cmml" xref="algx1.l2.m1.3.3.2.2"><csymbol cd="ambiguous" id="algx1.l2.m1.3.3.2.2.1.cmml" xref="algx1.l2.m1.3.3.2.2">subscript</csymbol><ci id="algx1.l2.m1.3.3.2.2.2.cmml" xref="algx1.l2.m1.3.3.2.2.2">𝒟</ci><ci id="algx1.l2.m1.3.3.2.2.3.cmml" xref="algx1.l2.m1.3.3.2.2.3">𝑃</ci></apply></set></annotation-xml><annotation encoding="application/x-tex" id="algx1.l2.m1.3c">\{\mathcal{D}_{1},\ldots,\mathcal{D}_{P}\}</annotation></semantics></math> (see <a href="#S2.SS4" title="2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">section</span> <span class="ltx_text ltx_ref_tag">2.4</span></a>).

</div>
<div id="algx1.l3" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l3.1.1.1" class="ltx_text" style="font-size:80%;">3:</span></span>Parameters of pretrained model used for initialization: <math id="algx1.l3.m1.1" class="ltx_Math" alttext="\bar{\theta}" display="inline"><semantics id="algx1.l3.m1.1a"><mover accent="true" id="algx1.l3.m1.1.1" xref="algx1.l3.m1.1.1.cmml"><mi id="algx1.l3.m1.1.1.2" xref="algx1.l3.m1.1.1.2.cmml">θ</mi><mo id="algx1.l3.m1.1.1.1" xref="algx1.l3.m1.1.1.1.cmml">¯</mo></mover><annotation-xml encoding="MathML-Content" id="algx1.l3.m1.1b"><apply id="algx1.l3.m1.1.1.cmml" xref="algx1.l3.m1.1.1"><ci id="algx1.l3.m1.1.1.1.cmml" xref="algx1.l3.m1.1.1.1">¯</ci><ci id="algx1.l3.m1.1.1.2.cmml" xref="algx1.l3.m1.1.1.2">𝜃</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l3.m1.1c">\bar{\theta}</annotation></semantics></math>. <math id="algx1.l3.m2.4" class="ltx_Math" alttext="\theta_{i}^{0}=\bar{\theta},i\in{1,\ldots,P}." display="inline"><semantics id="algx1.l3.m2.4a"><mrow id="algx1.l3.m2.4.4.1"><mrow id="algx1.l3.m2.4.4.1.1.2" xref="algx1.l3.m2.4.4.1.1.3.cmml"><mrow id="algx1.l3.m2.4.4.1.1.1.1" xref="algx1.l3.m2.4.4.1.1.1.1.cmml"><msubsup id="algx1.l3.m2.4.4.1.1.1.1.2" xref="algx1.l3.m2.4.4.1.1.1.1.2.cmml"><mi id="algx1.l3.m2.4.4.1.1.1.1.2.2.2" xref="algx1.l3.m2.4.4.1.1.1.1.2.2.2.cmml">θ</mi><mi id="algx1.l3.m2.4.4.1.1.1.1.2.2.3" xref="algx1.l3.m2.4.4.1.1.1.1.2.2.3.cmml">i</mi><mn id="algx1.l3.m2.4.4.1.1.1.1.2.3" xref="algx1.l3.m2.4.4.1.1.1.1.2.3.cmml">0</mn></msubsup><mo id="algx1.l3.m2.4.4.1.1.1.1.1" xref="algx1.l3.m2.4.4.1.1.1.1.1.cmml">=</mo><mover accent="true" id="algx1.l3.m2.4.4.1.1.1.1.3" xref="algx1.l3.m2.4.4.1.1.1.1.3.cmml"><mi id="algx1.l3.m2.4.4.1.1.1.1.3.2" xref="algx1.l3.m2.4.4.1.1.1.1.3.2.cmml">θ</mi><mo id="algx1.l3.m2.4.4.1.1.1.1.3.1" xref="algx1.l3.m2.4.4.1.1.1.1.3.1.cmml">¯</mo></mover></mrow><mo id="algx1.l3.m2.4.4.1.1.2.3" xref="algx1.l3.m2.4.4.1.1.3a.cmml">,</mo><mrow id="algx1.l3.m2.4.4.1.1.2.2" xref="algx1.l3.m2.4.4.1.1.2.2.cmml"><mi id="algx1.l3.m2.4.4.1.1.2.2.2" xref="algx1.l3.m2.4.4.1.1.2.2.2.cmml">i</mi><mo id="algx1.l3.m2.4.4.1.1.2.2.1" xref="algx1.l3.m2.4.4.1.1.2.2.1.cmml">∈</mo><mrow id="algx1.l3.m2.4.4.1.1.2.2.3.2" xref="algx1.l3.m2.4.4.1.1.2.2.3.1.cmml"><mn id="algx1.l3.m2.1.1" xref="algx1.l3.m2.1.1.cmml">1</mn><mo id="algx1.l3.m2.4.4.1.1.2.2.3.2.1" xref="algx1.l3.m2.4.4.1.1.2.2.3.1.cmml">,</mo><mi mathvariant="normal" id="algx1.l3.m2.2.2" xref="algx1.l3.m2.2.2.cmml">…</mi><mo id="algx1.l3.m2.4.4.1.1.2.2.3.2.2" xref="algx1.l3.m2.4.4.1.1.2.2.3.1.cmml">,</mo><mi id="algx1.l3.m2.3.3" xref="algx1.l3.m2.3.3.cmml">P</mi></mrow></mrow></mrow><mo lspace="0em" id="algx1.l3.m2.4.4.1.2">.</mo></mrow><annotation-xml encoding="MathML-Content" id="algx1.l3.m2.4b"><apply id="algx1.l3.m2.4.4.1.1.3.cmml" xref="algx1.l3.m2.4.4.1.1.2"><csymbol cd="ambiguous" id="algx1.l3.m2.4.4.1.1.3a.cmml" xref="algx1.l3.m2.4.4.1.1.2.3">formulae-sequence</csymbol><apply id="algx1.l3.m2.4.4.1.1.1.1.cmml" xref="algx1.l3.m2.4.4.1.1.1.1"><eq id="algx1.l3.m2.4.4.1.1.1.1.1.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.1"></eq><apply id="algx1.l3.m2.4.4.1.1.1.1.2.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2"><csymbol cd="ambiguous" id="algx1.l3.m2.4.4.1.1.1.1.2.1.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2">superscript</csymbol><apply id="algx1.l3.m2.4.4.1.1.1.1.2.2.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2"><csymbol cd="ambiguous" id="algx1.l3.m2.4.4.1.1.1.1.2.2.1.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2">subscript</csymbol><ci id="algx1.l3.m2.4.4.1.1.1.1.2.2.2.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2.2.2">𝜃</ci><ci id="algx1.l3.m2.4.4.1.1.1.1.2.2.3.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2.2.3">𝑖</ci></apply><cn type="integer" id="algx1.l3.m2.4.4.1.1.1.1.2.3.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.2.3">0</cn></apply><apply id="algx1.l3.m2.4.4.1.1.1.1.3.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.3"><ci id="algx1.l3.m2.4.4.1.1.1.1.3.1.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.3.1">¯</ci><ci id="algx1.l3.m2.4.4.1.1.1.1.3.2.cmml" xref="algx1.l3.m2.4.4.1.1.1.1.3.2">𝜃</ci></apply></apply><apply id="algx1.l3.m2.4.4.1.1.2.2.cmml" xref="algx1.l3.m2.4.4.1.1.2.2"><in id="algx1.l3.m2.4.4.1.1.2.2.1.cmml" xref="algx1.l3.m2.4.4.1.1.2.2.1"></in><ci id="algx1.l3.m2.4.4.1.1.2.2.2.cmml" xref="algx1.l3.m2.4.4.1.1.2.2.2">𝑖</ci><list id="algx1.l3.m2.4.4.1.1.2.2.3.1.cmml" xref="algx1.l3.m2.4.4.1.1.2.2.3.2"><cn type="integer" id="algx1.l3.m2.1.1.cmml" xref="algx1.l3.m2.1.1">1</cn><ci id="algx1.l3.m2.2.2.cmml" xref="algx1.l3.m2.2.2">…</ci><ci id="algx1.l3.m2.3.3.cmml" xref="algx1.l3.m2.3.3">𝑃</ci></list></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l3.m2.4c">\theta_{i}^{0}=\bar{\theta},i\in{1,\ldots,P}.</annotation></semantics></math>

</div>
<div id="algx1.l4" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l4.1.1.1" class="ltx_text" style="font-size:80%;">4:</span></span>Optimizers <span id="algx1.l4.2" class="ltx_text ltx_markedasmath ltx_font_typewriter">InnerOpt</span> and <span id="algx1.l4.3" class="ltx_text ltx_markedasmath ltx_font_typewriter">OuterOpt</span>

</div>
<div id="algx1.l5" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l5.2.1.1" class="ltx_text" style="font-size:80%;">5:</span></span><span id="algx1.l5.3" class="ltx_text ltx_font_bold">for</span> <span id="algx1.l5.1" class="ltx_text ltx_font_typewriter">outer step <math id="algx1.l5.1.m1.1" class="ltx_Math" alttext="t=1\ldots T" display="inline"><semantics id="algx1.l5.1.m1.1a"><mrow id="algx1.l5.1.m1.1.1" xref="algx1.l5.1.m1.1.1.cmml"><mi id="algx1.l5.1.m1.1.1.2" xref="algx1.l5.1.m1.1.1.2.cmml">t</mi><mo id="algx1.l5.1.m1.1.1.1" xref="algx1.l5.1.m1.1.1.1.cmml">=</mo><mrow id="algx1.l5.1.m1.1.1.3" xref="algx1.l5.1.m1.1.1.3.cmml"><mn id="algx1.l5.1.m1.1.1.3.2" xref="algx1.l5.1.m1.1.1.3.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="algx1.l5.1.m1.1.1.3.1" xref="algx1.l5.1.m1.1.1.3.1.cmml">​</mo><mi mathvariant="normal" id="algx1.l5.1.m1.1.1.3.3" xref="algx1.l5.1.m1.1.1.3.3.cmml">…</mi><mo lspace="0em" rspace="0em" id="algx1.l5.1.m1.1.1.3.1a" xref="algx1.l5.1.m1.1.1.3.1.cmml">​</mo><mi id="algx1.l5.1.m1.1.1.3.4" xref="algx1.l5.1.m1.1.1.3.4.cmml">T</mi></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l5.1.m1.1b"><apply id="algx1.l5.1.m1.1.1.cmml" xref="algx1.l5.1.m1.1.1"><eq id="algx1.l5.1.m1.1.1.1.cmml" xref="algx1.l5.1.m1.1.1.1"></eq><ci id="algx1.l5.1.m1.1.1.2.cmml" xref="algx1.l5.1.m1.1.1.2">𝑡</ci><apply id="algx1.l5.1.m1.1.1.3.cmml" xref="algx1.l5.1.m1.1.1.3"><times id="algx1.l5.1.m1.1.1.3.1.cmml" xref="algx1.l5.1.m1.1.1.3.1"></times><cn type="integer" id="algx1.l5.1.m1.1.1.3.2.cmml" xref="algx1.l5.1.m1.1.1.3.2">1</cn><ci id="algx1.l5.1.m1.1.1.3.3.cmml" xref="algx1.l5.1.m1.1.1.3.3">…</ci><ci id="algx1.l5.1.m1.1.1.3.4.cmml" xref="algx1.l5.1.m1.1.1.3.4">𝑇</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l5.1.m1.1c">t=1\ldots T</annotation></semantics></math></span> <span id="algx1.l5.4" class="ltx_text ltx_font_bold">do</span>

</div>
<div id="algx1.l6" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l6.1.1.1" class="ltx_text" style="font-size:80%;">6:</span></span>     (Optional, in this work done once during training) discriminatively re-shard data (see <a href="#S2.SS4.SSS2" title="2.4.2 Discriminative routing ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">section</span> <span class="ltx_text ltx_ref_tag">2.4.2</span></a>)

</div>
<div id="algx1.l7" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l7.2.1.1" class="ltx_text" style="font-size:80%;">7:</span></span>     <span id="algx1.l7.3" class="ltx_text ltx_font_bold">for</span> <span id="algx1.l7.1" class="ltx_text ltx_font_typewriter">worker <math id="algx1.l7.1.m1.1" class="ltx_Math" alttext="i=1\ldots P" display="inline"><semantics id="algx1.l7.1.m1.1a"><mrow id="algx1.l7.1.m1.1.1" xref="algx1.l7.1.m1.1.1.cmml"><mi id="algx1.l7.1.m1.1.1.2" xref="algx1.l7.1.m1.1.1.2.cmml">i</mi><mo id="algx1.l7.1.m1.1.1.1" xref="algx1.l7.1.m1.1.1.1.cmml">=</mo><mrow id="algx1.l7.1.m1.1.1.3" xref="algx1.l7.1.m1.1.1.3.cmml"><mn id="algx1.l7.1.m1.1.1.3.2" xref="algx1.l7.1.m1.1.1.3.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="algx1.l7.1.m1.1.1.3.1" xref="algx1.l7.1.m1.1.1.3.1.cmml">​</mo><mi mathvariant="normal" id="algx1.l7.1.m1.1.1.3.3" xref="algx1.l7.1.m1.1.1.3.3.cmml">…</mi><mo lspace="0em" rspace="0em" id="algx1.l7.1.m1.1.1.3.1a" xref="algx1.l7.1.m1.1.1.3.1.cmml">​</mo><mi id="algx1.l7.1.m1.1.1.3.4" xref="algx1.l7.1.m1.1.1.3.4.cmml">P</mi></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l7.1.m1.1b"><apply id="algx1.l7.1.m1.1.1.cmml" xref="algx1.l7.1.m1.1.1"><eq id="algx1.l7.1.m1.1.1.1.cmml" xref="algx1.l7.1.m1.1.1.1"></eq><ci id="algx1.l7.1.m1.1.1.2.cmml" xref="algx1.l7.1.m1.1.1.2">𝑖</ci><apply id="algx1.l7.1.m1.1.1.3.cmml" xref="algx1.l7.1.m1.1.1.3"><times id="algx1.l7.1.m1.1.1.3.1.cmml" xref="algx1.l7.1.m1.1.1.3.1"></times><cn type="integer" id="algx1.l7.1.m1.1.1.3.2.cmml" xref="algx1.l7.1.m1.1.1.3.2">1</cn><ci id="algx1.l7.1.m1.1.1.3.3.cmml" xref="algx1.l7.1.m1.1.1.3.3">…</ci><ci id="algx1.l7.1.m1.1.1.3.4.cmml" xref="algx1.l7.1.m1.1.1.3.4">𝑃</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l7.1.m1.1c">i=1\ldots P</annotation></semantics></math></span> <span id="algx1.l7.4" class="ltx_text ltx_font_bold">do</span>

</div>
<div id="algx1.l8" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l8.1.1.1" class="ltx_text" style="font-size:80%;">8:</span></span>         <math id="algx1.l8.m1.1" class="ltx_Math" alttext="\theta_{i}^{t}=\theta_{i}^{t-1}" display="inline"><semantics id="algx1.l8.m1.1a"><mrow id="algx1.l8.m1.1.1" xref="algx1.l8.m1.1.1.cmml"><msubsup id="algx1.l8.m1.1.1.2" xref="algx1.l8.m1.1.1.2.cmml"><mi id="algx1.l8.m1.1.1.2.2.2" xref="algx1.l8.m1.1.1.2.2.2.cmml">θ</mi><mi id="algx1.l8.m1.1.1.2.2.3" xref="algx1.l8.m1.1.1.2.2.3.cmml">i</mi><mi id="algx1.l8.m1.1.1.2.3" xref="algx1.l8.m1.1.1.2.3.cmml">t</mi></msubsup><mo id="algx1.l8.m1.1.1.1" xref="algx1.l8.m1.1.1.1.cmml">=</mo><msubsup id="algx1.l8.m1.1.1.3" xref="algx1.l8.m1.1.1.3.cmml"><mi id="algx1.l8.m1.1.1.3.2.2" xref="algx1.l8.m1.1.1.3.2.2.cmml">θ</mi><mi id="algx1.l8.m1.1.1.3.2.3" xref="algx1.l8.m1.1.1.3.2.3.cmml">i</mi><mrow id="algx1.l8.m1.1.1.3.3" xref="algx1.l8.m1.1.1.3.3.cmml"><mi id="algx1.l8.m1.1.1.3.3.2" xref="algx1.l8.m1.1.1.3.3.2.cmml">t</mi><mo id="algx1.l8.m1.1.1.3.3.1" xref="algx1.l8.m1.1.1.3.3.1.cmml">−</mo><mn id="algx1.l8.m1.1.1.3.3.3" xref="algx1.l8.m1.1.1.3.3.3.cmml">1</mn></mrow></msubsup></mrow><annotation-xml encoding="MathML-Content" id="algx1.l8.m1.1b"><apply id="algx1.l8.m1.1.1.cmml" xref="algx1.l8.m1.1.1"><eq id="algx1.l8.m1.1.1.1.cmml" xref="algx1.l8.m1.1.1.1"></eq><apply id="algx1.l8.m1.1.1.2.cmml" xref="algx1.l8.m1.1.1.2"><csymbol cd="ambiguous" id="algx1.l8.m1.1.1.2.1.cmml" xref="algx1.l8.m1.1.1.2">superscript</csymbol><apply id="algx1.l8.m1.1.1.2.2.cmml" xref="algx1.l8.m1.1.1.2"><csymbol cd="ambiguous" id="algx1.l8.m1.1.1.2.2.1.cmml" xref="algx1.l8.m1.1.1.2">subscript</csymbol><ci id="algx1.l8.m1.1.1.2.2.2.cmml" xref="algx1.l8.m1.1.1.2.2.2">𝜃</ci><ci id="algx1.l8.m1.1.1.2.2.3.cmml" xref="algx1.l8.m1.1.1.2.2.3">𝑖</ci></apply><ci id="algx1.l8.m1.1.1.2.3.cmml" xref="algx1.l8.m1.1.1.2.3">𝑡</ci></apply><apply id="algx1.l8.m1.1.1.3.cmml" xref="algx1.l8.m1.1.1.3"><csymbol cd="ambiguous" id="algx1.l8.m1.1.1.3.1.cmml" xref="algx1.l8.m1.1.1.3">superscript</csymbol><apply id="algx1.l8.m1.1.1.3.2.cmml" xref="algx1.l8.m1.1.1.3"><csymbol cd="ambiguous" id="algx1.l8.m1.1.1.3.2.1.cmml" xref="algx1.l8.m1.1.1.3">subscript</csymbol><ci id="algx1.l8.m1.1.1.3.2.2.cmml" xref="algx1.l8.m1.1.1.3.2.2">𝜃</ci><ci id="algx1.l8.m1.1.1.3.2.3.cmml" xref="algx1.l8.m1.1.1.3.2.3">𝑖</ci></apply><apply id="algx1.l8.m1.1.1.3.3.cmml" xref="algx1.l8.m1.1.1.3.3"><minus id="algx1.l8.m1.1.1.3.3.1.cmml" xref="algx1.l8.m1.1.1.3.3.1"></minus><ci id="algx1.l8.m1.1.1.3.3.2.cmml" xref="algx1.l8.m1.1.1.3.3.2">𝑡</ci><cn type="integer" id="algx1.l8.m1.1.1.3.3.3.cmml" xref="algx1.l8.m1.1.1.3.3.3">1</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l8.m1.1c">\theta_{i}^{t}=\theta_{i}^{t-1}</annotation></semantics></math>

</div>
<div id="algx1.l9" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l9.2.1.1" class="ltx_text" style="font-size:80%;">9:</span></span>         <span id="algx1.l9.3" class="ltx_text ltx_font_bold">for</span> <span id="algx1.l9.1" class="ltx_text ltx_font_typewriter">inner step <math id="algx1.l9.1.m1.1" class="ltx_Math" alttext="n=1\ldots\tau" display="inline"><semantics id="algx1.l9.1.m1.1a"><mrow id="algx1.l9.1.m1.1.1" xref="algx1.l9.1.m1.1.1.cmml"><mi id="algx1.l9.1.m1.1.1.2" xref="algx1.l9.1.m1.1.1.2.cmml">n</mi><mo id="algx1.l9.1.m1.1.1.1" xref="algx1.l9.1.m1.1.1.1.cmml">=</mo><mrow id="algx1.l9.1.m1.1.1.3" xref="algx1.l9.1.m1.1.1.3.cmml"><mn id="algx1.l9.1.m1.1.1.3.2" xref="algx1.l9.1.m1.1.1.3.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="algx1.l9.1.m1.1.1.3.1" xref="algx1.l9.1.m1.1.1.3.1.cmml">​</mo><mi mathvariant="normal" id="algx1.l9.1.m1.1.1.3.3" xref="algx1.l9.1.m1.1.1.3.3.cmml">…</mi><mo lspace="0em" rspace="0em" id="algx1.l9.1.m1.1.1.3.1a" xref="algx1.l9.1.m1.1.1.3.1.cmml">​</mo><mi id="algx1.l9.1.m1.1.1.3.4" xref="algx1.l9.1.m1.1.1.3.4.cmml">τ</mi></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l9.1.m1.1b"><apply id="algx1.l9.1.m1.1.1.cmml" xref="algx1.l9.1.m1.1.1"><eq id="algx1.l9.1.m1.1.1.1.cmml" xref="algx1.l9.1.m1.1.1.1"></eq><ci id="algx1.l9.1.m1.1.1.2.cmml" xref="algx1.l9.1.m1.1.1.2">𝑛</ci><apply id="algx1.l9.1.m1.1.1.3.cmml" xref="algx1.l9.1.m1.1.1.3"><times id="algx1.l9.1.m1.1.1.3.1.cmml" xref="algx1.l9.1.m1.1.1.3.1"></times><cn type="integer" id="algx1.l9.1.m1.1.1.3.2.cmml" xref="algx1.l9.1.m1.1.1.3.2">1</cn><ci id="algx1.l9.1.m1.1.1.3.3.cmml" xref="algx1.l9.1.m1.1.1.3.3">…</ci><ci id="algx1.l9.1.m1.1.1.3.4.cmml" xref="algx1.l9.1.m1.1.1.3.4">𝜏</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l9.1.m1.1c">n=1\ldots\tau</annotation></semantics></math> </span> <span id="algx1.l9.4" class="ltx_text ltx_font_bold">do</span> 
</div>
<div id="algx1.l10" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l10.2.1.1" class="ltx_text" style="font-size:80%;">10:</span></span>              <math id="algx1.l10.m1.1" class="ltx_Math" alttext="x\sim\mathcal{D}_{i}" display="inline"><semantics id="algx1.l10.m1.1a"><mrow id="algx1.l10.m1.1.1" xref="algx1.l10.m1.1.1.cmml"><mi id="algx1.l10.m1.1.1.2" xref="algx1.l10.m1.1.1.2.cmml">x</mi><mo id="algx1.l10.m1.1.1.1" xref="algx1.l10.m1.1.1.1.cmml">∼</mo><msub id="algx1.l10.m1.1.1.3" xref="algx1.l10.m1.1.1.3.cmml"><mi class="ltx_font_mathcaligraphic" id="algx1.l10.m1.1.1.3.2" xref="algx1.l10.m1.1.1.3.2.cmml">𝒟</mi><mi id="algx1.l10.m1.1.1.3.3" xref="algx1.l10.m1.1.1.3.3.cmml">i</mi></msub></mrow><annotation-xml encoding="MathML-Content" id="algx1.l10.m1.1b"><apply id="algx1.l10.m1.1.1.cmml" xref="algx1.l10.m1.1.1"><csymbol cd="latexml" id="algx1.l10.m1.1.1.1.cmml" xref="algx1.l10.m1.1.1.1">similar-to</csymbol><ci id="algx1.l10.m1.1.1.2.cmml" xref="algx1.l10.m1.1.1.2">𝑥</ci><apply id="algx1.l10.m1.1.1.3.cmml" xref="algx1.l10.m1.1.1.3"><csymbol cd="ambiguous" id="algx1.l10.m1.1.1.3.1.cmml" xref="algx1.l10.m1.1.1.3">subscript</csymbol><ci id="algx1.l10.m1.1.1.3.2.cmml" xref="algx1.l10.m1.1.1.3.2">𝒟</ci><ci id="algx1.l10.m1.1.1.3.3.cmml" xref="algx1.l10.m1.1.1.3.3">𝑖</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l10.m1.1c">x\sim\mathcal{D}_{i}</annotation></semantics></math>
<span id="algx1.l10.1" class="ltx_text ltx_inline-block" style="width:0.0pt;">                                                <math id="algx1.l10.1.m1.1" class="ltx_Math" alttext="\left.\begin{array}[]{@{}c@{}}\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\end{array}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\right\}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\begin{tabular}[]{l}{Inner Optimization}: parallel per path\end{tabular}" display="inline"><semantics id="algx1.l10.1.m1.1a"><mrow id="algx1.l10.1.m1.1.2" xref="algx1.l10.1.m1.1.2.cmml"><mrow id="algx1.l10.1.m1.1.2.2.2" xref="algx1.l10.1.m1.1.2.2.1.cmml"><mi id="algx1.l10.1.m1.1.2.2.2.1" xref="algx1.l10.1.m1.1.2.2.1.1.cmml"></mi><mtable rowspacing="0pt" id="algx1.l10.1.m1.1.1" xref="algx1.l10.1.m1.1.1.cmml"><mtr id="algx1.l10.1.m1.1.1a" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1b" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l10.1.m1.1.1c" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1d" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l10.1.m1.1.1e" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1f" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l10.1.m1.1.1g" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1h" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l10.1.m1.1.1i" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1j" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l10.1.m1.1.1k" xref="algx1.l10.1.m1.1.1.cmml"><mtd id="algx1.l10.1.m1.1.1l" xref="algx1.l10.1.m1.1.1.cmml"></mtd></mtr></mtable><mo id="algx1.l10.1.m1.1.2.2.2.2" xref="algx1.l10.1.m1.1.2.2.1.1.cmml">}</mo></mrow><mo lspace="0em" rspace="0em" id="algx1.l10.1.m1.1.2.1" xref="algx1.l10.1.m1.1.2.1.cmml">​</mo><mtable id="algx1.l10.1.m1.1.2.3" xref="algx1.l10.1.m1.1.2.3.cmml"><mtr id="algx1.l10.1.m1.1.2.3a" xref="algx1.l10.1.m1.1.2.3.cmml"><mtd class="ltx_align_left" columnalign="left" id="algx1.l10.1.m1.1.2.3b" xref="algx1.l10.1.m1.1.2.3.cmml"><mtext class="ltx_mathvariant_bold" mathcolor="#2267EB" id="algx1.l10.1.m1.1.2.3.1.1.1" xref="algx1.l10.1.m1.1.2.3.1.1.1a.cmml">Inner Optimization</mtext><mtext mathcolor="#2267EB" id="algx1.l10.1.m1.1.2.3.1.1.2" xref="algx1.l10.1.m1.1.2.3.1.1.2a.cmml">: parallel per path</mtext></mtd></mtr></mtable></mrow><annotation-xml encoding="MathML-Content" id="algx1.l10.1.m1.1b"><apply id="algx1.l10.1.m1.1.2.cmml" xref="algx1.l10.1.m1.1.2"><times id="algx1.l10.1.m1.1.2.1.cmml" xref="algx1.l10.1.m1.1.2.1"></times><apply id="algx1.l10.1.m1.1.2.2.1.cmml" xref="algx1.l10.1.m1.1.2.2.2"><csymbol cd="latexml" id="algx1.l10.1.m1.1.2.2.1.1.cmml" xref="algx1.l10.1.m1.1.2.2.2.1">cases</csymbol><matrix id="algx1.l10.1.m1.1.1.cmml" xref="algx1.l10.1.m1.1.1"><matrixrow id="algx1.l10.1.m1.1.1a.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1b.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1c.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l10.1.m1.1.1d.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1e.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1f.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l10.1.m1.1.1g.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1h.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1i.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l10.1.m1.1.1j.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1k.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1l.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l10.1.m1.1.1m.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1n.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1o.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l10.1.m1.1.1p.cmml" xref="algx1.l10.1.m1.1.1"><cerror id="algx1.l10.1.m1.1.1q.cmml" xref="algx1.l10.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.1r.cmml" xref="algx1.l10.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow></matrix></apply><matrix id="algx1.l10.1.m1.1.2.3.cmml" xref="algx1.l10.1.m1.1.2.3"><matrixrow id="algx1.l10.1.m1.1.2.3a.cmml" xref="algx1.l10.1.m1.1.2.3"><cerror id="algx1.l10.1.m1.1.2.3b.cmml" xref="algx1.l10.1.m1.1.2.3"><csymbol cd="ambiguous" id="algx1.l10.1.m1.1.2.3c.cmml" xref="algx1.l10.1.m1.1.2.3">fragments</csymbol><ci id="algx1.l10.1.m1.1.2.3.1.1.1a.cmml" xref="algx1.l10.1.m1.1.2.3.1.1.1"><mtext class="ltx_mathvariant_bold" mathcolor="#2267EB" id="algx1.l10.1.m1.1.2.3.1.1.1.cmml" xref="algx1.l10.1.m1.1.2.3.1.1.1">Inner Optimization</mtext></ci><ci id="algx1.l10.1.m1.1.2.3.1.1.2a.cmml" xref="algx1.l10.1.m1.1.2.3.1.1.2"><mtext mathcolor="#2267EB" id="algx1.l10.1.m1.1.2.3.1.1.2.cmml" xref="algx1.l10.1.m1.1.2.3.1.1.2">: parallel per path</mtext></ci></cerror></matrixrow></matrix></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l10.1.m1.1c">\left.\begin{array}[]{@{}c@{}}\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\end{array}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\right\}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\begin{tabular}[]{l}{Inner Optimization}: parallel per path\end{tabular}</annotation></semantics></math></span>

</div>
<div id="algx1.l11" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l11.1.1.1" class="ltx_text" style="font-size:80%;">11:</span></span>              <math id="algx1.l11.m1.2" class="ltx_Math" alttext="\mathcal{L}\leftarrow f(x,\theta_{i}^{t})" display="inline"><semantics id="algx1.l11.m1.2a"><mrow id="algx1.l11.m1.2.2" xref="algx1.l11.m1.2.2.cmml"><mi class="ltx_font_mathcaligraphic" id="algx1.l11.m1.2.2.3" xref="algx1.l11.m1.2.2.3.cmml">ℒ</mi><mo stretchy="false" id="algx1.l11.m1.2.2.2" xref="algx1.l11.m1.2.2.2.cmml">←</mo><mrow id="algx1.l11.m1.2.2.1" xref="algx1.l11.m1.2.2.1.cmml"><mi id="algx1.l11.m1.2.2.1.3" xref="algx1.l11.m1.2.2.1.3.cmml">f</mi><mo lspace="0em" rspace="0em" id="algx1.l11.m1.2.2.1.2" xref="algx1.l11.m1.2.2.1.2.cmml">​</mo><mrow id="algx1.l11.m1.2.2.1.1.1" xref="algx1.l11.m1.2.2.1.1.2.cmml"><mo stretchy="false" id="algx1.l11.m1.2.2.1.1.1.2" xref="algx1.l11.m1.2.2.1.1.2.cmml">(</mo><mi id="algx1.l11.m1.1.1" xref="algx1.l11.m1.1.1.cmml">x</mi><mo id="algx1.l11.m1.2.2.1.1.1.3" xref="algx1.l11.m1.2.2.1.1.2.cmml">,</mo><msubsup id="algx1.l11.m1.2.2.1.1.1.1" xref="algx1.l11.m1.2.2.1.1.1.1.cmml"><mi id="algx1.l11.m1.2.2.1.1.1.1.2.2" xref="algx1.l11.m1.2.2.1.1.1.1.2.2.cmml">θ</mi><mi id="algx1.l11.m1.2.2.1.1.1.1.2.3" xref="algx1.l11.m1.2.2.1.1.1.1.2.3.cmml">i</mi><mi id="algx1.l11.m1.2.2.1.1.1.1.3" xref="algx1.l11.m1.2.2.1.1.1.1.3.cmml">t</mi></msubsup><mo stretchy="false" id="algx1.l11.m1.2.2.1.1.1.4" xref="algx1.l11.m1.2.2.1.1.2.cmml">)</mo></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l11.m1.2b"><apply id="algx1.l11.m1.2.2.cmml" xref="algx1.l11.m1.2.2"><ci id="algx1.l11.m1.2.2.2.cmml" xref="algx1.l11.m1.2.2.2">←</ci><ci id="algx1.l11.m1.2.2.3.cmml" xref="algx1.l11.m1.2.2.3">ℒ</ci><apply id="algx1.l11.m1.2.2.1.cmml" xref="algx1.l11.m1.2.2.1"><times id="algx1.l11.m1.2.2.1.2.cmml" xref="algx1.l11.m1.2.2.1.2"></times><ci id="algx1.l11.m1.2.2.1.3.cmml" xref="algx1.l11.m1.2.2.1.3">𝑓</ci><interval closure="open" id="algx1.l11.m1.2.2.1.1.2.cmml" xref="algx1.l11.m1.2.2.1.1.1"><ci id="algx1.l11.m1.1.1.cmml" xref="algx1.l11.m1.1.1">𝑥</ci><apply id="algx1.l11.m1.2.2.1.1.1.1.cmml" xref="algx1.l11.m1.2.2.1.1.1.1"><csymbol cd="ambiguous" id="algx1.l11.m1.2.2.1.1.1.1.1.cmml" xref="algx1.l11.m1.2.2.1.1.1.1">superscript</csymbol><apply id="algx1.l11.m1.2.2.1.1.1.1.2.cmml" xref="algx1.l11.m1.2.2.1.1.1.1"><csymbol cd="ambiguous" id="algx1.l11.m1.2.2.1.1.1.1.2.1.cmml" xref="algx1.l11.m1.2.2.1.1.1.1">subscript</csymbol><ci id="algx1.l11.m1.2.2.1.1.1.1.2.2.cmml" xref="algx1.l11.m1.2.2.1.1.1.1.2.2">𝜃</ci><ci id="algx1.l11.m1.2.2.1.1.1.1.2.3.cmml" xref="algx1.l11.m1.2.2.1.1.1.1.2.3">𝑖</ci></apply><ci id="algx1.l11.m1.2.2.1.1.1.1.3.cmml" xref="algx1.l11.m1.2.2.1.1.1.1.3">𝑡</ci></apply></interval></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l11.m1.2c">\mathcal{L}\leftarrow f(x,\theta_{i}^{t})</annotation></semantics></math>

</div>
<div id="algx1.l12" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l12.1.1.1" class="ltx_text" style="font-size:80%;">12:</span></span>              <math id="algx1.l12.m1.2" class="ltx_Math" alttext="\theta_{i}^{t}\leftarrow\texttt{InnerOpt}(\theta_{i}^{t},\nabla_{\mathcal{L}})" display="inline"><semantics id="algx1.l12.m1.2a"><mrow id="algx1.l12.m1.2.2" xref="algx1.l12.m1.2.2.cmml"><msubsup id="algx1.l12.m1.2.2.4" xref="algx1.l12.m1.2.2.4.cmml"><mi id="algx1.l12.m1.2.2.4.2.2" xref="algx1.l12.m1.2.2.4.2.2.cmml">θ</mi><mi id="algx1.l12.m1.2.2.4.2.3" xref="algx1.l12.m1.2.2.4.2.3.cmml">i</mi><mi id="algx1.l12.m1.2.2.4.3" xref="algx1.l12.m1.2.2.4.3.cmml">t</mi></msubsup><mo stretchy="false" id="algx1.l12.m1.2.2.3" xref="algx1.l12.m1.2.2.3.cmml">←</mo><mrow id="algx1.l12.m1.2.2.2" xref="algx1.l12.m1.2.2.2.cmml"><mtext class="ltx_mathvariant_monospace" id="algx1.l12.m1.2.2.2.4" xref="algx1.l12.m1.2.2.2.4a.cmml">InnerOpt</mtext><mo lspace="0em" rspace="0em" id="algx1.l12.m1.2.2.2.3" xref="algx1.l12.m1.2.2.2.3.cmml">​</mo><mrow id="algx1.l12.m1.2.2.2.2.2" xref="algx1.l12.m1.2.2.2.2.3.cmml"><mo stretchy="false" id="algx1.l12.m1.2.2.2.2.2.3" xref="algx1.l12.m1.2.2.2.2.3.cmml">(</mo><msubsup id="algx1.l12.m1.1.1.1.1.1.1" xref="algx1.l12.m1.1.1.1.1.1.1.cmml"><mi id="algx1.l12.m1.1.1.1.1.1.1.2.2" xref="algx1.l12.m1.1.1.1.1.1.1.2.2.cmml">θ</mi><mi id="algx1.l12.m1.1.1.1.1.1.1.2.3" xref="algx1.l12.m1.1.1.1.1.1.1.2.3.cmml">i</mi><mi id="algx1.l12.m1.1.1.1.1.1.1.3" xref="algx1.l12.m1.1.1.1.1.1.1.3.cmml">t</mi></msubsup><mo id="algx1.l12.m1.2.2.2.2.2.4" xref="algx1.l12.m1.2.2.2.2.3.cmml">,</mo><msub id="algx1.l12.m1.2.2.2.2.2.2" xref="algx1.l12.m1.2.2.2.2.2.2.cmml"><mo id="algx1.l12.m1.2.2.2.2.2.2.2" xref="algx1.l12.m1.2.2.2.2.2.2.2.cmml">∇</mo><mi class="ltx_font_mathcaligraphic" id="algx1.l12.m1.2.2.2.2.2.2.3" xref="algx1.l12.m1.2.2.2.2.2.2.3.cmml">ℒ</mi></msub><mo stretchy="false" id="algx1.l12.m1.2.2.2.2.2.5" xref="algx1.l12.m1.2.2.2.2.3.cmml">)</mo></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l12.m1.2b"><apply id="algx1.l12.m1.2.2.cmml" xref="algx1.l12.m1.2.2"><ci id="algx1.l12.m1.2.2.3.cmml" xref="algx1.l12.m1.2.2.3">←</ci><apply id="algx1.l12.m1.2.2.4.cmml" xref="algx1.l12.m1.2.2.4"><csymbol cd="ambiguous" id="algx1.l12.m1.2.2.4.1.cmml" xref="algx1.l12.m1.2.2.4">superscript</csymbol><apply id="algx1.l12.m1.2.2.4.2.cmml" xref="algx1.l12.m1.2.2.4"><csymbol cd="ambiguous" id="algx1.l12.m1.2.2.4.2.1.cmml" xref="algx1.l12.m1.2.2.4">subscript</csymbol><ci id="algx1.l12.m1.2.2.4.2.2.cmml" xref="algx1.l12.m1.2.2.4.2.2">𝜃</ci><ci id="algx1.l12.m1.2.2.4.2.3.cmml" xref="algx1.l12.m1.2.2.4.2.3">𝑖</ci></apply><ci id="algx1.l12.m1.2.2.4.3.cmml" xref="algx1.l12.m1.2.2.4.3">𝑡</ci></apply><apply id="algx1.l12.m1.2.2.2.cmml" xref="algx1.l12.m1.2.2.2"><times id="algx1.l12.m1.2.2.2.3.cmml" xref="algx1.l12.m1.2.2.2.3"></times><ci id="algx1.l12.m1.2.2.2.4a.cmml" xref="algx1.l12.m1.2.2.2.4"><mtext class="ltx_mathvariant_monospace" id="algx1.l12.m1.2.2.2.4.cmml" xref="algx1.l12.m1.2.2.2.4">InnerOpt</mtext></ci><interval closure="open" id="algx1.l12.m1.2.2.2.2.3.cmml" xref="algx1.l12.m1.2.2.2.2.2"><apply id="algx1.l12.m1.1.1.1.1.1.1.cmml" xref="algx1.l12.m1.1.1.1.1.1.1"><csymbol cd="ambiguous" id="algx1.l12.m1.1.1.1.1.1.1.1.cmml" xref="algx1.l12.m1.1.1.1.1.1.1">superscript</csymbol><apply id="algx1.l12.m1.1.1.1.1.1.1.2.cmml" xref="algx1.l12.m1.1.1.1.1.1.1"><csymbol cd="ambiguous" id="algx1.l12.m1.1.1.1.1.1.1.2.1.cmml" xref="algx1.l12.m1.1.1.1.1.1.1">subscript</csymbol><ci id="algx1.l12.m1.1.1.1.1.1.1.2.2.cmml" xref="algx1.l12.m1.1.1.1.1.1.1.2.2">𝜃</ci><ci id="algx1.l12.m1.1.1.1.1.1.1.2.3.cmml" xref="algx1.l12.m1.1.1.1.1.1.1.2.3">𝑖</ci></apply><ci id="algx1.l12.m1.1.1.1.1.1.1.3.cmml" xref="algx1.l12.m1.1.1.1.1.1.1.3">𝑡</ci></apply><apply id="algx1.l12.m1.2.2.2.2.2.2.cmml" xref="algx1.l12.m1.2.2.2.2.2.2"><csymbol cd="ambiguous" id="algx1.l12.m1.2.2.2.2.2.2.1.cmml" xref="algx1.l12.m1.2.2.2.2.2.2">subscript</csymbol><ci id="algx1.l12.m1.2.2.2.2.2.2.2.cmml" xref="algx1.l12.m1.2.2.2.2.2.2.2">∇</ci><ci id="algx1.l12.m1.2.2.2.2.2.2.3.cmml" xref="algx1.l12.m1.2.2.2.2.2.2.3">ℒ</ci></apply></interval></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l12.m1.2c">\theta_{i}^{t}\leftarrow\texttt{InnerOpt}(\theta_{i}^{t},\nabla_{\mathcal{L}})</annotation></semantics></math>

</div>
<div id="algx1.l13" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l13.1.1.1" class="ltx_text" style="font-size:80%;">13:</span></span>         <span id="algx1.l13.2" class="ltx_text ltx_font_bold">end</span> <span id="algx1.l13.3" class="ltx_text ltx_font_bold">for</span>
</div>
<div id="algx1.l14" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l14.1.1.1" class="ltx_text" style="font-size:80%;">14:</span></span>     <span id="algx1.l14.2" class="ltx_text ltx_font_bold">end</span> <span id="algx1.l14.3" class="ltx_text ltx_font_bold">for</span>
</div>
<div id="algx1.l15" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l15.2.1.1" class="ltx_text" style="font-size:80%;">15:</span></span>     <span id="algx1.l15.3" class="ltx_text ltx_font_bold">for</span> <span id="algx1.l15.1" class="ltx_text ltx_font_typewriter">level <math id="algx1.l15.1.m1.1" class="ltx_Math" alttext="l=1\ldots L" display="inline"><semantics id="algx1.l15.1.m1.1a"><mrow id="algx1.l15.1.m1.1.1" xref="algx1.l15.1.m1.1.1.cmml"><mi id="algx1.l15.1.m1.1.1.2" xref="algx1.l15.1.m1.1.1.2.cmml">l</mi><mo id="algx1.l15.1.m1.1.1.1" xref="algx1.l15.1.m1.1.1.1.cmml">=</mo><mrow id="algx1.l15.1.m1.1.1.3" xref="algx1.l15.1.m1.1.1.3.cmml"><mn id="algx1.l15.1.m1.1.1.3.2" xref="algx1.l15.1.m1.1.1.3.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="algx1.l15.1.m1.1.1.3.1" xref="algx1.l15.1.m1.1.1.3.1.cmml">​</mo><mi mathvariant="normal" id="algx1.l15.1.m1.1.1.3.3" xref="algx1.l15.1.m1.1.1.3.3.cmml">…</mi><mo lspace="0em" rspace="0em" id="algx1.l15.1.m1.1.1.3.1a" xref="algx1.l15.1.m1.1.1.3.1.cmml">​</mo><mi id="algx1.l15.1.m1.1.1.3.4" xref="algx1.l15.1.m1.1.1.3.4.cmml">L</mi></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l15.1.m1.1b"><apply id="algx1.l15.1.m1.1.1.cmml" xref="algx1.l15.1.m1.1.1"><eq id="algx1.l15.1.m1.1.1.1.cmml" xref="algx1.l15.1.m1.1.1.1"></eq><ci id="algx1.l15.1.m1.1.1.2.cmml" xref="algx1.l15.1.m1.1.1.2">𝑙</ci><apply id="algx1.l15.1.m1.1.1.3.cmml" xref="algx1.l15.1.m1.1.1.3"><times id="algx1.l15.1.m1.1.1.3.1.cmml" xref="algx1.l15.1.m1.1.1.3.1"></times><cn type="integer" id="algx1.l15.1.m1.1.1.3.2.cmml" xref="algx1.l15.1.m1.1.1.3.2">1</cn><ci id="algx1.l15.1.m1.1.1.3.3.cmml" xref="algx1.l15.1.m1.1.1.3.3">…</ci><ci id="algx1.l15.1.m1.1.1.3.4.cmml" xref="algx1.l15.1.m1.1.1.3.4">𝐿</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l15.1.m1.1c">l=1\ldots L</annotation></semantics></math></span> <span id="algx1.l15.4" class="ltx_text ltx_font_bold">do</span>

</div>
<div id="algx1.l16" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l16.2.1.1" class="ltx_text" style="font-size:80%;">16:</span></span>         <span id="algx1.l16.3" class="ltx_text ltx_font_bold">for</span> <span id="algx1.l16.1" class="ltx_text ltx_font_typewriter">expert <math id="algx1.l16.1.m1.1" class="ltx_Math" alttext="e=1\ldots K_{l}" display="inline"><semantics id="algx1.l16.1.m1.1a"><mrow id="algx1.l16.1.m1.1.1" xref="algx1.l16.1.m1.1.1.cmml"><mi id="algx1.l16.1.m1.1.1.2" xref="algx1.l16.1.m1.1.1.2.cmml">e</mi><mo id="algx1.l16.1.m1.1.1.1" xref="algx1.l16.1.m1.1.1.1.cmml">=</mo><mrow id="algx1.l16.1.m1.1.1.3" xref="algx1.l16.1.m1.1.1.3.cmml"><mn id="algx1.l16.1.m1.1.1.3.2" xref="algx1.l16.1.m1.1.1.3.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="algx1.l16.1.m1.1.1.3.1" xref="algx1.l16.1.m1.1.1.3.1.cmml">​</mo><mi mathvariant="normal" id="algx1.l16.1.m1.1.1.3.3" xref="algx1.l16.1.m1.1.1.3.3.cmml">…</mi><mo lspace="0em" rspace="0em" id="algx1.l16.1.m1.1.1.3.1a" xref="algx1.l16.1.m1.1.1.3.1.cmml">​</mo><msub id="algx1.l16.1.m1.1.1.3.4" xref="algx1.l16.1.m1.1.1.3.4.cmml"><mi id="algx1.l16.1.m1.1.1.3.4.2" xref="algx1.l16.1.m1.1.1.3.4.2.cmml">K</mi><mi id="algx1.l16.1.m1.1.1.3.4.3" xref="algx1.l16.1.m1.1.1.3.4.3.cmml">l</mi></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l16.1.m1.1b"><apply id="algx1.l16.1.m1.1.1.cmml" xref="algx1.l16.1.m1.1.1"><eq id="algx1.l16.1.m1.1.1.1.cmml" xref="algx1.l16.1.m1.1.1.1"></eq><ci id="algx1.l16.1.m1.1.1.2.cmml" xref="algx1.l16.1.m1.1.1.2">𝑒</ci><apply id="algx1.l16.1.m1.1.1.3.cmml" xref="algx1.l16.1.m1.1.1.3"><times id="algx1.l16.1.m1.1.1.3.1.cmml" xref="algx1.l16.1.m1.1.1.3.1"></times><cn type="integer" id="algx1.l16.1.m1.1.1.3.2.cmml" xref="algx1.l16.1.m1.1.1.3.2">1</cn><ci id="algx1.l16.1.m1.1.1.3.3.cmml" xref="algx1.l16.1.m1.1.1.3.3">…</ci><apply id="algx1.l16.1.m1.1.1.3.4.cmml" xref="algx1.l16.1.m1.1.1.3.4"><csymbol cd="ambiguous" id="algx1.l16.1.m1.1.1.3.4.1.cmml" xref="algx1.l16.1.m1.1.1.3.4">subscript</csymbol><ci id="algx1.l16.1.m1.1.1.3.4.2.cmml" xref="algx1.l16.1.m1.1.1.3.4.2">𝐾</ci><ci id="algx1.l16.1.m1.1.1.3.4.3.cmml" xref="algx1.l16.1.m1.1.1.3.4.3">𝑙</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l16.1.m1.1c">e=1\ldots K_{l}</annotation></semantics></math></span> <span id="algx1.l16.4" class="ltx_text ltx_font_bold">do</span>

</div>
<div id="algx1.l17" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l17.2.1.1" class="ltx_text" style="font-size:80%;">17:</span></span>              <math id="algx1.l17.m1.11" class="ltx_Math" alttext="\Delta(l,e)^{t}\leftarrow\frac{1}{P_{l,e}}\sum_{i=1}^{P_{l,e}}(\theta(l,e)^{t-1}-\theta(l,e)_{i}^{t})" display="inline"><semantics id="algx1.l17.m1.11a"><mrow id="algx1.l17.m1.11.11" xref="algx1.l17.m1.11.11.cmml"><mrow id="algx1.l17.m1.11.11.3" xref="algx1.l17.m1.11.11.3.cmml"><mi mathvariant="normal" id="algx1.l17.m1.11.11.3.2" xref="algx1.l17.m1.11.11.3.2.cmml">Δ</mi><mo lspace="0em" rspace="0em" id="algx1.l17.m1.11.11.3.1" xref="algx1.l17.m1.11.11.3.1.cmml">​</mo><msup id="algx1.l17.m1.11.11.3.3" xref="algx1.l17.m1.11.11.3.3.cmml"><mrow id="algx1.l17.m1.11.11.3.3.2.2" xref="algx1.l17.m1.11.11.3.3.2.1.cmml"><mo stretchy="false" id="algx1.l17.m1.11.11.3.3.2.2.1" xref="algx1.l17.m1.11.11.3.3.2.1.cmml">(</mo><mi id="algx1.l17.m1.5.5" xref="algx1.l17.m1.5.5.cmml">l</mi><mo id="algx1.l17.m1.11.11.3.3.2.2.2" xref="algx1.l17.m1.11.11.3.3.2.1.cmml">,</mo><mi id="algx1.l17.m1.6.6" xref="algx1.l17.m1.6.6.cmml">e</mi><mo stretchy="false" id="algx1.l17.m1.11.11.3.3.2.2.3" xref="algx1.l17.m1.11.11.3.3.2.1.cmml">)</mo></mrow><mi id="algx1.l17.m1.11.11.3.3.3" xref="algx1.l17.m1.11.11.3.3.3.cmml">t</mi></msup></mrow><mo stretchy="false" id="algx1.l17.m1.11.11.2" xref="algx1.l17.m1.11.11.2.cmml">←</mo><mrow id="algx1.l17.m1.11.11.1" xref="algx1.l17.m1.11.11.1.cmml"><mfrac id="algx1.l17.m1.2.2" xref="algx1.l17.m1.2.2.cmml"><mn id="algx1.l17.m1.2.2.4" xref="algx1.l17.m1.2.2.4.cmml">1</mn><msub id="algx1.l17.m1.2.2.2" xref="algx1.l17.m1.2.2.2.cmml"><mi id="algx1.l17.m1.2.2.2.4" xref="algx1.l17.m1.2.2.2.4.cmml">P</mi><mrow id="algx1.l17.m1.2.2.2.2.2.4" xref="algx1.l17.m1.2.2.2.2.2.3.cmml"><mi id="algx1.l17.m1.1.1.1.1.1.1" xref="algx1.l17.m1.1.1.1.1.1.1.cmml">l</mi><mo id="algx1.l17.m1.2.2.2.2.2.4.1" xref="algx1.l17.m1.2.2.2.2.2.3.cmml">,</mo><mi id="algx1.l17.m1.2.2.2.2.2.2" xref="algx1.l17.m1.2.2.2.2.2.2.cmml">e</mi></mrow></msub></mfrac><mo lspace="0em" rspace="0em" id="algx1.l17.m1.11.11.1.2" xref="algx1.l17.m1.11.11.1.2.cmml">​</mo><mrow id="algx1.l17.m1.11.11.1.1" xref="algx1.l17.m1.11.11.1.1.cmml"><msubsup id="algx1.l17.m1.11.11.1.1.2" xref="algx1.l17.m1.11.11.1.1.2.cmml"><mo rspace="0em" id="algx1.l17.m1.11.11.1.1.2.2.2" xref="algx1.l17.m1.11.11.1.1.2.2.2.cmml">∑</mo><mrow id="algx1.l17.m1.11.11.1.1.2.2.3" xref="algx1.l17.m1.11.11.1.1.2.2.3.cmml"><mi id="algx1.l17.m1.11.11.1.1.2.2.3.2" xref="algx1.l17.m1.11.11.1.1.2.2.3.2.cmml">i</mi><mo id="algx1.l17.m1.11.11.1.1.2.2.3.1" xref="algx1.l17.m1.11.11.1.1.2.2.3.1.cmml">=</mo><mn id="algx1.l17.m1.11.11.1.1.2.2.3.3" xref="algx1.l17.m1.11.11.1.1.2.2.3.3.cmml">1</mn></mrow><msub id="algx1.l17.m1.4.4.2" xref="algx1.l17.m1.4.4.2.cmml"><mi id="algx1.l17.m1.4.4.2.4" xref="algx1.l17.m1.4.4.2.4.cmml">P</mi><mrow id="algx1.l17.m1.4.4.2.2.2.4" xref="algx1.l17.m1.4.4.2.2.2.3.cmml"><mi id="algx1.l17.m1.3.3.1.1.1.1" xref="algx1.l17.m1.3.3.1.1.1.1.cmml">l</mi><mo id="algx1.l17.m1.4.4.2.2.2.4.1" xref="algx1.l17.m1.4.4.2.2.2.3.cmml">,</mo><mi id="algx1.l17.m1.4.4.2.2.2.2" xref="algx1.l17.m1.4.4.2.2.2.2.cmml">e</mi></mrow></msub></msubsup><mrow id="algx1.l17.m1.11.11.1.1.1.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.cmml"><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.cmml">(</mo><mrow id="algx1.l17.m1.11.11.1.1.1.1.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.cmml"><mrow id="algx1.l17.m1.11.11.1.1.1.1.1.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.cmml"><mi id="algx1.l17.m1.11.11.1.1.1.1.1.2.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="algx1.l17.m1.11.11.1.1.1.1.1.2.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.1.cmml">​</mo><msup id="algx1.l17.m1.11.11.1.1.1.1.1.2.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.cmml"><mrow id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.1.cmml"><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.2.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.1.cmml">(</mo><mi id="algx1.l17.m1.7.7" xref="algx1.l17.m1.7.7.cmml">l</mi><mo id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.2.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.1.cmml">,</mo><mi id="algx1.l17.m1.8.8" xref="algx1.l17.m1.8.8.cmml">e</mi><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.2.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.1.cmml">)</mo></mrow><mrow id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.cmml"><mi id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.2.cmml">t</mi><mo id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.1.cmml">−</mo><mn id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.3.cmml">1</mn></mrow></msup></mrow><mo id="algx1.l17.m1.11.11.1.1.1.1.1.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.1.cmml">−</mo><mrow id="algx1.l17.m1.11.11.1.1.1.1.1.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.cmml"><mi id="algx1.l17.m1.11.11.1.1.1.1.1.3.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="algx1.l17.m1.11.11.1.1.1.1.1.3.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.1.cmml">​</mo><msubsup id="algx1.l17.m1.11.11.1.1.1.1.1.3.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.cmml"><mrow id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.1.cmml"><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.2.1" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.1.cmml">(</mo><mi id="algx1.l17.m1.9.9" xref="algx1.l17.m1.9.9.cmml">l</mi><mo id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.2.2" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.1.cmml">,</mo><mi id="algx1.l17.m1.10.10" xref="algx1.l17.m1.10.10.cmml">e</mi><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.2.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.1.cmml">)</mo></mrow><mi id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.3.cmml">i</mi><mi id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.3.cmml">t</mi></msubsup></mrow></mrow><mo stretchy="false" id="algx1.l17.m1.11.11.1.1.1.1.3" xref="algx1.l17.m1.11.11.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l17.m1.11b"><apply id="algx1.l17.m1.11.11.cmml" xref="algx1.l17.m1.11.11"><ci id="algx1.l17.m1.11.11.2.cmml" xref="algx1.l17.m1.11.11.2">←</ci><apply id="algx1.l17.m1.11.11.3.cmml" xref="algx1.l17.m1.11.11.3"><times id="algx1.l17.m1.11.11.3.1.cmml" xref="algx1.l17.m1.11.11.3.1"></times><ci id="algx1.l17.m1.11.11.3.2.cmml" xref="algx1.l17.m1.11.11.3.2">Δ</ci><apply id="algx1.l17.m1.11.11.3.3.cmml" xref="algx1.l17.m1.11.11.3.3"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.3.3.1.cmml" xref="algx1.l17.m1.11.11.3.3">superscript</csymbol><interval closure="open" id="algx1.l17.m1.11.11.3.3.2.1.cmml" xref="algx1.l17.m1.11.11.3.3.2.2"><ci id="algx1.l17.m1.5.5.cmml" xref="algx1.l17.m1.5.5">𝑙</ci><ci id="algx1.l17.m1.6.6.cmml" xref="algx1.l17.m1.6.6">𝑒</ci></interval><ci id="algx1.l17.m1.11.11.3.3.3.cmml" xref="algx1.l17.m1.11.11.3.3.3">𝑡</ci></apply></apply><apply id="algx1.l17.m1.11.11.1.cmml" xref="algx1.l17.m1.11.11.1"><times id="algx1.l17.m1.11.11.1.2.cmml" xref="algx1.l17.m1.11.11.1.2"></times><apply id="algx1.l17.m1.2.2.cmml" xref="algx1.l17.m1.2.2"><divide id="algx1.l17.m1.2.2.3.cmml" xref="algx1.l17.m1.2.2"></divide><cn type="integer" id="algx1.l17.m1.2.2.4.cmml" xref="algx1.l17.m1.2.2.4">1</cn><apply id="algx1.l17.m1.2.2.2.cmml" xref="algx1.l17.m1.2.2.2"><csymbol cd="ambiguous" id="algx1.l17.m1.2.2.2.3.cmml" xref="algx1.l17.m1.2.2.2">subscript</csymbol><ci id="algx1.l17.m1.2.2.2.4.cmml" xref="algx1.l17.m1.2.2.2.4">𝑃</ci><list id="algx1.l17.m1.2.2.2.2.2.3.cmml" xref="algx1.l17.m1.2.2.2.2.2.4"><ci id="algx1.l17.m1.1.1.1.1.1.1.cmml" xref="algx1.l17.m1.1.1.1.1.1.1">𝑙</ci><ci id="algx1.l17.m1.2.2.2.2.2.2.cmml" xref="algx1.l17.m1.2.2.2.2.2.2">𝑒</ci></list></apply></apply><apply id="algx1.l17.m1.11.11.1.1.cmml" xref="algx1.l17.m1.11.11.1.1"><apply id="algx1.l17.m1.11.11.1.1.2.cmml" xref="algx1.l17.m1.11.11.1.1.2"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.1.1.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.2">superscript</csymbol><apply id="algx1.l17.m1.11.11.1.1.2.2.cmml" xref="algx1.l17.m1.11.11.1.1.2"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.1.1.2.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.2">subscript</csymbol><sum id="algx1.l17.m1.11.11.1.1.2.2.2.cmml" xref="algx1.l17.m1.11.11.1.1.2.2.2"></sum><apply id="algx1.l17.m1.11.11.1.1.2.2.3.cmml" xref="algx1.l17.m1.11.11.1.1.2.2.3"><eq id="algx1.l17.m1.11.11.1.1.2.2.3.1.cmml" xref="algx1.l17.m1.11.11.1.1.2.2.3.1"></eq><ci id="algx1.l17.m1.11.11.1.1.2.2.3.2.cmml" xref="algx1.l17.m1.11.11.1.1.2.2.3.2">𝑖</ci><cn type="integer" id="algx1.l17.m1.11.11.1.1.2.2.3.3.cmml" xref="algx1.l17.m1.11.11.1.1.2.2.3.3">1</cn></apply></apply><apply id="algx1.l17.m1.4.4.2.cmml" xref="algx1.l17.m1.4.4.2"><csymbol cd="ambiguous" id="algx1.l17.m1.4.4.2.3.cmml" xref="algx1.l17.m1.4.4.2">subscript</csymbol><ci id="algx1.l17.m1.4.4.2.4.cmml" xref="algx1.l17.m1.4.4.2.4">𝑃</ci><list id="algx1.l17.m1.4.4.2.2.2.3.cmml" xref="algx1.l17.m1.4.4.2.2.2.4"><ci id="algx1.l17.m1.3.3.1.1.1.1.cmml" xref="algx1.l17.m1.3.3.1.1.1.1">𝑙</ci><ci id="algx1.l17.m1.4.4.2.2.2.2.cmml" xref="algx1.l17.m1.4.4.2.2.2.2">𝑒</ci></list></apply></apply><apply id="algx1.l17.m1.11.11.1.1.1.1.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1"><minus id="algx1.l17.m1.11.11.1.1.1.1.1.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.1"></minus><apply id="algx1.l17.m1.11.11.1.1.1.1.1.2.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2"><times id="algx1.l17.m1.11.11.1.1.1.1.1.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.1"></times><ci id="algx1.l17.m1.11.11.1.1.1.1.1.2.2.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.2">𝜃</ci><apply id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3">superscript</csymbol><interval closure="open" id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.2.2"><ci id="algx1.l17.m1.7.7.cmml" xref="algx1.l17.m1.7.7">𝑙</ci><ci id="algx1.l17.m1.8.8.cmml" xref="algx1.l17.m1.8.8">𝑒</ci></interval><apply id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3"><minus id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.1"></minus><ci id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.2.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.2">𝑡</ci><cn type="integer" id="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.2.3.3.3">1</cn></apply></apply></apply><apply id="algx1.l17.m1.11.11.1.1.1.1.1.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3"><times id="algx1.l17.m1.11.11.1.1.1.1.1.3.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.1"></times><ci id="algx1.l17.m1.11.11.1.1.1.1.1.3.2.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.2">𝜃</ci><apply id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3">superscript</csymbol><apply id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3"><csymbol cd="ambiguous" id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3">subscript</csymbol><interval closure="open" id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.1.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.2.2"><ci id="algx1.l17.m1.9.9.cmml" xref="algx1.l17.m1.9.9">𝑙</ci><ci id="algx1.l17.m1.10.10.cmml" xref="algx1.l17.m1.10.10">𝑒</ci></interval><ci id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.2.3">𝑖</ci></apply><ci id="algx1.l17.m1.11.11.1.1.1.1.1.3.3.3.cmml" xref="algx1.l17.m1.11.11.1.1.1.1.1.3.3.3">𝑡</ci></apply></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l17.m1.11c">\Delta(l,e)^{t}\leftarrow\frac{1}{P_{l,e}}\sum_{i=1}^{P_{l,e}}(\theta(l,e)^{t-1}-\theta(l,e)_{i}^{t})</annotation></semantics></math> <span id="algx1.l17.1" class="ltx_text ltx_inline-block" style="width:0.0pt;">    <math id="algx1.l17.1.m1.1" class="ltx_Math" alttext="\left.\begin{array}[]{@{}c@{}}\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\end{array}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\right\}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\begin{tabular}[]{l}{Outer Optimization}: parallel per module\end{tabular}" display="inline"><semantics id="algx1.l17.1.m1.1a"><mrow id="algx1.l17.1.m1.1.2" xref="algx1.l17.1.m1.1.2.cmml"><mrow id="algx1.l17.1.m1.1.2.2.2" xref="algx1.l17.1.m1.1.2.2.1.cmml"><mi id="algx1.l17.1.m1.1.2.2.2.1" xref="algx1.l17.1.m1.1.2.2.1.1.cmml"></mi><mtable rowspacing="0pt" id="algx1.l17.1.m1.1.1" xref="algx1.l17.1.m1.1.1.cmml"><mtr id="algx1.l17.1.m1.1.1a" xref="algx1.l17.1.m1.1.1.cmml"><mtd id="algx1.l17.1.m1.1.1b" xref="algx1.l17.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l17.1.m1.1.1c" xref="algx1.l17.1.m1.1.1.cmml"><mtd id="algx1.l17.1.m1.1.1d" xref="algx1.l17.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l17.1.m1.1.1e" xref="algx1.l17.1.m1.1.1.cmml"><mtd id="algx1.l17.1.m1.1.1f" xref="algx1.l17.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l17.1.m1.1.1g" xref="algx1.l17.1.m1.1.1.cmml"><mtd id="algx1.l17.1.m1.1.1h" xref="algx1.l17.1.m1.1.1.cmml"></mtd></mtr><mtr id="algx1.l17.1.m1.1.1i" xref="algx1.l17.1.m1.1.1.cmml"><mtd id="algx1.l17.1.m1.1.1j" xref="algx1.l17.1.m1.1.1.cmml"></mtd></mtr></mtable><mo id="algx1.l17.1.m1.1.2.2.2.2" xref="algx1.l17.1.m1.1.2.2.1.1.cmml">}</mo></mrow><mo lspace="0em" rspace="0em" id="algx1.l17.1.m1.1.2.1" xref="algx1.l17.1.m1.1.2.1.cmml">​</mo><mtable id="algx1.l17.1.m1.1.2.3" xref="algx1.l17.1.m1.1.2.3.cmml"><mtr id="algx1.l17.1.m1.1.2.3a" xref="algx1.l17.1.m1.1.2.3.cmml"><mtd class="ltx_align_left" columnalign="left" id="algx1.l17.1.m1.1.2.3b" xref="algx1.l17.1.m1.1.2.3.cmml"><mtext class="ltx_mathvariant_bold" mathcolor="#2267EB" id="algx1.l17.1.m1.1.2.3.1.1.1" xref="algx1.l17.1.m1.1.2.3.1.1.1a.cmml">Outer Optimization</mtext><mtext mathcolor="#2267EB" id="algx1.l17.1.m1.1.2.3.1.1.2" xref="algx1.l17.1.m1.1.2.3.1.1.2a.cmml">: parallel per module</mtext></mtd></mtr></mtable></mrow><annotation-xml encoding="MathML-Content" id="algx1.l17.1.m1.1b"><apply id="algx1.l17.1.m1.1.2.cmml" xref="algx1.l17.1.m1.1.2"><times id="algx1.l17.1.m1.1.2.1.cmml" xref="algx1.l17.1.m1.1.2.1"></times><apply id="algx1.l17.1.m1.1.2.2.1.cmml" xref="algx1.l17.1.m1.1.2.2.2"><csymbol cd="latexml" id="algx1.l17.1.m1.1.2.2.1.1.cmml" xref="algx1.l17.1.m1.1.2.2.2.1">cases</csymbol><matrix id="algx1.l17.1.m1.1.1.cmml" xref="algx1.l17.1.m1.1.1"><matrixrow id="algx1.l17.1.m1.1.1a.cmml" xref="algx1.l17.1.m1.1.1"><cerror id="algx1.l17.1.m1.1.1b.cmml" xref="algx1.l17.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.1c.cmml" xref="algx1.l17.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l17.1.m1.1.1d.cmml" xref="algx1.l17.1.m1.1.1"><cerror id="algx1.l17.1.m1.1.1e.cmml" xref="algx1.l17.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.1f.cmml" xref="algx1.l17.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l17.1.m1.1.1g.cmml" xref="algx1.l17.1.m1.1.1"><cerror id="algx1.l17.1.m1.1.1h.cmml" xref="algx1.l17.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.1i.cmml" xref="algx1.l17.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l17.1.m1.1.1j.cmml" xref="algx1.l17.1.m1.1.1"><cerror id="algx1.l17.1.m1.1.1k.cmml" xref="algx1.l17.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.1l.cmml" xref="algx1.l17.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow><matrixrow id="algx1.l17.1.m1.1.1m.cmml" xref="algx1.l17.1.m1.1.1"><cerror id="algx1.l17.1.m1.1.1n.cmml" xref="algx1.l17.1.m1.1.1"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.1o.cmml" xref="algx1.l17.1.m1.1.1">missing-subexpression</csymbol></cerror></matrixrow></matrix></apply><matrix id="algx1.l17.1.m1.1.2.3.cmml" xref="algx1.l17.1.m1.1.2.3"><matrixrow id="algx1.l17.1.m1.1.2.3a.cmml" xref="algx1.l17.1.m1.1.2.3"><cerror id="algx1.l17.1.m1.1.2.3b.cmml" xref="algx1.l17.1.m1.1.2.3"><csymbol cd="ambiguous" id="algx1.l17.1.m1.1.2.3c.cmml" xref="algx1.l17.1.m1.1.2.3">fragments</csymbol><ci id="algx1.l17.1.m1.1.2.3.1.1.1a.cmml" xref="algx1.l17.1.m1.1.2.3.1.1.1"><mtext class="ltx_mathvariant_bold" mathcolor="#2267EB" id="algx1.l17.1.m1.1.2.3.1.1.1.cmml" xref="algx1.l17.1.m1.1.2.3.1.1.1">Outer Optimization</mtext></ci><ci id="algx1.l17.1.m1.1.2.3.1.1.2a.cmml" xref="algx1.l17.1.m1.1.2.3.1.1.2"><mtext mathcolor="#2267EB" id="algx1.l17.1.m1.1.2.3.1.1.2.cmml" xref="algx1.l17.1.m1.1.2.3.1.1.2">: parallel per module</mtext></ci></cerror></matrixrow></matrix></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l17.1.m1.1c">\left.\begin{array}[]{@{}c@{}}\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\\
{}\hfil\end{array}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\right\}\color[rgb]{0.1328125,0.40234375,0.921875}\definecolor[named]{pgfstrokecolor}{rgb}{0.1328125,0.40234375,0.921875}\begin{tabular}[]{l}{Outer Optimization}: parallel per module\end{tabular}</annotation></semantics></math></span>

</div>
<div id="algx1.l18" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l18.1.1.1" class="ltx_text" style="font-size:80%;">18:</span></span>              <math id="algx1.l18.m1.8" class="ltx_Math" alttext="\theta(l,e)^{t}\leftarrow\texttt{OuterOpt}(\theta(l,e)^{t-1},\Delta(l,e)^{t})" display="inline"><semantics id="algx1.l18.m1.8a"><mrow id="algx1.l18.m1.8.8" xref="algx1.l18.m1.8.8.cmml"><mrow id="algx1.l18.m1.8.8.4" xref="algx1.l18.m1.8.8.4.cmml"><mi id="algx1.l18.m1.8.8.4.2" xref="algx1.l18.m1.8.8.4.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="algx1.l18.m1.8.8.4.1" xref="algx1.l18.m1.8.8.4.1.cmml">​</mo><msup id="algx1.l18.m1.8.8.4.3" xref="algx1.l18.m1.8.8.4.3.cmml"><mrow id="algx1.l18.m1.8.8.4.3.2.2" xref="algx1.l18.m1.8.8.4.3.2.1.cmml"><mo stretchy="false" id="algx1.l18.m1.8.8.4.3.2.2.1" xref="algx1.l18.m1.8.8.4.3.2.1.cmml">(</mo><mi id="algx1.l18.m1.1.1" xref="algx1.l18.m1.1.1.cmml">l</mi><mo id="algx1.l18.m1.8.8.4.3.2.2.2" xref="algx1.l18.m1.8.8.4.3.2.1.cmml">,</mo><mi id="algx1.l18.m1.2.2" xref="algx1.l18.m1.2.2.cmml">e</mi><mo stretchy="false" id="algx1.l18.m1.8.8.4.3.2.2.3" xref="algx1.l18.m1.8.8.4.3.2.1.cmml">)</mo></mrow><mi id="algx1.l18.m1.8.8.4.3.3" xref="algx1.l18.m1.8.8.4.3.3.cmml">t</mi></msup></mrow><mo stretchy="false" id="algx1.l18.m1.8.8.3" xref="algx1.l18.m1.8.8.3.cmml">←</mo><mrow id="algx1.l18.m1.8.8.2" xref="algx1.l18.m1.8.8.2.cmml"><mtext class="ltx_mathvariant_monospace" id="algx1.l18.m1.8.8.2.4" xref="algx1.l18.m1.8.8.2.4a.cmml">OuterOpt</mtext><mo lspace="0em" rspace="0em" id="algx1.l18.m1.8.8.2.3" xref="algx1.l18.m1.8.8.2.3.cmml">​</mo><mrow id="algx1.l18.m1.8.8.2.2.2" xref="algx1.l18.m1.8.8.2.2.3.cmml"><mo stretchy="false" id="algx1.l18.m1.8.8.2.2.2.3" xref="algx1.l18.m1.8.8.2.2.3.cmml">(</mo><mrow id="algx1.l18.m1.7.7.1.1.1.1" xref="algx1.l18.m1.7.7.1.1.1.1.cmml"><mi id="algx1.l18.m1.7.7.1.1.1.1.2" xref="algx1.l18.m1.7.7.1.1.1.1.2.cmml">θ</mi><mo lspace="0em" rspace="0em" id="algx1.l18.m1.7.7.1.1.1.1.1" xref="algx1.l18.m1.7.7.1.1.1.1.1.cmml">​</mo><msup id="algx1.l18.m1.7.7.1.1.1.1.3" xref="algx1.l18.m1.7.7.1.1.1.1.3.cmml"><mrow id="algx1.l18.m1.7.7.1.1.1.1.3.2.2" xref="algx1.l18.m1.7.7.1.1.1.1.3.2.1.cmml"><mo stretchy="false" id="algx1.l18.m1.7.7.1.1.1.1.3.2.2.1" xref="algx1.l18.m1.7.7.1.1.1.1.3.2.1.cmml">(</mo><mi id="algx1.l18.m1.3.3" xref="algx1.l18.m1.3.3.cmml">l</mi><mo id="algx1.l18.m1.7.7.1.1.1.1.3.2.2.2" xref="algx1.l18.m1.7.7.1.1.1.1.3.2.1.cmml">,</mo><mi id="algx1.l18.m1.4.4" xref="algx1.l18.m1.4.4.cmml">e</mi><mo stretchy="false" id="algx1.l18.m1.7.7.1.1.1.1.3.2.2.3" xref="algx1.l18.m1.7.7.1.1.1.1.3.2.1.cmml">)</mo></mrow><mrow id="algx1.l18.m1.7.7.1.1.1.1.3.3" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.cmml"><mi id="algx1.l18.m1.7.7.1.1.1.1.3.3.2" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.2.cmml">t</mi><mo id="algx1.l18.m1.7.7.1.1.1.1.3.3.1" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.1.cmml">−</mo><mn id="algx1.l18.m1.7.7.1.1.1.1.3.3.3" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.3.cmml">1</mn></mrow></msup></mrow><mo id="algx1.l18.m1.8.8.2.2.2.4" xref="algx1.l18.m1.8.8.2.2.3.cmml">,</mo><mrow id="algx1.l18.m1.8.8.2.2.2.2" xref="algx1.l18.m1.8.8.2.2.2.2.cmml"><mi mathvariant="normal" id="algx1.l18.m1.8.8.2.2.2.2.2" xref="algx1.l18.m1.8.8.2.2.2.2.2.cmml">Δ</mi><mo lspace="0em" rspace="0em" id="algx1.l18.m1.8.8.2.2.2.2.1" xref="algx1.l18.m1.8.8.2.2.2.2.1.cmml">​</mo><msup id="algx1.l18.m1.8.8.2.2.2.2.3" xref="algx1.l18.m1.8.8.2.2.2.2.3.cmml"><mrow id="algx1.l18.m1.8.8.2.2.2.2.3.2.2" xref="algx1.l18.m1.8.8.2.2.2.2.3.2.1.cmml"><mo stretchy="false" id="algx1.l18.m1.8.8.2.2.2.2.3.2.2.1" xref="algx1.l18.m1.8.8.2.2.2.2.3.2.1.cmml">(</mo><mi id="algx1.l18.m1.5.5" xref="algx1.l18.m1.5.5.cmml">l</mi><mo id="algx1.l18.m1.8.8.2.2.2.2.3.2.2.2" xref="algx1.l18.m1.8.8.2.2.2.2.3.2.1.cmml">,</mo><mi id="algx1.l18.m1.6.6" xref="algx1.l18.m1.6.6.cmml">e</mi><mo stretchy="false" id="algx1.l18.m1.8.8.2.2.2.2.3.2.2.3" xref="algx1.l18.m1.8.8.2.2.2.2.3.2.1.cmml">)</mo></mrow><mi id="algx1.l18.m1.8.8.2.2.2.2.3.3" xref="algx1.l18.m1.8.8.2.2.2.2.3.3.cmml">t</mi></msup></mrow><mo stretchy="false" id="algx1.l18.m1.8.8.2.2.2.5" xref="algx1.l18.m1.8.8.2.2.3.cmml">)</mo></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="algx1.l18.m1.8b"><apply id="algx1.l18.m1.8.8.cmml" xref="algx1.l18.m1.8.8"><ci id="algx1.l18.m1.8.8.3.cmml" xref="algx1.l18.m1.8.8.3">←</ci><apply id="algx1.l18.m1.8.8.4.cmml" xref="algx1.l18.m1.8.8.4"><times id="algx1.l18.m1.8.8.4.1.cmml" xref="algx1.l18.m1.8.8.4.1"></times><ci id="algx1.l18.m1.8.8.4.2.cmml" xref="algx1.l18.m1.8.8.4.2">𝜃</ci><apply id="algx1.l18.m1.8.8.4.3.cmml" xref="algx1.l18.m1.8.8.4.3"><csymbol cd="ambiguous" id="algx1.l18.m1.8.8.4.3.1.cmml" xref="algx1.l18.m1.8.8.4.3">superscript</csymbol><interval closure="open" id="algx1.l18.m1.8.8.4.3.2.1.cmml" xref="algx1.l18.m1.8.8.4.3.2.2"><ci id="algx1.l18.m1.1.1.cmml" xref="algx1.l18.m1.1.1">𝑙</ci><ci id="algx1.l18.m1.2.2.cmml" xref="algx1.l18.m1.2.2">𝑒</ci></interval><ci id="algx1.l18.m1.8.8.4.3.3.cmml" xref="algx1.l18.m1.8.8.4.3.3">𝑡</ci></apply></apply><apply id="algx1.l18.m1.8.8.2.cmml" xref="algx1.l18.m1.8.8.2"><times id="algx1.l18.m1.8.8.2.3.cmml" xref="algx1.l18.m1.8.8.2.3"></times><ci id="algx1.l18.m1.8.8.2.4a.cmml" xref="algx1.l18.m1.8.8.2.4"><mtext class="ltx_mathvariant_monospace" id="algx1.l18.m1.8.8.2.4.cmml" xref="algx1.l18.m1.8.8.2.4">OuterOpt</mtext></ci><interval closure="open" id="algx1.l18.m1.8.8.2.2.3.cmml" xref="algx1.l18.m1.8.8.2.2.2"><apply id="algx1.l18.m1.7.7.1.1.1.1.cmml" xref="algx1.l18.m1.7.7.1.1.1.1"><times id="algx1.l18.m1.7.7.1.1.1.1.1.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.1"></times><ci id="algx1.l18.m1.7.7.1.1.1.1.2.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.2">𝜃</ci><apply id="algx1.l18.m1.7.7.1.1.1.1.3.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3"><csymbol cd="ambiguous" id="algx1.l18.m1.7.7.1.1.1.1.3.1.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3">superscript</csymbol><interval closure="open" id="algx1.l18.m1.7.7.1.1.1.1.3.2.1.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3.2.2"><ci id="algx1.l18.m1.3.3.cmml" xref="algx1.l18.m1.3.3">𝑙</ci><ci id="algx1.l18.m1.4.4.cmml" xref="algx1.l18.m1.4.4">𝑒</ci></interval><apply id="algx1.l18.m1.7.7.1.1.1.1.3.3.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3.3"><minus id="algx1.l18.m1.7.7.1.1.1.1.3.3.1.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.1"></minus><ci id="algx1.l18.m1.7.7.1.1.1.1.3.3.2.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.2">𝑡</ci><cn type="integer" id="algx1.l18.m1.7.7.1.1.1.1.3.3.3.cmml" xref="algx1.l18.m1.7.7.1.1.1.1.3.3.3">1</cn></apply></apply></apply><apply id="algx1.l18.m1.8.8.2.2.2.2.cmml" xref="algx1.l18.m1.8.8.2.2.2.2"><times id="algx1.l18.m1.8.8.2.2.2.2.1.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.1"></times><ci id="algx1.l18.m1.8.8.2.2.2.2.2.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.2">Δ</ci><apply id="algx1.l18.m1.8.8.2.2.2.2.3.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.3"><csymbol cd="ambiguous" id="algx1.l18.m1.8.8.2.2.2.2.3.1.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.3">superscript</csymbol><interval closure="open" id="algx1.l18.m1.8.8.2.2.2.2.3.2.1.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.3.2.2"><ci id="algx1.l18.m1.5.5.cmml" xref="algx1.l18.m1.5.5">𝑙</ci><ci id="algx1.l18.m1.6.6.cmml" xref="algx1.l18.m1.6.6">𝑒</ci></interval><ci id="algx1.l18.m1.8.8.2.2.2.2.3.3.cmml" xref="algx1.l18.m1.8.8.2.2.2.2.3.3">𝑡</ci></apply></apply></interval></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="algx1.l18.m1.8c">\theta(l,e)^{t}\leftarrow\texttt{OuterOpt}(\theta(l,e)^{t-1},\Delta(l,e)^{t})</annotation></semantics></math> 
</div>
<div id="algx1.l19" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l19.1.1.1" class="ltx_text" style="font-size:80%;">19:</span></span>         <span id="algx1.l19.2" class="ltx_text ltx_font_bold">end</span> <span id="algx1.l19.3" class="ltx_text ltx_font_bold">for</span>
</div>
<div id="algx1.l20" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l20.1.1.1" class="ltx_text" style="font-size:80%;">20:</span></span>     <span id="algx1.l20.2" class="ltx_text ltx_font_bold">end</span> <span id="algx1.l20.3" class="ltx_text ltx_font_bold">for</span>
</div>
<div id="algx1.l21" class="ltx_listingline">
<span class="ltx_tag ltx_tag_listingline"><span id="algx1.l21.1.1.1" class="ltx_text" style="font-size:80%;">21:</span></span><span id="algx1.l21.2" class="ltx_text ltx_font_bold">end</span> <span id="algx1.l21.3" class="ltx_text ltx_font_bold">for</span>
</div>
</div>
</figure>
<section id="S2.SS6.SSS1" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.6.1 </span>Increasing Capacity</h4>

<figure id="S2.F5" class="ltx_figure"><img src="/html/2403.10616/assets/x6.png" id="S2.F5.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="332" height="222" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S2.F5.2.1.1" class="ltx_text" style="font-size:90%;">Figure 5</span>: </span><span id="S2.F5.3.2" class="ltx_text" style="font-size:90%;">DiPaCo with more capacity: In this example, level 3 modules are path specific, i.e., modules at that level are not shared by paths.</span></figcaption>
</figure>
<div id="S2.SS6.SSS1.p1" class="ltx_para">
<p id="S2.SS6.SSS1.p1.1" class="ltx_p">The more paths go through a module, the more opportunities for transfer learning across paths, but also the more constrained learning is and the less capacity the overall mixture has. DiPaCo gives a lot of flexibility in how to set this trade-off. Whether a level is shared by all paths or by a certain subset only does not change the overall framework. The only difference is in the choice of which subset of paths is used when computing the average in line <math id="S2.SS6.SSS1.p1.1.m1.1" class="ltx_Math" alttext="13" display="inline"><semantics id="S2.SS6.SSS1.p1.1.m1.1a"><mn id="S2.SS6.SSS1.p1.1.m1.1.1" xref="S2.SS6.SSS1.p1.1.m1.1.1.cmml">13</mn><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS1.p1.1.m1.1b"><cn type="integer" id="S2.SS6.SSS1.p1.1.m1.1.1.cmml" xref="S2.SS6.SSS1.p1.1.m1.1.1">13</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS1.p1.1.m1.1c">13</annotation></semantics></math> of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>.</p>
</div>
<div id="S2.SS6.SSS1.p2" class="ltx_para">
<p id="S2.SS6.SSS1.p2.1" class="ltx_p">One extreme choice is to allocate path-specific modules, as shown in <a href="#S2.F5" title="Figure 5 ‣ 2.6.1 Increasing Capacity ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 5</span></a>. In this case, the third level has a mixture with as many modules as there are paths. This is a particularly easy way to increase parameter count in DiPaCo.<span id="footnote1" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">1</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">1</sup><span class="ltx_tag ltx_tag_note">1</span>While there is no gradient averaging anymore for these path-specific modules, we still apply the outer optimization of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a> because empirically it improves convergence over the default optimizer.</span></span></span></p>
</div>
<div id="S2.SS6.SSS1.p3" class="ltx_para">
<p id="S2.SS6.SSS1.p3.1" class="ltx_p">Clearly, the proposed structure of DiPaCo is merely a first step. We consider it especially exciting to co-design the optimization methods and the architecture. In the next section, we discuss some of our initial approaches to adapting optimization methods for our setting.</p>
</div>
</section>
<section id="S2.SS6.SSS2" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.6.2 </span>Scaling the modular architecture</h4>

<div id="S2.SS6.SSS2.p1" class="ltx_para">
<p id="S2.SS6.SSS2.p1.3" class="ltx_p">DiPaCo is made of multiple levels, each with several expert modules. <span id="S2.SS6.SSS2.p1.3.1" class="ltx_text ltx_font_italic">e.g</span> a <math id="S2.SS6.SSS2.p1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S2.SS6.SSS2.p1.1.m1.1a"><mrow id="S2.SS6.SSS2.p1.1.m1.1.1" xref="S2.SS6.SSS2.p1.1.m1.1.1.cmml"><mn id="S2.SS6.SSS2.p1.1.m1.1.1.2" xref="S2.SS6.SSS2.p1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S2.SS6.SSS2.p1.1.m1.1.1.1" xref="S2.SS6.SSS2.p1.1.m1.1.1.1.cmml">×</mo><mn id="S2.SS6.SSS2.p1.1.m1.1.1.3" xref="S2.SS6.SSS2.p1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS2.p1.1.m1.1b"><apply id="S2.SS6.SSS2.p1.1.m1.1.1.cmml" xref="S2.SS6.SSS2.p1.1.m1.1.1"><times id="S2.SS6.SSS2.p1.1.m1.1.1.1.cmml" xref="S2.SS6.SSS2.p1.1.m1.1.1.1"></times><cn type="integer" id="S2.SS6.SSS2.p1.1.m1.1.1.2.cmml" xref="S2.SS6.SSS2.p1.1.m1.1.1.2">16</cn><cn type="integer" id="S2.SS6.SSS2.p1.1.m1.1.1.3.cmml" xref="S2.SS6.SSS2.p1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS2.p1.1.m1.1c">16\times 16</annotation></semantics></math> has two levels, which contains 16 modules on each level, producing in total 256 paths. It is easy to scale the architecture size by increasing the number of levels and modules per levels, <span id="S2.SS6.SSS2.p1.3.2" class="ltx_text ltx_font_italic">e.g</span> a <math id="S2.SS6.SSS2.p1.2.m2.1" class="ltx_Math" alttext="32\times 32\times 32" display="inline"><semantics id="S2.SS6.SSS2.p1.2.m2.1a"><mrow id="S2.SS6.SSS2.p1.2.m2.1.1" xref="S2.SS6.SSS2.p1.2.m2.1.1.cmml"><mn id="S2.SS6.SSS2.p1.2.m2.1.1.2" xref="S2.SS6.SSS2.p1.2.m2.1.1.2.cmml">32</mn><mo lspace="0.222em" rspace="0.222em" id="S2.SS6.SSS2.p1.2.m2.1.1.1" xref="S2.SS6.SSS2.p1.2.m2.1.1.1.cmml">×</mo><mn id="S2.SS6.SSS2.p1.2.m2.1.1.3" xref="S2.SS6.SSS2.p1.2.m2.1.1.3.cmml">32</mn><mo lspace="0.222em" rspace="0.222em" id="S2.SS6.SSS2.p1.2.m2.1.1.1a" xref="S2.SS6.SSS2.p1.2.m2.1.1.1.cmml">×</mo><mn id="S2.SS6.SSS2.p1.2.m2.1.1.4" xref="S2.SS6.SSS2.p1.2.m2.1.1.4.cmml">32</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS2.p1.2.m2.1b"><apply id="S2.SS6.SSS2.p1.2.m2.1.1.cmml" xref="S2.SS6.SSS2.p1.2.m2.1.1"><times id="S2.SS6.SSS2.p1.2.m2.1.1.1.cmml" xref="S2.SS6.SSS2.p1.2.m2.1.1.1"></times><cn type="integer" id="S2.SS6.SSS2.p1.2.m2.1.1.2.cmml" xref="S2.SS6.SSS2.p1.2.m2.1.1.2">32</cn><cn type="integer" id="S2.SS6.SSS2.p1.2.m2.1.1.3.cmml" xref="S2.SS6.SSS2.p1.2.m2.1.1.3">32</cn><cn type="integer" id="S2.SS6.SSS2.p1.2.m2.1.1.4.cmml" xref="S2.SS6.SSS2.p1.2.m2.1.1.4">32</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS2.p1.2.m2.1c">32\times 32\times 32</annotation></semantics></math> has <math id="S2.SS6.SSS2.p1.3.m3.2" class="ltx_Math" alttext="32{,}768" display="inline"><semantics id="S2.SS6.SSS2.p1.3.m3.2a"><mrow id="S2.SS6.SSS2.p1.3.m3.2.3.2" xref="S2.SS6.SSS2.p1.3.m3.2.3.1.cmml"><mn id="S2.SS6.SSS2.p1.3.m3.1.1" xref="S2.SS6.SSS2.p1.3.m3.1.1.cmml">32</mn><mo id="S2.SS6.SSS2.p1.3.m3.2.3.2.1" xref="S2.SS6.SSS2.p1.3.m3.2.3.1.cmml">,</mo><mn id="S2.SS6.SSS2.p1.3.m3.2.2" xref="S2.SS6.SSS2.p1.3.m3.2.2.cmml">768</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS2.p1.3.m3.2b"><list id="S2.SS6.SSS2.p1.3.m3.2.3.1.cmml" xref="S2.SS6.SSS2.p1.3.m3.2.3.2"><cn type="integer" id="S2.SS6.SSS2.p1.3.m3.1.1.cmml" xref="S2.SS6.SSS2.p1.3.m3.1.1">32</cn><cn type="integer" id="S2.SS6.SSS2.p1.3.m3.2.2.cmml" xref="S2.SS6.SSS2.p1.3.m3.2.2">768</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS2.p1.3.m3.2c">32{,}768</annotation></semantics></math> paths. It’s difficult to find enough devices for that number of paths, therefore only a subset of the paths can be trained at any moment in time. Potentially, at each start of inner optimization, we could sample a different subset of the paths. Doing so would allow us to scale DiPaCo to arbitrary large size.</p>
</div>
</section>
<section id="S2.SS6.SSS3" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">2.6.3 </span>Flat Mixture of Experts</h4>

<div id="S2.SS6.SSS3.p1" class="ltx_para">
<p id="S2.SS6.SSS3.p1.2" class="ltx_p">The extreme form of capacity increase as in <a href="#S2.SS6.SSS1" title="2.6.1 Increasing Capacity ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.6.1</span></a> would be to have each path be a completely independent network. We will refer to this model as a <span id="S2.SS6.SSS3.p1.2.1" class="ltx_text ltx_font_italic">flat</span> mixture of experts (flat MoE). In a flat MoE paths do not share <span id="S2.SS6.SSS3.p1.2.2" class="ltx_text ltx_font_italic">any</span> parameter. In the language of <a href="#S2.SS3" title="2.3 Notation ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 2.3</span></a>, there is only one level, and <math id="S2.SS6.SSS3.p1.1.m1.1" class="ltx_Math" alttext="B=B_{1}" display="inline"><semantics id="S2.SS6.SSS3.p1.1.m1.1a"><mrow id="S2.SS6.SSS3.p1.1.m1.1.1" xref="S2.SS6.SSS3.p1.1.m1.1.1.cmml"><mi id="S2.SS6.SSS3.p1.1.m1.1.1.2" xref="S2.SS6.SSS3.p1.1.m1.1.1.2.cmml">B</mi><mo id="S2.SS6.SSS3.p1.1.m1.1.1.1" xref="S2.SS6.SSS3.p1.1.m1.1.1.1.cmml">=</mo><msub id="S2.SS6.SSS3.p1.1.m1.1.1.3" xref="S2.SS6.SSS3.p1.1.m1.1.1.3.cmml"><mi id="S2.SS6.SSS3.p1.1.m1.1.1.3.2" xref="S2.SS6.SSS3.p1.1.m1.1.1.3.2.cmml">B</mi><mn id="S2.SS6.SSS3.p1.1.m1.1.1.3.3" xref="S2.SS6.SSS3.p1.1.m1.1.1.3.3.cmml">1</mn></msub></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS3.p1.1.m1.1b"><apply id="S2.SS6.SSS3.p1.1.m1.1.1.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1"><eq id="S2.SS6.SSS3.p1.1.m1.1.1.1.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.1"></eq><ci id="S2.SS6.SSS3.p1.1.m1.1.1.2.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.2">𝐵</ci><apply id="S2.SS6.SSS3.p1.1.m1.1.1.3.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.3"><csymbol cd="ambiguous" id="S2.SS6.SSS3.p1.1.m1.1.1.3.1.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.3">subscript</csymbol><ci id="S2.SS6.SSS3.p1.1.m1.1.1.3.2.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.3.2">𝐵</ci><cn type="integer" id="S2.SS6.SSS3.p1.1.m1.1.1.3.3.cmml" xref="S2.SS6.SSS3.p1.1.m1.1.1.3.3">1</cn></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS3.p1.1.m1.1c">B=B_{1}</annotation></semantics></math> is the entire network, and <math id="S2.SS6.SSS3.p1.2.m2.1" class="ltx_Math" alttext="K=K_{1}" display="inline"><semantics id="S2.SS6.SSS3.p1.2.m2.1a"><mrow id="S2.SS6.SSS3.p1.2.m2.1.1" xref="S2.SS6.SSS3.p1.2.m2.1.1.cmml"><mi id="S2.SS6.SSS3.p1.2.m2.1.1.2" xref="S2.SS6.SSS3.p1.2.m2.1.1.2.cmml">K</mi><mo id="S2.SS6.SSS3.p1.2.m2.1.1.1" xref="S2.SS6.SSS3.p1.2.m2.1.1.1.cmml">=</mo><msub id="S2.SS6.SSS3.p1.2.m2.1.1.3" xref="S2.SS6.SSS3.p1.2.m2.1.1.3.cmml"><mi id="S2.SS6.SSS3.p1.2.m2.1.1.3.2" xref="S2.SS6.SSS3.p1.2.m2.1.1.3.2.cmml">K</mi><mn id="S2.SS6.SSS3.p1.2.m2.1.1.3.3" xref="S2.SS6.SSS3.p1.2.m2.1.1.3.3.cmml">1</mn></msub></mrow><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS3.p1.2.m2.1b"><apply id="S2.SS6.SSS3.p1.2.m2.1.1.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1"><eq id="S2.SS6.SSS3.p1.2.m2.1.1.1.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.1"></eq><ci id="S2.SS6.SSS3.p1.2.m2.1.1.2.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.2">𝐾</ci><apply id="S2.SS6.SSS3.p1.2.m2.1.1.3.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.3"><csymbol cd="ambiguous" id="S2.SS6.SSS3.p1.2.m2.1.1.3.1.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.3">subscript</csymbol><ci id="S2.SS6.SSS3.p1.2.m2.1.1.3.2.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.3.2">𝐾</ci><cn type="integer" id="S2.SS6.SSS3.p1.2.m2.1.1.3.3.cmml" xref="S2.SS6.SSS3.p1.2.m2.1.1.3.3">1</cn></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS3.p1.2.m2.1c">K=K_{1}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS6.SSS3.p2" class="ltx_para">
<p id="S2.SS6.SSS3.p2.1" class="ltx_p">Flat MoE is a strong baseline for the compositional version of DiPaCo. In our experiments, we see that it can outperform architectures with shared modules when the number of shards is small compared to the size of the total amount of data.</p>
</div>
<div id="S2.SS6.SSS3.p3" class="ltx_para">
<p id="S2.SS6.SSS3.p3.1" class="ltx_p">With generative routing using <math id="S2.SS6.SSS3.p3.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS6.SSS3.p3.1.m1.1a"><mi id="S2.SS6.SSS3.p3.1.m1.1.1" xref="S2.SS6.SSS3.p3.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS6.SSS3.p3.1.m1.1b"><ci id="S2.SS6.SSS3.p3.1.m1.1.1.cmml" xref="S2.SS6.SSS3.p3.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS6.SSS3.p3.1.m1.1c">k</annotation></semantics></math>-means, this approach is essentially  <cite class="ltx_cite ltx_citemacro_citep">(Gross et al., <a href="#bib.bib19" title="" class="ltx_ref">2017</a>; Gururangan et al., <a href="#bib.bib20" title="" class="ltx_ref">2023</a>)</cite>; in our experiments below we will use discriminative routing as in the compositional models, as this consistently outperforms generative routing.</p>
</div>
</section>
</section>
<section id="S2.SS7" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.7 </span>Advanced Optimization Techniques</h3>

<section id="S2.SS7.SSS0.Px1" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Outer Gradient Norm Rescaling:</h5>

<div id="S2.SS7.SSS0.Px1.p1" class="ltx_para">
<p id="S2.SS7.SSS0.Px1.p1.1" class="ltx_p">In DiPaCo, each module can belong to a different number of paths. For instance, in a <math id="S2.SS7.SSS0.Px1.p1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S2.SS7.SSS0.Px1.p1.1.m1.1a"><mrow id="S2.SS7.SSS0.Px1.p1.1.m1.1.1" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.cmml"><mn id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.2" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.1" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.1.cmml">×</mo><mn id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.3" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S2.SS7.SSS0.Px1.p1.1.m1.1b"><apply id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.cmml" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1"><times id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.1.cmml" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.1"></times><cn type="integer" id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.2.cmml" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.2">16</cn><cn type="integer" id="S2.SS7.SSS0.Px1.p1.1.m1.1.1.3.cmml" xref="S2.SS7.SSS0.Px1.p1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS7.SSS0.Px1.p1.1.m1.1c">16\times 16</annotation></semantics></math> DiPaCo, 16 paths can contribute to a single module, while there could be a level that has path-specific modules which is not shared by any other.</p>
</div>
<div id="S2.SS7.SSS0.Px1.p2" class="ltx_para">
<p id="S2.SS7.SSS0.Px1.p2.1" class="ltx_p">Consequently, the average outer gradients, <math id="S2.SS7.SSS0.Px1.p2.1.m1.2" class="ltx_Math" alttext="\Delta(l,e)^{t}" display="inline"><semantics id="S2.SS7.SSS0.Px1.p2.1.m1.2a"><mrow id="S2.SS7.SSS0.Px1.p2.1.m1.2.3" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.cmml"><mi mathvariant="normal" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.2" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.2.cmml">Δ</mi><mo lspace="0em" rspace="0em" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.1" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.1.cmml">​</mo><msup id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.cmml"><mrow id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.2" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.1.cmml"><mo stretchy="false" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.2.1" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.1.cmml">(</mo><mi id="S2.SS7.SSS0.Px1.p2.1.m1.1.1" xref="S2.SS7.SSS0.Px1.p2.1.m1.1.1.cmml">l</mi><mo id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.2.2" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.1.cmml">,</mo><mi id="S2.SS7.SSS0.Px1.p2.1.m1.2.2" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.2.cmml">e</mi><mo stretchy="false" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.2.3" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.1.cmml">)</mo></mrow><mi id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.3" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.3.cmml">t</mi></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS7.SSS0.Px1.p2.1.m1.2b"><apply id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3"><times id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.1.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.1"></times><ci id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.2.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.2">Δ</ci><apply id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3"><csymbol cd="ambiguous" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.1.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3">superscript</csymbol><interval closure="open" id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.1.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.2.2"><ci id="S2.SS7.SSS0.Px1.p2.1.m1.1.1.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.1.1">𝑙</ci><ci id="S2.SS7.SSS0.Px1.p2.1.m1.2.2.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.2">𝑒</ci></interval><ci id="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.3.cmml" xref="S2.SS7.SSS0.Px1.p2.1.m1.2.3.3.3">𝑡</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS7.SSS0.Px1.p2.1.m1.2c">\Delta(l,e)^{t}</annotation></semantics></math> in <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>, have significant different norms for different modules. Using the intuition, and also empirical observation, that averaging across a larger number of paths is akin to using a larger batch size, we have rescaled the outer gradient norm by the square root of the number of paths going through a module.</p>
</div>
</section>
<section id="S2.SS7.SSS0.Px2" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Loss reweighing:</h5>

<div id="S2.SS7.SSS0.Px2.p1" class="ltx_para">
<p id="S2.SS7.SSS0.Px2.p1.1" class="ltx_p">In general, shards might have different amounts of data. If training of DiPaCo is distributed by using one worker per shard/path, then according to the uniform average of line <math id="S2.SS7.SSS0.Px2.p1.1.m1.1" class="ltx_Math" alttext="13" display="inline"><semantics id="S2.SS7.SSS0.Px2.p1.1.m1.1a"><mn id="S2.SS7.SSS0.Px2.p1.1.m1.1.1" xref="S2.SS7.SSS0.Px2.p1.1.m1.1.1.cmml">13</mn><annotation-xml encoding="MathML-Content" id="S2.SS7.SSS0.Px2.p1.1.m1.1b"><cn type="integer" id="S2.SS7.SSS0.Px2.p1.1.m1.1.1.cmml" xref="S2.SS7.SSS0.Px2.p1.1.m1.1.1">13</cn></annotation-xml><annotation encoding="application/x-tex" id="S2.SS7.SSS0.Px2.p1.1.m1.1c">13</annotation></semantics></math> in <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a> smaller shards are over-sampled.
To compute an unbiased estimate of the gradient, we therefore weigh the outer gradients proportionally to the shard size:</p>
<table id="S2.E2" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="S2.E2.m1.12" class="ltx_Math" alttext="\Delta^{(t)}_{l,e}\leftarrow\sum\alpha_{l,e}(\theta^{(t-1)}_{l,e}-{\theta_{l,e}^{(t)}})," display="block"><semantics id="S2.E2.m1.12a"><mrow id="S2.E2.m1.12.12.1" xref="S2.E2.m1.12.12.1.1.cmml"><mrow id="S2.E2.m1.12.12.1.1" xref="S2.E2.m1.12.12.1.1.cmml"><msubsup id="S2.E2.m1.12.12.1.1.3" xref="S2.E2.m1.12.12.1.1.3.cmml"><mi mathvariant="normal" id="S2.E2.m1.12.12.1.1.3.2.2" xref="S2.E2.m1.12.12.1.1.3.2.2.cmml">Δ</mi><mrow id="S2.E2.m1.3.3.2.4" xref="S2.E2.m1.3.3.2.3.cmml"><mi id="S2.E2.m1.2.2.1.1" xref="S2.E2.m1.2.2.1.1.cmml">l</mi><mo id="S2.E2.m1.3.3.2.4.1" xref="S2.E2.m1.3.3.2.3.cmml">,</mo><mi id="S2.E2.m1.3.3.2.2" xref="S2.E2.m1.3.3.2.2.cmml">e</mi></mrow><mrow id="S2.E2.m1.1.1.1.3" xref="S2.E2.m1.12.12.1.1.3.cmml"><mo stretchy="false" id="S2.E2.m1.1.1.1.3.1" xref="S2.E2.m1.12.12.1.1.3.cmml">(</mo><mi id="S2.E2.m1.1.1.1.1" xref="S2.E2.m1.1.1.1.1.cmml">t</mi><mo stretchy="false" id="S2.E2.m1.1.1.1.3.2" xref="S2.E2.m1.12.12.1.1.3.cmml">)</mo></mrow></msubsup><mo rspace="0.111em" stretchy="false" id="S2.E2.m1.12.12.1.1.2" xref="S2.E2.m1.12.12.1.1.2.cmml">←</mo><mrow id="S2.E2.m1.12.12.1.1.1" xref="S2.E2.m1.12.12.1.1.1.cmml"><mo movablelimits="false" id="S2.E2.m1.12.12.1.1.1.2" xref="S2.E2.m1.12.12.1.1.1.2.cmml">∑</mo><mrow id="S2.E2.m1.12.12.1.1.1.1" xref="S2.E2.m1.12.12.1.1.1.1.cmml"><msub id="S2.E2.m1.12.12.1.1.1.1.3" xref="S2.E2.m1.12.12.1.1.1.1.3.cmml"><mi id="S2.E2.m1.12.12.1.1.1.1.3.2" xref="S2.E2.m1.12.12.1.1.1.1.3.2.cmml">α</mi><mrow id="S2.E2.m1.5.5.2.4" xref="S2.E2.m1.5.5.2.3.cmml"><mi id="S2.E2.m1.4.4.1.1" xref="S2.E2.m1.4.4.1.1.cmml">l</mi><mo id="S2.E2.m1.5.5.2.4.1" xref="S2.E2.m1.5.5.2.3.cmml">,</mo><mi id="S2.E2.m1.5.5.2.2" xref="S2.E2.m1.5.5.2.2.cmml">e</mi></mrow></msub><mo lspace="0em" rspace="0em" id="S2.E2.m1.12.12.1.1.1.1.2" xref="S2.E2.m1.12.12.1.1.1.1.2.cmml">​</mo><mrow id="S2.E2.m1.12.12.1.1.1.1.1.1" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E2.m1.12.12.1.1.1.1.1.1.2" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E2.m1.12.12.1.1.1.1.1.1.1" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.cmml"><msubsup id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.cmml"><mi id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.2" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.2.cmml">θ</mi><mrow id="S2.E2.m1.8.8.2.4" xref="S2.E2.m1.8.8.2.3.cmml"><mi id="S2.E2.m1.7.7.1.1" xref="S2.E2.m1.7.7.1.1.cmml">l</mi><mo id="S2.E2.m1.8.8.2.4.1" xref="S2.E2.m1.8.8.2.3.cmml">,</mo><mi id="S2.E2.m1.8.8.2.2" xref="S2.E2.m1.8.8.2.2.cmml">e</mi></mrow><mrow id="S2.E2.m1.6.6.1.1" xref="S2.E2.m1.6.6.1.1.1.cmml"><mo stretchy="false" id="S2.E2.m1.6.6.1.1.2" xref="S2.E2.m1.6.6.1.1.1.cmml">(</mo><mrow id="S2.E2.m1.6.6.1.1.1" xref="S2.E2.m1.6.6.1.1.1.cmml"><mi id="S2.E2.m1.6.6.1.1.1.2" xref="S2.E2.m1.6.6.1.1.1.2.cmml">t</mi><mo id="S2.E2.m1.6.6.1.1.1.1" xref="S2.E2.m1.6.6.1.1.1.1.cmml">−</mo><mn id="S2.E2.m1.6.6.1.1.1.3" xref="S2.E2.m1.6.6.1.1.1.3.cmml">1</mn></mrow><mo stretchy="false" id="S2.E2.m1.6.6.1.1.3" xref="S2.E2.m1.6.6.1.1.1.cmml">)</mo></mrow></msubsup><mo id="S2.E2.m1.12.12.1.1.1.1.1.1.1.1" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.1.cmml">−</mo><msubsup id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.cmml"><mi id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.2" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.2.cmml">θ</mi><mrow id="S2.E2.m1.10.10.2.4" xref="S2.E2.m1.10.10.2.3.cmml"><mi id="S2.E2.m1.9.9.1.1" xref="S2.E2.m1.9.9.1.1.cmml">l</mi><mo id="S2.E2.m1.10.10.2.4.1" xref="S2.E2.m1.10.10.2.3.cmml">,</mo><mi id="S2.E2.m1.10.10.2.2" xref="S2.E2.m1.10.10.2.2.cmml">e</mi></mrow><mrow id="S2.E2.m1.11.11.1.3" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.cmml"><mo stretchy="false" id="S2.E2.m1.11.11.1.3.1" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.cmml">(</mo><mi id="S2.E2.m1.11.11.1.1" xref="S2.E2.m1.11.11.1.1.cmml">t</mi><mo stretchy="false" id="S2.E2.m1.11.11.1.3.2" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.cmml">)</mo></mrow></msubsup></mrow><mo stretchy="false" id="S2.E2.m1.12.12.1.1.1.1.1.1.3" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow></mrow><mo id="S2.E2.m1.12.12.1.2" xref="S2.E2.m1.12.12.1.1.cmml">,</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E2.m1.12b"><apply id="S2.E2.m1.12.12.1.1.cmml" xref="S2.E2.m1.12.12.1"><ci id="S2.E2.m1.12.12.1.1.2.cmml" xref="S2.E2.m1.12.12.1.1.2">←</ci><apply id="S2.E2.m1.12.12.1.1.3.cmml" xref="S2.E2.m1.12.12.1.1.3"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.3.1.cmml" xref="S2.E2.m1.12.12.1.1.3">subscript</csymbol><apply id="S2.E2.m1.12.12.1.1.3.2.cmml" xref="S2.E2.m1.12.12.1.1.3"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.3.2.1.cmml" xref="S2.E2.m1.12.12.1.1.3">superscript</csymbol><ci id="S2.E2.m1.12.12.1.1.3.2.2.cmml" xref="S2.E2.m1.12.12.1.1.3.2.2">Δ</ci><ci id="S2.E2.m1.1.1.1.1.cmml" xref="S2.E2.m1.1.1.1.1">𝑡</ci></apply><list id="S2.E2.m1.3.3.2.3.cmml" xref="S2.E2.m1.3.3.2.4"><ci id="S2.E2.m1.2.2.1.1.cmml" xref="S2.E2.m1.2.2.1.1">𝑙</ci><ci id="S2.E2.m1.3.3.2.2.cmml" xref="S2.E2.m1.3.3.2.2">𝑒</ci></list></apply><apply id="S2.E2.m1.12.12.1.1.1.cmml" xref="S2.E2.m1.12.12.1.1.1"><sum id="S2.E2.m1.12.12.1.1.1.2.cmml" xref="S2.E2.m1.12.12.1.1.1.2"></sum><apply id="S2.E2.m1.12.12.1.1.1.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1"><times id="S2.E2.m1.12.12.1.1.1.1.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.2"></times><apply id="S2.E2.m1.12.12.1.1.1.1.3.cmml" xref="S2.E2.m1.12.12.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.1.1.3.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.3">subscript</csymbol><ci id="S2.E2.m1.12.12.1.1.1.1.3.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.3.2">𝛼</ci><list id="S2.E2.m1.5.5.2.3.cmml" xref="S2.E2.m1.5.5.2.4"><ci id="S2.E2.m1.4.4.1.1.cmml" xref="S2.E2.m1.4.4.1.1">𝑙</ci><ci id="S2.E2.m1.5.5.2.2.cmml" xref="S2.E2.m1.5.5.2.2">𝑒</ci></list></apply><apply id="S2.E2.m1.12.12.1.1.1.1.1.1.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1"><minus id="S2.E2.m1.12.12.1.1.1.1.1.1.1.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.1"></minus><apply id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2">subscript</csymbol><apply id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2">superscript</csymbol><ci id="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.2.2.2">𝜃</ci><apply id="S2.E2.m1.6.6.1.1.1.cmml" xref="S2.E2.m1.6.6.1.1"><minus id="S2.E2.m1.6.6.1.1.1.1.cmml" xref="S2.E2.m1.6.6.1.1.1.1"></minus><ci id="S2.E2.m1.6.6.1.1.1.2.cmml" xref="S2.E2.m1.6.6.1.1.1.2">𝑡</ci><cn type="integer" id="S2.E2.m1.6.6.1.1.1.3.cmml" xref="S2.E2.m1.6.6.1.1.1.3">1</cn></apply></apply><list id="S2.E2.m1.8.8.2.3.cmml" xref="S2.E2.m1.8.8.2.4"><ci id="S2.E2.m1.7.7.1.1.cmml" xref="S2.E2.m1.7.7.1.1">𝑙</ci><ci id="S2.E2.m1.8.8.2.2.cmml" xref="S2.E2.m1.8.8.2.2">𝑒</ci></list></apply><apply id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3">superscript</csymbol><apply id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3">subscript</csymbol><ci id="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E2.m1.12.12.1.1.1.1.1.1.1.3.2.2">𝜃</ci><list id="S2.E2.m1.10.10.2.3.cmml" xref="S2.E2.m1.10.10.2.4"><ci id="S2.E2.m1.9.9.1.1.cmml" xref="S2.E2.m1.9.9.1.1">𝑙</ci><ci id="S2.E2.m1.10.10.2.2.cmml" xref="S2.E2.m1.10.10.2.2">𝑒</ci></list></apply><ci id="S2.E2.m1.11.11.1.1.cmml" xref="S2.E2.m1.11.11.1.1">𝑡</ci></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E2.m1.12c">\Delta^{(t)}_{l,e}\leftarrow\sum\alpha_{l,e}(\theta^{(t-1)}_{l,e}-{\theta_{l,e}^{(t)}}),</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
<td rowspan="1" class="ltx_eqn_cell ltx_eqn_eqno ltx_align_middle ltx_align_right"><span class="ltx_tag ltx_tag_equation ltx_align_right">(2)</span></td>
</tr></tbody>
</table>
<p id="S2.SS7.SSS0.Px2.p1.2" class="ltx_p">with:</p>
<table id="S2.E3" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="S2.E3.m1.9" class="ltx_Math" alttext="\alpha_{l,e}=\frac{|\mathcal{D}_{l,e}|}{\sum|\mathcal{D}_{l^{\prime},e^{\prime}}|}." display="block"><semantics id="S2.E3.m1.9a"><mrow id="S2.E3.m1.9.9.1" xref="S2.E3.m1.9.9.1.1.cmml"><mrow id="S2.E3.m1.9.9.1.1" xref="S2.E3.m1.9.9.1.1.cmml"><msub id="S2.E3.m1.9.9.1.1.2" xref="S2.E3.m1.9.9.1.1.2.cmml"><mi id="S2.E3.m1.9.9.1.1.2.2" xref="S2.E3.m1.9.9.1.1.2.2.cmml">α</mi><mrow id="S2.E3.m1.2.2.2.4" xref="S2.E3.m1.2.2.2.3.cmml"><mi id="S2.E3.m1.1.1.1.1" xref="S2.E3.m1.1.1.1.1.cmml">l</mi><mo id="S2.E3.m1.2.2.2.4.1" xref="S2.E3.m1.2.2.2.3.cmml">,</mo><mi id="S2.E3.m1.2.2.2.2" xref="S2.E3.m1.2.2.2.2.cmml">e</mi></mrow></msub><mo id="S2.E3.m1.9.9.1.1.1" xref="S2.E3.m1.9.9.1.1.1.cmml">=</mo><mfrac id="S2.E3.m1.8.8" xref="S2.E3.m1.8.8.cmml"><mrow id="S2.E3.m1.5.5.3.3" xref="S2.E3.m1.5.5.3.4.cmml"><mo stretchy="false" id="S2.E3.m1.5.5.3.3.2" xref="S2.E3.m1.5.5.3.4.1.cmml">|</mo><msub id="S2.E3.m1.5.5.3.3.1" xref="S2.E3.m1.5.5.3.3.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.E3.m1.5.5.3.3.1.2" xref="S2.E3.m1.5.5.3.3.1.2.cmml">𝒟</mi><mrow id="S2.E3.m1.4.4.2.2.2.4" xref="S2.E3.m1.4.4.2.2.2.3.cmml"><mi id="S2.E3.m1.3.3.1.1.1.1" xref="S2.E3.m1.3.3.1.1.1.1.cmml">l</mi><mo id="S2.E3.m1.4.4.2.2.2.4.1" xref="S2.E3.m1.4.4.2.2.2.3.cmml">,</mo><mi id="S2.E3.m1.4.4.2.2.2.2" xref="S2.E3.m1.4.4.2.2.2.2.cmml">e</mi></mrow></msub><mo stretchy="false" id="S2.E3.m1.5.5.3.3.3" xref="S2.E3.m1.5.5.3.4.1.cmml">|</mo></mrow><mrow id="S2.E3.m1.8.8.6" xref="S2.E3.m1.8.8.6.cmml"><mo rspace="0em" id="S2.E3.m1.8.8.6.4" xref="S2.E3.m1.8.8.6.4.cmml">∑</mo><mrow id="S2.E3.m1.8.8.6.3.1" xref="S2.E3.m1.8.8.6.3.2.cmml"><mo stretchy="false" id="S2.E3.m1.8.8.6.3.1.2" xref="S2.E3.m1.8.8.6.3.2.1.cmml">|</mo><msub id="S2.E3.m1.8.8.6.3.1.1" xref="S2.E3.m1.8.8.6.3.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.E3.m1.8.8.6.3.1.1.2" xref="S2.E3.m1.8.8.6.3.1.1.2.cmml">𝒟</mi><mrow id="S2.E3.m1.7.7.5.2.2.2" xref="S2.E3.m1.7.7.5.2.2.3.cmml"><msup id="S2.E3.m1.6.6.4.1.1.1.1" xref="S2.E3.m1.6.6.4.1.1.1.1.cmml"><mi id="S2.E3.m1.6.6.4.1.1.1.1.2" xref="S2.E3.m1.6.6.4.1.1.1.1.2.cmml">l</mi><mo id="S2.E3.m1.6.6.4.1.1.1.1.3" xref="S2.E3.m1.6.6.4.1.1.1.1.3.cmml">′</mo></msup><mo id="S2.E3.m1.7.7.5.2.2.2.3" xref="S2.E3.m1.7.7.5.2.2.3.cmml">,</mo><msup id="S2.E3.m1.7.7.5.2.2.2.2" xref="S2.E3.m1.7.7.5.2.2.2.2.cmml"><mi id="S2.E3.m1.7.7.5.2.2.2.2.2" xref="S2.E3.m1.7.7.5.2.2.2.2.2.cmml">e</mi><mo id="S2.E3.m1.7.7.5.2.2.2.2.3" xref="S2.E3.m1.7.7.5.2.2.2.2.3.cmml">′</mo></msup></mrow></msub><mo stretchy="false" id="S2.E3.m1.8.8.6.3.1.3" xref="S2.E3.m1.8.8.6.3.2.1.cmml">|</mo></mrow></mrow></mfrac></mrow><mo lspace="0em" id="S2.E3.m1.9.9.1.2" xref="S2.E3.m1.9.9.1.1.cmml">.</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E3.m1.9b"><apply id="S2.E3.m1.9.9.1.1.cmml" xref="S2.E3.m1.9.9.1"><eq id="S2.E3.m1.9.9.1.1.1.cmml" xref="S2.E3.m1.9.9.1.1.1"></eq><apply id="S2.E3.m1.9.9.1.1.2.cmml" xref="S2.E3.m1.9.9.1.1.2"><csymbol cd="ambiguous" id="S2.E3.m1.9.9.1.1.2.1.cmml" xref="S2.E3.m1.9.9.1.1.2">subscript</csymbol><ci id="S2.E3.m1.9.9.1.1.2.2.cmml" xref="S2.E3.m1.9.9.1.1.2.2">𝛼</ci><list id="S2.E3.m1.2.2.2.3.cmml" xref="S2.E3.m1.2.2.2.4"><ci id="S2.E3.m1.1.1.1.1.cmml" xref="S2.E3.m1.1.1.1.1">𝑙</ci><ci id="S2.E3.m1.2.2.2.2.cmml" xref="S2.E3.m1.2.2.2.2">𝑒</ci></list></apply><apply id="S2.E3.m1.8.8.cmml" xref="S2.E3.m1.8.8"><divide id="S2.E3.m1.8.8.7.cmml" xref="S2.E3.m1.8.8"></divide><apply id="S2.E3.m1.5.5.3.4.cmml" xref="S2.E3.m1.5.5.3.3"><abs id="S2.E3.m1.5.5.3.4.1.cmml" xref="S2.E3.m1.5.5.3.3.2"></abs><apply id="S2.E3.m1.5.5.3.3.1.cmml" xref="S2.E3.m1.5.5.3.3.1"><csymbol cd="ambiguous" id="S2.E3.m1.5.5.3.3.1.1.cmml" xref="S2.E3.m1.5.5.3.3.1">subscript</csymbol><ci id="S2.E3.m1.5.5.3.3.1.2.cmml" xref="S2.E3.m1.5.5.3.3.1.2">𝒟</ci><list id="S2.E3.m1.4.4.2.2.2.3.cmml" xref="S2.E3.m1.4.4.2.2.2.4"><ci id="S2.E3.m1.3.3.1.1.1.1.cmml" xref="S2.E3.m1.3.3.1.1.1.1">𝑙</ci><ci id="S2.E3.m1.4.4.2.2.2.2.cmml" xref="S2.E3.m1.4.4.2.2.2.2">𝑒</ci></list></apply></apply><apply id="S2.E3.m1.8.8.6.cmml" xref="S2.E3.m1.8.8.6"><sum id="S2.E3.m1.8.8.6.4.cmml" xref="S2.E3.m1.8.8.6.4"></sum><apply id="S2.E3.m1.8.8.6.3.2.cmml" xref="S2.E3.m1.8.8.6.3.1"><abs id="S2.E3.m1.8.8.6.3.2.1.cmml" xref="S2.E3.m1.8.8.6.3.1.2"></abs><apply id="S2.E3.m1.8.8.6.3.1.1.cmml" xref="S2.E3.m1.8.8.6.3.1.1"><csymbol cd="ambiguous" id="S2.E3.m1.8.8.6.3.1.1.1.cmml" xref="S2.E3.m1.8.8.6.3.1.1">subscript</csymbol><ci id="S2.E3.m1.8.8.6.3.1.1.2.cmml" xref="S2.E3.m1.8.8.6.3.1.1.2">𝒟</ci><list id="S2.E3.m1.7.7.5.2.2.3.cmml" xref="S2.E3.m1.7.7.5.2.2.2"><apply id="S2.E3.m1.6.6.4.1.1.1.1.cmml" xref="S2.E3.m1.6.6.4.1.1.1.1"><csymbol cd="ambiguous" id="S2.E3.m1.6.6.4.1.1.1.1.1.cmml" xref="S2.E3.m1.6.6.4.1.1.1.1">superscript</csymbol><ci id="S2.E3.m1.6.6.4.1.1.1.1.2.cmml" xref="S2.E3.m1.6.6.4.1.1.1.1.2">𝑙</ci><ci id="S2.E3.m1.6.6.4.1.1.1.1.3.cmml" xref="S2.E3.m1.6.6.4.1.1.1.1.3">′</ci></apply><apply id="S2.E3.m1.7.7.5.2.2.2.2.cmml" xref="S2.E3.m1.7.7.5.2.2.2.2"><csymbol cd="ambiguous" id="S2.E3.m1.7.7.5.2.2.2.2.1.cmml" xref="S2.E3.m1.7.7.5.2.2.2.2">superscript</csymbol><ci id="S2.E3.m1.7.7.5.2.2.2.2.2.cmml" xref="S2.E3.m1.7.7.5.2.2.2.2.2">𝑒</ci><ci id="S2.E3.m1.7.7.5.2.2.2.2.3.cmml" xref="S2.E3.m1.7.7.5.2.2.2.2.3">′</ci></apply></list></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E3.m1.9c">\alpha_{l,e}=\frac{|\mathcal{D}_{l,e}|}{\sum|\mathcal{D}_{l^{\prime},e^{\prime}}|}.</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
<td rowspan="1" class="ltx_eqn_cell ltx_eqn_eqno ltx_align_middle ltx_align_right"><span class="ltx_tag ltx_tag_equation ltx_align_right">(3)</span></td>
</tr></tbody>
</table>
</div>
</section>
<section id="S2.SS7.SSS0.Px3" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Early Stopping:</h5>

<div id="S2.SS7.SSS0.Px3.p1" class="ltx_para">
<p id="S2.SS7.SSS0.Px3.p1.1" class="ltx_p">By setting aside a small subset of training examples in each shard, we can perform path specific early stopping. In this case, for each path we selected the parameters that yield the lowest loss on the corresponding shard validation set. We found that early stopping improves generalization on small shards, e.g., when the number of shards is large.</p>
</div>
</section>
</section>
</section>
<section id="S3" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">3 </span>Infrastructure</h2>

<figure id="S3.F6" class="ltx_figure"><img src="/html/2403.10616/assets/x7.png" id="S3.F6.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="266" height="149" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S3.F6.8.1.1" class="ltx_text" style="font-size:90%;">Figure 6</span>: </span><span id="S3.F6.9.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Infrastructure<span id="S3.F6.9.2.1" class="ltx_text ltx_font_medium">. Training workers (in <span id="S3.F6.9.2.1.1" class="ltx_text" style="color:#FF8000;">orange</span>) responsible for performing the inner optimization fetch training task from the train task queue (in <span id="S3.F6.9.2.1.2" class="ltx_text" style="color:#800080;">purple</span>). Once the training checkpoints are saved to GFS, their paths and metadata (e.g. training step, path ID, phase ID) are written to a Spanner table (in <span id="S3.F6.9.2.1.3" class="ltx_text" style="color:#0000FF;">blue</span>) allowing easy look up of checkpoint path given metadata. The evaluation workers and the sharded outer optimization executors (in <span id="S3.F6.9.2.1.4" class="ltx_text" style="color:#FF0000;">red</span>) load training checkpoints once they become ready, as signaled by their appearance in the Spanner table. Finally, a monitoring worker (in <span id="S3.F6.9.2.1.5" class="ltx_text" style="color:#00FF00;">green</span>) periodically checks the health of workers and reboots them if necessary.</span></span></figcaption>
</figure>
<div id="S3.p1" class="ltx_para">
<p id="S3.p1.1" class="ltx_p">We now describe the infrastructure we designed to implement the framework detailed earlier. The infrastructure should achieve the following objectives: 1) reliably train the modular architecture via infrequent synchronization; 2) ensure easy scalability for scenarios involving training across multiple pods or data centers; and 3) provide continuous fault tolerance, allowing every component of the system to quickly recover from failed hosts and preemptions without halting training.</p>
</div>
<div id="S3.p2" class="ltx_para">
<p id="S3.p2.1" class="ltx_p">We illustrate the different components in <a href="#S3.F6" title="Figure 6 ‣ 3 Infrastructure ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 6</span></a>. The general training workflow is the following:</p>
<ol id="S3.I1" class="ltx_enumerate">
<li id="S3.I1.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S3.I1.i1.p1" class="ltx_para">
<p id="S3.I1.i1.p1.1" class="ltx_p">At the beginning of each training phase,
a list of training tasks, each including pairs of path and shard ids, is added to the training task queue maintained by the task scheduler server (in <span id="S3.I1.i1.p1.1.1" class="ltx_text" style="color:#800080;">purple</span>).</p>
</div>
</li>
<li id="S3.I1.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S3.I1.i2.p1" class="ltx_para">
<p id="S3.I1.i2.p1.1" class="ltx_p">When a training worker in the worker pool (in <span id="S3.I1.i2.p1.1.1" class="ltx_text" style="color:#FF8000;">orange</span>) becomes available, it fetches the next training task from the train task scheduler and performs inner optimization (L5-9 of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>) on accelerators. Once the inner optimization is finished, the checkpoint is saved to the Google’s distributed file system <cite class="ltx_cite ltx_citemacro_citep">(Ghemawat et al., <a href="#bib.bib16" title="" class="ltx_ref">2003</a>)</cite> and the training worker becomes available for the next training task. The path to the checkpoint, along with the metadata of the checkpoint (e.g., path ID, outer step ID, etc.), is recorded in a database table (shown in <span id="S3.I1.i2.p1.1.2" class="ltx_text" style="color:#0000FF;">blue</span>). This enables other components to query the checkpoint file path for a given path.</p>
</div>
</li>
<li id="S3.I1.i3" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">3.</span> 
<div id="S3.I1.i3.p1" class="ltx_para">
<p id="S3.I1.i3.p1.1" class="ltx_p">When a checkpoint is saved by a training worker (indicated in <span id="S3.I1.i3.p1.1.1" class="ltx_text" style="color:#008080;">teal</span>), evaluation tasks for that checkpoint are added to the eval task queue (highlighted in <span id="S3.I1.i3.p1.1.2" class="ltx_text" style="color:#FFFF00;">yellow</span>). The eval task queue is then consumed by the evaluation workers in the worker pool to carry out evaluation tasks.</p>
</div>
</li>
<li id="S3.I1.i4" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">4.</span> 
<div id="S3.I1.i4.p1" class="ltx_para">
<p id="S3.I1.i4.p1.1" class="ltx_p">The outer optimizer task scheduler (indicated in <span id="S3.I1.i4.p1.1.1" class="ltx_text" style="color:#00FFFF;">light blue</span>) distributes outer optimization tasks to sharded outer optimization executors (highlighted in <span id="S3.I1.i4.p1.1.2" class="ltx_text" style="color:#FF0000;">red</span>), each of which is responsible for the outer optimization of a shard of modules (e.g., a single module or a collection of modules). Each executor loads training checkpoints containing the corresponding modules as soon as they appear in the Spanner database table <cite class="ltx_cite ltx_citemacro_citep">(Corbett et al., <a href="#bib.bib8" title="" class="ltx_ref">2012</a>)</cite> and performs the parameter averaging and outer optimization (L13-14). The checkpoints of the updated module parameters are saved to the distributed file system, and their paths with metadata are written to the database table.</p>
</div>
</li>
<li id="S3.I1.i5" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">5.</span> 
<div id="S3.I1.i5.p1" class="ltx_para">
<p id="S3.I1.i5.p1.1" class="ltx_p">After completion of all the training tasks, the current training phase concludes, and the subsequent training phase commences and training tasks in the new phase can be started as soon as their corresponding modules finish the outer optimization step. See <a href="#S3.SS3" title="3.3 Outer Optimization Efficiency ‣ 3 Infrastructure ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3.3</span></a> for more detail.</p>
</div>
</li>
<li id="S3.I1.i6" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">6.</span> 
<div id="S3.I1.i6.p1" class="ltx_para">
<p id="S3.I1.i6.p1.1" class="ltx_p">Throughout training, a job status monitor (in <span id="S3.I1.i6.p1.1.1" class="ltx_text" style="color:#00FF00;">green</span>) periodically checks the health of workers and task queue servers, and restarts them if they become unresponsive.</p>
</div>
</li>
</ol>
</div>
<div id="S3.p3" class="ltx_para">
<p id="S3.p3.1" class="ltx_p">All the components are designed to be robust to failures and pre-emptions. We will discuss more details of each infrastructure component in the following subsections.</p>
</div>
<section id="S3.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">3.1 </span>Worker Pool</h3>

<div id="S3.SS1.p1" class="ltx_para">
<p id="S3.SS1.p1.1" class="ltx_p">We employ a producer-consumer design pattern to distribute the training tasks. Within each phase, training tasks are added to the task queue, each of which involves training a path for a specific number of steps from a given checkpoint. We maintain a pool of workers, where each worker retrieves training tasks from the task queue iteratively. Each training task is completely independent of other tasks, requiring no synchronization or communication among the workers.</p>
</div>
<div id="S3.SS1.p2" class="ltx_para">
<p id="S3.SS1.p2.1" class="ltx_p">In the event of worker failure or preemption, the fault-tolerant task queue server would return the task from the unavailable worker back to the task queue before reassigning it to another available worker. The task queue server also periodically checkpoints the current task queue, making it possible to recover from server failures or preemptions.</p>
</div>
<div id="S3.SS1.p3" class="ltx_para">
<p id="S3.SS1.p3.1" class="ltx_p">The key advantage of this design is that, thanks to the complete independence among workers, the system can continue making progress even if some workers become unavailable, as long as the worker pool is not empty. It is worth noting that the worker pool can contain <span id="S3.SS1.p3.1.1" class="ltx_text ltx_font_italic">heterogeneous</span> types of devices across <span id="S3.SS1.p3.1.2" class="ltx_text ltx_font_italic">different regions</span>, making it easy to scale up further. Additionally, the worker pool allows for elastic resource utilization by auto-scaling pool size according to resource availability.</p>
</div>
</section>
<section id="S3.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">3.2 </span>Task Queue System</h3>

<div id="S3.SS2.p1" class="ltx_para">
<p id="S3.SS2.p1.1" class="ltx_p">In practice, the aforementioned task queue system is designed to manage and distribute asynchronous tasks among the workers in the worker pool using remote procedure calls. It consists of three components: the task queue scheduler, the task queue server, and the task queue client. Training tasks are published by the task queue scheduler and sent to the task queue server, while each worker instantiates a task queue client to request tasks from the task queue.</p>
</div>
<div id="S3.SS2.p2" class="ltx_para">
<p id="S3.SS2.p2.1" class="ltx_p">In the multi-host training scenario, where multi-host parallelism is achieved by Single Program, Multiple Data (SPMD), initializing a task queue client on each host would result in each host pulling a different task from the task queue. We solve this issue by implementing a task queue client that synchronizes its results across all hosts. This is achieved by creating a real task queue client only on the first JAX host. All other hosts get the response from the first host by performing a blocking all-gather operation.</p>
</div>
<div id="S3.SS2.p3" class="ltx_para">
<p id="S3.SS2.p3.1" class="ltx_p">We also need to synchronize task queue write operation outside of the main Python thread in some places. For instance, evaluation tasks for a given checkpoint can only be enqueued after checkpointing on all hosts has finished. For this use case, we implemented a barrier, which blocks until each program running on their host have made a call with the same unique key.</p>
</div>
</section>
<section id="S3.SS3" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">3.3 </span>Outer Optimization Efficiency</h3>

<figure id="S3.F7" class="ltx_figure"><img src="/html/2403.10616/assets/x8.png" id="S3.F7.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="332" height="115" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S3.F7.3.1.1" class="ltx_text" style="font-size:90%;">Figure 7</span>: </span><span id="S3.F7.4.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Sharded Outer Optimization Executor<span id="S3.F7.4.2.1" class="ltx_text ltx_font_medium">: An example of outer optimization executors sharded by module, which significantly reduces the processing time and memory requirement.</span></span></figcaption>
</figure>
<div id="S3.SS3.p1" class="ltx_para">
<p id="S3.SS3.p1.2" class="ltx_p">Naïve implementation of outer optimization executor introduces significant overhead as we scale up the number of paths across the world. The following optimizations allow us to scale to hundreds of paths with average time per phase for outer update under 2 minutes (lines <math id="S3.SS3.p1.1.m1.1" class="ltx_Math" alttext="11" display="inline"><semantics id="S3.SS3.p1.1.m1.1a"><mn id="S3.SS3.p1.1.m1.1.1" xref="S3.SS3.p1.1.m1.1.1.cmml">11</mn><annotation-xml encoding="MathML-Content" id="S3.SS3.p1.1.m1.1b"><cn type="integer" id="S3.SS3.p1.1.m1.1.1.cmml" xref="S3.SS3.p1.1.m1.1.1">11</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS3.p1.1.m1.1c">11</annotation></semantics></math>-<math id="S3.SS3.p1.2.m2.1" class="ltx_Math" alttext="16" display="inline"><semantics id="S3.SS3.p1.2.m2.1a"><mn id="S3.SS3.p1.2.m2.1.1" xref="S3.SS3.p1.2.m2.1.1.cmml">16</mn><annotation-xml encoding="MathML-Content" id="S3.SS3.p1.2.m2.1b"><cn type="integer" id="S3.SS3.p1.2.m2.1.1.cmml" xref="S3.SS3.p1.2.m2.1.1">16</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS3.p1.2.m2.1c">16</annotation></semantics></math> of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>).</p>
</div>
<section id="S3.SS3.SSS0.Px1" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Online Parameter Gradient Averaging:</h5>

<div id="S3.SS3.SSS0.Px1.p1" class="ltx_para">
<p id="S3.SS3.SSS0.Px1.p1.1" class="ltx_p">As the training tasks within a phase may not finish at the same time, it is possible to reduce the parameter averaging time by loading and accumulating a training checkpoint to the current partial sum as soon as one becomes available. It is also possible to extend to an asynchronous update that doesn’t require to wait for all paths before doing an outer update <cite class="ltx_cite ltx_citemacro_citep">(Liu et al., <a href="#bib.bib27" title="" class="ltx_ref">2024</a>)</cite>.</p>
</div>
</section>
<section id="S3.SS3.SSS0.Px2" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Sharded Outer Optimization Executor:</h5>

<div id="S3.SS3.SSS0.Px2.p1" class="ltx_para">
<p id="S3.SS3.SSS0.Px2.p1.1" class="ltx_p">As our model architecture is modular, averaging meta-gradients of different modules can be performed independently. By distributing the parameter averaging across multiple servers, as illustrated in <a href="#S3.F7" title="Figure 7 ‣ 3.3 Outer Optimization Efficiency ‣ 3 Infrastructure ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 7</span></a>, we reduce the memory requirement of each averaging task, as well as the overhead of checkpoint loading and averaging. It also allows to reduce the training time as a training task can be started as soon as its corresponding modules have finished their outer update. As a consequence, the overall model is never materialized in a single location but always split across several servers.</p>
</div>
</section>
<section id="S3.SS3.SSS0.Px3" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Asynchronous Checkpoints Gathering:</h5>

<div id="S3.SS3.SSS0.Px3.p1" class="ltx_para">
<p id="S3.SS3.SSS0.Px3.p1.1" class="ltx_p">Each outer optimization executor loads as many checkpoints as the number of paths sharing its associated module. It is possible to load checkpoint from anywhere on earth from the distributed file system <cite class="ltx_cite ltx_citemacro_citep">(Ghemawat et al., <a href="#bib.bib16" title="" class="ltx_ref">2003</a>)</cite>. However, if the checkpoint, located on the server where training took place, is at a certain distance from the outer optimization executor, there is going to be a significant delay. In this case we launch in the background an Effingo process <cite class="ltx_cite ltx_citemacro_citep">(Google, <a href="#bib.bib17" title="" class="ltx_ref">2023</a>)</cite> to bring the checkpoint to a closer location.</p>
</div>
</section>
<section id="S3.SS3.SSS0.Px4" class="ltx_paragraph">
<h5 class="ltx_title ltx_title_paragraph">Miscellaneous Improvements:</h5>

<div id="S3.SS3.SSS0.Px4.p1" class="ltx_para">
<p id="S3.SS3.SSS0.Px4.p1.1" class="ltx_p">In addition to the previous speed improvements, we also a) cache the parameters of the outer optimizer <span id="S3.SS3.SSS0.Px4.p1.1.1" class="ltx_text ltx_font_typewriter">OuterOpt</span>, b) reuse the just-in-time compiled optimizer, and c) load multiple paths in parallel via a multi-threading queue.</p>
</div>
</section>
</section>
<section id="S3.SS4" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">3.4 </span>Backup Pool</h3>

<div id="S3.SS4.p1" class="ltx_para">
<p id="S3.SS4.p1.5" class="ltx_p">By default we create as many training workers (see in <span id="S3.SS4.p1.5.1" class="ltx_text" style="color:#FF8000;">orange</span> in <a href="#S3.F6" title="Figure 6 ‣ 3 Infrastructure ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 6</span></a>) <math id="S3.SS4.p1.1.m1.1" class="ltx_Math" alttext="\mathcal{W}_{i}" display="inline"><semantics id="S3.SS4.p1.1.m1.1a"><msub id="S3.SS4.p1.1.m1.1.1" xref="S3.SS4.p1.1.m1.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S3.SS4.p1.1.m1.1.1.2" xref="S3.SS4.p1.1.m1.1.1.2.cmml">𝒲</mi><mi id="S3.SS4.p1.1.m1.1.1.3" xref="S3.SS4.p1.1.m1.1.1.3.cmml">i</mi></msub><annotation-xml encoding="MathML-Content" id="S3.SS4.p1.1.m1.1b"><apply id="S3.SS4.p1.1.m1.1.1.cmml" xref="S3.SS4.p1.1.m1.1.1"><csymbol cd="ambiguous" id="S3.SS4.p1.1.m1.1.1.1.cmml" xref="S3.SS4.p1.1.m1.1.1">subscript</csymbol><ci id="S3.SS4.p1.1.m1.1.1.2.cmml" xref="S3.SS4.p1.1.m1.1.1.2">𝒲</ci><ci id="S3.SS4.p1.1.m1.1.1.3.cmml" xref="S3.SS4.p1.1.m1.1.1.3">𝑖</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p1.1.m1.1c">\mathcal{W}_{i}</annotation></semantics></math> as paths/shards of data, <math id="S3.SS4.p1.2.m2.2" class="ltx_Math" alttext="\forall i\in[1,P]" display="inline"><semantics id="S3.SS4.p1.2.m2.2a"><mrow id="S3.SS4.p1.2.m2.2.3" xref="S3.SS4.p1.2.m2.2.3.cmml"><mrow id="S3.SS4.p1.2.m2.2.3.2" xref="S3.SS4.p1.2.m2.2.3.2.cmml"><mo rspace="0.167em" id="S3.SS4.p1.2.m2.2.3.2.1" xref="S3.SS4.p1.2.m2.2.3.2.1.cmml">∀</mo><mi id="S3.SS4.p1.2.m2.2.3.2.2" xref="S3.SS4.p1.2.m2.2.3.2.2.cmml">i</mi></mrow><mo id="S3.SS4.p1.2.m2.2.3.1" xref="S3.SS4.p1.2.m2.2.3.1.cmml">∈</mo><mrow id="S3.SS4.p1.2.m2.2.3.3.2" xref="S3.SS4.p1.2.m2.2.3.3.1.cmml"><mo stretchy="false" id="S3.SS4.p1.2.m2.2.3.3.2.1" xref="S3.SS4.p1.2.m2.2.3.3.1.cmml">[</mo><mn id="S3.SS4.p1.2.m2.1.1" xref="S3.SS4.p1.2.m2.1.1.cmml">1</mn><mo id="S3.SS4.p1.2.m2.2.3.3.2.2" xref="S3.SS4.p1.2.m2.2.3.3.1.cmml">,</mo><mi id="S3.SS4.p1.2.m2.2.2" xref="S3.SS4.p1.2.m2.2.2.cmml">P</mi><mo stretchy="false" id="S3.SS4.p1.2.m2.2.3.3.2.3" xref="S3.SS4.p1.2.m2.2.3.3.1.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S3.SS4.p1.2.m2.2b"><apply id="S3.SS4.p1.2.m2.2.3.cmml" xref="S3.SS4.p1.2.m2.2.3"><in id="S3.SS4.p1.2.m2.2.3.1.cmml" xref="S3.SS4.p1.2.m2.2.3.1"></in><apply id="S3.SS4.p1.2.m2.2.3.2.cmml" xref="S3.SS4.p1.2.m2.2.3.2"><csymbol cd="latexml" id="S3.SS4.p1.2.m2.2.3.2.1.cmml" xref="S3.SS4.p1.2.m2.2.3.2.1">for-all</csymbol><ci id="S3.SS4.p1.2.m2.2.3.2.2.cmml" xref="S3.SS4.p1.2.m2.2.3.2.2">𝑖</ci></apply><interval closure="closed" id="S3.SS4.p1.2.m2.2.3.3.1.cmml" xref="S3.SS4.p1.2.m2.2.3.3.2"><cn type="integer" id="S3.SS4.p1.2.m2.1.1.cmml" xref="S3.SS4.p1.2.m2.1.1">1</cn><ci id="S3.SS4.p1.2.m2.2.2.cmml" xref="S3.SS4.p1.2.m2.2.2">𝑃</ci></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p1.2.m2.2c">\forall i\in[1,P]</annotation></semantics></math>. However as we scale the number of paths, the hardware requirement might become prohibitive. For instance, with 256 paths and 16 <span id="S3.SS4.p1.5.2" class="ltx_text ltx_font_typewriter">A100</span> GPUs per worker, we would need <math id="S3.SS4.p1.3.m3.1" class="ltx_Math" alttext="4\,096" display="inline"><semantics id="S3.SS4.p1.3.m3.1a"><mn id="S3.SS4.p1.3.m3.1.1" xref="S3.SS4.p1.3.m3.1.1.cmml">4 096</mn><annotation-xml encoding="MathML-Content" id="S3.SS4.p1.3.m3.1b"><cn type="integer" id="S3.SS4.p1.3.m3.1.1.cmml" xref="S3.SS4.p1.3.m3.1.1">4096</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p1.3.m3.1c">4\,096</annotation></semantics></math> GPUs. Therefore in practice, the number of workers is often less than the number of paths, and instead, we do multiple <span id="S3.SS4.p1.5.3" class="ltx_text ltx_font_italic">rounds</span> of training within an outer iteration step until all paths have been trained (L<math id="S3.SS4.p1.4.m4.1" class="ltx_Math" alttext="3" display="inline"><semantics id="S3.SS4.p1.4.m4.1a"><mn id="S3.SS4.p1.4.m4.1.1" xref="S3.SS4.p1.4.m4.1.1.cmml">3</mn><annotation-xml encoding="MathML-Content" id="S3.SS4.p1.4.m4.1b"><cn type="integer" id="S3.SS4.p1.4.m4.1.1.cmml" xref="S3.SS4.p1.4.m4.1.1">3</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p1.4.m4.1c">3</annotation></semantics></math>-<math id="S3.SS4.p1.5.m5.1" class="ltx_Math" alttext="10" display="inline"><semantics id="S3.SS4.p1.5.m5.1a"><mn id="S3.SS4.p1.5.m5.1.1" xref="S3.SS4.p1.5.m5.1.1.cmml">10</mn><annotation-xml encoding="MathML-Content" id="S3.SS4.p1.5.m5.1b"><cn type="integer" id="S3.SS4.p1.5.m5.1.1.cmml" xref="S3.SS4.p1.5.m5.1.1">10</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p1.5.m5.1c">10</annotation></semantics></math> of <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>).</p>
</div>
<div id="S3.SS4.p2" class="ltx_para">
<p id="S3.SS4.p2.2" class="ltx_p">To minimize the number of rounds, and therefore maximize speed, we create a backup pool of training workers. Namely, we spawn as many workers as there are paths across multiple accelerator types using a low-tier priority. As soon as an accelerator is available, it is snatched and used for an inner optimization phase. Although device with low-tier priority are preempted frequently, we can benefit from the backup pool since each training task only takes roughly a couple of minutes (the inner optimization L<math id="S3.SS4.p2.1.m1.1" class="ltx_Math" alttext="5" display="inline"><semantics id="S3.SS4.p2.1.m1.1a"><mn id="S3.SS4.p2.1.m1.1.1" xref="S3.SS4.p2.1.m1.1.1.cmml">5</mn><annotation-xml encoding="MathML-Content" id="S3.SS4.p2.1.m1.1b"><cn type="integer" id="S3.SS4.p2.1.m1.1.1.cmml" xref="S3.SS4.p2.1.m1.1.1">5</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p2.1.m1.1c">5</annotation></semantics></math>-<math id="S3.SS4.p2.2.m2.1" class="ltx_Math" alttext="9" display="inline"><semantics id="S3.SS4.p2.2.m2.1a"><mn id="S3.SS4.p2.2.m2.1.1" xref="S3.SS4.p2.2.m2.1.1.cmml">9</mn><annotation-xml encoding="MathML-Content" id="S3.SS4.p2.2.m2.1b"><cn type="integer" id="S3.SS4.p2.2.m2.1.1.cmml" xref="S3.SS4.p2.2.m2.1.1">9</cn></annotation-xml><annotation encoding="application/x-tex" id="S3.SS4.p2.2.m2.1c">9</annotation></semantics></math> in <a href="#alg1" title="Algorithm 1 ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Algorithm 1</span></a>).</p>
</div>
<div id="S3.SS4.p3" class="ltx_para">
<p id="S3.SS4.p3.1" class="ltx_p">Note that this backup pool is particularly useful when used with the path sampling idea mentioned in <a href="#S2.SS6.SSS2" title="2.6.2 Scaling the modular architecture ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.6.2</span></a>.</p>
</div>
</section>
</section>
<section id="S4" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">4 </span>Experiments</h2>

<div id="S4.p1" class="ltx_para">
<p id="S4.p1.6" class="ltx_p">We consider a language modeling task on the C4 dataset, derived from Common Crawl <cite class="ltx_cite ltx_citemacro_citep">(Raffel et al., <a href="#bib.bib35" title="" class="ltx_ref">2020</a>)</cite>, tokenized with a SentencePiece tokenizer <cite class="ltx_cite ltx_citemacro_citep">(Kudo and Richardson, <a href="#bib.bib24" title="" class="ltx_ref">2018</a>)</cite> with a vocabulary size of <math id="S4.p1.1.m1.2" class="ltx_Math" alttext="32{,}000" display="inline"><semantics id="S4.p1.1.m1.2a"><mrow id="S4.p1.1.m1.2.3.2" xref="S4.p1.1.m1.2.3.1.cmml"><mn id="S4.p1.1.m1.1.1" xref="S4.p1.1.m1.1.1.cmml">32</mn><mo id="S4.p1.1.m1.2.3.2.1" xref="S4.p1.1.m1.2.3.1.cmml">,</mo><mn id="S4.p1.1.m1.2.2" xref="S4.p1.1.m1.2.2.cmml">000</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p1.1.m1.2b"><list id="S4.p1.1.m1.2.3.1.cmml" xref="S4.p1.1.m1.2.3.2"><cn type="integer" id="S4.p1.1.m1.1.1.cmml" xref="S4.p1.1.m1.1.1">32</cn><cn type="integer" id="S4.p1.1.m1.2.2.cmml" xref="S4.p1.1.m1.2.2">000</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.1.m1.2c">32{,}000</annotation></semantics></math>. We report perplexity on the validation set against number of weight update steps used at training time, which is a close proxy for wall-clock time if all computations are done on the same accelerator type. The total number of weight update steps is set to <math id="S4.p1.2.m2.2" class="ltx_Math" alttext="88{,}000" display="inline"><semantics id="S4.p1.2.m2.2a"><mrow id="S4.p1.2.m2.2.3.2" xref="S4.p1.2.m2.2.3.1.cmml"><mn id="S4.p1.2.m2.1.1" xref="S4.p1.2.m2.1.1.cmml">88</mn><mo id="S4.p1.2.m2.2.3.2.1" xref="S4.p1.2.m2.2.3.1.cmml">,</mo><mn id="S4.p1.2.m2.2.2" xref="S4.p1.2.m2.2.2.cmml">000</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p1.2.m2.2b"><list id="S4.p1.2.m2.2.3.1.cmml" xref="S4.p1.2.m2.2.3.2"><cn type="integer" id="S4.p1.2.m2.1.1.cmml" xref="S4.p1.2.m2.1.1">88</cn><cn type="integer" id="S4.p1.2.m2.2.2.cmml" xref="S4.p1.2.m2.2.2">000</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.2.m2.2c">88{,}000</annotation></semantics></math>. All paths on DiPaCo have size <math id="S4.p1.3.m3.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.p1.3.m3.1a"><mn id="S4.p1.3.m3.1.1" xref="S4.p1.3.m3.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.p1.3.m3.1b"><cn type="integer" id="S4.p1.3.m3.1.1.cmml" xref="S4.p1.3.m3.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.3.m3.1c">150</annotation></semantics></math> million parameters, using a transformer with <math id="S4.p1.4.m4.1" class="ltx_Math" alttext="12" display="inline"><semantics id="S4.p1.4.m4.1a"><mn id="S4.p1.4.m4.1.1" xref="S4.p1.4.m4.1.1.cmml">12</mn><annotation-xml encoding="MathML-Content" id="S4.p1.4.m4.1b"><cn type="integer" id="S4.p1.4.m4.1.1.cmml" xref="S4.p1.4.m4.1.1">12</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.4.m4.1c">12</annotation></semantics></math> blocks, <math id="S4.p1.5.m5.1" class="ltx_Math" alttext="896" display="inline"><semantics id="S4.p1.5.m5.1a"><mn id="S4.p1.5.m5.1.1" xref="S4.p1.5.m5.1.1.cmml">896</mn><annotation-xml encoding="MathML-Content" id="S4.p1.5.m5.1b"><cn type="integer" id="S4.p1.5.m5.1.1.cmml" xref="S4.p1.5.m5.1.1">896</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.5.m5.1c">896</annotation></semantics></math> dimensional hidden states and <math id="S4.p1.6.m6.1" class="ltx_Math" alttext="16" display="inline"><semantics id="S4.p1.6.m6.1a"><mn id="S4.p1.6.m6.1.1" xref="S4.p1.6.m6.1.1.cmml">16</mn><annotation-xml encoding="MathML-Content" id="S4.p1.6.m6.1b"><cn type="integer" id="S4.p1.6.m6.1.1.cmml" xref="S4.p1.6.m6.1.1">16</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p1.6.m6.1c">16</annotation></semantics></math> heads.</p>
</div>
<div id="S4.p2" class="ltx_para">
<p id="S4.p2.4" class="ltx_p">We build modular networks of varying number of paths, and varying levels of module reuse. We compare to the performance of a dense transformer language model of size <math id="S4.p2.1.m1.1" class="ltx_Math" alttext="1.3" display="inline"><semantics id="S4.p2.1.m1.1a"><mn id="S4.p2.1.m1.1.1" xref="S4.p2.1.m1.1.1.cmml">1.3</mn><annotation-xml encoding="MathML-Content" id="S4.p2.1.m1.1b"><cn type="float" id="S4.p2.1.m1.1.1.cmml" xref="S4.p2.1.m1.1.1">1.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p2.1.m1.1c">1.3</annotation></semantics></math>B, which has <math id="S4.p2.2.m2.1" class="ltx_Math" alttext="24" display="inline"><semantics id="S4.p2.2.m2.1a"><mn id="S4.p2.2.m2.1.1" xref="S4.p2.2.m2.1.1.cmml">24</mn><annotation-xml encoding="MathML-Content" id="S4.p2.2.m2.1b"><cn type="integer" id="S4.p2.2.m2.1.1.cmml" xref="S4.p2.2.m2.1.1">24</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p2.2.m2.1c">24</annotation></semantics></math> blocks, <math id="S4.p2.3.m3.1" class="ltx_Math" alttext="2048" display="inline"><semantics id="S4.p2.3.m3.1a"><mn id="S4.p2.3.m3.1.1" xref="S4.p2.3.m3.1.1.cmml">2048</mn><annotation-xml encoding="MathML-Content" id="S4.p2.3.m3.1b"><cn type="integer" id="S4.p2.3.m3.1.1.cmml" xref="S4.p2.3.m3.1.1">2048</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p2.3.m3.1c">2048</annotation></semantics></math> dimensional hidden states and with the same number of heads; and a dense model that is the size of one DiPaCo path. These models are also trained for <math id="S4.p2.4.m4.2" class="ltx_Math" alttext="88{,}000" display="inline"><semantics id="S4.p2.4.m4.2a"><mrow id="S4.p2.4.m4.2.3.2" xref="S4.p2.4.m4.2.3.1.cmml"><mn id="S4.p2.4.m4.1.1" xref="S4.p2.4.m4.1.1.cmml">88</mn><mo id="S4.p2.4.m4.2.3.2.1" xref="S4.p2.4.m4.2.3.1.cmml">,</mo><mn id="S4.p2.4.m4.2.2" xref="S4.p2.4.m4.2.2.cmml">000</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p2.4.m4.2b"><list id="S4.p2.4.m4.2.3.1.cmml" xref="S4.p2.4.m4.2.3.2"><cn type="integer" id="S4.p2.4.m4.1.1.cmml" xref="S4.p2.4.m4.1.1">88</cn><cn type="integer" id="S4.p2.4.m4.2.2.cmml" xref="S4.p2.4.m4.2.2">000</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.p2.4.m4.2c">88{,}000</annotation></semantics></math> weight update steps.</p>
</div>
<div id="S4.p3" class="ltx_para">
<p id="S4.p3.2" class="ltx_p">We warn the reader that this comparison is <span id="S4.p3.2.1" class="ltx_text ltx_font_italic">not</span> standard in the literature, as weight updates for DiPaCo see more tokens and use more FLOPs when the number of paths is larger.
Nevertheless, while our modular networks is made of hundreds of paths, they are all trained in parallel, and thus our training wallclock time is 45% less than the <math id="S4.p3.1.m1.1" class="ltx_Math" alttext="1.3" display="inline"><semantics id="S4.p3.1.m1.1a"><mn id="S4.p3.1.m1.1.1" xref="S4.p3.1.m1.1.1.cmml">1.3</mn><annotation-xml encoding="MathML-Content" id="S4.p3.1.m1.1b"><cn type="float" id="S4.p3.1.m1.1.1.cmml" xref="S4.p3.1.m1.1.1">1.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p3.1.m1.1c">1.3</annotation></semantics></math>B dense counterpart, and roughly equivalent to the <math id="S4.p3.2.m2.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.p3.2.m2.1a"><mn id="S4.p3.2.m2.1.1" xref="S4.p3.2.m2.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.p3.2.m2.1b"><cn type="integer" id="S4.p3.2.m2.1.1.cmml" xref="S4.p3.2.m2.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p3.2.m2.1c">150</annotation></semantics></math>M parameter single-path counterpart.</p>
</div>
<div id="S4.p4" class="ltx_para">
<p id="S4.p4.4" class="ltx_p">In our experiments we have searched over relatively few hyper-parameters: mainly learning rate and value of Nesterov momentum. We use a sequence length of <math id="S4.p4.1.m1.2" class="ltx_Math" alttext="1{,}024" display="inline"><semantics id="S4.p4.1.m1.2a"><mrow id="S4.p4.1.m1.2.3.2" xref="S4.p4.1.m1.2.3.1.cmml"><mn id="S4.p4.1.m1.1.1" xref="S4.p4.1.m1.1.1.cmml">1</mn><mo id="S4.p4.1.m1.2.3.2.1" xref="S4.p4.1.m1.2.3.1.cmml">,</mo><mn id="S4.p4.1.m1.2.2" xref="S4.p4.1.m1.2.2.cmml">024</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p4.1.m1.2b"><list id="S4.p4.1.m1.2.3.1.cmml" xref="S4.p4.1.m1.2.3.2"><cn type="integer" id="S4.p4.1.m1.1.1.cmml" xref="S4.p4.1.m1.1.1">1</cn><cn type="integer" id="S4.p4.1.m1.2.2.cmml" xref="S4.p4.1.m1.2.2">024</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.p4.1.m1.2c">1{,}024</annotation></semantics></math> tokens and a batch size of <math id="S4.p4.2.m2.1" class="ltx_Math" alttext="512" display="inline"><semantics id="S4.p4.2.m2.1a"><mn id="S4.p4.2.m2.1.1" xref="S4.p4.2.m2.1.1.cmml">512</mn><annotation-xml encoding="MathML-Content" id="S4.p4.2.m2.1b"><cn type="integer" id="S4.p4.2.m2.1.1.cmml" xref="S4.p4.2.m2.1.1">512</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p4.2.m2.1c">512</annotation></semantics></math>. At eval time, we consider sequences of <math id="S4.p4.3.m3.2" class="ltx_Math" alttext="2{,}048" display="inline"><semantics id="S4.p4.3.m3.2a"><mrow id="S4.p4.3.m3.2.3.2" xref="S4.p4.3.m3.2.3.1.cmml"><mn id="S4.p4.3.m3.1.1" xref="S4.p4.3.m3.1.1.cmml">2</mn><mo id="S4.p4.3.m3.2.3.2.1" xref="S4.p4.3.m3.2.3.1.cmml">,</mo><mn id="S4.p4.3.m3.2.2" xref="S4.p4.3.m3.2.2.cmml">048</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p4.3.m3.2b"><list id="S4.p4.3.m3.2.3.1.cmml" xref="S4.p4.3.m3.2.3.2"><cn type="integer" id="S4.p4.3.m3.1.1.cmml" xref="S4.p4.3.m3.1.1">2</cn><cn type="integer" id="S4.p4.3.m3.2.2.cmml" xref="S4.p4.3.m3.2.2">048</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.p4.3.m3.2c">2{,}048</annotation></semantics></math> tokens. We use cosine learning rate scheduling with peak value of <math id="S4.p4.4.m4.1" class="ltx_Math" alttext="0.0004" display="inline"><semantics id="S4.p4.4.m4.1a"><mn id="S4.p4.4.m4.1.1" xref="S4.p4.4.m4.1.1.cmml">0.0004</mn><annotation-xml encoding="MathML-Content" id="S4.p4.4.m4.1b"><cn type="float" id="S4.p4.4.m4.1.1.cmml" xref="S4.p4.4.m4.1.1">0.0004</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.p4.4.m4.1c">0.0004</annotation></semantics></math>. We list all the hyperparameters in the appendix (<a href="#Sx2.SS1" title="7.1 Additional Details ‣ Supplementary Materials ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 7.1</span></a>).</p>
</div>
<div id="S4.p5" class="ltx_para">
<p id="S4.p5.1" class="ltx_p">All instances of DiPaCo use one phase of discriminative routing (see <a href="#S2.SS4.SSS2" title="2.4.2 Discriminative routing ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.4.2</span></a> and <a href="#Sx2.SS2.SSS1" title="7.2.1 Document level discriminative routing ‣ 7.2 Routing Details ‣ Supplementary Materials ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">7.2.1</span></a>). The <math id="S4.p5.1.m1.1" class="ltx_Math" alttext="16\times 16=256" display="inline"><semantics id="S4.p5.1.m1.1a"><mrow id="S4.p5.1.m1.1.1" xref="S4.p5.1.m1.1.1.cmml"><mrow id="S4.p5.1.m1.1.1.2" xref="S4.p5.1.m1.1.1.2.cmml"><mn id="S4.p5.1.m1.1.1.2.2" xref="S4.p5.1.m1.1.1.2.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.p5.1.m1.1.1.2.1" xref="S4.p5.1.m1.1.1.2.1.cmml">×</mo><mn id="S4.p5.1.m1.1.1.2.3" xref="S4.p5.1.m1.1.1.2.3.cmml">16</mn></mrow><mo id="S4.p5.1.m1.1.1.1" xref="S4.p5.1.m1.1.1.1.cmml">=</mo><mn id="S4.p5.1.m1.1.1.3" xref="S4.p5.1.m1.1.1.3.cmml">256</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.p5.1.m1.1b"><apply id="S4.p5.1.m1.1.1.cmml" xref="S4.p5.1.m1.1.1"><eq id="S4.p5.1.m1.1.1.1.cmml" xref="S4.p5.1.m1.1.1.1"></eq><apply id="S4.p5.1.m1.1.1.2.cmml" xref="S4.p5.1.m1.1.1.2"><times id="S4.p5.1.m1.1.1.2.1.cmml" xref="S4.p5.1.m1.1.1.2.1"></times><cn type="integer" id="S4.p5.1.m1.1.1.2.2.cmml" xref="S4.p5.1.m1.1.1.2.2">16</cn><cn type="integer" id="S4.p5.1.m1.1.1.2.3.cmml" xref="S4.p5.1.m1.1.1.2.3">16</cn></apply><cn type="integer" id="S4.p5.1.m1.1.1.3.cmml" xref="S4.p5.1.m1.1.1.3">256</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.p5.1.m1.1c">16\times 16=256</annotation></semantics></math> path DiPaCo uses top-2 overlapping shards (<a href="#S2.SS4.SSS4" title="2.4.4 Overlapping Shards ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.4.4</span></a>) at training time (this does not make training slower in wallclock or FLOPs, although it does increase the size of the shards).</p>
</div>
<section id="S4.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.1 </span>Comparison with 1B dense model</h3>

<figure id="S4.F8" class="ltx_figure"><img src="/html/2403.10616/assets/x9.png" id="S4.F8.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="664" height="478" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F8.10.3.1" class="ltx_text" style="font-size:90%;">Figure 8</span>: </span><span id="S4.F8.4.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Convergence curves of DiPaCo vs Dense Baseline<span id="S4.F8.4.2.2" class="ltx_text ltx_font_medium">: We first pretrain a <math id="S4.F8.3.1.1.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.F8.3.1.1.m1.1b"><mn id="S4.F8.3.1.1.m1.1.1" xref="S4.F8.3.1.1.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.F8.3.1.1.m1.1c"><cn type="integer" id="S4.F8.3.1.1.m1.1.1.cmml" xref="S4.F8.3.1.1.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.F8.3.1.1.m1.1d">150</annotation></semantics></math>M parameters model for 24k training steps (<span id="S4.F8.4.2.2.1" class="ltx_text" style="color:#6932E6;">purple</span>). We then finetune a 16x16 DiPaCo <math id="S4.F8.4.2.2.m2.1" class="ltx_Math" alttext="P=256" display="inline"><semantics id="S4.F8.4.2.2.m2.1b"><mrow id="S4.F8.4.2.2.m2.1.1" xref="S4.F8.4.2.2.m2.1.1.cmml"><mi id="S4.F8.4.2.2.m2.1.1.2" xref="S4.F8.4.2.2.m2.1.1.2.cmml">P</mi><mo id="S4.F8.4.2.2.m2.1.1.1" xref="S4.F8.4.2.2.m2.1.1.1.cmml">=</mo><mn id="S4.F8.4.2.2.m2.1.1.3" xref="S4.F8.4.2.2.m2.1.1.3.cmml">256</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.F8.4.2.2.m2.1c"><apply id="S4.F8.4.2.2.m2.1.1.cmml" xref="S4.F8.4.2.2.m2.1.1"><eq id="S4.F8.4.2.2.m2.1.1.1.cmml" xref="S4.F8.4.2.2.m2.1.1.1"></eq><ci id="S4.F8.4.2.2.m2.1.1.2.cmml" xref="S4.F8.4.2.2.m2.1.1.2">𝑃</ci><cn type="integer" id="S4.F8.4.2.2.m2.1.1.3.cmml" xref="S4.F8.4.2.2.m2.1.1.3">256</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.F8.4.2.2.m2.1d">P=256</annotation></semantics></math> (<span id="S4.F8.4.2.2.2" class="ltx_text" style="color:#E1144B;">red</span>). The remaining gap to the dense 1.3B parameters baseline (<span id="S4.F8.4.2.2.3" class="ltx_text" style="color:#FF5F19;">orange</span>) is reached by gating more frequently at test time.</span></span></figcaption>
</figure>
<div id="S4.SS1.p1" class="ltx_para">
<p id="S4.SS1.p1.5" class="ltx_p"><a href="#S4.F8" title="Figure 8 ‣ 4.1 Comparison with 1B dense model ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 8</span></a> compares the convergence curve of the 1.3B dense language model with the <math id="S4.SS1.p1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S4.SS1.p1.1.m1.1a"><mrow id="S4.SS1.p1.1.m1.1.1" xref="S4.SS1.p1.1.m1.1.1.cmml"><mn id="S4.SS1.p1.1.m1.1.1.2" xref="S4.SS1.p1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS1.p1.1.m1.1.1.1" xref="S4.SS1.p1.1.m1.1.1.1.cmml">×</mo><mn id="S4.SS1.p1.1.m1.1.1.3" xref="S4.SS1.p1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.1.m1.1b"><apply id="S4.SS1.p1.1.m1.1.1.cmml" xref="S4.SS1.p1.1.m1.1.1"><times id="S4.SS1.p1.1.m1.1.1.1.cmml" xref="S4.SS1.p1.1.m1.1.1.1"></times><cn type="integer" id="S4.SS1.p1.1.m1.1.1.2.cmml" xref="S4.SS1.p1.1.m1.1.1.2">16</cn><cn type="integer" id="S4.SS1.p1.1.m1.1.1.3.cmml" xref="S4.SS1.p1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.1.m1.1c">16\times 16</annotation></semantics></math> DiPaCo.
DiPaCo with <math id="S4.SS1.p1.2.m2.1" class="ltx_Math" alttext="256" display="inline"><semantics id="S4.SS1.p1.2.m2.1a"><mn id="S4.SS1.p1.2.m2.1.1" xref="S4.SS1.p1.2.m2.1.1.cmml">256</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.2.m2.1b"><cn type="integer" id="S4.SS1.p1.2.m2.1.1.cmml" xref="S4.SS1.p1.2.m2.1.1">256</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.2.m2.1c">256</annotation></semantics></math> paths of size <math id="S4.SS1.p1.3.m3.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.SS1.p1.3.m3.1a"><mn id="S4.SS1.p1.3.m3.1.1" xref="S4.SS1.p1.3.m3.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.3.m3.1b"><cn type="integer" id="S4.SS1.p1.3.m3.1.1.cmml" xref="S4.SS1.p1.3.m3.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.3.m3.1c">150</annotation></semantics></math>M partially sharing parameters, and one routing decision per document, is nearly sufficient to match the performance of a dense model of size <math id="S4.SS1.p1.4.m4.1" class="ltx_Math" alttext="1.3" display="inline"><semantics id="S4.SS1.p1.4.m4.1a"><mn id="S4.SS1.p1.4.m4.1.1" xref="S4.SS1.p1.4.m4.1.1.cmml">1.3</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.4.m4.1b"><cn type="float" id="S4.SS1.p1.4.m4.1.1.cmml" xref="S4.SS1.p1.4.m4.1.1">1.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.4.m4.1c">1.3</annotation></semantics></math>B.
The remaining performance gap with the dense <math id="S4.SS1.p1.5.m5.1" class="ltx_Math" alttext="1.3" display="inline"><semantics id="S4.SS1.p1.5.m5.1a"><mn id="S4.SS1.p1.5.m5.1.1" xref="S4.SS1.p1.5.m5.1.1.cmml">1.3</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.5.m5.1b"><cn type="float" id="S4.SS1.p1.5.m5.1.1.cmml" xref="S4.SS1.p1.5.m5.1.1">1.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.5.m5.1c">1.3</annotation></semantics></math>B model can be closed with early stopping and routing every 64 tokens at test time, see <a href="#S4.T3" title="Table 3 ‣ 4.4 Routing frequency during evaluation ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 3</span></a>.</p>
</div>
<div id="S4.SS1.p2" class="ltx_para">
<p id="S4.SS1.p2.4" class="ltx_p">We again warn the reader that the <math id="S4.SS1.p2.1.m1.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S4.SS1.p2.1.m1.1a"><mi id="S4.SS1.p2.1.m1.1.1" xref="S4.SS1.p2.1.m1.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S4.SS1.p2.1.m1.1b"><ci id="S4.SS1.p2.1.m1.1.1.cmml" xref="S4.SS1.p2.1.m1.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p2.1.m1.1c">x</annotation></semantics></math> axis in <a href="#S4.F8" title="Figure 8 ‣ 4.1 Comparison with 1B dense model ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 8</span></a> is not FLOP equivalent for the different models. For the <math id="S4.SS1.p2.2.m2.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.SS1.p2.2.m2.1a"><mn id="S4.SS1.p2.2.m2.1.1" xref="S4.SS1.p2.2.m2.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p2.2.m2.1b"><cn type="integer" id="S4.SS1.p2.2.m2.1.1.cmml" xref="S4.SS1.p2.2.m2.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p2.2.m2.1c">150</annotation></semantics></math>M baseline model and for DiPaCo, the <math id="S4.SS1.p2.3.m3.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S4.SS1.p2.3.m3.1a"><mi id="S4.SS1.p2.3.m3.1.1" xref="S4.SS1.p2.3.m3.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S4.SS1.p2.3.m3.1b"><ci id="S4.SS1.p2.3.m3.1.1.cmml" xref="S4.SS1.p2.3.m3.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p2.3.m3.1c">x</annotation></semantics></math> axis roughly corresponds to equivalent wall-clock (and the <math id="S4.SS1.p2.4.m4.1" class="ltx_Math" alttext="1.3" display="inline"><semantics id="S4.SS1.p2.4.m4.1a"><mn id="S4.SS1.p2.4.m4.1.1" xref="S4.SS1.p2.4.m4.1.1.cmml">1.3</mn><annotation-xml encoding="MathML-Content" id="S4.SS1.p2.4.m4.1b"><cn type="float" id="S4.SS1.p2.4.m4.1.1.cmml" xref="S4.SS1.p2.4.m4.1.1">1.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p2.4.m4.1c">1.3</annotation></semantics></math>B model is slower w.r.t. wall-clock for the same number of updates). On the other hand, DiPaCo uses many more FLOPs and sees more tokens per training step than the dense baselines. However, it is not trivial to parallelize training of the dense models to the scale of DiPaCo, see the comparisons in <a href="#S4.SS3" title="4.3 Parameter sharing between paths ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 4.3</span></a> and in particular <a href="#S4.T1" title="Table 1 ‣ 4.3 Parameter sharing between paths ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 1</span></a>, and in <cite class="ltx_cite ltx_citemacro_citep">(Douillard et al., <a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite>.</p>
</div>
</section>
<section id="S4.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.2 </span>Number of paths and parameters</h3>

<figure id="S4.F9" class="ltx_figure"><img src="/html/2403.10616/assets/x10.png" id="S4.F9.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="664" height="478" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F9.7.2.1" class="ltx_text" style="font-size:90%;">Figure 9</span>: </span><span id="S4.F9.2.1" class="ltx_text ltx_font_bold" style="font-size:90%;">Scaling number of paths in DiPaCo<span id="S4.F9.2.1.1" class="ltx_text ltx_font_medium">: We report validation perplexity of DiPaCo with different numbers of paths and parameters. We increase capacity by varying both the number of paths (from <span id="S4.F9.2.1.1.1" class="ltx_text" style="color:#6932E6;">8</span> to <span id="S4.F9.2.1.1.2" class="ltx_text" style="color:#E1144B;">256</span>) and by using path-specific modules. The size of a path (roughly equivalent to serving cost during deployment) is <math id="S4.F9.2.1.1.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.F9.2.1.1.m1.1b"><mn id="S4.F9.2.1.1.m1.1.1" xref="S4.F9.2.1.1.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.F9.2.1.1.m1.1c"><cn type="integer" id="S4.F9.2.1.1.m1.1.1.cmml" xref="S4.F9.2.1.1.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.F9.2.1.1.m1.1d">150</annotation></semantics></math>M parameters for all models in this figure.</span></span></figcaption>
</figure>
<div id="S4.SS2.p1" class="ltx_para">
<p id="S4.SS2.p1.8" class="ltx_p"><a href="#S4.F9" title="Figure 9 ‣ 4.2 Number of paths and parameters ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Figure 9</span></a> shows that generalization improves as more paths and more overall parameters are added to the mixture. Since the overhead of DiLoCo is minimal, the x axis could be replaced with wall-clock time. Of course, <math id="S4.SS2.p1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S4.SS2.p1.1.m1.1a"><mrow id="S4.SS2.p1.1.m1.1.1" xref="S4.SS2.p1.1.m1.1.1.cmml"><mn id="S4.SS2.p1.1.m1.1.1.2" xref="S4.SS2.p1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS2.p1.1.m1.1.1.1" xref="S4.SS2.p1.1.m1.1.1.1.cmml">×</mo><mn id="S4.SS2.p1.1.m1.1.1.3" xref="S4.SS2.p1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.1.m1.1b"><apply id="S4.SS2.p1.1.m1.1.1.cmml" xref="S4.SS2.p1.1.m1.1.1"><times id="S4.SS2.p1.1.m1.1.1.1.cmml" xref="S4.SS2.p1.1.m1.1.1.1"></times><cn type="integer" id="S4.SS2.p1.1.m1.1.1.2.cmml" xref="S4.SS2.p1.1.m1.1.1.2">16</cn><cn type="integer" id="S4.SS2.p1.1.m1.1.1.3.cmml" xref="S4.SS2.p1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.1.m1.1c">16\times 16</annotation></semantics></math> DiPaCo requires <math id="S4.SS2.p1.2.m2.1" class="ltx_Math" alttext="32" display="inline"><semantics id="S4.SS2.p1.2.m2.1a"><mn id="S4.SS2.p1.2.m2.1.1" xref="S4.SS2.p1.2.m2.1.1.cmml">32</mn><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.2.m2.1b"><cn type="integer" id="S4.SS2.p1.2.m2.1.1.cmml" xref="S4.SS2.p1.2.m2.1.1">32</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.2.m2.1c">32</annotation></semantics></math> times more compute than <math id="S4.SS2.p1.3.m3.1" class="ltx_Math" alttext="2\times 4" display="inline"><semantics id="S4.SS2.p1.3.m3.1a"><mrow id="S4.SS2.p1.3.m3.1.1" xref="S4.SS2.p1.3.m3.1.1.cmml"><mn id="S4.SS2.p1.3.m3.1.1.2" xref="S4.SS2.p1.3.m3.1.1.2.cmml">2</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS2.p1.3.m3.1.1.1" xref="S4.SS2.p1.3.m3.1.1.1.cmml">×</mo><mn id="S4.SS2.p1.3.m3.1.1.3" xref="S4.SS2.p1.3.m3.1.1.3.cmml">4</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.3.m3.1b"><apply id="S4.SS2.p1.3.m3.1.1.cmml" xref="S4.SS2.p1.3.m3.1.1"><times id="S4.SS2.p1.3.m3.1.1.1.cmml" xref="S4.SS2.p1.3.m3.1.1.1"></times><cn type="integer" id="S4.SS2.p1.3.m3.1.1.2.cmml" xref="S4.SS2.p1.3.m3.1.1.2">2</cn><cn type="integer" id="S4.SS2.p1.3.m3.1.1.3.cmml" xref="S4.SS2.p1.3.m3.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.3.m3.1c">2\times 4</annotation></semantics></math> DiPaCo, but these additional paths can be trained in parallel with very little communication. In this experiment the inner optimization consists of <math id="S4.SS2.p1.4.m4.1" class="ltx_Math" alttext="\tau=150" display="inline"><semantics id="S4.SS2.p1.4.m4.1a"><mrow id="S4.SS2.p1.4.m4.1.1" xref="S4.SS2.p1.4.m4.1.1.cmml"><mi id="S4.SS2.p1.4.m4.1.1.2" xref="S4.SS2.p1.4.m4.1.1.2.cmml">τ</mi><mo id="S4.SS2.p1.4.m4.1.1.1" xref="S4.SS2.p1.4.m4.1.1.1.cmml">=</mo><mn id="S4.SS2.p1.4.m4.1.1.3" xref="S4.SS2.p1.4.m4.1.1.3.cmml">150</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.4.m4.1b"><apply id="S4.SS2.p1.4.m4.1.1.cmml" xref="S4.SS2.p1.4.m4.1.1"><eq id="S4.SS2.p1.4.m4.1.1.1.cmml" xref="S4.SS2.p1.4.m4.1.1.1"></eq><ci id="S4.SS2.p1.4.m4.1.1.2.cmml" xref="S4.SS2.p1.4.m4.1.1.2">𝜏</ci><cn type="integer" id="S4.SS2.p1.4.m4.1.1.3.cmml" xref="S4.SS2.p1.4.m4.1.1.3">150</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.4.m4.1c">\tau=150</annotation></semantics></math> steps. For all experiments with <span id="S4.SS2.p1.8.1" class="ltx_text ltx_font_italic">path-specific modules</span>, the transformer blocks <math id="S4.SS2.p1.5.m5.1" class="ltx_Math" alttext="0" display="inline"><semantics id="S4.SS2.p1.5.m5.1a"><mn id="S4.SS2.p1.5.m5.1.1" xref="S4.SS2.p1.5.m5.1.1.cmml">0</mn><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.5.m5.1b"><cn type="integer" id="S4.SS2.p1.5.m5.1.1.cmml" xref="S4.SS2.p1.5.m5.1.1">0</cn></annotation-xml></semantics></math>, <math id="S4.SS2.p1.6.m6.1" class="ltx_Math" alttext="5" display="inline"><semantics id="S4.SS2.p1.6.m6.1a"><mn id="S4.SS2.p1.6.m6.1.1" xref="S4.SS2.p1.6.m6.1.1.cmml">5</mn><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.6.m6.1b"><cn type="integer" id="S4.SS2.p1.6.m6.1.1.cmml" xref="S4.SS2.p1.6.m6.1.1">5</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.6.m6.1c">5</annotation></semantics></math>, <math id="S4.SS2.p1.7.m7.1" class="ltx_Math" alttext="6" display="inline"><semantics id="S4.SS2.p1.7.m7.1a"><mn id="S4.SS2.p1.7.m7.1.1" xref="S4.SS2.p1.7.m7.1.1.cmml">6</mn><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.7.m7.1b"><cn type="integer" id="S4.SS2.p1.7.m7.1.1.cmml" xref="S4.SS2.p1.7.m7.1.1">6</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.7.m7.1c">6</annotation></semantics></math> <math id="S4.SS2.p1.8.m8.1" class="ltx_Math" alttext="11" display="inline"><semantics id="S4.SS2.p1.8.m8.1a"><mn id="S4.SS2.p1.8.m8.1.1" xref="S4.SS2.p1.8.m8.1.1.cmml">11</mn><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.8.m8.1b"><cn type="integer" id="S4.SS2.p1.8.m8.1.1.cmml" xref="S4.SS2.p1.8.m8.1.1">11</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.8.m8.1c">11</annotation></semantics></math>, and the embedding matrix are not communicated across paths.</p>
</div>
</section>
<section id="S4.SS3" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.3 </span>Parameter sharing between paths</h3>

<figure id="S4.T1" class="ltx_table">
<table id="S4.T1.35" class="ltx_tabular ltx_centering ltx_guessed_headers ltx_align_middle">
<tbody class="ltx_tbody">
<tr id="S4.T1.35.36.1" class="ltx_tr">
<th id="S4.T1.35.36.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_tt">Model</th>
<td id="S4.T1.35.36.1.2" class="ltx_td ltx_align_center ltx_border_tt">Time</td>
<td id="S4.T1.35.36.1.3" class="ltx_td ltx_align_center ltx_border_tt">Compute and Data</td>
<td id="S4.T1.35.36.1.4" class="ltx_td ltx_align_center ltx_border_tt">Total Parameters</td>
<td id="S4.T1.35.36.1.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_tt">Validation PPL</td>
</tr>
<tr id="S4.T1.3.3" class="ltx_tr">
<th id="S4.T1.3.3.4" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t">Baseline</th>
<td id="S4.T1.1.1.1" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.1.1.1.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.1.1.1.m1.1a"><mrow id="S4.T1.1.1.1.m1.1b"><mn id="S4.T1.1.1.1.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.1.1.1.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.1.1.1.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.2.2.2" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.2.2.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.2.2.2.m1.1a"><mrow id="S4.T1.2.2.2.m1.1b"><mn id="S4.T1.2.2.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.2.2.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.2.2.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.3.3.3" class="ltx_td ltx_align_center ltx_border_t">
<math id="S4.T1.3.3.3.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.T1.3.3.3.m1.1a"><mn id="S4.T1.3.3.3.m1.1.1" xref="S4.T1.3.3.3.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.T1.3.3.3.m1.1b"><cn type="integer" id="S4.T1.3.3.3.m1.1.1.cmml" xref="S4.T1.3.3.3.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.3.3.3.m1.1c">150</annotation></semantics></math>M</td>
<td id="S4.T1.3.3.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">16.23</td>
</tr>
<tr id="S4.T1.7.7" class="ltx_tr">
<th id="S4.T1.4.4.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t">DiLoCo <math id="S4.T1.4.4.1.m1.1" class="ltx_Math" alttext="P=8" display="inline"><semantics id="S4.T1.4.4.1.m1.1a"><mrow id="S4.T1.4.4.1.m1.1.1" xref="S4.T1.4.4.1.m1.1.1.cmml"><mi id="S4.T1.4.4.1.m1.1.1.2" xref="S4.T1.4.4.1.m1.1.1.2.cmml">P</mi><mo id="S4.T1.4.4.1.m1.1.1.1" xref="S4.T1.4.4.1.m1.1.1.1.cmml">=</mo><mn id="S4.T1.4.4.1.m1.1.1.3" xref="S4.T1.4.4.1.m1.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.4.4.1.m1.1b"><apply id="S4.T1.4.4.1.m1.1.1.cmml" xref="S4.T1.4.4.1.m1.1.1"><eq id="S4.T1.4.4.1.m1.1.1.1.cmml" xref="S4.T1.4.4.1.m1.1.1.1"></eq><ci id="S4.T1.4.4.1.m1.1.1.2.cmml" xref="S4.T1.4.4.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T1.4.4.1.m1.1.1.3.cmml" xref="S4.T1.4.4.1.m1.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.4.4.1.m1.1c">P=8</annotation></semantics></math>
</th>
<td id="S4.T1.5.5.2" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.5.5.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.5.5.2.m1.1a"><mrow id="S4.T1.5.5.2.m1.1b"><mn id="S4.T1.5.5.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.5.5.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.5.5.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.6.6.3" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.6.6.3.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.6.6.3.m1.1a"><mrow id="S4.T1.6.6.3.m1.1b"><mn id="S4.T1.6.6.3.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.6.6.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.6.6.3.m1.1c">8\times</annotation></semantics></math></td>
<td id="S4.T1.7.7.4" class="ltx_td ltx_align_center ltx_border_t">
<math id="S4.T1.7.7.4.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.T1.7.7.4.m1.1a"><mn id="S4.T1.7.7.4.m1.1.1" xref="S4.T1.7.7.4.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.T1.7.7.4.m1.1b"><cn type="integer" id="S4.T1.7.7.4.m1.1.1.cmml" xref="S4.T1.7.7.4.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.7.7.4.m1.1c">150</annotation></semantics></math>M</td>
<td id="S4.T1.7.7.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">15.02</td>
</tr>
<tr id="S4.T1.11.11" class="ltx_tr">
<th id="S4.T1.8.8.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">DiLoCo <math id="S4.T1.8.8.1.m1.1" class="ltx_Math" alttext="P=64" display="inline"><semantics id="S4.T1.8.8.1.m1.1a"><mrow id="S4.T1.8.8.1.m1.1.1" xref="S4.T1.8.8.1.m1.1.1.cmml"><mi id="S4.T1.8.8.1.m1.1.1.2" xref="S4.T1.8.8.1.m1.1.1.2.cmml">P</mi><mo id="S4.T1.8.8.1.m1.1.1.1" xref="S4.T1.8.8.1.m1.1.1.1.cmml">=</mo><mn id="S4.T1.8.8.1.m1.1.1.3" xref="S4.T1.8.8.1.m1.1.1.3.cmml">64</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.8.8.1.m1.1b"><apply id="S4.T1.8.8.1.m1.1.1.cmml" xref="S4.T1.8.8.1.m1.1.1"><eq id="S4.T1.8.8.1.m1.1.1.1.cmml" xref="S4.T1.8.8.1.m1.1.1.1"></eq><ci id="S4.T1.8.8.1.m1.1.1.2.cmml" xref="S4.T1.8.8.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T1.8.8.1.m1.1.1.3.cmml" xref="S4.T1.8.8.1.m1.1.1.3">64</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.8.8.1.m1.1c">P=64</annotation></semantics></math>
</th>
<td id="S4.T1.9.9.2" class="ltx_td ltx_align_center"><math id="S4.T1.9.9.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.9.9.2.m1.1a"><mrow id="S4.T1.9.9.2.m1.1b"><mn id="S4.T1.9.9.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.9.9.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.9.9.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.10.10.3" class="ltx_td ltx_align_center"><math id="S4.T1.10.10.3.m1.1" class="ltx_math_unparsed" alttext="64\times" display="inline"><semantics id="S4.T1.10.10.3.m1.1a"><mrow id="S4.T1.10.10.3.m1.1b"><mn id="S4.T1.10.10.3.m1.1.1">64</mn><mo lspace="0.222em" id="S4.T1.10.10.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.10.10.3.m1.1c">64\times</annotation></semantics></math></td>
<td id="S4.T1.11.11.4" class="ltx_td ltx_align_center">
<math id="S4.T1.11.11.4.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.T1.11.11.4.m1.1a"><mn id="S4.T1.11.11.4.m1.1.1" xref="S4.T1.11.11.4.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.T1.11.11.4.m1.1b"><cn type="integer" id="S4.T1.11.11.4.m1.1.1.cmml" xref="S4.T1.11.11.4.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.11.11.4.m1.1c">150</annotation></semantics></math>M</td>
<td id="S4.T1.11.11.5" class="ltx_td ltx_nopad_r ltx_align_center">14.96</td>
</tr>
<tr id="S4.T1.15.15" class="ltx_tr">
<th id="S4.T1.12.12.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t">Flat MoE <math id="S4.T1.12.12.1.m1.1" class="ltx_Math" alttext="P=8" display="inline"><semantics id="S4.T1.12.12.1.m1.1a"><mrow id="S4.T1.12.12.1.m1.1.1" xref="S4.T1.12.12.1.m1.1.1.cmml"><mi id="S4.T1.12.12.1.m1.1.1.2" xref="S4.T1.12.12.1.m1.1.1.2.cmml">P</mi><mo id="S4.T1.12.12.1.m1.1.1.1" xref="S4.T1.12.12.1.m1.1.1.1.cmml">=</mo><mn id="S4.T1.12.12.1.m1.1.1.3" xref="S4.T1.12.12.1.m1.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.12.12.1.m1.1b"><apply id="S4.T1.12.12.1.m1.1.1.cmml" xref="S4.T1.12.12.1.m1.1.1"><eq id="S4.T1.12.12.1.m1.1.1.1.cmml" xref="S4.T1.12.12.1.m1.1.1.1"></eq><ci id="S4.T1.12.12.1.m1.1.1.2.cmml" xref="S4.T1.12.12.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T1.12.12.1.m1.1.1.3.cmml" xref="S4.T1.12.12.1.m1.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.12.12.1.m1.1c">P=8</annotation></semantics></math>
</th>
<td id="S4.T1.13.13.2" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.13.13.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.13.13.2.m1.1a"><mrow id="S4.T1.13.13.2.m1.1b"><mn id="S4.T1.13.13.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.13.13.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.13.13.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.14.14.3" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.14.14.3.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.14.14.3.m1.1a"><mrow id="S4.T1.14.14.3.m1.1b"><mn id="S4.T1.14.14.3.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.14.14.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.14.14.3.m1.1c">8\times</annotation></semantics></math></td>
<td id="S4.T1.15.15.4" class="ltx_td ltx_align_center ltx_border_t">
<math id="S4.T1.15.15.4.m1.1" class="ltx_Math" alttext="1.2" display="inline"><semantics id="S4.T1.15.15.4.m1.1a"><mn id="S4.T1.15.15.4.m1.1.1" xref="S4.T1.15.15.4.m1.1.1.cmml">1.2</mn><annotation-xml encoding="MathML-Content" id="S4.T1.15.15.4.m1.1b"><cn type="float" id="S4.T1.15.15.4.m1.1.1.cmml" xref="S4.T1.15.15.4.m1.1.1">1.2</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.15.15.4.m1.1c">1.2</annotation></semantics></math>B</td>
<td id="S4.T1.15.15.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">14.62</td>
</tr>
<tr id="S4.T1.19.19" class="ltx_tr">
<th id="S4.T1.16.16.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Flat MoE <math id="S4.T1.16.16.1.m1.1" class="ltx_Math" alttext="P=64" display="inline"><semantics id="S4.T1.16.16.1.m1.1a"><mrow id="S4.T1.16.16.1.m1.1.1" xref="S4.T1.16.16.1.m1.1.1.cmml"><mi id="S4.T1.16.16.1.m1.1.1.2" xref="S4.T1.16.16.1.m1.1.1.2.cmml">P</mi><mo id="S4.T1.16.16.1.m1.1.1.1" xref="S4.T1.16.16.1.m1.1.1.1.cmml">=</mo><mn id="S4.T1.16.16.1.m1.1.1.3" xref="S4.T1.16.16.1.m1.1.1.3.cmml">64</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.16.16.1.m1.1b"><apply id="S4.T1.16.16.1.m1.1.1.cmml" xref="S4.T1.16.16.1.m1.1.1"><eq id="S4.T1.16.16.1.m1.1.1.1.cmml" xref="S4.T1.16.16.1.m1.1.1.1"></eq><ci id="S4.T1.16.16.1.m1.1.1.2.cmml" xref="S4.T1.16.16.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T1.16.16.1.m1.1.1.3.cmml" xref="S4.T1.16.16.1.m1.1.1.3">64</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.16.16.1.m1.1c">P=64</annotation></semantics></math>
</th>
<td id="S4.T1.17.17.2" class="ltx_td ltx_align_center"><math id="S4.T1.17.17.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.17.17.2.m1.1a"><mrow id="S4.T1.17.17.2.m1.1b"><mn id="S4.T1.17.17.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.17.17.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.17.17.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.18.18.3" class="ltx_td ltx_align_center"><math id="S4.T1.18.18.3.m1.1" class="ltx_math_unparsed" alttext="64\times" display="inline"><semantics id="S4.T1.18.18.3.m1.1a"><mrow id="S4.T1.18.18.3.m1.1b"><mn id="S4.T1.18.18.3.m1.1.1">64</mn><mo lspace="0.222em" id="S4.T1.18.18.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.18.18.3.m1.1c">64\times</annotation></semantics></math></td>
<td id="S4.T1.19.19.4" class="ltx_td ltx_align_center">
<math id="S4.T1.19.19.4.m1.1" class="ltx_Math" alttext="9.6" display="inline"><semantics id="S4.T1.19.19.4.m1.1a"><mn id="S4.T1.19.19.4.m1.1.1" xref="S4.T1.19.19.4.m1.1.1.cmml">9.6</mn><annotation-xml encoding="MathML-Content" id="S4.T1.19.19.4.m1.1b"><cn type="float" id="S4.T1.19.19.4.m1.1.1.cmml" xref="S4.T1.19.19.4.m1.1.1">9.6</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.19.19.4.m1.1c">9.6</annotation></semantics></math>B</td>
<td id="S4.T1.19.19.5" class="ltx_td ltx_nopad_r ltx_align_center">12.76</td>
</tr>
<tr id="S4.T1.23.23" class="ltx_tr">
<th id="S4.T1.20.20.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t">DiPaCo <math id="S4.T1.20.20.1.m1.1" class="ltx_Math" alttext="2\times 4" display="inline"><semantics id="S4.T1.20.20.1.m1.1a"><mrow id="S4.T1.20.20.1.m1.1.1" xref="S4.T1.20.20.1.m1.1.1.cmml"><mn id="S4.T1.20.20.1.m1.1.1.2" xref="S4.T1.20.20.1.m1.1.1.2.cmml">2</mn><mo lspace="0.222em" rspace="0.222em" id="S4.T1.20.20.1.m1.1.1.1" xref="S4.T1.20.20.1.m1.1.1.1.cmml">×</mo><mn id="S4.T1.20.20.1.m1.1.1.3" xref="S4.T1.20.20.1.m1.1.1.3.cmml">4</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.20.20.1.m1.1b"><apply id="S4.T1.20.20.1.m1.1.1.cmml" xref="S4.T1.20.20.1.m1.1.1"><times id="S4.T1.20.20.1.m1.1.1.1.cmml" xref="S4.T1.20.20.1.m1.1.1.1"></times><cn type="integer" id="S4.T1.20.20.1.m1.1.1.2.cmml" xref="S4.T1.20.20.1.m1.1.1.2">2</cn><cn type="integer" id="S4.T1.20.20.1.m1.1.1.3.cmml" xref="S4.T1.20.20.1.m1.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.20.20.1.m1.1c">2\times 4</annotation></semantics></math>
</th>
<td id="S4.T1.21.21.2" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.21.21.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.21.21.2.m1.1a"><mrow id="S4.T1.21.21.2.m1.1b"><mn id="S4.T1.21.21.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.21.21.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.21.21.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.22.22.3" class="ltx_td ltx_align_center ltx_border_t"><math id="S4.T1.22.22.3.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.22.22.3.m1.1a"><mrow id="S4.T1.22.22.3.m1.1b"><mn id="S4.T1.22.22.3.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.22.22.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.22.22.3.m1.1c">8\times</annotation></semantics></math></td>
<td id="S4.T1.23.23.4" class="ltx_td ltx_align_center ltx_border_t">
<math id="S4.T1.23.23.4.m1.1" class="ltx_Math" alttext="0.5" display="inline"><semantics id="S4.T1.23.23.4.m1.1a"><mn id="S4.T1.23.23.4.m1.1.1" xref="S4.T1.23.23.4.m1.1.1.cmml">0.5</mn><annotation-xml encoding="MathML-Content" id="S4.T1.23.23.4.m1.1b"><cn type="float" id="S4.T1.23.23.4.m1.1.1.cmml" xref="S4.T1.23.23.4.m1.1.1">0.5</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.23.23.4.m1.1c">0.5</annotation></semantics></math>B</td>
<td id="S4.T1.23.23.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">14.86</td>
</tr>
<tr id="S4.T1.27.27" class="ltx_tr">
<th id="S4.T1.24.24.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">DiPaCo <math id="S4.T1.24.24.1.m1.1" class="ltx_Math" alttext="8\times 8" display="inline"><semantics id="S4.T1.24.24.1.m1.1a"><mrow id="S4.T1.24.24.1.m1.1.1" xref="S4.T1.24.24.1.m1.1.1.cmml"><mn id="S4.T1.24.24.1.m1.1.1.2" xref="S4.T1.24.24.1.m1.1.1.2.cmml">8</mn><mo lspace="0.222em" rspace="0.222em" id="S4.T1.24.24.1.m1.1.1.1" xref="S4.T1.24.24.1.m1.1.1.1.cmml">×</mo><mn id="S4.T1.24.24.1.m1.1.1.3" xref="S4.T1.24.24.1.m1.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.24.24.1.m1.1b"><apply id="S4.T1.24.24.1.m1.1.1.cmml" xref="S4.T1.24.24.1.m1.1.1"><times id="S4.T1.24.24.1.m1.1.1.1.cmml" xref="S4.T1.24.24.1.m1.1.1.1"></times><cn type="integer" id="S4.T1.24.24.1.m1.1.1.2.cmml" xref="S4.T1.24.24.1.m1.1.1.2">8</cn><cn type="integer" id="S4.T1.24.24.1.m1.1.1.3.cmml" xref="S4.T1.24.24.1.m1.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.24.24.1.m1.1c">8\times 8</annotation></semantics></math>
</th>
<td id="S4.T1.25.25.2" class="ltx_td ltx_align_center"><math id="S4.T1.25.25.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.25.25.2.m1.1a"><mrow id="S4.T1.25.25.2.m1.1b"><mn id="S4.T1.25.25.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.25.25.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.25.25.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.26.26.3" class="ltx_td ltx_align_center"><math id="S4.T1.26.26.3.m1.1" class="ltx_math_unparsed" alttext="64\times" display="inline"><semantics id="S4.T1.26.26.3.m1.1a"><mrow id="S4.T1.26.26.3.m1.1b"><mn id="S4.T1.26.26.3.m1.1.1">64</mn><mo lspace="0.222em" id="S4.T1.26.26.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.26.26.3.m1.1c">64\times</annotation></semantics></math></td>
<td id="S4.T1.27.27.4" class="ltx_td ltx_align_center">
<math id="S4.T1.27.27.4.m1.1" class="ltx_Math" alttext="1.2" display="inline"><semantics id="S4.T1.27.27.4.m1.1a"><mn id="S4.T1.27.27.4.m1.1.1" xref="S4.T1.27.27.4.m1.1.1.cmml">1.2</mn><annotation-xml encoding="MathML-Content" id="S4.T1.27.27.4.m1.1b"><cn type="float" id="S4.T1.27.27.4.m1.1.1.cmml" xref="S4.T1.27.27.4.m1.1.1">1.2</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.27.27.4.m1.1c">1.2</annotation></semantics></math>B</td>
<td id="S4.T1.27.27.5" class="ltx_td ltx_nopad_r ltx_align_center">13.37</td>
</tr>
<tr id="S4.T1.31.31" class="ltx_tr">
<th id="S4.T1.28.28.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">DiPaCo <math id="S4.T1.28.28.1.m1.1" class="ltx_Math" alttext="8\times 8+" display="inline"><semantics id="S4.T1.28.28.1.m1.1a"><mrow id="S4.T1.28.28.1.m1.1.1" xref="S4.T1.28.28.1.m1.1.1.cmml"><mrow id="S4.T1.28.28.1.m1.1.1.2" xref="S4.T1.28.28.1.m1.1.1.2.cmml"><mn id="S4.T1.28.28.1.m1.1.1.2.2" xref="S4.T1.28.28.1.m1.1.1.2.2.cmml">8</mn><mo lspace="0.222em" rspace="0.222em" id="S4.T1.28.28.1.m1.1.1.2.1" xref="S4.T1.28.28.1.m1.1.1.2.1.cmml">×</mo><mn id="S4.T1.28.28.1.m1.1.1.2.3" xref="S4.T1.28.28.1.m1.1.1.2.3.cmml">8</mn></mrow><mo id="S4.T1.28.28.1.m1.1.1.3" xref="S4.T1.28.28.1.m1.1.1.3.cmml">+</mo></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.28.28.1.m1.1b"><apply id="S4.T1.28.28.1.m1.1.1.cmml" xref="S4.T1.28.28.1.m1.1.1"><csymbol cd="latexml" id="S4.T1.28.28.1.m1.1.1.1.cmml" xref="S4.T1.28.28.1.m1.1.1">limit-from</csymbol><apply id="S4.T1.28.28.1.m1.1.1.2.cmml" xref="S4.T1.28.28.1.m1.1.1.2"><times id="S4.T1.28.28.1.m1.1.1.2.1.cmml" xref="S4.T1.28.28.1.m1.1.1.2.1"></times><cn type="integer" id="S4.T1.28.28.1.m1.1.1.2.2.cmml" xref="S4.T1.28.28.1.m1.1.1.2.2">8</cn><cn type="integer" id="S4.T1.28.28.1.m1.1.1.2.3.cmml" xref="S4.T1.28.28.1.m1.1.1.2.3">8</cn></apply><plus id="S4.T1.28.28.1.m1.1.1.3.cmml" xref="S4.T1.28.28.1.m1.1.1.3"></plus></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.28.28.1.m1.1c">8\times 8+</annotation></semantics></math> Path Specific Modules</th>
<td id="S4.T1.29.29.2" class="ltx_td ltx_align_center"><math id="S4.T1.29.29.2.m1.1" class="ltx_math_unparsed" alttext="1\times" display="inline"><semantics id="S4.T1.29.29.2.m1.1a"><mrow id="S4.T1.29.29.2.m1.1b"><mn id="S4.T1.29.29.2.m1.1.1">1</mn><mo lspace="0.222em" id="S4.T1.29.29.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.29.29.2.m1.1c">1\times</annotation></semantics></math></td>
<td id="S4.T1.30.30.3" class="ltx_td ltx_align_center"><math id="S4.T1.30.30.3.m1.1" class="ltx_math_unparsed" alttext="64\times" display="inline"><semantics id="S4.T1.30.30.3.m1.1a"><mrow id="S4.T1.30.30.3.m1.1b"><mn id="S4.T1.30.30.3.m1.1.1">64</mn><mo lspace="0.222em" id="S4.T1.30.30.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.30.30.3.m1.1c">64\times</annotation></semantics></math></td>
<td id="S4.T1.31.31.4" class="ltx_td ltx_align_center">
<math id="S4.T1.31.31.4.m1.1" class="ltx_Math" alttext="5.3" display="inline"><semantics id="S4.T1.31.31.4.m1.1a"><mn id="S4.T1.31.31.4.m1.1.1" xref="S4.T1.31.31.4.m1.1.1.cmml">5.3</mn><annotation-xml encoding="MathML-Content" id="S4.T1.31.31.4.m1.1b"><cn type="float" id="S4.T1.31.31.4.m1.1.1.cmml" xref="S4.T1.31.31.4.m1.1.1">5.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.31.31.4.m1.1c">5.3</annotation></semantics></math>B</td>
<td id="S4.T1.31.31.5" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T1.31.31.5.1" class="ltx_text ltx_font_bold">12.70</span></td>
</tr>
<tr id="S4.T1.35.35" class="ltx_tr">
<th id="S4.T1.32.32.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_bb ltx_border_r ltx_border_tt">Baseline, <math id="S4.T1.32.32.1.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.32.32.1.m1.1a"><mrow id="S4.T1.32.32.1.m1.1b"><mn id="S4.T1.32.32.1.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.32.32.1.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.32.32.1.m1.1c">8\times</annotation></semantics></math> steps</th>
<td id="S4.T1.33.33.2" class="ltx_td ltx_align_center ltx_border_bb ltx_border_tt"><math id="S4.T1.33.33.2.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.33.33.2.m1.1a"><mrow id="S4.T1.33.33.2.m1.1b"><mn id="S4.T1.33.33.2.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.33.33.2.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.33.33.2.m1.1c">8\times</annotation></semantics></math></td>
<td id="S4.T1.34.34.3" class="ltx_td ltx_align_center ltx_border_bb ltx_border_tt"><math id="S4.T1.34.34.3.m1.1" class="ltx_math_unparsed" alttext="8\times" display="inline"><semantics id="S4.T1.34.34.3.m1.1a"><mrow id="S4.T1.34.34.3.m1.1b"><mn id="S4.T1.34.34.3.m1.1.1">8</mn><mo lspace="0.222em" id="S4.T1.34.34.3.m1.1.2">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T1.34.34.3.m1.1c">8\times</annotation></semantics></math></td>
<td id="S4.T1.35.35.4" class="ltx_td ltx_align_center ltx_border_bb ltx_border_tt">
<math id="S4.T1.35.35.4.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.T1.35.35.4.m1.1a"><mn id="S4.T1.35.35.4.m1.1.1" xref="S4.T1.35.35.4.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.T1.35.35.4.m1.1b"><cn type="integer" id="S4.T1.35.35.4.m1.1.1.cmml" xref="S4.T1.35.35.4.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.35.35.4.m1.1c">150</annotation></semantics></math>M</td>
<td id="S4.T1.35.35.5" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_bb ltx_border_tt">14.72</td>
</tr>
</tbody>
</table>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T1.38.1.1" class="ltx_text" style="font-size:90%;">Table 1</span>: </span><span id="S4.T1.39.2" class="ltx_text ltx_font_bold" style="font-size:90%;">DiPaCo vs. Flat MoE vs. DiLoCo<span id="S4.T1.39.2.1" class="ltx_text ltx_font_medium">: We compare DiLoCo and Flat MoE (<a href="#S2.SS6.SSS3" title="2.6.3 Flat Mixture of Experts ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.6.3</span></a>) vs. DiPaCo. We also compare against the baseline trained for the same FLOPs as the distributed models.
In DiLoCo all paths are collapsed together at each outer optimization step. In Flat MoE, all paths are independent. In DiPaCo, only a subset of the parameters are “collapsed” using information from a subset of paths. Because of the distributed framework, DiPaCo is trained in the same wallclock time as DiLoCo, which is roughly the wallclock of the baseline.</span></span></figcaption>
</figure>
<div id="S4.SS3.p1" class="ltx_para">
<p id="S4.SS3.p1.1" class="ltx_p">At one end of the spectrum, as described in <a href="#S2.SS6.SSS3" title="2.6.3 Flat Mixture of Experts ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.6.3</span></a>, we use paths that share no parameters, and can be trained fully independently. Continuing the trend from <a href="#S4.SS2" title="4.2 Number of paths and parameters ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsection 4.2</span></a>, where there is monotone improvement in validation performance as parameters are added, we see improvements by unsharing all parameters for small number of paths. Using 64 unshared paths we can approach a <math id="S4.SS3.p1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S4.SS3.p1.1.m1.1a"><mrow id="S4.SS3.p1.1.m1.1.1" xref="S4.SS3.p1.1.m1.1.1.cmml"><mn id="S4.SS3.p1.1.m1.1.1.2" xref="S4.SS3.p1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS3.p1.1.m1.1.1.1" xref="S4.SS3.p1.1.m1.1.1.1.cmml">×</mo><mn id="S4.SS3.p1.1.m1.1.1.3" xref="S4.SS3.p1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS3.p1.1.m1.1b"><apply id="S4.SS3.p1.1.m1.1.1.cmml" xref="S4.SS3.p1.1.m1.1.1"><times id="S4.SS3.p1.1.m1.1.1.1.cmml" xref="S4.SS3.p1.1.m1.1.1.1"></times><cn type="integer" id="S4.SS3.p1.1.m1.1.1.2.cmml" xref="S4.SS3.p1.1.m1.1.1.2">16</cn><cn type="integer" id="S4.SS3.p1.1.m1.1.1.3.cmml" xref="S4.SS3.p1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS3.p1.1.m1.1c">16\times 16</annotation></semantics></math> DiPaCo that has no unsharing.</p>
</div>
<div id="S4.SS3.p2" class="ltx_para">
<p id="S4.SS3.p2.1" class="ltx_p">However, this model starts over-fitting once the number of shards (and corresponding experts) is further increased as shown in <a href="#S4.T2" title="Table 2 ‣ 4.4 Routing frequency during evaluation ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 2</span></a>. We are able to recover some performance by overlapping shards, but still we find that we cannot naively scale the number of fully independent paths. We consider approaches for more sophisticated overlapping of shards and path branching to be a promising direction for future work.</p>
</div>
<div id="S4.SS3.p3" class="ltx_para">
<p id="S4.SS3.p3.2" class="ltx_p">At the other end of the spectrum, paths can share all parameters, effectively collapsing back into a single path and recovering the performance of the original DiLoCo <cite class="ltx_cite ltx_citemacro_citep">(Douillard et al., <a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite>.
We can see that the base model trained with DiLoCo, or base model overtrained for many steps, although receiving the same number of FLOPs and tokens as the larger DiPaCo models, do not have the capacity to make use of these extra FLOPs. Of course a properly sized baseline would do much better; and indeed the 1.3B model above already has better perplexity. However, DiPaCo consists of paths that each require only <math id="S4.SS3.p3.1.m1.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.SS3.p3.1.m1.1a"><mn id="S4.SS3.p3.1.m1.1.1" xref="S4.SS3.p3.1.m1.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.SS3.p3.1.m1.1b"><cn type="integer" id="S4.SS3.p3.1.m1.1.1.cmml" xref="S4.SS3.p3.1.m1.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS3.p3.1.m1.1c">150</annotation></semantics></math>M size island of compute both in train and testing, and its distributed training algorithm is built into the architecture, so that only paths of <math id="S4.SS3.p3.2.m2.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.SS3.p3.2.m2.1a"><mn id="S4.SS3.p3.2.m2.1.1" xref="S4.SS3.p3.2.m2.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.SS3.p3.2.m2.1b"><cn type="integer" id="S4.SS3.p3.2.m2.1.1.cmml" xref="S4.SS3.p3.2.m2.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS3.p3.2.m2.1c">150</annotation></semantics></math>M parameters ever need to be materialized on any training or deployment worker. We compare all these variants in <a href="#S4.T1" title="Table 1 ‣ 4.3 Parameter sharing between paths ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 1</span></a>.</p>
</div>
</section>
<section id="S4.SS4" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.4 </span>Routing frequency during evaluation</h3>

<div id="S4.SS4.p1" class="ltx_para">
<p id="S4.SS4.p1.9" class="ltx_p">We demonstrate the results of routing more frequently at test time in <a href="#S4.T3" title="Table 3 ‣ 4.4 Routing frequency during evaluation ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 3</span></a>. We observe a significant improvement of <math id="S4.SS4.p1.1.m1.1" class="ltx_Math" alttext="0.74" display="inline"><semantics id="S4.SS4.p1.1.m1.1a"><mn id="S4.SS4.p1.1.m1.1.1" xref="S4.SS4.p1.1.m1.1.1.cmml">0.74</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.1.m1.1b"><cn type="float" id="S4.SS4.p1.1.m1.1.1.cmml" xref="S4.SS4.p1.1.m1.1.1">0.74</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.1.m1.1c">0.74</annotation></semantics></math> perplexity points going from routing once per sequence of <math id="S4.SS4.p1.2.m2.2" class="ltx_Math" alttext="1{,}024" display="inline"><semantics id="S4.SS4.p1.2.m2.2a"><mrow id="S4.SS4.p1.2.m2.2.3.2" xref="S4.SS4.p1.2.m2.2.3.1.cmml"><mn id="S4.SS4.p1.2.m2.1.1" xref="S4.SS4.p1.2.m2.1.1.cmml">1</mn><mo id="S4.SS4.p1.2.m2.2.3.2.1" xref="S4.SS4.p1.2.m2.2.3.1.cmml">,</mo><mn id="S4.SS4.p1.2.m2.2.2" xref="S4.SS4.p1.2.m2.2.2.cmml">024</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.2.m2.2b"><list id="S4.SS4.p1.2.m2.2.3.1.cmml" xref="S4.SS4.p1.2.m2.2.3.2"><cn type="integer" id="S4.SS4.p1.2.m2.1.1.cmml" xref="S4.SS4.p1.2.m2.1.1">1</cn><cn type="integer" id="S4.SS4.p1.2.m2.2.2.cmml" xref="S4.SS4.p1.2.m2.2.2">024</cn></list></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.2.m2.2c">1{,}024</annotation></semantics></math> tokens to every <math id="S4.SS4.p1.3.m3.1" class="ltx_Math" alttext="128" display="inline"><semantics id="S4.SS4.p1.3.m3.1a"><mn id="S4.SS4.p1.3.m3.1.1" xref="S4.SS4.p1.3.m3.1.1.cmml">128</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.3.m3.1b"><cn type="integer" id="S4.SS4.p1.3.m3.1.1.cmml" xref="S4.SS4.p1.3.m3.1.1">128</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.3.m3.1c">128</annotation></semantics></math> tokens. Further increasing the granularity of the re-routing consistently improve results, although more marginally so.
This shows that at training time we can route at the sequence level to enable pre-sharding for more efficient learning, and then at evaluation time we can gain back performance by routing more frequently. By routing tokens potentially to a different path every 64 tokens, DiPaCo reaches <math id="S4.SS4.p1.4.m4.1" class="ltx_Math" alttext="11.38" display="inline"><semantics id="S4.SS4.p1.4.m4.1a"><mn id="S4.SS4.p1.4.m4.1.1" xref="S4.SS4.p1.4.m4.1.1.cmml">11.38</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.4.m4.1b"><cn type="float" id="S4.SS4.p1.4.m4.1.1.cmml" xref="S4.SS4.p1.4.m4.1.1">11.38</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.4.m4.1c">11.38</annotation></semantics></math> perplexity, matching the performance of a <math id="S4.SS4.p1.5.m5.1" class="ltx_Math" alttext="1" display="inline"><semantics id="S4.SS4.p1.5.m5.1a"><mn id="S4.SS4.p1.5.m5.1.1" xref="S4.SS4.p1.5.m5.1.1.cmml">1</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.5.m5.1b"><cn type="integer" id="S4.SS4.p1.5.m5.1.1.cmml" xref="S4.SS4.p1.5.m5.1.1">1</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.5.m5.1c">1</annotation></semantics></math>B dense model (<math id="S4.SS4.p1.6.m6.1" class="ltx_Math" alttext="11.41" display="inline"><semantics id="S4.SS4.p1.6.m6.1a"><mn id="S4.SS4.p1.6.m6.1.1" xref="S4.SS4.p1.6.m6.1.1.cmml">11.41</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.6.m6.1b"><cn type="float" id="S4.SS4.p1.6.m6.1.1.cmml" xref="S4.SS4.p1.6.m6.1.1">11.41</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.6.m6.1c">11.41</annotation></semantics></math>). Since at any time DiPaCo uses a single path of size <math id="S4.SS4.p1.7.m7.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.SS4.p1.7.m7.1a"><mn id="S4.SS4.p1.7.m7.1.1" xref="S4.SS4.p1.7.m7.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.7.m7.1b"><cn type="integer" id="S4.SS4.p1.7.m7.1.1.cmml" xref="S4.SS4.p1.7.m7.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.7.m7.1c">150</annotation></semantics></math>M parameters, it matches the performance of the <math id="S4.SS4.p1.8.m8.1" class="ltx_Math" alttext="1" display="inline"><semantics id="S4.SS4.p1.8.m8.1a"><mn id="S4.SS4.p1.8.m8.1.1" xref="S4.SS4.p1.8.m8.1.1.cmml">1</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.8.m8.1b"><cn type="integer" id="S4.SS4.p1.8.m8.1.1.cmml" xref="S4.SS4.p1.8.m8.1.1">1</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.8.m8.1c">1</annotation></semantics></math>B model using <math id="S4.SS4.p1.9.m9.1" class="ltx_Math" alttext="6" display="inline"><semantics id="S4.SS4.p1.9.m9.1a"><mn id="S4.SS4.p1.9.m9.1.1" xref="S4.SS4.p1.9.m9.1.1.cmml">6</mn><annotation-xml encoding="MathML-Content" id="S4.SS4.p1.9.m9.1b"><cn type="integer" id="S4.SS4.p1.9.m9.1.1.cmml" xref="S4.SS4.p1.9.m9.1.1">6</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p1.9.m9.1c">6</annotation></semantics></math> times less parameters, and therefore, requiring less compute at inference time.</p>
</div>
<figure id="S4.T2" class="ltx_table">
<table id="S4.T2.3" class="ltx_tabular ltx_centering ltx_guessed_headers ltx_align_middle">
<tbody class="ltx_tbody">
<tr id="S4.T2.3.4.1" class="ltx_tr">
<th id="S4.T2.3.4.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_tt"># fully independent paths</th>
<td id="S4.T2.3.4.1.2" class="ltx_td ltx_align_center ltx_border_tt">Validation PPL</td>
</tr>
<tr id="S4.T2.1.1" class="ltx_tr">
<th id="S4.T2.1.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t"><math id="S4.T2.1.1.1.m1.1" class="ltx_Math" alttext="P=8" display="inline"><semantics id="S4.T2.1.1.1.m1.1a"><mrow id="S4.T2.1.1.1.m1.1.1" xref="S4.T2.1.1.1.m1.1.1.cmml"><mi id="S4.T2.1.1.1.m1.1.1.2" xref="S4.T2.1.1.1.m1.1.1.2.cmml">P</mi><mo id="S4.T2.1.1.1.m1.1.1.1" xref="S4.T2.1.1.1.m1.1.1.1.cmml">=</mo><mn id="S4.T2.1.1.1.m1.1.1.3" xref="S4.T2.1.1.1.m1.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T2.1.1.1.m1.1b"><apply id="S4.T2.1.1.1.m1.1.1.cmml" xref="S4.T2.1.1.1.m1.1.1"><eq id="S4.T2.1.1.1.m1.1.1.1.cmml" xref="S4.T2.1.1.1.m1.1.1.1"></eq><ci id="S4.T2.1.1.1.m1.1.1.2.cmml" xref="S4.T2.1.1.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T2.1.1.1.m1.1.1.3.cmml" xref="S4.T2.1.1.1.m1.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.1.1.1.m1.1c">P=8</annotation></semantics></math></th>
<td id="S4.T2.1.1.2" class="ltx_td ltx_align_center ltx_border_t">14.6</td>
</tr>
<tr id="S4.T2.2.2" class="ltx_tr">
<th id="S4.T2.2.2.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r"><math id="S4.T2.2.2.1.m1.1" class="ltx_Math" alttext="P=16" display="inline"><semantics id="S4.T2.2.2.1.m1.1a"><mrow id="S4.T2.2.2.1.m1.1.1" xref="S4.T2.2.2.1.m1.1.1.cmml"><mi id="S4.T2.2.2.1.m1.1.1.2" xref="S4.T2.2.2.1.m1.1.1.2.cmml">P</mi><mo id="S4.T2.2.2.1.m1.1.1.1" xref="S4.T2.2.2.1.m1.1.1.1.cmml">=</mo><mn id="S4.T2.2.2.1.m1.1.1.3" xref="S4.T2.2.2.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T2.2.2.1.m1.1b"><apply id="S4.T2.2.2.1.m1.1.1.cmml" xref="S4.T2.2.2.1.m1.1.1"><eq id="S4.T2.2.2.1.m1.1.1.1.cmml" xref="S4.T2.2.2.1.m1.1.1.1"></eq><ci id="S4.T2.2.2.1.m1.1.1.2.cmml" xref="S4.T2.2.2.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T2.2.2.1.m1.1.1.3.cmml" xref="S4.T2.2.2.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.2.2.1.m1.1c">P=16</annotation></semantics></math></th>
<td id="S4.T2.2.2.2" class="ltx_td ltx_align_center">13.9</td>
</tr>
<tr id="S4.T2.3.3" class="ltx_tr">
<th id="S4.T2.3.3.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_bb ltx_border_r"><math id="S4.T2.3.3.1.m1.1" class="ltx_Math" alttext="P=256" display="inline"><semantics id="S4.T2.3.3.1.m1.1a"><mrow id="S4.T2.3.3.1.m1.1.1" xref="S4.T2.3.3.1.m1.1.1.cmml"><mi id="S4.T2.3.3.1.m1.1.1.2" xref="S4.T2.3.3.1.m1.1.1.2.cmml">P</mi><mo id="S4.T2.3.3.1.m1.1.1.1" xref="S4.T2.3.3.1.m1.1.1.1.cmml">=</mo><mn id="S4.T2.3.3.1.m1.1.1.3" xref="S4.T2.3.3.1.m1.1.1.3.cmml">256</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T2.3.3.1.m1.1b"><apply id="S4.T2.3.3.1.m1.1.1.cmml" xref="S4.T2.3.3.1.m1.1.1"><eq id="S4.T2.3.3.1.m1.1.1.1.cmml" xref="S4.T2.3.3.1.m1.1.1.1"></eq><ci id="S4.T2.3.3.1.m1.1.1.2.cmml" xref="S4.T2.3.3.1.m1.1.1.2">𝑃</ci><cn type="integer" id="S4.T2.3.3.1.m1.1.1.3.cmml" xref="S4.T2.3.3.1.m1.1.1.3">256</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.3.3.1.m1.1c">P=256</annotation></semantics></math></th>
<td id="S4.T2.3.3.2" class="ltx_td ltx_align_center ltx_border_bb">14.2</td>
</tr>
</tbody>
</table>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T2.20.8.1" class="ltx_text" style="font-size:90%;">Table 2</span>: </span><span id="S4.T2.17.7" class="ltx_text ltx_font_bold" style="font-size:90%;">Flat MoE (independent paths) overfits as number of paths increase<span id="S4.T2.17.7.7" class="ltx_text ltx_font_medium">: Validation perplexity of
DiPaCo with no shared modules, as described in <a href="#S2.SS6.SSS3" title="2.6.3 Flat Mixture of Experts ‣ 2.6 DiPaCo ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.6.3</span></a> after <math id="S4.T2.11.1.1.m1.1" class="ltx_Math" alttext="10" display="inline"><semantics id="S4.T2.11.1.1.m1.1b"><mn id="S4.T2.11.1.1.m1.1.1" xref="S4.T2.11.1.1.m1.1.1.cmml">10</mn><annotation-xml encoding="MathML-Content" id="S4.T2.11.1.1.m1.1c"><cn type="integer" id="S4.T2.11.1.1.m1.1.1.cmml" xref="S4.T2.11.1.1.m1.1.1">10</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.11.1.1.m1.1d">10</annotation></semantics></math>K steps of training. With overlapping shards as described in <a href="#S2.SS4.SSS4" title="2.4.4 Overlapping Shards ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">subsubsection 2.4.4</span></a> and early stopping, we can improve the <math id="S4.T2.12.2.2.m2.1" class="ltx_Math" alttext="P=256" display="inline"><semantics id="S4.T2.12.2.2.m2.1b"><mrow id="S4.T2.12.2.2.m2.1.1" xref="S4.T2.12.2.2.m2.1.1.cmml"><mi id="S4.T2.12.2.2.m2.1.1.2" xref="S4.T2.12.2.2.m2.1.1.2.cmml">P</mi><mo id="S4.T2.12.2.2.m2.1.1.1" xref="S4.T2.12.2.2.m2.1.1.1.cmml">=</mo><mn id="S4.T2.12.2.2.m2.1.1.3" xref="S4.T2.12.2.2.m2.1.1.3.cmml">256</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T2.12.2.2.m2.1c"><apply id="S4.T2.12.2.2.m2.1.1.cmml" xref="S4.T2.12.2.2.m2.1.1"><eq id="S4.T2.12.2.2.m2.1.1.1.cmml" xref="S4.T2.12.2.2.m2.1.1.1"></eq><ci id="S4.T2.12.2.2.m2.1.1.2.cmml" xref="S4.T2.12.2.2.m2.1.1.2">𝑃</ci><cn type="integer" id="S4.T2.12.2.2.m2.1.1.3.cmml" xref="S4.T2.12.2.2.m2.1.1.3">256</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.12.2.2.m2.1d">P=256</annotation></semantics></math> result to <math id="S4.T2.13.3.3.m3.1" class="ltx_Math" alttext="13.6" display="inline"><semantics id="S4.T2.13.3.3.m3.1b"><mn id="S4.T2.13.3.3.m3.1.1" xref="S4.T2.13.3.3.m3.1.1.cmml">13.6</mn><annotation-xml encoding="MathML-Content" id="S4.T2.13.3.3.m3.1c"><cn type="float" id="S4.T2.13.3.3.m3.1.1.cmml" xref="S4.T2.13.3.3.m3.1.1">13.6</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.13.3.3.m3.1d">13.6</annotation></semantics></math>, but still the models overfit at <math id="S4.T2.14.4.4.m4.1" class="ltx_Math" alttext="64" display="inline"><semantics id="S4.T2.14.4.4.m4.1b"><mn id="S4.T2.14.4.4.m4.1.1" xref="S4.T2.14.4.4.m4.1.1.cmml">64</mn><annotation-xml encoding="MathML-Content" id="S4.T2.14.4.4.m4.1c"><cn type="integer" id="S4.T2.14.4.4.m4.1.1.cmml" xref="S4.T2.14.4.4.m4.1.1">64</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.14.4.4.m4.1d">64</annotation></semantics></math>K steps with <math id="S4.T2.15.5.5.m5.1" class="ltx_Math" alttext="256" display="inline"><semantics id="S4.T2.15.5.5.m5.1b"><mn id="S4.T2.15.5.5.m5.1.1" xref="S4.T2.15.5.5.m5.1.1.cmml">256</mn><annotation-xml encoding="MathML-Content" id="S4.T2.15.5.5.m5.1c"><cn type="integer" id="S4.T2.15.5.5.m5.1.1.cmml" xref="S4.T2.15.5.5.m5.1.1">256</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.15.5.5.m5.1d">256</annotation></semantics></math> paths. In contrast, there is no overfitting at <math id="S4.T2.16.6.6.m6.1" class="ltx_Math" alttext="256" display="inline"><semantics id="S4.T2.16.6.6.m6.1b"><mn id="S4.T2.16.6.6.m6.1.1" xref="S4.T2.16.6.6.m6.1.1.cmml">256</mn><annotation-xml encoding="MathML-Content" id="S4.T2.16.6.6.m6.1c"><cn type="integer" id="S4.T2.16.6.6.m6.1.1.cmml" xref="S4.T2.16.6.6.m6.1.1">256</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.16.6.6.m6.1d">256</annotation></semantics></math> paths with overlapping shards with a <math id="S4.T2.17.7.7.m7.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S4.T2.17.7.7.m7.1b"><mrow id="S4.T2.17.7.7.m7.1.1" xref="S4.T2.17.7.7.m7.1.1.cmml"><mn id="S4.T2.17.7.7.m7.1.1.2" xref="S4.T2.17.7.7.m7.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.T2.17.7.7.m7.1.1.1" xref="S4.T2.17.7.7.m7.1.1.1.cmml">×</mo><mn id="S4.T2.17.7.7.m7.1.1.3" xref="S4.T2.17.7.7.m7.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T2.17.7.7.m7.1c"><apply id="S4.T2.17.7.7.m7.1.1.cmml" xref="S4.T2.17.7.7.m7.1.1"><times id="S4.T2.17.7.7.m7.1.1.1.cmml" xref="S4.T2.17.7.7.m7.1.1.1"></times><cn type="integer" id="S4.T2.17.7.7.m7.1.1.2.cmml" xref="S4.T2.17.7.7.m7.1.1.2">16</cn><cn type="integer" id="S4.T2.17.7.7.m7.1.1.3.cmml" xref="S4.T2.17.7.7.m7.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T2.17.7.7.m7.1d">16\times 16</annotation></semantics></math> DiPaCo.</span></span></figcaption>
</figure>
<figure id="S4.T3" class="ltx_table">
<table id="S4.T3.16" class="ltx_tabular ltx_centering ltx_align_middle">
<tbody class="ltx_tbody">
<tr id="S4.T3.16.1.1" class="ltx_tr">
<td id="S4.T3.16.1.1.1" class="ltx_td ltx_align_center ltx_border_tt">Early Stopping</td>
<td id="S4.T3.16.1.1.2" class="ltx_td ltx_align_center ltx_border_r ltx_border_tt">Route Every</td>
<td id="S4.T3.16.1.1.3" class="ltx_td ltx_nopad_r ltx_align_right ltx_border_tt">Perplexity</td>
</tr>
<tr id="S4.T3.16.2.2" class="ltx_tr">
<td id="S4.T3.16.2.2.1" class="ltx_td ltx_align_center ltx_border_t">✗</td>
<td id="S4.T3.16.2.2.2" class="ltx_td ltx_align_center ltx_border_r ltx_border_t" rowspan="2"><span id="S4.T3.16.2.2.2.1" class="ltx_text">Once per sequence</span></td>
<td id="S4.T3.16.2.2.3" class="ltx_td ltx_nopad_r ltx_align_right ltx_border_t">12.39</td>
</tr>
<tr id="S4.T3.16.3.3" class="ltx_tr">
<td id="S4.T3.16.3.3.1" class="ltx_td ltx_align_center">✓</td>
<td id="S4.T3.16.3.3.2" class="ltx_td ltx_nopad_r ltx_align_right">12.22</td>
</tr>
<tr id="S4.T3.16.4.4" class="ltx_tr">
<td id="S4.T3.16.4.4.1" class="ltx_td ltx_align_center ltx_border_bb ltx_border_t" rowspan="4"><span id="S4.T3.16.4.4.1.1" class="ltx_text">✓</span></td>
<td id="S4.T3.16.4.4.2" class="ltx_td ltx_align_center ltx_border_r ltx_border_t">128</td>
<td id="S4.T3.16.4.4.3" class="ltx_td ltx_nopad_r ltx_align_right ltx_border_t">11.48</td>
</tr>
<tr id="S4.T3.16.5.5" class="ltx_tr">
<td id="S4.T3.16.5.5.1" class="ltx_td ltx_align_center ltx_border_r"><span id="S4.T3.16.5.5.1.1" class="ltx_text ltx_font_bold">64</span></td>
<td id="S4.T3.16.5.5.2" class="ltx_td ltx_nopad_r ltx_align_right"><span id="S4.T3.16.5.5.2.1" class="ltx_text ltx_font_bold">11.38</span></td>
</tr>
<tr id="S4.T3.16.6.6" class="ltx_tr">
<td id="S4.T3.16.6.6.1" class="ltx_td ltx_align_center ltx_border_r">32</td>
<td id="S4.T3.16.6.6.2" class="ltx_td ltx_nopad_r ltx_align_right">11.31</td>
</tr>
<tr id="S4.T3.16.7.7" class="ltx_tr">
<td id="S4.T3.16.7.7.1" class="ltx_td ltx_align_center ltx_border_bb ltx_border_r">16</td>
<td id="S4.T3.16.7.7.2" class="ltx_td ltx_nopad_r ltx_align_right ltx_border_bb">11.26</td>
</tr>
</tbody>
</table>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T3.18.8.1" class="ltx_text" style="font-size:90%;">Table 3</span>: </span><span id="S4.T3.14.7" class="ltx_text ltx_font_bold" style="font-size:90%;">Frequent Routing at Eval-Time<span id="S4.T3.14.7.7" class="ltx_text ltx_font_medium">: Validation perplexity of a <math id="S4.T3.8.1.1.m1.1" class="ltx_Math" alttext="16\times 16" display="inline"><semantics id="S4.T3.8.1.1.m1.1b"><mrow id="S4.T3.8.1.1.m1.1.1" xref="S4.T3.8.1.1.m1.1.1.cmml"><mn id="S4.T3.8.1.1.m1.1.1.2" xref="S4.T3.8.1.1.m1.1.1.2.cmml">16</mn><mo lspace="0.222em" rspace="0.222em" id="S4.T3.8.1.1.m1.1.1.1" xref="S4.T3.8.1.1.m1.1.1.1.cmml">×</mo><mn id="S4.T3.8.1.1.m1.1.1.3" xref="S4.T3.8.1.1.m1.1.1.3.cmml">16</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T3.8.1.1.m1.1c"><apply id="S4.T3.8.1.1.m1.1.1.cmml" xref="S4.T3.8.1.1.m1.1.1"><times id="S4.T3.8.1.1.m1.1.1.1.cmml" xref="S4.T3.8.1.1.m1.1.1.1"></times><cn type="integer" id="S4.T3.8.1.1.m1.1.1.2.cmml" xref="S4.T3.8.1.1.m1.1.1.2">16</cn><cn type="integer" id="S4.T3.8.1.1.m1.1.1.3.cmml" xref="S4.T3.8.1.1.m1.1.1.3">16</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.8.1.1.m1.1d">16\times 16</annotation></semantics></math> DiPaCo with <math id="S4.T3.9.2.2.m2.1" class="ltx_Math" alttext="P=256" display="inline"><semantics id="S4.T3.9.2.2.m2.1b"><mrow id="S4.T3.9.2.2.m2.1.1" xref="S4.T3.9.2.2.m2.1.1.cmml"><mi id="S4.T3.9.2.2.m2.1.1.2" xref="S4.T3.9.2.2.m2.1.1.2.cmml">P</mi><mo id="S4.T3.9.2.2.m2.1.1.1" xref="S4.T3.9.2.2.m2.1.1.1.cmml">=</mo><mn id="S4.T3.9.2.2.m2.1.1.3" xref="S4.T3.9.2.2.m2.1.1.3.cmml">256</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.T3.9.2.2.m2.1c"><apply id="S4.T3.9.2.2.m2.1.1.cmml" xref="S4.T3.9.2.2.m2.1.1"><eq id="S4.T3.9.2.2.m2.1.1.1.cmml" xref="S4.T3.9.2.2.m2.1.1.1"></eq><ci id="S4.T3.9.2.2.m2.1.1.2.cmml" xref="S4.T3.9.2.2.m2.1.1.2">𝑃</ci><cn type="integer" id="S4.T3.9.2.2.m2.1.1.3.cmml" xref="S4.T3.9.2.2.m2.1.1.3">256</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.9.2.2.m2.1d">P=256</annotation></semantics></math>. Despite using paths of size <math id="S4.T3.10.3.3.m3.1" class="ltx_Math" alttext="150" display="inline"><semantics id="S4.T3.10.3.3.m3.1b"><mn id="S4.T3.10.3.3.m3.1.1" xref="S4.T3.10.3.3.m3.1.1.cmml">150</mn><annotation-xml encoding="MathML-Content" id="S4.T3.10.3.3.m3.1c"><cn type="integer" id="S4.T3.10.3.3.m3.1.1.cmml" xref="S4.T3.10.3.3.m3.1.1">150</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.10.3.3.m3.1d">150</annotation></semantics></math>M (<math id="S4.T3.11.4.4.m4.1" class="ltx_math_unparsed" alttext="&gt;6\times" display="inline"><semantics id="S4.T3.11.4.4.m4.1b"><mrow id="S4.T3.11.4.4.m4.1c"><mo id="S4.T3.11.4.4.m4.1.1">&gt;</mo><mn id="S4.T3.11.4.4.m4.1.2">6</mn><mo lspace="0.222em" id="S4.T3.11.4.4.m4.1.3">×</mo></mrow><annotation encoding="application/x-tex" id="S4.T3.11.4.4.m4.1d">&gt;6\times</annotation></semantics></math> fewer parameters), we match the performance of a dense <math id="S4.T3.12.5.5.m5.1" class="ltx_Math" alttext="1" display="inline"><semantics id="S4.T3.12.5.5.m5.1b"><mn id="S4.T3.12.5.5.m5.1.1" xref="S4.T3.12.5.5.m5.1.1.cmml">1</mn><annotation-xml encoding="MathML-Content" id="S4.T3.12.5.5.m5.1c"><cn type="integer" id="S4.T3.12.5.5.m5.1.1.cmml" xref="S4.T3.12.5.5.m5.1.1">1</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.12.5.5.m5.1d">1</annotation></semantics></math>B model (<math id="S4.T3.13.6.6.m6.1" class="ltx_Math" alttext="11.41" display="inline"><semantics id="S4.T3.13.6.6.m6.1b"><mn id="S4.T3.13.6.6.m6.1.1" xref="S4.T3.13.6.6.m6.1.1.cmml">11.41</mn><annotation-xml encoding="MathML-Content" id="S4.T3.13.6.6.m6.1c"><cn type="float" id="S4.T3.13.6.6.m6.1.1.cmml" xref="S4.T3.13.6.6.m6.1.1">11.41</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.13.6.6.m6.1d">11.41</annotation></semantics></math>) by potentially re-routing to a different path every <math id="S4.T3.14.7.7.m7.1" class="ltx_Math" alttext="64" display="inline"><semantics id="S4.T3.14.7.7.m7.1b"><mn id="S4.T3.14.7.7.m7.1.1" xref="S4.T3.14.7.7.m7.1.1.cmml">64</mn><annotation-xml encoding="MathML-Content" id="S4.T3.14.7.7.m7.1c"><cn type="integer" id="S4.T3.14.7.7.m7.1.1.cmml" xref="S4.T3.14.7.7.m7.1.1">64</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.14.7.7.m7.1d">64</annotation></semantics></math> tokens, although the model may choose to route to the exact same path that it selected previously.</span></span></figcaption>
</figure>
</section>
<section id="S4.SS5" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.5 </span>Parameter synchronization with DiLoCo vs. true gradients</h3>

<div id="S4.SS5.p1" class="ltx_para">
<p id="S4.SS5.p1.6" class="ltx_p">Finally, we ablate the effect of using the partially synchronous optimization DiLoCo to understand whether we lose performance by communicating less frequently at training time. In the setting where all devices are actually collocated and there is no constraint in terms of communication, we can train the DiPaCo model in a fully synchronous manner without using DiLoCo: At every step, each path computes gradients on <span id="S4.SS5.p1.6.1" class="ltx_text ltx_font_bold">its own batch of data</span> from its own data shard; Gradients across all paths are then exchanged and aggregated module by module; Finally, the model performs one step of AdamW update with the aggregated gradient. To our surprise, we did not observe a significant performance gap between the fully synchronous training setting and the partially synchronous training setting. More specifically, DiPaCo trained with DiLoCo slightly outperforms their fully-synchronously-trained version by <math id="S4.SS5.p1.1.m1.1" class="ltx_Math" alttext="0.3" display="inline"><semantics id="S4.SS5.p1.1.m1.1a"><mn id="S4.SS5.p1.1.m1.1.1" xref="S4.SS5.p1.1.m1.1.1.cmml">0.3</mn><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.1.m1.1b"><cn type="float" id="S4.SS5.p1.1.m1.1.1.cmml" xref="S4.SS5.p1.1.m1.1.1">0.3</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.1.m1.1c">0.3</annotation></semantics></math> and <math id="S4.SS5.p1.2.m2.1" class="ltx_Math" alttext="0.6" display="inline"><semantics id="S4.SS5.p1.2.m2.1a"><mn id="S4.SS5.p1.2.m2.1.1" xref="S4.SS5.p1.2.m2.1.1.cmml">0.6</mn><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.2.m2.1b"><cn type="float" id="S4.SS5.p1.2.m2.1.1.cmml" xref="S4.SS5.p1.2.m2.1.1">0.6</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.2.m2.1c">0.6</annotation></semantics></math> perplexity points when using a <math id="S4.SS5.p1.3.m3.1" class="ltx_Math" alttext="2\times 2" display="inline"><semantics id="S4.SS5.p1.3.m3.1a"><mrow id="S4.SS5.p1.3.m3.1.1" xref="S4.SS5.p1.3.m3.1.1.cmml"><mn id="S4.SS5.p1.3.m3.1.1.2" xref="S4.SS5.p1.3.m3.1.1.2.cmml">2</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS5.p1.3.m3.1.1.1" xref="S4.SS5.p1.3.m3.1.1.1.cmml">×</mo><mn id="S4.SS5.p1.3.m3.1.1.3" xref="S4.SS5.p1.3.m3.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.3.m3.1b"><apply id="S4.SS5.p1.3.m3.1.1.cmml" xref="S4.SS5.p1.3.m3.1.1"><times id="S4.SS5.p1.3.m3.1.1.1.cmml" xref="S4.SS5.p1.3.m3.1.1.1"></times><cn type="integer" id="S4.SS5.p1.3.m3.1.1.2.cmml" xref="S4.SS5.p1.3.m3.1.1.2">2</cn><cn type="integer" id="S4.SS5.p1.3.m3.1.1.3.cmml" xref="S4.SS5.p1.3.m3.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.3.m3.1c">2\times 2</annotation></semantics></math> and <math id="S4.SS5.p1.4.m4.1" class="ltx_Math" alttext="4\times 4" display="inline"><semantics id="S4.SS5.p1.4.m4.1a"><mrow id="S4.SS5.p1.4.m4.1.1" xref="S4.SS5.p1.4.m4.1.1.cmml"><mn id="S4.SS5.p1.4.m4.1.1.2" xref="S4.SS5.p1.4.m4.1.1.2.cmml">4</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS5.p1.4.m4.1.1.1" xref="S4.SS5.p1.4.m4.1.1.1.cmml">×</mo><mn id="S4.SS5.p1.4.m4.1.1.3" xref="S4.SS5.p1.4.m4.1.1.3.cmml">4</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.4.m4.1b"><apply id="S4.SS5.p1.4.m4.1.1.cmml" xref="S4.SS5.p1.4.m4.1.1"><times id="S4.SS5.p1.4.m4.1.1.1.cmml" xref="S4.SS5.p1.4.m4.1.1.1"></times><cn type="integer" id="S4.SS5.p1.4.m4.1.1.2.cmml" xref="S4.SS5.p1.4.m4.1.1.2">4</cn><cn type="integer" id="S4.SS5.p1.4.m4.1.1.3.cmml" xref="S4.SS5.p1.4.m4.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.4.m4.1c">4\times 4</annotation></semantics></math> architecture, respectively. At <math id="S4.SS5.p1.5.m5.1" class="ltx_Math" alttext="8\times 8" display="inline"><semantics id="S4.SS5.p1.5.m5.1a"><mrow id="S4.SS5.p1.5.m5.1.1" xref="S4.SS5.p1.5.m5.1.1.cmml"><mn id="S4.SS5.p1.5.m5.1.1.2" xref="S4.SS5.p1.5.m5.1.1.2.cmml">8</mn><mo lspace="0.222em" rspace="0.222em" id="S4.SS5.p1.5.m5.1.1.1" xref="S4.SS5.p1.5.m5.1.1.1.cmml">×</mo><mn id="S4.SS5.p1.5.m5.1.1.3" xref="S4.SS5.p1.5.m5.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.5.m5.1b"><apply id="S4.SS5.p1.5.m5.1.1.cmml" xref="S4.SS5.p1.5.m5.1.1"><times id="S4.SS5.p1.5.m5.1.1.1.cmml" xref="S4.SS5.p1.5.m5.1.1.1"></times><cn type="integer" id="S4.SS5.p1.5.m5.1.1.2.cmml" xref="S4.SS5.p1.5.m5.1.1.2">8</cn><cn type="integer" id="S4.SS5.p1.5.m5.1.1.3.cmml" xref="S4.SS5.p1.5.m5.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.5.m5.1c">8\times 8</annotation></semantics></math> DiPaCo trained fully synchronously reaches better perplexity by only <math id="S4.SS5.p1.6.m6.1" class="ltx_Math" alttext="0.1" display="inline"><semantics id="S4.SS5.p1.6.m6.1a"><mn id="S4.SS5.p1.6.m6.1.1" xref="S4.SS5.p1.6.m6.1.1.cmml">0.1</mn><annotation-xml encoding="MathML-Content" id="S4.SS5.p1.6.m6.1b"><cn type="float" id="S4.SS5.p1.6.m6.1.1.cmml" xref="S4.SS5.p1.6.m6.1.1">0.1</cn></annotation-xml><annotation encoding="application/x-tex" id="S4.SS5.p1.6.m6.1c">0.1</annotation></semantics></math> perplexity despite communicating hundred of times more. This suggests that DiLoCo is an effective distributed optimization algorithm for DiPaCo.</p>
</div>
</section>
</section>
<section id="S5" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">5 </span>Related Work</h2>

<div id="S5.p1" class="ltx_para">
<p id="S5.p1.1" class="ltx_p">As mentioned in the introduction, this work shares the same motivation and intuitions expressed in Pathways <cite class="ltx_cite ltx_citemacro_citep">(Dean, <a href="#bib.bib10" title="" class="ltx_ref">2021</a>)</cite>. Unlike the Pathways framework <cite class="ltx_cite ltx_citemacro_citep">(Barham et al., <a href="#bib.bib4" title="" class="ltx_ref">2022</a>)</cite> which supports training of general modular multimodal multitask asynchronous systems, we propose a particular instantiation of a modular system that supports such kind of distributed training.
Our approach also shares motivations and intuitions with <cite class="ltx_cite ltx_citemacro_cite">Ryabinin and Gusev (<a href="#bib.bib38" title="" class="ltx_ref">2020</a>); Borzunov et al. (<a href="#bib.bib6" title="" class="ltx_ref">2022</a>)</cite>. The key difference in this work is that each worker trains a <span id="S5.p1.1.1" class="ltx_text ltx_font_italic">path</span> through modules, rather than a module.</p>
</div>
<section id="S5.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">5.1 </span>Modularity</h3>

<div id="S5.SS1.p1" class="ltx_para">
<p id="S5.SS1.p1.1" class="ltx_p">There is a large body of literature on modularity, although this is not always framed in such terms.
At one end of the spectrum there are approaches that use a very large number of very tiny modules.
For instance, RETRO <cite class="ltx_cite ltx_citemacro_citep">(Borgeaud et al., <a href="#bib.bib5" title="" class="ltx_ref">2021</a>)</cite> is one such approach, whereby there is one module per token n-gram of a large retrieval dataset, and the modules are merely vectors of biases.
At the other end of the spectrum there are approaches that use very few (e.g., two) but very large modules. For instance, Flamingo <cite class="ltx_cite ltx_citemacro_citep">(Alayrac et al., <a href="#bib.bib1" title="" class="ltx_ref">2022</a>)</cite> is a vision language model composed by two main modules, a pre-trained visual encoder and a pre-trained language model. Similarly, <cite class="ltx_cite ltx_citemacro_cite">Dalmia et al. (<a href="#bib.bib9" title="" class="ltx_ref">2023</a>)</cite> proposed another approach where modules are pre-trained language models and speech recognition encoders. The authors showed how such modules could be plugged in different ways to enable 0-shot learning of entirely new tasks, such as producing a speech recognizer that translates into another language.</p>
</div>
<div id="S5.SS1.p2" class="ltx_para">
<p id="S5.SS1.p2.1" class="ltx_p">These lines of work are representative of the vast majority of the literature on modularity.
There is however also some work where the number of modules and their size is intermediate.
For instance, <cite class="ltx_cite ltx_citemacro_cite">Purushwalkam et al. (<a href="#bib.bib32" title="" class="ltx_ref">2019</a>)</cite> showed how softly gated modules could be used in visual question answering to handle compositional tasks. In multi-task cross-lingual transfer, <cite class="ltx_cite ltx_citemacro_cite">Pfeiffer et al. (<a href="#bib.bib30" title="" class="ltx_ref">2020</a>)</cite> used instead a network where modules are very small adapters, and the vast majority of parameters are shared across paths.
We refer the reader to the survey by <cite class="ltx_cite ltx_citemacro_cite">Pfeiffer et al. (<a href="#bib.bib31" title="" class="ltx_ref">2023</a>)</cite> for additional pointers, and focus next on modularity via mixture models.</p>
</div>
</section>
<section id="S5.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">5.2 </span>Mixture of Experts</h3>

<div id="S5.SS2.p1" class="ltx_para">
<p id="S5.SS2.p1.1" class="ltx_p">Together with ensembling and boosting, mixture of experts (MoE) are an early approach to modularity in the machine learning literature <cite class="ltx_cite ltx_citemacro_citep">(Jacobs et al., <a href="#bib.bib21" title="" class="ltx_ref">1991</a>; Jordan and Jacobs, <a href="#bib.bib22" title="" class="ltx_ref">1994</a>)</cite>. Unlike ensembling, mixture models could be more efficient if an input is processed only by a few experts. Unlike boosting, training of mixture models is efficient because it can be parallelized, while in boosting the process is inherently sequential.</p>
</div>
<div id="S5.SS2.p2" class="ltx_para">
<p id="S5.SS2.p2.1" class="ltx_p">While early work considered the problem of hierarchical gating of flat mixtures <cite class="ltx_cite ltx_citemacro_citep">(Jordan and Jacobs, <a href="#bib.bib22" title="" class="ltx_ref">1994</a>)</cite>, mixtures at multiple levels appeared only more recently. For instance, <cite class="ltx_cite ltx_citemacro_cite">Eigen et al. (<a href="#bib.bib14" title="" class="ltx_ref">2014</a>)</cite> proposed a two-level mixture where paths were softly mixed and trained jointly via stochastic gradient descent. In that paper, the authors recognized the difficulty of learning non-degenerate routing functions and discussed some workarounds. The issue of how to learn a good routing function has been a major topic of research in the field since then.</p>
</div>
<div id="S5.SS2.p3" class="ltx_para">
<p id="S5.SS2.p3.1" class="ltx_p">In their seminal work, <cite class="ltx_cite ltx_citemacro_cite">Shazeer et al. (<a href="#bib.bib39" title="" class="ltx_ref">2017</a>)</cite> proposed a very large mixture of experts LSTM model for sequence modeling tasks. Most works MoE for sequence-modeling works that followed <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib25" title="" class="ltx_ref">2021</a>; Fedus et al., <a href="#bib.bib15" title="" class="ltx_ref">2021</a>; Artetxe et al., <a href="#bib.bib3" title="" class="ltx_ref">2021</a>; Clark et al., <a href="#bib.bib7" title="" class="ltx_ref">2022</a>)</cite> have used a recipe whereby FFN layers of transformers are replaced by mixtures. These MoEs operate at the token level, and training operates similarly to dense model training: all paths are synchronized at every gradient descent step.
Token-MoE are currently state of the art with respect to training FLOP efficiency, but require even more co-located accelerators than the equivalent-activated dense model for training.</p>
</div>
<div id="S5.SS2.p4" class="ltx_para">
<p id="S5.SS2.p4.1" class="ltx_p">In contrast, <cite class="ltx_cite ltx_citemacro_citep">(Gururangan et al., <a href="#bib.bib20" title="" class="ltx_ref">2023</a>)</cite> trains experts independently using a document level router; this approach had been used in computer vision by <cite class="ltx_cite ltx_citemacro_citet">Gross et al. (<a href="#bib.bib19" title="" class="ltx_ref">2017</a>)</cite>, and it also appeared in the federated learning literature <cite class="ltx_cite ltx_citemacro_citep">(Reisser et al., <a href="#bib.bib37" title="" class="ltx_ref">2021</a>)</cite>. Our flat MoE baseline is closely related to the approach in <cite class="ltx_cite ltx_citemacro_citep">(Gururangan et al., <a href="#bib.bib20" title="" class="ltx_ref">2023</a>)</cite>, except we used a discriminative router as opposed to <math id="S5.SS2.p4.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S5.SS2.p4.1.m1.1a"><mi id="S5.SS2.p4.1.m1.1.1" xref="S5.SS2.p4.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S5.SS2.p4.1.m1.1b"><ci id="S5.SS2.p4.1.m1.1.1.cmml" xref="S5.SS2.p4.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S5.SS2.p4.1.m1.1c">k</annotation></semantics></math>-Means, and we used the first 32 tokens of a document (and the features from the base Transformer LM) as the input to the router, as opposed to SVD of tf-idf vectors.</p>
</div>
<div id="S5.SS2.p5" class="ltx_para">
<p id="S5.SS2.p5.1" class="ltx_p">Finally, the idea to iteratively shard the data by domain to train models in parallel, followed by a model merging phase has also been proposed by <cite class="ltx_cite ltx_citemacro_citet">Li et al. (<a href="#bib.bib26" title="" class="ltx_ref">2022</a>)</cite>. In this work, we take this approach a step further by learning a mixture of experts model, and by automatically (and even discriminatively) sharding the data.</p>
</div>
</section>
</section>
<section id="S6" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">6 </span>Limitations</h2>

<div id="S6.p1" class="ltx_para">
<p id="S6.p1.1" class="ltx_p">The most salient limitation to DiPaCo is with respect to FLOP efficiency. In this work we made no effort to optimize its FLOP efficiency; and in the results presented, DiPaCo is significantly less FLOP efficient per evaluation perplexity than a standard dense compute optimal model.</p>
</div>
<div id="S6.p2" class="ltx_para">
<p id="S6.p2.1" class="ltx_p">This work also does not study scaling laws of DiPaCo. We work at a single path-size scale, and on one dataset.</p>
</div>
<div id="S6.p3" class="ltx_para">
<p id="S6.p3.1" class="ltx_p">We also made no effort in this work to optimize deployment of DiPaCo. In particular, if we were to naively route more frequently during deployment as in Section (<a href="#S2.SS4.SSS3" title="2.4.3 Routing More Frequently at Test Time ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.4.3</span></a>), we would need to re-compute the <math id="S6.p3.1.m1.1" class="ltx_Math" alttext="KV" display="inline"><semantics id="S6.p3.1.m1.1a"><mrow id="S6.p3.1.m1.1.1" xref="S6.p3.1.m1.1.1.cmml"><mi id="S6.p3.1.m1.1.1.2" xref="S6.p3.1.m1.1.1.2.cmml">K</mi><mo lspace="0em" rspace="0em" id="S6.p3.1.m1.1.1.1" xref="S6.p3.1.m1.1.1.1.cmml">​</mo><mi id="S6.p3.1.m1.1.1.3" xref="S6.p3.1.m1.1.1.3.cmml">V</mi></mrow><annotation-xml encoding="MathML-Content" id="S6.p3.1.m1.1b"><apply id="S6.p3.1.m1.1.1.cmml" xref="S6.p3.1.m1.1.1"><times id="S6.p3.1.m1.1.1.1.cmml" xref="S6.p3.1.m1.1.1.1"></times><ci id="S6.p3.1.m1.1.1.2.cmml" xref="S6.p3.1.m1.1.1.2">𝐾</ci><ci id="S6.p3.1.m1.1.1.3.cmml" xref="S6.p3.1.m1.1.1.3">𝑉</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S6.p3.1.m1.1c">KV</annotation></semantics></math>-cache for each query after each re-route.</p>
</div>
<div id="S6.p4" class="ltx_para">
<p id="S6.p4.1" class="ltx_p">The severity of these limitations depends on the constraints of the training and deployment infrastructure, as well as on the task at hand.
For instance, they may matter less when the computing infrastructure consists of hundreds of small and poorly connected islands of devices for tasks where sequences can be mostly associated to one (or very few) domains.
In general, addressing these limitations constitutes avenue of future research.</p>
</div>
</section>
<section id="S7" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">
<span class="ltx_tag ltx_tag_section">7 </span>Conclusions and future work</h2>

<div id="S7.p1" class="ltx_para">
<p id="S7.p1.1" class="ltx_p">In this work, we architect an ML model as a union of modules, and we define input-output mappings that are paths through these modules. We show that we can train paths (almost) independently, and (virtually) recombine these paths back into the desired large model.
In particular, we were able to show that distributed training by path can match the performance of a 1B dense language model for a similar wall clock time.</p>
</div>
<div id="S7.p2" class="ltx_para">
<p id="S7.p2.1" class="ltx_p">The design of the algorithm and architecture is guided by practical trade-offs. In our work we assume that communication is costly but compute is not. In other words, the compute that counts the most is the one of co-located devices that host a path. As long as communication between such relatively small islands of devices is limited, we can scale by distributing computation across a large number of such islands.</p>
</div>
<div id="S7.p3" class="ltx_para">
<p id="S7.p3.1" class="ltx_p">We introduced two key ingredients to operate under these constraints: 1) we route coarsely at the sequence level and pre-shard data by path, and 2) we extend DiLoCo to update parameters of modules shared by multiple paths. We believe that by combining these two ideas we have made a step towards a potentially much more scalable paradigm to train large-models.</p>
</div>
<div id="S7.p4" class="ltx_para">
<p id="S7.p4.1" class="ltx_p">There are several avenues of future work. In our investigations we made no effort to control the FLOP efficiency of DiPaCo, focusing on wall-clock efficiency. Naturally, there are several straightforward design choices that could increase FLOP efficiency, for example allowing some paths to co-locate. In addition, our approaches to sharding are fairly simple, and there are many interesting possibilities for more sophisticated sharding. DiPaCo was designed with continual learning in mind, and we would like to apply it in that setting eventually. Last but certainly not least, scaling up is of utmost interest.
Shared with  <cite class="ltx_cite ltx_citemacro_cite">Dean (<a href="#bib.bib10" title="" class="ltx_ref">2021</a>); Raffel (<a href="#bib.bib33" title="" class="ltx_ref">2023a</a>)</cite>, our long-term dream is to further refine this approach and produce a never-ending, community-driven, modular learning system that can be used by everyone to compose new predictors out of existing modules, and thus efficiently develop entirely new models and capabilities in a positive feedback loop.</p>
</div>
<div id="S7.p5" class="ltx_para">
<span id="S7.p5.1" class="ltx_ERROR undefined">\nobibliography</span>
<p id="S7.p5.2" class="ltx_p">*</p>
</div>
</section>
<section id="bib" class="ltx_bibliography" lang="en">
<h2 class="ltx_title ltx_title_bibliography">References</h2>

<ul class="ltx_biblist">
<li id="bib.bib1" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Alayrac et al. (2022)</span>
<span class="ltx_bibblock">
J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, R. Ring, E. Rutherford, S. Cabi, T. Han, Z. Gong, S. Samangooei, M. Monteiro, J. Menick, S. Borgeaud, A. Brock, A. Nematzadeh, S. Sharifzadeh, M. Binkowski, R. Barreira, O. Vinyals, A. Zisserman, and K. Simonyan.

</span>
<span class="ltx_bibblock">Flamingo: a visual language model for few-shot learning.

</span>
<span class="ltx_bibblock">In <em id="bib.bib1.1.1" class="ltx_emph ltx_font_italic">NeurIPS</em>, 2022.

</span>
</li>
<li id="bib.bib2" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Anil et al. (2023)</span>
<span class="ltx_bibblock">
R. Anil, S. Borgeaud, Y. Wu, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, K. Millican, D. Silver, S. Petrov, M. Johnson, I. Antonoglou, J. Schrittwieser, A. Glaese, J. Chen, E. Pitler, T. Lillicrap, A. Lazaridou, O. Firat, J. Molloy, M. Isard, P. R. Barham, T. Hennigan, B. Lee, F. Viola, M. Reynolds, Y. Xu, R. Doherty, E. Collins, C. Meyer, E. Rutherford, E. Moreira, K. Ayoub, M. Goel, G. Tucker, E. Piqueras, M. Krikun, I. Barr, N. Savinov, I. Danihelka, B. Roelofs, A. White, A. Andreassen, T. von Glehn, L. Yagati, M. Kazemi, L. Gonzalez, M. Khalman, J. Sygnowski, A. Frechette, C. Smith, L. Culp, L. Proleev, Y. Luan, X. Chen, J. Lottes, N. Schucher, F. Lebron, A. Rrustemi, N. Clay, P. Crone, T. Kocisky, J. Zhao, B. Perz, D. Yu, H. Howard, A. Bloniarz, J. W. Rae, H. Lu, L. Sifre, M. Maggioni, F. Alcober, D. Garrette, M. Barnes, S. Thakoor, J. Austin, G. Barth-Maron, W. Wong, R. Joshi, R. Chaabouni, D. Fatiha, A. Ahuja, R. Liu, Y. Li, S. Cogan, J. Chen, C. Jia, C. Gu, Q. Zhang,
J. Grimstad, A. J. Hartman, M. Chadwick, G. S. Tomar, X. Garcia, E. Senter, E. Taropa, T. S. Pillai, J. Devlin, M. Laskin, D. de Las Casas, D. Valter, C. Tao, L. Blanco, A. P. Badia, D. Reitter, M. Chen, J. Brennan, C. Rivera, S. Brin, S. Iqbal, G. Surita, J. Labanowski, A. Rao, S. Winkler, E. Parisotto, Y. Gu, K. Olszewska, Y. Zhang, R. Addanki, A. Miech, A. Louis, L. E. Shafey, D. Teplyashin, G. Brown, E. Catt, N. Attaluri, J. Balaguer, J. Xiang, P. Wang, Z. Ashwood, A. Briukhov, A. Webson, S. Ganapathy, S. Sanghavi, A. Kannan, M.-W. Chang, A. Stjerngren, J. Djolonga, Y. Sun, A. Bapna, M. Aitchison, P. Pejman, H. Michalewski, T. Yu, C. Wang, J. Love, J. Ahn, D. Bloxwich, K. Han, P. Humphreys, T. Sellam, J. Bradbury, V. Godbole, S. Samangooei, B. Damoc, A. Kaskasoli, S. M. R. Arnold, V. Vasudevan, S. Agrawal, J. Riesa, D. Lepikhin, R. Tanburn, S. Srinivasan, H. Lim, S. Hodkinson, P. Shyam, J. Ferret, S. Hand, A. Garg, T. L. Paine, J. Li, Y. Li, M. Giang, A. Neitz, Z. Abbas, S. York, M. Reid, E. Cole,
A. Chowdhery, D. Das, D. Rogozińska, V. Nikolaev, P. Sprechmann, Z. Nado, L. Zilka, F. Prost, L. He, M. Monteiro, G. Mishra, C. Welty, J. Newlan, D. Jia, M. Allamanis, C. H. Hu, R. de Liedekerke, J. Gilmer, C. Saroufim, S. Rijhwani, S. Hou, D. Shrivastava, A. Baddepudi, A. Goldin, A. Ozturel, A. Cassirer, Y. Xu, D. Sohn, D. Sachan, R. K. Amplayo, C. Swanson, D. Petrova, S. Narayan, A. Guez, S. Brahma, J. Landon, M. Patel, R. Zhao, K. Villela, L. Wang, W. Jia, M. Rahtz, M. Giménez, L. Yeung, H. Lin, J. Keeling, P. Georgiev, D. Mincu, B. Wu, S. Haykal, R. Saputro, K. Vodrahalli, J. Qin, Z. Cankara, A. Sharma, N. Fernando, W. Hawkins, B. Neyshabur, S. Kim, A. Hutter, P. Agrawal, A. Castro-Ros, G. van den Driessche, T. Wang, F. Yang, S. yiin Chang, P. Komarek, R. McIlroy, M. Lučić, G. Zhang, W. Farhan, M. Sharman, P. Natsev, P. Michel, Y. Cheng, Y. Bansal, S. Qiao, K. Cao, S. Shakeri, C. Butterfield, J. Chung, P. K. Rubenstein, S. Agrawal, A. Mensch, K. Soparkar, K. Lenc, T. Chung, A. Pope, L. Maggiore,
J. Kay, P. Jhakra, S. Wang, J. Maynez, M. Phuong, T. Tobin, A. Tacchetti, M. Trebacz, K. Robinson, Y. Katariya, S. Riedel, P. Bailey, K. Xiao, N. Ghelani, L. Aroyo, A. Slone, N. Houlsby, X. Xiong, Z. Yang, E. Gribovskaya, J. Adler, M. Wirth, L. Lee, M. Li, T. Kagohara, J. Pavagadhi, S. Bridgers, A. Bortsova, S. Ghemawat, Z. Ahmed, T. Liu, R. Powell, V. Bolina, M. Iinuma, P. Zablotskaia, J. Besley, D.-W. Chung, T. Dozat, R. Comanescu, X. Si, J. Greer, G. Su, M. Polacek, R. L. Kaufman, S. Tokumine, H. Hu, E. Buchatskaya, Y. Miao, M. Elhawaty, A. Siddhant, N. Tomasev, J. Xing, C. Greer, H. Miller, S. Ashraf, A. Roy, Z. Zhang, A. Ma, A. Filos, M. Besta, R. Blevins, T. Klimenko, C.-K. Yeh, S. Changpinyo, J. Mu, O. Chang, M. Pajarskas, C. Muir, V. Cohen, C. L. Lan, K. Haridasan, A. Marathe, S. Hansen, S. Douglas, R. Samuel, M. Wang, S. Austin, C. Lan, J. Jiang, J. Chiu, J. A. Lorenzo, L. L. Sjösund, S. Cevey, Z. Gleicher, T. Avrahami, A. Boral, H. Srinivasan, V. Selo, R. May, K. Aisopos, L. Hussenot, L. B.
Soares, K. Baumli, M. B. Chang, A. Recasens, B. Caine, A. Pritzel, F. Pavetic, F. Pardo, A. Gergely, J. Frye, V. Ramasesh, D. Horgan, K. Badola, N. Kassner, S. Roy, E. Dyer, V. Campos, A. Tomala, Y. Tang, D. E. Badawy, E. White, B. Mustafa, O. Lang, A. Jindal, S. Vikram, Z. Gong, S. Caelles, R. Hemsley, G. Thornton, F. Feng, W. Stokowiec, C. Zheng, P. Thacker, Çağlar Ünlü, Z. Zhang, M. Saleh, J. Svensson, M. Bileschi, P. Patil, A. Anand, R. Ring, K. Tsihlas, A. Vezer, M. Selvi, T. Shevlane, M. Rodriguez, T. Kwiatkowski, S. Daruki, K. Rong, A. Dafoe, N. FitzGerald, K. Gu-Lemberg, M. Khan, L. A. Hendricks, M. Pellat, V. Feinberg, J. Cobon-Kerr, T. Sainath, M. Rauh, S. H. Hashemi, R. Ives, Y. Hasson, Y. Li, E. Noland, Y. Cao, N. Byrd, L. Hou, Q. Wang, T. Sottiaux, M. Paganini, J.-B. Lespiau, A. Moufarek, S. Hassan, K. Shivakumar, J. van Amersfoort, A. Mandhane, P. Joshi, A. Goyal, M. Tung, A. Brock, H. Sheahan, V. Misra, C. Li, N. Rakićević, M. Dehghani, F. Liu, S. Mittal, J. Oh, S. Noury, E. Sezener,
F. Huot, M. Lamm, N. D. Cao, C. Chen, G. Elsayed, E. Chi, M. Mahdieh, I. Tenney, N. Hua, I. Petrychenko, P. Kane, D. Scandinaro, R. Jain, J. Uesato, R. Datta, A. Sadovsky, O. Bunyan, D. Rabiej, S. Wu, J. Zhang, G. Vasudevan, E. Leurent, M. Alnahlawi, I. Georgescu, N. Wei, I. Zheng, B. Chan, P. G. Rabinovitch, P. Stanczyk, Y. Zhang, D. Steiner, S. Naskar, M. Azzam, M. Johnson, A. Paszke, C.-C. Chiu, J. S. Elias, A. Mohiuddin, F. Muhammad, J. Miao, A. Lee, N. Vieillard, S. Potluri, J. Park, E. Davoodi, J. Zhang, J. Stanway, D. Garmon, A. Karmarkar, Z. Dong, J. Lee, A. Kumar, L. Zhou, J. Evens, W. Isaac, Z. Chen, J. Jia, A. Levskaya, Z. Zhu, C. Gorgolewski, P. Grabowski, Y. Mao, A. Magni, K. Yao, J. Snaider, N. Casagrande, P. Suganthan, E. Palmer, G. Irving, E. Loper, M. Faruqui, I. Arkatkar, N. Chen, I. Shafran, M. Fink, A. Castaño, I. Giannoumis, W. Kim, M. Rybiński, A. Sreevatsa, J. Prendki, D. Soergel, A. Goedeckemeyer, W. Gierke, M. Jafari, M. Gaba, J. Wiesner, D. G. Wright, Y. Wei, H. Vashisht,
Y. Kulizhskaya, J. Hoover, M. Le, L. Li, C. Iwuanyanwu, L. Liu, K. Ramirez, A. Khorlin, A. Cui, T. LIN, M. Georgiev, M. Wu, R. Aguilar, K. Pallo, A. Chakladar, A. Repina, X. Wu, T. van der Weide, P. Ponnapalli, C. Kaplan, J. Simsa, S. Li, O. Dousse, F. Yang, J. Piper, N. Ie, M. Lui, R. Pasumarthi, N. Lintz, A. Vijayakumar, L. N. Thiet, D. Andor, P. Valenzuela, C. Paduraru, D. Peng, K. Lee, S. Zhang, S. Greene, D. D. Nguyen, P. Kurylowicz, S. Velury, S. Krause, C. Hardin, L. Dixon, L. Janzer, K. Choo, Z. Feng, B. Zhang, A. Singhal, T. Latkar, M. Zhang, Q. Le, E. A. Abellan, D. Du, D. McKinnon, N. Antropova, T. Bolukbasi, O. Keller, D. Reid, D. Finchelstein, M. A. Raad, R. Crocker, P. Hawkins, R. Dadashi, C. Gaffney, S. Lall, K. Franko, E. Filonov, A. Bulanova, R. Leblond, V. Yadav, S. Chung, H. Askham, L. C. Cobo, K. Xu, F. Fischer, J. Xu, C. Sorokin, C. Alberti, C.-C. Lin, C. Evans, H. Zhou, A. Dimitriev, H. Forbes, D. Banarse, Z. Tung, J. Liu, M. Omernick, C. Bishop, C. Kumar, R. Sterneck, R. Foley,
R. Jain, S. Mishra, J. Xia, T. Bos, G. Cideron, E. Amid, F. Piccinno, X. Wang, P. Banzal, P. Gurita, H. Noga, P. Shah, D. J. Mankowitz, A. Polozov, N. Kushman, V. Krakovna, S. Brown, M. Bateni, D. Duan, V. Firoiu, M. Thotakuri, T. Natan, A. Mohananey, M. Geist, S. Mudgal, S. Girgin, H. Li, J. Ye, O. Roval, R. Tojo, M. Kwong, J. Lee-Thorp, C. Yew, Q. Yuan, S. Bagri, D. Sinopalnikov, S. Ramos, J. Mellor, A. Sharma, A. Severyn, J. Lai, K. Wu, H.-T. Cheng, D. Miller, N. Sonnerat, D. Vnukov, R. Greig, J. Beattie, E. Caveness, L. Bai, J. Eisenschlos, A. Korchemniy, T. Tsai, M. Jasarevic, W. Kong, P. Dao, Z. Zheng, F. Liu, F. Yang, R. Zhu, M. Geller, T. H. Teh, J. Sanmiya, E. Gladchenko, N. Trdin, A. Sozanschi, D. Toyama, E. Rosen, S. Tavakkol, L. Xue, C. Elkind, O. Woodman, J. Carpenter, G. Papamakarios, R. Kemp, S. Kafle, T. Grunina, R. Sinha, A. Talbert, A. Goyal, D. Wu, D. Owusu-Afriyie, C. Du, C. Thornton, J. Pont-Tuset, P. Narayana, J. Li, S. Fatehi, J. Wieting, O. Ajmeri, B. Uria, T. Zhu, Y. Ko, L. Knight,
A. Héliou, N. Niu, S. Gu, C. Pang, D. Tran, Y. Li, N. Levine, A. Stolovich, N. Kalb, R. Santamaria-Fernandez, S. Goenka, W. Yustalim, R. Strudel, A. Elqursh, B. Lakshminarayanan, C. Deck, S. Upadhyay, H. Lee, M. Dusenberry, Z. Li, X. Wang, K. Levin, R. Hoffmann, D. Holtmann-Rice, O. Bachem, S. Yue, S. Arora, E. Malmi, D. Mirylenka, Q. Tan, C. Koh, S. H. Yeganeh, S. Põder, S. Zheng, F. Pongetti, M. Tariq, Y. Sun, L. Ionita, M. Seyedhosseini, P. Tafti, R. Kotikalapudi, Z. Liu, A. Gulati, J. Liu, X. Ye, B. Chrzaszcz, L. Wang, N. Sethi, T. Li, B. Brown, S. Singh, W. Fan, A. Parisi, J. Stanton, C. Kuang, V. Koverkathu, C. A. Choquette-Choo, Y. Li, T. Lu, A. Ittycheriah, P. Shroff, P. Sun, M. Varadarajan, S. Bahargam, R. Willoughby, D. Gaddy, I. Dasgupta, G. Desjardins, M. Cornero, B. Robenek, B. Mittal, B. Albrecht, A. Shenoy, F. Moiseev, H. Jacobsson, A. Ghaffarkhah, M. Rivière, A. Walton, C. Crepy, A. Parrish, Y. Liu, Z. Zhou, C. Farabet, C. Radebaugh, P. Srinivasan, C. van der Salm, A. Fidjeland,
S. Scellato, E. Latorre-Chimoto, H. Klimczak-Plucińska, D. Bridson, D. de Cesare, T. Hudson, P. Mendolicchio, L. Walker, A. Morris, I. Penchev, M. Mauger, A. Guseynov, A. Reid, S. Odoom, L. Loher, V. Cotruta, M. Yenugula, D. Grewe, A. Petrushkina, T. Duerig, A. Sanchez, S. Yadlowsky, A. Shen, A. Globerson, A. Kurzrok, L. Webb, S. Dua, D. Li, P. Lahoti, S. Bhupatiraju, D. Hurt, H. Qureshi, A. Agarwal, T. Shani, M. Eyal, A. Khare, S. R. Belle, L. Wang, C. Tekur, M. S. Kale, J. Wei, R. Sang, B. Saeta, T. Liechty, Y. Sun, Y. Zhao, S. Lee, P. Nayak, D. Fritz, M. R. Vuyyuru, J. Aslanides, N. Vyas, M. Wicke, X. Ma, T. Bilal, E. Eltyshev, D. Balle, N. Martin, H. Cate, J. Manyika, K. Amiri, Y. Kim, X. Xiong, K. Kang, F. Luisier, N. Tripuraneni, D. Madras, M. Guo, A. Waters, O. Wang, J. Ainslie, J. Baldridge, H. Zhang, G. Pruthi, J. Bauer, F. Yang, R. Mansour, J. Gelman, Y. Xu, G. Polovets, J. Liu, H. Cai, W. Chen, X. Sheng, E. Xue, S. Ozair, A. Yu, C. Angermueller, X. Li, W. Wang, J. Wiesinger, E. Koukoumidis,
Y. Tian, A. Iyer, M. Gurumurthy, M. Goldenson, P. Shah, M. Blake, H. Yu, A. Urbanowicz, J. Palomaki, C. Fernando, K. Brooks, K. Durden, H. Mehta, N. Momchev, E. Rahimtoroghi, M. Georgaki, A. Raul, S. Ruder, M. Redshaw, J. Lee, K. Jalan, D. Li, G. Perng, B. Hechtman, P. Schuh, M. Nasr, M. Chen, K. Milan, V. Mikulik, T. Strohman, J. Franco, T. Green, D. Hassabis, K. Kavukcuoglu, J. Dean, and O. Vinyals.

</span>
<span class="ltx_bibblock">Gemini: A family of highly capable multimodal models, 2023.

</span>
</li>
<li id="bib.bib3" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Artetxe et al. (2021)</span>
<span class="ltx_bibblock">
M. Artetxe, S. Bhosale, N. Goyal, T. Mihaylov, M. Ott, S. Shleifer, X. V. Lin, J. Du, S. Iyer, R. Pasunuru, et al.

</span>
<span class="ltx_bibblock">Efficient large scale language modeling with mixtures of experts.

</span>
<span class="ltx_bibblock"><em id="bib.bib3.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:2112.10684</em>, 2021.

</span>
</li>
<li id="bib.bib4" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Barham et al. (2022)</span>
<span class="ltx_bibblock">
P. Barham, A. Chowdhery, J. Dean, S. Ghemawat, S. Hand, D. Hurt, M. Isard, H. Lim, R. Pang, S. Roy, B. Saeta, P. Schuh, R. Sepassi, L. E. Shafey, C. A. Thekkath, and Y. Wu.

</span>
<span class="ltx_bibblock">Pathways: Asynchronous distributed dataflow for ml, 2022.

</span>
</li>
<li id="bib.bib5" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Borgeaud et al. (2021)</span>
<span class="ltx_bibblock">
S. Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, G. van den Driessche, J. Lespiau, B. Damoc, A. Clark, D. de Las Casas, A. Guy, J. Menick, R. Ring, T. Hennigan, S. Huang, L. Maggiore, C. Jones, A. Cassirer, A. Brock, M. Paganini, G. Irving, O. Vinyals, S. Osindero, K. Simonyan, J. W. Rae, E. Elsen, and L. Sifre.

</span>
<span class="ltx_bibblock">Improving language models by retrieving from trillions of tokens, 2021.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://arxiv.org/abs/2112.04426" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://arxiv.org/abs/2112.04426</a>.

</span>
</li>
<li id="bib.bib6" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Borzunov et al. (2022)</span>
<span class="ltx_bibblock">
A. Borzunov, D. Baranchuk, T. Dettmers, M. Ryabinin, Y. Belkada, A. Chumachenko, P. Samygin, and C. Raffel.

</span>
<span class="ltx_bibblock">Petals: Collaborative inference and fine-tuning of large models.

</span>
<span class="ltx_bibblock"><em id="bib.bib6.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:2209.01188</em>, 2022.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://arxiv.org/abs/2209.01188" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://arxiv.org/abs/2209.01188</a>.

</span>
</li>
<li id="bib.bib7" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Clark et al. (2022)</span>
<span class="ltx_bibblock">
A. Clark, D. D. L. Casas, A. Guy, A. Mensch, M. Paganini, J. Hoffmann, B. Damoc, B. Hechtman, T. Cai, S. Borgeaud, G. B. V. D. Driessche, E. Rutherford, T. Hennigan, M. J. Johnson, A. Cassirer, C. Jones, E. Buchatskaya, D. Budden, L. Sifre, S. Osindero, O. Vinyals, M. Ranzato, J. Rae, E. Elsen, K. Kavukcuoglu, and K. Simonyan.

</span>
<span class="ltx_bibblock">Unified scaling laws for routed language models.

</span>
<span class="ltx_bibblock">In <em id="bib.bib7.1.1" class="ltx_emph ltx_font_italic">Conference International Conference on Machine Learning</em>, 2022.

</span>
</li>
<li id="bib.bib8" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Corbett et al. (2012)</span>
<span class="ltx_bibblock">
J. C. Corbett, J. Dean, M. Epstein, A. Fikes, C. Frost, J. Furman, S. Ghemawat, A. Gubarev, C. Heiser, P. Hochschild, W. Hsieh, S. Kanthak, E. Kogan, H. Li, A. Lloyd, S. Melnik, D. Mwaura, D. Nagle, S. Quinlan, R. Rao, L. Rolig, Y. Saito, M. Szymaniak, C. Taylor, R. Wang, and D. Woodford.

</span>
<span class="ltx_bibblock">Spanner: Google’s globally-distributed database.

</span>
<span class="ltx_bibblock">In <em id="bib.bib8.1.1" class="ltx_emph ltx_font_italic">OSDI</em>, 2012.

</span>
</li>
<li id="bib.bib9" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Dalmia et al. (2023)</span>
<span class="ltx_bibblock">
S. Dalmia, D. Okhonko, M. Lewis, S. Edunov, S. Watanabe, F. Metze, L. Zettlemoyer, and A. Mohamed.

</span>
<span class="ltx_bibblock">Legonn: Building modular encoder-decoder models, 2023.

</span>
</li>
<li id="bib.bib10" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Dean (2021)</span>
<span class="ltx_bibblock">
J. Dean.

</span>
<span class="ltx_bibblock">Introducing pathways: A next-generation ai architecture.

</span>
<span class="ltx_bibblock"><a target="_blank" href="https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/</a>, 2021.

</span>
<span class="ltx_bibblock">Google blog post.

</span>
</li>
<li id="bib.bib11" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Dean et al. (2012)</span>
<span class="ltx_bibblock">
J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, M. Mao, M. a. Ranzato, A. Senior, P. Tucker, K. Yang, Q. Le, and A. Ng.

</span>
<span class="ltx_bibblock">Large scale distributed deep networks.

</span>
<span class="ltx_bibblock">In F. Pereira, C. Burges, L. Bottou, and K. Weinberger, editors, <em id="bib.bib11.1.1" class="ltx_emph ltx_font_italic">Advances in Neural Information Processing Systems</em>, volume 25. Curran Associates, Inc., 2012.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://proceedings.neurips.cc/paper_files/paper/2012/file/6aca97005c68f1206823815f66102863-Paper.pdf" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://proceedings.neurips.cc/paper_files/paper/2012/file/6aca97005c68f1206823815f66102863-Paper.pdf</a>.

</span>
</li>
<li id="bib.bib12" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Dempster et al. (1977)</span>
<span class="ltx_bibblock">
A. P. Dempster, N. M. Laird, and D. B. Rubin.

</span>
<span class="ltx_bibblock">Maximum likelihood from incomplete data via the em algorithm.

</span>
<span class="ltx_bibblock"><em id="bib.bib12.1.1" class="ltx_emph ltx_font_italic">Journal of the royal statistical society. Series B (methodological)</em>, pages 1–38, 1977.

</span>
</li>
<li id="bib.bib13" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Douillard et al. (2023)</span>
<span class="ltx_bibblock">
A. Douillard, Q. Feng, A. A. Rusu, R. Chhaparia, Y. Donchev, A. Kuncoro, M. Ranzato, A. Szlam, and J. Shen.

</span>
<span class="ltx_bibblock">Diloco: Distributed low-communication training of language models, 2023.

</span>
</li>
<li id="bib.bib14" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Eigen et al. (2014)</span>
<span class="ltx_bibblock">
D. Eigen, I. Sutskever, and M. Ranzato.

</span>
<span class="ltx_bibblock">Learning factored representations in a deep mixture of experts.

</span>
<span class="ltx_bibblock">In <em id="bib.bib14.1.1" class="ltx_emph ltx_font_italic">Workshop at the International Conference on Learning Representations</em>, 2014.

</span>
</li>
<li id="bib.bib15" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Fedus et al. (2021)</span>
<span class="ltx_bibblock">
W. Fedus, B. Zoph, and N. Shazeer.

</span>
<span class="ltx_bibblock">Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.

</span>
<span class="ltx_bibblock"><em id="bib.bib15.1.1" class="ltx_emph ltx_font_italic">CoRR</em>, abs/2101.03961, 2021.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://arxiv.org/abs/2101.03961" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://arxiv.org/abs/2101.03961</a>.

</span>
</li>
<li id="bib.bib16" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Ghemawat et al. (2003)</span>
<span class="ltx_bibblock">
S. Ghemawat, H. Gobioff, and S.-T. Leung.

</span>
<span class="ltx_bibblock">The google file system.

</span>
<span class="ltx_bibblock">In <em id="bib.bib16.1.1" class="ltx_emph ltx_font_italic">SOSP</em>, 2003.

</span>
</li>
<li id="bib.bib17" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Google (2023)</span>
<span class="ltx_bibblock">
Google.

</span>
<span class="ltx_bibblock">Effingo: the internal google copy service moving data at scalen.

</span>
<span class="ltx_bibblock"><a target="_blank" href="https://cloud.google.com/blog/products/storage-data-transfer/inside-googles-internal-effingo-data-copy-service/" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://cloud.google.com/blog/products/storage-data-transfer/inside-googles-internal-effingo-data-copy-service/</a>, 2023.

</span>
<span class="ltx_bibblock">Accessed: 2014-02-08.

</span>
</li>
<li id="bib.bib18" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Goyal et al. (2017)</span>
<span class="ltx_bibblock">
P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He.

</span>
<span class="ltx_bibblock">Accurate, large minibatch sgd: Training imagenet in 1 hour.

</span>
<span class="ltx_bibblock"><em id="bib.bib18.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1706.02677</em>, 2017.

</span>
</li>
<li id="bib.bib19" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Gross et al. (2017)</span>
<span class="ltx_bibblock">
S. Gross, M. Ranzato, and A. Szlam.

</span>
<span class="ltx_bibblock">Hard mixtures of experts for large scale weakly supervised vision.

</span>
<span class="ltx_bibblock">In <em id="bib.bib19.1.1" class="ltx_emph ltx_font_italic">CVPR</em>, 2017.

</span>
</li>
<li id="bib.bib20" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Gururangan et al. (2023)</span>
<span class="ltx_bibblock">
S. Gururangan, M. Li, M. Lewis, W. Shi, T. Althoff, N. A. Smith, and L. Zettlemoyer.

</span>
<span class="ltx_bibblock">Scaling expert language models with unsupervised domain discovery.

</span>
<span class="ltx_bibblock"><em id="bib.bib20.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:2303.14177</em>, 2023.

</span>
</li>
<li id="bib.bib21" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Jacobs et al. (1991)</span>
<span class="ltx_bibblock">
R. A. Jacobs, M. I. Jordan, S. Nowlan, and G. E. Hinton.

</span>
<span class="ltx_bibblock">Adaptive mixtures of local experts.

</span>
<span class="ltx_bibblock"><em id="bib.bib21.1.1" class="ltx_emph ltx_font_italic">Neural Computation</em>, 3:1––12, 1991.

</span>
</li>
<li id="bib.bib22" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Jordan and Jacobs (1994)</span>
<span class="ltx_bibblock">
M. I. Jordan and R. A. Jacobs.

</span>
<span class="ltx_bibblock">Hierarchical mixtures of experts and the em algorithm.

</span>
<span class="ltx_bibblock"><em id="bib.bib22.1.1" class="ltx_emph ltx_font_italic">Neural Computation</em>, 6:181–214, 1994.

</span>
</li>
<li id="bib.bib23" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Kingma and Ba (2014)</span>
<span class="ltx_bibblock">
D. P. Kingma and J. Ba.

</span>
<span class="ltx_bibblock">Adam: A method for stochastic optimization.

</span>
<span class="ltx_bibblock"><em id="bib.bib23.1.1" class="ltx_emph ltx_font_italic">Proceedings of the 3rd International Conference on Learning Representations (ICLR)</em>, 2014.

</span>
</li>
<li id="bib.bib24" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Kudo and Richardson (2018)</span>
<span class="ltx_bibblock">
T. Kudo and J. Richardson.

</span>
<span class="ltx_bibblock">Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing.

</span>
<span class="ltx_bibblock">In <em id="bib.bib24.1.1" class="ltx_emph ltx_font_italic">EMNLP</em>, 2018.

</span>
</li>
<li id="bib.bib25" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Lepikhin et al. (2021)</span>
<span class="ltx_bibblock">
D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen.

</span>
<span class="ltx_bibblock">{GS}hard: Scaling giant models with conditional computation and automatic sharding.

</span>
<span class="ltx_bibblock">In <em id="bib.bib25.1.1" class="ltx_emph ltx_font_italic">International Conference on Learning Representations</em>, 2021.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://openreview.net/forum?id=qrwe7XHTmYb" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://openreview.net/forum?id=qrwe7XHTmYb</a>.

</span>
</li>
<li id="bib.bib26" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Li et al. (2022)</span>
<span class="ltx_bibblock">
M. Li, S. Gururangan, T. Dettmers, M. Lewis, T. Althoff, N. A. Smith, and L. Zettlemoyer.

</span>
<span class="ltx_bibblock">Branch-train-merge: Embarrassingly parallel training of expert language models.

</span>
<span class="ltx_bibblock">In <em id="bib.bib26.1.1" class="ltx_emph ltx_font_italic">First Workshop on Interpolation Regularizers and Beyond at NeurIPS 2022</em>, 2022.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://openreview.net/forum?id=SQgVgE2Sq4" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://openreview.net/forum?id=SQgVgE2Sq4</a>.

</span>
</li>
<li id="bib.bib27" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Liu et al. (2024)</span>
<span class="ltx_bibblock">
B. Liu, R. Chhaparia, A. Douillard, S. Kale, A. A. Rusu, J. Shen, A. Szlam, and M. Ranzato.

</span>
<span class="ltx_bibblock">Asynchronous local-sgd training for language modeling, 2024.

</span>
</li>
<li id="bib.bib28" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Narayanan et al. (2020)</span>
<span class="ltx_bibblock">
D. Narayanan, A. Phanishayee, K. Shi, X. Chen, and M. Zaharia.

</span>
<span class="ltx_bibblock">Memory-efficient pipeline-parallel DNN training.

</span>
<span class="ltx_bibblock"><em id="bib.bib28.1.1" class="ltx_emph ltx_font_italic">CoRR</em>, abs/2006.09503, 2020.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://arxiv.org/abs/2006.09503" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://arxiv.org/abs/2006.09503</a>.

</span>
</li>
<li id="bib.bib29" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">OpenAI et al. (2023)</span>
<span class="ltx_bibblock">
OpenAI, :, J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, R. Avila, I. Babuschkin, S. Balaji, V. Balcom, P. Baltescu, H. Bao, M. Bavarian, J. Belgum, I. Bello, J. Berdine, G. Bernadett-Shapiro, C. Berner, L. Bogdonoff, O. Boiko, M. Boyd, A.-L. Brakman, G. Brockman, T. Brooks, M. Brundage, K. Button, T. Cai, R. Campbell, A. Cann, B. Carey, C. Carlson, R. Carmichael, B. Chan, C. Chang, F. Chantzis, D. Chen, S. Chen, R. Chen, J. Chen, M. Chen, B. Chess, C. Cho, C. Chu, H. W. Chung, D. Cummings, J. Currier, Y. Dai, C. Decareaux, T. Degry, N. Deutsch, D. Deville, A. Dhar, D. Dohan, S. Dowling, S. Dunning, A. Ecoffet, A. Eleti, T. Eloundou, D. Farhi, L. Fedus, N. Felix, S. P. Fishman, J. Forte, I. Fulford, L. Gao, E. Georges, C. Gibson, V. Goel, T. Gogineni, G. Goh, R. Gontijo-Lopes, J. Gordon, M. Grafstein, S. Gray, R. Greene, J. Gross, S. S. Gu, Y. Guo, C. Hallacy, J. Han, J. Harris, Y. He, M. Heaton, J. Heidecke, C. Hesse, A. Hickey,
W. Hickey, P. Hoeschele, B. Houghton, K. Hsu, S. Hu, X. Hu, J. Huizinga, S. Jain, S. Jain, J. Jang, A. Jiang, R. Jiang, H. Jin, D. Jin, S. Jomoto, B. Jonn, H. Jun, T. Kaftan, Łukasz Kaiser, A. Kamali, I. Kanitscheider, N. S. Keskar, T. Khan, L. Kilpatrick, J. W. Kim, C. Kim, Y. Kim, H. Kirchner, J. Kiros, M. Knight, D. Kokotajlo, Łukasz Kondraciuk, A. Kondrich, A. Konstantinidis, K. Kosic, G. Krueger, V. Kuo, M. Lampe, I. Lan, T. Lee, J. Leike, J. Leung, D. Levy, C. M. Li, R. Lim, M. Lin, S. Lin, M. Litwin, T. Lopez, R. Lowe, P. Lue, A. Makanju, K. Malfacini, S. Manning, T. Markov, Y. Markovski, B. Martin, K. Mayer, A. Mayne, B. McGrew, S. M. McKinney, C. McLeavey, P. McMillan, J. McNeil, D. Medina, A. Mehta, J. Menick, L. Metz, A. Mishchenko, P. Mishkin, V. Monaco, E. Morikawa, D. Mossing, T. Mu, M. Murati, O. Murk, D. Mély, A. Nair, R. Nakano, R. Nayak, A. Neelakantan, R. Ngo, H. Noh, L. Ouyang, C. O’Keefe, J. Pachocki, A. Paino, J. Palermo, A. Pantuliano, G. Parascandolo, J. Parish, E. Parparita,
A. Passos, M. Pavlov, A. Peng, A. Perelman, F. de Avila Belbute Peres, M. Petrov, H. P. de Oliveira Pinto, Michael, Pokorny, M. Pokrass, V. Pong, T. Powell, A. Power, B. Power, E. Proehl, R. Puri, A. Radford, J. Rae, A. Ramesh, C. Raymond, F. Real, K. Rimbach, C. Ross, B. Rotsted, H. Roussez, N. Ryder, M. Saltarelli, T. Sanders, S. Santurkar, G. Sastry, H. Schmidt, D. Schnurr, J. Schulman, D. Selsam, K. Sheppard, T. Sherbakov, J. Shieh, S. Shoker, P. Shyam, S. Sidor, E. Sigler, M. Simens, J. Sitkin, K. Slama, I. Sohl, B. Sokolowsky, Y. Song, N. Staudacher, F. P. Such, N. Summers, I. Sutskever, J. Tang, N. Tezak, M. Thompson, P. Tillet, A. Tootoonchian, E. Tseng, P. Tuggle, N. Turley, J. Tworek, J. F. C. Uribe, A. Vallone, A. Vijayvergiya, C. Voss, C. Wainwright, J. J. Wang, A. Wang, B. Wang, J. Ward, J. Wei, C. Weinmann, A. Welihinda, P. Welinder, J. Weng, L. Weng, M. Wiethoff, D. Willner, C. Winter, S. Wolrich, H. Wong, L. Workman, S. Wu, J. Wu, M. Wu, K. Xiao, T. Xu, S. Yoo, K. Yu, Q. Yuan, W. Zaremba,
R. Zellers, C. Zhang, M. Zhang, S. Zhao, T. Zheng, J. Zhuang, W. Zhuk, and B. Zoph.

</span>
<span class="ltx_bibblock">Gpt-4 technical report, 2023.

</span>
</li>
<li id="bib.bib30" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Pfeiffer et al. (2020)</span>
<span class="ltx_bibblock">
J. Pfeiffer, I. Vulić, I. Gurevych, and S. Ruder.

</span>
<span class="ltx_bibblock">Mad-x: An adapter-based framework for multi-task cross-lingual transfer.

</span>
<span class="ltx_bibblock">In <em id="bib.bib30.1.1" class="ltx_emph ltx_font_italic">EMNLP</em>, 2020.

</span>
</li>
<li id="bib.bib31" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Pfeiffer et al. (2023)</span>
<span class="ltx_bibblock">
J. Pfeiffer, S. Ruder, I. Vulić, and E. M. Ponti.

</span>
<span class="ltx_bibblock">Modular deep learning, 2023.

</span>
</li>
<li id="bib.bib32" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Purushwalkam et al. (2019)</span>
<span class="ltx_bibblock">
S. Purushwalkam, M. Nickel, A. Gupta, and M. Ranzato.

</span>
<span class="ltx_bibblock">Task-driven modular networks for zero-shot compositional learning.

</span>
<span class="ltx_bibblock">In <em id="bib.bib32.1.1" class="ltx_emph ltx_font_italic">ICCV</em>, 2019.

</span>
</li>
<li id="bib.bib33" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Raffel (2023a)</span>
<span class="ltx_bibblock">
C. Raffel.

</span>
<span class="ltx_bibblock">Building machine learning models like open source software.

</span>
<span class="ltx_bibblock"><em id="bib.bib33.1.1" class="ltx_emph ltx_font_italic">Communications of the ACM</em>, 66(2):38–40, 2023a.

</span>
</li>
<li id="bib.bib34" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Raffel (2023b)</span>
<span class="ltx_bibblock">
C. Raffel.

</span>
<span class="ltx_bibblock">Building machine learning models like open source software.

</span>
<span class="ltx_bibblock"><em id="bib.bib34.1.1" class="ltx_emph ltx_font_italic">Commun. ACM</em>, 66(2):38–40, jan 2023b.

</span>
<span class="ltx_bibblock">ISSN 0001-0782.

</span>
<span class="ltx_bibblock"><a target="_blank" href="https:/doi.org/10.1145/3545111" title="" class="ltx_ref">10.1145/3545111</a>.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://doi.org/10.1145/3545111" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://doi.org/10.1145/3545111</a>.

</span>
</li>
<li id="bib.bib35" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Raffel et al. (2020)</span>
<span class="ltx_bibblock">
C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu.

</span>
<span class="ltx_bibblock">Exploring the limits of transfer learning with a unified text-to-text transformer.

</span>
<span class="ltx_bibblock"><em id="bib.bib35.1.1" class="ltx_emph ltx_font_italic">Journal of Machine Learning Research</em>, 21(1):5485–5551, 2020.

</span>
</li>
<li id="bib.bib36" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Reddi et al. (2021)</span>
<span class="ltx_bibblock">
S. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečný, S. Kumar, and H. B. McMahan.

</span>
<span class="ltx_bibblock">Adaptive federated optimization, 2021.

</span>
</li>
<li id="bib.bib37" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Reisser et al. (2021)</span>
<span class="ltx_bibblock">
M. Reisser, C. Louizos, E. Gavves, and M. Welling.

</span>
<span class="ltx_bibblock">Federated mixture of experts, 2021.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://openreview.net/forum?id=YgrdmztE4OY" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://openreview.net/forum?id=YgrdmztE4OY</a>.

</span>
</li>
<li id="bib.bib38" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Ryabinin and Gusev (2020)</span>
<span class="ltx_bibblock">
M. Ryabinin and A. Gusev.

</span>
<span class="ltx_bibblock">Towards crowdsourced training of large neural networks using decentralized mixture-of-experts.

</span>
<span class="ltx_bibblock"><em id="bib.bib38.1.1" class="ltx_emph ltx_font_italic">Advances in Neural Information Processing Systems</em>, 33:3659–3672, 2020.

</span>
</li>
<li id="bib.bib39" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Shazeer et al. (2017)</span>
<span class="ltx_bibblock">
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean.

</span>
<span class="ltx_bibblock">Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.

</span>
<span class="ltx_bibblock">In <em id="bib.bib39.1.1" class="ltx_emph ltx_font_italic">International Conference on Learning Representations</em>, 2017.

</span>
<span class="ltx_bibblock">URL <a target="_blank" href="https://openreview.net/forum?id=B1ckMDqlg" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://openreview.net/forum?id=B1ckMDqlg</a>.

</span>
</li>
<li id="bib.bib40" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Sutskever et al. (2013)</span>
<span class="ltx_bibblock">
I. Sutskever, J. Martens, G. Dahl, and G. Hinton.

</span>
<span class="ltx_bibblock">On the importance of initialization and momentum in deep learning.

</span>
<span class="ltx_bibblock"><em id="bib.bib40.1.1" class="ltx_emph ltx_font_italic">International Conference on Machine Learning (ICML)</em>, 2013.

</span>
</li>
<li id="bib.bib41" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Touvron et al. (2023)</span>
<span class="ltx_bibblock">
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom.

</span>
<span class="ltx_bibblock">Llama 2: Open foundation and fine-tuned chat models, 2023.

</span>
</li>
</ul>
</section>
<section id="Sx1" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">Acknowledgements</h2>

<div id="Sx1.p1" class="ltx_para">
<p id="Sx1.p1.1" class="ltx_p">We would like to thank Owen He, Diego de las Casas, Ross Hemsley, Elena Gribovskaya, Yee Whye Teh, Yutian Chen, Alexandre Galashov, Bogdan Damoc, Shreya Pathak, Amal Rannen-Triki, Alek Andreev, David Budden, Jörg Bornschein for their valuable feedback throughout the development of the project.
We would also like to thank Jeff Dean, Satinder Baveja, and Raia Hadsell for their suggestions, feedback and support. Finally, we would like to thank Kitty Stacpoole and Guy Scully for their program management support.</p>
</div>
<div class="ltx_pagination ltx_role_newpage"></div>
</section>
<section id="Sx2" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">Supplementary Materials</h2>

<section id="Sx2.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">7.1 </span>Additional Details</h3>

<figure id="Sx2.T4" class="ltx_table">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="Sx2.T4.5.1.1" class="ltx_text" style="font-size:90%;">Table 4</span>: </span><span id="Sx2.T4.6.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Hyperparameter for Inner Optimization<span id="Sx2.T4.6.2.1" class="ltx_text ltx_font_medium">.</span></span></figcaption>
<div id="Sx2.T4.2" class="ltx_inline-block ltx_align_center ltx_transformed_outer" style="width:433.6pt;height:338.2pt;vertical-align:-0.0pt;"><span class="ltx_transformed_inner" style="transform:translate(113.0pt,-88.1pt) scale(2.0877743729746,2.0877743729746) ;">
<table id="Sx2.T4.2.2" class="ltx_tabular ltx_guessed_headers ltx_align_middle">
<tbody class="ltx_tbody">
<tr id="Sx2.T4.2.2.3.1" class="ltx_tr">
<th id="Sx2.T4.2.2.3.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_tt">Hyperparameter</th>
<td id="Sx2.T4.2.2.3.1.2" class="ltx_td ltx_align_center ltx_border_tt">150M model</td>
<td id="Sx2.T4.2.2.3.1.3" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_tt">1B model</td>
</tr>
<tr id="Sx2.T4.2.2.4.2" class="ltx_tr">
<th id="Sx2.T4.2.2.4.2.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t"># Layers</th>
<td id="Sx2.T4.2.2.4.2.2" class="ltx_td ltx_align_center ltx_border_t">12</td>
<td id="Sx2.T4.2.2.4.2.3" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">24</td>
</tr>
<tr id="Sx2.T4.2.2.5.3" class="ltx_tr">
<th id="Sx2.T4.2.2.5.3.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Hidden dimension</th>
<td id="Sx2.T4.2.2.5.3.2" class="ltx_td ltx_align_center">896</td>
<td id="Sx2.T4.2.2.5.3.3" class="ltx_td ltx_nopad_r ltx_align_center">2048</td>
</tr>
<tr id="Sx2.T4.2.2.6.4" class="ltx_tr">
<th id="Sx2.T4.2.2.6.4.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r"># Heads</th>
<td id="Sx2.T4.2.2.6.4.2" class="ltx_td ltx_align_center">16</td>
<td id="Sx2.T4.2.2.6.4.3" class="ltx_td ltx_nopad_r ltx_align_center">16</td>
</tr>
<tr id="Sx2.T4.2.2.7.5" class="ltx_tr">
<th id="Sx2.T4.2.2.7.5.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Key/value size</th>
<td id="Sx2.T4.2.2.7.5.2" class="ltx_td ltx_align_center">64</td>
<td id="Sx2.T4.2.2.7.5.3" class="ltx_td ltx_nopad_r ltx_align_center">128</td>
</tr>
<tr id="Sx2.T4.2.2.8.6" class="ltx_tr">
<th id="Sx2.T4.2.2.8.6.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Weight decay</th>
<td id="Sx2.T4.2.2.8.6.2" class="ltx_td ltx_align_center">0.1</td>
<td id="Sx2.T4.2.2.8.6.3" class="ltx_td ltx_nopad_r ltx_align_center">0.1</td>
</tr>
<tr id="Sx2.T4.2.2.2" class="ltx_tr">
<th id="Sx2.T4.2.2.2.3" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Max learning rate</th>
<td id="Sx2.T4.1.1.1.1" class="ltx_td ltx_align_center"><math id="Sx2.T4.1.1.1.1.m1.1" class="ltx_Math" alttext="4e^{-4}" display="inline"><semantics id="Sx2.T4.1.1.1.1.m1.1a"><mrow id="Sx2.T4.1.1.1.1.m1.1.1" xref="Sx2.T4.1.1.1.1.m1.1.1.cmml"><mn id="Sx2.T4.1.1.1.1.m1.1.1.2" xref="Sx2.T4.1.1.1.1.m1.1.1.2.cmml">4</mn><mo lspace="0em" rspace="0em" id="Sx2.T4.1.1.1.1.m1.1.1.1" xref="Sx2.T4.1.1.1.1.m1.1.1.1.cmml">​</mo><msup id="Sx2.T4.1.1.1.1.m1.1.1.3" xref="Sx2.T4.1.1.1.1.m1.1.1.3.cmml"><mi id="Sx2.T4.1.1.1.1.m1.1.1.3.2" xref="Sx2.T4.1.1.1.1.m1.1.1.3.2.cmml">e</mi><mrow id="Sx2.T4.1.1.1.1.m1.1.1.3.3" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3.cmml"><mo id="Sx2.T4.1.1.1.1.m1.1.1.3.3a" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3.cmml">−</mo><mn id="Sx2.T4.1.1.1.1.m1.1.1.3.3.2" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3.2.cmml">4</mn></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="Sx2.T4.1.1.1.1.m1.1b"><apply id="Sx2.T4.1.1.1.1.m1.1.1.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1"><times id="Sx2.T4.1.1.1.1.m1.1.1.1.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.1"></times><cn type="integer" id="Sx2.T4.1.1.1.1.m1.1.1.2.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.2">4</cn><apply id="Sx2.T4.1.1.1.1.m1.1.1.3.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3"><csymbol cd="ambiguous" id="Sx2.T4.1.1.1.1.m1.1.1.3.1.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3">superscript</csymbol><ci id="Sx2.T4.1.1.1.1.m1.1.1.3.2.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3.2">𝑒</ci><apply id="Sx2.T4.1.1.1.1.m1.1.1.3.3.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3"><minus id="Sx2.T4.1.1.1.1.m1.1.1.3.3.1.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3"></minus><cn type="integer" id="Sx2.T4.1.1.1.1.m1.1.1.3.3.2.cmml" xref="Sx2.T4.1.1.1.1.m1.1.1.3.3.2">4</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T4.1.1.1.1.m1.1c">4e^{-4}</annotation></semantics></math></td>
<td id="Sx2.T4.2.2.2.2" class="ltx_td ltx_nopad_r ltx_align_center"><math id="Sx2.T4.2.2.2.2.m1.1" class="ltx_Math" alttext="2e^{-4}" display="inline"><semantics id="Sx2.T4.2.2.2.2.m1.1a"><mrow id="Sx2.T4.2.2.2.2.m1.1.1" xref="Sx2.T4.2.2.2.2.m1.1.1.cmml"><mn id="Sx2.T4.2.2.2.2.m1.1.1.2" xref="Sx2.T4.2.2.2.2.m1.1.1.2.cmml">2</mn><mo lspace="0em" rspace="0em" id="Sx2.T4.2.2.2.2.m1.1.1.1" xref="Sx2.T4.2.2.2.2.m1.1.1.1.cmml">​</mo><msup id="Sx2.T4.2.2.2.2.m1.1.1.3" xref="Sx2.T4.2.2.2.2.m1.1.1.3.cmml"><mi id="Sx2.T4.2.2.2.2.m1.1.1.3.2" xref="Sx2.T4.2.2.2.2.m1.1.1.3.2.cmml">e</mi><mrow id="Sx2.T4.2.2.2.2.m1.1.1.3.3" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3.cmml"><mo id="Sx2.T4.2.2.2.2.m1.1.1.3.3a" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3.cmml">−</mo><mn id="Sx2.T4.2.2.2.2.m1.1.1.3.3.2" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3.2.cmml">4</mn></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="Sx2.T4.2.2.2.2.m1.1b"><apply id="Sx2.T4.2.2.2.2.m1.1.1.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1"><times id="Sx2.T4.2.2.2.2.m1.1.1.1.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.1"></times><cn type="integer" id="Sx2.T4.2.2.2.2.m1.1.1.2.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.2">2</cn><apply id="Sx2.T4.2.2.2.2.m1.1.1.3.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3"><csymbol cd="ambiguous" id="Sx2.T4.2.2.2.2.m1.1.1.3.1.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3">superscript</csymbol><ci id="Sx2.T4.2.2.2.2.m1.1.1.3.2.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3.2">𝑒</ci><apply id="Sx2.T4.2.2.2.2.m1.1.1.3.3.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3"><minus id="Sx2.T4.2.2.2.2.m1.1.1.3.3.1.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3"></minus><cn type="integer" id="Sx2.T4.2.2.2.2.m1.1.1.3.3.2.cmml" xref="Sx2.T4.2.2.2.2.m1.1.1.3.3.2">4</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T4.2.2.2.2.m1.1c">2e^{-4}</annotation></semantics></math></td>
</tr>
<tr id="Sx2.T4.2.2.9.7" class="ltx_tr">
<th id="Sx2.T4.2.2.9.7.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Batch size</th>
<td id="Sx2.T4.2.2.9.7.2" class="ltx_td ltx_align_center">512</td>
<td id="Sx2.T4.2.2.9.7.3" class="ltx_td ltx_nopad_r ltx_align_center">512</td>
</tr>
<tr id="Sx2.T4.2.2.10.8" class="ltx_tr">
<th id="Sx2.T4.2.2.10.8.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_bb ltx_border_r"># Warmup steps</th>
<td id="Sx2.T4.2.2.10.8.2" class="ltx_td ltx_align_center ltx_border_bb">1000</td>
<td id="Sx2.T4.2.2.10.8.3" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_bb">1000</td>
</tr>
</tbody>
</table>
</span></div>
</figure>
<div id="Sx2.SS1.p1" class="ltx_para">
<p id="Sx2.SS1.p1.2" class="ltx_p">We list in <a href="#Sx2.T4" title="Table 4 ‣ 7.1 Additional Details ‣ Supplementary Materials ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 4</span></a> the hyperparameters used in the inner optimization. The <span id="Sx2.SS1.p1.2.1" class="ltx_text ltx_font_italic">150M model</span> hyperparameters are used for both the dense baselines of that size and for the path optimization of DiPaCo. In all our experiments, we used as outer optimizer Nesterov <cite class="ltx_cite ltx_citemacro_citep">(Sutskever et al., <a href="#bib.bib40" title="" class="ltx_ref">2013</a>)</cite> with outer learning rate of <math id="Sx2.SS1.p1.1.m1.1" class="ltx_Math" alttext="0.7" display="inline"><semantics id="Sx2.SS1.p1.1.m1.1a"><mn id="Sx2.SS1.p1.1.m1.1.1" xref="Sx2.SS1.p1.1.m1.1.1.cmml">0.7</mn><annotation-xml encoding="MathML-Content" id="Sx2.SS1.p1.1.m1.1b"><cn type="float" id="Sx2.SS1.p1.1.m1.1.1.cmml" xref="Sx2.SS1.p1.1.m1.1.1">0.7</cn></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS1.p1.1.m1.1c">0.7</annotation></semantics></math> and outer momentum of <math id="Sx2.SS1.p1.2.m2.1" class="ltx_Math" alttext="0.9" display="inline"><semantics id="Sx2.SS1.p1.2.m2.1a"><mn id="Sx2.SS1.p1.2.m2.1.1" xref="Sx2.SS1.p1.2.m2.1.1.cmml">0.9</mn><annotation-xml encoding="MathML-Content" id="Sx2.SS1.p1.2.m2.1b"><cn type="float" id="Sx2.SS1.p1.2.m2.1.1.cmml" xref="Sx2.SS1.p1.2.m2.1.1">0.9</cn></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS1.p1.2.m2.1c">0.9</annotation></semantics></math>, following the same recipe introduced by <cite class="ltx_cite ltx_citemacro_citet">Douillard et al. (<a href="#bib.bib13" title="" class="ltx_ref">2023</a>)</cite>.</p>
</div>
</section>
<section id="Sx2.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">7.2 </span>Routing Details</h3>

<div id="Sx2.SS2.p1" class="ltx_para">
<p id="Sx2.SS2.p1.1" class="ltx_p">We first discuss the details of discriminative coarse routing, as used in training. We then discuss the details of more frequent routing, as optionally used in evaluation.</p>
</div>
<section id="Sx2.SS2.SSS1" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">7.2.1 </span>Document level discriminative routing</h4>

<div id="Sx2.SS2.SSS1.p1" class="ltx_para">
<p id="Sx2.SS2.SSS1.p1.1" class="ltx_p">For discriminative routing, we assume that we have</p>
<ul id="Sx2.I1" class="ltx_itemize">
<li id="Sx2.I1.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="Sx2.I1.i1.p1" class="ltx_para">
<p id="Sx2.I1.i1.p1.1" class="ltx_p">Trained an initial language model that can be used as an initialization for finetuning experts.</p>
</div>
</li>
<li id="Sx2.I1.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="Sx2.I1.i2.p1" class="ltx_para">
<p id="Sx2.I1.i2.p1.1" class="ltx_p">Set up a feature extractor that gets a feature for each document to feed to the router.</p>
</div>
</li>
</ul>
<p id="Sx2.SS2.SSS1.p1.2" class="ltx_p">We alternate between two stages:</p>
<ol id="Sx2.I2" class="ltx_enumerate">
<li id="Sx2.I2.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="Sx2.I2.i1.p1" class="ltx_para">
<p id="Sx2.I2.i1.p1.1" class="ltx_p">Train a gater, perhaps based on the PPL of current experts on an unused portion of the training set,</p>
</div>
</li>
<li id="Sx2.I2.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="Sx2.I2.i2.p1" class="ltx_para">
<p id="Sx2.I2.i2.p1.1" class="ltx_p">Train experts based on a sharding from current gater.</p>
</div>
</li>
</ol>
</div>
<div id="Sx2.SS2.SSS1.p2" class="ltx_para">
<p id="Sx2.SS2.SSS1.p2.1" class="ltx_p">In this work, the feature for the router is always the average of the hidden state from the last transformer block from the initial LM over the first 32 tokens of a document; let us denote this as <math id="Sx2.SS2.SSS1.p2.1.m1.1" class="ltx_Math" alttext="g(\text{document})" display="inline"><semantics id="Sx2.SS2.SSS1.p2.1.m1.1a"><mrow id="Sx2.SS2.SSS1.p2.1.m1.1.2" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.cmml"><mi id="Sx2.SS2.SSS1.p2.1.m1.1.2.2" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.2.cmml">g</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p2.1.m1.1.2.1" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.1.cmml">​</mo><mrow id="Sx2.SS2.SSS1.p2.1.m1.1.2.3.2" xref="Sx2.SS2.SSS1.p2.1.m1.1.1a.cmml"><mo stretchy="false" id="Sx2.SS2.SSS1.p2.1.m1.1.2.3.2.1" xref="Sx2.SS2.SSS1.p2.1.m1.1.1a.cmml">(</mo><mtext id="Sx2.SS2.SSS1.p2.1.m1.1.1" xref="Sx2.SS2.SSS1.p2.1.m1.1.1.cmml">document</mtext><mo stretchy="false" id="Sx2.SS2.SSS1.p2.1.m1.1.2.3.2.2" xref="Sx2.SS2.SSS1.p2.1.m1.1.1a.cmml">)</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p2.1.m1.1b"><apply id="Sx2.SS2.SSS1.p2.1.m1.1.2.cmml" xref="Sx2.SS2.SSS1.p2.1.m1.1.2"><times id="Sx2.SS2.SSS1.p2.1.m1.1.2.1.cmml" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.1"></times><ci id="Sx2.SS2.SSS1.p2.1.m1.1.2.2.cmml" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.2">𝑔</ci><ci id="Sx2.SS2.SSS1.p2.1.m1.1.1a.cmml" xref="Sx2.SS2.SSS1.p2.1.m1.1.2.3.2"><mtext id="Sx2.SS2.SSS1.p2.1.m1.1.1.cmml" xref="Sx2.SS2.SSS1.p2.1.m1.1.1">document</mtext></ci></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p2.1.m1.1c">g(\text{document})</annotation></semantics></math>. We also reserve a small subset of documents (from here on called the “router data”) for making decisions about router quality and training the router; we use .005 of C4 for router data.</p>
</div>
<div id="Sx2.SS2.SSS1.p3" class="ltx_para">
<p id="Sx2.SS2.SSS1.p3.1" class="ltx_p">Our initial router is always constructed via k-means; and we train the first set of paths using that router (and sharding by argmin to cluster centroid).</p>
</div>
<div id="Sx2.SS2.SSS1.p4" class="ltx_para">
<p id="Sx2.SS2.SSS1.p4.5" class="ltx_p">To train the discriminative router, we first compute the summed auto-regressive log-likelihood of each document in the router data if encoded by each path. That is, given <math id="Sx2.SS2.SSS1.p4.1.m1.1" class="ltx_Math" alttext="n" display="inline"><semantics id="Sx2.SS2.SSS1.p4.1.m1.1a"><mi id="Sx2.SS2.SSS1.p4.1.m1.1.1" xref="Sx2.SS2.SSS1.p4.1.m1.1.1.cmml">n</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.1.m1.1b"><ci id="Sx2.SS2.SSS1.p4.1.m1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.1.m1.1.1">𝑛</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.1.m1.1c">n</annotation></semantics></math> documents of length <math id="Sx2.SS2.SSS1.p4.2.m2.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS1.p4.2.m2.1a"><mi id="Sx2.SS2.SSS1.p4.2.m2.1.1" xref="Sx2.SS2.SSS1.p4.2.m2.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.2.m2.1b"><ci id="Sx2.SS2.SSS1.p4.2.m2.1.1.cmml" xref="Sx2.SS2.SSS1.p4.2.m2.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.2.m2.1c">L</annotation></semantics></math> and paths <math id="Sx2.SS2.SSS1.p4.3.m3.3" class="ltx_Math" alttext="f_{1},...,f_{K}" display="inline"><semantics id="Sx2.SS2.SSS1.p4.3.m3.3a"><mrow id="Sx2.SS2.SSS1.p4.3.m3.3.3.2" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.3.cmml"><msub id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.cmml"><mi id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.2" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.2.cmml">f</mi><mn id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.3" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.3.cmml">1</mn></msub><mo id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.3" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.3.cmml">,</mo><mi mathvariant="normal" id="Sx2.SS2.SSS1.p4.3.m3.1.1" xref="Sx2.SS2.SSS1.p4.3.m3.1.1.cmml">…</mi><mo id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.4" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.3.cmml">,</mo><msub id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.cmml"><mi id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.2" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.2.cmml">f</mi><mi id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.3" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.3.cmml">K</mi></msub></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.3.m3.3b"><list id="Sx2.SS2.SSS1.p4.3.m3.3.3.3.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2"><apply id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1">subscript</csymbol><ci id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.2.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.2">𝑓</ci><cn type="integer" id="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.3.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.2.2.1.1.3">1</cn></apply><ci id="Sx2.SS2.SSS1.p4.3.m3.1.1.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.1.1">…</ci><apply id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.1.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2">subscript</csymbol><ci id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.2.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.2">𝑓</ci><ci id="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.3.cmml" xref="Sx2.SS2.SSS1.p4.3.m3.3.3.2.2.3">𝐾</ci></apply></list></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.3.m3.3c">f_{1},...,f_{K}</annotation></semantics></math> we get an <math id="Sx2.SS2.SSS1.p4.4.m4.1" class="ltx_Math" alttext="n\times L\times K" display="inline"><semantics id="Sx2.SS2.SSS1.p4.4.m4.1a"><mrow id="Sx2.SS2.SSS1.p4.4.m4.1.1" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.cmml"><mi id="Sx2.SS2.SSS1.p4.4.m4.1.1.2" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.2.cmml">n</mi><mo lspace="0.222em" rspace="0.222em" id="Sx2.SS2.SSS1.p4.4.m4.1.1.1" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.1.cmml">×</mo><mi id="Sx2.SS2.SSS1.p4.4.m4.1.1.3" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.3.cmml">L</mi><mo lspace="0.222em" rspace="0.222em" id="Sx2.SS2.SSS1.p4.4.m4.1.1.1a" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.1.cmml">×</mo><mi id="Sx2.SS2.SSS1.p4.4.m4.1.1.4" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.4.cmml">K</mi></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.4.m4.1b"><apply id="Sx2.SS2.SSS1.p4.4.m4.1.1.cmml" xref="Sx2.SS2.SSS1.p4.4.m4.1.1"><times id="Sx2.SS2.SSS1.p4.4.m4.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.1"></times><ci id="Sx2.SS2.SSS1.p4.4.m4.1.1.2.cmml" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.2">𝑛</ci><ci id="Sx2.SS2.SSS1.p4.4.m4.1.1.3.cmml" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.3">𝐿</ci><ci id="Sx2.SS2.SSS1.p4.4.m4.1.1.4.cmml" xref="Sx2.SS2.SSS1.p4.4.m4.1.1.4">𝐾</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.4.m4.1c">n\times L\times K</annotation></semantics></math> array of scores <math id="Sx2.SS2.SSS1.p4.5.m5.1" class="ltx_Math" alttext="S_{ijp}" display="inline"><semantics id="Sx2.SS2.SSS1.p4.5.m5.1a"><msub id="Sx2.SS2.SSS1.p4.5.m5.1.1" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.cmml"><mi id="Sx2.SS2.SSS1.p4.5.m5.1.1.2" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.2.cmml">S</mi><mrow id="Sx2.SS2.SSS1.p4.5.m5.1.1.3" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.cmml"><mi id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.2" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.3" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1a" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.4" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.4.cmml">p</mi></mrow></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.5.m5.1b"><apply id="Sx2.SS2.SSS1.p4.5.m5.1.1.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p4.5.m5.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1">subscript</csymbol><ci id="Sx2.SS2.SSS1.p4.5.m5.1.1.2.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.2">𝑆</ci><apply id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3"><times id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.1"></times><ci id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.2.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.2">𝑖</ci><ci id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.3.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.3">𝑗</ci><ci id="Sx2.SS2.SSS1.p4.5.m5.1.1.3.4.cmml" xref="Sx2.SS2.SSS1.p4.5.m5.1.1.3.4">𝑝</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.5.m5.1c">S_{ijp}</annotation></semantics></math> where</p>
<table id="Sx2.Ex1" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="Sx2.Ex1.m1.8" class="ltx_Math" alttext="S_{ijp}=\log p_{f_{i}}(t_{l,j}|t_{l-1,j},...,t_{1,j})," display="block"><semantics id="Sx2.Ex1.m1.8a"><mrow id="Sx2.Ex1.m1.8.8.1" xref="Sx2.Ex1.m1.8.8.1.1.cmml"><mrow id="Sx2.Ex1.m1.8.8.1.1" xref="Sx2.Ex1.m1.8.8.1.1.cmml"><msub id="Sx2.Ex1.m1.8.8.1.1.3" xref="Sx2.Ex1.m1.8.8.1.1.3.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.3.2" xref="Sx2.Ex1.m1.8.8.1.1.3.2.cmml">S</mi><mrow id="Sx2.Ex1.m1.8.8.1.1.3.3" xref="Sx2.Ex1.m1.8.8.1.1.3.3.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.3.3.2" xref="Sx2.Ex1.m1.8.8.1.1.3.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex1.m1.8.8.1.1.3.3.1" xref="Sx2.Ex1.m1.8.8.1.1.3.3.1.cmml">​</mo><mi id="Sx2.Ex1.m1.8.8.1.1.3.3.3" xref="Sx2.Ex1.m1.8.8.1.1.3.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex1.m1.8.8.1.1.3.3.1a" xref="Sx2.Ex1.m1.8.8.1.1.3.3.1.cmml">​</mo><mi id="Sx2.Ex1.m1.8.8.1.1.3.3.4" xref="Sx2.Ex1.m1.8.8.1.1.3.3.4.cmml">p</mi></mrow></msub><mo id="Sx2.Ex1.m1.8.8.1.1.2" xref="Sx2.Ex1.m1.8.8.1.1.2.cmml">=</mo><mrow id="Sx2.Ex1.m1.8.8.1.1.1" xref="Sx2.Ex1.m1.8.8.1.1.1.cmml"><mrow id="Sx2.Ex1.m1.8.8.1.1.1.3" xref="Sx2.Ex1.m1.8.8.1.1.1.3.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.3.1" xref="Sx2.Ex1.m1.8.8.1.1.1.3.1.cmml">log</mi><mo lspace="0.167em" id="Sx2.Ex1.m1.8.8.1.1.1.3a" xref="Sx2.Ex1.m1.8.8.1.1.1.3.cmml">⁡</mo><msub id="Sx2.Ex1.m1.8.8.1.1.1.3.2" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.3.2.2" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.2.cmml">p</mi><msub id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.2" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.2.cmml">f</mi><mi id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.3" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.3.cmml">i</mi></msub></msub></mrow><mo lspace="0em" rspace="0em" id="Sx2.Ex1.m1.8.8.1.1.1.2" xref="Sx2.Ex1.m1.8.8.1.1.1.2.cmml">​</mo><mrow id="Sx2.Ex1.m1.8.8.1.1.1.1.1" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.cmml"><mo stretchy="false" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.cmml">(</mo><mrow id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.cmml"><msub id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.2.cmml">t</mi><mrow id="Sx2.Ex1.m1.2.2.2.4" xref="Sx2.Ex1.m1.2.2.2.3.cmml"><mi id="Sx2.Ex1.m1.1.1.1.1" xref="Sx2.Ex1.m1.1.1.1.1.cmml">l</mi><mo id="Sx2.Ex1.m1.2.2.2.4.1" xref="Sx2.Ex1.m1.2.2.2.3.cmml">,</mo><mi id="Sx2.Ex1.m1.2.2.2.2" xref="Sx2.Ex1.m1.2.2.2.2.cmml">j</mi></mrow></msub><mo fence="false" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.3" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.3.cmml">|</mo><mrow id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.3.cmml"><msub id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.2.cmml">t</mi><mrow id="Sx2.Ex1.m1.4.4.2.2" xref="Sx2.Ex1.m1.4.4.2.3.cmml"><mrow id="Sx2.Ex1.m1.4.4.2.2.1" xref="Sx2.Ex1.m1.4.4.2.2.1.cmml"><mi id="Sx2.Ex1.m1.4.4.2.2.1.2" xref="Sx2.Ex1.m1.4.4.2.2.1.2.cmml">l</mi><mo id="Sx2.Ex1.m1.4.4.2.2.1.1" xref="Sx2.Ex1.m1.4.4.2.2.1.1.cmml">−</mo><mn id="Sx2.Ex1.m1.4.4.2.2.1.3" xref="Sx2.Ex1.m1.4.4.2.2.1.3.cmml">1</mn></mrow><mo id="Sx2.Ex1.m1.4.4.2.2.2" xref="Sx2.Ex1.m1.4.4.2.3.cmml">,</mo><mi id="Sx2.Ex1.m1.3.3.1.1" xref="Sx2.Ex1.m1.3.3.1.1.cmml">j</mi></mrow></msub><mo id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.3" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.3.cmml">,</mo><mi mathvariant="normal" id="Sx2.Ex1.m1.7.7" xref="Sx2.Ex1.m1.7.7.cmml">…</mi><mo id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.4" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.3.cmml">,</mo><msub id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.cmml"><mi id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.2" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.2.cmml">t</mi><mrow id="Sx2.Ex1.m1.6.6.2.4" xref="Sx2.Ex1.m1.6.6.2.3.cmml"><mn id="Sx2.Ex1.m1.5.5.1.1" xref="Sx2.Ex1.m1.5.5.1.1.cmml">1</mn><mo id="Sx2.Ex1.m1.6.6.2.4.1" xref="Sx2.Ex1.m1.6.6.2.3.cmml">,</mo><mi id="Sx2.Ex1.m1.6.6.2.2" xref="Sx2.Ex1.m1.6.6.2.2.cmml">j</mi></mrow></msub></mrow></mrow><mo stretchy="false" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.3" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow><mo id="Sx2.Ex1.m1.8.8.1.2" xref="Sx2.Ex1.m1.8.8.1.1.cmml">,</mo></mrow><annotation-xml encoding="MathML-Content" id="Sx2.Ex1.m1.8b"><apply id="Sx2.Ex1.m1.8.8.1.1.cmml" xref="Sx2.Ex1.m1.8.8.1"><eq id="Sx2.Ex1.m1.8.8.1.1.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.2"></eq><apply id="Sx2.Ex1.m1.8.8.1.1.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.3.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.3.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.2">𝑆</ci><apply id="Sx2.Ex1.m1.8.8.1.1.3.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.3"><times id="Sx2.Ex1.m1.8.8.1.1.3.3.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.3.1"></times><ci id="Sx2.Ex1.m1.8.8.1.1.3.3.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.3.2">𝑖</ci><ci id="Sx2.Ex1.m1.8.8.1.1.3.3.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.3.3">𝑗</ci><ci id="Sx2.Ex1.m1.8.8.1.1.3.3.4.cmml" xref="Sx2.Ex1.m1.8.8.1.1.3.3.4">𝑝</ci></apply></apply><apply id="Sx2.Ex1.m1.8.8.1.1.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1"><times id="Sx2.Ex1.m1.8.8.1.1.1.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.2"></times><apply id="Sx2.Ex1.m1.8.8.1.1.1.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3"><log id="Sx2.Ex1.m1.8.8.1.1.1.3.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.1"></log><apply id="Sx2.Ex1.m1.8.8.1.1.1.3.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.1.3.2.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.1.3.2.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.2">𝑝</ci><apply id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.2">𝑓</ci><ci id="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.3.2.3.3">𝑖</ci></apply></apply></apply><apply id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1"><csymbol cd="latexml" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.3">conditional</csymbol><apply id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.4.2">𝑡</ci><list id="Sx2.Ex1.m1.2.2.2.3.cmml" xref="Sx2.Ex1.m1.2.2.2.4"><ci id="Sx2.Ex1.m1.1.1.1.1.cmml" xref="Sx2.Ex1.m1.1.1.1.1">𝑙</ci><ci id="Sx2.Ex1.m1.2.2.2.2.cmml" xref="Sx2.Ex1.m1.2.2.2.2">𝑗</ci></list></apply><list id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.3.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2"><apply id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.1.1.1.2">𝑡</ci><list id="Sx2.Ex1.m1.4.4.2.3.cmml" xref="Sx2.Ex1.m1.4.4.2.2"><apply id="Sx2.Ex1.m1.4.4.2.2.1.cmml" xref="Sx2.Ex1.m1.4.4.2.2.1"><minus id="Sx2.Ex1.m1.4.4.2.2.1.1.cmml" xref="Sx2.Ex1.m1.4.4.2.2.1.1"></minus><ci id="Sx2.Ex1.m1.4.4.2.2.1.2.cmml" xref="Sx2.Ex1.m1.4.4.2.2.1.2">𝑙</ci><cn type="integer" id="Sx2.Ex1.m1.4.4.2.2.1.3.cmml" xref="Sx2.Ex1.m1.4.4.2.2.1.3">1</cn></apply><ci id="Sx2.Ex1.m1.3.3.1.1.cmml" xref="Sx2.Ex1.m1.3.3.1.1">𝑗</ci></list></apply><ci id="Sx2.Ex1.m1.7.7.cmml" xref="Sx2.Ex1.m1.7.7">…</ci><apply id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2"><csymbol cd="ambiguous" id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.1.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2">subscript</csymbol><ci id="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.2.cmml" xref="Sx2.Ex1.m1.8.8.1.1.1.1.1.1.2.2.2.2">𝑡</ci><list id="Sx2.Ex1.m1.6.6.2.3.cmml" xref="Sx2.Ex1.m1.6.6.2.4"><cn type="integer" id="Sx2.Ex1.m1.5.5.1.1.cmml" xref="Sx2.Ex1.m1.5.5.1.1">1</cn><ci id="Sx2.Ex1.m1.6.6.2.2.cmml" xref="Sx2.Ex1.m1.6.6.2.2">𝑗</ci></list></apply></list></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.Ex1.m1.8c">S_{ijp}=\log p_{f_{i}}(t_{l,j}|t_{l-1,j},...,t_{1,j}),</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
</tr></tbody>
</table>
<p id="Sx2.SS2.SSS1.p4.9" class="ltx_p">and where <math id="Sx2.SS2.SSS1.p4.6.m1.2" class="ltx_Math" alttext="t_{i,j}" display="inline"><semantics id="Sx2.SS2.SSS1.p4.6.m1.2a"><msub id="Sx2.SS2.SSS1.p4.6.m1.2.3" xref="Sx2.SS2.SSS1.p4.6.m1.2.3.cmml"><mi id="Sx2.SS2.SSS1.p4.6.m1.2.3.2" xref="Sx2.SS2.SSS1.p4.6.m1.2.3.2.cmml">t</mi><mrow id="Sx2.SS2.SSS1.p4.6.m1.2.2.2.4" xref="Sx2.SS2.SSS1.p4.6.m1.2.2.2.3.cmml"><mi id="Sx2.SS2.SSS1.p4.6.m1.1.1.1.1" xref="Sx2.SS2.SSS1.p4.6.m1.1.1.1.1.cmml">i</mi><mo id="Sx2.SS2.SSS1.p4.6.m1.2.2.2.4.1" xref="Sx2.SS2.SSS1.p4.6.m1.2.2.2.3.cmml">,</mo><mi id="Sx2.SS2.SSS1.p4.6.m1.2.2.2.2" xref="Sx2.SS2.SSS1.p4.6.m1.2.2.2.2.cmml">j</mi></mrow></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.6.m1.2b"><apply id="Sx2.SS2.SSS1.p4.6.m1.2.3.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.2.3"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p4.6.m1.2.3.1.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.2.3">subscript</csymbol><ci id="Sx2.SS2.SSS1.p4.6.m1.2.3.2.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.2.3.2">𝑡</ci><list id="Sx2.SS2.SSS1.p4.6.m1.2.2.2.3.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.2.2.2.4"><ci id="Sx2.SS2.SSS1.p4.6.m1.1.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.1.1.1.1">𝑖</ci><ci id="Sx2.SS2.SSS1.p4.6.m1.2.2.2.2.cmml" xref="Sx2.SS2.SSS1.p4.6.m1.2.2.2.2">𝑗</ci></list></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.6.m1.2c">t_{i,j}</annotation></semantics></math> is the <math id="Sx2.SS2.SSS1.p4.7.m2.1" class="ltx_Math" alttext="j" display="inline"><semantics id="Sx2.SS2.SSS1.p4.7.m2.1a"><mi id="Sx2.SS2.SSS1.p4.7.m2.1.1" xref="Sx2.SS2.SSS1.p4.7.m2.1.1.cmml">j</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.7.m2.1b"><ci id="Sx2.SS2.SSS1.p4.7.m2.1.1.cmml" xref="Sx2.SS2.SSS1.p4.7.m2.1.1">𝑗</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.7.m2.1c">j</annotation></semantics></math>th token in the <math id="Sx2.SS2.SSS1.p4.8.m3.1" class="ltx_Math" alttext="i" display="inline"><semantics id="Sx2.SS2.SSS1.p4.8.m3.1a"><mi id="Sx2.SS2.SSS1.p4.8.m3.1.1" xref="Sx2.SS2.SSS1.p4.8.m3.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.8.m3.1b"><ci id="Sx2.SS2.SSS1.p4.8.m3.1.1.cmml" xref="Sx2.SS2.SSS1.p4.8.m3.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.8.m3.1c">i</annotation></semantics></math>th document.
The router is always trained using a <math id="Sx2.SS2.SSS1.p4.9.m4.1" class="ltx_Math" alttext="K" display="inline"><semantics id="Sx2.SS2.SSS1.p4.9.m4.1a"><mi id="Sx2.SS2.SSS1.p4.9.m4.1.1" xref="Sx2.SS2.SSS1.p4.9.m4.1.1.cmml">K</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.9.m4.1b"><ci id="Sx2.SS2.SSS1.p4.9.m4.1.1.cmml" xref="Sx2.SS2.SSS1.p4.9.m4.1.1">𝐾</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.9.m4.1c">K</annotation></semantics></math> class linear logistic classifier with</p>
<table id="Sx2.Ex2" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="Sx2.Ex2.m1.1" class="ltx_Math" alttext="\text{argmax}_{p}\sum_{j=1}^{L}S_{ijp}" display="block"><semantics id="Sx2.Ex2.m1.1a"><mrow id="Sx2.Ex2.m1.1.1" xref="Sx2.Ex2.m1.1.1.cmml"><msub id="Sx2.Ex2.m1.1.1.2" xref="Sx2.Ex2.m1.1.1.2.cmml"><mtext id="Sx2.Ex2.m1.1.1.2.2" xref="Sx2.Ex2.m1.1.1.2.2a.cmml">argmax</mtext><mi id="Sx2.Ex2.m1.1.1.2.3" xref="Sx2.Ex2.m1.1.1.2.3.cmml">p</mi></msub><mo lspace="0em" rspace="0em" id="Sx2.Ex2.m1.1.1.1" xref="Sx2.Ex2.m1.1.1.1.cmml">​</mo><mrow id="Sx2.Ex2.m1.1.1.3" xref="Sx2.Ex2.m1.1.1.3.cmml"><munderover id="Sx2.Ex2.m1.1.1.3.1" xref="Sx2.Ex2.m1.1.1.3.1.cmml"><mo movablelimits="false" id="Sx2.Ex2.m1.1.1.3.1.2.2" xref="Sx2.Ex2.m1.1.1.3.1.2.2.cmml">∑</mo><mrow id="Sx2.Ex2.m1.1.1.3.1.2.3" xref="Sx2.Ex2.m1.1.1.3.1.2.3.cmml"><mi id="Sx2.Ex2.m1.1.1.3.1.2.3.2" xref="Sx2.Ex2.m1.1.1.3.1.2.3.2.cmml">j</mi><mo id="Sx2.Ex2.m1.1.1.3.1.2.3.1" xref="Sx2.Ex2.m1.1.1.3.1.2.3.1.cmml">=</mo><mn id="Sx2.Ex2.m1.1.1.3.1.2.3.3" xref="Sx2.Ex2.m1.1.1.3.1.2.3.3.cmml">1</mn></mrow><mi id="Sx2.Ex2.m1.1.1.3.1.3" xref="Sx2.Ex2.m1.1.1.3.1.3.cmml">L</mi></munderover><msub id="Sx2.Ex2.m1.1.1.3.2" xref="Sx2.Ex2.m1.1.1.3.2.cmml"><mi id="Sx2.Ex2.m1.1.1.3.2.2" xref="Sx2.Ex2.m1.1.1.3.2.2.cmml">S</mi><mrow id="Sx2.Ex2.m1.1.1.3.2.3" xref="Sx2.Ex2.m1.1.1.3.2.3.cmml"><mi id="Sx2.Ex2.m1.1.1.3.2.3.2" xref="Sx2.Ex2.m1.1.1.3.2.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex2.m1.1.1.3.2.3.1" xref="Sx2.Ex2.m1.1.1.3.2.3.1.cmml">​</mo><mi id="Sx2.Ex2.m1.1.1.3.2.3.3" xref="Sx2.Ex2.m1.1.1.3.2.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex2.m1.1.1.3.2.3.1a" xref="Sx2.Ex2.m1.1.1.3.2.3.1.cmml">​</mo><mi id="Sx2.Ex2.m1.1.1.3.2.3.4" xref="Sx2.Ex2.m1.1.1.3.2.3.4.cmml">p</mi></mrow></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.Ex2.m1.1b"><apply id="Sx2.Ex2.m1.1.1.cmml" xref="Sx2.Ex2.m1.1.1"><times id="Sx2.Ex2.m1.1.1.1.cmml" xref="Sx2.Ex2.m1.1.1.1"></times><apply id="Sx2.Ex2.m1.1.1.2.cmml" xref="Sx2.Ex2.m1.1.1.2"><csymbol cd="ambiguous" id="Sx2.Ex2.m1.1.1.2.1.cmml" xref="Sx2.Ex2.m1.1.1.2">subscript</csymbol><ci id="Sx2.Ex2.m1.1.1.2.2a.cmml" xref="Sx2.Ex2.m1.1.1.2.2"><mtext id="Sx2.Ex2.m1.1.1.2.2.cmml" xref="Sx2.Ex2.m1.1.1.2.2">argmax</mtext></ci><ci id="Sx2.Ex2.m1.1.1.2.3.cmml" xref="Sx2.Ex2.m1.1.1.2.3">𝑝</ci></apply><apply id="Sx2.Ex2.m1.1.1.3.cmml" xref="Sx2.Ex2.m1.1.1.3"><apply id="Sx2.Ex2.m1.1.1.3.1.cmml" xref="Sx2.Ex2.m1.1.1.3.1"><csymbol cd="ambiguous" id="Sx2.Ex2.m1.1.1.3.1.1.cmml" xref="Sx2.Ex2.m1.1.1.3.1">superscript</csymbol><apply id="Sx2.Ex2.m1.1.1.3.1.2.cmml" xref="Sx2.Ex2.m1.1.1.3.1"><csymbol cd="ambiguous" id="Sx2.Ex2.m1.1.1.3.1.2.1.cmml" xref="Sx2.Ex2.m1.1.1.3.1">subscript</csymbol><sum id="Sx2.Ex2.m1.1.1.3.1.2.2.cmml" xref="Sx2.Ex2.m1.1.1.3.1.2.2"></sum><apply id="Sx2.Ex2.m1.1.1.3.1.2.3.cmml" xref="Sx2.Ex2.m1.1.1.3.1.2.3"><eq id="Sx2.Ex2.m1.1.1.3.1.2.3.1.cmml" xref="Sx2.Ex2.m1.1.1.3.1.2.3.1"></eq><ci id="Sx2.Ex2.m1.1.1.3.1.2.3.2.cmml" xref="Sx2.Ex2.m1.1.1.3.1.2.3.2">𝑗</ci><cn type="integer" id="Sx2.Ex2.m1.1.1.3.1.2.3.3.cmml" xref="Sx2.Ex2.m1.1.1.3.1.2.3.3">1</cn></apply></apply><ci id="Sx2.Ex2.m1.1.1.3.1.3.cmml" xref="Sx2.Ex2.m1.1.1.3.1.3">𝐿</ci></apply><apply id="Sx2.Ex2.m1.1.1.3.2.cmml" xref="Sx2.Ex2.m1.1.1.3.2"><csymbol cd="ambiguous" id="Sx2.Ex2.m1.1.1.3.2.1.cmml" xref="Sx2.Ex2.m1.1.1.3.2">subscript</csymbol><ci id="Sx2.Ex2.m1.1.1.3.2.2.cmml" xref="Sx2.Ex2.m1.1.1.3.2.2">𝑆</ci><apply id="Sx2.Ex2.m1.1.1.3.2.3.cmml" xref="Sx2.Ex2.m1.1.1.3.2.3"><times id="Sx2.Ex2.m1.1.1.3.2.3.1.cmml" xref="Sx2.Ex2.m1.1.1.3.2.3.1"></times><ci id="Sx2.Ex2.m1.1.1.3.2.3.2.cmml" xref="Sx2.Ex2.m1.1.1.3.2.3.2">𝑖</ci><ci id="Sx2.Ex2.m1.1.1.3.2.3.3.cmml" xref="Sx2.Ex2.m1.1.1.3.2.3.3">𝑗</ci><ci id="Sx2.Ex2.m1.1.1.3.2.3.4.cmml" xref="Sx2.Ex2.m1.1.1.3.2.3.4">𝑝</ci></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.Ex2.m1.1c">\text{argmax}_{p}\sum_{j=1}^{L}S_{ijp}</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
</tr></tbody>
</table>
<p id="Sx2.SS2.SSS1.p4.10" class="ltx_p">as the target and <math id="Sx2.SS2.SSS1.p4.10.m1.1" class="ltx_Math" alttext="g(\text{document}_{i})" display="inline"><semantics id="Sx2.SS2.SSS1.p4.10.m1.1a"><mrow id="Sx2.SS2.SSS1.p4.10.m1.1.1" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.cmml"><mi id="Sx2.SS2.SSS1.p4.10.m1.1.1.3" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.3.cmml">g</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p4.10.m1.1.1.2" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.2.cmml">​</mo><mrow id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.cmml"><mo stretchy="false" id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.2" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.cmml">(</mo><msub id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.cmml"><mtext id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2a.cmml">document</mtext><mi id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.3" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.3.cmml">i</mi></msub><mo stretchy="false" id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.3" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.cmml">)</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p4.10.m1.1b"><apply id="Sx2.SS2.SSS1.p4.10.m1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1"><times id="Sx2.SS2.SSS1.p4.10.m1.1.1.2.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.2"></times><ci id="Sx2.SS2.SSS1.p4.10.m1.1.1.3.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.3">𝑔</ci><apply id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.1.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1">subscript</csymbol><ci id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2a.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2"><mtext id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.2">document</mtext></ci><ci id="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.3.cmml" xref="Sx2.SS2.SSS1.p4.10.m1.1.1.1.1.1.3">𝑖</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p4.10.m1.1c">g(\text{document}_{i})</annotation></semantics></math> as the feature.</p>
</div>
<div id="Sx2.SS2.SSS1.p5" class="ltx_para">
<p id="Sx2.SS2.SSS1.p5.1" class="ltx_p">With large numbers of paths, we saw that the paths that were assigned the fewest documents via <math id="Sx2.SS2.SSS1.p5.1.m1.1" class="ltx_Math" alttext="\text{argmax}_{p}\sum_{j}S_{ijp}" display="inline"><semantics id="Sx2.SS2.SSS1.p5.1.m1.1a"><mrow id="Sx2.SS2.SSS1.p5.1.m1.1.1" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.cmml"><msub id="Sx2.SS2.SSS1.p5.1.m1.1.1.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.cmml"><mtext id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2a.cmml">argmax</mtext><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.3" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.3.cmml">p</mi></msub><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p5.1.m1.1.1.1" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.1.cmml">​</mo><mrow id="Sx2.SS2.SSS1.p5.1.m1.1.1.3" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.cmml"><msub id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.cmml"><mo id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.2.cmml">∑</mo><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.3" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.3.cmml">j</mi></msub><msub id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.cmml"><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.2.cmml">S</mi><mrow id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.cmml"><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.2" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.3" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1a" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.4" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.4.cmml">p</mi></mrow></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS1.p5.1.m1.1b"><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1"><times id="Sx2.SS2.SSS1.p5.1.m1.1.1.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.1"></times><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2">subscript</csymbol><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2a.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2"><mtext id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.2">argmax</mtext></ci><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.2.3.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.2.3">𝑝</ci></apply><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3"><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1">subscript</csymbol><sum id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.2"></sum><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.3.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.1.3">𝑗</ci></apply><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2"><csymbol cd="ambiguous" id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2">subscript</csymbol><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.2">𝑆</ci><apply id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3"><times id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.1"></times><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.2.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.2">𝑖</ci><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.3.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.3">𝑗</ci><ci id="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.4.cmml" xref="Sx2.SS2.SSS1.p5.1.m1.1.1.3.2.3.4">𝑝</ci></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS1.p5.1.m1.1c">\text{argmax}_{p}\sum_{j}S_{ijp}</annotation></semantics></math> were even more under-represented by the output of the trained logistic regressor, and it would often happen that a path was empty when using the output of the regressor. To remedy this, we trained a bias term to match the target document-to-path distribution.</p>
</div>
<div id="Sx2.SS2.SSS1.p6" class="ltx_para">
<p id="Sx2.SS2.SSS1.p6.1" class="ltx_p">We find that</p>
<ul id="Sx2.I3" class="ltx_itemize">
<li id="Sx2.I3.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="Sx2.I3.i1.p1" class="ltx_para">
<p id="Sx2.I3.i1.p1.1" class="ltx_p">One step of training the discriminative router leads to a significant improvement in validation perplexity, more alternating steps lead to further minor improvements.</p>
</div>
</li>
<li id="Sx2.I3.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="Sx2.I3.i2.p1" class="ltx_para">
<p id="Sx2.I3.i2.p1.1" class="ltx_p">Gains are larger with more paths.</p>
</div>
</li>
</ul>
</div>
<figure id="Sx2.F10" class="ltx_figure"><img src="/html/2403.10616/assets/number_of_experts.png" id="Sx2.F10.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="580" height="190" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="Sx2.F10.2.1.1" class="ltx_text" style="font-size:90%;">Figure 10</span>: </span><span id="Sx2.F10.3.2" class="ltx_text" style="font-size:90%;">Generative and discriminative flat MoE models with varying number of paths, three alternating discriminative phases. The “branching” structure of the figure is due to the ancestry of the models: at the far left, all share an initial dense ancestor. At 16K steps, each discriminative-generative pair of experiments shares a generative ancestor, this is continued (restarting the cosine learning rate decay) with the same router (“generative”) or with a sequence of new routers (“discriminative”)</span></figcaption>
</figure>
<figure id="Sx2.F11" class="ltx_figure"><img src="/html/2403.10616/assets/num_steps.png" id="Sx2.F11.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="580" height="190" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="Sx2.F11.4.2.1" class="ltx_text" style="font-size:90%;">Figure 11</span>: </span><span id="Sx2.F11.2.1" class="ltx_text" style="font-size:90%;">Validation perplexity vs number of alternating minimization steps, 16 path flat MoE. As the number of alternating steps increases, the results improve, but each step gives less improvement: <math id="Sx2.F11.2.1.m1.1" class="ltx_Math" alttext="14.0\rightarrow 13.38\rightarrow 13.36\rightarrow 13.25" display="inline"><semantics id="Sx2.F11.2.1.m1.1b"><mrow id="Sx2.F11.2.1.m1.1.1" xref="Sx2.F11.2.1.m1.1.1.cmml"><mn id="Sx2.F11.2.1.m1.1.1.2" xref="Sx2.F11.2.1.m1.1.1.2.cmml">14.0</mn><mo stretchy="false" id="Sx2.F11.2.1.m1.1.1.3" xref="Sx2.F11.2.1.m1.1.1.3.cmml">→</mo><mn id="Sx2.F11.2.1.m1.1.1.4" xref="Sx2.F11.2.1.m1.1.1.4.cmml">13.38</mn><mo stretchy="false" id="Sx2.F11.2.1.m1.1.1.5" xref="Sx2.F11.2.1.m1.1.1.5.cmml">→</mo><mn id="Sx2.F11.2.1.m1.1.1.6" xref="Sx2.F11.2.1.m1.1.1.6.cmml">13.36</mn><mo stretchy="false" id="Sx2.F11.2.1.m1.1.1.7" xref="Sx2.F11.2.1.m1.1.1.7.cmml">→</mo><mn id="Sx2.F11.2.1.m1.1.1.8" xref="Sx2.F11.2.1.m1.1.1.8.cmml">13.25</mn></mrow><annotation-xml encoding="MathML-Content" id="Sx2.F11.2.1.m1.1c"><apply id="Sx2.F11.2.1.m1.1.1.cmml" xref="Sx2.F11.2.1.m1.1.1"><and id="Sx2.F11.2.1.m1.1.1a.cmml" xref="Sx2.F11.2.1.m1.1.1"></and><apply id="Sx2.F11.2.1.m1.1.1b.cmml" xref="Sx2.F11.2.1.m1.1.1"><ci id="Sx2.F11.2.1.m1.1.1.3.cmml" xref="Sx2.F11.2.1.m1.1.1.3">→</ci><cn type="float" id="Sx2.F11.2.1.m1.1.1.2.cmml" xref="Sx2.F11.2.1.m1.1.1.2">14.0</cn><cn type="float" id="Sx2.F11.2.1.m1.1.1.4.cmml" xref="Sx2.F11.2.1.m1.1.1.4">13.38</cn></apply><apply id="Sx2.F11.2.1.m1.1.1c.cmml" xref="Sx2.F11.2.1.m1.1.1"><ci id="Sx2.F11.2.1.m1.1.1.5.cmml" xref="Sx2.F11.2.1.m1.1.1.5">→</ci><share href="#Sx2.F11.2.1.m1.1.1.4.cmml" id="Sx2.F11.2.1.m1.1.1d.cmml" xref="Sx2.F11.2.1.m1.1.1"></share><cn type="float" id="Sx2.F11.2.1.m1.1.1.6.cmml" xref="Sx2.F11.2.1.m1.1.1.6">13.36</cn></apply><apply id="Sx2.F11.2.1.m1.1.1e.cmml" xref="Sx2.F11.2.1.m1.1.1"><ci id="Sx2.F11.2.1.m1.1.1.7.cmml" xref="Sx2.F11.2.1.m1.1.1.7">→</ci><share href="#Sx2.F11.2.1.m1.1.1.6.cmml" id="Sx2.F11.2.1.m1.1.1f.cmml" xref="Sx2.F11.2.1.m1.1.1"></share><cn type="float" id="Sx2.F11.2.1.m1.1.1.8.cmml" xref="Sx2.F11.2.1.m1.1.1.8">13.25</cn></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.F11.2.1.m1.1d">14.0\rightarrow 13.38\rightarrow 13.36\rightarrow 13.25</annotation></semantics></math> PPL for phases 0, 1, 2, 3 respectively. All discriminative routing results in the main text use one alternating minimization phase.</span></figcaption>
</figure>
</section>
<section id="Sx2.SS2.SSS2" class="ltx_subsubsection">
<h4 class="ltx_title ltx_title_subsubsection">
<span class="ltx_tag ltx_tag_subsubsection">7.2.2 </span>More frequent routing during evaluation</h4>

<div id="Sx2.SS2.SSS2.p1" class="ltx_para">
<p id="Sx2.SS2.SSS2.p1.4" class="ltx_p">As before, we have the array <math id="Sx2.SS2.SSS2.p1.1.m1.1" class="ltx_Math" alttext="S_{ijp}" display="inline"><semantics id="Sx2.SS2.SSS2.p1.1.m1.1a"><msub id="Sx2.SS2.SSS2.p1.1.m1.1.1" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.cmml"><mi id="Sx2.SS2.SSS2.p1.1.m1.1.1.2" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.2.cmml">S</mi><mrow id="Sx2.SS2.SSS2.p1.1.m1.1.1.3" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.cmml"><mi id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.2" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.3" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1a" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.4" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.4.cmml">p</mi></mrow></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.1.m1.1b"><apply id="Sx2.SS2.SSS2.p1.1.m1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1"><csymbol cd="ambiguous" id="Sx2.SS2.SSS2.p1.1.m1.1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1">subscript</csymbol><ci id="Sx2.SS2.SSS2.p1.1.m1.1.1.2.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.2">𝑆</ci><apply id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3"><times id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.1"></times><ci id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.2.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.2">𝑖</ci><ci id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.3.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.3">𝑗</ci><ci id="Sx2.SS2.SSS2.p1.1.m1.1.1.3.4.cmml" xref="Sx2.SS2.SSS2.p1.1.m1.1.1.3.4">𝑝</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.1.m1.1c">S_{ijp}</annotation></semantics></math> giving the scores of each token in the router data set of documents. We fix a number <math id="Sx2.SS2.SSS2.p1.2.m2.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS2.p1.2.m2.1a"><mi id="Sx2.SS2.SSS2.p1.2.m2.1.1" xref="Sx2.SS2.SSS2.p1.2.m2.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.2.m2.1b"><ci id="Sx2.SS2.SSS2.p1.2.m2.1.1.cmml" xref="Sx2.SS2.SSS2.p1.2.m2.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.2.m2.1c">L</annotation></semantics></math>, and build a sequence transduction training set where the inputs are the token sequences <math id="Sx2.SS2.SSS2.p1.3.m3.2" class="ltx_Math" alttext="t_{i,j}" display="inline"><semantics id="Sx2.SS2.SSS2.p1.3.m3.2a"><msub id="Sx2.SS2.SSS2.p1.3.m3.2.3" xref="Sx2.SS2.SSS2.p1.3.m3.2.3.cmml"><mi id="Sx2.SS2.SSS2.p1.3.m3.2.3.2" xref="Sx2.SS2.SSS2.p1.3.m3.2.3.2.cmml">t</mi><mrow id="Sx2.SS2.SSS2.p1.3.m3.2.2.2.4" xref="Sx2.SS2.SSS2.p1.3.m3.2.2.2.3.cmml"><mi id="Sx2.SS2.SSS2.p1.3.m3.1.1.1.1" xref="Sx2.SS2.SSS2.p1.3.m3.1.1.1.1.cmml">i</mi><mo id="Sx2.SS2.SSS2.p1.3.m3.2.2.2.4.1" xref="Sx2.SS2.SSS2.p1.3.m3.2.2.2.3.cmml">,</mo><mi id="Sx2.SS2.SSS2.p1.3.m3.2.2.2.2" xref="Sx2.SS2.SSS2.p1.3.m3.2.2.2.2.cmml">j</mi></mrow></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.3.m3.2b"><apply id="Sx2.SS2.SSS2.p1.3.m3.2.3.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.2.3"><csymbol cd="ambiguous" id="Sx2.SS2.SSS2.p1.3.m3.2.3.1.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.2.3">subscript</csymbol><ci id="Sx2.SS2.SSS2.p1.3.m3.2.3.2.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.2.3.2">𝑡</ci><list id="Sx2.SS2.SSS2.p1.3.m3.2.2.2.3.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.2.2.2.4"><ci id="Sx2.SS2.SSS2.p1.3.m3.1.1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.1.1.1.1">𝑖</ci><ci id="Sx2.SS2.SSS2.p1.3.m3.2.2.2.2.cmml" xref="Sx2.SS2.SSS2.p1.3.m3.2.2.2.2">𝑗</ci></list></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.3.m3.2c">t_{i,j}</annotation></semantics></math> from the reserved portion of the training documents, and the outputs <math id="Sx2.SS2.SSS2.p1.4.m4.1" class="ltx_Math" alttext="T=T_{ij}" display="inline"><semantics id="Sx2.SS2.SSS2.p1.4.m4.1a"><mrow id="Sx2.SS2.SSS2.p1.4.m4.1.1" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.cmml"><mi id="Sx2.SS2.SSS2.p1.4.m4.1.1.2" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.2.cmml">T</mi><mo id="Sx2.SS2.SSS2.p1.4.m4.1.1.1" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.1.cmml">=</mo><msub id="Sx2.SS2.SSS2.p1.4.m4.1.1.3" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.cmml"><mi id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.2" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.2.cmml">T</mi><mrow id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.cmml"><mi id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.2" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.1" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.1.cmml">​</mo><mi id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.3" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.3.cmml">j</mi></mrow></msub></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.4.m4.1b"><apply id="Sx2.SS2.SSS2.p1.4.m4.1.1.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1"><eq id="Sx2.SS2.SSS2.p1.4.m4.1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.1"></eq><ci id="Sx2.SS2.SSS2.p1.4.m4.1.1.2.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.2">𝑇</ci><apply id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3"><csymbol cd="ambiguous" id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.1.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3">subscript</csymbol><ci id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.2.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.2">𝑇</ci><apply id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3"><times id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.1.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.1"></times><ci id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.2.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.2">𝑖</ci><ci id="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.3.cmml" xref="Sx2.SS2.SSS2.p1.4.m4.1.1.3.3.3">𝑗</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.4.m4.1c">T=T_{ij}</annotation></semantics></math> are
defined by</p>
<table id="Sx2.Ex3" class="ltx_equation ltx_eqn_table">

<tbody><tr class="ltx_equation ltx_eqn_row ltx_align_baseline">
<td class="ltx_eqn_cell ltx_eqn_center_padleft"></td>
<td class="ltx_eqn_cell ltx_align_center"><math id="Sx2.Ex3.m1.1" class="ltx_Math" alttext="T_{ij}=\text{argmax}_{p}\sum_{j}^{j^{\prime}}S_{ijp}," display="block"><semantics id="Sx2.Ex3.m1.1a"><mrow id="Sx2.Ex3.m1.1.1.1" xref="Sx2.Ex3.m1.1.1.1.1.cmml"><mrow id="Sx2.Ex3.m1.1.1.1.1" xref="Sx2.Ex3.m1.1.1.1.1.cmml"><msub id="Sx2.Ex3.m1.1.1.1.1.2" xref="Sx2.Ex3.m1.1.1.1.1.2.cmml"><mi id="Sx2.Ex3.m1.1.1.1.1.2.2" xref="Sx2.Ex3.m1.1.1.1.1.2.2.cmml">T</mi><mrow id="Sx2.Ex3.m1.1.1.1.1.2.3" xref="Sx2.Ex3.m1.1.1.1.1.2.3.cmml"><mi id="Sx2.Ex3.m1.1.1.1.1.2.3.2" xref="Sx2.Ex3.m1.1.1.1.1.2.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex3.m1.1.1.1.1.2.3.1" xref="Sx2.Ex3.m1.1.1.1.1.2.3.1.cmml">​</mo><mi id="Sx2.Ex3.m1.1.1.1.1.2.3.3" xref="Sx2.Ex3.m1.1.1.1.1.2.3.3.cmml">j</mi></mrow></msub><mo id="Sx2.Ex3.m1.1.1.1.1.1" xref="Sx2.Ex3.m1.1.1.1.1.1.cmml">=</mo><mrow id="Sx2.Ex3.m1.1.1.1.1.3" xref="Sx2.Ex3.m1.1.1.1.1.3.cmml"><msub id="Sx2.Ex3.m1.1.1.1.1.3.2" xref="Sx2.Ex3.m1.1.1.1.1.3.2.cmml"><mtext id="Sx2.Ex3.m1.1.1.1.1.3.2.2" xref="Sx2.Ex3.m1.1.1.1.1.3.2.2a.cmml">argmax</mtext><mi id="Sx2.Ex3.m1.1.1.1.1.3.2.3" xref="Sx2.Ex3.m1.1.1.1.1.3.2.3.cmml">p</mi></msub><mo lspace="0em" rspace="0em" id="Sx2.Ex3.m1.1.1.1.1.3.1" xref="Sx2.Ex3.m1.1.1.1.1.3.1.cmml">​</mo><mrow id="Sx2.Ex3.m1.1.1.1.1.3.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.cmml"><munderover id="Sx2.Ex3.m1.1.1.1.1.3.3.1" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.cmml"><mo movablelimits="false" id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.2" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.2.cmml">∑</mo><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.3.cmml">j</mi><msup id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.cmml"><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.2" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.2.cmml">j</mi><mo id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.3.cmml">′</mo></msup></munderover><msub id="Sx2.Ex3.m1.1.1.1.1.3.3.2" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.cmml"><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.2.2" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.2.cmml">S</mi><mrow id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.cmml"><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.2" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.2.cmml">i</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1.cmml">​</mo><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.3" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.3.cmml">j</mi><mo lspace="0em" rspace="0em" id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1a" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1.cmml">​</mo><mi id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.4" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.4.cmml">p</mi></mrow></msub></mrow></mrow></mrow><mo id="Sx2.Ex3.m1.1.1.1.2" xref="Sx2.Ex3.m1.1.1.1.1.cmml">,</mo></mrow><annotation-xml encoding="MathML-Content" id="Sx2.Ex3.m1.1b"><apply id="Sx2.Ex3.m1.1.1.1.1.cmml" xref="Sx2.Ex3.m1.1.1.1"><eq id="Sx2.Ex3.m1.1.1.1.1.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.1"></eq><apply id="Sx2.Ex3.m1.1.1.1.1.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.2.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2">subscript</csymbol><ci id="Sx2.Ex3.m1.1.1.1.1.2.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2.2">𝑇</ci><apply id="Sx2.Ex3.m1.1.1.1.1.2.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2.3"><times id="Sx2.Ex3.m1.1.1.1.1.2.3.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2.3.1"></times><ci id="Sx2.Ex3.m1.1.1.1.1.2.3.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2.3.2">𝑖</ci><ci id="Sx2.Ex3.m1.1.1.1.1.2.3.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.2.3.3">𝑗</ci></apply></apply><apply id="Sx2.Ex3.m1.1.1.1.1.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3"><times id="Sx2.Ex3.m1.1.1.1.1.3.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.1"></times><apply id="Sx2.Ex3.m1.1.1.1.1.3.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.2"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.3.2.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.2">subscript</csymbol><ci id="Sx2.Ex3.m1.1.1.1.1.3.2.2a.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.2.2"><mtext id="Sx2.Ex3.m1.1.1.1.1.3.2.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.2.2">argmax</mtext></ci><ci id="Sx2.Ex3.m1.1.1.1.1.3.2.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.2.3">𝑝</ci></apply><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3"><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.3.3.1.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1">superscript</csymbol><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1">subscript</csymbol><sum id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.2"></sum><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.2.3">𝑗</ci></apply><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3">superscript</csymbol><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.2">𝑗</ci><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.1.3.3">′</ci></apply></apply><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2"><csymbol cd="ambiguous" id="Sx2.Ex3.m1.1.1.1.1.3.3.2.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2">subscript</csymbol><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.2.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.2">𝑆</ci><apply id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3"><times id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.1"></times><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.2.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.2">𝑖</ci><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.3.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.3">𝑗</ci><ci id="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.4.cmml" xref="Sx2.Ex3.m1.1.1.1.1.3.3.2.3.4">𝑝</ci></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.Ex3.m1.1c">T_{ij}=\text{argmax}_{p}\sum_{j}^{j^{\prime}}S_{ijp},</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
</tr></tbody>
</table>
<p id="Sx2.SS2.SSS2.p1.9" class="ltx_p">where <math id="Sx2.SS2.SSS2.p1.5.m1.3" class="ltx_Math" alttext="j^{\prime}=\min(\text{length of document},j+L-1)" display="inline"><semantics id="Sx2.SS2.SSS2.p1.5.m1.3a"><mrow id="Sx2.SS2.SSS2.p1.5.m1.3.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.cmml"><msup id="Sx2.SS2.SSS2.p1.5.m1.3.3.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3.cmml"><mi id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.2" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3.2.cmml">j</mi><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3.3.cmml">′</mo></msup><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.2" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.2.cmml">=</mo><mrow id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml"><mi id="Sx2.SS2.SSS2.p1.5.m1.1.1" xref="Sx2.SS2.SSS2.p1.5.m1.1.1.cmml">min</mi><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1a" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml">⁡</mo><mrow id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml"><mo stretchy="false" id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.2" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml">(</mo><mtext id="Sx2.SS2.SSS2.p1.5.m1.2.2" xref="Sx2.SS2.SSS2.p1.5.m1.2.2a.cmml">length of document</mtext><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml">,</mo><mrow id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.cmml"><mrow id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.cmml"><mi id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.2" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.2.cmml">j</mi><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.1" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.1.cmml">+</mo><mi id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.3.cmml">L</mi></mrow><mo id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.1" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.1.cmml">−</mo><mn id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.3" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.3.cmml">1</mn></mrow><mo stretchy="false" id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.4" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml">)</mo></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.5.m1.3b"><apply id="Sx2.SS2.SSS2.p1.5.m1.3.3.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3"><eq id="Sx2.SS2.SSS2.p1.5.m1.3.3.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.2"></eq><apply id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3"><csymbol cd="ambiguous" id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.1.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3">superscript</csymbol><ci id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3.2">𝑗</ci><ci id="Sx2.SS2.SSS2.p1.5.m1.3.3.3.3.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.3.3">′</ci></apply><apply id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1"><min id="Sx2.SS2.SSS2.p1.5.m1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.1.1"></min><ci id="Sx2.SS2.SSS2.p1.5.m1.2.2a.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.2.2"><mtext id="Sx2.SS2.SSS2.p1.5.m1.2.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.2.2">length of document</mtext></ci><apply id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1"><minus id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.1.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.1"></minus><apply id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2"><plus id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.1.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.1"></plus><ci id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.2.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.2">𝑗</ci><ci id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.3.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.2.3">𝐿</ci></apply><cn type="integer" id="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.3.cmml" xref="Sx2.SS2.SSS2.p1.5.m1.3.3.1.1.1.1.3">1</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.5.m1.3c">j^{\prime}=\min(\text{length of document},j+L-1)</annotation></semantics></math>, and <math id="Sx2.SS2.SSS2.p1.6.m2.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS2.p1.6.m2.1a"><mi id="Sx2.SS2.SSS2.p1.6.m2.1.1" xref="Sx2.SS2.SSS2.p1.6.m2.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.6.m2.1b"><ci id="Sx2.SS2.SSS2.p1.6.m2.1.1.cmml" xref="Sx2.SS2.SSS2.p1.6.m2.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.6.m2.1c">L</annotation></semantics></math> is a window size. Once we have <math id="Sx2.SS2.SSS2.p1.7.m3.1" class="ltx_Math" alttext="T" display="inline"><semantics id="Sx2.SS2.SSS2.p1.7.m3.1a"><mi id="Sx2.SS2.SSS2.p1.7.m3.1.1" xref="Sx2.SS2.SSS2.p1.7.m3.1.1.cmml">T</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.7.m3.1b"><ci id="Sx2.SS2.SSS2.p1.7.m3.1.1.cmml" xref="Sx2.SS2.SSS2.p1.7.m3.1.1">𝑇</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.7.m3.1c">T</annotation></semantics></math>, we can train a transformer to transduce <math id="Sx2.SS2.SSS2.p1.8.m4.1" class="ltx_Math" alttext="T" display="inline"><semantics id="Sx2.SS2.SSS2.p1.8.m4.1a"><mi id="Sx2.SS2.SSS2.p1.8.m4.1.1" xref="Sx2.SS2.SSS2.p1.8.m4.1.1.cmml">T</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.8.m4.1b"><ci id="Sx2.SS2.SSS2.p1.8.m4.1.1.cmml" xref="Sx2.SS2.SSS2.p1.8.m4.1.1">𝑇</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.8.m4.1c">T</annotation></semantics></math> from <math id="Sx2.SS2.SSS2.p1.9.m5.1" class="ltx_Math" alttext="t" display="inline"><semantics id="Sx2.SS2.SSS2.p1.9.m5.1a"><mi id="Sx2.SS2.SSS2.p1.9.m5.1.1" xref="Sx2.SS2.SSS2.p1.9.m5.1.1.cmml">t</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p1.9.m5.1b"><ci id="Sx2.SS2.SSS2.p1.9.m5.1.1.cmml" xref="Sx2.SS2.SSS2.p1.9.m5.1.1">𝑡</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p1.9.m5.1c">t</annotation></semantics></math>. We finetune this router from the same model that we fork to train the paths, using the same hyperparameters as our standard LM training (except no warmup), and it converges after a 2K steps. The learned router can then make a decision at any token.</p>
</div>
<div id="Sx2.SS2.SSS2.p2" class="ltx_para">
<p id="Sx2.SS2.SSS2.p2.4" class="ltx_p">We found that choosing <math id="Sx2.SS2.SSS2.p2.1.m1.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS2.p2.1.m1.1a"><mi id="Sx2.SS2.SSS2.p2.1.m1.1.1" xref="Sx2.SS2.SSS2.p2.1.m1.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p2.1.m1.1b"><ci id="Sx2.SS2.SSS2.p2.1.m1.1.1.cmml" xref="Sx2.SS2.SSS2.p2.1.m1.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p2.1.m1.1c">L</annotation></semantics></math> to be the same size as the number of tokens before re-routing (as in Table <a href="#S4.T3" title="Table 3 ‣ 4.4 Routing frequency during evaluation ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>) led to the best results, but was only marginally better than choosing <math id="Sx2.SS2.SSS2.p2.2.m2.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS2.p2.2.m2.1a"><mi id="Sx2.SS2.SSS2.p2.2.m2.1.1" xref="Sx2.SS2.SSS2.p2.2.m2.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p2.2.m2.1b"><ci id="Sx2.SS2.SSS2.p2.2.m2.1.1.cmml" xref="Sx2.SS2.SSS2.p2.2.m2.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p2.2.m2.1c">L</annotation></semantics></math> large enough so the window always went to the end of the document. The results in Table <a href="#S4.T3" title="Table 3 ‣ 4.4 Routing frequency during evaluation ‣ 4 Experiments ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a> use this choice of <math id="Sx2.SS2.SSS2.p2.3.m3.1" class="ltx_Math" alttext="L" display="inline"><semantics id="Sx2.SS2.SSS2.p2.3.m3.1a"><mi id="Sx2.SS2.SSS2.p2.3.m3.1.1" xref="Sx2.SS2.SSS2.p2.3.m3.1.1.cmml">L</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p2.3.m3.1b"><ci id="Sx2.SS2.SSS2.p2.3.m3.1.1.cmml" xref="Sx2.SS2.SSS2.p2.3.m3.1.1">𝐿</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p2.3.m3.1c">L</annotation></semantics></math> (the length of the whole sequence, <math id="Sx2.SS2.SSS2.p2.4.m4.1" class="ltx_Math" alttext="1024" display="inline"><semantics id="Sx2.SS2.SSS2.p2.4.m4.1a"><mn id="Sx2.SS2.SSS2.p2.4.m4.1.1" xref="Sx2.SS2.SSS2.p2.4.m4.1.1.cmml">1024</mn><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p2.4.m4.1b"><cn type="integer" id="Sx2.SS2.SSS2.p2.4.m4.1.1.cmml" xref="Sx2.SS2.SSS2.p2.4.m4.1.1">1024</cn></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p2.4.m4.1c">1024</annotation></semantics></math>). We then can use the same learned router for any choice of evaluation-time frequency.</p>
</div>
<div id="Sx2.SS2.SSS2.p3" class="ltx_para">
<p id="Sx2.SS2.SSS2.p3.3" class="ltx_p">In <a href="#Sx2.T5" title="Table 5 ‣ 7.3 More sophisticated generative routing ‣ Supplementary Materials ‣ DiPaCo: Distributed Path Composition" class="ltx_ref ltx_refmacro_autoref"><span class="ltx_text ltx_ref_tag">Table 5</span></a> we show the effect of different <span id="Sx2.SS2.SSS2.p3.3.1" class="ltx_text ltx_font_italic">sharding methods</span> on a <math id="Sx2.SS2.SSS2.p3.1.m1.1" class="ltx_Math" alttext="8\times 8" display="inline"><semantics id="Sx2.SS2.SSS2.p3.1.m1.1a"><mrow id="Sx2.SS2.SSS2.p3.1.m1.1.1" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.cmml"><mn id="Sx2.SS2.SSS2.p3.1.m1.1.1.2" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.2.cmml">8</mn><mo lspace="0.222em" rspace="0.222em" id="Sx2.SS2.SSS2.p3.1.m1.1.1.1" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.1.cmml">×</mo><mn id="Sx2.SS2.SSS2.p3.1.m1.1.1.3" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.3.cmml">8</mn></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p3.1.m1.1b"><apply id="Sx2.SS2.SSS2.p3.1.m1.1.1.cmml" xref="Sx2.SS2.SSS2.p3.1.m1.1.1"><times id="Sx2.SS2.SSS2.p3.1.m1.1.1.1.cmml" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.1"></times><cn type="integer" id="Sx2.SS2.SSS2.p3.1.m1.1.1.2.cmml" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.2">8</cn><cn type="integer" id="Sx2.SS2.SSS2.p3.1.m1.1.1.3.cmml" xref="Sx2.SS2.SSS2.p3.1.m1.1.1.3">8</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p3.1.m1.1c">8\times 8</annotation></semantics></math> DiPaCo. We observe that sharding makes a significant difference, with the discriminative variant yielding the best generalization, with an absolute gain of <math id="Sx2.SS2.SSS2.p3.2.m2.1" class="ltx_Math" alttext="0.7" display="inline"><semantics id="Sx2.SS2.SSS2.p3.2.m2.1a"><mn id="Sx2.SS2.SSS2.p3.2.m2.1.1" xref="Sx2.SS2.SSS2.p3.2.m2.1.1.cmml">0.7</mn><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p3.2.m2.1b"><cn type="float" id="Sx2.SS2.SSS2.p3.2.m2.1.1.cmml" xref="Sx2.SS2.SSS2.p3.2.m2.1.1">0.7</cn></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p3.2.m2.1c">0.7</annotation></semantics></math> perplexity points over the generative <math id="Sx2.SS2.SSS2.p3.3.m3.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS2.SSS2.p3.3.m3.1a"><mi id="Sx2.SS2.SSS2.p3.3.m3.1.1" xref="Sx2.SS2.SSS2.p3.3.m3.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS2.SSS2.p3.3.m3.1b"><ci id="Sx2.SS2.SSS2.p3.3.m3.1.1.cmml" xref="Sx2.SS2.SSS2.p3.3.m3.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS2.SSS2.p3.3.m3.1c">k</annotation></semantics></math>-means variant. We found that longer training runs yield even greater gain from the discriminative gating.</p>
</div>
</section>
</section>
<section id="Sx2.SS3" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">7.3 </span>More sophisticated generative routing</h3>

<figure id="Sx2.T5" class="ltx_table">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="Sx2.T5.9.3.1" class="ltx_text" style="font-size:90%;">Table 5</span>: </span><span id="Sx2.T5.4.2" class="ltx_text ltx_font_bold" style="font-size:90%;">Sharding Impact on 8x8 DiPaCo<span id="Sx2.T5.4.2.2" class="ltx_text ltx_font_medium">: Validation perplexity after 32 outer optimization steps, each consisting of 62 inner optimization steps with a 8x8 DiPaCo with <math id="Sx2.T5.3.1.1.m1.1" class="ltx_Math" alttext="P=64" display="inline"><semantics id="Sx2.T5.3.1.1.m1.1b"><mrow id="Sx2.T5.3.1.1.m1.1.1" xref="Sx2.T5.3.1.1.m1.1.1.cmml"><mi id="Sx2.T5.3.1.1.m1.1.1.2" xref="Sx2.T5.3.1.1.m1.1.1.2.cmml">P</mi><mo id="Sx2.T5.3.1.1.m1.1.1.1" xref="Sx2.T5.3.1.1.m1.1.1.1.cmml">=</mo><mn id="Sx2.T5.3.1.1.m1.1.1.3" xref="Sx2.T5.3.1.1.m1.1.1.3.cmml">64</mn></mrow><annotation-xml encoding="MathML-Content" id="Sx2.T5.3.1.1.m1.1c"><apply id="Sx2.T5.3.1.1.m1.1.1.cmml" xref="Sx2.T5.3.1.1.m1.1.1"><eq id="Sx2.T5.3.1.1.m1.1.1.1.cmml" xref="Sx2.T5.3.1.1.m1.1.1.1"></eq><ci id="Sx2.T5.3.1.1.m1.1.1.2.cmml" xref="Sx2.T5.3.1.1.m1.1.1.2">𝑃</ci><cn type="integer" id="Sx2.T5.3.1.1.m1.1.1.3.cmml" xref="Sx2.T5.3.1.1.m1.1.1.3">64</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T5.3.1.1.m1.1d">P=64</annotation></semantics></math>. Longer training further increases the gap between methods. The discriminative method is based on product <math id="Sx2.T5.4.2.2.m2.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.T5.4.2.2.m2.1b"><mi id="Sx2.T5.4.2.2.m2.1.1" xref="Sx2.T5.4.2.2.m2.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.T5.4.2.2.m2.1c"><ci id="Sx2.T5.4.2.2.m2.1.1.cmml" xref="Sx2.T5.4.2.2.m2.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T5.4.2.2.m2.1d">k</annotation></semantics></math>-Means.</span></span></figcaption>
<table id="Sx2.T5.6" class="ltx_tabular ltx_centering ltx_guessed_headers ltx_align_middle">
<thead class="ltx_thead">
<tr id="Sx2.T5.6.3.1" class="ltx_tr">
<th id="Sx2.T5.6.3.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_column ltx_th_row ltx_border_r ltx_border_tt">Sharding</th>
<th id="Sx2.T5.6.3.1.2" class="ltx_td ltx_nopad_r ltx_align_center ltx_th ltx_th_column ltx_border_tt">Validation Perplexity</th>
</tr>
</thead>
<tbody class="ltx_tbody">
<tr id="Sx2.T5.5.1" class="ltx_tr">
<th id="Sx2.T5.5.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r ltx_border_t">
<math id="Sx2.T5.5.1.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.T5.5.1.1.m1.1a"><mi id="Sx2.T5.5.1.1.m1.1.1" xref="Sx2.T5.5.1.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.T5.5.1.1.m1.1b"><ci id="Sx2.T5.5.1.1.m1.1.1.cmml" xref="Sx2.T5.5.1.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T5.5.1.1.m1.1c">k</annotation></semantics></math>-Means</th>
<td id="Sx2.T5.5.1.2" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_t">17.2</td>
</tr>
<tr id="Sx2.T5.6.2" class="ltx_tr">
<th id="Sx2.T5.6.2.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_r">Product <math id="Sx2.T5.6.2.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.T5.6.2.1.m1.1a"><mi id="Sx2.T5.6.2.1.m1.1.1" xref="Sx2.T5.6.2.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.T5.6.2.1.m1.1b"><ci id="Sx2.T5.6.2.1.m1.1.1.cmml" xref="Sx2.T5.6.2.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.T5.6.2.1.m1.1c">k</annotation></semantics></math>-Means</th>
<td id="Sx2.T5.6.2.2" class="ltx_td ltx_nopad_r ltx_align_center">16.8</td>
</tr>
<tr id="Sx2.T5.6.4.1" class="ltx_tr">
<th id="Sx2.T5.6.4.1.1" class="ltx_td ltx_align_left ltx_th ltx_th_row ltx_border_bb ltx_border_r">Discriminative</th>
<td id="Sx2.T5.6.4.1.2" class="ltx_td ltx_nopad_r ltx_align_center ltx_border_bb">16.5</td>
</tr>
</tbody>
</table>
</figure>
<div id="Sx2.SS3.p1" class="ltx_para">
<p id="Sx2.SS3.p1.19" class="ltx_p">We explored Product <math id="Sx2.SS3.p1.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.1.m1.1a"><mi id="Sx2.SS3.p1.1.m1.1.1" xref="Sx2.SS3.p1.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.1.m1.1b"><ci id="Sx2.SS3.p1.1.m1.1.1.cmml" xref="Sx2.SS3.p1.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.1.m1.1c">k</annotation></semantics></math>-Means as generative router for DiPaCo.
When <math id="Sx2.SS3.p1.2.m2.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.2.m2.1a"><mi id="Sx2.SS3.p1.2.m2.1.1" xref="Sx2.SS3.p1.2.m2.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.2.m2.1b"><ci id="Sx2.SS3.p1.2.m2.1.1.cmml" xref="Sx2.SS3.p1.2.m2.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.2.m2.1c">k</annotation></semantics></math> is large, <math id="Sx2.SS3.p1.3.m3.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.3.m3.1a"><mi id="Sx2.SS3.p1.3.m3.1.1" xref="Sx2.SS3.p1.3.m3.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.3.m3.1b"><ci id="Sx2.SS3.p1.3.m3.1.1.cmml" xref="Sx2.SS3.p1.3.m3.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.3.m3.1c">k</annotation></semantics></math>-Means can become inefficient, because some clusters are assigned few examples. Moreover, in our use case, the structure of the shards does not match the structure of the modules.
Product <math id="Sx2.SS3.p1.4.m4.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.4.m4.1a"><mi id="Sx2.SS3.p1.4.m4.1.1" xref="Sx2.SS3.p1.4.m4.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.4.m4.1b"><ci id="Sx2.SS3.p1.4.m4.1.1.cmml" xref="Sx2.SS3.p1.4.m4.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.4.m4.1c">k</annotation></semantics></math>-Means is a way to mitigate these issues. For a 2 level DiPaCo, the feature <math id="Sx2.SS3.p1.5.m5.1" class="ltx_Math" alttext="z\in\mathcal{R}^{D}" display="inline"><semantics id="Sx2.SS3.p1.5.m5.1a"><mrow id="Sx2.SS3.p1.5.m5.1.1" xref="Sx2.SS3.p1.5.m5.1.1.cmml"><mi id="Sx2.SS3.p1.5.m5.1.1.2" xref="Sx2.SS3.p1.5.m5.1.1.2.cmml">z</mi><mo id="Sx2.SS3.p1.5.m5.1.1.1" xref="Sx2.SS3.p1.5.m5.1.1.1.cmml">∈</mo><msup id="Sx2.SS3.p1.5.m5.1.1.3" xref="Sx2.SS3.p1.5.m5.1.1.3.cmml"><mi class="ltx_font_mathcaligraphic" id="Sx2.SS3.p1.5.m5.1.1.3.2" xref="Sx2.SS3.p1.5.m5.1.1.3.2.cmml">ℛ</mi><mi id="Sx2.SS3.p1.5.m5.1.1.3.3" xref="Sx2.SS3.p1.5.m5.1.1.3.3.cmml">D</mi></msup></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.5.m5.1b"><apply id="Sx2.SS3.p1.5.m5.1.1.cmml" xref="Sx2.SS3.p1.5.m5.1.1"><in id="Sx2.SS3.p1.5.m5.1.1.1.cmml" xref="Sx2.SS3.p1.5.m5.1.1.1"></in><ci id="Sx2.SS3.p1.5.m5.1.1.2.cmml" xref="Sx2.SS3.p1.5.m5.1.1.2">𝑧</ci><apply id="Sx2.SS3.p1.5.m5.1.1.3.cmml" xref="Sx2.SS3.p1.5.m5.1.1.3"><csymbol cd="ambiguous" id="Sx2.SS3.p1.5.m5.1.1.3.1.cmml" xref="Sx2.SS3.p1.5.m5.1.1.3">superscript</csymbol><ci id="Sx2.SS3.p1.5.m5.1.1.3.2.cmml" xref="Sx2.SS3.p1.5.m5.1.1.3.2">ℛ</ci><ci id="Sx2.SS3.p1.5.m5.1.1.3.3.cmml" xref="Sx2.SS3.p1.5.m5.1.1.3.3">𝐷</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.5.m5.1c">z\in\mathcal{R}^{D}</annotation></semantics></math> is divided in two groups: <math id="Sx2.SS3.p1.6.m6.2" class="ltx_Math" alttext="z=[z_{1},z_{2}]" display="inline"><semantics id="Sx2.SS3.p1.6.m6.2a"><mrow id="Sx2.SS3.p1.6.m6.2.2" xref="Sx2.SS3.p1.6.m6.2.2.cmml"><mi id="Sx2.SS3.p1.6.m6.2.2.4" xref="Sx2.SS3.p1.6.m6.2.2.4.cmml">z</mi><mo id="Sx2.SS3.p1.6.m6.2.2.3" xref="Sx2.SS3.p1.6.m6.2.2.3.cmml">=</mo><mrow id="Sx2.SS3.p1.6.m6.2.2.2.2" xref="Sx2.SS3.p1.6.m6.2.2.2.3.cmml"><mo stretchy="false" id="Sx2.SS3.p1.6.m6.2.2.2.2.3" xref="Sx2.SS3.p1.6.m6.2.2.2.3.cmml">[</mo><msub id="Sx2.SS3.p1.6.m6.1.1.1.1.1" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1.cmml"><mi id="Sx2.SS3.p1.6.m6.1.1.1.1.1.2" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1.2.cmml">z</mi><mn id="Sx2.SS3.p1.6.m6.1.1.1.1.1.3" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1.3.cmml">1</mn></msub><mo id="Sx2.SS3.p1.6.m6.2.2.2.2.4" xref="Sx2.SS3.p1.6.m6.2.2.2.3.cmml">,</mo><msub id="Sx2.SS3.p1.6.m6.2.2.2.2.2" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2.cmml"><mi id="Sx2.SS3.p1.6.m6.2.2.2.2.2.2" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2.2.cmml">z</mi><mn id="Sx2.SS3.p1.6.m6.2.2.2.2.2.3" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2.3.cmml">2</mn></msub><mo stretchy="false" id="Sx2.SS3.p1.6.m6.2.2.2.2.5" xref="Sx2.SS3.p1.6.m6.2.2.2.3.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.6.m6.2b"><apply id="Sx2.SS3.p1.6.m6.2.2.cmml" xref="Sx2.SS3.p1.6.m6.2.2"><eq id="Sx2.SS3.p1.6.m6.2.2.3.cmml" xref="Sx2.SS3.p1.6.m6.2.2.3"></eq><ci id="Sx2.SS3.p1.6.m6.2.2.4.cmml" xref="Sx2.SS3.p1.6.m6.2.2.4">𝑧</ci><interval closure="closed" id="Sx2.SS3.p1.6.m6.2.2.2.3.cmml" xref="Sx2.SS3.p1.6.m6.2.2.2.2"><apply id="Sx2.SS3.p1.6.m6.1.1.1.1.1.cmml" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.6.m6.1.1.1.1.1.1.cmml" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1">subscript</csymbol><ci id="Sx2.SS3.p1.6.m6.1.1.1.1.1.2.cmml" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.6.m6.1.1.1.1.1.3.cmml" xref="Sx2.SS3.p1.6.m6.1.1.1.1.1.3">1</cn></apply><apply id="Sx2.SS3.p1.6.m6.2.2.2.2.2.cmml" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2"><csymbol cd="ambiguous" id="Sx2.SS3.p1.6.m6.2.2.2.2.2.1.cmml" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2">subscript</csymbol><ci id="Sx2.SS3.p1.6.m6.2.2.2.2.2.2.cmml" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.6.m6.2.2.2.2.2.3.cmml" xref="Sx2.SS3.p1.6.m6.2.2.2.2.2.3">2</cn></apply></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.6.m6.2c">z=[z_{1},z_{2}]</annotation></semantics></math>, with <math id="Sx2.SS3.p1.7.m7.1" class="ltx_Math" alttext="z_{1}" display="inline"><semantics id="Sx2.SS3.p1.7.m7.1a"><msub id="Sx2.SS3.p1.7.m7.1.1" xref="Sx2.SS3.p1.7.m7.1.1.cmml"><mi id="Sx2.SS3.p1.7.m7.1.1.2" xref="Sx2.SS3.p1.7.m7.1.1.2.cmml">z</mi><mn id="Sx2.SS3.p1.7.m7.1.1.3" xref="Sx2.SS3.p1.7.m7.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.7.m7.1b"><apply id="Sx2.SS3.p1.7.m7.1.1.cmml" xref="Sx2.SS3.p1.7.m7.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.7.m7.1.1.1.cmml" xref="Sx2.SS3.p1.7.m7.1.1">subscript</csymbol><ci id="Sx2.SS3.p1.7.m7.1.1.2.cmml" xref="Sx2.SS3.p1.7.m7.1.1.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.7.m7.1.1.3.cmml" xref="Sx2.SS3.p1.7.m7.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.7.m7.1c">z_{1}</annotation></semantics></math> comprising the first <math id="Sx2.SS3.p1.8.m8.1" class="ltx_Math" alttext="D/2" display="inline"><semantics id="Sx2.SS3.p1.8.m8.1a"><mrow id="Sx2.SS3.p1.8.m8.1.1" xref="Sx2.SS3.p1.8.m8.1.1.cmml"><mi id="Sx2.SS3.p1.8.m8.1.1.2" xref="Sx2.SS3.p1.8.m8.1.1.2.cmml">D</mi><mo id="Sx2.SS3.p1.8.m8.1.1.1" xref="Sx2.SS3.p1.8.m8.1.1.1.cmml">/</mo><mn id="Sx2.SS3.p1.8.m8.1.1.3" xref="Sx2.SS3.p1.8.m8.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.8.m8.1b"><apply id="Sx2.SS3.p1.8.m8.1.1.cmml" xref="Sx2.SS3.p1.8.m8.1.1"><divide id="Sx2.SS3.p1.8.m8.1.1.1.cmml" xref="Sx2.SS3.p1.8.m8.1.1.1"></divide><ci id="Sx2.SS3.p1.8.m8.1.1.2.cmml" xref="Sx2.SS3.p1.8.m8.1.1.2">𝐷</ci><cn type="integer" id="Sx2.SS3.p1.8.m8.1.1.3.cmml" xref="Sx2.SS3.p1.8.m8.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.8.m8.1c">D/2</annotation></semantics></math> dimensions, and <math id="Sx2.SS3.p1.9.m9.1" class="ltx_Math" alttext="z_{2}" display="inline"><semantics id="Sx2.SS3.p1.9.m9.1a"><msub id="Sx2.SS3.p1.9.m9.1.1" xref="Sx2.SS3.p1.9.m9.1.1.cmml"><mi id="Sx2.SS3.p1.9.m9.1.1.2" xref="Sx2.SS3.p1.9.m9.1.1.2.cmml">z</mi><mn id="Sx2.SS3.p1.9.m9.1.1.3" xref="Sx2.SS3.p1.9.m9.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.9.m9.1b"><apply id="Sx2.SS3.p1.9.m9.1.1.cmml" xref="Sx2.SS3.p1.9.m9.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.9.m9.1.1.1.cmml" xref="Sx2.SS3.p1.9.m9.1.1">subscript</csymbol><ci id="Sx2.SS3.p1.9.m9.1.1.2.cmml" xref="Sx2.SS3.p1.9.m9.1.1.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.9.m9.1.1.3.cmml" xref="Sx2.SS3.p1.9.m9.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.9.m9.1c">z_{2}</annotation></semantics></math> the second set of <math id="Sx2.SS3.p1.10.m10.1" class="ltx_Math" alttext="D/2" display="inline"><semantics id="Sx2.SS3.p1.10.m10.1a"><mrow id="Sx2.SS3.p1.10.m10.1.1" xref="Sx2.SS3.p1.10.m10.1.1.cmml"><mi id="Sx2.SS3.p1.10.m10.1.1.2" xref="Sx2.SS3.p1.10.m10.1.1.2.cmml">D</mi><mo id="Sx2.SS3.p1.10.m10.1.1.1" xref="Sx2.SS3.p1.10.m10.1.1.1.cmml">/</mo><mn id="Sx2.SS3.p1.10.m10.1.1.3" xref="Sx2.SS3.p1.10.m10.1.1.3.cmml">2</mn></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.10.m10.1b"><apply id="Sx2.SS3.p1.10.m10.1.1.cmml" xref="Sx2.SS3.p1.10.m10.1.1"><divide id="Sx2.SS3.p1.10.m10.1.1.1.cmml" xref="Sx2.SS3.p1.10.m10.1.1.1"></divide><ci id="Sx2.SS3.p1.10.m10.1.1.2.cmml" xref="Sx2.SS3.p1.10.m10.1.1.2">𝐷</ci><cn type="integer" id="Sx2.SS3.p1.10.m10.1.1.3.cmml" xref="Sx2.SS3.p1.10.m10.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.10.m10.1c">D/2</annotation></semantics></math> dimensions.<span id="footnote2" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">2</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">2</sup><span class="ltx_tag ltx_tag_note">2</span>We could create more groups, but we only used two in this work.</span></span></span> We then perform <math id="Sx2.SS3.p1.11.m11.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.11.m11.1a"><mi id="Sx2.SS3.p1.11.m11.1.1" xref="Sx2.SS3.p1.11.m11.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.11.m11.1b"><ci id="Sx2.SS3.p1.11.m11.1.1.cmml" xref="Sx2.SS3.p1.11.m11.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.11.m11.1c">k</annotation></semantics></math>-means on each set of features independently, and assign each sequence <math id="Sx2.SS3.p1.12.m12.1" class="ltx_Math" alttext="z" display="inline"><semantics id="Sx2.SS3.p1.12.m12.1a"><mi id="Sx2.SS3.p1.12.m12.1.1" xref="Sx2.SS3.p1.12.m12.1.1.cmml">z</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.12.m12.1b"><ci id="Sx2.SS3.p1.12.m12.1.1.cmml" xref="Sx2.SS3.p1.12.m12.1.1">𝑧</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.12.m12.1c">z</annotation></semantics></math> to the pair <math id="Sx2.SS3.p1.13.m13.2" class="ltx_Math" alttext="(i,j)" display="inline"><semantics id="Sx2.SS3.p1.13.m13.2a"><mrow id="Sx2.SS3.p1.13.m13.2.3.2" xref="Sx2.SS3.p1.13.m13.2.3.1.cmml"><mo stretchy="false" id="Sx2.SS3.p1.13.m13.2.3.2.1" xref="Sx2.SS3.p1.13.m13.2.3.1.cmml">(</mo><mi id="Sx2.SS3.p1.13.m13.1.1" xref="Sx2.SS3.p1.13.m13.1.1.cmml">i</mi><mo id="Sx2.SS3.p1.13.m13.2.3.2.2" xref="Sx2.SS3.p1.13.m13.2.3.1.cmml">,</mo><mi id="Sx2.SS3.p1.13.m13.2.2" xref="Sx2.SS3.p1.13.m13.2.2.cmml">j</mi><mo stretchy="false" id="Sx2.SS3.p1.13.m13.2.3.2.3" xref="Sx2.SS3.p1.13.m13.2.3.1.cmml">)</mo></mrow><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.13.m13.2b"><interval closure="open" id="Sx2.SS3.p1.13.m13.2.3.1.cmml" xref="Sx2.SS3.p1.13.m13.2.3.2"><ci id="Sx2.SS3.p1.13.m13.1.1.cmml" xref="Sx2.SS3.p1.13.m13.1.1">𝑖</ci><ci id="Sx2.SS3.p1.13.m13.2.2.cmml" xref="Sx2.SS3.p1.13.m13.2.2">𝑗</ci></interval></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.13.m13.2c">(i,j)</annotation></semantics></math>. The first index refers to the assignment of <math id="Sx2.SS3.p1.14.m14.1" class="ltx_Math" alttext="z_{1}" display="inline"><semantics id="Sx2.SS3.p1.14.m14.1a"><msub id="Sx2.SS3.p1.14.m14.1.1" xref="Sx2.SS3.p1.14.m14.1.1.cmml"><mi id="Sx2.SS3.p1.14.m14.1.1.2" xref="Sx2.SS3.p1.14.m14.1.1.2.cmml">z</mi><mn id="Sx2.SS3.p1.14.m14.1.1.3" xref="Sx2.SS3.p1.14.m14.1.1.3.cmml">1</mn></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.14.m14.1b"><apply id="Sx2.SS3.p1.14.m14.1.1.cmml" xref="Sx2.SS3.p1.14.m14.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.14.m14.1.1.1.cmml" xref="Sx2.SS3.p1.14.m14.1.1">subscript</csymbol><ci id="Sx2.SS3.p1.14.m14.1.1.2.cmml" xref="Sx2.SS3.p1.14.m14.1.1.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.14.m14.1.1.3.cmml" xref="Sx2.SS3.p1.14.m14.1.1.3">1</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.14.m14.1c">z_{1}</annotation></semantics></math>, and the second index to the assignment of <math id="Sx2.SS3.p1.15.m15.1" class="ltx_Math" alttext="z_{2}" display="inline"><semantics id="Sx2.SS3.p1.15.m15.1a"><msub id="Sx2.SS3.p1.15.m15.1.1" xref="Sx2.SS3.p1.15.m15.1.1.cmml"><mi id="Sx2.SS3.p1.15.m15.1.1.2" xref="Sx2.SS3.p1.15.m15.1.1.2.cmml">z</mi><mn id="Sx2.SS3.p1.15.m15.1.1.3" xref="Sx2.SS3.p1.15.m15.1.1.3.cmml">2</mn></msub><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.15.m15.1b"><apply id="Sx2.SS3.p1.15.m15.1.1.cmml" xref="Sx2.SS3.p1.15.m15.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.15.m15.1.1.1.cmml" xref="Sx2.SS3.p1.15.m15.1.1">subscript</csymbol><ci id="Sx2.SS3.p1.15.m15.1.1.2.cmml" xref="Sx2.SS3.p1.15.m15.1.1.2">𝑧</ci><cn type="integer" id="Sx2.SS3.p1.15.m15.1.1.3.cmml" xref="Sx2.SS3.p1.15.m15.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.15.m15.1c">z_{2}</annotation></semantics></math> as per equation <a href="#S2.E1" title="Equation 1 ‣ 2.4.1 Generative routing ‣ 2.4 Coarse Routing ‣ 2 Approach ‣ DiPaCo: Distributed Path Composition" class="ltx_ref"><span class="ltx_text ltx_ref_tag">1</span></a>. Although each assignment takes only <math id="Sx2.SS3.p1.16.m16.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.16.m16.1a"><mi id="Sx2.SS3.p1.16.m16.1.1" xref="Sx2.SS3.p1.16.m16.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.16.m16.1b"><ci id="Sx2.SS3.p1.16.m16.1.1.cmml" xref="Sx2.SS3.p1.16.m16.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.16.m16.1c">k</annotation></semantics></math> possible values, the number of possible unique pair-assignments is <math id="Sx2.SS3.p1.17.m17.1" class="ltx_Math" alttext="k^{2}" display="inline"><semantics id="Sx2.SS3.p1.17.m17.1a"><msup id="Sx2.SS3.p1.17.m17.1.1" xref="Sx2.SS3.p1.17.m17.1.1.cmml"><mi id="Sx2.SS3.p1.17.m17.1.1.2" xref="Sx2.SS3.p1.17.m17.1.1.2.cmml">k</mi><mn id="Sx2.SS3.p1.17.m17.1.1.3" xref="Sx2.SS3.p1.17.m17.1.1.3.cmml">2</mn></msup><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.17.m17.1b"><apply id="Sx2.SS3.p1.17.m17.1.1.cmml" xref="Sx2.SS3.p1.17.m17.1.1"><csymbol cd="ambiguous" id="Sx2.SS3.p1.17.m17.1.1.1.cmml" xref="Sx2.SS3.p1.17.m17.1.1">superscript</csymbol><ci id="Sx2.SS3.p1.17.m17.1.1.2.cmml" xref="Sx2.SS3.p1.17.m17.1.1.2">𝑘</ci><cn type="integer" id="Sx2.SS3.p1.17.m17.1.1.3.cmml" xref="Sx2.SS3.p1.17.m17.1.1.3">2</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.17.m17.1c">k^{2}</annotation></semantics></math>. Moreover, the cost of assignment grows with the square root of the number of total pair assignments. We used this approach when experimenting with hundreds of paths; while the results were an improvement over simple <math id="Sx2.SS3.p1.18.m18.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.18.m18.1a"><mi id="Sx2.SS3.p1.18.m18.1.1" xref="Sx2.SS3.p1.18.m18.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.18.m18.1b"><ci id="Sx2.SS3.p1.18.m18.1.1.cmml" xref="Sx2.SS3.p1.18.m18.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.18.m18.1c">k</annotation></semantics></math>-Means, discriminative routing was better still; and using the product <math id="Sx2.SS3.p1.19.m19.1" class="ltx_Math" alttext="k" display="inline"><semantics id="Sx2.SS3.p1.19.m19.1a"><mi id="Sx2.SS3.p1.19.m19.1.1" xref="Sx2.SS3.p1.19.m19.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="Sx2.SS3.p1.19.m19.1b"><ci id="Sx2.SS3.p1.19.m19.1.1.cmml" xref="Sx2.SS3.p1.19.m19.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="Sx2.SS3.p1.19.m19.1c">k</annotation></semantics></math>-means as a router for the initial module training before the discriminative step did not significantly improve results. Note that the discriminative routing adapts to the structure of the DiPaCo, because the paths are used to score documents to find the targets. Nevertheless, we consider alternative generative routing approaches important for further study, as discriminative scaling by scoring all paths cannot scale to large numbers of paths.</p>
</div>
<div class="ltx_pagination ltx_role_newpage"></div>
</section>
</section>
</article>
</div>
<div class="ar5iv-footer"><a href="/html/2403.10615" class="ar5iv-nav-button ar5iv-nav-button-prev">◄</a>
    <a class="ar5iv-home-button" href="/"><img height="40" alt="ar5iv homepage" src="/assets/ar5iv.png"></a>
    <a href="/feeling_lucky" class="ar5iv-text-button">Feeling<br>lucky?</a>
    <a href="/log/2403.10616" class="ar5iv-text-button ar5iv-severity-error">Conversion<br>report</a>
    <a class="ar5iv-text-button" target="_blank" href="https://github.com/dginev/ar5iv/issues/new?template=improve-article--arxiv-id-.md&title=Improve+article+2403.10616">Report<br>an issue</a>
    <a href="https://arxiv.org/abs/2403.10616" class="ar5iv-text-button arxiv-ui-theme">View&nbsp;original<br>on&nbsp;arXiv</a><a href="/html/2403.10617" class="ar5iv-nav-button ar5iv-nav-button-next">►</a>
</div><footer class="ltx_page_footer">
<a class="ar5iv-toggle-color-scheme" href="javascript:toggleColorScheme()" title="Toggle ar5iv color scheme"><span class="color-scheme-icon"></span></a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/license" target="_blank">Copyright</a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/policies/privacy_policy" target="_blank">Privacy Policy</a>

<div class="ltx_page_logo">Generated  on Fri Apr  5 17:24:29 2024 by <a target="_blank" href="http://dlmf.nist.gov/LaTeXML/" class="ltx_LaTeXML_logo"><span style="letter-spacing:-0.2em; margin-right:0.1em;">L<span class="ltx_font_smallcaps" style="position:relative; bottom:2.2pt;">a</span>T<span class="ltx_font_smallcaps" style="font-size:120%;position:relative; bottom:-0.2ex;">e</span></span><span style="font-size:90%; position:relative; bottom:-0.2ex;">XML</span><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAsAAAAOCAYAAAD5YeaVAAAAAXNSR0IArs4c6QAAAAZiS0dEAP8A/wD/oL2nkwAAAAlwSFlzAAALEwAACxMBAJqcGAAAAAd0SU1FB9wKExQZLWTEaOUAAAAddEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIFRoZSBHSU1Q72QlbgAAAdpJREFUKM9tkL+L2nAARz9fPZNCKFapUn8kyI0e4iRHSR1Kb8ng0lJw6FYHFwv2LwhOpcWxTjeUunYqOmqd6hEoRDhtDWdA8ApRYsSUCDHNt5ul13vz4w0vWCgUnnEc975arX6ORqN3VqtVZbfbTQC4uEHANM3jSqXymFI6yWazP2KxWAXAL9zCUa1Wy2tXVxheKA9YNoR8Pt+aTqe4FVVVvz05O6MBhqUIBGk8Hn8HAOVy+T+XLJfLS4ZhTiRJgqIoVBRFIoric47jPnmeB1mW/9rr9ZpSSn3Lsmir1fJZlqWlUonKsvwWwD8ymc/nXwVBeLjf7xEKhdBut9Hr9WgmkyGEkJwsy5eHG5vN5g0AKIoCAEgkEkin0wQAfN9/cXPdheu6P33fBwB4ngcAcByHJpPJl+fn54mD3Gg0NrquXxeLRQAAwzAYj8cwTZPwPH9/sVg8PXweDAauqqr2cDjEer1GJBLBZDJBs9mE4zjwfZ85lAGg2+06hmGgXq+j3+/DsixYlgVN03a9Xu8jgCNCyIegIAgx13Vfd7vdu+FweG8YRkjXdWy329+dTgeSJD3ieZ7RNO0VAXAPwDEAO5VKndi2fWrb9jWl9Esul6PZbDY9Go1OZ7PZ9z/lyuD3OozU2wAAAABJRU5ErkJggg==" alt="Mascot Sammy"></a>
</div></footer>
</div>

    <script>
      var canMathML = typeof(MathMLElement) == "function";
      if (!canMathML) {
        var body = document.querySelector("body");
        body.firstElementChild.setAttribute('style', 'opacity: 0;');
        var loading = document.createElement("div");
        loading.setAttribute("id", "mathjax-loading-spinner");
        var message = document.createElement("div");
        message.setAttribute("id", "mathjax-loading-message");
        message.innerText = "Typesetting Equations...";
        body.prepend(loading);
        body.prepend(message);

        var el = document.createElement("script");
        el.src = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js";
        document.querySelector("head").appendChild(el);

        window.MathJax = {
          startup: {
            pageReady: () => {
              return MathJax.startup.defaultPageReady().then(() => {
                body.removeChild(loading);
                body.removeChild(message);
                body.firstElementChild.removeAttribute('style');
              }); } } };
      }
    </script>
    <script>
    // Auxiliary function, building the preview feature when
    // an inline citation is clicked
    function clicked_cite(e) {
      e.preventDefault();
      let cite = this.closest('.ltx_cite');
      let next = cite.nextSibling;
      if (next && next.nodeType == Node.ELEMENT_NODE && next.getAttribute('class') == "ar5iv-bibitem-preview") {
        next.remove();
        return; }
      // Before adding a preview modal,
      // cleanup older previews, in case they're still open
      document.querySelectorAll('span.ar5iv-bibitem-preview').forEach(function(node) {
        node.remove();
      })

      // Create the preview
      preview = document.createElement('span');
      preview.setAttribute('class','ar5iv-bibitem-preview');
      let target = document.getElementById(this.getAttribute('href').slice(1));
      target.childNodes.forEach(function (child) {
        preview.append(child.cloneNode(true));
      });
      let close_x = document.createElement('button');
      close_x.setAttribute("aria-label","Close modal for bibliography item preview");
      close_x.textContent = "×";
      close_x.setAttribute('class', 'ar5iv-button-close-preview');
      close_x.setAttribute('onclick','this.parentNode.remove()');
      preview.append(close_x);
      preview.querySelectorAll('.ltx_tag_bibitem').forEach(function(node) {
        node.remove();
      });
      cite.parentNode.insertBefore(preview, cite.nextSibling);
      return;
    }
    // Global Document initialization:
    // - assign the preview feature to all inline citation links
    document.querySelectorAll(".ltx_cite .ltx_ref").forEach(function (link) {
      link.addEventListener("click", clicked_cite);
    });
    </script>
    </body>
</html>

2403.10616
Progress in machine learning (ML)  has been fueled by scaling neural network models. This scaling has been enabled by ever more heroic feats of engineering, necessary for accommodating ML approaches that require high bandwidth communication between devices working in parallel.
In this work, we propose a co-designed modular architecture and training approach for ML models, dubbed DIstributed PAth COmposition (DiPaCo) . During training, DiPaCo distributes computation by paths through a set of shared modules. Together with a Local-SGD inspired optimization (DiLoCo)  that keeps modules in sync with drastically reduced communication,
Our approach facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions. At inference time, only a single path needs to be executed for each input, without the need for any model compression. We
consider this approach as a first prototype towards a new paradigm of large-scale learning, one that is less synchronous and more modular. Our experiments on the widely used C4 benchmark show that, for the same amount of training steps but less wall-clock time, DiPaCo exceeds the performance of a 1 billion-parameter dense transformer language model by choosing one of 256 possible paths, each with a size of 150 million parameters.
abstract
DiPaCo: Distributed Path Composition
Abstract
2403.10616
1IntroductionProgress in machine learning and AI has been driven by spending more FLOPs on larger neural network models, trained on bigger data sets. This scaling has been accomplished via data and model parallelismand pipeliningto distribute computation, enabling the concurrent use of a large number of devices. Although model architectureshave also been used to allow computational parallelism, and optimization procedures to prefer larger batches(again allowing more data parallelism) , the current training paradigm has not fundamentally changed model architecture or optimization procedure to facilitate distributed training.
State of the art models are still essentially monoliths, and their optimization requires exchanges of parameters, gradients, and activations at every step of the learning process.This approach incurs engineering and infrastructure challenges associated with provisioning and managing the large number of tightly interconnected devices required for
the lengthy training process. The training process itself is often restarted for each new model release, essentially discarding much of the computation for training the last model. Moreover, training monoliths incurs human-organizational challenges, as it is difficult to localize the effects to the final model of changes to any step in the process (beyond data preparation) .
In particular, it is difficult to leverage the potential of the greater ML community, rather than a single organization.
Because of these challenges, it may become more difficult to continue to scale with the current approach.As inand depicted inFigure 1, we envision an alternative approach where models are scalable both in terms of ability to ingest and train on large amounts of data and in terms of enabling collaboration with many contributors, and can be continuously updated and expanded, by virtue of modular design. Modular ML may have many other benefits, seefor further discussion and references.In this work, we take a step towards this scalable modular ML paradigm,
proposing an architecture and training algorithm, DIstributed PAths COmposition (DiPaCo) .
DiPaCo’s architecture and optimization have been co-designed to reduce communication and enable better scaling. The high level idea is todistribute computation by path; here, a “path” means a sequence of modules that define an input-output function.
Paths are small relative to the entire model, thus requiring only a handful of tightly connected devices to train or evaluate.
During both training and deployment, a query is routed to a replica of a path rather than a replica of the whole model; in other words, the DiPaCo architecture is sparsely activated. We give a detailed exposition of DiPaCo insection 2, and the infrastructure we used to implement it insection 3.Insection 4we demonstrate the feasibility of DiPaCo by training a language model on the C4 datasetwith paths of size 150 million parameters, matching the performance in terms of validation perplexity of a 1.3 billion model, but with 45% less wall clock training time. While the dense 1.3B system required the use of all co-located devices, DiPaCo uses 256 islands of compute, each of which is one-eighth the number of devices used to train the baseline. Neither during train or evaluation time is it necessary to co-locate the full DiPaCo model.Figure 1:Long-term Goal: Ultimately, we envision a modular network where different components,paths, are optimized for different tasks,, each designed by different researchers. The paths, trained on any available hardware type, communicate infrequently across the world, exchanging useful information and enabling new forms of composition.
introduction
DiPaCo: Distributed Path Composition
1Introduction
2403.10616
2ApproachIn this section we give a detailed description of DiPaCo. We start by describing the assumptions and setting in which we will work.2.1SettingOur goal in this work is to demonstrate scalable ML models in the setting of many smaller islands of compute, as opposed to one tightly connected compute landmass. Accordingly, we assume:1. Training compute (FLOPs)  is relatively cheap2. Communication is relatively expensiveThese assumptions arenotrealistic in the current ML training paradigm.
As discussed above, devices are usually co-located and run by a single organization, FLOPs are expensive, constrained by the purchase and installation of many accelerators, but communication between devices is relatively cheap because of the co-location.Nevertheless, in our view,
these assumptions
may be realistic in the near future, if our compute requirements (or model sizes)  grow beyond what can be reasonably co-located, and so communication costs become relatively more expensive.
Conversely, progress in distributed training of ML models may allow simpler infrastructure build-outs, leading eventually to more available compute. As it is, infrastructure is designed around the standard approach to training large monoliths; and ML models are architected to take advantage of the current infrastructure and training approaches. This feedback loop may be leading the community to a spurious local minimum where compute is more constrained than it needs to be.In addition (and closely related)  to the two assumptions above, we also assume that both during train and evaluation, we cannot instantiate models as large as we would like to have on any single compute island.We will work in the context of language modeling (LM) .
Based on the above assumptions, we will evaluate models by evaluation perplexity (PPL)  against wall-clock time, keeping evaluation complexity in mind . We choose this metric because it encourages exploring how to distribute model training, and more generally, points to in our opinion plausible new paradigms for ML.2.2Overview of SystemThe core idea of DiPaCo is to train a sparsely-activated modular system where data and computation are distributed by the choice of path through the modules. There are two key ideas to making this work:Coarse Routing:Sparsely routed Mixture of Experts (MoE)  have shown great results in language modeling. The standard approach in a transformer MoE LM is to make a routing decision at each token, based on the feature at each routed layer. In contrast, in this work, during training we will route once per document and offline, as in.Routing once per document allows batching computation across all tokens of a sequence, without the need to swap modules in and out as a sequence is processed. This in turn allows parameters to be distributed across distant workers.
In addition, instead of learning the router along with the model, we will compute routing decisionsoffline. This enables pre-sharding data by path, before the start of training. This is critical to distribute training across paths, as each worker can now train a path by processing its own shard of data, using its own set of parameters. Insubsection 2.4, we will describe in detail how we approach coarse routing.DiLoCo to keep modules in syncPaths cannot be trained completely independently, because some modules might be shared across multiple paths. To support module sharing across paths, we use DiLoCofor low communication data parallelism, seesubsection 2.5. If there arepaths (assigned toworkers)  that share module, each corresponding worker performs SGD on its own shard of data, and everyfew hundredsteps, workers average the difference of the parameters of modulebefore and after the local SGD phase. These averages are then used to update a global parameter vector which is then re-distributed across the workers to sync them, a procedure described insubsection 2.6.With these two choices,
neither at training nor at test time does the entire network (collection of paths)  need to be materialized together, as shown inFigure 1.2.3NotationIn this section we introduce the basic notation used throughout this work.
We assume that we have a base model architecture with parameters. We take a partition of the parameter indices intosubsets.
For each, we choose a numberthat represents the number of possible choices for the parameters in. We will call a set of parameters associated to aa “module” or, as in, an “expert”. In this terminology,is the number of distinct modules associated to.For a simple example, consider a four-layer fully connected network. We might take the first two layers to beand the third and fourth layers to be, so. If we choosefor, then we have amodules in, andmodules in, for a total ofmodules. Any choice of module fromand module formdefines a neural network; as in, we call each of thesepossible networks a “path”, seeFigure 2.Figure 2:An illustration of the first example from Section2.3. Alayer neural network, with blockconsisting of the firstlayers andconsisting of the nextlayers. Each block haschoices of module (each with its own parameters) , represented by different colors. On the left, we show all of thepossible paths. On the right, we show a single path.In this example, and most of the cases we will consider in the rest of this paper, the “blocks”determine contiguous sub-networks ofthat can act as an input-output mapping, but this is not necessary. We could just as easily letbe layersandandbe layersand; orbe all the biases in the network, andbe all the linear parameters. Nevertheless, we will callthe “level” of the module.In order to transform an inputto an output, we need to figure out which parameters will be used to operate on. The function that takesto a choice of parameters will be called the “router”, and is denoted by; thushas the form, where eachindexes one of thechoices of parameters for, andmaps an input to a path. We will also often collapse the sequenceinto a single number, where. If the training dataset is denoted by, the subset of data that is routed to pathwill be called the-th “shard”.When the router maps anto a distribution overrather to a deterministic choice, we say that the router is “soft”. In this work, we will consider “hard” routers instead, where the choice is single valued and deterministic.Finally, it is useful to refer to module parameters. Assume that the-th path goes through moduleat level, with,.
Since during training paths are not fully synchronized, we denote the local copy of the parameters of modulein pathat iterationwith. The collection of parameters used by the-th path (across all levels)  is denoted by. When modules across paths are synchronized, the parameter values of moduleis the same across all paths, and therefore we denote the global copy of the parameters of moduleat iterationby omitting the path index, as in. Similarly, we denote local and global gradients of modulewithand, respectively.2.4Coarse RoutingIn this section we describe how we route sequences.
The basic intuition is to use some context to decide which path is most suitable to process a given sequence. Since we focus on language modeling, we use as context the firsttokens of each sequence. At training time, we then use the remaining tokens for learning the model parameters. At test time, we use the remaining tokens to compute perplexity on a held-out validation set. All methods, including dense baselines, will be evaluated in the same way, by calculating perplexity using all but the firsttokens of each sequence, which was used to determine the routing decision of our models.Notice that the outcome of routing is assigning each sequence to a certain path in the network. Collectively, all sequences that are assigned to a particular path form a shard of the original dataset. In this work, there is a one-to-one association between paths and corresponding shards; pathis associated to shard.Next, we describe how we use this prefix oftokens to route, as well as other variants of routing.Figure 3:routing More Frequently at Test-Time: At training time (left panel) , the router selects the pathusing the prefix. We train the chosen path on the whole sequence using the usual language modeling loss. At test time (right panel) , the path selected by the router given the prefix is used to score the next chunk of tokens. Then, we re-use the router to choose the most likely path given the new chunk. This process repeats until the whole sequence has been scored.2.4.1Generative routingIn generative routing, the decision about which shard to assign a sequence to is not informed by the task at hand, namely language modeling. Instead, the choice is based on minimizing feature reconstruction error. Given a representationof the firsttokens of a sequence (context) , we perform-Means on the featuresof each sequence, and then we use the-means assignment algorithm to shard the data intoshards. Ifare theprototypes learned by-means, the sequence with prefixis assigned to shardvia:(1) 2.4.2Discriminative routingGenerative routing is agnostic of the language modeling task. In discriminative routing the sharding takes into account how well experts perform on each sequence. The training set is split into two parts. The first and largest part is used to train the paths on data sharded using a generative sharding appproach. The second and much smaller part is used to evaluate paths on each sequence. Letbe the log-likelihood score of the pathon the sequence associated with prefix. The index of the path attaining the largest likelihood score,, is used as label for. We then train a logistic regression classifier to directly predictfrom the features; this trained classifier is used as.
Finally, we use such classifier to re-shard the entire training and validation datasets.This process is an approximation to Expectation Maximization, a coordinate descent method alternating between updating the parameters of the model (paths) , and latent variables (path assignments) . It is an approximation because, for simplicity, we use as label the top scoring path as opposed to the full posterior distribution over path scores. Of course, this re-assignment process can be repeated for as many times as desired.We invite those readers interested in diving deeper in the routing mechanism to readsubsection 7.2of the Appendix for further details.2.4.3Routing More Frequently at Test TimeRouting only once per sequence is critical to distribute training across paths. However, at test time we can afford routing more frequently. LikeFigure 3shows, we can divide a sequence at test time into chunks ofconsecutive tokens, e.g.,. We can then feed a router with the-th chunk to predict the best path to use at the-th chunk, using a process similar to the one described insubsubsection 2.4.2. The cost of switching paths is small, since the router is run infrequently and in scoring (as opposed to auto-regressive generation)  mode. Moreover, only text needs to be communicated to the router and next selected path.2.4.4Overlapping ShardsAbove, each sequence was associated to one and only one shard; in other words, we pre-sharded the dataset into a set ofdisjoint partitions.When the number of shardsis large, the mixture might overfit, particularly when paths share few parameters.
To combat this issue we might assign each sequence not just to one shard but to its“closest” shards.
This makes shards overlap with each other, smoothing out the boundaries among themselves, which may improve generalization when the number of shards is relatively large. This can be doneindependentlyat train time and at evaluation time; that is, one might overlap shards only during train, at eval, or both, or neither. At train time, overlapping the shards limits the specialization of each path, and increases the storage necessary for each shard, but otherwise does not have a computational cost, as each path is still trained independently on its shard. However, overlapping the shards at evaluation means that the model has to be forwarded through multiple paths, which increases the evaluation computational cost.In this work, when we overlap the shards in training, we use the top-2 choices, while we never overlap the shards at evaluation. We leave to future work other approaches to overlap shards.2.5DiLoCo: ReviewIn this section, we summarize DiLoCoas it lays the foundation of our distributed optimization algorithm, which will be described in the next section. Together with coarse routing ofsubsection 2.4, DiLoCo is a critical component of DiPaCo.DiLoCo optimizes a dense model acrossworkers. First, the original dataset is sharded intoshards, and each shard is associated to a worker. Second, all workers start from the same copy of the model parameters (typically a pretrained model) , and perform(inner)  optimization steps over their own shard independently. Third, each worker sends to a central CPU server the difference in parameter space between their new parameters and their initial ones. We refer to these differences asouter gradients. The central server averages these outer gradients into a single update vector. Finally, the global parameter vector is updated using an (outer)  optimizer, and the new global parameter vector is re-dispatched to all workers for another phase of local training. The process repeats for as many rounds as desired.In language modeling applications using transformers, the inner and outer optimizers that have been shown to be most effective are respectively AdamWand Nesterov momentum. Empirically in prior work,has been set up to 64 andup to thousand. Note that other alternatives, such as FedOpt, are compatible with this framework.2.6DiPaCoFigure 4:DiPaCo: (left)  The dataset is pre-sharded intoshards,(here) . (middle)  Compact view of aDiPaCo, which is never instantiated. In this toy illustration, there are three levels. Level 2 and 3 have a mixture with two modules each. Level 1 has a single module shared by all paths. (right) 
We associate each shardto apath. In this toy illustration, a path is the composition of three neural network blocks. The color refers to the id of a module. The figure shows the modular network unrolled across the four paths. These are trained by using DiLoCo which requires communicating only every few hundred steps. In this example, module 2a (inred)  is shared by pathsand. Workers associated to paths might use different hardware types (different kind of GPUs or TPUs)  and might be placed in far away geographic areas.Insubsection 2.4, we described methods to pre-shard data, and insubsection 2.5we described how a model can be trained across several workers that communicated infrequently.
This section combines these two approaches towards the efficient training of a composable mixture, a network whose levels are replaced by a mixture of experts. We dub such architecture and training algorithm DiPaCo. In our experiments below, a level is composed of several consecutive transformer blocks.In the toy illustration ofFigure 4there are three levels,,, and. There is only one module (equivalently, one set of parameters)  in, and it is shared across all paths. There are two modules in(meaning that there will be two distinct parameter vectors for the module at that level) ; the first is shared by pathsand, and the second by pathsand. Similarly, at the third levelandshare parameters, andandshare parameters.The resultingDiPaCo haspaths in total (as shown on the middle panel ofFigure 4) . However, the full model need never be fully instantiated, neither during training nor testing. Only paths are realized (as shown on the right panel ofFigure 4) , and these are trained in parallel with infrequent communication.At training time DiLoCo is applied at the level of modules, as opposed to the entire network; see lines 11-16 ofAlgorithm 1. The outer optimizer updates independently the parameters of each module after receiving the corresponding outer gradients from workers processing paths passing through that particular module. For instance and with reference toFigure 4, the outer gradient of module 3b is calculated by averaging the gradients of module 3b of the worker processingand. Such outer gradient is then used to update the parameters of module 3b.
A worker does not need to be powerful enough (in terms of memory and compute)  to host the entire model, but just a single path. Like in DiLoCo work, the inner optimizer is AdamW and the outer optimizer is Nesterov momentum.
Similarly, at test time, the paths are instantiated, and served independently, with text routed to each path via a router.Algorithm 1DiPaCo (see notation insection2.3) 1:Num. levels, num. experts per level, pathswithand. Letbe the number of paths going through moduleat level.2:Pre-sharded training set into(seesection2.4) .3:Parameters of pretrained model used for initialization:.4:OptimizersInnerOptandOuterOpt5:forouter stepdo6:(Optional, in this work done once during training)  discriminatively re-shard data (seesection2.4.2) 7:forworkerdo8:9:forinner stepdo10:11:12:13:endfor14:endfor15:forleveldo16:forexpertdo17:18:19:endfor20:endfor21:endfor2.6.1Increasing CapacityFigure 5:DiPaCo with more capacity: In this example, level 3 modules are path specific, i.e., modules at that level are not shared by paths.The more paths go through a module, the more opportunities for transfer learning across paths, but also the more constrained learning is and the less capacity the overall mixture has. DiPaCo gives a lot of flexibility in how to set this trade-off. Whether a level is shared by all paths or by a certain subset only does not change the overall framework. The only difference is in the choice of which subset of paths is used when computing the average in lineofAlgorithm 1.One extreme choice is to allocate path-specific modules, as shown inFigure 5. In this case, the third level has a mixture with as many modules as there are paths. This is a particularly easy way to increase parameter count in DiPaCo.111While there is no gradient averaging anymore for these path-specific modules, we still apply the outer optimization ofAlgorithm 1because empirically it improves convergence over the default optimizer.Clearly, the proposed structure of DiPaCo is merely a first step. We consider it especially exciting to co-design the optimization methods and the architecture. In the next section, we discuss some of our initial approaches to adapting optimization methods for our setting.2.6.2Scaling the modular architectureDiPaCo is made of multiple levels, each with several expert modules.e.gahas two levels, which contains 16 modules on each level, producing in total 256 paths. It is easy to scale the architecture size by increasing the number of levels and modules per levels,e.gahaspaths. It’s difficult to find enough devices for that number of paths, therefore only a subset of the paths can be trained at any moment in time. Potentially, at each start of inner optimization, we could sample a different subset of the paths. Doing so would allow us to scale DiPaCo to arbitrary large size.2.6.3Flat Mixture of ExpertsThe extreme form of capacity increase as insubsubsection 2.6.1would be to have each path be a completely independent network. We will refer to this model as aflatmixture of experts (flat MoE) . In a flat MoE paths do not shareanyparameter. In the language ofsubsection 2.3, there is only one level, andis the entire network, and.Flat MoE is a strong baseline for the compositional version of DiPaCo. In our experiments, we see that it can outperform architectures with shared modules when the number of shards is small compared to the size of the total amount of data.With generative routing using-means, this approach is essentially; in our experiments below we will use discriminative routing as in the compositional models, as this consistently outperforms generative routing.2.7Advanced Optimization TechniquesOuter Gradient Norm Rescaling:In DiPaCo, each module can belong to a different number of paths. For instance, in aDiPaCo, 16 paths can contribute to a single module, while there could be a level that has path-specific modules which is not shared by any other.Consequently, the average outer gradients,inAlgorithm 1, have significant different norms for different modules. Using the intuition, and also empirical observation, that averaging across a larger number of paths is akin to using a larger batch size, we have rescaled the outer gradient norm by the square root of the number of paths going through a module.Loss reweighing:In general, shards might have different amounts of data. If training of DiPaCo is distributed by using one worker per shard/path, then according to the uniform average of lineinAlgorithm 1smaller shards are over-sampled.
To compute an unbiased estimate of the gradient, we therefore weigh the outer gradients proportionally to the shard size:(2) with:(3) Early Stopping:By setting aside a small subset of training examples in each shard, we can perform path specific early stopping. In this case, for each path we selected the parameters that yield the lowest loss on the corresponding shard validation set. We found that early stopping improves generalization on small shards, e.g., when the number of shards is large.
body
DiPaCo: Distributed Path Composition
2Approach
2403.10616
3InfrastructureFigure 6:Infrastructure. Training workers (inorange)  responsible for performing the inner optimization fetch training task from the train task queue (inpurple) . Once the training checkpoints are saved to GFS, their paths and metadata (e.g. training step, path ID, phase ID)  are written to a Spanner table (inblue)  allowing easy look up of checkpoint path given metadata. The evaluation workers and the sharded outer optimization executors (inred)  load training checkpoints once they become ready, as signaled by their appearance in the Spanner table. Finally, a monitoring worker (ingreen)  periodically checks the health of workers and reboots them if necessary.We now describe the infrastructure we designed to implement the framework detailed earlier. The infrastructure should achieve the following objectives: 1)  reliably train the modular architecture via infrequent synchronization; 2)  ensure easy scalability for scenarios involving training across multiple pods or data centers; and 3)  provide continuous fault tolerance, allowing every component of the system to quickly recover from failed hosts and preemptions without halting training.We illustrate the different components inFigure 6. The general training workflow is the following:1.At the beginning of each training phase,
a list of training tasks, each including pairs of path and shard ids, is added to the training task queue maintained by the task scheduler server (inpurple) .2.When a training worker in the worker pool (inorange)  becomes available, it fetches the next training task from the train task scheduler and performs inner optimization (L5-9 ofAlgorithm 1)  on accelerators. Once the inner optimization is finished, the checkpoint is saved to the Google’s distributed file systemand the training worker becomes available for the next training task. The path to the checkpoint, along with the metadata of the checkpoint (e.g., path ID, outer step ID, etc.) , is recorded in a database table (shown inblue) . This enables other components to query the checkpoint file path for a given path.3.When a checkpoint is saved by a training worker (indicated inteal) , evaluation tasks for that checkpoint are added to the eval task queue (highlighted inyellow) . The eval task queue is then consumed by the evaluation workers in the worker pool to carry out evaluation tasks.4.The outer optimizer task scheduler (indicated inlight blue)  distributes outer optimization tasks to sharded outer optimization executors (highlighted inred) , each of which is responsible for the outer optimization of a shard of modules (e.g., a single module or a collection of modules) . Each executor loads training checkpoints containing the corresponding modules as soon as they appear in the Spanner database tableand performs the parameter averaging and outer optimization (L13-14) . The checkpoints of the updated module parameters are saved to the distributed file system, and their paths with metadata are written to the database table.5.After completion of all the training tasks, the current training phase concludes, and the subsequent training phase commences and training tasks in the new phase can be started as soon as their corresponding modules finish the outer optimization step. See3.3for more detail.6.Throughout training, a job status monitor (ingreen)  periodically checks the health of workers and task queue servers, and restarts them if they become unresponsive.All the components are designed to be robust to failures and pre-emptions. We will discuss more details of each infrastructure component in the following subsections.3.1Worker PoolWe employ a producer-consumer design pattern to distribute the training tasks. Within each phase, training tasks are added to the task queue, each of which involves training a path for a specific number of steps from a given checkpoint. We maintain a pool of workers, where each worker retrieves training tasks from the task queue iteratively. Each training task is completely independent of other tasks, requiring no synchronization or communication among the workers.In the event of worker failure or preemption, the fault-tolerant task queue server would return the task from the unavailable worker back to the task queue before reassigning it to another available worker. The task queue server also periodically checkpoints the current task queue, making it possible to recover from server failures or preemptions.The key advantage of this design is that, thanks to the complete independence among workers, the system can continue making progress even if some workers become unavailable, as long as the worker pool is not empty. It is worth noting that the worker pool can containheterogeneoustypes of devices acrossdifferent regions, making it easy to scale up further. Additionally, the worker pool allows for elastic resource utilization by auto-scaling pool size according to resource availability.3.2Task Queue SystemIn practice, the aforementioned task queue system is designed to manage and distribute asynchronous tasks among the workers in the worker pool using remote procedure calls. It consists of three components: the task queue scheduler, the task queue server, and the task queue client. Training tasks are published by the task queue scheduler and sent to the task queue server, while each worker instantiates a task queue client to request tasks from the task queue.In the multi-host training scenario, where multi-host parallelism is achieved by Single Program, Multiple Data (SPMD) , initializing a task queue client on each host would result in each host pulling a different task from the task queue. We solve this issue by implementing a task queue client that synchronizes its results across all hosts. This is achieved by creating a real task queue client only on the first JAX host. All other hosts get the response from the first host by performing a blocking all-gather operation.We also need to synchronize task queue write operation outside of the main Python thread in some places. For instance, evaluation tasks for a given checkpoint can only be enqueued after checkpointing on all hosts has finished. For this use case, we implemented a barrier, which blocks until each program running on their host have made a call with the same unique key.3.3Outer Optimization EfficiencyFigure 7:Sharded Outer Optimization Executor: An example of outer optimization executors sharded by module, which significantly reduces the processing time and memory requirement.Naïve implementation of outer optimization executor introduces significant overhead as we scale up the number of paths across the world. The following optimizations allow us to scale to hundreds of paths with average time per phase for outer update under 2 minutes (lines-ofAlgorithm 1) .Online Parameter Gradient Averaging:As the training tasks within a phase may not finish at the same time, it is possible to reduce the parameter averaging time by loading and accumulating a training checkpoint to the current partial sum as soon as one becomes available. It is also possible to extend to an asynchronous update that doesn’t require to wait for all paths before doing an outer update.Sharded Outer Optimization Executor:As our model architecture is modular, averaging meta-gradients of different modules can be performed independently. By distributing the parameter averaging across multiple servers, as illustrated inFigure 7, we reduce the memory requirement of each averaging task, as well as the overhead of checkpoint loading and averaging. It also allows to reduce the training time as a training task can be started as soon as its corresponding modules have finished their outer update. As a consequence, the overall model is never materialized in a single location but always split across several servers.Asynchronous Checkpoints Gathering:Each outer optimization executor loads as many checkpoints as the number of paths sharing its associated module. It is possible to load checkpoint from anywhere on earth from the distributed file system. However, if the checkpoint, located on the server where training took place, is at a certain distance from the outer optimization executor, there is going to be a significant delay. In this case we launch in the background an Effingo processto bring the checkpoint to a closer location.Miscellaneous Improvements:In addition to the previous speed improvements, we also a)  cache the parameters of the outer optimizerOuterOpt, b)  reuse the just-in-time compiled optimizer, and c)  load multiple paths in parallel via a multi-threading queue.3.4Backup PoolBy default we create as many training workers (see inorangeinFigure 6) as paths/shards of data,. However as we scale the number of paths, the hardware requirement might become prohibitive. For instance, with 256 paths and 16A100GPUs per worker, we would needGPUs. Therefore in practice, the number of workers is often less than the number of paths, and instead, we do multipleroundsof training within an outer iteration step until all paths have been trained (L-ofAlgorithm 1) .To minimize the number of rounds, and therefore maximize speed, we create a backup pool of training workers. Namely, we spawn as many workers as there are paths across multiple accelerator types using a low-tier priority. As soon as an accelerator is available, it is snatched and used for an inner optimization phase. Although device with low-tier priority are preempted frequently, we can benefit from the backup pool since each training task only takes roughly a couple of minutes (the inner optimization L-inAlgorithm 1) .Note that this backup pool is particularly useful when used with the path sampling idea mentioned insubsubsection 2.6.2.
body
DiPaCo: Distributed Path Composition
3Infrastructure
2403.10616
4ExperimentsWe consider a language modeling task on the C4 dataset, derived from Common Crawl, tokenized with a SentencePiece tokenizerwith a vocabulary size of. We report perplexity on the validation set against number of weight update steps used at training time, which is a close proxy for wall-clock time if all computations are done on the same accelerator type. The total number of weight update steps is set to. All paths on DiPaCo have sizemillion parameters, using a transformer withblocks,dimensional hidden states andheads.We build modular networks of varying number of paths, and varying levels of module reuse. We compare to the performance of a dense transformer language model of sizeB, which hasblocks,dimensional hidden states and with the same number of heads; and a dense model that is the size of one DiPaCo path. These models are also trained forweight update steps.We warn the reader that this comparison isnotstandard in the literature, as weight updates for DiPaCo see more tokens and use more FLOPs when the number of paths is larger.
Nevertheless, while our modular networks is made of hundreds of paths, they are all trained in parallel, and thus our training wallclock time is 45% less than theB dense counterpart, and roughly equivalent to theM parameter single-path counterpart.In our experiments we have searched over relatively few hyper-parameters: mainly learning rate and value of Nesterov momentum. We use a sequence length oftokens and a batch size of. At eval time, we consider sequences oftokens. We use cosine learning rate scheduling with peak value of. We list all the hyperparameters in the appendix (subsection 7.1) .All instances of DiPaCo use one phase of discriminative routing (see2.4.2and7.2.1) . Thepath DiPaCo uses top-2 overlapping shards (2.4.4)  at training time (this does not make training slower in wallclock or FLOPs, although it does increase the size of the shards) .4.1Comparison with 1B dense modelFigure 8:Convergence curves of DiPaCo vs Dense Baseline: We first pretrain aM parameters model for 24k training steps (purple) . We then finetune a 16x16 DiPaCo(red) . The remaining gap to the dense 1.3B parameters baseline (orange)  is reached by gating more frequently at test time.Figure 8compares the convergence curve of the 1.3B dense language model with theDiPaCo.
DiPaCo withpaths of sizeM partially sharing parameters, and one routing decision per document, is nearly sufficient to match the performance of a dense model of sizeB.
The remaining performance gap with the denseB model can be closed with early stopping and routing every 64 tokens at test time, seeTable 3.We again warn the reader that theaxis inFigure 8is not FLOP equivalent for the different models. For theM baseline model and for DiPaCo, theaxis roughly corresponds to equivalent wall-clock (and theB model is slower w.r.t. wall-clock for the same number of updates) . On the other hand, DiPaCo uses many more FLOPs and sees more tokens per training step than the dense baselines. However, it is not trivial to parallelize training of the dense models to the scale of DiPaCo, see the comparisons insubsection 4.3and in particularTable 1, and in.4.2Number of paths and parametersFigure 9:Scaling number of paths in DiPaCo: We report validation perplexity of DiPaCo with different numbers of paths and parameters. We increase capacity by varying both the number of paths (from8to256)  and by using path-specific modules. The size of a path (roughly equivalent to serving cost during deployment)  isM parameters for all models in this figure.Figure 9shows that generalization improves as more paths and more overall parameters are added to the mixture. Since the overhead of DiLoCo is minimal, the x axis could be replaced with wall-clock time. Of course,DiPaCo requirestimes more compute thanDiPaCo, but these additional paths can be trained in parallel with very little communication. In this experiment the inner optimization consists ofsteps. For all experiments withpath-specific modules, the transformer blocks,,, and the embedding matrix are not communicated across paths.4.3Parameter sharing between pathsModelTimeCompute and DataTotal ParametersValidation PPLBaselineM16.23DiLoCoM15.02DiLoCoM14.96Flat MoEB14.62Flat MoEB12.76DiPaCoB14.86DiPaCoB13.37DiPaCoPath Specific ModulesB12.70Baseline,stepsM14.72Table 1:DiPaCo vs. Flat MoE vs. DiLoCo: We compare DiLoCo and Flat MoE (2.6.3)  vs. DiPaCo. We also compare against the baseline trained for the same FLOPs as the distributed models.
In DiLoCo all paths are collapsed together at each outer optimization step. In Flat MoE, all paths are independent. In DiPaCo, only a subset of the parameters are “collapsed” using information from a subset of paths. Because of the distributed framework, DiPaCo is trained in the same wallclock time as DiLoCo, which is roughly the wallclock of the baseline.At one end of the spectrum, as described insubsubsection 2.6.3, we use paths that share no parameters, and can be trained fully independently. Continuing the trend fromsubsection 4.2, where there is monotone improvement in validation performance as parameters are added, we see improvements by unsharing all parameters for small number of paths. Using 64 unshared paths we can approach aDiPaCo that has no unsharing.However, this model starts over-fitting once the number of shards (and corresponding experts)  is further increased as shown inTable 2. We are able to recover some performance by overlapping shards, but still we find that we cannot naively scale the number of fully independent paths. We consider approaches for more sophisticated overlapping of shards and path branching to be a promising direction for future work.At the other end of the spectrum, paths can share all parameters, effectively collapsing back into a single path and recovering the performance of the original DiLoCo.
We can see that the base model trained with DiLoCo, or base model overtrained for many steps, although receiving the same number of FLOPs and tokens as the larger DiPaCo models, do not have the capacity to make use of these extra FLOPs. Of course a properly sized baseline would do much better; and indeed the 1.3B model above already has better perplexity. However, DiPaCo consists of paths that each require onlyM size island of compute both in train and testing, and its distributed training algorithm is built into the architecture, so that only paths ofM parameters ever need to be materialized on any training or deployment worker. We compare all these variants inTable 1.4.4Routing frequency during evaluationWe demonstrate the results of routing more frequently at test time inTable 3. We observe a significant improvement ofperplexity points going from routing once per sequence oftokens to everytokens. Further increasing the granularity of the re-routing consistently improve results, although more marginally so.
This shows that at training time we can route at the sequence level to enable pre-sharding for more efficient learning, and then at evaluation time we can gain back performance by routing more frequently. By routing tokens potentially to a different path every 64 tokens, DiPaCo reachesperplexity, matching the performance of aB dense model () . Since at any time DiPaCo uses a single path of sizeM parameters, it matches the performance of theB model usingtimes less parameters, and therefore, requiring less compute at inference time.# fully independent pathsValidation PPL14.613.914.2Table 2:Flat MoE (independent paths)  overfits as number of paths increase: Validation perplexity of
DiPaCo with no shared modules, as described insubsubsection 2.6.3afterK steps of training. With overlapping shards as described insubsubsection 2.4.4and early stopping, we can improve theresult to, but still the models overfit atK steps withpaths. In contrast, there is no overfitting atpaths with overlapping shards with aDiPaCo.Early StoppingRoute EveryPerplexity✗Once per sequence12.39✓12.22✓12811.486411.383211.311611.26Table 3:Frequent Routing at Eval-Time: Validation perplexity of aDiPaCo with. Despite using paths of sizeM (fewer parameters) , we match the performance of a denseB model ()  by potentially re-routing to a different path everytokens, although the model may choose to route to the exact same path that it selected previously.4.5Parameter synchronization with DiLoCo vs. true gradientsFinally, we ablate the effect of using the partially synchronous optimization DiLoCo to understand whether we lose performance by communicating less frequently at training time. In the setting where all devices are actually collocated and there is no constraint in terms of communication, we can train the DiPaCo model in a fully synchronous manner without using DiLoCo: At every step, each path computes gradients onits own batch of datafrom its own data shard; Gradients across all paths are then exchanged and aggregated module by module; Finally, the model performs one step of AdamW update with the aggregated gradient. To our surprise, we did not observe a significant performance gap between the fully synchronous training setting and the partially synchronous training setting. More specifically, DiPaCo trained with DiLoCo slightly outperforms their fully-synchronously-trained version byandperplexity points when using aandarchitecture, respectively. AtDiPaCo trained fully synchronously reaches better perplexity by onlyperplexity despite communicating hundred of times more. This suggests that DiLoCo is an effective distributed optimization algorithm for DiPaCo.
body
DiPaCo: Distributed Path Composition
4Experiments
2403.10616
5Related WorkAs mentioned in the introduction, this work shares the same motivation and intuitions expressed in Pathways. Unlike the Pathways frameworkwhich supports training of general modular multimodal multitask asynchronous systems, we propose a particular instantiation of a modular system that supports such kind of distributed training.
Our approach also shares motivations and intuitions with. The key difference in this work is that each worker trains apaththrough modules, rather than a module.5.1ModularityThere is a large body of literature on modularity, although this is not always framed in such terms.
At one end of the spectrum there are approaches that use a very large number of very tiny modules.
For instance, RETROis one such approach, whereby there is one module per token n-gram of a large retrieval dataset, and the modules are merely vectors of biases.
At the other end of the spectrum there are approaches that use very few (e.g., two)  but very large modules. For instance, Flamingois a vision language model composed by two main modules, a pre-trained visual encoder and a pre-trained language model. Similarly,proposed another approach where modules are pre-trained language models and speech recognition encoders. The authors showed how such modules could be plugged in different ways to enable 0-shot learning of entirely new tasks, such as producing a speech recognizer that translates into another language.These lines of work are representative of the vast majority of the literature on modularity.
There is however also some work where the number of modules and their size is intermediate.
For instance,showed how softly gated modules could be used in visual question answering to handle compositional tasks. In multi-task cross-lingual transfer,used instead a network where modules are very small adapters, and the vast majority of parameters are shared across paths.
We refer the reader to the survey byfor additional pointers, and focus next on modularity via mixture models.5.2Mixture of ExpertsTogether with ensembling and boosting, mixture of experts (MoE)  are an early approach to modularity in the machine learning literature. Unlike ensembling, mixture models could be more efficient if an input is processed only by a few experts. Unlike boosting, training of mixture models is efficient because it can be parallelized, while in boosting the process is inherently sequential.While early work considered the problem of hierarchical gating of flat mixtures, mixtures at multiple levels appeared only more recently. For instance,proposed a two-level mixture where paths were softly mixed and trained jointly via stochastic gradient descent. In that paper, the authors recognized the difficulty of learning non-degenerate routing functions and discussed some workarounds. The issue of how to learn a good routing function has been a major topic of research in the field since then.In their seminal work,proposed a very large mixture of experts LSTM model for sequence modeling tasks. Most works MoE for sequence-modeling works that followedhave used a recipe whereby FFN layers of transformers are replaced by mixtures. These MoEs operate at the token level, and training operates similarly to dense model training: all paths are synchronized at every gradient descent step.
Token-MoE are currently state of the art with respect to training FLOP efficiency, but require even more co-located accelerators than the equivalent-activated dense model for training.In contrast,trains experts independently using a document level router; this approach had been used in computer vision by, and it also appeared in the federated learning literature. Our flat MoE baseline is closely related to the approach in, except we used a discriminative router as opposed to-Means, and we used the first 32 tokens of a document (and the features from the base Transformer LM)  as the input to the router, as opposed to SVD of tf-idf vectors.Finally, the idea to iteratively shard the data by domain to train models in parallel, followed by a model merging phase has also been proposed by. In this work, we take this approach a step further by learning a mixture of experts model, and by automatically (and even discriminatively)  sharding the data.
body
DiPaCo: Distributed Path Composition
5Related Work
2403.10616
6LimitationsThe most salient limitation to DiPaCo is with respect to FLOP efficiency. In this work we made no effort to optimize its FLOP efficiency; and in the results presented, DiPaCo is significantly less FLOP efficient per evaluation perplexity than a standard dense compute optimal model.This work also does not study scaling laws of DiPaCo. We work at a single path-size scale, and on one dataset.We also made no effort in this work to optimize deployment of DiPaCo. In particular, if we were to naively route more frequently during deployment as in Section (2.4.3) , we would need to re-compute the-cache for each query after each re-route.The severity of these limitations depends on the constraints of the training and deployment infrastructure, as well as on the task at hand.
For instance, they may matter less when the computing infrastructure consists of hundreds of small and poorly connected islands of devices for tasks where sequences can be mostly associated to one (or very few)  domains.
In general, addressing these limitations constitutes avenue of future research.
body
DiPaCo: Distributed Path Composition
6Limitations
2403.10616
7Conclusions and future workIn this work, we architect an ML model as a union of modules, and we define input-output mappings that are paths through these modules. We show that we can train paths (almost)  independently, and (virtually)  recombine these paths back into the desired large model.
In particular, we were able to show that distributed training by path can match the performance of a 1B dense language model for a similar wall clock time.The design of the algorithm and architecture is guided by practical trade-offs. In our work we assume that communication is costly but compute is not. In other words, the compute that counts the most is the one of co-located devices that host a path. As long as communication between such relatively small islands of devices is limited, we can scale by distributing computation across a large number of such islands.We introduced two key ingredients to operate under these constraints: 1)  we route coarsely at the sequence level and pre-shard data by path, and 2)  we extend DiLoCo to update parameters of modules shared by multiple paths. We believe that by combining these two ideas we have made a step towards a potentially much more scalable paradigm to train large-models.There are several avenues of future work. In our investigations we made no effort to control the FLOP efficiency of DiPaCo, focusing on wall-clock efficiency. Naturally, there are several straightforward design choices that could increase FLOP efficiency, for example allowing some paths to co-locate. In addition, our approaches to sharding are fairly simple, and there are many interesting possibilities for more sophisticated sharding. DiPaCo was designed with continual learning in mind, and we would like to apply it in that setting eventually. Last but certainly not least, scaling up is of utmost interest.
Shared with, our long-term dream is to further refine this approach and produce a never-ending, community-driven, modular learning system that can be used by everyone to compose new predictors out of existing modules, and thus efficiently develop entirely new models and capabilities in a positive feedback loop.\nobibliography*
body
DiPaCo: Distributed Path Composition
7Conclusions and future work
2403.10616
AcknowledgementsWe would like to thank Owen He, Diego de las Casas, Ross Hemsley, Elena Gribovskaya, Yee Whye Teh, Yutian Chen, Alexandre Galashov, Bogdan Damoc, Shreya Pathak, Amal Rannen-Triki, Alek Andreev, David Budden, Jörg Bornschein for their valuable feedback throughout the development of the project.
We would also like to thank Jeff Dean, Satinder Baveja, and Raia Hadsell for their suggestions, feedback and support. Finally, we would like to thank Kitty Stacpoole and Guy Scully for their program management support.
body
DiPaCo: Distributed Path Composition
Acknowledgements
2403.10616
Supplementary Materials7.1Additional DetailsTable 4:Hyperparameter for Inner Optimization.Hyperparameter150M model1B model# Layers1224Hidden dimension8962048# Heads1616Key/value size64128Weight decay0.10.1Max learning rateBatch size512512# Warmup steps10001000We list inTable 4the hyperparameters used in the inner optimization. The150M modelhyperparameters are used for both the dense baselines of that size and for the path optimization of DiPaCo. In all our experiments, we used as outer optimizer Nesterovwith outer learning rate ofand outer momentum of, following the same recipe introduced by.7.2Routing DetailsWe first discuss the details of discriminative coarse routing, as used in training. We then discuss the details of more frequent routing, as optionally used in evaluation.7.2.1Document level discriminative routingFor discriminative routing, we assume that we have•Trained an initial language model that can be used as an initialization for finetuning experts.•Set up a feature extractor that gets a feature for each document to feed to the router.We alternate between two stages:1.Train a gater, perhaps based on the PPL of current experts on an unused portion of the training set,2.Train experts based on a sharding from current gater.In this work, the feature for the router is always the average of the hidden state from the last transformer block from the initial LM over the first 32 tokens of a document; let us denote this as. We also reserve a small subset of documents (from here on called the “router data”)  for making decisions about router quality and training the router; we use .005 of C4 for router data.Our initial router is always constructed via k-means; and we train the first set of paths using that router (and sharding by argmin to cluster centroid) .To train the discriminative router, we first compute the summed auto-regressive log-likelihood of each document in the router data if encoded by each path. That is, givendocuments of lengthand pathswe get anarray of scoreswhereand whereis theth token in theth document.
The router is always trained using aclass linear logistic classifier withas the target andas the feature.With large numbers of paths, we saw that the paths that were assigned the fewest documents viawere even more under-represented by the output of the trained logistic regressor, and it would often happen that a path was empty when using the output of the regressor. To remedy this, we trained a bias term to match the target document-to-path distribution.We find that•One step of training the discriminative router leads to a significant improvement in validation perplexity, more alternating steps lead to further minor improvements.•Gains are larger with more paths.Figure 10:Generative and discriminative flat MoE models with varying number of paths, three alternating discriminative phases. The “branching” structure of the figure is due to the ancestry of the models: at the far left, all share an initial dense ancestor. At 16K steps, each discriminative-generative pair of experiments shares a generative ancestor, this is continued (restarting the cosine learning rate decay)  with the same router (“generative”)  or with a sequence of new routers (“discriminative”) Figure 11:Validation perplexity vs number of alternating minimization steps, 16 path flat MoE. As the number of alternating steps increases, the results improve, but each step gives less improvement:PPL for phases 0, 1, 2, 3 respectively. All discriminative routing results in the main text use one alternating minimization phase.7.2.2More frequent routing during evaluationAs before, we have the arraygiving the scores of each token in the router data set of documents. We fix a number, and build a sequence transduction training set where the inputs are the token sequencesfrom the reserved portion of the training documents, and the outputsare
defined bywhere, andis a window size. Once we have, we can train a transformer to transducefrom. We finetune this router from the same model that we fork to train the paths, using the same hyperparameters as our standard LM training (except no warmup) , and it converges after a 2K steps. The learned router can then make a decision at any token.We found that choosingto be the same size as the number of tokens before re-routing (as in Table3)  led to the best results, but was only marginally better than choosinglarge enough so the window always went to the end of the document. The results in Table3use this choice of(the length of the whole sequence,) . We then can use the same learned router for any choice of evaluation-time frequency.InTable 5we show the effect of differentsharding methodson aDiPaCo. We observe that sharding makes a significant difference, with the discriminative variant yielding the best generalization, with an absolute gain ofperplexity points over the generative-means variant. We found that longer training runs yield even greater gain from the discriminative gating.7.3More sophisticated generative routingTable 5:Sharding Impact on 8x8 DiPaCo: Validation perplexity after 32 outer optimization steps, each consisting of 62 inner optimization steps with a 8x8 DiPaCo with. Longer training further increases the gap between methods. The discriminative method is based on product-Means.ShardingValidation Perplexity-Means17.2Product-Means16.8Discriminative16.5We explored Product-Means as generative router for DiPaCo.
Whenis large,-Means can become inefficient, because some clusters are assigned few examples. Moreover, in our use case, the structure of the shards does not match the structure of the modules.
Product-Means is a way to mitigate these issues. For a 2 level DiPaCo, the featureis divided in two groups:, withcomprising the firstdimensions, andthe second set ofdimensions.222We could create more groups, but we only used two in this work.We then perform-means on each set of features independently, and assign each sequenceto the pair. The first index refers to the assignment of, and the second index to the assignment ofas per equation1. Although each assignment takes onlypossible values, the number of possible unique pair-assignments is. Moreover, the cost of assignment grows with the square root of the number of total pair assignments. We used this approach when experimenting with hundreds of paths; while the results were an improvement over simple-Means, discriminative routing was better still; and using the product-means as a router for the initial module training before the discriminative step did not significantly improve results. Note that the discriminative routing adapts to the structure of the DiPaCo, because the paths are used to score documents to find the targets. Nevertheless, we consider alternative generative routing approaches important for further study, as discriminative scaling by scoring all paths cannot scale to large numbers of paths.
conclusion
DiPaCo: Distributed Path Composition
Supplementary Materials