<!DOCTYPE html><html lang="en">
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8">
<title>[2407.21770] MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts</title><meta property="og:description" content="We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing exp…">
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="twitter:card" content="summary">
<meta name="twitter:title" content="MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts">
<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="MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts">
<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/2407.21770">

<!--Generated on Mon Aug  5 15:48:33 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">
<p id="p1.1" class="ltx_p">1]Meta FAIR
<span id="p1.1.1" class="ltx_ERROR undefined">\contribution</span>[*]Joint first authors</p>
</div>
<h1 class="ltx_title ltx_title_document">MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts</h1>
<div class="ltx_authors">
<span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Xi Victoria Lin
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Akshat Shrivastava
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Liang Luo
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Srinivasan Iyer
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Mike Lewis
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Gargi Ghosh
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">            Luke Zettlemoyer
</span></span>
<span class="ltx_author_before">  </span><span class="ltx_creator ltx_role_author">
<span class="ltx_personname">Armen Aghajanyan
</span><span class="ltx_author_notes">
<span class="ltx_contact ltx_role_affiliation">[
</span>
<span class="ltx_contact ltx_role_email"><a href="mailto:victorialin@meta.com">victorialin@meta.com</a>
</span>
<span class="ltx_contact ltx_role_email"><a href="mailto:akshats@meta.com">akshats@meta.com</a>
</span>
<span class="ltx_contact ltx_role_email"><a href="mailto:armen.ag@live.com">armen.ag@live.com</a>
</span></span></span>
</div>

<div class="ltx_abstract">
<h6 class="ltx_title ltx_title_abstract">Abstract</h6>
<p id="id8.8" class="ltx_p">We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7<math id="id1.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id1.1.m1.1a"><mo id="id1.1.m1.1.1" xref="id1.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id1.1.m1.1b"><times id="id1.1.m1.1.1.cmml" xref="id1.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id1.1.m1.1c">\times</annotation></semantics></math> overall, with 2.6<math id="id2.2.m2.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id2.2.m2.1a"><mo id="id2.2.m2.1.1" xref="id2.2.m2.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id2.2.m2.1b"><times id="id2.2.m2.1.1.cmml" xref="id2.2.m2.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id2.2.m2.1c">\times</annotation></semantics></math> for text and 5.2<math id="id3.3.m3.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id3.3.m3.1a"><mo id="id3.3.m3.1.1" xref="id3.3.m3.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id3.3.m3.1b"><times id="id3.3.m3.1.1.cmml" xref="id3.3.m3.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id3.3.m3.1c">\times</annotation></semantics></math> for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3× overall FLOPs savings (3<math id="id4.4.m4.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id4.4.m4.1a"><mo id="id4.4.m4.1.1" xref="id4.4.m4.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id4.4.m4.1b"><times id="id4.4.m4.1.1.cmml" xref="id4.4.m4.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id4.4.m4.1c">\times</annotation></semantics></math> for text, 2.8<math id="id5.5.m5.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id5.5.m5.1a"><mo id="id5.5.m5.1.1" xref="id5.5.m5.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id5.5.m5.1b"><times id="id5.5.m5.1.1.cmml" xref="id5.5.m5.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id5.5.m5.1c">\times</annotation></semantics></math> for image). Combining MoMa with mixture-of-depths (MoD) further improves pretraining FLOPs savings to 4.2<math id="id6.6.m6.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id6.6.m6.1a"><mo id="id6.6.m6.1.1" xref="id6.6.m6.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id6.6.m6.1b"><times id="id6.6.m6.1.1.cmml" xref="id6.6.m6.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id6.6.m6.1c">\times</annotation></semantics></math> overall (text: 3.4<math id="id7.7.m7.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id7.7.m7.1a"><mo id="id7.7.m7.1.1" xref="id7.7.m7.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id7.7.m7.1b"><times id="id7.7.m7.1.1.cmml" xref="id7.7.m7.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id7.7.m7.1c">\times</annotation></semantics></math>, image: 5.3<math id="id8.8.m8.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="id8.8.m8.1a"><mo id="id8.8.m8.1.1" xref="id8.8.m8.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="id8.8.m8.1b"><times id="id8.8.m8.1.1.cmml" xref="id8.8.m8.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="id8.8.m8.1c">\times</annotation></semantics></math>),
although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa’s potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.
</p>
</div>
<div id="p2" class="ltx_para">
<span id="p2.1" class="ltx_ERROR undefined">\correspondence</span>
<p id="p2.2" class="ltx_p">, , 
<span id="p2.2.1" class="ltx_text" lang="en"></span></p>
</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">Auto-regressive mixed-modal foundation models have shown significant promise in applications requiring the processing of mixed-modal inputs and the generation of mixed-modal outputs <cite class="ltx_cite ltx_citemacro_citep">(Gemini et al., <a href="#bib.bib1" title="" class="ltx_ref">2023</a>, <a href="#bib.bib2" title="" class="ltx_ref">2024</a>; OpenAI et al., <a href="#bib.bib3" title="" class="ltx_ref">2024</a>; Lu et al., <a href="#bib.bib4" title="" class="ltx_ref">2023</a>; Chameleon Team, <a href="#bib.bib5" title="" class="ltx_ref">2024</a>)</cite>. These models have demonstrated remarkable capabilities in tasks ranging from visual question answering to multimodal content generation, pushing the boundaries of AI’s ability to understand and interact with our inherently multimodal world.</p>
</div>
<div id="S1.p2" class="ltx_para">
<p id="S1.p2.1" class="ltx_p">While a popular architecture design for mixed-modal foundation models involves fusing modality-specific encoders or decoders <cite class="ltx_cite ltx_citemacro_citep">(Gemini et al., <a href="#bib.bib1" title="" class="ltx_ref">2023</a>, <a href="#bib.bib2" title="" class="ltx_ref">2024</a>; Lu et al., <a href="#bib.bib4" title="" class="ltx_ref">2023</a>; OpenAI et al., <a href="#bib.bib3" title="" class="ltx_ref">2024</a>)</cite>, this approach can limit the model’s ability to integrate information across modalities and generate content with interleaved modalities. To address this limitation, <span id="S1.p2.1.1" class="ltx_text ltx_font_typewriter">Chameleon</span> <cite class="ltx_cite ltx_citemacro_citep">(Chameleon Team, <a href="#bib.bib5" title="" class="ltx_ref">2024</a>)</cite> recently introduced a single transformer architecture with a next-token prediction objective to model mixed-modal sequences consisting of discrete image and text tokens, allowing for seamless reasoning and generation across modalities.
<span id="S1.p2.1.2" class="ltx_text ltx_font_typewriter">Chameleon</span>, pretrained on approximately 10 trillion mixed-modal tokens, has demonstrated broad vision and language capabilities across various downstream tasks. Notably, it surpasses commercial baselines such as Gemini 1.0 Pro<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><a target="_blank" href="https://ai.google.dev/" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://ai.google.dev/</a></span></span></span> and GPT-4V<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><span id="footnote2.1" class="ltx_text ltx_font_typewriter">gpt-4-vision-preview</span>: <a target="_blank" href="https://openai.com/api/" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://openai.com/api/</a></span></span></span> in generating mixed-modal long-form responses <cite class="ltx_cite ltx_citemacro_citep">(Chameleon Team, <a href="#bib.bib5" title="" class="ltx_ref">2024</a>)</cite>. However, scaling such mixed-modal early-fusion foundation models to greater capacities presents significant computational challenges.</p>
</div>
<div id="S1.p3" class="ltx_para">
<p id="S1.p3.1" class="ltx_p">To address these challenges, we investigate the application of routed sparse architectures <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib6" title="" class="ltx_ref">2020</a>; Fedus et al., <a href="#bib.bib7" title="" class="ltx_ref">2022</a>; Clark et al., <a href="#bib.bib8" title="" class="ltx_ref">2022</a>; Jiang et al., <a href="#bib.bib9" title="" class="ltx_ref">2024</a>; Raposo et al., <a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>. These architectures have previously shown effectiveness in scaling language and vision-specific foundation models, as well as multimodal contrastive learning <cite class="ltx_cite ltx_citemacro_citep">(Mustafa et al., <a href="#bib.bib11" title="" class="ltx_ref">2022</a>)</cite>. However, their application to mixed-modal early-fusion models presents unique opportunities and challenges.</p>
</div>
<div id="S1.p4" class="ltx_para">
<p id="S1.p4.1" class="ltx_p">The insight driving our approach is the inherent heterogeneity of modalities: text and image tokens possess distinct information densities and redundancy patterns <cite class="ltx_cite ltx_citemacro_citep">(Liang et al., <a href="#bib.bib12" title="" class="ltx_ref">2023</a>)</cite>. While we integrate these tokens into a unified early-fusion architecture, we propose further optimizing this framework by incorporating modality-specific modules. This concept, which we term <em id="S1.p4.1.1" class="ltx_emph ltx_font_italic">modality-aware sparsity</em> (MaS), enables models to better capture features specific to each modality while maintaining strong cross-modality integration through partial parameter sharing and attention mechanisms. Previous work such as VLMo <cite class="ltx_cite ltx_citemacro_citep">(Bao et al., <a href="#bib.bib13" title="" class="ltx_ref">2022</a>)</cite>, BEiT-3 <cite class="ltx_cite ltx_citemacro_citep">(Wang et al., <a href="#bib.bib14" title="" class="ltx_ref">2022a</a>)</cite> and VL-MoE <cite class="ltx_cite ltx_citemacro_citep">(Shen et al., <a href="#bib.bib15" title="" class="ltx_ref">2023</a>)</cite> have applied mixture-of-modality-experts (MoME) to train vision-language encoders and masked language models. We apply this type of sparse component to the mixed-modal, early-fusion language modeling setting.</p>
</div>
<div id="S1.p5" class="ltx_para">
<p id="S1.p5.1" class="ltx_p">Specifically, we adopt <span id="S1.p5.1.1" class="ltx_text ltx_font_typewriter">Chameleon</span> as the base transformer architecture and introduce sparsity along two dimensions:</p>
</div>
<div id="S1.p6" class="ltx_para">
<ol id="S1.I1" class="ltx_enumerate">
<li id="S1.I1.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S1.I1.i1.p1" class="ltx_para">
<p id="S1.I1.i1.p1.1" class="ltx_p">Width: We apply <em id="S1.I1.i1.p1.1.1" class="ltx_emph ltx_font_italic">mixture-of-experts</em> (MoE) <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib6" title="" class="ltx_ref">2020</a>)</cite> where tokens are routed across a set of feed-forward blocks (experts) at each layer. Crucially, we divide these experts into modality-specific groups, with each group processing only tokens of its designated modality while learned routing is done within each group for semantics-based adaptivity. We dub this approach <em id="S1.I1.i1.p1.1.2" class="ltx_emph ltx_font_italic">mixture of modality-aware experts</em> (MoMa).</p>
</div>
</li>
<li id="S1.I1.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S1.I1.i2.p1" class="ltx_para">
<p id="S1.I1.i2.p1.1" class="ltx_p">Depth: We apply <em id="S1.I1.i2.p1.1.1" class="ltx_emph ltx_font_italic">mixture-of-depths</em> (MoD) <cite class="ltx_cite ltx_citemacro_citep">(Raposo et al., <a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite> where tokens can selectively skip the attention and feed-forward computation at certain layers.</p>
</div>
</li>
</ol>
</div>
<div id="S1.p7" class="ltx_para">
<p id="S1.p7.1" class="ltx_p">For both dimensions, we adopt expert-choice routing <cite class="ltx_cite ltx_citemacro_citep">(Zhou et al., <a href="#bib.bib16" title="" class="ltx_ref">2022</a>)</cite> to ensure load balancing and maintain a static computation graph <cite class="ltx_cite ltx_citemacro_citep">(Paszke et al., <a href="#bib.bib17" title="" class="ltx_ref">2019</a>)</cite>, both important for high training throughput.</p>
</div>
<div id="S1.p8" class="ltx_para">
<p id="S1.p8.9" class="ltx_p">We conduct extensive FLOPs-controlled experiments comparing our proposed architecture with the dense baseline and multiple sparse variations. With a 1-trillion-token training budget, <span id="S1.p8.9.1" class="ltx_text ltx_font_typewriter">Chameleon-MoMa</span> 1.4B employing 4 text experts and 4 image experts achieves a substantial 3.7<math id="S1.p8.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.1.m1.1a"><mo id="S1.p8.1.m1.1.1" xref="S1.p8.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.1.m1.1b"><times id="S1.p8.1.m1.1.1.cmml" xref="S1.p8.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.1.m1.1c">\times</annotation></semantics></math> overall FLOPs savings (text: 2.6<math id="S1.p8.2.m2.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.2.m2.1a"><mo id="S1.p8.2.m2.1.1" xref="S1.p8.2.m2.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.2.m2.1b"><times id="S1.p8.2.m2.1.1.cmml" xref="S1.p8.2.m2.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.2.m2.1c">\times</annotation></semantics></math>, image: 5.2<math id="S1.p8.3.m3.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.3.m3.1a"><mo id="S1.p8.3.m3.1.1" xref="S1.p8.3.m3.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.3.m3.1b"><times id="S1.p8.3.m3.1.1.cmml" xref="S1.p8.3.m3.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.3.m3.1c">\times</annotation></semantics></math>) compared to the 1.4B compute-matched dense baseline as measured by pre-training loss,
while maintaining a relatively modest -17% throughput reduction.
In contrast, the standard expert-choice MoE with 8 mixed-modal experts achieves 3<math id="S1.p8.4.m4.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.4.m4.1a"><mo id="S1.p8.4.m4.1.1" xref="S1.p8.4.m4.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.4.m4.1b"><times id="S1.p8.4.m4.1.1.cmml" xref="S1.p8.4.m4.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.4.m4.1c">\times</annotation></semantics></math> FLOPs savings (text: 3<math id="S1.p8.5.m5.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.5.m5.1a"><mo id="S1.p8.5.m5.1.1" xref="S1.p8.5.m5.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.5.m5.1b"><times id="S1.p8.5.m5.1.1.cmml" xref="S1.p8.5.m5.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.5.m5.1c">\times</annotation></semantics></math>, image: 2.8<math id="S1.p8.6.m6.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.6.m6.1a"><mo id="S1.p8.6.m6.1.1" xref="S1.p8.6.m6.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.6.m6.1b"><times id="S1.p8.6.m6.1.1.cmml" xref="S1.p8.6.m6.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.6.m6.1c">\times</annotation></semantics></math>) under the same setting, with -9% throughput reduction.
Further combination with MoD (<span id="S1.p8.9.2" class="ltx_text ltx_font_typewriter">Chameleon-MoMaD</span>) boosts the FLOPs savings to 4.2<math id="S1.p8.7.m7.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.7.m7.1a"><mo id="S1.p8.7.m7.1.1" xref="S1.p8.7.m7.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.7.m7.1b"><times id="S1.p8.7.m7.1.1.cmml" xref="S1.p8.7.m7.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.7.m7.1c">\times</annotation></semantics></math> overall (text: 3.4<math id="S1.p8.8.m8.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.8.m8.1a"><mo id="S1.p8.8.m8.1.1" xref="S1.p8.8.m8.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.8.m8.1b"><times id="S1.p8.8.m8.1.1.cmml" xref="S1.p8.8.m8.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.8.m8.1c">\times</annotation></semantics></math>, image: 5.3<math id="S1.p8.9.m9.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S1.p8.9.m9.1a"><mo id="S1.p8.9.m9.1.1" xref="S1.p8.9.m9.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S1.p8.9.m9.1b"><times id="S1.p8.9.m9.1.1.cmml" xref="S1.p8.9.m9.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S1.p8.9.m9.1c">\times</annotation></semantics></math>) as measured by pre-training loss. However, the auto-regressive inference performance of the MoD model is compromised due to its sensitivity to routing accuracy, resulting in subpar performance compared to <span id="S1.p8.9.3" class="ltx_text ltx_font_typewriter">Chameleon-MoMa</span> (§<a href="#S4.SS6" title="4.6 Inference-time Performance ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.6</span></a>).</p>
</div>
<div id="S1.p9" class="ltx_para">
<p id="S1.p9.1" class="ltx_p">We also demonstrate that it is possible to enhance the performance of <span id="S1.p9.1.1" class="ltx_text ltx_font_typewriter">Chameleon-MoMa</span> using a modality-untied upcycling technique <cite class="ltx_cite ltx_citemacro_citep">(Komatsuzaki et al., <a href="#bib.bib18" title="" class="ltx_ref">2023</a>)</cite>. In particular, we can initialize the model with a seed sparse architecture consisting of 1 expert per modality, and effectively improve router learning for each modality. By training the seed model for only 10k steps, the resulting model achieves additional FLOP reduction while maintaining or improving performance.</p>
</div>
<div id="S1.p10" class="ltx_para">
<p id="S1.p10.1" class="ltx_p">Our main contributions can be summarized as follows:</p>
</div>
<div id="S1.p11" class="ltx_para">
<ol id="S1.I2" class="ltx_enumerate">
<li id="S1.I2.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S1.I2.i1.p1" class="ltx_para">
<p id="S1.I2.i1.p1.1" class="ltx_p">We introduce mixture of modality-aware experts (MoMa) for mixed-modal early-fusion models, allowing for more efficient parameter allocation and utilization.</p>
</div>
</li>
<li id="S1.I2.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S1.I2.i2.p1" class="ltx_para">
<p id="S1.I2.i2.p1.1" class="ltx_p">We investigated sparse scaling along both the width and depth dimensions of transformers, demonstrating substantial speedup of pre-training loss convergence can be realized in each direction. However, effectively combining these two approaches for better performance in an autor-regressive inference setup remains an open research challenge. </p>
</div>
</li>
<li id="S1.I2.i3" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">3.</span> 
<div id="S1.I2.i3.p1" class="ltx_para">
<p id="S1.I2.i3.p1.1" class="ltx_p">We present a modality-untied upcycling technique that further enhances the efficiency of our sparse architecture.</p>
</div>
</li>
<li id="S1.I2.i4" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">4.</span> 
<div id="S1.I2.i4.p1" class="ltx_para">
<p id="S1.I2.i4.p1.1" class="ltx_p">We provide extensive empirical results and analysis, offering insights into our approach’s scaling behavior and efficiency gains.</p>
</div>
</li>
</ol>
</div>
</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>Model</h2>

<figure id="S2.F1" class="ltx_figure"><img src="/html/2407.21770/assets/x1.png" id="S2.F1.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="472" height="192" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S2.F1.2.1.1" class="ltx_text" style="font-size:90%;">Figure 1</span>: </span><span id="S2.F1.3.2" class="ltx_text" style="font-size:90%;">Overview of our proposed multimodal early-fusion architecture. (a) The early-fusion architecture processes interleaved text and image data as a unified token sequence using a transformer model. (b) Detailed architecture of the mixture of modality-aware experts (MoMa) transformer block, where layer norms and the residual connection of attention layer are omitted for readability. The input sequence of tokens are first routed to different expert groups based on their modality, undergo learned routing within each group, and then have their output restored to a tensor of the original input shape of the FFN module.
</span></figcaption>
</figure>
<section id="S2.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">2.1 </span>Early Fusion</h3>

<div id="S2.SS1.p1" class="ltx_para">
<p id="S2.SS1.p1.1" class="ltx_p">Our model builds upon the early fusion architecture introduced by <span id="S2.SS1.p1.1.1" class="ltx_text ltx_font_typewriter">Chameleon</span> <cite class="ltx_cite ltx_citemacro_citep">(Chameleon Team, <a href="#bib.bib5" title="" class="ltx_ref">2024</a>)</cite>, which represents images and text as a series of discrete tokens within a unified transformer.
The core of <span id="S2.SS1.p1.1.2" class="ltx_text ltx_font_typewriter">Chameleon</span> is a transformer-based model that applies self-attention mechanisms over the combined sequence of image and text tokens. This allows the model to capture complex relationships between and within modalities. The model is trained using a next-token prediction objective, learning to generate both text and image tokens autoregressively.</p>
</div>
<div id="S2.SS1.p2" class="ltx_para">
<p id="S2.SS1.p2.1" class="ltx_p">In <span id="S2.SS1.p2.1.1" class="ltx_text ltx_font_typewriter">Chameleon</span>, images are tokenized using a learned image tokenizer that encodes a 512 <math id="S2.SS1.p2.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S2.SS1.p2.1.m1.1a"><mo id="S2.SS1.p2.1.m1.1.1" xref="S2.SS1.p2.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S2.SS1.p2.1.m1.1b"><times id="S2.SS1.p2.1.m1.1.1.cmml" xref="S2.SS1.p2.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S2.SS1.p2.1.m1.1c">\times</annotation></semantics></math> 512 image into 1024 discrete tokens from a codebook of size 8192. Text is tokenized using a BPE tokenizer with a vocabulary size of 65,536, which includes the 8192 image codebook tokens. This unified tokenization scheme enables the model to process arbitrary sequences of interleaved image and text tokens.</p>
</div>
<div id="S2.SS1.p3" class="ltx_para">
<p id="S2.SS1.p3.1" class="ltx_p">By adopting this early fusion approach, our model inherits several key advantages:</p>
<ol id="S2.I1" class="ltx_enumerate">
<li id="S2.I1.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S2.I1.i1.p1" class="ltx_para">
<p id="S2.I1.i1.p1.1" class="ltx_p"><em id="S2.I1.i1.p1.1.1" class="ltx_emph ltx_font_italic">Unified representation:</em> The model learns a shared representation space for images and text, facilitating cross-modal reasoning and generation.</p>
</div>
</li>
<li id="S2.I1.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S2.I1.i2.p1" class="ltx_para">
<p id="S2.I1.i2.p1.1" class="ltx_p"><em id="S2.I1.i2.p1.1.1" class="ltx_emph ltx_font_italic">Flexibility:</em> The architecture can handle arbitrary sequences of images and text, enabling diverse multimodal tasks such as image captioning, visual dialogue, and mixed-modal document generation.</p>
</div>
</li>
<li id="S2.I1.i3" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">3.</span> 
<div id="S2.I1.i3.p1" class="ltx_para">
<p id="S2.I1.i3.p1.1" class="ltx_p"><em id="S2.I1.i3.p1.1.1" class="ltx_emph ltx_font_italic">Scalability:</em> The token-based approach allows for uniform processing of both text and image data, enabling efficient scaling to larger model sizes and diverse datasets. This is evidenced by <span id="S2.I1.i3.p1.1.2" class="ltx_text ltx_font_typewriter">Chameleon</span>’s successful training on approximately 10 trillion mixed-modal tokens.</p>
</div>
</li>
<li id="S2.I1.i4" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">4.</span> 
<div id="S2.I1.i4.p1" class="ltx_para">
<p id="S2.I1.i4.p1.1" class="ltx_p"><em id="S2.I1.i4.p1.1.1" class="ltx_emph ltx_font_italic">End-to-end learning:</em> The entire auto-regressive model, is trained end-to-end, allowing for joint optimization of the representation and task-specific performance.</p>
</div>
</li>
</ol>
</div>
<div id="S2.SS1.p4" class="ltx_para">
<p id="S2.SS1.p4.1" class="ltx_p">Building on this foundation (Figure <a href="#S2.F1" title="Figure 1 ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">1</span></a>a), our work introduces modality-aware sparsity techniques to further enhance the efficiency and performance of early fusion models, as detailed in the following sections. These techniques aim to address the computational challenges of scaling early fusion models while maintaining their powerful cross-modal reasoning capabilities.</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>Width Scaling: Mixture of Modality-Aware Experts</h3>

<div id="S2.SS2.p1" class="ltx_para">
<p id="S2.SS2.p1.1" class="ltx_p">We propose a width scaling approach that incorporates modality-aware block sparsity in the feed-forward module, extending the standard mixture-of-experts (MoE) architecture <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib6" title="" class="ltx_ref">2020</a>; Fedus et al., <a href="#bib.bib7" title="" class="ltx_ref">2022</a>; Wang et al., <a href="#bib.bib19" title="" class="ltx_ref">2022b</a>)</cite>. The insight driving this approach is that tokens of different modalities have distinct characteristics and information densities. By creating separate expert groups for each modality, we allow the model to develop specialized processing pathways while maintaining the ability to integrate information across modalities.</p>
</div>
<div id="S2.SS2.p2" class="ltx_para">
<p id="S2.SS2.p2.1" class="ltx_p">We describe the key components in the mixture of modality-aware experts (MoMa) formulation below (Figure <a href="#S2.F1" title="Figure 1 ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">1</span></a>b).</p>
</div>
<section id="S2.SS2.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Modality-Specific Expert Groups.</h4>

<div id="S2.SS2.SSS0.Px1.p1" class="ltx_para">
<p id="S2.SS2.SSS0.Px1.p1.1" class="ltx_p">We divide the experts in each MoE layer into distinct groups, each specialized in processing tokens from a specific modality <math id="S2.SS2.SSS0.Px1.p1.1.m1.4" class="ltx_Math" alttext="\mathcal{E}=\{E_{1}^{\mathcal{T}}\ldots,E_{m}^{\mathcal{T}},E_{1}^{\mathcal{I}}\ldots,E_{n}^{\mathcal{T}}\}" display="inline"><semantics id="S2.SS2.SSS0.Px1.p1.1.m1.4a"><mrow id="S2.SS2.SSS0.Px1.p1.1.m1.4.4" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.6" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.6.cmml">ℰ</mi><mo id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.5" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.5.cmml">=</mo><mrow id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml"><mo stretchy="false" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.5" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml">{</mo><mrow id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.cmml"><msubsup id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.cmml"><mi id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.2.cmml">E</mi><mn id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.3.cmml">1</mn><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.3.cmml">𝒯</mi></msubsup><mo lspace="0em" rspace="0em" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.1" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.1.cmml">​</mo><mi mathvariant="normal" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.3.cmml">…</mi></mrow><mo id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.6" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml">,</mo><msubsup id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.cmml"><mi id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.2.cmml">E</mi><mi id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.3.cmml">m</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.3.cmml">𝒯</mi></msubsup><mo id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.7" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml">,</mo><mrow id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.cmml"><msubsup id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.cmml"><mi id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.2.cmml">E</mi><mn id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.3.cmml">1</mn><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.3.cmml">ℐ</mi></msubsup><mo lspace="0em" rspace="0em" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.1" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.1.cmml">​</mo><mi mathvariant="normal" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.3.cmml">…</mi></mrow><mo id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.8" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml">,</mo><msubsup id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.cmml"><mi id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.2" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.2.cmml">E</mi><mi id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.3.cmml">n</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.3" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.3.cmml">𝒯</mi></msubsup><mo stretchy="false" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.9" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml">}</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px1.p1.1.m1.4b"><apply id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4"><eq id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.5.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.5"></eq><ci id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.6.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.6">ℰ</ci><set id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.5.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4"><apply id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1"><times id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.1"></times><apply id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2">superscript</csymbol><apply id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2">subscript</csymbol><ci id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.2">𝐸</ci><cn type="integer" id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.2.3">1</cn></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.2.3">𝒯</ci></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.1.1.1.1.1.3">…</ci></apply><apply id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2">superscript</csymbol><apply id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2">subscript</csymbol><ci id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.2">𝐸</ci><ci id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.2.3">𝑚</ci></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.2.2.2.2.2.3">𝒯</ci></apply><apply id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3"><times id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.1"></times><apply id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2">superscript</csymbol><apply id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2">subscript</csymbol><ci id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.2">𝐸</ci><cn type="integer" id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.2.3">1</cn></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.2.3">ℐ</ci></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.3.3.3.3.3.3">…</ci></apply><apply id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4">superscript</csymbol><apply id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.1.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4">subscript</csymbol><ci id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.2.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.2">𝐸</ci><ci id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.2.3">𝑛</ci></apply><ci id="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.3.cmml" xref="S2.SS2.SSS0.Px1.p1.1.m1.4.4.4.4.4.3">𝒯</ci></apply></set></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px1.p1.1.m1.4c">\mathcal{E}=\{E_{1}^{\mathcal{T}}\ldots,E_{m}^{\mathcal{T}},E_{1}^{\mathcal{I}}\ldots,E_{n}^{\mathcal{T}}\}</annotation></semantics></math>: one group for processing text tokens and another for image tokens<span id="footnote3" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">3</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">3</sup><span class="ltx_tag ltx_tag_note">3</span>While we focus on text and image modalities in this work, our formulation is designed to be easily extensible to accommodate an arbitrary number of modalities.</span></span></span>. This separation allows each group to specialize in features relevant to its respective modality.</p>
</div>
<div id="S2.SS2.SSS0.Px1.p2" class="ltx_para">
<p id="S2.SS2.SSS0.Px1.p2.1" class="ltx_p">By implementing modality-aware block sparsity, we aim to achieve several benefits:</p>
<ul id="S2.I2" class="ltx_itemize">
<li id="S2.I2.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S2.I2.i1.p1" class="ltx_para">
<p id="S2.I2.i1.p1.1" class="ltx_p">Improved Efficiency: by routing tokens to modality-specific experts, we reduce the computational overhead of processing tokens with experts not specialized for their modality.</p>
</div>
</li>
<li id="S2.I2.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S2.I2.i2.p1" class="ltx_para">
<p id="S2.I2.i2.p1.1" class="ltx_p">Enhanced Specialization: modality-specific expert groups can develop more refined features relevant to their respective modalities.</p>
</div>
</li>
<li id="S2.I2.i3" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S2.I2.i3.p1" class="ltx_para">
<p id="S2.I2.i3.p1.1" class="ltx_p">Maintained Cross-Modal Integration: despite the separation into modality-specific groups, the model still can integrate information across modalities through the shared self-attention mechanisms in non-MoE layers.</p>
</div>
</li>
</ul>
</div>
</section>
<section id="S2.SS2.SSS0.Px2" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Hierarchical Routing.</h4>

<div id="S2.SS2.SSS0.Px2.p1" class="ltx_para">
<p id="S2.SS2.SSS0.Px2.p1.1" class="ltx_p">We adopt a token-based routing mechanism <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib6" title="" class="ltx_ref">2020</a>; Fedus et al., <a href="#bib.bib7" title="" class="ltx_ref">2022</a>; Jiang et al., <a href="#bib.bib9" title="" class="ltx_ref">2024</a>)</cite>. For an input token <math id="S2.SS2.SSS0.Px2.p1.1.m1.1" class="ltx_Math" alttext="x" 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">x</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">x</annotation></semantics></math>, our routing mechanism operates in two stages.</p>
</div>
<div id="S2.SS2.SSS0.Px2.p2" class="ltx_para">
<ol id="S2.I3" class="ltx_enumerate">
<li id="S2.I3.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S2.I3.i1.p1" class="ltx_para">
<p id="S2.I3.i1.p1.2" class="ltx_p">Modality-aware routing: tokens are first routed to their corresponding modality-specific expert group based on their modality (text <math id="S2.I3.i1.p1.1.m1.1" class="ltx_Math" alttext="\mathcal{T}" display="inline"><semantics id="S2.I3.i1.p1.1.m1.1a"><mi class="ltx_font_mathcaligraphic" id="S2.I3.i1.p1.1.m1.1.1" xref="S2.I3.i1.p1.1.m1.1.1.cmml">𝒯</mi><annotation-xml encoding="MathML-Content" id="S2.I3.i1.p1.1.m1.1b"><ci id="S2.I3.i1.p1.1.m1.1.1.cmml" xref="S2.I3.i1.p1.1.m1.1.1">𝒯</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.I3.i1.p1.1.m1.1c">\mathcal{T}</annotation></semantics></math> or image <math id="S2.I3.i1.p1.2.m2.1" class="ltx_Math" alttext="\mathcal{I}" display="inline"><semantics id="S2.I3.i1.p1.2.m2.1a"><mi class="ltx_font_mathcaligraphic" id="S2.I3.i1.p1.2.m2.1.1" xref="S2.I3.i1.p1.2.m2.1.1.cmml">ℐ</mi><annotation-xml encoding="MathML-Content" id="S2.I3.i1.p1.2.m2.1b"><ci id="S2.I3.i1.p1.2.m2.1.1.cmml" xref="S2.I3.i1.p1.2.m2.1.1">ℐ</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.I3.i1.p1.2.m2.1c">\mathcal{I}</annotation></semantics></math>).</p>
</div>
</li>
<li id="S2.I3.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S2.I3.i2.p1" class="ltx_para">
<p id="S2.I3.i2.p1.2" class="ltx_p">Intra-modality routing: within each modality-specific expert group <math id="S2.I3.i2.p1.1.m1.1" class="ltx_Math" alttext="\mathcal{E}^{\mathcal{M}}" display="inline"><semantics id="S2.I3.i2.p1.1.m1.1a"><msup id="S2.I3.i2.p1.1.m1.1.1" xref="S2.I3.i2.p1.1.m1.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.I3.i2.p1.1.m1.1.1.2" xref="S2.I3.i2.p1.1.m1.1.1.2.cmml">ℰ</mi><mi class="ltx_font_mathcaligraphic" id="S2.I3.i2.p1.1.m1.1.1.3" xref="S2.I3.i2.p1.1.m1.1.1.3.cmml">ℳ</mi></msup><annotation-xml encoding="MathML-Content" id="S2.I3.i2.p1.1.m1.1b"><apply id="S2.I3.i2.p1.1.m1.1.1.cmml" xref="S2.I3.i2.p1.1.m1.1.1"><csymbol cd="ambiguous" id="S2.I3.i2.p1.1.m1.1.1.1.cmml" xref="S2.I3.i2.p1.1.m1.1.1">superscript</csymbol><ci id="S2.I3.i2.p1.1.m1.1.1.2.cmml" xref="S2.I3.i2.p1.1.m1.1.1.2">ℰ</ci><ci id="S2.I3.i2.p1.1.m1.1.1.3.cmml" xref="S2.I3.i2.p1.1.m1.1.1.3">ℳ</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.I3.i2.p1.1.m1.1c">\mathcal{E}^{\mathcal{M}}</annotation></semantics></math>, tokens are then routed to specific experts using a learned routing function. Specifically, we use a projection matrix <math id="S2.I3.i2.p1.2.m2.3" class="ltx_Math" alttext="W_{g}^{\mathcal{M}}\in\mathbb{R}^{d\times|\mathcal{E}^{\mathcal{M}}|},d=\text{transformer hidden dimension}" display="inline"><semantics id="S2.I3.i2.p1.2.m2.3a"><mrow id="S2.I3.i2.p1.2.m2.3.3.2" xref="S2.I3.i2.p1.2.m2.3.3.3.cmml"><mrow id="S2.I3.i2.p1.2.m2.2.2.1.1" xref="S2.I3.i2.p1.2.m2.2.2.1.1.cmml"><msubsup id="S2.I3.i2.p1.2.m2.2.2.1.1.2" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.cmml"><mi id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.2" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.2.cmml">W</mi><mi id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.3" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.3.cmml">g</mi><mi class="ltx_font_mathcaligraphic" id="S2.I3.i2.p1.2.m2.2.2.1.1.2.3" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.3.cmml">ℳ</mi></msubsup><mo id="S2.I3.i2.p1.2.m2.2.2.1.1.1" xref="S2.I3.i2.p1.2.m2.2.2.1.1.1.cmml">∈</mo><msup id="S2.I3.i2.p1.2.m2.2.2.1.1.3" xref="S2.I3.i2.p1.2.m2.2.2.1.1.3.cmml"><mi id="S2.I3.i2.p1.2.m2.2.2.1.1.3.2" xref="S2.I3.i2.p1.2.m2.2.2.1.1.3.2.cmml">ℝ</mi><mrow id="S2.I3.i2.p1.2.m2.1.1.1" xref="S2.I3.i2.p1.2.m2.1.1.1.cmml"><mi id="S2.I3.i2.p1.2.m2.1.1.1.3" xref="S2.I3.i2.p1.2.m2.1.1.1.3.cmml">d</mi><mo lspace="0.222em" rspace="0.222em" id="S2.I3.i2.p1.2.m2.1.1.1.2" xref="S2.I3.i2.p1.2.m2.1.1.1.2.cmml">×</mo><mrow id="S2.I3.i2.p1.2.m2.1.1.1.1.1" xref="S2.I3.i2.p1.2.m2.1.1.1.1.2.cmml"><mo stretchy="false" id="S2.I3.i2.p1.2.m2.1.1.1.1.1.2" xref="S2.I3.i2.p1.2.m2.1.1.1.1.2.1.cmml">|</mo><msup id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.2" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.2.cmml">ℰ</mi><mi class="ltx_font_mathcaligraphic" id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.3" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.3.cmml">ℳ</mi></msup><mo stretchy="false" id="S2.I3.i2.p1.2.m2.1.1.1.1.1.3" xref="S2.I3.i2.p1.2.m2.1.1.1.1.2.1.cmml">|</mo></mrow></mrow></msup></mrow><mo id="S2.I3.i2.p1.2.m2.3.3.2.3" xref="S2.I3.i2.p1.2.m2.3.3.3a.cmml">,</mo><mrow id="S2.I3.i2.p1.2.m2.3.3.2.2" xref="S2.I3.i2.p1.2.m2.3.3.2.2.cmml"><mi id="S2.I3.i2.p1.2.m2.3.3.2.2.2" xref="S2.I3.i2.p1.2.m2.3.3.2.2.2.cmml">d</mi><mo id="S2.I3.i2.p1.2.m2.3.3.2.2.1" xref="S2.I3.i2.p1.2.m2.3.3.2.2.1.cmml">=</mo><mtext id="S2.I3.i2.p1.2.m2.3.3.2.2.3" xref="S2.I3.i2.p1.2.m2.3.3.2.2.3a.cmml">transformer hidden dimension</mtext></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.I3.i2.p1.2.m2.3b"><apply id="S2.I3.i2.p1.2.m2.3.3.3.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2"><csymbol cd="ambiguous" id="S2.I3.i2.p1.2.m2.3.3.3a.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.3">formulae-sequence</csymbol><apply id="S2.I3.i2.p1.2.m2.2.2.1.1.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1"><in id="S2.I3.i2.p1.2.m2.2.2.1.1.1.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.1"></in><apply id="S2.I3.i2.p1.2.m2.2.2.1.1.2.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2"><csymbol cd="ambiguous" id="S2.I3.i2.p1.2.m2.2.2.1.1.2.1.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2">superscript</csymbol><apply id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2"><csymbol cd="ambiguous" id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.1.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2">subscript</csymbol><ci id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.2.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.2">𝑊</ci><ci id="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.3.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.2.3">𝑔</ci></apply><ci id="S2.I3.i2.p1.2.m2.2.2.1.1.2.3.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.2.3">ℳ</ci></apply><apply id="S2.I3.i2.p1.2.m2.2.2.1.1.3.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.3"><csymbol cd="ambiguous" id="S2.I3.i2.p1.2.m2.2.2.1.1.3.1.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.3">superscript</csymbol><ci id="S2.I3.i2.p1.2.m2.2.2.1.1.3.2.cmml" xref="S2.I3.i2.p1.2.m2.2.2.1.1.3.2">ℝ</ci><apply id="S2.I3.i2.p1.2.m2.1.1.1.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1"><times id="S2.I3.i2.p1.2.m2.1.1.1.2.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.2"></times><ci id="S2.I3.i2.p1.2.m2.1.1.1.3.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.3">𝑑</ci><apply id="S2.I3.i2.p1.2.m2.1.1.1.1.2.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1"><abs id="S2.I3.i2.p1.2.m2.1.1.1.1.2.1.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.2"></abs><apply id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1"><csymbol cd="ambiguous" id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.1.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1">superscript</csymbol><ci id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.2.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.2">ℰ</ci><ci id="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.3.cmml" xref="S2.I3.i2.p1.2.m2.1.1.1.1.1.1.3">ℳ</ci></apply></apply></apply></apply></apply><apply id="S2.I3.i2.p1.2.m2.3.3.2.2.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.2"><eq id="S2.I3.i2.p1.2.m2.3.3.2.2.1.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.2.1"></eq><ci id="S2.I3.i2.p1.2.m2.3.3.2.2.2.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.2.2">𝑑</ci><ci id="S2.I3.i2.p1.2.m2.3.3.2.2.3a.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.2.3"><mtext id="S2.I3.i2.p1.2.m2.3.3.2.2.3.cmml" xref="S2.I3.i2.p1.2.m2.3.3.2.2.3">transformer hidden dimension</mtext></ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.I3.i2.p1.2.m2.3c">W_{g}^{\mathcal{M}}\in\mathbb{R}^{d\times|\mathcal{E}^{\mathcal{M}}|},d=\text{transformer hidden dimension}</annotation></semantics></math>, to compute the token-to-expert affinity scores, following other established MoE formulations <cite class="ltx_cite ltx_citemacro_citep">(Lepikhin et al., <a href="#bib.bib6" title="" class="ltx_ref">2020</a>; Fedus et al., <a href="#bib.bib7" title="" class="ltx_ref">2022</a>; Jiang et al., <a href="#bib.bib9" title="" class="ltx_ref">2024</a>)</cite>.</p>
</div>
</li>
</ol>
</div>
</section>
<section id="S2.SS2.SSS0.Px3" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Expert choice.</h4>

<div id="S2.SS2.SSS0.Px3.p1" class="ltx_para">
<p id="S2.SS2.SSS0.Px3.p1.5" class="ltx_p">Within each modality group, we implemented expert-choice (EC) routing <cite class="ltx_cite ltx_citemacro_citep">(Zhou et al., <a href="#bib.bib16" title="" class="ltx_ref">2022</a>)</cite>, where each expert has a fixed bucket size and processes the top-<math id="S2.SS2.SSS0.Px3.p1.1.m1.1" class="ltx_Math" alttext="k_{e}" display="inline"><semantics id="S2.SS2.SSS0.Px3.p1.1.m1.1a"><msub id="S2.SS2.SSS0.Px3.p1.1.m1.1.1" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1.cmml"><mi id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.2" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1.2.cmml">k</mi><mi id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.3" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1.3.cmml">e</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p1.1.m1.1b"><apply id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1">subscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1.2">𝑘</ci><ci id="S2.SS2.SSS0.Px3.p1.1.m1.1.1.3.cmml" xref="S2.SS2.SSS0.Px3.p1.1.m1.1.1.3">𝑒</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p1.1.m1.1c">k_{e}</annotation></semantics></math> tokens in a batch. We set <math id="S2.SS2.SSS0.Px3.p1.2.m2.1" class="ltx_Math" alttext="k_{e}=b^{\mathcal{M}}\cdot c_{e}" display="inline"><semantics id="S2.SS2.SSS0.Px3.p1.2.m2.1a"><mrow id="S2.SS2.SSS0.Px3.p1.2.m2.1.1" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.cmml"><msub id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.cmml"><mi id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.2" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.2.cmml">k</mi><mi id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.3" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.3.cmml">e</mi></msub><mo id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.1" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.1.cmml">=</mo><mrow id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.cmml"><msup id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.cmml"><mi id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.2" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.2.cmml">b</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.3" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.3.cmml">ℳ</mi></msup><mo lspace="0.222em" rspace="0.222em" id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.1" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.1.cmml">⋅</mo><msub id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.cmml"><mi id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.2" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.2.cmml">c</mi><mi id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.3" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.3.cmml">e</mi></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p1.2.m2.1b"><apply id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1"><eq id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.1"></eq><apply id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2">subscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.2.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.2">𝑘</ci><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.3.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.2.3">𝑒</ci></apply><apply id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3"><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.1">⋅</ci><apply id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2">superscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.2.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.2">𝑏</ci><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.3.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.2.3">ℳ</ci></apply><apply id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.1.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3">subscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.2.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.2">𝑐</ci><ci id="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.3.cmml" xref="S2.SS2.SSS0.Px3.p1.2.m2.1.1.3.3.3">𝑒</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p1.2.m2.1c">k_{e}=b^{\mathcal{M}}\cdot c_{e}</annotation></semantics></math>,
where <math id="S2.SS2.SSS0.Px3.p1.3.m3.1" class="ltx_Math" alttext="c_{e}=\frac{1}{|\mathcal{E}^{\mathcal{M}}|}" display="inline"><semantics id="S2.SS2.SSS0.Px3.p1.3.m3.1a"><mrow id="S2.SS2.SSS0.Px3.p1.3.m3.1.2" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.cmml"><msub id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.cmml"><mi id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.2" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.2.cmml">c</mi><mi id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.3" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.3.cmml">e</mi></msub><mo id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.1" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.1.cmml">=</mo><mfrac id="S2.SS2.SSS0.Px3.p1.3.m3.1.1" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.cmml"><mn id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.3" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.3.cmml">1</mn><mrow id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.2.cmml"><mo stretchy="false" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.2" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.2.1.cmml">|</mo><msup id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.2" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.2.cmml">ℰ</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.3" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.3.cmml">ℳ</mi></msup><mo stretchy="false" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.3" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.2.1.cmml">|</mo></mrow></mfrac></mrow><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p1.3.m3.1b"><apply id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2"><eq id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.1"></eq><apply id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2">subscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.2">𝑐</ci><ci id="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.3.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.2.2.3">𝑒</ci></apply><apply id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1"><divide id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1"></divide><cn type="integer" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.3.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.3">1</cn><apply id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1"><abs id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.2.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.2"></abs><apply id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1">superscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.2">ℰ</ci><ci id="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.3.cmml" xref="S2.SS2.SSS0.Px3.p1.3.m3.1.1.1.1.1.3">ℳ</ci></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p1.3.m3.1c">c_{e}=\frac{1}{|\mathcal{E}^{\mathcal{M}}|}</annotation></semantics></math> is a capacity factor and <math id="S2.SS2.SSS0.Px3.p1.4.m4.1" class="ltx_Math" alttext="b^{\mathcal{M}}" display="inline"><semantics id="S2.SS2.SSS0.Px3.p1.4.m4.1a"><msup id="S2.SS2.SSS0.Px3.p1.4.m4.1.1" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1.cmml"><mi id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.2" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1.2.cmml">b</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.3" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1.3.cmml">ℳ</mi></msup><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p1.4.m4.1b"><apply id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1"><csymbol cd="ambiguous" id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1">superscript</csymbol><ci id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.2.cmml" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1.2">𝑏</ci><ci id="S2.SS2.SSS0.Px3.p1.4.m4.1.1.3.cmml" xref="S2.SS2.SSS0.Px3.p1.4.m4.1.1.3">ℳ</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p1.4.m4.1c">b^{\mathcal{M}}</annotation></semantics></math> is the total number of tokens of modality <math id="S2.SS2.SSS0.Px3.p1.5.m5.1" class="ltx_Math" alttext="\mathcal{M}" display="inline"><semantics id="S2.SS2.SSS0.Px3.p1.5.m5.1a"><mi class="ltx_font_mathcaligraphic" id="S2.SS2.SSS0.Px3.p1.5.m5.1.1" xref="S2.SS2.SSS0.Px3.p1.5.m5.1.1.cmml">ℳ</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p1.5.m5.1b"><ci id="S2.SS2.SSS0.Px3.p1.5.m5.1.1.cmml" xref="S2.SS2.SSS0.Px3.p1.5.m5.1.1">ℳ</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p1.5.m5.1c">\mathcal{M}</annotation></semantics></math> in a batch.
As a result, each token can be routed to a variable number of experts.</p>
</div>
<div id="S2.SS2.SSS0.Px3.p2" class="ltx_para">
<p id="S2.SS2.SSS0.Px3.p2.1" class="ltx_p">EC routing ensures balanced expert utilization during training and eliminates needing a separate load-balancing loss term. Maintaining a single loss term promotes more stable optimization and a more straightforward training process.
However, EC routing compromises causality in auto-regressive language modeling, as each expert selects the top tokens to process in a batch by comparing token scores across all tokens. We use a combination of two techniques to address this issue and enable expert-choice training for auto-regressive LMs.</p>
<ol id="S2.I4" class="ltx_enumerate">
<li id="S2.I4.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">1.</span> 
<div id="S2.I4.i1.p1" class="ltx_para">
<p id="S2.I4.i1.p1.1" class="ltx_p">We use Sigmoid as the non-linearity in the router scoring function, enabling independent calculation of token-to-expert affinity scores for each token.</p>
</div>
</li>
<li id="S2.I4.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">2.</span> 
<div id="S2.I4.i2.p1" class="ltx_para">
<p id="S2.I4.i2.p1.1" class="ltx_p">We introduce auxiliary routers, inspired by <cite class="ltx_cite ltx_citemacro_cite">Raposo et al. (<a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>, which predict the likelihood of an expert selecting a token solely based on its hidden state representation. These routers are trained after the main model training is completed and employed during inference to ensure causality. We discuss the details of the auxiliary routers in §<a href="#S2.SS3" title="2.3 Mixture-of-Depths ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.3</span></a>.</p>
</div>
</li>
</ol>
</div>
<div id="S2.SS2.SSS0.Px3.p3" class="ltx_para">
<p id="S2.SS2.SSS0.Px3.p3.1" class="ltx_p">In summary, the MoMa module for an input token <math id="S2.SS2.SSS0.Px3.p3.1.m1.1" class="ltx_Math" alttext="x" display="inline"><semantics id="S2.SS2.SSS0.Px3.p3.1.m1.1a"><mi id="S2.SS2.SSS0.Px3.p3.1.m1.1.1" xref="S2.SS2.SSS0.Px3.p3.1.m1.1.1.cmml">x</mi><annotation-xml encoding="MathML-Content" id="S2.SS2.SSS0.Px3.p3.1.m1.1b"><ci id="S2.SS2.SSS0.Px3.p3.1.m1.1.1.cmml" xref="S2.SS2.SSS0.Px3.p3.1.m1.1.1">𝑥</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS2.SSS0.Px3.p3.1.m1.1c">x</annotation></semantics></math> can be formally defined as:</p>
</div>
<div id="S2.SS2.SSS0.Px3.p4" class="ltx_para">
<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="y=\begin{cases}\sum_{j=1}^{m}g^{\mathcal{T}}(x)_{j}\cdot{\text{FFN}^{\mathcal{T}}_{\text{SwiGLU}}}_{j}(x)&amp;\text{if }x\in\mathcal{T}\\
\sum_{j=1}^{n}g^{\mathcal{I}}(x)_{j}\cdot{\text{FFN}^{\mathcal{I}}_{\text{SwiGLU}}}_{j}(x)&amp;\text{if }x\in\mathcal{I}\end{cases}" display="block"><semantics id="S2.E1.m1.4a"><mrow id="S2.E1.m1.4.5" xref="S2.E1.m1.4.5.cmml"><mi id="S2.E1.m1.4.5.2" xref="S2.E1.m1.4.5.2.cmml">y</mi><mo id="S2.E1.m1.4.5.1" xref="S2.E1.m1.4.5.1.cmml">=</mo><mrow id="S2.E1.m1.4.4" xref="S2.E1.m1.4.5.3.1.cmml"><mo id="S2.E1.m1.4.4.5" xref="S2.E1.m1.4.5.3.1.1.cmml">{</mo><mtable columnspacing="5pt" displaystyle="true" rowspacing="0pt" id="S2.E1.m1.4.4.4" xref="S2.E1.m1.4.5.3.1.cmml"><mtr id="S2.E1.m1.4.4.4a" xref="S2.E1.m1.4.5.3.1.cmml"><mtd class="ltx_align_left" columnalign="left" id="S2.E1.m1.4.4.4b" xref="S2.E1.m1.4.5.3.1.cmml"><mrow id="S2.E1.m1.1.1.1.1.1.1" xref="S2.E1.m1.1.1.1.1.1.1.cmml"><mstyle displaystyle="false" id="S2.E1.m1.1.1.1.1.1.1.3" xref="S2.E1.m1.1.1.1.1.1.1.3.cmml"><msubsup id="S2.E1.m1.1.1.1.1.1.1.3a" xref="S2.E1.m1.1.1.1.1.1.1.3.cmml"><mo id="S2.E1.m1.1.1.1.1.1.1.3.2.2" xref="S2.E1.m1.1.1.1.1.1.1.3.2.2.cmml">∑</mo><mrow id="S2.E1.m1.1.1.1.1.1.1.3.2.3" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.cmml"><mi id="S2.E1.m1.1.1.1.1.1.1.3.2.3.2" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.2.cmml">j</mi><mo id="S2.E1.m1.1.1.1.1.1.1.3.2.3.1" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.1.cmml">=</mo><mn id="S2.E1.m1.1.1.1.1.1.1.3.2.3.3" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.3.cmml">1</mn></mrow><mi id="S2.E1.m1.1.1.1.1.1.1.3.3" xref="S2.E1.m1.1.1.1.1.1.1.3.3.cmml">m</mi></msubsup></mstyle><mrow id="S2.E1.m1.1.1.1.1.1.1.4" xref="S2.E1.m1.1.1.1.1.1.1.4.cmml"><mrow id="S2.E1.m1.1.1.1.1.1.1.4.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.cmml"><mrow id="S2.E1.m1.1.1.1.1.1.1.4.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.cmml"><msup id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.cmml"><mi id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.2.cmml">g</mi><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.3.cmml">𝒯</mi></msup><mo lspace="0em" rspace="0em" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.1" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.1.cmml">​</mo><msub id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.cmml"><mrow id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.cmml"><mo stretchy="false" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.2.2.1" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.cmml">(</mo><mi id="S2.E1.m1.1.1.1.1.1.1.1" xref="S2.E1.m1.1.1.1.1.1.1.1.cmml">x</mi><mo rspace="0.055em" stretchy="false" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.2.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.cmml">)</mo></mrow><mi id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.3.cmml">j</mi></msub></mrow><mo rspace="0.222em" id="S2.E1.m1.1.1.1.1.1.1.4.2.1" xref="S2.E1.m1.1.1.1.1.1.1.4.2.1.cmml">⋅</mo><mmultiscripts id="S2.E1.m1.1.1.1.1.1.1.4.2.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.cmml"><mtext id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2a.cmml">FFN</mtext><mtext id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3a.cmml">SwiGLU</mtext><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.3.cmml">𝒯</mi><mi id="S2.E1.m1.1.1.1.1.1.1.4.2.3.3" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.3.cmml">j</mi><mrow id="S2.E1.m1.1.1.1.1.1.1.4.2.3a" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.cmml"></mrow></mmultiscripts></mrow><mo lspace="0em" rspace="0em" id="S2.E1.m1.1.1.1.1.1.1.4.1" xref="S2.E1.m1.1.1.1.1.1.1.4.1.cmml">​</mo><mrow id="S2.E1.m1.1.1.1.1.1.1.4.3.2" xref="S2.E1.m1.1.1.1.1.1.1.4.cmml"><mo stretchy="false" id="S2.E1.m1.1.1.1.1.1.1.4.3.2.1" xref="S2.E1.m1.1.1.1.1.1.1.4.cmml">(</mo><mi id="S2.E1.m1.1.1.1.1.1.1.2" xref="S2.E1.m1.1.1.1.1.1.1.2.cmml">x</mi><mo stretchy="false" id="S2.E1.m1.1.1.1.1.1.1.4.3.2.2" xref="S2.E1.m1.1.1.1.1.1.1.4.cmml">)</mo></mrow></mrow></mrow></mtd><mtd class="ltx_align_left" columnalign="left" id="S2.E1.m1.4.4.4c" xref="S2.E1.m1.4.5.3.1.cmml"><mrow id="S2.E1.m1.2.2.2.2.2.1" xref="S2.E1.m1.2.2.2.2.2.1.cmml"><mrow id="S2.E1.m1.2.2.2.2.2.1.2" xref="S2.E1.m1.2.2.2.2.2.1.2.cmml"><mtext id="S2.E1.m1.2.2.2.2.2.1.2.2" xref="S2.E1.m1.2.2.2.2.2.1.2.2a.cmml">if </mtext><mo lspace="0em" rspace="0em" id="S2.E1.m1.2.2.2.2.2.1.2.1" xref="S2.E1.m1.2.2.2.2.2.1.2.1.cmml">​</mo><mi id="S2.E1.m1.2.2.2.2.2.1.2.3" xref="S2.E1.m1.2.2.2.2.2.1.2.3.cmml">x</mi></mrow><mo id="S2.E1.m1.2.2.2.2.2.1.1" xref="S2.E1.m1.2.2.2.2.2.1.1.cmml">∈</mo><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.2.2.2.2.2.1.3" xref="S2.E1.m1.2.2.2.2.2.1.3.cmml">𝒯</mi></mrow></mtd></mtr><mtr id="S2.E1.m1.4.4.4d" xref="S2.E1.m1.4.5.3.1.cmml"><mtd class="ltx_align_left" columnalign="left" id="S2.E1.m1.4.4.4e" xref="S2.E1.m1.4.5.3.1.cmml"><mrow id="S2.E1.m1.3.3.3.3.1.1" xref="S2.E1.m1.3.3.3.3.1.1.cmml"><mstyle displaystyle="false" id="S2.E1.m1.3.3.3.3.1.1.3" xref="S2.E1.m1.3.3.3.3.1.1.3.cmml"><msubsup id="S2.E1.m1.3.3.3.3.1.1.3a" xref="S2.E1.m1.3.3.3.3.1.1.3.cmml"><mo id="S2.E1.m1.3.3.3.3.1.1.3.2.2" xref="S2.E1.m1.3.3.3.3.1.1.3.2.2.cmml">∑</mo><mrow id="S2.E1.m1.3.3.3.3.1.1.3.2.3" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.cmml"><mi id="S2.E1.m1.3.3.3.3.1.1.3.2.3.2" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.2.cmml">j</mi><mo id="S2.E1.m1.3.3.3.3.1.1.3.2.3.1" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.1.cmml">=</mo><mn id="S2.E1.m1.3.3.3.3.1.1.3.2.3.3" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.3.cmml">1</mn></mrow><mi id="S2.E1.m1.3.3.3.3.1.1.3.3" xref="S2.E1.m1.3.3.3.3.1.1.3.3.cmml">n</mi></msubsup></mstyle><mrow id="S2.E1.m1.3.3.3.3.1.1.4" xref="S2.E1.m1.3.3.3.3.1.1.4.cmml"><mrow id="S2.E1.m1.3.3.3.3.1.1.4.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.cmml"><mrow id="S2.E1.m1.3.3.3.3.1.1.4.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.cmml"><msup id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.cmml"><mi id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.2.cmml">g</mi><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.3.cmml">ℐ</mi></msup><mo lspace="0em" rspace="0em" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.1" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.1.cmml">​</mo><msub id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.cmml"><mrow id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.cmml"><mo stretchy="false" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.2.2.1" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.cmml">(</mo><mi id="S2.E1.m1.3.3.3.3.1.1.1" xref="S2.E1.m1.3.3.3.3.1.1.1.cmml">x</mi><mo rspace="0.055em" stretchy="false" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.2.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.cmml">)</mo></mrow><mi id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.3.cmml">j</mi></msub></mrow><mo rspace="0.222em" id="S2.E1.m1.3.3.3.3.1.1.4.2.1" xref="S2.E1.m1.3.3.3.3.1.1.4.2.1.cmml">⋅</mo><mmultiscripts id="S2.E1.m1.3.3.3.3.1.1.4.2.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.cmml"><mtext id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2a.cmml">FFN</mtext><mtext id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3a.cmml">SwiGLU</mtext><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.3.cmml">ℐ</mi><mi id="S2.E1.m1.3.3.3.3.1.1.4.2.3.3" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.3.cmml">j</mi><mrow id="S2.E1.m1.3.3.3.3.1.1.4.2.3a" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.cmml"></mrow></mmultiscripts></mrow><mo lspace="0em" rspace="0em" id="S2.E1.m1.3.3.3.3.1.1.4.1" xref="S2.E1.m1.3.3.3.3.1.1.4.1.cmml">​</mo><mrow id="S2.E1.m1.3.3.3.3.1.1.4.3.2" xref="S2.E1.m1.3.3.3.3.1.1.4.cmml"><mo stretchy="false" id="S2.E1.m1.3.3.3.3.1.1.4.3.2.1" xref="S2.E1.m1.3.3.3.3.1.1.4.cmml">(</mo><mi id="S2.E1.m1.3.3.3.3.1.1.2" xref="S2.E1.m1.3.3.3.3.1.1.2.cmml">x</mi><mo stretchy="false" id="S2.E1.m1.3.3.3.3.1.1.4.3.2.2" xref="S2.E1.m1.3.3.3.3.1.1.4.cmml">)</mo></mrow></mrow></mrow></mtd><mtd class="ltx_align_left" columnalign="left" id="S2.E1.m1.4.4.4f" xref="S2.E1.m1.4.5.3.1.cmml"><mrow id="S2.E1.m1.4.4.4.4.2.1" xref="S2.E1.m1.4.4.4.4.2.1.cmml"><mrow id="S2.E1.m1.4.4.4.4.2.1.2" xref="S2.E1.m1.4.4.4.4.2.1.2.cmml"><mtext id="S2.E1.m1.4.4.4.4.2.1.2.2" xref="S2.E1.m1.4.4.4.4.2.1.2.2a.cmml">if </mtext><mo lspace="0em" rspace="0em" id="S2.E1.m1.4.4.4.4.2.1.2.1" xref="S2.E1.m1.4.4.4.4.2.1.2.1.cmml">​</mo><mi id="S2.E1.m1.4.4.4.4.2.1.2.3" xref="S2.E1.m1.4.4.4.4.2.1.2.3.cmml">x</mi></mrow><mo id="S2.E1.m1.4.4.4.4.2.1.1" xref="S2.E1.m1.4.4.4.4.2.1.1.cmml">∈</mo><mi class="ltx_font_mathcaligraphic" id="S2.E1.m1.4.4.4.4.2.1.3" xref="S2.E1.m1.4.4.4.4.2.1.3.cmml">ℐ</mi></mrow></mtd></mtr></mtable></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.E1.m1.4b"><apply id="S2.E1.m1.4.5.cmml" xref="S2.E1.m1.4.5"><eq id="S2.E1.m1.4.5.1.cmml" xref="S2.E1.m1.4.5.1"></eq><ci id="S2.E1.m1.4.5.2.cmml" xref="S2.E1.m1.4.5.2">𝑦</ci><apply id="S2.E1.m1.4.5.3.1.cmml" xref="S2.E1.m1.4.4"><csymbol cd="latexml" id="S2.E1.m1.4.5.3.1.1.cmml" xref="S2.E1.m1.4.4.5">cases</csymbol><apply id="S2.E1.m1.1.1.1.1.1.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1"><apply id="S2.E1.m1.1.1.1.1.1.1.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.3.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3">superscript</csymbol><apply id="S2.E1.m1.1.1.1.1.1.1.3.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3">subscript</csymbol><sum id="S2.E1.m1.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.2.2"></sum><apply id="S2.E1.m1.1.1.1.1.1.1.3.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3"><eq id="S2.E1.m1.1.1.1.1.1.1.3.2.3.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.1"></eq><ci id="S2.E1.m1.1.1.1.1.1.1.3.2.3.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.2">𝑗</ci><cn type="integer" id="S2.E1.m1.1.1.1.1.1.1.3.2.3.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.2.3.3">1</cn></apply></apply><ci id="S2.E1.m1.1.1.1.1.1.1.3.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.3.3">𝑚</ci></apply><apply id="S2.E1.m1.1.1.1.1.1.1.4.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4"><times id="S2.E1.m1.1.1.1.1.1.1.4.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.1"></times><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2"><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.1">⋅</ci><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2"><times id="S2.E1.m1.1.1.1.1.1.1.4.2.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.1"></times><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2">superscript</csymbol><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.2">𝑔</ci><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.2.3">𝒯</ci></apply><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3">subscript</csymbol><ci id="S2.E1.m1.1.1.1.1.1.1.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.1">𝑥</ci><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.2.3.3">𝑗</ci></apply></apply><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.4.2.3.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3">subscript</csymbol><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3">subscript</csymbol><apply id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.1.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3">superscript</csymbol><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2a.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2"><mtext id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.2">FFN</mtext></ci><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.2.3">𝒯</ci></apply><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3a.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3"><mtext mathsize="70%" id="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.2.3">SwiGLU</mtext></ci></apply><ci id="S2.E1.m1.1.1.1.1.1.1.4.2.3.3.cmml" xref="S2.E1.m1.1.1.1.1.1.1.4.2.3.3">𝑗</ci></apply></apply><ci id="S2.E1.m1.1.1.1.1.1.1.2.cmml" xref="S2.E1.m1.1.1.1.1.1.1.2">𝑥</ci></apply></apply><apply id="S2.E1.m1.2.2.2.2.2.1.cmml" xref="S2.E1.m1.2.2.2.2.2.1"><in id="S2.E1.m1.2.2.2.2.2.1.1.cmml" xref="S2.E1.m1.2.2.2.2.2.1.1"></in><apply id="S2.E1.m1.2.2.2.2.2.1.2.cmml" xref="S2.E1.m1.2.2.2.2.2.1.2"><times id="S2.E1.m1.2.2.2.2.2.1.2.1.cmml" xref="S2.E1.m1.2.2.2.2.2.1.2.1"></times><ci id="S2.E1.m1.2.2.2.2.2.1.2.2a.cmml" xref="S2.E1.m1.2.2.2.2.2.1.2.2"><mtext id="S2.E1.m1.2.2.2.2.2.1.2.2.cmml" xref="S2.E1.m1.2.2.2.2.2.1.2.2">if </mtext></ci><ci id="S2.E1.m1.2.2.2.2.2.1.2.3.cmml" xref="S2.E1.m1.2.2.2.2.2.1.2.3">𝑥</ci></apply><ci id="S2.E1.m1.2.2.2.2.2.1.3.cmml" xref="S2.E1.m1.2.2.2.2.2.1.3">𝒯</ci></apply><apply id="S2.E1.m1.3.3.3.3.1.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1"><apply id="S2.E1.m1.3.3.3.3.1.1.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.3.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3">superscript</csymbol><apply id="S2.E1.m1.3.3.3.3.1.1.3.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.3.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3">subscript</csymbol><sum id="S2.E1.m1.3.3.3.3.1.1.3.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.2.2"></sum><apply id="S2.E1.m1.3.3.3.3.1.1.3.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3"><eq id="S2.E1.m1.3.3.3.3.1.1.3.2.3.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.1"></eq><ci id="S2.E1.m1.3.3.3.3.1.1.3.2.3.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.2">𝑗</ci><cn type="integer" id="S2.E1.m1.3.3.3.3.1.1.3.2.3.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.2.3.3">1</cn></apply></apply><ci id="S2.E1.m1.3.3.3.3.1.1.3.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.3.3">𝑛</ci></apply><apply id="S2.E1.m1.3.3.3.3.1.1.4.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4"><times id="S2.E1.m1.3.3.3.3.1.1.4.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.1"></times><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2"><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.1">⋅</ci><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2"><times id="S2.E1.m1.3.3.3.3.1.1.4.2.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.1"></times><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2">superscript</csymbol><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.2">𝑔</ci><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.2.3">ℐ</ci></apply><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3">subscript</csymbol><ci id="S2.E1.m1.3.3.3.3.1.1.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.1">𝑥</ci><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.2.3.3">𝑗</ci></apply></apply><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.4.2.3.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3">subscript</csymbol><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3">subscript</csymbol><apply id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3"><csymbol cd="ambiguous" id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.1.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3">superscript</csymbol><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2a.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2"><mtext id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.2">FFN</mtext></ci><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.2.3">ℐ</ci></apply><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3a.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3"><mtext mathsize="70%" id="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.2.3">SwiGLU</mtext></ci></apply><ci id="S2.E1.m1.3.3.3.3.1.1.4.2.3.3.cmml" xref="S2.E1.m1.3.3.3.3.1.1.4.2.3.3">𝑗</ci></apply></apply><ci id="S2.E1.m1.3.3.3.3.1.1.2.cmml" xref="S2.E1.m1.3.3.3.3.1.1.2">𝑥</ci></apply></apply><apply id="S2.E1.m1.4.4.4.4.2.1.cmml" xref="S2.E1.m1.4.4.4.4.2.1"><in id="S2.E1.m1.4.4.4.4.2.1.1.cmml" xref="S2.E1.m1.4.4.4.4.2.1.1"></in><apply id="S2.E1.m1.4.4.4.4.2.1.2.cmml" xref="S2.E1.m1.4.4.4.4.2.1.2"><times id="S2.E1.m1.4.4.4.4.2.1.2.1.cmml" xref="S2.E1.m1.4.4.4.4.2.1.2.1"></times><ci id="S2.E1.m1.4.4.4.4.2.1.2.2a.cmml" xref="S2.E1.m1.4.4.4.4.2.1.2.2"><mtext id="S2.E1.m1.4.4.4.4.2.1.2.2.cmml" xref="S2.E1.m1.4.4.4.4.2.1.2.2">if </mtext></ci><ci id="S2.E1.m1.4.4.4.4.2.1.2.3.cmml" xref="S2.E1.m1.4.4.4.4.2.1.2.3">𝑥</ci></apply><ci id="S2.E1.m1.4.4.4.4.2.1.3.cmml" xref="S2.E1.m1.4.4.4.4.2.1.3">ℐ</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E1.m1.4c">y=\begin{cases}\sum_{j=1}^{m}g^{\mathcal{T}}(x)_{j}\cdot{\text{FFN}^{\mathcal{T}}_{\text{SwiGLU}}}_{j}(x)&amp;\text{if }x\in\mathcal{T}\\
\sum_{j=1}^{n}g^{\mathcal{I}}(x)_{j}\cdot{\text{FFN}^{\mathcal{I}}_{\text{SwiGLU}}}_{j}(x)&amp;\text{if }x\in\mathcal{I}\end{cases}</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>
<div id="S2.SS2.SSS0.Px3.p5" class="ltx_para">
<table id="S2.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="S2.Ex1.m1.5" class="ltx_Math" alttext="g^{\mathcal{M}}(x)_{j}=\begin{cases}\text{Sigmoid}(x\cdot W^{\mathcal{M}}_{g})_{j}&amp;\text{if }x\text{ is selected by }E_{j}\\
0&amp;\text{otherwise}\end{cases}" display="block"><semantics id="S2.Ex1.m1.5a"><mrow id="S2.Ex1.m1.5.6" xref="S2.Ex1.m1.5.6.cmml"><mrow id="S2.Ex1.m1.5.6.2" xref="S2.Ex1.m1.5.6.2.cmml"><msup id="S2.Ex1.m1.5.6.2.2" xref="S2.Ex1.m1.5.6.2.2.cmml"><mi id="S2.Ex1.m1.5.6.2.2.2" xref="S2.Ex1.m1.5.6.2.2.2.cmml">g</mi><mi class="ltx_font_mathcaligraphic" id="S2.Ex1.m1.5.6.2.2.3" xref="S2.Ex1.m1.5.6.2.2.3.cmml">ℳ</mi></msup><mo lspace="0em" rspace="0em" id="S2.Ex1.m1.5.6.2.1" xref="S2.Ex1.m1.5.6.2.1.cmml">​</mo><msub id="S2.Ex1.m1.5.6.2.3" xref="S2.Ex1.m1.5.6.2.3.cmml"><mrow id="S2.Ex1.m1.5.6.2.3.2.2" xref="S2.Ex1.m1.5.6.2.3.cmml"><mo stretchy="false" id="S2.Ex1.m1.5.6.2.3.2.2.1" xref="S2.Ex1.m1.5.6.2.3.cmml">(</mo><mi id="S2.Ex1.m1.5.5" xref="S2.Ex1.m1.5.5.cmml">x</mi><mo stretchy="false" id="S2.Ex1.m1.5.6.2.3.2.2.2" xref="S2.Ex1.m1.5.6.2.3.cmml">)</mo></mrow><mi id="S2.Ex1.m1.5.6.2.3.3" xref="S2.Ex1.m1.5.6.2.3.3.cmml">j</mi></msub></mrow><mo id="S2.Ex1.m1.5.6.1" xref="S2.Ex1.m1.5.6.1.cmml">=</mo><mrow id="S2.Ex1.m1.4.4" xref="S2.Ex1.m1.5.6.3.1.cmml"><mo id="S2.Ex1.m1.4.4.5" xref="S2.Ex1.m1.5.6.3.1.1.cmml">{</mo><mtable columnspacing="5pt" displaystyle="true" rowspacing="0pt" id="S2.Ex1.m1.4.4.4" xref="S2.Ex1.m1.5.6.3.1.cmml"><mtr id="S2.Ex1.m1.4.4.4a" xref="S2.Ex1.m1.5.6.3.1.cmml"><mtd class="ltx_align_left" columnalign="left" id="S2.Ex1.m1.4.4.4b" xref="S2.Ex1.m1.5.6.3.1.cmml"><mrow id="S2.Ex1.m1.1.1.1.1.1.1" xref="S2.Ex1.m1.1.1.1.1.1.1.cmml"><mtext id="S2.Ex1.m1.1.1.1.1.1.1.3" xref="S2.Ex1.m1.1.1.1.1.1.1.3a.cmml">Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.Ex1.m1.1.1.1.1.1.1.2" xref="S2.Ex1.m1.1.1.1.1.1.1.2.cmml">​</mo><msub id="S2.Ex1.m1.1.1.1.1.1.1.1" xref="S2.Ex1.m1.1.1.1.1.1.1.1.cmml"><mrow id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.2" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.cmml"><mi id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.2" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.2.cmml">x</mi><mo lspace="0.222em" rspace="0.222em" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.1" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.1.cmml">⋅</mo><msubsup id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.cmml"><mi id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.2" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml">W</mi><mi id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.3" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.3.cmml">g</mi><mi class="ltx_font_mathcaligraphic" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.3" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.3.cmml">ℳ</mi></msubsup></mrow><mo stretchy="false" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.3" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.cmml">)</mo></mrow><mi id="S2.Ex1.m1.1.1.1.1.1.1.1.3" xref="S2.Ex1.m1.1.1.1.1.1.1.1.3.cmml">j</mi></msub></mrow></mtd><mtd class="ltx_align_left" columnalign="left" id="S2.Ex1.m1.4.4.4c" xref="S2.Ex1.m1.5.6.3.1.cmml"><mrow id="S2.Ex1.m1.2.2.2.2.2.1" xref="S2.Ex1.m1.2.2.2.2.2.1.cmml"><mtext id="S2.Ex1.m1.2.2.2.2.2.1.2" xref="S2.Ex1.m1.2.2.2.2.2.1.2a.cmml">if </mtext><mo lspace="0em" rspace="0em" id="S2.Ex1.m1.2.2.2.2.2.1.1" xref="S2.Ex1.m1.2.2.2.2.2.1.1.cmml">​</mo><mi id="S2.Ex1.m1.2.2.2.2.2.1.3" xref="S2.Ex1.m1.2.2.2.2.2.1.3.cmml">x</mi><mo lspace="0em" rspace="0em" id="S2.Ex1.m1.2.2.2.2.2.1.1a" xref="S2.Ex1.m1.2.2.2.2.2.1.1.cmml">​</mo><mtext id="S2.Ex1.m1.2.2.2.2.2.1.4" xref="S2.Ex1.m1.2.2.2.2.2.1.4a.cmml"> is selected by </mtext><mo lspace="0em" rspace="0em" id="S2.Ex1.m1.2.2.2.2.2.1.1b" xref="S2.Ex1.m1.2.2.2.2.2.1.1.cmml">​</mo><msub id="S2.Ex1.m1.2.2.2.2.2.1.5" xref="S2.Ex1.m1.2.2.2.2.2.1.5.cmml"><mi id="S2.Ex1.m1.2.2.2.2.2.1.5.2" xref="S2.Ex1.m1.2.2.2.2.2.1.5.2.cmml">E</mi><mi id="S2.Ex1.m1.2.2.2.2.2.1.5.3" xref="S2.Ex1.m1.2.2.2.2.2.1.5.3.cmml">j</mi></msub></mrow></mtd></mtr><mtr id="S2.Ex1.m1.4.4.4d" xref="S2.Ex1.m1.5.6.3.1.cmml"><mtd class="ltx_align_left" columnalign="left" id="S2.Ex1.m1.4.4.4e" xref="S2.Ex1.m1.5.6.3.1.cmml"><mn id="S2.Ex1.m1.3.3.3.3.1.1" xref="S2.Ex1.m1.3.3.3.3.1.1.cmml">0</mn></mtd><mtd class="ltx_align_left" columnalign="left" id="S2.Ex1.m1.4.4.4f" xref="S2.Ex1.m1.5.6.3.1.cmml"><mtext id="S2.Ex1.m1.4.4.4.4.2.1" xref="S2.Ex1.m1.4.4.4.4.2.1a.cmml">otherwise</mtext></mtd></mtr></mtable></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.Ex1.m1.5b"><apply id="S2.Ex1.m1.5.6.cmml" xref="S2.Ex1.m1.5.6"><eq id="S2.Ex1.m1.5.6.1.cmml" xref="S2.Ex1.m1.5.6.1"></eq><apply id="S2.Ex1.m1.5.6.2.cmml" xref="S2.Ex1.m1.5.6.2"><times id="S2.Ex1.m1.5.6.2.1.cmml" xref="S2.Ex1.m1.5.6.2.1"></times><apply id="S2.Ex1.m1.5.6.2.2.cmml" xref="S2.Ex1.m1.5.6.2.2"><csymbol cd="ambiguous" id="S2.Ex1.m1.5.6.2.2.1.cmml" xref="S2.Ex1.m1.5.6.2.2">superscript</csymbol><ci id="S2.Ex1.m1.5.6.2.2.2.cmml" xref="S2.Ex1.m1.5.6.2.2.2">𝑔</ci><ci id="S2.Ex1.m1.5.6.2.2.3.cmml" xref="S2.Ex1.m1.5.6.2.2.3">ℳ</ci></apply><apply id="S2.Ex1.m1.5.6.2.3.cmml" xref="S2.Ex1.m1.5.6.2.3"><csymbol cd="ambiguous" id="S2.Ex1.m1.5.6.2.3.1.cmml" xref="S2.Ex1.m1.5.6.2.3">subscript</csymbol><ci id="S2.Ex1.m1.5.5.cmml" xref="S2.Ex1.m1.5.5">𝑥</ci><ci id="S2.Ex1.m1.5.6.2.3.3.cmml" xref="S2.Ex1.m1.5.6.2.3.3">𝑗</ci></apply></apply><apply id="S2.Ex1.m1.5.6.3.1.cmml" xref="S2.Ex1.m1.4.4"><csymbol cd="latexml" id="S2.Ex1.m1.5.6.3.1.1.cmml" xref="S2.Ex1.m1.4.4.5">cases</csymbol><apply id="S2.Ex1.m1.1.1.1.1.1.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1"><times id="S2.Ex1.m1.1.1.1.1.1.1.2.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.2"></times><ci id="S2.Ex1.m1.1.1.1.1.1.1.3a.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.3"><mtext id="S2.Ex1.m1.1.1.1.1.1.1.3.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.3">Sigmoid</mtext></ci><apply id="S2.Ex1.m1.1.1.1.1.1.1.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1"><csymbol cd="ambiguous" id="S2.Ex1.m1.1.1.1.1.1.1.1.2.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1">subscript</csymbol><apply id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1"><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.1">⋅</ci><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.2.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.2">𝑥</ci><apply id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3">subscript</csymbol><apply id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.1.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3">superscript</csymbol><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.2">𝑊</ci><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.3.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.2.3">ℳ</ci></apply><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.3.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.1.1.1.3.3">𝑔</ci></apply></apply><ci id="S2.Ex1.m1.1.1.1.1.1.1.1.3.cmml" xref="S2.Ex1.m1.1.1.1.1.1.1.1.3">𝑗</ci></apply></apply><apply id="S2.Ex1.m1.2.2.2.2.2.1.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1"><times id="S2.Ex1.m1.2.2.2.2.2.1.1.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.1"></times><ci id="S2.Ex1.m1.2.2.2.2.2.1.2a.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.2"><mtext id="S2.Ex1.m1.2.2.2.2.2.1.2.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.2">if </mtext></ci><ci id="S2.Ex1.m1.2.2.2.2.2.1.3.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.3">𝑥</ci><ci id="S2.Ex1.m1.2.2.2.2.2.1.4a.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.4"><mtext id="S2.Ex1.m1.2.2.2.2.2.1.4.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.4"> is selected by </mtext></ci><apply id="S2.Ex1.m1.2.2.2.2.2.1.5.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.5"><csymbol cd="ambiguous" id="S2.Ex1.m1.2.2.2.2.2.1.5.1.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.5">subscript</csymbol><ci id="S2.Ex1.m1.2.2.2.2.2.1.5.2.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.5.2">𝐸</ci><ci id="S2.Ex1.m1.2.2.2.2.2.1.5.3.cmml" xref="S2.Ex1.m1.2.2.2.2.2.1.5.3">𝑗</ci></apply></apply><cn type="integer" id="S2.Ex1.m1.3.3.3.3.1.1.cmml" xref="S2.Ex1.m1.3.3.3.3.1.1">0</cn><ci id="S2.Ex1.m1.4.4.4.4.2.1a.cmml" xref="S2.Ex1.m1.4.4.4.4.2.1"><mtext id="S2.Ex1.m1.4.4.4.4.2.1.cmml" xref="S2.Ex1.m1.4.4.4.4.2.1">otherwise</mtext></ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.Ex1.m1.5c">g^{\mathcal{M}}(x)_{j}=\begin{cases}\text{Sigmoid}(x\cdot W^{\mathcal{M}}_{g})_{j}&amp;\text{if }x\text{ is selected by }E_{j}\\
0&amp;\text{otherwise}\end{cases}</annotation></semantics></math></td>
<td class="ltx_eqn_cell ltx_eqn_center_padright"></td>
</tr></tbody>
</table>
</div>
<div id="S2.SS2.SSS0.Px3.p6" class="ltx_para">
<p id="S2.SS2.SSS0.Px3.p6.1" class="ltx_p">We further apply residual connection and the Swin Transformer normalization <cite class="ltx_cite ltx_citemacro_citep">(Liu et al., <a href="#bib.bib20" title="" class="ltx_ref">2022a</a>)</cite> post MoMa computation. Our experiments, detailed in later sections, demonstrate that MoMa significantly improves efficiency and performance compared to dense baselines and standard MoE architectures.</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>Mixture-of-Depths</h3>

<figure id="S2.F2" class="ltx_figure"><img src="/html/2407.21770/assets/figures/moma_arch_mod.png" id="S2.F2.g1" class="ltx_graphics ltx_centering ltx_img_portrait" width="168" height="274" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S2.F2.2.1.1" class="ltx_text" style="font-size:90%;">Figure 2</span>: </span><span id="S2.F2.3.2" class="ltx_text" style="font-size:90%;">Architecture of transformer layer consisting of MoMa combined with mixture-of-depths (MoD).</span></figcaption>
</figure>
<div id="S2.SS3.p1" class="ltx_para">
<p id="S2.SS3.p1.1" class="ltx_p">We further investigate introducing sparsity in the depth dimension. Prior work explores sparisty in depth through either stochastic layer drop <cite class="ltx_cite ltx_citemacro_citep">(Elhoushi et al., <a href="#bib.bib21" title="" class="ltx_ref">2024</a>)</cite> or through learnable routers <cite class="ltx_cite ltx_citemacro_citep">(Raposo et al., <a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>. We focus on learnable routers, and incorporate the recently proposed mixture-of-depths (MoD) technique <cite class="ltx_cite ltx_citemacro_citep">(Raposo et al., <a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>.
Specifically, in each MoD layer, we integrate MoD prior to any mixture-of-experts (MoE) routing, ensuring it is applied to the full batch before modality split (Figure <a href="#S2.F2" title="Figure 2 ‣ 2.3 Mixture-of-Depths ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2</span></a>).<span id="footnote4" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">4</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">4</sup><span class="ltx_tag ltx_tag_note">4</span>This corresponding to the staged MoDE implementation presented in <cite class="ltx_cite ltx_citemacro_cite">Raposo et al. (<a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>. We chose this implementation because it also introduces sparsity in attention, whereas the integrated MoDE implementation proposed in the same paper always computes full attention.</span></span></span></p>
</div>
<div id="S2.SS3.p2" class="ltx_para">
<p id="S2.SS3.p2.1" class="ltx_p">Following <cite class="ltx_cite ltx_citemacro_cite">Raposo et al. (<a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>, for each MoD layer, we use a projection matrix <math id="S2.SS3.p2.1.m1.1" class="ltx_Math" alttext="\hat{W_{g}}\in\mathbb{R}^{d\times 1}" display="inline"><semantics id="S2.SS3.p2.1.m1.1a"><mrow id="S2.SS3.p2.1.m1.1.1" xref="S2.SS3.p2.1.m1.1.1.cmml"><mover accent="true" id="S2.SS3.p2.1.m1.1.1.2" xref="S2.SS3.p2.1.m1.1.1.2.cmml"><msub id="S2.SS3.p2.1.m1.1.1.2.2" xref="S2.SS3.p2.1.m1.1.1.2.2.cmml"><mi id="S2.SS3.p2.1.m1.1.1.2.2.2" xref="S2.SS3.p2.1.m1.1.1.2.2.2.cmml">W</mi><mi id="S2.SS3.p2.1.m1.1.1.2.2.3" xref="S2.SS3.p2.1.m1.1.1.2.2.3.cmml">g</mi></msub><mo id="S2.SS3.p2.1.m1.1.1.2.1" xref="S2.SS3.p2.1.m1.1.1.2.1.cmml">^</mo></mover><mo id="S2.SS3.p2.1.m1.1.1.1" xref="S2.SS3.p2.1.m1.1.1.1.cmml">∈</mo><msup id="S2.SS3.p2.1.m1.1.1.3" xref="S2.SS3.p2.1.m1.1.1.3.cmml"><mi id="S2.SS3.p2.1.m1.1.1.3.2" xref="S2.SS3.p2.1.m1.1.1.3.2.cmml">ℝ</mi><mrow id="S2.SS3.p2.1.m1.1.1.3.3" xref="S2.SS3.p2.1.m1.1.1.3.3.cmml"><mi id="S2.SS3.p2.1.m1.1.1.3.3.2" xref="S2.SS3.p2.1.m1.1.1.3.3.2.cmml">d</mi><mo lspace="0.222em" rspace="0.222em" id="S2.SS3.p2.1.m1.1.1.3.3.1" xref="S2.SS3.p2.1.m1.1.1.3.3.1.cmml">×</mo><mn id="S2.SS3.p2.1.m1.1.1.3.3.3" xref="S2.SS3.p2.1.m1.1.1.3.3.3.cmml">1</mn></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.1.m1.1b"><apply id="S2.SS3.p2.1.m1.1.1.cmml" xref="S2.SS3.p2.1.m1.1.1"><in id="S2.SS3.p2.1.m1.1.1.1.cmml" xref="S2.SS3.p2.1.m1.1.1.1"></in><apply id="S2.SS3.p2.1.m1.1.1.2.cmml" xref="S2.SS3.p2.1.m1.1.1.2"><ci id="S2.SS3.p2.1.m1.1.1.2.1.cmml" xref="S2.SS3.p2.1.m1.1.1.2.1">^</ci><apply id="S2.SS3.p2.1.m1.1.1.2.2.cmml" xref="S2.SS3.p2.1.m1.1.1.2.2"><csymbol cd="ambiguous" id="S2.SS3.p2.1.m1.1.1.2.2.1.cmml" xref="S2.SS3.p2.1.m1.1.1.2.2">subscript</csymbol><ci id="S2.SS3.p2.1.m1.1.1.2.2.2.cmml" xref="S2.SS3.p2.1.m1.1.1.2.2.2">𝑊</ci><ci id="S2.SS3.p2.1.m1.1.1.2.2.3.cmml" xref="S2.SS3.p2.1.m1.1.1.2.2.3">𝑔</ci></apply></apply><apply id="S2.SS3.p2.1.m1.1.1.3.cmml" xref="S2.SS3.p2.1.m1.1.1.3"><csymbol cd="ambiguous" id="S2.SS3.p2.1.m1.1.1.3.1.cmml" xref="S2.SS3.p2.1.m1.1.1.3">superscript</csymbol><ci id="S2.SS3.p2.1.m1.1.1.3.2.cmml" xref="S2.SS3.p2.1.m1.1.1.3.2">ℝ</ci><apply id="S2.SS3.p2.1.m1.1.1.3.3.cmml" xref="S2.SS3.p2.1.m1.1.1.3.3"><times id="S2.SS3.p2.1.m1.1.1.3.3.1.cmml" xref="S2.SS3.p2.1.m1.1.1.3.3.1"></times><ci id="S2.SS3.p2.1.m1.1.1.3.3.2.cmml" xref="S2.SS3.p2.1.m1.1.1.3.3.2">𝑑</ci><cn type="integer" id="S2.SS3.p2.1.m1.1.1.3.3.3.cmml" xref="S2.SS3.p2.1.m1.1.1.3.3.3">1</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.1.m1.1c">\hat{W_{g}}\in\mathbb{R}^{d\times 1}</annotation></semantics></math> to compute the token-to-layer affinity score, followed by a Sigmoid non-linearity</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.2" class="ltx_Math" alttext="\hat{g}(x)=\text{Sigmoid}(x\cdot\hat{W_{g}})." display="block"><semantics id="S2.E2.m1.2a"><mrow id="S2.E2.m1.2.2.1" xref="S2.E2.m1.2.2.1.1.cmml"><mrow id="S2.E2.m1.2.2.1.1" xref="S2.E2.m1.2.2.1.1.cmml"><mrow id="S2.E2.m1.2.2.1.1.3" xref="S2.E2.m1.2.2.1.1.3.cmml"><mover accent="true" id="S2.E2.m1.2.2.1.1.3.2" xref="S2.E2.m1.2.2.1.1.3.2.cmml"><mi id="S2.E2.m1.2.2.1.1.3.2.2" xref="S2.E2.m1.2.2.1.1.3.2.2.cmml">g</mi><mo id="S2.E2.m1.2.2.1.1.3.2.1" xref="S2.E2.m1.2.2.1.1.3.2.1.cmml">^</mo></mover><mo lspace="0em" rspace="0em" id="S2.E2.m1.2.2.1.1.3.1" xref="S2.E2.m1.2.2.1.1.3.1.cmml">​</mo><mrow id="S2.E2.m1.2.2.1.1.3.3.2" xref="S2.E2.m1.2.2.1.1.3.cmml"><mo stretchy="false" id="S2.E2.m1.2.2.1.1.3.3.2.1" xref="S2.E2.m1.2.2.1.1.3.cmml">(</mo><mi id="S2.E2.m1.1.1" xref="S2.E2.m1.1.1.cmml">x</mi><mo stretchy="false" id="S2.E2.m1.2.2.1.1.3.3.2.2" xref="S2.E2.m1.2.2.1.1.3.cmml">)</mo></mrow></mrow><mo id="S2.E2.m1.2.2.1.1.2" xref="S2.E2.m1.2.2.1.1.2.cmml">=</mo><mrow id="S2.E2.m1.2.2.1.1.1" xref="S2.E2.m1.2.2.1.1.1.cmml"><mtext id="S2.E2.m1.2.2.1.1.1.3" xref="S2.E2.m1.2.2.1.1.1.3a.cmml">Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.E2.m1.2.2.1.1.1.2" xref="S2.E2.m1.2.2.1.1.1.2.cmml">​</mo><mrow id="S2.E2.m1.2.2.1.1.1.1.1" xref="S2.E2.m1.2.2.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E2.m1.2.2.1.1.1.1.1.2" xref="S2.E2.m1.2.2.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E2.m1.2.2.1.1.1.1.1.1" xref="S2.E2.m1.2.2.1.1.1.1.1.1.cmml"><mi id="S2.E2.m1.2.2.1.1.1.1.1.1.2" xref="S2.E2.m1.2.2.1.1.1.1.1.1.2.cmml">x</mi><mo lspace="0.222em" rspace="0.222em" id="S2.E2.m1.2.2.1.1.1.1.1.1.1" xref="S2.E2.m1.2.2.1.1.1.1.1.1.1.cmml">⋅</mo><mover accent="true" id="S2.E2.m1.2.2.1.1.1.1.1.1.3" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.cmml"><msub id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.cmml"><mi id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.2" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.2.cmml">W</mi><mi id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.3" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.3.cmml">g</mi></msub><mo id="S2.E2.m1.2.2.1.1.1.1.1.1.3.1" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.1.cmml">^</mo></mover></mrow><mo stretchy="false" id="S2.E2.m1.2.2.1.1.1.1.1.3" xref="S2.E2.m1.2.2.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow><mo lspace="0em" id="S2.E2.m1.2.2.1.2" xref="S2.E2.m1.2.2.1.1.cmml">.</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E2.m1.2b"><apply id="S2.E2.m1.2.2.1.1.cmml" xref="S2.E2.m1.2.2.1"><eq id="S2.E2.m1.2.2.1.1.2.cmml" xref="S2.E2.m1.2.2.1.1.2"></eq><apply id="S2.E2.m1.2.2.1.1.3.cmml" xref="S2.E2.m1.2.2.1.1.3"><times id="S2.E2.m1.2.2.1.1.3.1.cmml" xref="S2.E2.m1.2.2.1.1.3.1"></times><apply id="S2.E2.m1.2.2.1.1.3.2.cmml" xref="S2.E2.m1.2.2.1.1.3.2"><ci id="S2.E2.m1.2.2.1.1.3.2.1.cmml" xref="S2.E2.m1.2.2.1.1.3.2.1">^</ci><ci id="S2.E2.m1.2.2.1.1.3.2.2.cmml" xref="S2.E2.m1.2.2.1.1.3.2.2">𝑔</ci></apply><ci id="S2.E2.m1.1.1.cmml" xref="S2.E2.m1.1.1">𝑥</ci></apply><apply id="S2.E2.m1.2.2.1.1.1.cmml" xref="S2.E2.m1.2.2.1.1.1"><times id="S2.E2.m1.2.2.1.1.1.2.cmml" xref="S2.E2.m1.2.2.1.1.1.2"></times><ci id="S2.E2.m1.2.2.1.1.1.3a.cmml" xref="S2.E2.m1.2.2.1.1.1.3"><mtext id="S2.E2.m1.2.2.1.1.1.3.cmml" xref="S2.E2.m1.2.2.1.1.1.3">Sigmoid</mtext></ci><apply id="S2.E2.m1.2.2.1.1.1.1.1.1.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1"><ci id="S2.E2.m1.2.2.1.1.1.1.1.1.1.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.1">⋅</ci><ci id="S2.E2.m1.2.2.1.1.1.1.1.1.2.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.2">𝑥</ci><apply id="S2.E2.m1.2.2.1.1.1.1.1.1.3.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3"><ci id="S2.E2.m1.2.2.1.1.1.1.1.1.3.1.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.1">^</ci><apply id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2"><csymbol cd="ambiguous" id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2">subscript</csymbol><ci id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.2">𝑊</ci><ci id="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.3.cmml" xref="S2.E2.m1.2.2.1.1.1.1.1.1.3.2.3">𝑔</ci></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E2.m1.2c">\hat{g}(x)=\text{Sigmoid}(x\cdot\hat{W_{g}}).</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.SS3.p2.5" class="ltx_p">Analogous to expert-choice routing in MoE, we set a fixed capacity of <math id="S2.SS3.p2.2.m1.1" class="ltx_Math" alttext="k_{d}" display="inline"><semantics id="S2.SS3.p2.2.m1.1a"><msub id="S2.SS3.p2.2.m1.1.1" xref="S2.SS3.p2.2.m1.1.1.cmml"><mi id="S2.SS3.p2.2.m1.1.1.2" xref="S2.SS3.p2.2.m1.1.1.2.cmml">k</mi><mi id="S2.SS3.p2.2.m1.1.1.3" xref="S2.SS3.p2.2.m1.1.1.3.cmml">d</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.2.m1.1b"><apply id="S2.SS3.p2.2.m1.1.1.cmml" xref="S2.SS3.p2.2.m1.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.2.m1.1.1.1.cmml" xref="S2.SS3.p2.2.m1.1.1">subscript</csymbol><ci id="S2.SS3.p2.2.m1.1.1.2.cmml" xref="S2.SS3.p2.2.m1.1.1.2">𝑘</ci><ci id="S2.SS3.p2.2.m1.1.1.3.cmml" xref="S2.SS3.p2.2.m1.1.1.3">𝑑</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.2.m1.1c">k_{d}</annotation></semantics></math> tokens, selected from the top-scoring tokens in a batch. We set <math id="S2.SS3.p2.3.m2.1" class="ltx_Math" alttext="k_{d}=b\cdot c_{d}" display="inline"><semantics id="S2.SS3.p2.3.m2.1a"><mrow id="S2.SS3.p2.3.m2.1.1" xref="S2.SS3.p2.3.m2.1.1.cmml"><msub id="S2.SS3.p2.3.m2.1.1.2" xref="S2.SS3.p2.3.m2.1.1.2.cmml"><mi id="S2.SS3.p2.3.m2.1.1.2.2" xref="S2.SS3.p2.3.m2.1.1.2.2.cmml">k</mi><mi id="S2.SS3.p2.3.m2.1.1.2.3" xref="S2.SS3.p2.3.m2.1.1.2.3.cmml">d</mi></msub><mo id="S2.SS3.p2.3.m2.1.1.1" xref="S2.SS3.p2.3.m2.1.1.1.cmml">=</mo><mrow id="S2.SS3.p2.3.m2.1.1.3" xref="S2.SS3.p2.3.m2.1.1.3.cmml"><mi id="S2.SS3.p2.3.m2.1.1.3.2" xref="S2.SS3.p2.3.m2.1.1.3.2.cmml">b</mi><mo lspace="0.222em" rspace="0.222em" id="S2.SS3.p2.3.m2.1.1.3.1" xref="S2.SS3.p2.3.m2.1.1.3.1.cmml">⋅</mo><msub id="S2.SS3.p2.3.m2.1.1.3.3" xref="S2.SS3.p2.3.m2.1.1.3.3.cmml"><mi id="S2.SS3.p2.3.m2.1.1.3.3.2" xref="S2.SS3.p2.3.m2.1.1.3.3.2.cmml">c</mi><mi id="S2.SS3.p2.3.m2.1.1.3.3.3" xref="S2.SS3.p2.3.m2.1.1.3.3.3.cmml">d</mi></msub></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.3.m2.1b"><apply id="S2.SS3.p2.3.m2.1.1.cmml" xref="S2.SS3.p2.3.m2.1.1"><eq id="S2.SS3.p2.3.m2.1.1.1.cmml" xref="S2.SS3.p2.3.m2.1.1.1"></eq><apply id="S2.SS3.p2.3.m2.1.1.2.cmml" xref="S2.SS3.p2.3.m2.1.1.2"><csymbol cd="ambiguous" id="S2.SS3.p2.3.m2.1.1.2.1.cmml" xref="S2.SS3.p2.3.m2.1.1.2">subscript</csymbol><ci id="S2.SS3.p2.3.m2.1.1.2.2.cmml" xref="S2.SS3.p2.3.m2.1.1.2.2">𝑘</ci><ci id="S2.SS3.p2.3.m2.1.1.2.3.cmml" xref="S2.SS3.p2.3.m2.1.1.2.3">𝑑</ci></apply><apply id="S2.SS3.p2.3.m2.1.1.3.cmml" xref="S2.SS3.p2.3.m2.1.1.3"><ci id="S2.SS3.p2.3.m2.1.1.3.1.cmml" xref="S2.SS3.p2.3.m2.1.1.3.1">⋅</ci><ci id="S2.SS3.p2.3.m2.1.1.3.2.cmml" xref="S2.SS3.p2.3.m2.1.1.3.2">𝑏</ci><apply id="S2.SS3.p2.3.m2.1.1.3.3.cmml" xref="S2.SS3.p2.3.m2.1.1.3.3"><csymbol cd="ambiguous" id="S2.SS3.p2.3.m2.1.1.3.3.1.cmml" xref="S2.SS3.p2.3.m2.1.1.3.3">subscript</csymbol><ci id="S2.SS3.p2.3.m2.1.1.3.3.2.cmml" xref="S2.SS3.p2.3.m2.1.1.3.3.2">𝑐</ci><ci id="S2.SS3.p2.3.m2.1.1.3.3.3.cmml" xref="S2.SS3.p2.3.m2.1.1.3.3.3">𝑑</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.3.m2.1c">k_{d}=b\cdot c_{d}</annotation></semantics></math>, where <math id="S2.SS3.p2.4.m3.2" class="ltx_Math" alttext="c_{d}\in(0,1]" display="inline"><semantics id="S2.SS3.p2.4.m3.2a"><mrow id="S2.SS3.p2.4.m3.2.3" xref="S2.SS3.p2.4.m3.2.3.cmml"><msub id="S2.SS3.p2.4.m3.2.3.2" xref="S2.SS3.p2.4.m3.2.3.2.cmml"><mi id="S2.SS3.p2.4.m3.2.3.2.2" xref="S2.SS3.p2.4.m3.2.3.2.2.cmml">c</mi><mi id="S2.SS3.p2.4.m3.2.3.2.3" xref="S2.SS3.p2.4.m3.2.3.2.3.cmml">d</mi></msub><mo id="S2.SS3.p2.4.m3.2.3.1" xref="S2.SS3.p2.4.m3.2.3.1.cmml">∈</mo><mrow id="S2.SS3.p2.4.m3.2.3.3.2" xref="S2.SS3.p2.4.m3.2.3.3.1.cmml"><mo stretchy="false" id="S2.SS3.p2.4.m3.2.3.3.2.1" xref="S2.SS3.p2.4.m3.2.3.3.1.cmml">(</mo><mn id="S2.SS3.p2.4.m3.1.1" xref="S2.SS3.p2.4.m3.1.1.cmml">0</mn><mo id="S2.SS3.p2.4.m3.2.3.3.2.2" xref="S2.SS3.p2.4.m3.2.3.3.1.cmml">,</mo><mn id="S2.SS3.p2.4.m3.2.2" xref="S2.SS3.p2.4.m3.2.2.cmml">1</mn><mo stretchy="false" id="S2.SS3.p2.4.m3.2.3.3.2.3" xref="S2.SS3.p2.4.m3.2.3.3.1.cmml">]</mo></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.4.m3.2b"><apply id="S2.SS3.p2.4.m3.2.3.cmml" xref="S2.SS3.p2.4.m3.2.3"><in id="S2.SS3.p2.4.m3.2.3.1.cmml" xref="S2.SS3.p2.4.m3.2.3.1"></in><apply id="S2.SS3.p2.4.m3.2.3.2.cmml" xref="S2.SS3.p2.4.m3.2.3.2"><csymbol cd="ambiguous" id="S2.SS3.p2.4.m3.2.3.2.1.cmml" xref="S2.SS3.p2.4.m3.2.3.2">subscript</csymbol><ci id="S2.SS3.p2.4.m3.2.3.2.2.cmml" xref="S2.SS3.p2.4.m3.2.3.2.2">𝑐</ci><ci id="S2.SS3.p2.4.m3.2.3.2.3.cmml" xref="S2.SS3.p2.4.m3.2.3.2.3">𝑑</ci></apply><interval closure="open-closed" id="S2.SS3.p2.4.m3.2.3.3.1.cmml" xref="S2.SS3.p2.4.m3.2.3.3.2"><cn type="integer" id="S2.SS3.p2.4.m3.1.1.cmml" xref="S2.SS3.p2.4.m3.1.1">0</cn><cn type="integer" id="S2.SS3.p2.4.m3.2.2.cmml" xref="S2.SS3.p2.4.m3.2.2">1</cn></interval></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.4.m3.2c">c_{d}\in(0,1]</annotation></semantics></math> is a capacity factor chosen empirically. In practice, we first fixed the MoD layer interval and capacity factor <math id="S2.SS3.p2.5.m4.1" class="ltx_Math" alttext="c_{d}" display="inline"><semantics id="S2.SS3.p2.5.m4.1a"><msub id="S2.SS3.p2.5.m4.1.1" xref="S2.SS3.p2.5.m4.1.1.cmml"><mi id="S2.SS3.p2.5.m4.1.1.2" xref="S2.SS3.p2.5.m4.1.1.2.cmml">c</mi><mi id="S2.SS3.p2.5.m4.1.1.3" xref="S2.SS3.p2.5.m4.1.1.3.cmml">d</mi></msub><annotation-xml encoding="MathML-Content" id="S2.SS3.p2.5.m4.1b"><apply id="S2.SS3.p2.5.m4.1.1.cmml" xref="S2.SS3.p2.5.m4.1.1"><csymbol cd="ambiguous" id="S2.SS3.p2.5.m4.1.1.1.cmml" xref="S2.SS3.p2.5.m4.1.1">subscript</csymbol><ci id="S2.SS3.p2.5.m4.1.1.2.cmml" xref="S2.SS3.p2.5.m4.1.1.2">𝑐</ci><ci id="S2.SS3.p2.5.m4.1.1.3.cmml" xref="S2.SS3.p2.5.m4.1.1.3">𝑑</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS3.p2.5.m4.1c">c_{d}</annotation></semantics></math>, then adjust the total number of transformer layers to ensure the resulting architecture has a comparable FLOPs per token to the base architecture.</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>Inference</h3>

<div id="S2.SS4.p1" class="ltx_para">
<p id="S2.SS4.p1.1" class="ltx_p">We cannot directly apply the expert-choice routing for MoE and layer-choice routing for MoD during inference time, as the top-<math id="S2.SS4.p1.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.p1.1.m1.1a"><mi id="S2.SS4.p1.1.m1.1.1" xref="S2.SS4.p1.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.p1.1.m1.1b"><ci id="S2.SS4.p1.1.m1.1.1.cmml" xref="S2.SS4.p1.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p1.1.m1.1c">k</annotation></semantics></math> token selection within a batch breaks causality. Inspired by <cite class="ltx_cite ltx_citemacro_cite">Raposo et al. (<a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>, to ensure causality during inference, we introduce auxiliary routers, to predict the likelihood of a token being selected by an expert or a layer solely based on its hidden representation.</p>
</div>
<div id="S2.SS4.p2" class="ltx_para">
<p id="S2.SS4.p2.3" class="ltx_p">Formally, for each MoE layer, we introduce an auxiliary router</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.2" class="ltx_Math" alttext="g^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot W^{\mathcal{M}}_{a1})\cdot W^{\mathcal{M}}_{a2})," display="block"><semantics id="S2.E3.m1.2a"><mrow id="S2.E3.m1.2.2.1" xref="S2.E3.m1.2.2.1.1.cmml"><mrow id="S2.E3.m1.2.2.1.1" xref="S2.E3.m1.2.2.1.1.cmml"><mrow id="S2.E3.m1.2.2.1.1.3" xref="S2.E3.m1.2.2.1.1.3.cmml"><msubsup id="S2.E3.m1.2.2.1.1.3.2" xref="S2.E3.m1.2.2.1.1.3.2.cmml"><mi id="S2.E3.m1.2.2.1.1.3.2.2.2" xref="S2.E3.m1.2.2.1.1.3.2.2.2.cmml">g</mi><mtext id="S2.E3.m1.2.2.1.1.3.2.3" xref="S2.E3.m1.2.2.1.1.3.2.3a.cmml">aux</mtext><mi class="ltx_font_mathcaligraphic" id="S2.E3.m1.2.2.1.1.3.2.2.3" xref="S2.E3.m1.2.2.1.1.3.2.2.3.cmml">ℳ</mi></msubsup><mo lspace="0em" rspace="0em" id="S2.E3.m1.2.2.1.1.3.1" xref="S2.E3.m1.2.2.1.1.3.1.cmml">​</mo><mrow id="S2.E3.m1.2.2.1.1.3.3.2" xref="S2.E3.m1.2.2.1.1.3.cmml"><mo stretchy="false" id="S2.E3.m1.2.2.1.1.3.3.2.1" xref="S2.E3.m1.2.2.1.1.3.cmml">(</mo><mi id="S2.E3.m1.1.1" xref="S2.E3.m1.1.1.cmml">x</mi><mo stretchy="false" id="S2.E3.m1.2.2.1.1.3.3.2.2" xref="S2.E3.m1.2.2.1.1.3.cmml">)</mo></mrow></mrow><mo id="S2.E3.m1.2.2.1.1.2" xref="S2.E3.m1.2.2.1.1.2.cmml">=</mo><mrow id="S2.E3.m1.2.2.1.1.1" xref="S2.E3.m1.2.2.1.1.1.cmml"><mtext id="S2.E3.m1.2.2.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.3a.cmml">Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.E3.m1.2.2.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.2.cmml">​</mo><mrow id="S2.E3.m1.2.2.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E3.m1.2.2.1.1.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.cmml"><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.cmml"><mtext id="S2.E3.m1.2.2.1.1.1.1.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.3a.cmml">SiLU</mtext><mo lspace="0em" rspace="0em" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.2.cmml">​</mo><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml"><mi id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.2.cmml">x</mi><mo lspace="0.222em" rspace="0.222em" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.1.cmml">⋅</mo><msubsup id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.cmml"><mi id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml">W</mi><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.cmml"><mi id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1.cmml">​</mo><mn id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3.cmml">1</mn></mrow><mi class="ltx_font_mathcaligraphic" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.3.cmml">ℳ</mi></msubsup></mrow><mo rspace="0.055em" stretchy="false" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml">)</mo></mrow></mrow><mo rspace="0.222em" id="S2.E3.m1.2.2.1.1.1.1.1.1.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.2.cmml">⋅</mo><msubsup id="S2.E3.m1.2.2.1.1.1.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.cmml"><mi id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.2.cmml">W</mi><mrow id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.cmml"><mi id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.2" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.1" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.1.cmml">​</mo><mn id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.3.cmml">2</mn></mrow><mi class="ltx_font_mathcaligraphic" id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.3.cmml">ℳ</mi></msubsup></mrow><mo stretchy="false" id="S2.E3.m1.2.2.1.1.1.1.1.3" xref="S2.E3.m1.2.2.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow><mo id="S2.E3.m1.2.2.1.2" xref="S2.E3.m1.2.2.1.1.cmml">,</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E3.m1.2b"><apply id="S2.E3.m1.2.2.1.1.cmml" xref="S2.E3.m1.2.2.1"><eq id="S2.E3.m1.2.2.1.1.2.cmml" xref="S2.E3.m1.2.2.1.1.2"></eq><apply id="S2.E3.m1.2.2.1.1.3.cmml" xref="S2.E3.m1.2.2.1.1.3"><times id="S2.E3.m1.2.2.1.1.3.1.cmml" xref="S2.E3.m1.2.2.1.1.3.1"></times><apply id="S2.E3.m1.2.2.1.1.3.2.cmml" xref="S2.E3.m1.2.2.1.1.3.2"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.3.2.1.cmml" xref="S2.E3.m1.2.2.1.1.3.2">subscript</csymbol><apply id="S2.E3.m1.2.2.1.1.3.2.2.cmml" xref="S2.E3.m1.2.2.1.1.3.2"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.3.2.2.1.cmml" xref="S2.E3.m1.2.2.1.1.3.2">superscript</csymbol><ci id="S2.E3.m1.2.2.1.1.3.2.2.2.cmml" xref="S2.E3.m1.2.2.1.1.3.2.2.2">𝑔</ci><ci id="S2.E3.m1.2.2.1.1.3.2.2.3.cmml" xref="S2.E3.m1.2.2.1.1.3.2.2.3">ℳ</ci></apply><ci id="S2.E3.m1.2.2.1.1.3.2.3a.cmml" xref="S2.E3.m1.2.2.1.1.3.2.3"><mtext mathsize="70%" id="S2.E3.m1.2.2.1.1.3.2.3.cmml" xref="S2.E3.m1.2.2.1.1.3.2.3">aux</mtext></ci></apply><ci id="S2.E3.m1.1.1.cmml" xref="S2.E3.m1.1.1">𝑥</ci></apply><apply id="S2.E3.m1.2.2.1.1.1.cmml" xref="S2.E3.m1.2.2.1.1.1"><times id="S2.E3.m1.2.2.1.1.1.2.cmml" xref="S2.E3.m1.2.2.1.1.1.2"></times><ci id="S2.E3.m1.2.2.1.1.1.3a.cmml" xref="S2.E3.m1.2.2.1.1.1.3"><mtext id="S2.E3.m1.2.2.1.1.1.3.cmml" xref="S2.E3.m1.2.2.1.1.1.3">Sigmoid</mtext></ci><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1"><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.2">⋅</ci><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1"><times id="S2.E3.m1.2.2.1.1.1.1.1.1.1.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.2"></times><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.3a.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.3"><mtext id="S2.E3.m1.2.2.1.1.1.1.1.1.1.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.3">SiLU</mtext></ci><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1"><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.1">⋅</ci><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.2">𝑥</ci><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3">subscript</csymbol><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3">superscript</csymbol><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2">𝑊</ci><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.3">ℳ</ci></apply><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3"><times id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1"></times><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2">𝑎</ci><cn type="integer" id="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3">1</cn></apply></apply></apply></apply><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.1.1.1.1.3.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3">subscript</csymbol><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3">superscript</csymbol><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.2">𝑊</ci><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.2.3">ℳ</ci></apply><apply id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3"><times id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.1.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.1"></times><ci id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.2.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.2">𝑎</ci><cn type="integer" id="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.3.cmml" xref="S2.E3.m1.2.2.1.1.1.1.1.1.3.3.3">2</cn></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E3.m1.2c">g^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot W^{\mathcal{M}}_{a1})\cdot W^{\mathcal{M}}_{a2}),</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>
<p id="S2.SS4.p2.2" class="ltx_p">where <math id="S2.SS4.p2.1.m1.1" class="ltx_Math" alttext="W^{\mathcal{M}}_{a1}\in\mathbb{R}^{d\times(d/2)}" display="inline"><semantics id="S2.SS4.p2.1.m1.1a"><mrow id="S2.SS4.p2.1.m1.1.2" xref="S2.SS4.p2.1.m1.1.2.cmml"><msubsup id="S2.SS4.p2.1.m1.1.2.2" xref="S2.SS4.p2.1.m1.1.2.2.cmml"><mi id="S2.SS4.p2.1.m1.1.2.2.2.2" xref="S2.SS4.p2.1.m1.1.2.2.2.2.cmml">W</mi><mrow id="S2.SS4.p2.1.m1.1.2.2.3" xref="S2.SS4.p2.1.m1.1.2.2.3.cmml"><mi id="S2.SS4.p2.1.m1.1.2.2.3.2" xref="S2.SS4.p2.1.m1.1.2.2.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.SS4.p2.1.m1.1.2.2.3.1" xref="S2.SS4.p2.1.m1.1.2.2.3.1.cmml">​</mo><mn id="S2.SS4.p2.1.m1.1.2.2.3.3" xref="S2.SS4.p2.1.m1.1.2.2.3.3.cmml">1</mn></mrow><mi class="ltx_font_mathcaligraphic" id="S2.SS4.p2.1.m1.1.2.2.2.3" xref="S2.SS4.p2.1.m1.1.2.2.2.3.cmml">ℳ</mi></msubsup><mo id="S2.SS4.p2.1.m1.1.2.1" xref="S2.SS4.p2.1.m1.1.2.1.cmml">∈</mo><msup id="S2.SS4.p2.1.m1.1.2.3" xref="S2.SS4.p2.1.m1.1.2.3.cmml"><mi id="S2.SS4.p2.1.m1.1.2.3.2" xref="S2.SS4.p2.1.m1.1.2.3.2.cmml">ℝ</mi><mrow id="S2.SS4.p2.1.m1.1.1.1" xref="S2.SS4.p2.1.m1.1.1.1.cmml"><mi id="S2.SS4.p2.1.m1.1.1.1.3" xref="S2.SS4.p2.1.m1.1.1.1.3.cmml">d</mi><mo lspace="0.222em" rspace="0.222em" id="S2.SS4.p2.1.m1.1.1.1.2" xref="S2.SS4.p2.1.m1.1.1.1.2.cmml">×</mo><mrow id="S2.SS4.p2.1.m1.1.1.1.1.1" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.p2.1.m1.1.1.1.1.1.2" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.SS4.p2.1.m1.1.1.1.1.1.1" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.cmml"><mi id="S2.SS4.p2.1.m1.1.1.1.1.1.1.2" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.2.cmml">d</mi><mo id="S2.SS4.p2.1.m1.1.1.1.1.1.1.1" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.1.cmml">/</mo><mn id="S2.SS4.p2.1.m1.1.1.1.1.1.1.3" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.3.cmml">2</mn></mrow><mo stretchy="false" id="S2.SS4.p2.1.m1.1.1.1.1.1.3" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.p2.1.m1.1b"><apply id="S2.SS4.p2.1.m1.1.2.cmml" xref="S2.SS4.p2.1.m1.1.2"><in id="S2.SS4.p2.1.m1.1.2.1.cmml" xref="S2.SS4.p2.1.m1.1.2.1"></in><apply id="S2.SS4.p2.1.m1.1.2.2.cmml" xref="S2.SS4.p2.1.m1.1.2.2"><csymbol cd="ambiguous" id="S2.SS4.p2.1.m1.1.2.2.1.cmml" xref="S2.SS4.p2.1.m1.1.2.2">subscript</csymbol><apply id="S2.SS4.p2.1.m1.1.2.2.2.cmml" xref="S2.SS4.p2.1.m1.1.2.2"><csymbol cd="ambiguous" id="S2.SS4.p2.1.m1.1.2.2.2.1.cmml" xref="S2.SS4.p2.1.m1.1.2.2">superscript</csymbol><ci id="S2.SS4.p2.1.m1.1.2.2.2.2.cmml" xref="S2.SS4.p2.1.m1.1.2.2.2.2">𝑊</ci><ci id="S2.SS4.p2.1.m1.1.2.2.2.3.cmml" xref="S2.SS4.p2.1.m1.1.2.2.2.3">ℳ</ci></apply><apply id="S2.SS4.p2.1.m1.1.2.2.3.cmml" xref="S2.SS4.p2.1.m1.1.2.2.3"><times id="S2.SS4.p2.1.m1.1.2.2.3.1.cmml" xref="S2.SS4.p2.1.m1.1.2.2.3.1"></times><ci id="S2.SS4.p2.1.m1.1.2.2.3.2.cmml" xref="S2.SS4.p2.1.m1.1.2.2.3.2">𝑎</ci><cn type="integer" id="S2.SS4.p2.1.m1.1.2.2.3.3.cmml" xref="S2.SS4.p2.1.m1.1.2.2.3.3">1</cn></apply></apply><apply id="S2.SS4.p2.1.m1.1.2.3.cmml" xref="S2.SS4.p2.1.m1.1.2.3"><csymbol cd="ambiguous" id="S2.SS4.p2.1.m1.1.2.3.1.cmml" xref="S2.SS4.p2.1.m1.1.2.3">superscript</csymbol><ci id="S2.SS4.p2.1.m1.1.2.3.2.cmml" xref="S2.SS4.p2.1.m1.1.2.3.2">ℝ</ci><apply id="S2.SS4.p2.1.m1.1.1.1.cmml" xref="S2.SS4.p2.1.m1.1.1.1"><times id="S2.SS4.p2.1.m1.1.1.1.2.cmml" xref="S2.SS4.p2.1.m1.1.1.1.2"></times><ci id="S2.SS4.p2.1.m1.1.1.1.3.cmml" xref="S2.SS4.p2.1.m1.1.1.1.3">𝑑</ci><apply id="S2.SS4.p2.1.m1.1.1.1.1.1.1.cmml" xref="S2.SS4.p2.1.m1.1.1.1.1.1"><divide id="S2.SS4.p2.1.m1.1.1.1.1.1.1.1.cmml" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.1"></divide><ci id="S2.SS4.p2.1.m1.1.1.1.1.1.1.2.cmml" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.2">𝑑</ci><cn type="integer" id="S2.SS4.p2.1.m1.1.1.1.1.1.1.3.cmml" xref="S2.SS4.p2.1.m1.1.1.1.1.1.1.3">2</cn></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p2.1.m1.1c">W^{\mathcal{M}}_{a1}\in\mathbb{R}^{d\times(d/2)}</annotation></semantics></math> and <math id="S2.SS4.p2.2.m2.2" class="ltx_Math" alttext="W^{\mathcal{M}}_{a2}\in\mathbb{R}^{(d/2)\times|\mathcal{E}^{\mathcal{M}}|}" display="inline"><semantics id="S2.SS4.p2.2.m2.2a"><mrow id="S2.SS4.p2.2.m2.2.3" xref="S2.SS4.p2.2.m2.2.3.cmml"><msubsup id="S2.SS4.p2.2.m2.2.3.2" xref="S2.SS4.p2.2.m2.2.3.2.cmml"><mi id="S2.SS4.p2.2.m2.2.3.2.2.2" xref="S2.SS4.p2.2.m2.2.3.2.2.2.cmml">W</mi><mrow id="S2.SS4.p2.2.m2.2.3.2.3" xref="S2.SS4.p2.2.m2.2.3.2.3.cmml"><mi id="S2.SS4.p2.2.m2.2.3.2.3.2" xref="S2.SS4.p2.2.m2.2.3.2.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.SS4.p2.2.m2.2.3.2.3.1" xref="S2.SS4.p2.2.m2.2.3.2.3.1.cmml">​</mo><mn id="S2.SS4.p2.2.m2.2.3.2.3.3" xref="S2.SS4.p2.2.m2.2.3.2.3.3.cmml">2</mn></mrow><mi class="ltx_font_mathcaligraphic" id="S2.SS4.p2.2.m2.2.3.2.2.3" xref="S2.SS4.p2.2.m2.2.3.2.2.3.cmml">ℳ</mi></msubsup><mo id="S2.SS4.p2.2.m2.2.3.1" xref="S2.SS4.p2.2.m2.2.3.1.cmml">∈</mo><msup id="S2.SS4.p2.2.m2.2.3.3" xref="S2.SS4.p2.2.m2.2.3.3.cmml"><mi id="S2.SS4.p2.2.m2.2.3.3.2" xref="S2.SS4.p2.2.m2.2.3.3.2.cmml">ℝ</mi><mrow id="S2.SS4.p2.2.m2.2.2.2" xref="S2.SS4.p2.2.m2.2.2.2.cmml"><mrow id="S2.SS4.p2.2.m2.1.1.1.1.1" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.p2.2.m2.1.1.1.1.1.2" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.SS4.p2.2.m2.1.1.1.1.1.1" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.cmml"><mi id="S2.SS4.p2.2.m2.1.1.1.1.1.1.2" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.2.cmml">d</mi><mo id="S2.SS4.p2.2.m2.1.1.1.1.1.1.1" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.1.cmml">/</mo><mn id="S2.SS4.p2.2.m2.1.1.1.1.1.1.3" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.3.cmml">2</mn></mrow><mo rspace="0.055em" stretchy="false" id="S2.SS4.p2.2.m2.1.1.1.1.1.3" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.cmml">)</mo></mrow><mo rspace="0.222em" id="S2.SS4.p2.2.m2.2.2.2.3" xref="S2.SS4.p2.2.m2.2.2.2.3.cmml">×</mo><mrow id="S2.SS4.p2.2.m2.2.2.2.2.1" xref="S2.SS4.p2.2.m2.2.2.2.2.2.cmml"><mo stretchy="false" id="S2.SS4.p2.2.m2.2.2.2.2.1.2" xref="S2.SS4.p2.2.m2.2.2.2.2.2.1.cmml">|</mo><msup id="S2.SS4.p2.2.m2.2.2.2.2.1.1" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S2.SS4.p2.2.m2.2.2.2.2.1.1.2" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1.2.cmml">ℰ</mi><mi class="ltx_font_mathcaligraphic" id="S2.SS4.p2.2.m2.2.2.2.2.1.1.3" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1.3.cmml">ℳ</mi></msup><mo stretchy="false" id="S2.SS4.p2.2.m2.2.2.2.2.1.3" xref="S2.SS4.p2.2.m2.2.2.2.2.2.1.cmml">|</mo></mrow></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.p2.2.m2.2b"><apply id="S2.SS4.p2.2.m2.2.3.cmml" xref="S2.SS4.p2.2.m2.2.3"><in id="S2.SS4.p2.2.m2.2.3.1.cmml" xref="S2.SS4.p2.2.m2.2.3.1"></in><apply id="S2.SS4.p2.2.m2.2.3.2.cmml" xref="S2.SS4.p2.2.m2.2.3.2"><csymbol cd="ambiguous" id="S2.SS4.p2.2.m2.2.3.2.1.cmml" xref="S2.SS4.p2.2.m2.2.3.2">subscript</csymbol><apply id="S2.SS4.p2.2.m2.2.3.2.2.cmml" xref="S2.SS4.p2.2.m2.2.3.2"><csymbol cd="ambiguous" id="S2.SS4.p2.2.m2.2.3.2.2.1.cmml" xref="S2.SS4.p2.2.m2.2.3.2">superscript</csymbol><ci id="S2.SS4.p2.2.m2.2.3.2.2.2.cmml" xref="S2.SS4.p2.2.m2.2.3.2.2.2">𝑊</ci><ci id="S2.SS4.p2.2.m2.2.3.2.2.3.cmml" xref="S2.SS4.p2.2.m2.2.3.2.2.3">ℳ</ci></apply><apply id="S2.SS4.p2.2.m2.2.3.2.3.cmml" xref="S2.SS4.p2.2.m2.2.3.2.3"><times id="S2.SS4.p2.2.m2.2.3.2.3.1.cmml" xref="S2.SS4.p2.2.m2.2.3.2.3.1"></times><ci id="S2.SS4.p2.2.m2.2.3.2.3.2.cmml" xref="S2.SS4.p2.2.m2.2.3.2.3.2">𝑎</ci><cn type="integer" id="S2.SS4.p2.2.m2.2.3.2.3.3.cmml" xref="S2.SS4.p2.2.m2.2.3.2.3.3">2</cn></apply></apply><apply id="S2.SS4.p2.2.m2.2.3.3.cmml" xref="S2.SS4.p2.2.m2.2.3.3"><csymbol cd="ambiguous" id="S2.SS4.p2.2.m2.2.3.3.1.cmml" xref="S2.SS4.p2.2.m2.2.3.3">superscript</csymbol><ci id="S2.SS4.p2.2.m2.2.3.3.2.cmml" xref="S2.SS4.p2.2.m2.2.3.3.2">ℝ</ci><apply id="S2.SS4.p2.2.m2.2.2.2.cmml" xref="S2.SS4.p2.2.m2.2.2.2"><times id="S2.SS4.p2.2.m2.2.2.2.3.cmml" xref="S2.SS4.p2.2.m2.2.2.2.3"></times><apply id="S2.SS4.p2.2.m2.1.1.1.1.1.1.cmml" xref="S2.SS4.p2.2.m2.1.1.1.1.1"><divide id="S2.SS4.p2.2.m2.1.1.1.1.1.1.1.cmml" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.1"></divide><ci id="S2.SS4.p2.2.m2.1.1.1.1.1.1.2.cmml" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.2">𝑑</ci><cn type="integer" id="S2.SS4.p2.2.m2.1.1.1.1.1.1.3.cmml" xref="S2.SS4.p2.2.m2.1.1.1.1.1.1.3">2</cn></apply><apply id="S2.SS4.p2.2.m2.2.2.2.2.2.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1"><abs id="S2.SS4.p2.2.m2.2.2.2.2.2.1.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1.2"></abs><apply id="S2.SS4.p2.2.m2.2.2.2.2.1.1.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1"><csymbol cd="ambiguous" id="S2.SS4.p2.2.m2.2.2.2.2.1.1.1.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1">superscript</csymbol><ci id="S2.SS4.p2.2.m2.2.2.2.2.1.1.2.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1.2">ℰ</ci><ci id="S2.SS4.p2.2.m2.2.2.2.2.1.1.3.cmml" xref="S2.SS4.p2.2.m2.2.2.2.2.1.1.3">ℳ</ci></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p2.2.m2.2c">W^{\mathcal{M}}_{a2}\in\mathbb{R}^{(d/2)\times|\mathcal{E}^{\mathcal{M}}|}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS4.p3" class="ltx_para">
<p id="S2.SS4.p3.3" class="ltx_p">For each MoD layer, we introduce an auxiliary router</p>
<table id="S2.E4" 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.E4.m1.2" class="ltx_Math" alttext="\hat{g}^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot\hat{W}_{a1})\cdot\hat{W}_{a2})," display="block"><semantics id="S2.E4.m1.2a"><mrow id="S2.E4.m1.2.2.1" xref="S2.E4.m1.2.2.1.1.cmml"><mrow id="S2.E4.m1.2.2.1.1" xref="S2.E4.m1.2.2.1.1.cmml"><mrow id="S2.E4.m1.2.2.1.1.3" xref="S2.E4.m1.2.2.1.1.3.cmml"><msubsup id="S2.E4.m1.2.2.1.1.3.2" xref="S2.E4.m1.2.2.1.1.3.2.cmml"><mover accent="true" id="S2.E4.m1.2.2.1.1.3.2.2.2" xref="S2.E4.m1.2.2.1.1.3.2.2.2.cmml"><mi id="S2.E4.m1.2.2.1.1.3.2.2.2.2" xref="S2.E4.m1.2.2.1.1.3.2.2.2.2.cmml">g</mi><mo id="S2.E4.m1.2.2.1.1.3.2.2.2.1" xref="S2.E4.m1.2.2.1.1.3.2.2.2.1.cmml">^</mo></mover><mtext id="S2.E4.m1.2.2.1.1.3.2.3" xref="S2.E4.m1.2.2.1.1.3.2.3a.cmml">aux</mtext><mi class="ltx_font_mathcaligraphic" id="S2.E4.m1.2.2.1.1.3.2.2.3" xref="S2.E4.m1.2.2.1.1.3.2.2.3.cmml">ℳ</mi></msubsup><mo lspace="0em" rspace="0em" id="S2.E4.m1.2.2.1.1.3.1" xref="S2.E4.m1.2.2.1.1.3.1.cmml">​</mo><mrow id="S2.E4.m1.2.2.1.1.3.3.2" xref="S2.E4.m1.2.2.1.1.3.cmml"><mo stretchy="false" id="S2.E4.m1.2.2.1.1.3.3.2.1" xref="S2.E4.m1.2.2.1.1.3.cmml">(</mo><mi id="S2.E4.m1.1.1" xref="S2.E4.m1.1.1.cmml">x</mi><mo stretchy="false" id="S2.E4.m1.2.2.1.1.3.3.2.2" xref="S2.E4.m1.2.2.1.1.3.cmml">)</mo></mrow></mrow><mo id="S2.E4.m1.2.2.1.1.2" xref="S2.E4.m1.2.2.1.1.2.cmml">=</mo><mrow id="S2.E4.m1.2.2.1.1.1" xref="S2.E4.m1.2.2.1.1.1.cmml"><mtext id="S2.E4.m1.2.2.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.3a.cmml">Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.E4.m1.2.2.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.2.cmml">​</mo><mrow id="S2.E4.m1.2.2.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E4.m1.2.2.1.1.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.cmml"><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.cmml"><mtext id="S2.E4.m1.2.2.1.1.1.1.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.3a.cmml">SiLU</mtext><mo lspace="0em" rspace="0em" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.2.cmml">​</mo><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml"><mi id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.2.cmml">x</mi><mo lspace="0.222em" rspace="0.222em" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.1.cmml">⋅</mo><msub id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.cmml"><mover accent="true" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.cmml"><mi id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml">W</mi><mo id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.1.cmml">^</mo></mover><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.cmml"><mi id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1.cmml">​</mo><mn id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3.cmml">1</mn></mrow></msub></mrow><mo rspace="0.055em" stretchy="false" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml">)</mo></mrow></mrow><mo rspace="0.222em" id="S2.E4.m1.2.2.1.1.1.1.1.1.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.2.cmml">⋅</mo><msub id="S2.E4.m1.2.2.1.1.1.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.cmml"><mover accent="true" id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.cmml"><mi id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.2.cmml">W</mi><mo id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.1.cmml">^</mo></mover><mrow id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.cmml"><mi id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.2" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.1" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.1.cmml">​</mo><mn id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.3.cmml">2</mn></mrow></msub></mrow><mo stretchy="false" id="S2.E4.m1.2.2.1.1.1.1.1.3" xref="S2.E4.m1.2.2.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></mrow><mo id="S2.E4.m1.2.2.1.2" xref="S2.E4.m1.2.2.1.1.cmml">,</mo></mrow><annotation-xml encoding="MathML-Content" id="S2.E4.m1.2b"><apply id="S2.E4.m1.2.2.1.1.cmml" xref="S2.E4.m1.2.2.1"><eq id="S2.E4.m1.2.2.1.1.2.cmml" xref="S2.E4.m1.2.2.1.1.2"></eq><apply id="S2.E4.m1.2.2.1.1.3.cmml" xref="S2.E4.m1.2.2.1.1.3"><times id="S2.E4.m1.2.2.1.1.3.1.cmml" xref="S2.E4.m1.2.2.1.1.3.1"></times><apply id="S2.E4.m1.2.2.1.1.3.2.cmml" xref="S2.E4.m1.2.2.1.1.3.2"><csymbol cd="ambiguous" id="S2.E4.m1.2.2.1.1.3.2.1.cmml" xref="S2.E4.m1.2.2.1.1.3.2">subscript</csymbol><apply id="S2.E4.m1.2.2.1.1.3.2.2.cmml" xref="S2.E4.m1.2.2.1.1.3.2"><csymbol cd="ambiguous" id="S2.E4.m1.2.2.1.1.3.2.2.1.cmml" xref="S2.E4.m1.2.2.1.1.3.2">superscript</csymbol><apply id="S2.E4.m1.2.2.1.1.3.2.2.2.cmml" xref="S2.E4.m1.2.2.1.1.3.2.2.2"><ci id="S2.E4.m1.2.2.1.1.3.2.2.2.1.cmml" xref="S2.E4.m1.2.2.1.1.3.2.2.2.1">^</ci><ci id="S2.E4.m1.2.2.1.1.3.2.2.2.2.cmml" xref="S2.E4.m1.2.2.1.1.3.2.2.2.2">𝑔</ci></apply><ci id="S2.E4.m1.2.2.1.1.3.2.2.3.cmml" xref="S2.E4.m1.2.2.1.1.3.2.2.3">ℳ</ci></apply><ci id="S2.E4.m1.2.2.1.1.3.2.3a.cmml" xref="S2.E4.m1.2.2.1.1.3.2.3"><mtext mathsize="70%" id="S2.E4.m1.2.2.1.1.3.2.3.cmml" xref="S2.E4.m1.2.2.1.1.3.2.3">aux</mtext></ci></apply><ci id="S2.E4.m1.1.1.cmml" xref="S2.E4.m1.1.1">𝑥</ci></apply><apply id="S2.E4.m1.2.2.1.1.1.cmml" xref="S2.E4.m1.2.2.1.1.1"><times id="S2.E4.m1.2.2.1.1.1.2.cmml" xref="S2.E4.m1.2.2.1.1.1.2"></times><ci id="S2.E4.m1.2.2.1.1.1.3a.cmml" xref="S2.E4.m1.2.2.1.1.1.3"><mtext id="S2.E4.m1.2.2.1.1.1.3.cmml" xref="S2.E4.m1.2.2.1.1.1.3">Sigmoid</mtext></ci><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1"><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.2">⋅</ci><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1"><times id="S2.E4.m1.2.2.1.1.1.1.1.1.1.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.2"></times><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.3a.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.3"><mtext id="S2.E4.m1.2.2.1.1.1.1.1.1.1.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.3">SiLU</mtext></ci><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1"><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.1">⋅</ci><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.2">𝑥</ci><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3">subscript</csymbol><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2"><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.1">^</ci><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.2.2">𝑊</ci></apply><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3"><times id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.1"></times><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.2">𝑎</ci><cn type="integer" id="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.1.1.1.1.3.3.3">1</cn></apply></apply></apply></apply><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E4.m1.2.2.1.1.1.1.1.1.3.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3">subscript</csymbol><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2"><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.1">^</ci><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.2.2">𝑊</ci></apply><apply id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3"><times id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.1.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.1"></times><ci id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.2.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.2">𝑎</ci><cn type="integer" id="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.3.cmml" xref="S2.E4.m1.2.2.1.1.1.1.1.1.3.3.3">2</cn></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E4.m1.2c">\hat{g}^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot\hat{W}_{a1})\cdot\hat{W}_{a2}),</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">(4)</span></td>
</tr></tbody>
</table>
<p id="S2.SS4.p3.2" class="ltx_p">where <math id="S2.SS4.p3.1.m1.1" class="ltx_Math" alttext="\hat{W}_{a1}\in\mathbb{R}^{d\times(d/2)}" display="inline"><semantics id="S2.SS4.p3.1.m1.1a"><mrow id="S2.SS4.p3.1.m1.1.2" xref="S2.SS4.p3.1.m1.1.2.cmml"><msub id="S2.SS4.p3.1.m1.1.2.2" xref="S2.SS4.p3.1.m1.1.2.2.cmml"><mover accent="true" id="S2.SS4.p3.1.m1.1.2.2.2" xref="S2.SS4.p3.1.m1.1.2.2.2.cmml"><mi id="S2.SS4.p3.1.m1.1.2.2.2.2" xref="S2.SS4.p3.1.m1.1.2.2.2.2.cmml">W</mi><mo id="S2.SS4.p3.1.m1.1.2.2.2.1" xref="S2.SS4.p3.1.m1.1.2.2.2.1.cmml">^</mo></mover><mrow id="S2.SS4.p3.1.m1.1.2.2.3" xref="S2.SS4.p3.1.m1.1.2.2.3.cmml"><mi id="S2.SS4.p3.1.m1.1.2.2.3.2" xref="S2.SS4.p3.1.m1.1.2.2.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.SS4.p3.1.m1.1.2.2.3.1" xref="S2.SS4.p3.1.m1.1.2.2.3.1.cmml">​</mo><mn id="S2.SS4.p3.1.m1.1.2.2.3.3" xref="S2.SS4.p3.1.m1.1.2.2.3.3.cmml">1</mn></mrow></msub><mo id="S2.SS4.p3.1.m1.1.2.1" xref="S2.SS4.p3.1.m1.1.2.1.cmml">∈</mo><msup id="S2.SS4.p3.1.m1.1.2.3" xref="S2.SS4.p3.1.m1.1.2.3.cmml"><mi id="S2.SS4.p3.1.m1.1.2.3.2" xref="S2.SS4.p3.1.m1.1.2.3.2.cmml">ℝ</mi><mrow id="S2.SS4.p3.1.m1.1.1.1" xref="S2.SS4.p3.1.m1.1.1.1.cmml"><mi id="S2.SS4.p3.1.m1.1.1.1.3" xref="S2.SS4.p3.1.m1.1.1.1.3.cmml">d</mi><mo lspace="0.222em" rspace="0.222em" id="S2.SS4.p3.1.m1.1.1.1.2" xref="S2.SS4.p3.1.m1.1.1.1.2.cmml">×</mo><mrow id="S2.SS4.p3.1.m1.1.1.1.1.1" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.p3.1.m1.1.1.1.1.1.2" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.SS4.p3.1.m1.1.1.1.1.1.1" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.cmml"><mi id="S2.SS4.p3.1.m1.1.1.1.1.1.1.2" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.2.cmml">d</mi><mo id="S2.SS4.p3.1.m1.1.1.1.1.1.1.1" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.1.cmml">/</mo><mn id="S2.SS4.p3.1.m1.1.1.1.1.1.1.3" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.3.cmml">2</mn></mrow><mo stretchy="false" id="S2.SS4.p3.1.m1.1.1.1.1.1.3" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.cmml">)</mo></mrow></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.p3.1.m1.1b"><apply id="S2.SS4.p3.1.m1.1.2.cmml" xref="S2.SS4.p3.1.m1.1.2"><in id="S2.SS4.p3.1.m1.1.2.1.cmml" xref="S2.SS4.p3.1.m1.1.2.1"></in><apply id="S2.SS4.p3.1.m1.1.2.2.cmml" xref="S2.SS4.p3.1.m1.1.2.2"><csymbol cd="ambiguous" id="S2.SS4.p3.1.m1.1.2.2.1.cmml" xref="S2.SS4.p3.1.m1.1.2.2">subscript</csymbol><apply id="S2.SS4.p3.1.m1.1.2.2.2.cmml" xref="S2.SS4.p3.1.m1.1.2.2.2"><ci id="S2.SS4.p3.1.m1.1.2.2.2.1.cmml" xref="S2.SS4.p3.1.m1.1.2.2.2.1">^</ci><ci id="S2.SS4.p3.1.m1.1.2.2.2.2.cmml" xref="S2.SS4.p3.1.m1.1.2.2.2.2">𝑊</ci></apply><apply id="S2.SS4.p3.1.m1.1.2.2.3.cmml" xref="S2.SS4.p3.1.m1.1.2.2.3"><times id="S2.SS4.p3.1.m1.1.2.2.3.1.cmml" xref="S2.SS4.p3.1.m1.1.2.2.3.1"></times><ci id="S2.SS4.p3.1.m1.1.2.2.3.2.cmml" xref="S2.SS4.p3.1.m1.1.2.2.3.2">𝑎</ci><cn type="integer" id="S2.SS4.p3.1.m1.1.2.2.3.3.cmml" xref="S2.SS4.p3.1.m1.1.2.2.3.3">1</cn></apply></apply><apply id="S2.SS4.p3.1.m1.1.2.3.cmml" xref="S2.SS4.p3.1.m1.1.2.3"><csymbol cd="ambiguous" id="S2.SS4.p3.1.m1.1.2.3.1.cmml" xref="S2.SS4.p3.1.m1.1.2.3">superscript</csymbol><ci id="S2.SS4.p3.1.m1.1.2.3.2.cmml" xref="S2.SS4.p3.1.m1.1.2.3.2">ℝ</ci><apply id="S2.SS4.p3.1.m1.1.1.1.cmml" xref="S2.SS4.p3.1.m1.1.1.1"><times id="S2.SS4.p3.1.m1.1.1.1.2.cmml" xref="S2.SS4.p3.1.m1.1.1.1.2"></times><ci id="S2.SS4.p3.1.m1.1.1.1.3.cmml" xref="S2.SS4.p3.1.m1.1.1.1.3">𝑑</ci><apply id="S2.SS4.p3.1.m1.1.1.1.1.1.1.cmml" xref="S2.SS4.p3.1.m1.1.1.1.1.1"><divide id="S2.SS4.p3.1.m1.1.1.1.1.1.1.1.cmml" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.1"></divide><ci id="S2.SS4.p3.1.m1.1.1.1.1.1.1.2.cmml" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.2">𝑑</ci><cn type="integer" id="S2.SS4.p3.1.m1.1.1.1.1.1.1.3.cmml" xref="S2.SS4.p3.1.m1.1.1.1.1.1.1.3">2</cn></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p3.1.m1.1c">\hat{W}_{a1}\in\mathbb{R}^{d\times(d/2)}</annotation></semantics></math> and <math id="S2.SS4.p3.2.m2.1" class="ltx_Math" alttext="\hat{W}_{a2}\in\mathbb{R}^{(d/2)\times 1}" display="inline"><semantics id="S2.SS4.p3.2.m2.1a"><mrow id="S2.SS4.p3.2.m2.1.2" xref="S2.SS4.p3.2.m2.1.2.cmml"><msub id="S2.SS4.p3.2.m2.1.2.2" xref="S2.SS4.p3.2.m2.1.2.2.cmml"><mover accent="true" id="S2.SS4.p3.2.m2.1.2.2.2" xref="S2.SS4.p3.2.m2.1.2.2.2.cmml"><mi id="S2.SS4.p3.2.m2.1.2.2.2.2" xref="S2.SS4.p3.2.m2.1.2.2.2.2.cmml">W</mi><mo id="S2.SS4.p3.2.m2.1.2.2.2.1" xref="S2.SS4.p3.2.m2.1.2.2.2.1.cmml">^</mo></mover><mrow id="S2.SS4.p3.2.m2.1.2.2.3" xref="S2.SS4.p3.2.m2.1.2.2.3.cmml"><mi id="S2.SS4.p3.2.m2.1.2.2.3.2" xref="S2.SS4.p3.2.m2.1.2.2.3.2.cmml">a</mi><mo lspace="0em" rspace="0em" id="S2.SS4.p3.2.m2.1.2.2.3.1" xref="S2.SS4.p3.2.m2.1.2.2.3.1.cmml">​</mo><mn id="S2.SS4.p3.2.m2.1.2.2.3.3" xref="S2.SS4.p3.2.m2.1.2.2.3.3.cmml">2</mn></mrow></msub><mo id="S2.SS4.p3.2.m2.1.2.1" xref="S2.SS4.p3.2.m2.1.2.1.cmml">∈</mo><msup id="S2.SS4.p3.2.m2.1.2.3" xref="S2.SS4.p3.2.m2.1.2.3.cmml"><mi id="S2.SS4.p3.2.m2.1.2.3.2" xref="S2.SS4.p3.2.m2.1.2.3.2.cmml">ℝ</mi><mrow id="S2.SS4.p3.2.m2.1.1.1" xref="S2.SS4.p3.2.m2.1.1.1.cmml"><mrow id="S2.SS4.p3.2.m2.1.1.1.1.1" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.cmml"><mo stretchy="false" id="S2.SS4.p3.2.m2.1.1.1.1.1.2" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.cmml">(</mo><mrow id="S2.SS4.p3.2.m2.1.1.1.1.1.1" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.cmml"><mi id="S2.SS4.p3.2.m2.1.1.1.1.1.1.2" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.2.cmml">d</mi><mo id="S2.SS4.p3.2.m2.1.1.1.1.1.1.1" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.1.cmml">/</mo><mn id="S2.SS4.p3.2.m2.1.1.1.1.1.1.3" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.3.cmml">2</mn></mrow><mo rspace="0.055em" stretchy="false" id="S2.SS4.p3.2.m2.1.1.1.1.1.3" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.cmml">)</mo></mrow><mo rspace="0.222em" id="S2.SS4.p3.2.m2.1.1.1.2" xref="S2.SS4.p3.2.m2.1.1.1.2.cmml">×</mo><mn id="S2.SS4.p3.2.m2.1.1.1.3" xref="S2.SS4.p3.2.m2.1.1.1.3.cmml">1</mn></mrow></msup></mrow><annotation-xml encoding="MathML-Content" id="S2.SS4.p3.2.m2.1b"><apply id="S2.SS4.p3.2.m2.1.2.cmml" xref="S2.SS4.p3.2.m2.1.2"><in id="S2.SS4.p3.2.m2.1.2.1.cmml" xref="S2.SS4.p3.2.m2.1.2.1"></in><apply id="S2.SS4.p3.2.m2.1.2.2.cmml" xref="S2.SS4.p3.2.m2.1.2.2"><csymbol cd="ambiguous" id="S2.SS4.p3.2.m2.1.2.2.1.cmml" xref="S2.SS4.p3.2.m2.1.2.2">subscript</csymbol><apply id="S2.SS4.p3.2.m2.1.2.2.2.cmml" xref="S2.SS4.p3.2.m2.1.2.2.2"><ci id="S2.SS4.p3.2.m2.1.2.2.2.1.cmml" xref="S2.SS4.p3.2.m2.1.2.2.2.1">^</ci><ci id="S2.SS4.p3.2.m2.1.2.2.2.2.cmml" xref="S2.SS4.p3.2.m2.1.2.2.2.2">𝑊</ci></apply><apply id="S2.SS4.p3.2.m2.1.2.2.3.cmml" xref="S2.SS4.p3.2.m2.1.2.2.3"><times id="S2.SS4.p3.2.m2.1.2.2.3.1.cmml" xref="S2.SS4.p3.2.m2.1.2.2.3.1"></times><ci id="S2.SS4.p3.2.m2.1.2.2.3.2.cmml" xref="S2.SS4.p3.2.m2.1.2.2.3.2">𝑎</ci><cn type="integer" id="S2.SS4.p3.2.m2.1.2.2.3.3.cmml" xref="S2.SS4.p3.2.m2.1.2.2.3.3">2</cn></apply></apply><apply id="S2.SS4.p3.2.m2.1.2.3.cmml" xref="S2.SS4.p3.2.m2.1.2.3"><csymbol cd="ambiguous" id="S2.SS4.p3.2.m2.1.2.3.1.cmml" xref="S2.SS4.p3.2.m2.1.2.3">superscript</csymbol><ci id="S2.SS4.p3.2.m2.1.2.3.2.cmml" xref="S2.SS4.p3.2.m2.1.2.3.2">ℝ</ci><apply id="S2.SS4.p3.2.m2.1.1.1.cmml" xref="S2.SS4.p3.2.m2.1.1.1"><times id="S2.SS4.p3.2.m2.1.1.1.2.cmml" xref="S2.SS4.p3.2.m2.1.1.1.2"></times><apply id="S2.SS4.p3.2.m2.1.1.1.1.1.1.cmml" xref="S2.SS4.p3.2.m2.1.1.1.1.1"><divide id="S2.SS4.p3.2.m2.1.1.1.1.1.1.1.cmml" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.1"></divide><ci id="S2.SS4.p3.2.m2.1.1.1.1.1.1.2.cmml" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.2">𝑑</ci><cn type="integer" id="S2.SS4.p3.2.m2.1.1.1.1.1.1.3.cmml" xref="S2.SS4.p3.2.m2.1.1.1.1.1.1.3">2</cn></apply><cn type="integer" id="S2.SS4.p3.2.m2.1.1.1.3.cmml" xref="S2.SS4.p3.2.m2.1.1.1.3">1</cn></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p3.2.m2.1c">\hat{W}_{a2}\in\mathbb{R}^{(d/2)\times 1}</annotation></semantics></math>.</p>
</div>
<div id="S2.SS4.p4" class="ltx_para">
<p id="S2.SS4.p4.1" class="ltx_p">We employ a two-stage training approach, where the main model and auxiliary routers are trained separately. First, we train the main model to convergence. Then, we train the auxiliary routers using binary cross-entropy loss, supervised by the ground-truth top-<math id="S2.SS4.p4.1.m1.1" class="ltx_Math" alttext="k" display="inline"><semantics id="S2.SS4.p4.1.m1.1a"><mi id="S2.SS4.p4.1.m1.1.1" xref="S2.SS4.p4.1.m1.1.1.cmml">k</mi><annotation-xml encoding="MathML-Content" id="S2.SS4.p4.1.m1.1b"><ci id="S2.SS4.p4.1.m1.1.1.cmml" xref="S2.SS4.p4.1.m1.1.1">𝑘</ci></annotation-xml><annotation encoding="application/x-tex" id="S2.SS4.p4.1.m1.1c">k</annotation></semantics></math> routing assignments computed over an entire batch. At inference time, the main routers are only used to generate the weight values, and tokens are selectively routed to an expert or layer based on thresholding on the auxiliary router score (&gt; 0.5).</p>
</div>
</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>Upcycling</h3>

<div id="S2.SS5.p1" class="ltx_para">
<p id="S2.SS5.p1.1" class="ltx_p">Training MoE architectures with learnable routers from scratch presents unique challenges in optimizing both the representation space and routing mechanism <cite class="ltx_cite ltx_citemacro_citep">(Xue et al., <a href="#bib.bib22" title="" class="ltx_ref">2024</a>)</cite>.
The insight we identified is that MoE routers are responsible for partitioning the representation space for each expert. However, this representation space is sub-optimal in the early stages of model training, leading to a sub-optimally trained routing function.</p>
</div>
<div id="S2.SS5.p2" class="ltx_para">
<p id="S2.SS5.p2.1" class="ltx_p">To address this limitation of router training, we propose an <span id="S2.SS5.p2.1.1" class="ltx_text ltx_font_bold">upcycling</span> approach, inspired by <cite class="ltx_cite ltx_citemacro_cite">Komatsuzaki et al. (<a href="#bib.bib18" title="" class="ltx_ref">2023</a>)</cite>.
Specifically, we begin by training an architecture consisting of 1 FFN expert per modality. After a predetermined number of steps, we upcycle this model by converting each modality-specific FFN into an expert-choice MoE module, initializing each expert with the expert trained from the first stage.
We reset the learning rate scheduler while preserving the data loader state from the previous stage, thereby ensuring the second stage training is exposed to refreshed data.</p>
</div>
<div id="S2.SS5.p3" class="ltx_para">
<p id="S2.SS5.p3.3" class="ltx_p">To promote expert specialization, we augment the MoE routing function with Gumbel noise <cite class="ltx_cite ltx_citemacro_citep">(Liu et al., <a href="#bib.bib23" title="" class="ltx_ref">2022b</a>; Geng et al., <a href="#bib.bib24" title="" class="ltx_ref">2020</a>)</cite>, allowing our router to differentiably sample experts. This is expressed in Equation <a href="#S2.E5" title="Equation 5 ‣ 2.5 Upcycling ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">5</span></a>:</p>
<table id="S2.E5" 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.E5.m1.2" class="ltx_Math" alttext="\text{Gumbel-Sigmoid}(x)=\text{Sigmoid}(x+G^{{}^{\prime}}-G^{{}^{\prime\prime}})" display="block"><semantics id="S2.E5.m1.2a"><mrow id="S2.E5.m1.2.2" xref="S2.E5.m1.2.2.cmml"><mrow id="S2.E5.m1.2.2.3" xref="S2.E5.m1.2.2.3.cmml"><mtext id="S2.E5.m1.2.2.3.2" xref="S2.E5.m1.2.2.3.2a.cmml">Gumbel-Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.E5.m1.2.2.3.1" xref="S2.E5.m1.2.2.3.1.cmml">​</mo><mrow id="S2.E5.m1.2.2.3.3.2" xref="S2.E5.m1.2.2.3.cmml"><mo stretchy="false" id="S2.E5.m1.2.2.3.3.2.1" xref="S2.E5.m1.2.2.3.cmml">(</mo><mi id="S2.E5.m1.1.1" xref="S2.E5.m1.1.1.cmml">x</mi><mo stretchy="false" id="S2.E5.m1.2.2.3.3.2.2" xref="S2.E5.m1.2.2.3.cmml">)</mo></mrow></mrow><mo id="S2.E5.m1.2.2.2" xref="S2.E5.m1.2.2.2.cmml">=</mo><mrow id="S2.E5.m1.2.2.1" xref="S2.E5.m1.2.2.1.cmml"><mtext id="S2.E5.m1.2.2.1.3" xref="S2.E5.m1.2.2.1.3a.cmml">Sigmoid</mtext><mo lspace="0em" rspace="0em" id="S2.E5.m1.2.2.1.2" xref="S2.E5.m1.2.2.1.2.cmml">​</mo><mrow id="S2.E5.m1.2.2.1.1.1" xref="S2.E5.m1.2.2.1.1.1.1.cmml"><mo stretchy="false" id="S2.E5.m1.2.2.1.1.1.2" xref="S2.E5.m1.2.2.1.1.1.1.cmml">(</mo><mrow id="S2.E5.m1.2.2.1.1.1.1" xref="S2.E5.m1.2.2.1.1.1.1.cmml"><mrow id="S2.E5.m1.2.2.1.1.1.1.2" xref="S2.E5.m1.2.2.1.1.1.1.2.cmml"><mi id="S2.E5.m1.2.2.1.1.1.1.2.2" xref="S2.E5.m1.2.2.1.1.1.1.2.2.cmml">x</mi><mo id="S2.E5.m1.2.2.1.1.1.1.2.1" xref="S2.E5.m1.2.2.1.1.1.1.2.1.cmml">+</mo><msup id="S2.E5.m1.2.2.1.1.1.1.2.3" xref="S2.E5.m1.2.2.1.1.1.1.2.3.cmml"><mi id="S2.E5.m1.2.2.1.1.1.1.2.3.2" xref="S2.E5.m1.2.2.1.1.1.1.2.3.2.cmml">G</mi><msup id="S2.E5.m1.2.2.1.1.1.1.2.3.3" xref="S2.E5.m1.2.2.1.1.1.1.2.3.3.cmml"><mi id="S2.E5.m1.2.2.1.1.1.1.2.3.3a" xref="S2.E5.m1.2.2.1.1.1.1.2.3.3.cmml"></mi><mo id="S2.E5.m1.2.2.1.1.1.1.2.3.3.1" xref="S2.E5.m1.2.2.1.1.1.1.2.3.3.1.cmml">′</mo></msup></msup></mrow><mo id="S2.E5.m1.2.2.1.1.1.1.1" xref="S2.E5.m1.2.2.1.1.1.1.1.cmml">−</mo><msup id="S2.E5.m1.2.2.1.1.1.1.3" xref="S2.E5.m1.2.2.1.1.1.1.3.cmml"><mi id="S2.E5.m1.2.2.1.1.1.1.3.2" xref="S2.E5.m1.2.2.1.1.1.1.3.2.cmml">G</mi><msup id="S2.E5.m1.2.2.1.1.1.1.3.3" xref="S2.E5.m1.2.2.1.1.1.1.3.3.cmml"><mi id="S2.E5.m1.2.2.1.1.1.1.3.3a" xref="S2.E5.m1.2.2.1.1.1.1.3.3.cmml"></mi><mo id="S2.E5.m1.2.2.1.1.1.1.3.3.1" xref="S2.E5.m1.2.2.1.1.1.1.3.3.1.cmml">′′</mo></msup></msup></mrow><mo stretchy="false" id="S2.E5.m1.2.2.1.1.1.3" xref="S2.E5.m1.2.2.1.1.1.1.cmml">)</mo></mrow></mrow></mrow><annotation-xml encoding="MathML-Content" id="S2.E5.m1.2b"><apply id="S2.E5.m1.2.2.cmml" xref="S2.E5.m1.2.2"><eq id="S2.E5.m1.2.2.2.cmml" xref="S2.E5.m1.2.2.2"></eq><apply id="S2.E5.m1.2.2.3.cmml" xref="S2.E5.m1.2.2.3"><times id="S2.E5.m1.2.2.3.1.cmml" xref="S2.E5.m1.2.2.3.1"></times><ci id="S2.E5.m1.2.2.3.2a.cmml" xref="S2.E5.m1.2.2.3.2"><mtext id="S2.E5.m1.2.2.3.2.cmml" xref="S2.E5.m1.2.2.3.2">Gumbel-Sigmoid</mtext></ci><ci id="S2.E5.m1.1.1.cmml" xref="S2.E5.m1.1.1">𝑥</ci></apply><apply id="S2.E5.m1.2.2.1.cmml" xref="S2.E5.m1.2.2.1"><times id="S2.E5.m1.2.2.1.2.cmml" xref="S2.E5.m1.2.2.1.2"></times><ci id="S2.E5.m1.2.2.1.3a.cmml" xref="S2.E5.m1.2.2.1.3"><mtext id="S2.E5.m1.2.2.1.3.cmml" xref="S2.E5.m1.2.2.1.3">Sigmoid</mtext></ci><apply id="S2.E5.m1.2.2.1.1.1.1.cmml" xref="S2.E5.m1.2.2.1.1.1"><minus id="S2.E5.m1.2.2.1.1.1.1.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.1"></minus><apply id="S2.E5.m1.2.2.1.1.1.1.2.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2"><plus id="S2.E5.m1.2.2.1.1.1.1.2.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.1"></plus><ci id="S2.E5.m1.2.2.1.1.1.1.2.2.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.2">𝑥</ci><apply id="S2.E5.m1.2.2.1.1.1.1.2.3.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.3"><csymbol cd="ambiguous" id="S2.E5.m1.2.2.1.1.1.1.2.3.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.3">superscript</csymbol><ci id="S2.E5.m1.2.2.1.1.1.1.2.3.2.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.3.2">𝐺</ci><apply id="S2.E5.m1.2.2.1.1.1.1.2.3.3.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.3.3"><ci id="S2.E5.m1.2.2.1.1.1.1.2.3.3.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.2.3.3.1">′</ci></apply></apply></apply><apply id="S2.E5.m1.2.2.1.1.1.1.3.cmml" xref="S2.E5.m1.2.2.1.1.1.1.3"><csymbol cd="ambiguous" id="S2.E5.m1.2.2.1.1.1.1.3.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.3">superscript</csymbol><ci id="S2.E5.m1.2.2.1.1.1.1.3.2.cmml" xref="S2.E5.m1.2.2.1.1.1.1.3.2">𝐺</ci><apply id="S2.E5.m1.2.2.1.1.1.1.3.3.cmml" xref="S2.E5.m1.2.2.1.1.1.1.3.3"><ci id="S2.E5.m1.2.2.1.1.1.1.3.3.1.cmml" xref="S2.E5.m1.2.2.1.1.1.1.3.3.1">′′</ci></apply></apply></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.E5.m1.2c">\text{Gumbel-Sigmoid}(x)=\text{Sigmoid}(x+G^{{}^{\prime}}-G^{{}^{\prime\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">(5)</span></td>
</tr></tbody>
</table>
<p id="S2.SS5.p3.2" class="ltx_p">where <math id="S2.SS5.p3.1.m1.1" class="ltx_Math" alttext="G^{{}^{\prime}}" display="inline"><semantics id="S2.SS5.p3.1.m1.1a"><msup id="S2.SS5.p3.1.m1.1.1" xref="S2.SS5.p3.1.m1.1.1.cmml"><mi id="S2.SS5.p3.1.m1.1.1.2" xref="S2.SS5.p3.1.m1.1.1.2.cmml">G</mi><msup id="S2.SS5.p3.1.m1.1.1.3" xref="S2.SS5.p3.1.m1.1.1.3.cmml"><mi id="S2.SS5.p3.1.m1.1.1.3a" xref="S2.SS5.p3.1.m1.1.1.3.cmml"></mi><mo id="S2.SS5.p3.1.m1.1.1.3.1" xref="S2.SS5.p3.1.m1.1.1.3.1.cmml">′</mo></msup></msup><annotation-xml encoding="MathML-Content" id="S2.SS5.p3.1.m1.1b"><apply id="S2.SS5.p3.1.m1.1.1.cmml" xref="S2.SS5.p3.1.m1.1.1"><csymbol cd="ambiguous" id="S2.SS5.p3.1.m1.1.1.1.cmml" xref="S2.SS5.p3.1.m1.1.1">superscript</csymbol><ci id="S2.SS5.p3.1.m1.1.1.2.cmml" xref="S2.SS5.p3.1.m1.1.1.2">𝐺</ci><apply id="S2.SS5.p3.1.m1.1.1.3.cmml" xref="S2.SS5.p3.1.m1.1.1.3"><ci id="S2.SS5.p3.1.m1.1.1.3.1.cmml" xref="S2.SS5.p3.1.m1.1.1.3.1">′</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p3.1.m1.1c">G^{{}^{\prime}}</annotation></semantics></math> and <math id="S2.SS5.p3.2.m2.1" class="ltx_Math" alttext="G^{{}^{\prime\prime}}" display="inline"><semantics id="S2.SS5.p3.2.m2.1a"><msup id="S2.SS5.p3.2.m2.1.1" xref="S2.SS5.p3.2.m2.1.1.cmml"><mi id="S2.SS5.p3.2.m2.1.1.2" xref="S2.SS5.p3.2.m2.1.1.2.cmml">G</mi><msup id="S2.SS5.p3.2.m2.1.1.3" xref="S2.SS5.p3.2.m2.1.1.3.cmml"><mi id="S2.SS5.p3.2.m2.1.1.3a" xref="S2.SS5.p3.2.m2.1.1.3.cmml"></mi><mo id="S2.SS5.p3.2.m2.1.1.3.1" xref="S2.SS5.p3.2.m2.1.1.3.1.cmml">′′</mo></msup></msup><annotation-xml encoding="MathML-Content" id="S2.SS5.p3.2.m2.1b"><apply id="S2.SS5.p3.2.m2.1.1.cmml" xref="S2.SS5.p3.2.m2.1.1"><csymbol cd="ambiguous" id="S2.SS5.p3.2.m2.1.1.1.cmml" xref="S2.SS5.p3.2.m2.1.1">superscript</csymbol><ci id="S2.SS5.p3.2.m2.1.1.2.cmml" xref="S2.SS5.p3.2.m2.1.1.2">𝐺</ci><apply id="S2.SS5.p3.2.m2.1.1.3.cmml" xref="S2.SS5.p3.2.m2.1.1.3"><ci id="S2.SS5.p3.2.m2.1.1.3.1.cmml" xref="S2.SS5.p3.2.m2.1.1.3.1">′′</ci></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S2.SS5.p3.2.m2.1c">G^{{}^{\prime\prime}}</annotation></semantics></math> are independent Gumbel noise samples.</p>
</div>
<div id="S2.SS5.p4" class="ltx_para">
<p id="S2.SS5.p4.1" class="ltx_p">The upcycling approach, combined with the Gumbel-Sigmoid technique, allows us to overcome the limitations of learned routers and achieve improved performance in our modality-aware sparse architecture.
In practice, we found that a short period of 10k-20k steps in the first-stage training significantly improves model training efficiency and stability, aligning with the findings of <cite class="ltx_cite ltx_citemacro_citet">Xue et al. (<a href="#bib.bib22" title="" class="ltx_ref">2024</a>)</cite>.</p>
</div>
</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>Efficiency Optimization</h2>

<div id="S3.p1" class="ltx_para">
<p id="S3.p1.1" class="ltx_p">To facilitate the distributed training of mixture of modality-aware experts (MoMa), we employ Fully Sharded Data Parallel (FSDP) <cite class="ltx_cite ltx_citemacro_citep">(Zhao et al., <a href="#bib.bib25" title="" class="ltx_ref">2023</a>)</cite>. However, this approach presents unique efficiency challenges compared to vanilla MoEs. This section discusses these challenges and our strategies to address them.</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>Load Balancing</h3>

<div id="S3.SS1.p1" class="ltx_para">
<p id="S3.SS1.p1.1" class="ltx_p">Without constraints, load imbalance can occur in our system because the ratio of text to image tokens can vary significantly across different GPUs and iterations. Imbalances can create a cascading straggler effect, delaying weight prefetching for subsequent layers and gradient releases from previous layers. This effectively bounds the training latency by the maximum time required to process text and image experts across all GPUs within a batch.</p>
</div>
<div id="S3.SS1.p2" class="ltx_para">
<p id="S3.SS1.p2.1" class="ltx_p">To mitigate these issues, we developed a balanced data mix that aligns the text-to-image data ratio with the expert ratio on each GPU. This approach ensures load balancing in expectation. While alternative rebalancing algorithms, such as token redistribution at each FFN layer, are possible, they may introduce additional communication overhead.</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>Efficient Expert Execution</h3>

<div id="S3.SS2.p1" class="ltx_para">
<p id="S3.SS2.p1.1" class="ltx_p">We explored several strategies to efficiently execute experts for different modalities. The first approach restricts to homogeneous experts across modalities and prohibits routing of text tokens to image experts and vice versa. This method allows processing of all tokens and modalities simultaneously, provided all experts share the same token count. Alternatively, we could enhance execution efficiency by employing block sparsity <cite class="ltx_cite ltx_citemacro_citep">(Gale et al., <a href="#bib.bib26" title="" class="ltx_ref">2023</a>)</cite>, which offers similar benefits to the first approach without requiring perfect expert token balance. Additionally, we considered running experts from different modalities sequentially when the number of modalities is limited. This approach allows better overlap of computation from previous modality experts with weight prefetching for subsequent modality experts, alleviating memory pressure. It also removes assumptions about expert load balance.</p>
</div>
<div id="S3.SS2.p2" class="ltx_para">
<p id="S3.SS2.p2.1" class="ltx_p">Given that our experiments involve a sufficiently large token count per GPU, hardware utilization is not a major concern even with multiple batched matrix multiplications. Thus, we find the sequential approach to be a clean and flexible choice for our experimental environment at the current scale.</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>Additional Optimizations</h3>

<div id="S3.SS3.p1" class="ltx_para">
<p id="S3.SS3.p1.1" class="ltx_p">We implemented several optimizations to further enhance throughput. These include generic optimizations such as gradient communication quantization and automatic GPU kernel fusion, as well as graph optimizations via <span id="S3.SS3.p1.1.1" class="ltx_text ltx_font_typewriter">torch.compile</span> <cite class="ltx_cite ltx_citemacro_citep">(Ansel et al., <a href="#bib.bib27" title="" class="ltx_ref">2024</a>)</cite>. Additionally, we developed MoMa-specific optimizations, including the reuse of modality token indices across different layers to minimize device synchronization between CPU and GPU<span id="footnote5" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">5</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">5</sup><span class="ltx_tag ltx_tag_note">5</span>This optimization is incompatible with MoD in its current form and was not used in our final experiments to ensure fair cross-comparison across model variations. However, it can be modified to eliminate device synchronization with permutation while tracking live token counts in each MoD layer.</span></span></span>. We also consolidated per-layer stats communication and moved these operations off the critical path of training.</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>

<section id="S4.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.1 </span>Setup</h3>

<div id="S4.SS1.p1" class="ltx_para">
<p id="S4.SS1.p1.4" class="ltx_p">We use the same pre-training dataset and preprocessing as <cite class="ltx_cite ltx_citemacro_citet">Chameleon Team (<a href="#bib.bib5" title="" class="ltx_ref">2024</a>)</cite>. To assess scaling performance, we train all models with over 1 trillion tokens. Unless specified otherwise, we employ a sequence length of 4096 tokens and a model parallel size of 1. Our training regimen includes a peak learning rate of <math id="S4.SS1.p1.1.m1.1" class="ltx_Math" alttext="1e-4" 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"><mrow id="S4.SS1.p1.1.m1.1.1.2" xref="S4.SS1.p1.1.m1.1.1.2.cmml"><mn id="S4.SS1.p1.1.m1.1.1.2.2" xref="S4.SS1.p1.1.m1.1.1.2.2.cmml">1</mn><mo lspace="0em" rspace="0em" id="S4.SS1.p1.1.m1.1.1.2.1" xref="S4.SS1.p1.1.m1.1.1.2.1.cmml">​</mo><mi id="S4.SS1.p1.1.m1.1.1.2.3" xref="S4.SS1.p1.1.m1.1.1.2.3.cmml">e</mi></mrow><mo 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">4</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"><minus id="S4.SS1.p1.1.m1.1.1.1.cmml" xref="S4.SS1.p1.1.m1.1.1.1"></minus><apply id="S4.SS1.p1.1.m1.1.1.2.cmml" xref="S4.SS1.p1.1.m1.1.1.2"><times id="S4.SS1.p1.1.m1.1.1.2.1.cmml" xref="S4.SS1.p1.1.m1.1.1.2.1"></times><cn type="integer" id="S4.SS1.p1.1.m1.1.1.2.2.cmml" xref="S4.SS1.p1.1.m1.1.1.2.2">1</cn><ci id="S4.SS1.p1.1.m1.1.1.2.3.cmml" xref="S4.SS1.p1.1.m1.1.1.2.3">𝑒</ci></apply><cn type="integer" id="S4.SS1.p1.1.m1.1.1.3.cmml" xref="S4.SS1.p1.1.m1.1.1.3">4</cn></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.1.m1.1c">1e-4</annotation></semantics></math>, a 4000-step warm-up period, and linear annealing of the learning rate to 1% of its peak value.
For all MoE architectures, we implement MoE in every layer, setting each expert’s training capacity <math id="S4.SS1.p1.2.m2.1" class="ltx_Math" alttext="k_{e}" display="inline"><semantics id="S4.SS1.p1.2.m2.1a"><msub id="S4.SS1.p1.2.m2.1.1" xref="S4.SS1.p1.2.m2.1.1.cmml"><mi id="S4.SS1.p1.2.m2.1.1.2" xref="S4.SS1.p1.2.m2.1.1.2.cmml">k</mi><mi id="S4.SS1.p1.2.m2.1.1.3" xref="S4.SS1.p1.2.m2.1.1.3.cmml">e</mi></msub><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.2.m2.1b"><apply id="S4.SS1.p1.2.m2.1.1.cmml" xref="S4.SS1.p1.2.m2.1.1"><csymbol cd="ambiguous" id="S4.SS1.p1.2.m2.1.1.1.cmml" xref="S4.SS1.p1.2.m2.1.1">subscript</csymbol><ci id="S4.SS1.p1.2.m2.1.1.2.cmml" xref="S4.SS1.p1.2.m2.1.1.2">𝑘</ci><ci id="S4.SS1.p1.2.m2.1.1.3.cmml" xref="S4.SS1.p1.2.m2.1.1.3">𝑒</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.2.m2.1c">k_{e}</annotation></semantics></math> to <math id="S4.SS1.p1.3.m3.1" class="ltx_Math" alttext="\frac{b^{\mathcal{M}}}{|\mathcal{E}^{\mathcal{M}}|}" display="inline"><semantics id="S4.SS1.p1.3.m3.1a"><mfrac id="S4.SS1.p1.3.m3.1.1" xref="S4.SS1.p1.3.m3.1.1.cmml"><msup id="S4.SS1.p1.3.m3.1.1.3" xref="S4.SS1.p1.3.m3.1.1.3.cmml"><mi id="S4.SS1.p1.3.m3.1.1.3.2" xref="S4.SS1.p1.3.m3.1.1.3.2.cmml">b</mi><mi class="ltx_font_mathcaligraphic" id="S4.SS1.p1.3.m3.1.1.3.3" xref="S4.SS1.p1.3.m3.1.1.3.3.cmml">ℳ</mi></msup><mrow id="S4.SS1.p1.3.m3.1.1.1.1" xref="S4.SS1.p1.3.m3.1.1.1.2.cmml"><mo stretchy="false" id="S4.SS1.p1.3.m3.1.1.1.1.2" xref="S4.SS1.p1.3.m3.1.1.1.2.1.cmml">|</mo><msup id="S4.SS1.p1.3.m3.1.1.1.1.1" xref="S4.SS1.p1.3.m3.1.1.1.1.1.cmml"><mi class="ltx_font_mathcaligraphic" id="S4.SS1.p1.3.m3.1.1.1.1.1.2" xref="S4.SS1.p1.3.m3.1.1.1.1.1.2.cmml">ℰ</mi><mi class="ltx_font_mathcaligraphic" id="S4.SS1.p1.3.m3.1.1.1.1.1.3" xref="S4.SS1.p1.3.m3.1.1.1.1.1.3.cmml">ℳ</mi></msup><mo stretchy="false" id="S4.SS1.p1.3.m3.1.1.1.1.3" xref="S4.SS1.p1.3.m3.1.1.1.2.1.cmml">|</mo></mrow></mfrac><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.3.m3.1b"><apply id="S4.SS1.p1.3.m3.1.1.cmml" xref="S4.SS1.p1.3.m3.1.1"><divide id="S4.SS1.p1.3.m3.1.1.2.cmml" xref="S4.SS1.p1.3.m3.1.1"></divide><apply id="S4.SS1.p1.3.m3.1.1.3.cmml" xref="S4.SS1.p1.3.m3.1.1.3"><csymbol cd="ambiguous" id="S4.SS1.p1.3.m3.1.1.3.1.cmml" xref="S4.SS1.p1.3.m3.1.1.3">superscript</csymbol><ci id="S4.SS1.p1.3.m3.1.1.3.2.cmml" xref="S4.SS1.p1.3.m3.1.1.3.2">𝑏</ci><ci id="S4.SS1.p1.3.m3.1.1.3.3.cmml" xref="S4.SS1.p1.3.m3.1.1.3.3">ℳ</ci></apply><apply id="S4.SS1.p1.3.m3.1.1.1.2.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1"><abs id="S4.SS1.p1.3.m3.1.1.1.2.1.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1.2"></abs><apply id="S4.SS1.p1.3.m3.1.1.1.1.1.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1.1"><csymbol cd="ambiguous" id="S4.SS1.p1.3.m3.1.1.1.1.1.1.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1.1">superscript</csymbol><ci id="S4.SS1.p1.3.m3.1.1.1.1.1.2.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1.1.2">ℰ</ci><ci id="S4.SS1.p1.3.m3.1.1.1.1.1.3.cmml" xref="S4.SS1.p1.3.m3.1.1.1.1.1.3">ℳ</ci></apply></apply></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.3.m3.1c">\frac{b^{\mathcal{M}}}{|\mathcal{E}^{\mathcal{M}}|}</annotation></semantics></math> to maintain FLOPs per token comparable to the base dense model. In MoD architectures, we implement MoD in alternating layers, beginning with layer 0, using a layer capacity factor <math id="S4.SS1.p1.4.m4.1" class="ltx_Math" alttext="c_{d}" display="inline"><semantics id="S4.SS1.p1.4.m4.1a"><msub id="S4.SS1.p1.4.m4.1.1" xref="S4.SS1.p1.4.m4.1.1.cmml"><mi id="S4.SS1.p1.4.m4.1.1.2" xref="S4.SS1.p1.4.m4.1.1.2.cmml">c</mi><mi id="S4.SS1.p1.4.m4.1.1.3" xref="S4.SS1.p1.4.m4.1.1.3.cmml">d</mi></msub><annotation-xml encoding="MathML-Content" id="S4.SS1.p1.4.m4.1b"><apply id="S4.SS1.p1.4.m4.1.1.cmml" xref="S4.SS1.p1.4.m4.1.1"><csymbol cd="ambiguous" id="S4.SS1.p1.4.m4.1.1.1.cmml" xref="S4.SS1.p1.4.m4.1.1">subscript</csymbol><ci id="S4.SS1.p1.4.m4.1.1.2.cmml" xref="S4.SS1.p1.4.m4.1.1.2">𝑐</ci><ci id="S4.SS1.p1.4.m4.1.1.3.cmml" xref="S4.SS1.p1.4.m4.1.1.3">𝑑</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS1.p1.4.m4.1c">c_{d}</annotation></semantics></math> of 25%. To achieve FLOPs parity with the base dense model, we increase the total layer count while maintaining a constant hidden dimension. Table <a href="#S4.T1" title="Table 1 ‣ 4.1 Setup ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">1</span></a> provides detailed configurations for our dense and sparse models. Additional pre-training details are available in Appendix <a href="#A1.SS1" title="A.1 Pre-training ‣ Appendix A Additional Implementation Details ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">A.1</span></a>.</p>
</div>
<div id="S4.SS1.p2" class="ltx_para">
<p id="S4.SS1.p2.1" class="ltx_p">For model comparison, we report training losses. Given that our 1-trillion-token training budget covers less than one epoch of our extensive pre-training data, we use training loss as a proxy for validation performance. Our use of expert-choice routing in both MoE and MoD modules introduces a caveat: the training loss calculation compromises causality, as token selection considers the top proportion from a batch, including future tokens. We address this in §<a href="#S4.SS6" title="4.6 Inference-time Performance ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.6</span></a> by reporting test-time performances using auxiliary routers, demonstrating result generalization to validation settings and causal scenarios.</p>
</div>
<figure id="S4.T1" class="ltx_table">
<table id="S4.T1.12.12" class="ltx_tabular ltx_centering ltx_align_middle">
<tr id="S4.T1.12.12.13" class="ltx_tr">
<td id="S4.T1.12.12.13.1" class="ltx_td ltx_align_center ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;" colspan="5"><span id="S4.T1.12.12.13.1.1" class="ltx_text" style="font-size:90%;">Dense</span></td>
<td id="S4.T1.12.12.13.2" class="ltx_td ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.13.3" class="ltx_td ltx_align_center ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;" colspan="4"><span id="S4.T1.12.12.13.3.1" class="ltx_text" style="font-size:90%;">MoE</span></td>
<td id="S4.T1.12.12.13.4" class="ltx_td ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.13.5" class="ltx_td ltx_align_center ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;" colspan="6"><span id="S4.T1.12.12.13.5.1" class="ltx_text" style="font-size:90%;">MoD</span></td>
<td id="S4.T1.12.12.13.6" class="ltx_td ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.13.7" class="ltx_td ltx_align_center ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.13.7.1" class="ltx_text" style="font-size:90%;">MoDE</span></td>
<td id="S4.T1.12.12.13.8" class="ltx_td ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.13.9" class="ltx_td ltx_border_tt" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
</tr>
<tr id="S4.T1.12.12.12" class="ltx_tr">
<td id="S4.T1.12.12.12.13" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><em id="S4.T1.12.12.12.13.1" class="ltx_emph ltx_font_italic" style="font-size:90%;">params</em></td>
<td id="S4.T1.1.1.1.1" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.1.1.1.1.m1.1" class="ltx_Math" alttext="l" display="inline"><semantics id="S4.T1.1.1.1.1.m1.1a"><mi mathsize="90%" id="S4.T1.1.1.1.1.m1.1.1" xref="S4.T1.1.1.1.1.m1.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="S4.T1.1.1.1.1.m1.1b"><ci id="S4.T1.1.1.1.1.m1.1.1.cmml" xref="S4.T1.1.1.1.1.m1.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.1.1.1.1.m1.1c">l</annotation></semantics></math></td>
<td id="S4.T1.2.2.2.2" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.2.2.2.2.m1.1" class="ltx_Math" alttext="d" display="inline"><semantics id="S4.T1.2.2.2.2.m1.1a"><mi mathsize="90%" id="S4.T1.2.2.2.2.m1.1.1" xref="S4.T1.2.2.2.2.m1.1.1.cmml">d</mi><annotation-xml encoding="MathML-Content" id="S4.T1.2.2.2.2.m1.1b"><ci id="S4.T1.2.2.2.2.m1.1.1.cmml" xref="S4.T1.2.2.2.2.m1.1.1">𝑑</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.2.2.2.2.m1.1c">d</annotation></semantics></math></td>
<td id="S4.T1.3.3.3.3" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.3.3.3.3.m1.1" class="ltx_Math" alttext="ffn" display="inline"><semantics id="S4.T1.3.3.3.3.m1.1a"><mrow id="S4.T1.3.3.3.3.m1.1.1" xref="S4.T1.3.3.3.3.m1.1.1.cmml"><mi mathsize="90%" id="S4.T1.3.3.3.3.m1.1.1.2" xref="S4.T1.3.3.3.3.m1.1.1.2.cmml">f</mi><mo lspace="0em" rspace="0em" id="S4.T1.3.3.3.3.m1.1.1.1" xref="S4.T1.3.3.3.3.m1.1.1.1.cmml">​</mo><mi mathsize="90%" id="S4.T1.3.3.3.3.m1.1.1.3" xref="S4.T1.3.3.3.3.m1.1.1.3.cmml">f</mi><mo lspace="0em" rspace="0em" id="S4.T1.3.3.3.3.m1.1.1.1a" xref="S4.T1.3.3.3.3.m1.1.1.1.cmml">​</mo><mi mathsize="90%" id="S4.T1.3.3.3.3.m1.1.1.4" xref="S4.T1.3.3.3.3.m1.1.1.4.cmml">n</mi></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.3.3.3.3.m1.1b"><apply id="S4.T1.3.3.3.3.m1.1.1.cmml" xref="S4.T1.3.3.3.3.m1.1.1"><times id="S4.T1.3.3.3.3.m1.1.1.1.cmml" xref="S4.T1.3.3.3.3.m1.1.1.1"></times><ci id="S4.T1.3.3.3.3.m1.1.1.2.cmml" xref="S4.T1.3.3.3.3.m1.1.1.2">𝑓</ci><ci id="S4.T1.3.3.3.3.m1.1.1.3.cmml" xref="S4.T1.3.3.3.3.m1.1.1.3">𝑓</ci><ci id="S4.T1.3.3.3.3.m1.1.1.4.cmml" xref="S4.T1.3.3.3.3.m1.1.1.4">𝑛</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.3.3.3.3.m1.1c">ffn</annotation></semantics></math></td>
<td id="S4.T1.4.4.4.4" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.4.4.4.4.m1.1" class="ltx_Math" alttext="h" display="inline"><semantics id="S4.T1.4.4.4.4.m1.1a"><mi mathsize="90%" id="S4.T1.4.4.4.4.m1.1.1" xref="S4.T1.4.4.4.4.m1.1.1.cmml">h</mi><annotation-xml encoding="MathML-Content" id="S4.T1.4.4.4.4.m1.1b"><ci id="S4.T1.4.4.4.4.m1.1.1.cmml" xref="S4.T1.4.4.4.4.m1.1.1">ℎ</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.4.4.4.4.m1.1c">h</annotation></semantics></math></td>
<td id="S4.T1.12.12.12.14" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.12.15" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><em id="S4.T1.12.12.12.15.1" class="ltx_emph ltx_font_italic" style="font-size:90%;">params</em></td>
<td id="S4.T1.5.5.5.5" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.5.5.5.5.m1.1" class="ltx_Math" alttext="e" display="inline"><semantics id="S4.T1.5.5.5.5.m1.1a"><mi mathsize="90%" id="S4.T1.5.5.5.5.m1.1.1" xref="S4.T1.5.5.5.5.m1.1.1.cmml">e</mi><annotation-xml encoding="MathML-Content" id="S4.T1.5.5.5.5.m1.1b"><ci id="S4.T1.5.5.5.5.m1.1.1.cmml" xref="S4.T1.5.5.5.5.m1.1.1">𝑒</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.5.5.5.5.m1.1c">e</annotation></semantics></math></td>
<td id="S4.T1.6.6.6.6" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.6.6.6.6.m1.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S4.T1.6.6.6.6.m1.1a"><mi mathsize="90%" id="S4.T1.6.6.6.6.m1.1.1" xref="S4.T1.6.6.6.6.m1.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S4.T1.6.6.6.6.m1.1b"><ci id="S4.T1.6.6.6.6.m1.1.1.cmml" xref="S4.T1.6.6.6.6.m1.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.6.6.6.6.m1.1c">i</annotation></semantics></math></td>
<td id="S4.T1.7.7.7.7" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.7.7.7.7.m1.1" class="ltx_Math" alttext="c_{e}" display="inline"><semantics id="S4.T1.7.7.7.7.m1.1a"><msub id="S4.T1.7.7.7.7.m1.1.1" xref="S4.T1.7.7.7.7.m1.1.1.cmml"><mi mathsize="90%" id="S4.T1.7.7.7.7.m1.1.1.2" xref="S4.T1.7.7.7.7.m1.1.1.2.cmml">c</mi><mi mathsize="90%" id="S4.T1.7.7.7.7.m1.1.1.3" xref="S4.T1.7.7.7.7.m1.1.1.3.cmml">e</mi></msub><annotation-xml encoding="MathML-Content" id="S4.T1.7.7.7.7.m1.1b"><apply id="S4.T1.7.7.7.7.m1.1.1.cmml" xref="S4.T1.7.7.7.7.m1.1.1"><csymbol cd="ambiguous" id="S4.T1.7.7.7.7.m1.1.1.1.cmml" xref="S4.T1.7.7.7.7.m1.1.1">subscript</csymbol><ci id="S4.T1.7.7.7.7.m1.1.1.2.cmml" xref="S4.T1.7.7.7.7.m1.1.1.2">𝑐</ci><ci id="S4.T1.7.7.7.7.m1.1.1.3.cmml" xref="S4.T1.7.7.7.7.m1.1.1.3">𝑒</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.7.7.7.7.m1.1c">c_{e}</annotation></semantics></math></td>
<td id="S4.T1.12.12.12.16" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.12.17" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><em id="S4.T1.12.12.12.17.1" class="ltx_emph ltx_font_italic" style="font-size:90%;">params</em></td>
<td id="S4.T1.8.8.8.8" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.8.8.8.8.m1.1" class="ltx_Math" alttext="l" display="inline"><semantics id="S4.T1.8.8.8.8.m1.1a"><mi mathsize="90%" id="S4.T1.8.8.8.8.m1.1.1" xref="S4.T1.8.8.8.8.m1.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="S4.T1.8.8.8.8.m1.1b"><ci id="S4.T1.8.8.8.8.m1.1.1.cmml" xref="S4.T1.8.8.8.8.m1.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.8.8.8.8.m1.1c">l</annotation></semantics></math></td>
<td id="S4.T1.9.9.9.9" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.9.9.9.9.m1.1" class="ltx_Math" alttext="d" display="inline"><semantics id="S4.T1.9.9.9.9.m1.1a"><mi mathsize="90%" id="S4.T1.9.9.9.9.m1.1.1" xref="S4.T1.9.9.9.9.m1.1.1.cmml">d</mi><annotation-xml encoding="MathML-Content" id="S4.T1.9.9.9.9.m1.1b"><ci id="S4.T1.9.9.9.9.m1.1.1.cmml" xref="S4.T1.9.9.9.9.m1.1.1">𝑑</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.9.9.9.9.m1.1c">d</annotation></semantics></math></td>
<td id="S4.T1.10.10.10.10" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.10.10.10.10.m1.1" class="ltx_Math" alttext="h" display="inline"><semantics id="S4.T1.10.10.10.10.m1.1a"><mi mathsize="90%" id="S4.T1.10.10.10.10.m1.1.1" xref="S4.T1.10.10.10.10.m1.1.1.cmml">h</mi><annotation-xml encoding="MathML-Content" id="S4.T1.10.10.10.10.m1.1b"><ci id="S4.T1.10.10.10.10.m1.1.1.cmml" xref="S4.T1.10.10.10.10.m1.1.1">ℎ</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.10.10.10.10.m1.1c">h</annotation></semantics></math></td>
<td id="S4.T1.11.11.11.11" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.11.11.11.11.m1.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S4.T1.11.11.11.11.m1.1a"><mi mathsize="90%" id="S4.T1.11.11.11.11.m1.1.1" xref="S4.T1.11.11.11.11.m1.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S4.T1.11.11.11.11.m1.1b"><ci id="S4.T1.11.11.11.11.m1.1.1.cmml" xref="S4.T1.11.11.11.11.m1.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.11.11.11.11.m1.1c">i</annotation></semantics></math></td>
<td id="S4.T1.12.12.12.12" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><math id="S4.T1.12.12.12.12.m1.1" class="ltx_Math" alttext="c_{d}" display="inline"><semantics id="S4.T1.12.12.12.12.m1.1a"><msub id="S4.T1.12.12.12.12.m1.1.1" xref="S4.T1.12.12.12.12.m1.1.1.cmml"><mi mathsize="90%" id="S4.T1.12.12.12.12.m1.1.1.2" xref="S4.T1.12.12.12.12.m1.1.1.2.cmml">c</mi><mi mathsize="90%" id="S4.T1.12.12.12.12.m1.1.1.3" xref="S4.T1.12.12.12.12.m1.1.1.3.cmml">d</mi></msub><annotation-xml encoding="MathML-Content" id="S4.T1.12.12.12.12.m1.1b"><apply id="S4.T1.12.12.12.12.m1.1.1.cmml" xref="S4.T1.12.12.12.12.m1.1.1"><csymbol cd="ambiguous" id="S4.T1.12.12.12.12.m1.1.1.1.cmml" xref="S4.T1.12.12.12.12.m1.1.1">subscript</csymbol><ci id="S4.T1.12.12.12.12.m1.1.1.2.cmml" xref="S4.T1.12.12.12.12.m1.1.1.2">𝑐</ci><ci id="S4.T1.12.12.12.12.m1.1.1.3.cmml" xref="S4.T1.12.12.12.12.m1.1.1.3">𝑑</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.12.12.12.12.m1.1c">c_{d}</annotation></semantics></math></td>
<td id="S4.T1.12.12.12.18" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.12.19" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><em id="S4.T1.12.12.12.19.1" class="ltx_emph ltx_font_italic" style="font-size:90%;">params</em></td>
<td id="S4.T1.12.12.12.20" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.12.21" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
</tr>
<tr id="S4.T1.12.12.14" class="ltx_tr">
<td id="S4.T1.12.12.14.1" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.1.1" class="ltx_text" style="font-size:90%;">90M</span></td>
<td id="S4.T1.12.12.14.2" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.2.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.14.3" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.3.1" class="ltx_text" style="font-size:90%;">512</span></td>
<td id="S4.T1.12.12.14.4" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.4.1" class="ltx_text" style="font-size:90%;">2048</span></td>
<td id="S4.T1.12.12.14.5" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.5.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.14.6" class="ltx_td ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.14.7" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.7.1" class="ltx_text" style="font-size:90%;">210M</span></td>
<td id="S4.T1.12.12.14.8" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.8.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.14.9" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.9.1" class="ltx_text" style="font-size:90%;">1</span></td>
<td id="S4.T1.12.12.14.10" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.10.1" class="ltx_text" style="font-size:90%;">0.125</span></td>
<td id="S4.T1.12.12.14.11" class="ltx_td ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.14.12" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.12.1" class="ltx_text" style="font-size:90%;">110M</span></td>
<td id="S4.T1.12.12.14.13" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.13.1" class="ltx_text" style="font-size:90%;">14</span></td>
<td id="S4.T1.12.12.14.14" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.14.1" class="ltx_text" style="font-size:90%;">512</span></td>
<td id="S4.T1.12.12.14.15" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.15.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.14.16" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.16.1" class="ltx_text" style="font-size:90%;">2</span></td>
<td id="S4.T1.12.12.14.17" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.17.1" class="ltx_text" style="font-size:90%;">0.25</span></td>
<td id="S4.T1.12.12.14.18" class="ltx_td ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.14.19" class="ltx_td ltx_align_center ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.14.19.1" class="ltx_text" style="font-size:90%;">317M</span></td>
<td id="S4.T1.12.12.14.20" class="ltx_td ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.14.21" class="ltx_td ltx_border_t" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
</tr>
<tr id="S4.T1.12.12.15" class="ltx_tr">
<td id="S4.T1.12.12.15.1" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.1.1" class="ltx_text" style="font-size:90%;">435M</span></td>
<td id="S4.T1.12.12.15.2" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.2.1" class="ltx_text" style="font-size:90%;">24</span></td>
<td id="S4.T1.12.12.15.3" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.3.1" class="ltx_text" style="font-size:90%;">1024</span></td>
<td id="S4.T1.12.12.15.4" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.4.1" class="ltx_text" style="font-size:90%;">4096</span></td>
<td id="S4.T1.12.12.15.5" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.5.1" class="ltx_text" style="font-size:90%;">16</span></td>
<td id="S4.T1.12.12.15.6" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.15.7" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.7.1" class="ltx_text" style="font-size:90%;">1.9B</span></td>
<td id="S4.T1.12.12.15.8" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.8.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.15.9" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.9.1" class="ltx_text" style="font-size:90%;">1</span></td>
<td id="S4.T1.12.12.15.10" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.10.1" class="ltx_text" style="font-size:90%;">0.125</span></td>
<td id="S4.T1.12.12.15.11" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.15.12" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.12.1" class="ltx_text" style="font-size:90%;">635M</span></td>
<td id="S4.T1.12.12.15.13" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.13.1" class="ltx_text" style="font-size:90%;">40</span></td>
<td id="S4.T1.12.12.15.14" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.14.1" class="ltx_text" style="font-size:90%;">1024</span></td>
<td id="S4.T1.12.12.15.15" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.15.1" class="ltx_text" style="font-size:90%;">16</span></td>
<td id="S4.T1.12.12.15.16" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.16.1" class="ltx_text" style="font-size:90%;">2</span></td>
<td id="S4.T1.12.12.15.17" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.17.1" class="ltx_text" style="font-size:90%;">0.25</span></td>
<td id="S4.T1.12.12.15.18" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.15.19" class="ltx_td ltx_align_center" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.15.19.1" class="ltx_text" style="font-size:90%;">3B</span></td>
<td id="S4.T1.12.12.15.20" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.15.21" class="ltx_td" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
</tr>
<tr id="S4.T1.12.12.16" class="ltx_tr">
<td id="S4.T1.12.12.16.1" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.1.1" class="ltx_text" style="font-size:90%;">1.4B</span></td>
<td id="S4.T1.12.12.16.2" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.2.1" class="ltx_text" style="font-size:90%;">24</span></td>
<td id="S4.T1.12.12.16.3" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.3.1" class="ltx_text" style="font-size:90%;">2048</span></td>
<td id="S4.T1.12.12.16.4" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.4.1" class="ltx_text" style="font-size:90%;">8192</span></td>
<td id="S4.T1.12.12.16.5" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.5.1" class="ltx_text" style="font-size:90%;">16</span></td>
<td id="S4.T1.12.12.16.6" class="ltx_td ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.16.7" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.7.1" class="ltx_text" style="font-size:90%;">7.1B</span></td>
<td id="S4.T1.12.12.16.8" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.8.1" class="ltx_text" style="font-size:90%;">8</span></td>
<td id="S4.T1.12.12.16.9" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.9.1" class="ltx_text" style="font-size:90%;">1</span></td>
<td id="S4.T1.12.12.16.10" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.10.1" class="ltx_text" style="font-size:90%;">0.125</span></td>
<td id="S4.T1.12.12.16.11" class="ltx_td ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.16.12" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.12.1" class="ltx_text" style="font-size:90%;">2.3B</span></td>
<td id="S4.T1.12.12.16.13" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.13.1" class="ltx_text" style="font-size:90%;">32</span></td>
<td id="S4.T1.12.12.16.14" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.14.1" class="ltx_text" style="font-size:90%;">2304</span></td>
<td id="S4.T1.12.12.16.15" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.15.1" class="ltx_text" style="font-size:90%;">18</span></td>
<td id="S4.T1.12.12.16.16" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.16.1" class="ltx_text" style="font-size:90%;">2</span></td>
<td id="S4.T1.12.12.16.17" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.17.1" class="ltx_text" style="font-size:90%;">0.25</span></td>
<td id="S4.T1.12.12.16.18" class="ltx_td ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.16.19" class="ltx_td ltx_align_center ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"><span id="S4.T1.12.12.16.19.1" class="ltx_text" style="font-size:90%;">12B</span></td>
<td id="S4.T1.12.12.16.20" class="ltx_td ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
<td id="S4.T1.12.12.16.21" class="ltx_td ltx_border_bb" style="padding-left:3.5pt;padding-right:3.5pt;"></td>
</tr>
</table>
<figcaption class="ltx_caption"><span class="ltx_tag ltx_tag_table"><span id="S4.T1.29.8.1" class="ltx_text" style="font-size:90%;">Table 1</span>: </span><span id="S4.T1.26.7" class="ltx_text" style="font-size:90%;">Specifications of dense and sparse architectures used in our experiments. Architectures listed in the same row are compute-matched with the same amount of active FLOPs per token. <em id="S4.T1.26.7.1" class="ltx_emph ltx_font_italic">params</em>: total number of parameters; <math id="S4.T1.20.1.m1.1" class="ltx_Math" alttext="l" display="inline"><semantics id="S4.T1.20.1.m1.1b"><mi id="S4.T1.20.1.m1.1.1" xref="S4.T1.20.1.m1.1.1.cmml">l</mi><annotation-xml encoding="MathML-Content" id="S4.T1.20.1.m1.1c"><ci id="S4.T1.20.1.m1.1.1.cmml" xref="S4.T1.20.1.m1.1.1">𝑙</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.20.1.m1.1d">l</annotation></semantics></math>: number of layers; <math id="S4.T1.21.2.m2.1" class="ltx_Math" alttext="d" display="inline"><semantics id="S4.T1.21.2.m2.1b"><mi id="S4.T1.21.2.m2.1.1" xref="S4.T1.21.2.m2.1.1.cmml">d</mi><annotation-xml encoding="MathML-Content" id="S4.T1.21.2.m2.1c"><ci id="S4.T1.21.2.m2.1.1.cmml" xref="S4.T1.21.2.m2.1.1">𝑑</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.21.2.m2.1d">d</annotation></semantics></math>: transformer hidden dimension; <math id="S4.T1.22.3.m3.1" class="ltx_Math" alttext="ffn" display="inline"><semantics id="S4.T1.22.3.m3.1b"><mrow id="S4.T1.22.3.m3.1.1" xref="S4.T1.22.3.m3.1.1.cmml"><mi id="S4.T1.22.3.m3.1.1.2" xref="S4.T1.22.3.m3.1.1.2.cmml">f</mi><mo lspace="0em" rspace="0em" id="S4.T1.22.3.m3.1.1.1" xref="S4.T1.22.3.m3.1.1.1.cmml">​</mo><mi id="S4.T1.22.3.m3.1.1.3" xref="S4.T1.22.3.m3.1.1.3.cmml">f</mi><mo lspace="0em" rspace="0em" id="S4.T1.22.3.m3.1.1.1b" xref="S4.T1.22.3.m3.1.1.1.cmml">​</mo><mi id="S4.T1.22.3.m3.1.1.4" xref="S4.T1.22.3.m3.1.1.4.cmml">n</mi></mrow><annotation-xml encoding="MathML-Content" id="S4.T1.22.3.m3.1c"><apply id="S4.T1.22.3.m3.1.1.cmml" xref="S4.T1.22.3.m3.1.1"><times id="S4.T1.22.3.m3.1.1.1.cmml" xref="S4.T1.22.3.m3.1.1.1"></times><ci id="S4.T1.22.3.m3.1.1.2.cmml" xref="S4.T1.22.3.m3.1.1.2">𝑓</ci><ci id="S4.T1.22.3.m3.1.1.3.cmml" xref="S4.T1.22.3.m3.1.1.3">𝑓</ci><ci id="S4.T1.22.3.m3.1.1.4.cmml" xref="S4.T1.22.3.m3.1.1.4">𝑛</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.22.3.m3.1d">ffn</annotation></semantics></math>: feed-forward module hidden dimension; <math id="S4.T1.23.4.m4.1" class="ltx_Math" alttext="h" display="inline"><semantics id="S4.T1.23.4.m4.1b"><mi id="S4.T1.23.4.m4.1.1" xref="S4.T1.23.4.m4.1.1.cmml">h</mi><annotation-xml encoding="MathML-Content" id="S4.T1.23.4.m4.1c"><ci id="S4.T1.23.4.m4.1.1.cmml" xref="S4.T1.23.4.m4.1.1">ℎ</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.23.4.m4.1d">h</annotation></semantics></math>: number of attention heads; <math id="S4.T1.24.5.m5.1" class="ltx_Math" alttext="e" display="inline"><semantics id="S4.T1.24.5.m5.1b"><mi id="S4.T1.24.5.m5.1.1" xref="S4.T1.24.5.m5.1.1.cmml">e</mi><annotation-xml encoding="MathML-Content" id="S4.T1.24.5.m5.1c"><ci id="S4.T1.24.5.m5.1.1.cmml" xref="S4.T1.24.5.m5.1.1">𝑒</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.24.5.m5.1d">e</annotation></semantics></math>: number of experts; <math id="S4.T1.25.6.m6.1" class="ltx_Math" alttext="i" display="inline"><semantics id="S4.T1.25.6.m6.1b"><mi id="S4.T1.25.6.m6.1.1" xref="S4.T1.25.6.m6.1.1.cmml">i</mi><annotation-xml encoding="MathML-Content" id="S4.T1.25.6.m6.1c"><ci id="S4.T1.25.6.m6.1.1.cmml" xref="S4.T1.25.6.m6.1.1">𝑖</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.25.6.m6.1d">i</annotation></semantics></math>: MoE/D layer interval; <math id="S4.T1.26.7.m7.1" class="ltx_Math" alttext="c" display="inline"><semantics id="S4.T1.26.7.m7.1b"><mi id="S4.T1.26.7.m7.1.1" xref="S4.T1.26.7.m7.1.1.cmml">c</mi><annotation-xml encoding="MathML-Content" id="S4.T1.26.7.m7.1c"><ci id="S4.T1.26.7.m7.1.1.cmml" xref="S4.T1.26.7.m7.1.1">𝑐</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T1.26.7.m7.1d">c</annotation></semantics></math>: MoE expert capacity/MoD layer capacity.</span></figcaption>
</figure>
<figure id="S4.F3" class="ltx_figure"><img src="/html/2407.21770/assets/figures/main_scaling_curves.png" id="S4.F3.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="598" height="289" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F3.5.1.1" class="ltx_text" style="font-size:90%;">Figure 3</span>: </span><span id="S4.F3.6.2" class="ltx_text" style="font-size:90%;">Scaling of performance with compute. We consider three dense model sizes – 90M, 435M and 1.4B parameters – and their compute-matched sparse variations. We highlight three configurations of mixture-of-experts (MoE): 1) using 1 text expert and 1 image expert (<span id="S4.F3.6.2.1" class="ltx_text ltx_font_typewriter">moe_1t1i</span>), 2) using 4 text experts and 4 image experts (<span id="S4.F3.6.2.2" class="ltx_text ltx_font_typewriter">moe_4t4i</span>), and 3) using 8 mixed-modal experts (<span id="S4.F3.6.2.3" class="ltx_text ltx_font_typewriter">moe_8x</span>) in addition to the dense baseline. All models were pre-trained and annealed to 1 trillion tokens.
When controlling for FLOPs per token, all sparse architectures demonstrate better scaling properties compared to the dense baseline. In particular, architectures with modality-specific experts consistently outperforms architectures with only mixed-modal parameters, with particular large benefits shown on the image modality.</span></figcaption>
</figure>
</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>Scaling of Performance with Compute</h3>

<div id="S4.SS2.p1" class="ltx_para">
<p id="S4.SS2.p1.2" class="ltx_p">We present the scaling performance of various models across multiple compute levels, with FLOPs matched to three dense model sizes: 90M, 435M, and 1.4B parameters. We report two key metrics: (1) training loss and (2) pre-training speed-up factor <math id="S4.SS2.p1.1.m1.1" class="ltx_Math" alttext="\eta" display="inline"><semantics id="S4.SS2.p1.1.m1.1a"><mi id="S4.SS2.p1.1.m1.1.1" xref="S4.SS2.p1.1.m1.1.1.cmml">η</mi><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.1.m1.1b"><ci id="S4.SS2.p1.1.m1.1.1.cmml" xref="S4.SS2.p1.1.m1.1.1">𝜂</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.1.m1.1c">\eta</annotation></semantics></math>, which indicates that a sparse model can match the pre-training loss of the iso-FLOPs dense baseline using only <math id="S4.SS2.p1.2.m2.1" class="ltx_Math" alttext="1/\eta" display="inline"><semantics id="S4.SS2.p1.2.m2.1a"><mrow id="S4.SS2.p1.2.m2.1.1" xref="S4.SS2.p1.2.m2.1.1.cmml"><mn id="S4.SS2.p1.2.m2.1.1.2" xref="S4.SS2.p1.2.m2.1.1.2.cmml">1</mn><mo id="S4.SS2.p1.2.m2.1.1.1" xref="S4.SS2.p1.2.m2.1.1.1.cmml">/</mo><mi id="S4.SS2.p1.2.m2.1.1.3" xref="S4.SS2.p1.2.m2.1.1.3.cmml">η</mi></mrow><annotation-xml encoding="MathML-Content" id="S4.SS2.p1.2.m2.1b"><apply id="S4.SS2.p1.2.m2.1.1.cmml" xref="S4.SS2.p1.2.m2.1.1"><divide id="S4.SS2.p1.2.m2.1.1.1.cmml" xref="S4.SS2.p1.2.m2.1.1.1"></divide><cn type="integer" id="S4.SS2.p1.2.m2.1.1.2.cmml" xref="S4.SS2.p1.2.m2.1.1.2">1</cn><ci id="S4.SS2.p1.2.m2.1.1.3.cmml" xref="S4.SS2.p1.2.m2.1.1.3">𝜂</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS2.p1.2.m2.1c">1/\eta</annotation></semantics></math> of total FLOPs<span id="footnote6" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">6</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">6</sup><span class="ltx_tag ltx_tag_note">6</span>Our definition of <math id="footnote6.m1.1" class="ltx_Math" alttext="\eta" display="inline"><semantics id="footnote6.m1.1b"><mi id="footnote6.m1.1.1" xref="footnote6.m1.1.1.cmml">η</mi><annotation-xml encoding="MathML-Content" id="footnote6.m1.1c"><ci id="footnote6.m1.1.1.cmml" xref="footnote6.m1.1.1">𝜂</ci></annotation-xml><annotation encoding="application/x-tex" id="footnote6.m1.1d">\eta</annotation></semantics></math> is analogous to the speed-up factor proposed by <cite class="ltx_cite ltx_citemacro_cite">Artetxe et al. (<a href="#bib.bib28" title="" class="ltx_ref">2021</a>)</cite>, but is defined in terms of pre-training loss whereas the original definition uses validation perplexity.</span></span></span>.</p>
</div>
<figure id="S4.F4" class="ltx_figure"><img src="/html/2407.21770/assets/figures/combine_depth_and_width_scaling.png" id="S4.F4.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="598" height="290" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F4.2.1.1" class="ltx_text" style="font-size:90%;">Figure 4</span>: </span><span id="S4.F4.3.2" class="ltx_text" style="font-size:90%;">Effect of combining with mixture-of-depths (MoD). We consider the sparse model variations compute-matched to three dense model sizes – 90M, 435M and 1.4B parameters. Combining MoD with MoE in general lead to better pre-training loss convergence across model sizes and configurations.</span></figcaption>
</figure>
<section id="S4.SS2.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Modality untying.</h4>

<div id="S4.SS2.SSS0.Px1.p1" class="ltx_para">
<p id="S4.SS2.SSS0.Px1.p1.1" class="ltx_p">Introducing modality-specific expert groups improves pre-training efficiency across different scales, with particularly large benefits for the image modality. As shown in Figure <a href="#S4.F3" title="Figure 3 ‣ 4.1 Setup ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>, the <span id="S4.SS2.SSS0.Px1.p1.1.1" class="ltx_text ltx_font_typewriter">moe_1t1i</span> configuration, using 1 image expert and 1 text expert, significantly outperforms the dense baseline. The image loss of <span id="S4.SS2.SSS0.Px1.p1.1.2" class="ltx_text ltx_font_typewriter">moe_1t1i</span> nearly matches that of the regular expert-choice MoE model using 8 mix-modal experts (<span id="S4.SS2.SSS0.Px1.p1.1.3" class="ltx_text ltx_font_typewriter">moe_8x</span>), whereas the text loss remains substantially higher, indicating parameter untying has a disproportionate impact on the image modality. Scaling the number of experts within each modality group further improve the model performance. The <span id="S4.SS2.SSS0.Px1.p1.1.4" class="ltx_text ltx_font_typewriter">moe_4t4i</span> configuration, using 4 text experts and 4 image experts, consistently outperforms <span id="S4.SS2.SSS0.Px1.p1.1.5" class="ltx_text ltx_font_typewriter">moe_8x</span> across different scales, with a substantial margin in image loss. However, this comparison reveals a slight regression in the text modality, suggesting that processing text modality across more experts may be beneficial.</p>
</div>
</section>
<section id="S4.SS2.SSS0.Px2" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Mixture-of-Depths-and-Experts.</h4>

<div id="S4.SS2.SSS0.Px2.p1" class="ltx_para">
<p id="S4.SS2.SSS0.Px2.p1.1" class="ltx_p">We observe accelerated convergence of training loss when adopting both MoE and MoD, as well as their combinations. As shown in Figure <a href="#S4.F4" title="Figure 4 ‣ 4.2 Scaling of Performance with Compute ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4</span></a>, adding MoD to the <span id="S4.SS2.SSS0.Px2.p1.1.1" class="ltx_text ltx_font_typewriter">moe_1t1i</span> architecture (<span id="S4.SS2.SSS0.Px2.p1.1.2" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span>) significantly improves model performance across various model sizes. Moreover, <span id="S4.SS2.SSS0.Px2.p1.1.3" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> performs on par with or better than <span id="S4.SS2.SSS0.Px2.p1.1.4" class="ltx_text ltx_font_typewriter">moe_4t4i</span> across model sizes and modalities, indicating that introducing sparsity over the depth dimension can also effectively improve training efficiency. On the other hand, we observe diminishing returns when stacking MoD and MoE.
Adding MoD to the <span id="S4.SS2.SSS0.Px2.p1.1.5" class="ltx_text ltx_font_typewriter">moe_4t4i</span> architecture results in only a mild performance boost compared to <span id="S4.SS2.SSS0.Px2.p1.1.6" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> and <span id="S4.SS2.SSS0.Px2.p1.1.7" class="ltx_text ltx_font_typewriter">moe_4t4i</span>. The improvement is also more visible for the text modality, while the image improvement is less significant. These findings suggest that future research may explore the combination of width and depth scaling to further enhance text modality performance. In contrast, improving image modality performance will require additional exploration of alternative approaches.</p>
</div>
</section>
</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>Scaling Number of Experts</h3>

<figure id="S4.F5" class="ltx_figure">
<div class="ltx_flex_figure">
<div class="ltx_flex_cell ltx_flex_size_1">
<figure id="S4.F5.sf1" class="ltx_figure ltx_figure_panel ltx_align_center"><img src="/html/2407.21770/assets/figures/scaling_num_experts_balanced.png" id="S4.F5.sf1.g1" class="ltx_graphics ltx_img_landscape" width="598" height="165" alt="Refer to caption">
<figcaption class="ltx_caption"><span class="ltx_tag ltx_tag_figure"><span id="S4.F5.sf1.2.1.1" class="ltx_text" style="font-size:90%;">(a)</span> </span><span id="S4.F5.sf1.3.2" class="ltx_text" style="font-size:90%;">The same number of experts are allocated to each modality</span></figcaption>
</figure>
</div>
<div class="ltx_flex_break"></div>
<div class="ltx_flex_cell ltx_flex_size_1">
<figure id="S4.F5.sf2" class="ltx_figure ltx_figure_panel ltx_align_center"><img src="/html/2407.21770/assets/figures/scaling_num_experts_imbalanced.png" id="S4.F5.sf2.g1" class="ltx_graphics ltx_img_landscape" width="598" height="165" alt="Refer to caption">
<figcaption class="ltx_caption"><span class="ltx_tag ltx_tag_figure"><span id="S4.F5.sf2.2.1.1" class="ltx_text" style="font-size:90%;">(b)</span> </span><span id="S4.F5.sf2.3.2" class="ltx_text" style="font-size:90%;">Uneven number of experts are allocated to each modality</span></figcaption>
</figure>
</div>
</div>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F5.2.1.1" class="ltx_text" style="font-size:90%;">Figure 5</span>: </span><span id="S4.F5.3.2" class="ltx_text" style="font-size:90%;">Training curves of 435M compute-matched MoE architectures with increasing number experts with various expert allocations.</span></figcaption>
</figure>
<div id="S4.SS3.p1" class="ltx_para">
<p id="S4.SS3.p1.1" class="ltx_p">We conducted further ablation experiments to investigate the impact of scaling the number of experts. We explored two scenarios: allocating an equal number of experts per modality (balanced expert allocation) and allocating different numbers of experts per modality. We included the dense model and expert-choice MoE with 8 mixed-modal experts (<span id="S4.SS3.p1.1.1" class="ltx_text ltx_font_typewriter">moe_8x</span>) as baselines.</p>
</div>
<section id="S4.SS3.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Balanced expert allocation.</h4>

<div id="S4.SS3.SSS0.Px1.p1" class="ltx_para">
<p id="S4.SS3.SSS0.Px1.p1.1" class="ltx_p">Figure <a href="#S4.F5.sf1" title="Figure 5(a) ‣ Figure 5 ‣ 4.3 Scaling Number of Experts ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">5(a)</span></a> demonstrates that training loss consistently improves as the number of experts increases. Text and image losses exhibit distinct scaling patterns. While text loss improves steadily with each doubling of expert count, image loss shows diminishing returns as the number of experts increases from 2 to 4. This suggests that the intrinsic characteristics of each modality lead to different behaviors in sparse modeling. This finding aligns with our previous observation in §<a href="#S4.SS2" title="4.2 Scaling of Performance with Compute ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.2</span></a>, where untying modalities over 8 experts significantly enhanced image performance but not text performance, indicating that the text modality may benefit more from a larger number of experts.</p>
</div>
</section>
<section id="S4.SS3.SSS0.Px2" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Imbalanced expert allocation.</h4>

<div id="S4.SS3.SSS0.Px2.p1" class="ltx_para">
<p id="S4.SS3.SSS0.Px2.p1.1" class="ltx_p">Given the diminishing returns observed when allocating more than 2 experts to the image modality, we experimented with configurations allocating fewer experts to the image modality and more to the text modality. Figure <a href="#S4.F5.sf2" title="Figure 5(b) ‣ Figure 5 ‣ 4.3 Scaling Number of Experts ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">5(b)</span></a> compares three configurations with the same total number of experts (8), varying the allocation between text and image modalities (<span id="S4.SS3.SSS0.Px2.p1.1.1" class="ltx_text ltx_font_typewriter">moe_7t1i</span>, <span id="S4.SS3.SSS0.Px2.p1.1.2" class="ltx_text ltx_font_typewriter">moe_6t2i</span> and <span id="S4.SS3.SSS0.Px2.p1.1.3" class="ltx_text ltx_font_typewriter">moe_4t4i</span>). We observed that allocating more experts to a modality generally enhances its performance. However, assigning more than 4 experts to the text modality also yields diminishing returns. The overall training losses of the three configurations converge to similar levels. Considering the better load balancing of balanced expert assignment with our pre-training data mixture, we selected <span id="S4.SS3.SSS0.Px2.p1.1.4" class="ltx_text ltx_font_typewriter">moe_4t4i</span> as our optimal MoE configuration, despite <span id="S4.SS3.SSS0.Px2.p1.1.5" class="ltx_text ltx_font_typewriter">moe_6t2i</span> having a slightly lower average loss. We leave the design of sparse architectures that can effectively leverage the intrinsic characteristics of different modalities to future work.</p>
</div>
</section>
</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>Upcycling</h3>

<div id="S4.SS4.p1" class="ltx_para">
<p id="S4.SS4.p1.1" class="ltx_p">We investigate further the influence of upcycling as described in §<a href="#S2.SS5" title="2.5 Upcycling ‣ 2 Model ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2.5</span></a>. Specifically, we examine the 2.3B MoD model, comparing the training dynamics of <span id="S4.SS4.p1.1.1" class="ltx_text ltx_font_typewriter">mod_moe_4t4i</span> when trained from scratch versus when initialized from an <span id="S4.SS4.p1.1.2" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> checkpoint. To ensure a fair comparison, we adjust the data loader and training steps to account for the number of training steps already completed by <span id="S4.SS4.p1.1.3" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span>, thereby maintaining equivalent training FLOPs. We ablate initializing the model from 10k and 20k steps.</p>
</div>
<div id="S4.SS4.p2" class="ltx_para">
<p id="S4.SS4.p2.2" class="ltx_p">Figure <a href="#S4.F6" title="Figure 6 ‣ Optimal Upcycling Period. ‣ 4.4 Upcycling ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">6</span></a> presents a comparison of the training curves for the various model variations. The training curves for the upcycled models have been adjusted to account for the computational cost of the first stage. We experimented with two seed checkpoints: <span id="S4.SS4.p2.2.1" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> trained for 10k and 20k steps, respectively. Our results show that upcycling further enhances model training, yielding a 1.2<math id="S4.SS4.p2.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S4.SS4.p2.1.m1.1a"><mo id="S4.SS4.p2.1.m1.1.1" xref="S4.SS4.p2.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S4.SS4.p2.1.m1.1b"><times id="S4.SS4.p2.1.m1.1.1.cmml" xref="S4.SS4.p2.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p2.1.m1.1c">\times</annotation></semantics></math> FLOPs gain with 10k steps in the first stage and a 1.16<math id="S4.SS4.p2.2.m2.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S4.SS4.p2.2.m2.1a"><mo id="S4.SS4.p2.2.m2.1.1" xref="S4.SS4.p2.2.m2.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S4.SS4.p2.2.m2.1b"><times id="S4.SS4.p2.2.m2.1.1.cmml" xref="S4.SS4.p2.2.m2.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.p2.2.m2.1c">\times</annotation></semantics></math> FLOPs gain with 20k steps. We observe the performance gap between upcycled models and the from-scratch model widens throughout training.</p>
</div>
<section id="S4.SS4.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Optimal Upcycling Period.</h4>

<div id="S4.SS4.SSS0.Px1.p1" class="ltx_para">
<p id="S4.SS4.SSS0.Px1.p1.2" class="ltx_p">As shown in Figure <a href="#S4.F6" title="Figure 6 ‣ Optimal Upcycling Period. ‣ 4.4 Upcycling ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">6</span></a>, both initializations outperform the base MoE model, with 10k steps providing a 1.2<math id="S4.SS4.SSS0.Px1.p1.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S4.SS4.SSS0.Px1.p1.1.m1.1a"><mo id="S4.SS4.SSS0.Px1.p1.1.m1.1.1" xref="S4.SS4.SSS0.Px1.p1.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S4.SS4.SSS0.Px1.p1.1.m1.1b"><times id="S4.SS4.SSS0.Px1.p1.1.m1.1.1.cmml" xref="S4.SS4.SSS0.Px1.p1.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.SSS0.Px1.p1.1.m1.1c">\times</annotation></semantics></math> speedup and 20k steps providing a 1.16<math id="S4.SS4.SSS0.Px1.p1.2.m2.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S4.SS4.SSS0.Px1.p1.2.m2.1a"><mo id="S4.SS4.SSS0.Px1.p1.2.m2.1.1" xref="S4.SS4.SSS0.Px1.p1.2.m2.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S4.SS4.SSS0.Px1.p1.2.m2.1b"><times id="S4.SS4.SSS0.Px1.p1.2.m2.1.1.cmml" xref="S4.SS4.SSS0.Px1.p1.2.m2.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S4.SS4.SSS0.Px1.p1.2.m2.1c">\times</annotation></semantics></math> speedup compared to the model trained from scratch. This suggests that a sweet spot likely exists for upcycling, where undertraining the seed model offers minimal benefits for router convergence, while overtraining it impedes future specialization. Based on these findings, we recommend upcycling from 10k steps. However, we hypothesize that the optimal number of upcycle steps may change when training beyond 1T tokens, and we leave this exploration for future research on in-depth upcycling of MoMa.</p>
</div>
<figure id="S4.F6" class="ltx_figure"><img src="/html/2407.21770/assets/figures/upcycling_23b_finished.png" id="S4.F6.g1" class="ltx_graphics ltx_centering ltx_img_landscape" width="598" height="194" alt="Refer to caption">
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_figure"><span id="S4.F6.5.1.1" class="ltx_text" style="font-size:90%;">Figure 6</span>: </span><span id="S4.F6.6.2" class="ltx_text" style="font-size:90%;">Upcycling experiments. We compare the training curves of 2.3B <span id="S4.F6.6.2.1" class="ltx_text ltx_font_typewriter">mod_moe_4t4i</span> starting from three different initialization points: (1) scratch, (2) 2.3B <span id="S4.F6.6.2.2" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> trained for 10k steps, and (3) 2.3B <span id="S4.F6.6.2.3" class="ltx_text ltx_font_typewriter">mod_moe_1t1i</span> trained for 20k steps. All curves are FLOP adjusted to be equivalent, by shifting the upcycled models by the number of steps the dense model has been trained on.</span></figcaption>
</figure>
</section>
</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>Throughput Analysis</h3>

<div id="S4.SS5.p1" class="ltx_para">
<p id="S4.SS5.p1.1" class="ltx_p">Sparse models often cannot immediately lead to performance gains due to their added dynamism and associated data balancing issues (§<a href="#S3" title="3 Efficiency Optimization ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>). To quantify the impact of our proposals on training efficiency, we conducted a controlled experiment comparing the training throughput of various architectures, including a combination of MoMa, mixture-of-experts (MoE), and mixture-of-depths (MoD).
We compare architectures FLOPs-matched to the 435M dense baseline. We performed our experiments on 256 A100 GPUs with a sequence length of 4096 and a batch size of 6 per GPU and summarized the model throughput in Table <a href="#S4.T2" title="Table 2 ‣ 4.5 Throughput Analysis ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">2</span></a>.</p>
</div>
<figure id="S4.T2" class="ltx_table">
<table id="S4.T2.2" class="ltx_tabular ltx_centering ltx_align_middle">
<tr id="S4.T2.2.1" class="ltx_tr">
<td id="S4.T2.2.1.1" class="ltx_td ltx_align_left ltx_border_tt">Configuration</td>
<td id="S4.T2.2.1.2" class="ltx_td ltx_align_right ltx_border_tt">WPS</td>
<td id="S4.T2.2.1.3" class="ltx_td ltx_align_right ltx_border_tt">% diff</td>
</tr>
<tr id="S4.T2.2.2" class="ltx_tr">
<td id="S4.T2.2.2.1" class="ltx_td ltx_align_left ltx_border_t">Dense</td>
<td id="S4.T2.2.2.2" class="ltx_td ltx_align_right ltx_border_t">31,970</td>
<td id="S4.T2.2.2.3" class="ltx_td ltx_align_right ltx_border_t">–</td>
</tr>
<tr id="S4.T2.2.3" class="ltx_tr">
<td id="S4.T2.2.3.1" class="ltx_td ltx_align_left">MoMa 8X</td>
<td id="S4.T2.2.3.2" class="ltx_td ltx_align_right">28,990</td>
<td id="S4.T2.2.3.3" class="ltx_td ltx_align_right">-9%</td>
</tr>
<tr id="S4.T2.2.4" class="ltx_tr">
<td id="S4.T2.2.4.1" class="ltx_td ltx_align_left">MoMa 1t1i</td>
<td id="S4.T2.2.4.2" class="ltx_td ltx_align_right">30,196</td>
<td id="S4.T2.2.4.3" class="ltx_td ltx_align_right">-6%</td>
</tr>
<tr id="S4.T2.2.5" class="ltx_tr">
<td id="S4.T2.2.5.1" class="ltx_td ltx_align_left">MoMa 4t4i</td>
<td id="S4.T2.2.5.2" class="ltx_td ltx_align_right">26,424</td>
<td id="S4.T2.2.5.3" class="ltx_td ltx_align_right">-17%</td>
</tr>
<tr id="S4.T2.2.6" class="ltx_tr">
<td id="S4.T2.2.6.1" class="ltx_td ltx_align_left">MoD MoMa 1t1i</td>
<td id="S4.T2.2.6.2" class="ltx_td ltx_align_right">25,367</td>
<td id="S4.T2.2.6.3" class="ltx_td ltx_align_right">-21%</td>
</tr>
<tr id="S4.T2.2.7" class="ltx_tr">
<td id="S4.T2.2.7.1" class="ltx_td ltx_align_left ltx_border_bb">MoD MoMa 4t4i</td>
<td id="S4.T2.2.7.2" class="ltx_td ltx_align_right ltx_border_bb">22,114</td>
<td id="S4.T2.2.7.3" class="ltx_td ltx_align_right ltx_border_bb">-31%</td>
</tr>
</table>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T2.3.1.1" class="ltx_text" style="font-size:90%;">Table 2</span>: </span><span id="S4.T2.4.2" class="ltx_text" style="font-size:90%;">Training throughput (wps/GPU) of various architectures compared with the same flops as a 435M dense model on 256 A100 GPUs. Modality-conditioned sparsity offers a good quality-throughput tradeoff compared to the dense model, and exhibited reasonable scalability with more experts. On the other hand, while MoD variants achieves the best absolute loss, they generally incurs larger overheads due to additional dynamism and imbalance.</span></figcaption>
</figure>
<div id="S4.SS5.p2" class="ltx_para">
<p id="S4.SS5.p2.1" class="ltx_p">Comparing expert-choice MoE (<span id="S4.SS5.p2.1.1" class="ltx_text ltx_font_typewriter">moe_8x</span>) and dense models, introducing sparsity does incur overheads (-9%). This loss in throughput likely comes from the need to compute routing decisions and synchronize gradients for all experts despite being FLOPs-equivalent. On the other hand, as discussed in §<a href="#S3" title="3 Efficiency Optimization ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>, running experts sequentially by modality (<span id="S4.SS5.p2.1.2" class="ltx_text ltx_font_typewriter">moe_1t1i</span>) does not suffer from large execution overheads, and we can attribute most of the loss (-6%) to computing token indices for each modality, which can be amortized by precomputing the indices and share with each transformer layer.</p>
</div>
<div id="S4.SS5.p3" class="ltx_para">
<p id="S4.SS5.p3.1" class="ltx_p">When combining modality conditioned feed-forward with learned routing (<span id="S4.SS5.p3.1.1" class="ltx_text ltx_font_typewriter">moe_4t4i</span>), we observe a smooth degradation in throughput as the number of experts increases, incurring additional 11% overheads with 8 experts, which is on par with the throughput loss of 9% when transitioning from the dense model to MoE with 8 experts.</p>
</div>
<div id="S4.SS5.p4" class="ltx_para">
<p id="S4.SS5.p4.1" class="ltx_p">While combining MoD and MoE achieves the best training loss in our experiments, introducing MoD triggers an estimated throughput loss of 15% . This is because MoD architectures require an additional router on the depth dimension, introducing complexity and potential bottlenecks. Furthermore, when combined with MoMa, MoD may exacerbate system imbalance due to the varying active tokens for each modality at each layer, deviating from the predefined token mix ratio in the dataset. This deviation undermines the assumption of efficient MoMa execution, particularly with larger world sizes.
To mitigate this issue, we can force the MoD routers to admit tokens based on the predefined token mix ratio. However, the impact on model quality requires further investigation, which we leave for future work.</p>
</div>
</section>
<section id="S4.SS6" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">4.6 </span>Inference-time Performance</h3>

<div id="S4.SS6.p1" class="ltx_para">
<p id="S4.SS6.p1.1" class="ltx_p">We evaluate our models (1.4B dense, 1.4B <span id="S4.SS6.p1.1.1" class="ltx_text ltx_font_typewriter">moe_1t1i</span>, 1.4B <span id="S4.SS6.p1.1.2" class="ltx_text ltx_font_typewriter">moe_8x</span>, 1.4B <span id="S4.SS6.p1.1.3" class="ltx_text ltx_font_typewriter">moe_4t4i</span>, 2.3B <span id="S4.SS6.p1.1.4" class="ltx_text ltx_font_typewriter">mod_moe_4t4i</span>) on held-out language modeling data and downstream tasks. We report perplexity on subsets of OBELICS <cite class="ltx_cite ltx_citemacro_citep">(Laurençon et al., <a href="#bib.bib29" title="" class="ltx_ref">2023</a>)</cite> and Shutterstock<span id="footnote7" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">7</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">7</sup><span class="ltx_tag ltx_tag_note">7</span><a target="_blank" href="https://www.shuitterstock.com/" title="" class="ltx_ref ltx_url ltx_font_typewriter">https://www.shuitterstock.com/</a></span></span></span>, which are held-out from our pre-training dataset. In addition, we report 0-shot performance on a set of commonsense reasoning of tasks commonly used for benchmarking pre-trained language models. We also selected several vision-language task datasets and report the perplexity of the ground truth output in these datasets for cross model comparison, where we format the examples in a zero-shot manner.<span id="footnote8" class="ltx_note ltx_role_footnote"><sup class="ltx_note_mark">8</sup><span class="ltx_note_outer"><span class="ltx_note_content"><sup class="ltx_note_mark">8</sup><span class="ltx_tag ltx_tag_note">8</span>We do not report the generation evaluation metrics for these tasks, as our experiment checkpoints are under-trained and have not achieved sufficient image understanding and generation capabilities to produce meaningful results on downstream tasks.</span></span></span></p>
<ul id="S4.I1" class="ltx_itemize">
<li id="S4.I1.i1" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S4.I1.i1.p1" class="ltx_para">
<p id="S4.I1.i1.p1.1" class="ltx_p"><span id="S4.I1.i1.p1.1.1" class="ltx_text ltx_font_bold">Commonsense Reasoning:</span> We evaluated on a set of text-only tasks commonly used for benchmarking the commonsense reasoning capabilities of language models: <span id="S4.I1.i1.p1.1.2" class="ltx_text ltx_font_typewriter">PIQA</span> <cite class="ltx_cite ltx_citemacro_citep">(Bisk et al., <a href="#bib.bib30" title="" class="ltx_ref">2020</a>)</cite>, <span id="S4.I1.i1.p1.1.3" class="ltx_text ltx_font_typewriter">SIQA</span> <cite class="ltx_cite ltx_citemacro_citep">(Sap et al., <a href="#bib.bib31" title="" class="ltx_ref">2019</a>)</cite>, <span id="S4.I1.i1.p1.1.4" class="ltx_text ltx_font_typewriter">HellaSwag</span> <cite class="ltx_cite ltx_citemacro_citep">(Zellers et al., <a href="#bib.bib32" title="" class="ltx_ref">2019</a>)</cite>, <span id="S4.I1.i1.p1.1.5" class="ltx_text ltx_font_typewriter">WinoGrande</span> <cite class="ltx_cite ltx_citemacro_citep">(Sakaguchi et al., <a href="#bib.bib33" title="" class="ltx_ref">2021</a>)</cite>, <span id="S4.I1.i1.p1.1.6" class="ltx_text ltx_font_typewriter">ARC-Easy</span> <cite class="ltx_cite ltx_citemacro_citep">(Clark et al., <a href="#bib.bib34" title="" class="ltx_ref">2018</a>)</cite>, <span id="S4.I1.i1.p1.1.7" class="ltx_text ltx_font_typewriter">ARC-Challenge</span> <cite class="ltx_cite ltx_citemacro_citep">(Clark et al., <a href="#bib.bib34" title="" class="ltx_ref">2018</a>)</cite>, <span id="S4.I1.i1.p1.1.8" class="ltx_text ltx_font_typewriter">OpenBookQA</span> <cite class="ltx_cite ltx_citemacro_citep">(Mihaylov et al., <a href="#bib.bib35" title="" class="ltx_ref">2018</a>)</cite>, and <span id="S4.I1.i1.p1.1.9" class="ltx_text ltx_font_typewriter">BoolQ</span> <cite class="ltx_cite ltx_citemacro_citep">(Clark et al., <a href="#bib.bib36" title="" class="ltx_ref">2019</a>)</cite>. We score the prompt with each candidate’s answer and compute accuracy using the candidate with the highest score.</p>
</div>
</li>
<li id="S4.I1.i2" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S4.I1.i2.p1" class="ltx_para">
<p id="S4.I1.i2.p1.1" class="ltx_p"><span id="S4.I1.i2.p1.1.1" class="ltx_text ltx_font_bold">Image Captioning:</span> We take the Karpathy test split of MS-COCO <cite class="ltx_cite ltx_citemacro_citep">(Lin et al., <a href="#bib.bib37" title="" class="ltx_ref">2014</a>)</cite> and the Karpathy test split of Flickr30k <cite class="ltx_cite ltx_citemacro_citep">(Plummer et al., <a href="#bib.bib38" title="" class="ltx_ref">2015</a>)</cite>, and report text-to-image and image-to-text conditional perplexity using these two datasets.</p>
</div>
</li>
<li id="S4.I1.i3" class="ltx_item" style="list-style-type:none;">
<span class="ltx_tag ltx_tag_item">•</span> 
<div id="S4.I1.i3.p1" class="ltx_para">
<p id="S4.I1.i3.p1.1" class="ltx_p"><span id="S4.I1.i3.p1.1.1" class="ltx_text ltx_font_bold">Visual Question Answering:</span> We report the perplexity of ground truth answers on the test-dev split of VQA-v2 <cite class="ltx_cite ltx_citemacro_citep">(Goyal et al., <a href="#bib.bib39" title="" class="ltx_ref">2017</a>)</cite>.</p>
</div>
</li>
</ul>
</div>
<figure id="S4.T3" class="ltx_table">
<div id="S4.T3.3" class="ltx_inline-block ltx_align_center ltx_transformed_outer" style="width:518.6pt;height:123.3pt;vertical-align:-0.8pt;"><span class="ltx_transformed_inner" style="transform:translate(-56.9pt,13.4pt) scale(0.82,0.82) ;">
<table id="S4.T3.3.3" class="ltx_tabular ltx_align_middle">
<tr id="S4.T3.3.3.3" class="ltx_tr">
<td id="S4.T3.3.3.3.4" class="ltx_td ltx_border_tt"></td>
<td id="S4.T3.1.1.1.1" class="ltx_td ltx_align_center ltx_border_tt" colspan="3"><span id="S4.T3.1.1.1.1.1" class="ltx_text ltx_font_bold">Image to text (PPL <math id="S4.T3.1.1.1.1.1.m1.1" class="ltx_Math" alttext="\downarrow" display="inline"><semantics id="S4.T3.1.1.1.1.1.m1.1a"><mo stretchy="false" id="S4.T3.1.1.1.1.1.m1.1.1" xref="S4.T3.1.1.1.1.1.m1.1.1.cmml">↓</mo><annotation-xml encoding="MathML-Content" id="S4.T3.1.1.1.1.1.m1.1b"><ci id="S4.T3.1.1.1.1.1.m1.1.1.cmml" xref="S4.T3.1.1.1.1.1.m1.1.1">↓</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.1.1.1.1.1.m1.1c">\downarrow</annotation></semantics></math>)</span></td>
<td id="S4.T3.2.2.2.2" class="ltx_td ltx_align_center ltx_border_tt" colspan="2"><span id="S4.T3.2.2.2.2.1" class="ltx_text ltx_font_bold">Text to image (PPL <math id="S4.T3.2.2.2.2.1.m1.1" class="ltx_Math" alttext="\downarrow" display="inline"><semantics id="S4.T3.2.2.2.2.1.m1.1a"><mo stretchy="false" id="S4.T3.2.2.2.2.1.m1.1.1" xref="S4.T3.2.2.2.2.1.m1.1.1.cmml">↓</mo><annotation-xml encoding="MathML-Content" id="S4.T3.2.2.2.2.1.m1.1b"><ci id="S4.T3.2.2.2.2.1.m1.1.1.cmml" xref="S4.T3.2.2.2.2.1.m1.1.1">↓</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.2.2.2.2.1.m1.1c">\downarrow</annotation></semantics></math>)</span></td>
<td id="S4.T3.3.3.3.3" class="ltx_td ltx_align_center ltx_border_tt" colspan="6"><span id="S4.T3.3.3.3.3.1" class="ltx_text ltx_font_bold">Interleaved (PPL <math id="S4.T3.3.3.3.3.1.m1.1" class="ltx_Math" alttext="\downarrow" display="inline"><semantics id="S4.T3.3.3.3.3.1.m1.1a"><mo stretchy="false" id="S4.T3.3.3.3.3.1.m1.1.1" xref="S4.T3.3.3.3.3.1.m1.1.1.cmml">↓</mo><annotation-xml encoding="MathML-Content" id="S4.T3.3.3.3.3.1.m1.1b"><ci id="S4.T3.3.3.3.3.1.m1.1.1.cmml" xref="S4.T3.3.3.3.3.1.m1.1.1">↓</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.T3.3.3.3.3.1.m1.1c">\downarrow</annotation></semantics></math>)</span></td>
</tr>
<tr id="S4.T3.3.3.4" class="ltx_tr">
<td id="S4.T3.3.3.4.1" class="ltx_td ltx_align_left ltx_border_t"><span id="S4.T3.3.3.4.1.1" class="ltx_text ltx_font_bold">Model</span></td>
<td id="S4.T3.3.3.4.2" class="ltx_td ltx_align_right ltx_border_t"><span id="S4.T3.3.3.4.2.1" class="ltx_text ltx_font_bold">COCO</span></td>
<td id="S4.T3.3.3.4.3" class="ltx_td ltx_align_right ltx_border_t"><span id="S4.T3.3.3.4.3.1" class="ltx_text ltx_font_bold">Flickr</span></td>
<td id="S4.T3.3.3.4.4" class="ltx_td ltx_align_right ltx_border_t"><span id="S4.T3.3.3.4.4.1" class="ltx_text ltx_font_bold">VQAV2</span></td>
<td id="S4.T3.3.3.4.5" class="ltx_td ltx_align_right ltx_border_t"><span id="S4.T3.3.3.4.5.1" class="ltx_text ltx_font_bold">COCO</span></td>
<td id="S4.T3.3.3.4.6" class="ltx_td ltx_align_right ltx_border_t"><span id="S4.T3.3.3.4.6.1" class="ltx_text ltx_font_bold">Flickr</span></td>
<td id="S4.T3.3.3.4.7" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.7.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.7.2" class="ltx_text">
<span id="S4.T3.3.3.4.7.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.7.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.7.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.7.2.1.1.1.1" class="ltx_text ltx_font_bold">Obelics</span></span></span>
<span id="S4.T3.3.3.4.7.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.7.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.7.2.1.2.1.1" class="ltx_text ltx_font_bold">Overall</span></span></span>
</span></span><span id="S4.T3.3.3.4.7.3" class="ltx_text"></span>
</td>
<td id="S4.T3.3.3.4.8" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.8.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.8.2" class="ltx_text">
<span id="S4.T3.3.3.4.8.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.8.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.8.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.8.2.1.1.1.1" class="ltx_text ltx_font_bold">Obelics</span></span></span>
<span id="S4.T3.3.3.4.8.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.8.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.8.2.1.2.1.1" class="ltx_text ltx_font_bold">Text</span></span></span>
</span></span><span id="S4.T3.3.3.4.8.3" class="ltx_text"></span>
</td>
<td id="S4.T3.3.3.4.9" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.9.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.9.2" class="ltx_text">
<span id="S4.T3.3.3.4.9.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.9.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.9.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.9.2.1.1.1.1" class="ltx_text ltx_font_bold">Obelics</span></span></span>
<span id="S4.T3.3.3.4.9.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.9.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.9.2.1.2.1.1" class="ltx_text ltx_font_bold">Image</span></span></span>
</span></span><span id="S4.T3.3.3.4.9.3" class="ltx_text"></span>
</td>
<td id="S4.T3.3.3.4.10" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.10.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.10.2" class="ltx_text">
<span id="S4.T3.3.3.4.10.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.10.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.10.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.10.2.1.1.1.1" class="ltx_text ltx_font_bold">SSTK</span></span></span>
<span id="S4.T3.3.3.4.10.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.10.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.10.2.1.2.1.1" class="ltx_text ltx_font_bold">Overall</span></span></span>
</span></span><span id="S4.T3.3.3.4.10.3" class="ltx_text"></span>
</td>
<td id="S4.T3.3.3.4.11" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.11.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.11.2" class="ltx_text">
<span id="S4.T3.3.3.4.11.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.11.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.11.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.11.2.1.1.1.1" class="ltx_text ltx_font_bold">SSTK</span></span></span>
<span id="S4.T3.3.3.4.11.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.11.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.11.2.1.2.1.1" class="ltx_text ltx_font_bold">Text</span></span></span>
</span></span><span id="S4.T3.3.3.4.11.3" class="ltx_text"></span>
</td>
<td id="S4.T3.3.3.4.12" class="ltx_td ltx_align_center ltx_border_t">
<span id="S4.T3.3.3.4.12.1" class="ltx_text"></span> <span id="S4.T3.3.3.4.12.2" class="ltx_text">
<span id="S4.T3.3.3.4.12.2.1" class="ltx_tabular ltx_align_middle">
<span id="S4.T3.3.3.4.12.2.1.1" class="ltx_tr">
<span id="S4.T3.3.3.4.12.2.1.1.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.12.2.1.1.1.1" class="ltx_text ltx_font_bold">SSTK</span></span></span>
<span id="S4.T3.3.3.4.12.2.1.2" class="ltx_tr">
<span id="S4.T3.3.3.4.12.2.1.2.1" class="ltx_td ltx_nopad_r ltx_align_center"><span id="S4.T3.3.3.4.12.2.1.2.1.1" class="ltx_text ltx_font_bold">Image</span></span></span>
</span></span><span id="S4.T3.3.3.4.12.3" class="ltx_text"></span>
</td>
</tr>
<tr id="S4.T3.3.3.5" class="ltx_tr">
<td id="S4.T3.3.3.5.1" class="ltx_td ltx_align_left ltx_border_t">1.4B Dense</td>
<td id="S4.T3.3.3.5.2" class="ltx_td ltx_align_right ltx_border_t">21.7</td>
<td id="S4.T3.3.3.5.3" class="ltx_td ltx_align_right ltx_border_t">28.6</td>
<td id="S4.T3.3.3.5.4" class="ltx_td ltx_align_right ltx_border_t">20.0</td>
<td id="S4.T3.3.3.5.5" class="ltx_td ltx_align_right ltx_border_t">458.3</td>
<td id="S4.T3.3.3.5.6" class="ltx_td ltx_align_right ltx_border_t">559.4</td>
<td id="S4.T3.3.3.5.7" class="ltx_td ltx_align_right ltx_border_t">44.6</td>
<td id="S4.T3.3.3.5.8" class="ltx_td ltx_align_right ltx_border_t">14.3</td>
<td id="S4.T3.3.3.5.9" class="ltx_td ltx_align_right ltx_border_t">64.7</td>
<td id="S4.T3.3.3.5.10" class="ltx_td ltx_align_right ltx_border_t">238.5</td>
<td id="S4.T3.3.3.5.11" class="ltx_td ltx_align_right ltx_border_t">48.7</td>
<td id="S4.T3.3.3.5.12" class="ltx_td ltx_align_right ltx_border_t">245.8</td>
</tr>
<tr id="S4.T3.3.3.6" class="ltx_tr">
<td id="S4.T3.3.3.6.1" class="ltx_td ltx_align_left">1.4B MoMa 1t1i</td>
<td id="S4.T3.3.3.6.2" class="ltx_td ltx_align_right">21.9</td>
<td id="S4.T3.3.3.6.3" class="ltx_td ltx_align_right">29.3</td>
<td id="S4.T3.3.3.6.4" class="ltx_td ltx_align_right">19.8</td>
<td id="S4.T3.3.3.6.5" class="ltx_td ltx_align_right">416.7</td>
<td id="S4.T3.3.3.6.6" class="ltx_td ltx_align_right">508.6</td>
<td id="S4.T3.3.3.6.7" class="ltx_td ltx_align_right">41.5</td>
<td id="S4.T3.3.3.6.8" class="ltx_td ltx_align_right">13.8</td>
<td id="S4.T3.3.3.6.9" class="ltx_td ltx_align_right">59.3</td>
<td id="S4.T3.3.3.6.10" class="ltx_td ltx_align_right">208.1</td>
<td id="S4.T3.3.3.6.11" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.6.11.1" class="ltx_text ltx_font_bold">7.6</span></td>
<td id="S4.T3.3.3.6.12" class="ltx_td ltx_align_right">221.5</td>
</tr>
<tr id="S4.T3.3.3.7" class="ltx_tr">
<td id="S4.T3.3.3.7.1" class="ltx_td ltx_align_left">1.4B MoE 8x</td>
<td id="S4.T3.3.3.7.2" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.7.2.1" class="ltx_text ltx_font_bold">18.9</span></td>
<td id="S4.T3.3.3.7.3" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.7.3.1" class="ltx_text ltx_font_bold">24.2</span></td>
<td id="S4.T3.3.3.7.4" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.7.4.1" class="ltx_text ltx_font_bold">18.9</span></td>
<td id="S4.T3.3.3.7.5" class="ltx_td ltx_align_right">426.5</td>
<td id="S4.T3.3.3.7.6" class="ltx_td ltx_align_right">518.5</td>
<td id="S4.T3.3.3.7.7" class="ltx_td ltx_align_right">40.4</td>
<td id="S4.T3.3.3.7.8" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.7.8.1" class="ltx_text ltx_font_bold">12.5</span></td>
<td id="S4.T3.3.3.7.9" class="ltx_td ltx_align_right">59.2</td>
<td id="S4.T3.3.3.7.10" class="ltx_td ltx_align_right">218.8</td>
<td id="S4.T3.3.3.7.11" class="ltx_td ltx_align_right">33.7</td>
<td id="S4.T3.3.3.7.12" class="ltx_td ltx_align_right">226.6</td>
</tr>
<tr id="S4.T3.3.3.8" class="ltx_tr">
<td id="S4.T3.3.3.8.1" class="ltx_td ltx_align_left">1.4B MoMa 4t4i</td>
<td id="S4.T3.3.3.8.2" class="ltx_td ltx_align_right">20.9</td>
<td id="S4.T3.3.3.8.3" class="ltx_td ltx_align_right">26.6</td>
<td id="S4.T3.3.3.8.4" class="ltx_td ltx_align_right">19.8</td>
<td id="S4.T3.3.3.8.5" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.5.1" class="ltx_text ltx_font_bold">392.5</span></td>
<td id="S4.T3.3.3.8.6" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.6.1" class="ltx_text ltx_font_bold">479.6</span></td>
<td id="S4.T3.3.3.8.7" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.7.1" class="ltx_text ltx_font_bold">39.0</span></td>
<td id="S4.T3.3.3.8.8" class="ltx_td ltx_align_right">13.0</td>
<td id="S4.T3.3.3.8.9" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.9.1" class="ltx_text ltx_font_bold">55.8</span></td>
<td id="S4.T3.3.3.8.10" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.10.1" class="ltx_text ltx_font_bold">194.7</span></td>
<td id="S4.T3.3.3.8.11" class="ltx_td ltx_align_right">10.5</td>
<td id="S4.T3.3.3.8.12" class="ltx_td ltx_align_right"><span id="S4.T3.3.3.8.12.1" class="ltx_text ltx_font_bold">205.7</span></td>
</tr>
<tr id="S4.T3.3.3.9" class="ltx_tr">
<td id="S4.T3.3.3.9.1" class="ltx_td ltx_align_left ltx_border_bb">
<span id="S4.T3.3.3.9.1.1" class="ltx_ERROR undefined">\hdashline</span>2.3B MoD MoMa 4t4i</td>
<td id="S4.T3.3.3.9.2" class="ltx_td ltx_align_right ltx_border_bb">22.9</td>
<td id="S4.T3.3.3.9.3" class="ltx_td ltx_align_right ltx_border_bb">28.0</td>
<td id="S4.T3.3.3.9.4" class="ltx_td ltx_align_right ltx_border_bb">20.5</td>
<td id="S4.T3.3.3.9.5" class="ltx_td ltx_align_right ltx_border_bb">428.1</td>
<td id="S4.T3.3.3.9.6" class="ltx_td ltx_align_right ltx_border_bb">518.0</td>
<td id="S4.T3.3.3.9.7" class="ltx_td ltx_align_right ltx_border_bb">41.4</td>
<td id="S4.T3.3.3.9.8" class="ltx_td ltx_align_right ltx_border_bb">13.2</td>
<td id="S4.T3.3.3.9.9" class="ltx_td ltx_align_right ltx_border_bb">60.0</td>
<td id="S4.T3.3.3.9.10" class="ltx_td ltx_align_right ltx_border_bb">215.1</td>
<td id="S4.T3.3.3.9.11" class="ltx_td ltx_align_right ltx_border_bb">15.0</td>
<td id="S4.T3.3.3.9.12" class="ltx_td ltx_align_right ltx_border_bb">226.1</td>
</tr>
</table>
</span></div>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T3.5.1.1" class="ltx_text" style="font-size:90%;">Table 3</span>: </span><span id="S4.T3.6.2" class="ltx_text" style="font-size:90%;">Interleaved data modeling performance.</span></figcaption>
</figure>
<figure id="S4.T4" class="ltx_table">
<div id="S4.T4.2" class="ltx_inline-block ltx_align_center ltx_transformed_outer" style="width:456.8pt;height:86.4pt;vertical-align:-0.9pt;"><span class="ltx_transformed_inner" style="transform:translate(-12.0pt,2.3pt) scale(0.95,0.95) ;">
<table id="S4.T4.2.1" class="ltx_tabular ltx_align_middle">
<tr id="S4.T4.2.1.1" class="ltx_tr">
<td id="S4.T4.2.1.1.1" class="ltx_td ltx_align_left ltx_border_tt"><span id="S4.T4.2.1.1.1.1" class="ltx_text ltx_font_bold">Model</span></td>
<td id="S4.T4.2.1.1.2" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.2.1" class="ltx_text ltx_font_bold">BoolQ</span></td>
<td id="S4.T4.2.1.1.3" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.3.1" class="ltx_text ltx_font_bold">PiQA</span></td>
<td id="S4.T4.2.1.1.4" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.4.1" class="ltx_text ltx_font_bold">SiQA</span></td>
<td id="S4.T4.2.1.1.5" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.5.1" class="ltx_text ltx_font_bold">Winogrande1.1</span></td>
<td id="S4.T4.2.1.1.6" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.6.1" class="ltx_text ltx_font_bold">OBQA</span></td>
<td id="S4.T4.2.1.1.7" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.7.1" class="ltx_text ltx_font_bold">Hellaswag</span></td>
<td id="S4.T4.2.1.1.8" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.8.1" class="ltx_text ltx_font_bold">Arc-E</span></td>
<td id="S4.T4.2.1.1.9" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.9.1" class="ltx_text ltx_font_bold">Arc-C</span></td>
<td id="S4.T4.2.1.1.10" class="ltx_td ltx_align_center ltx_border_tt"><span id="S4.T4.2.1.1.10.1" class="ltx_text ltx_font_bold">Avg</span></td>
</tr>
<tr id="S4.T4.2.1.2" class="ltx_tr">
<td id="S4.T4.2.1.2.1" class="ltx_td ltx_align_left ltx_border_t">1.4B Dense</td>
<td id="S4.T4.2.1.2.2" class="ltx_td ltx_align_center ltx_border_t">61.8</td>
<td id="S4.T4.2.1.2.3" class="ltx_td ltx_align_center ltx_border_t">72.1</td>
<td id="S4.T4.2.1.2.4" class="ltx_td ltx_align_center ltx_border_t">44.3</td>
<td id="S4.T4.2.1.2.5" class="ltx_td ltx_align_center ltx_border_t">54.2</td>
<td id="S4.T4.2.1.2.6" class="ltx_td ltx_align_center ltx_border_t">42.0</td>
<td id="S4.T4.2.1.2.7" class="ltx_td ltx_align_center ltx_border_t">51.0</td>
<td id="S4.T4.2.1.2.8" class="ltx_td ltx_align_center ltx_border_t">59.8</td>
<td id="S4.T4.2.1.2.9" class="ltx_td ltx_align_center ltx_border_t">34.3</td>
<td id="S4.T4.2.1.2.10" class="ltx_td ltx_align_center ltx_border_t">52.4</td>
</tr>
<tr id="S4.T4.2.1.3" class="ltx_tr">
<td id="S4.T4.2.1.3.1" class="ltx_td ltx_align_left">1.4B MoMa 1t1i</td>
<td id="S4.T4.2.1.3.2" class="ltx_td ltx_align_center">61.5</td>
<td id="S4.T4.2.1.3.3" class="ltx_td ltx_align_center">71.3</td>
<td id="S4.T4.2.1.3.4" class="ltx_td ltx_align_center">44.4</td>
<td id="S4.T4.2.1.3.5" class="ltx_td ltx_align_center">54.7</td>
<td id="S4.T4.2.1.3.6" class="ltx_td ltx_align_center">46.0</td>
<td id="S4.T4.2.1.3.7" class="ltx_td ltx_align_center">52.8</td>
<td id="S4.T4.2.1.3.8" class="ltx_td ltx_align_center">61.1</td>
<td id="S4.T4.2.1.3.9" class="ltx_td ltx_align_center">34.9</td>
<td id="S4.T4.2.1.3.10" class="ltx_td ltx_align_center">53.3</td>
</tr>
<tr id="S4.T4.2.1.4" class="ltx_tr">
<td id="S4.T4.2.1.4.1" class="ltx_td ltx_align_left">1.4B MoE 8x</td>
<td id="S4.T4.2.1.4.2" class="ltx_td ltx_align_center">62.2</td>
<td id="S4.T4.2.1.4.3" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.3.1" class="ltx_text ltx_font_bold">73.5</span></td>
<td id="S4.T4.2.1.4.4" class="ltx_td ltx_align_center">44.4</td>
<td id="S4.T4.2.1.4.5" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.5.1" class="ltx_text ltx_font_bold">58.3</span></td>
<td id="S4.T4.2.1.4.6" class="ltx_td ltx_align_center">42.2</td>
<td id="S4.T4.2.1.4.7" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.7.1" class="ltx_text ltx_font_bold">58.2</span></td>
<td id="S4.T4.2.1.4.8" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.8.1" class="ltx_text ltx_font_bold">62.6</span></td>
<td id="S4.T4.2.1.4.9" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.9.1" class="ltx_text ltx_font_bold">36.9</span></td>
<td id="S4.T4.2.1.4.10" class="ltx_td ltx_align_center"><span id="S4.T4.2.1.4.10.1" class="ltx_text ltx_font_bold">54.8</span></td>
</tr>
<tr id="S4.T4.2.1.5" class="ltx_tr">
<td id="S4.T4.2.1.5.1" class="ltx_td ltx_align_left ltx_border_bb">1.4B MoMa 4t4i</td>
<td id="S4.T4.2.1.5.2" class="ltx_td ltx_align_center ltx_border_bb"><span id="S4.T4.2.1.5.2.1" class="ltx_text ltx_font_bold">62.3</span></td>
<td id="S4.T4.2.1.5.3" class="ltx_td ltx_align_center ltx_border_bb">72.5</td>
<td id="S4.T4.2.1.5.4" class="ltx_td ltx_align_center ltx_border_bb"><span id="S4.T4.2.1.5.4.1" class="ltx_text ltx_font_bold">44.6</span></td>
<td id="S4.T4.2.1.5.5" class="ltx_td ltx_align_center ltx_border_bb">55.9</td>
<td id="S4.T4.2.1.5.6" class="ltx_td ltx_align_center ltx_border_bb"><span id="S4.T4.2.1.5.6.1" class="ltx_text ltx_font_bold">47.0</span></td>
<td id="S4.T4.2.1.5.7" class="ltx_td ltx_align_center ltx_border_bb">55.2</td>
<td id="S4.T4.2.1.5.8" class="ltx_td ltx_align_center ltx_border_bb">60.7</td>
<td id="S4.T4.2.1.5.9" class="ltx_td ltx_align_center ltx_border_bb">35.7</td>
<td id="S4.T4.2.1.5.10" class="ltx_td ltx_align_center ltx_border_bb">54.2</td>
</tr>
</table>
</span></div>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="S4.T4.3.1.1" class="ltx_text" style="font-size:90%;">Table 4</span>: </span><span id="S4.T4.4.2" class="ltx_text" style="font-size:90%;">Text-to-text commonsense reasoning task performance.</span></figcaption>
</figure>
<div id="S4.SS6.p2" class="ltx_para">
<p id="S4.SS6.p2.1" class="ltx_p">For all sparse models, we perform the second-stage training of auxiliary routers and use the auxiliary router for causal inference. The details of auxiliary router training can be found in Appendix <a href="#A1.SS2" title="A.2 Auxiliary Router Training ‣ Appendix A Additional Implementation Details ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">A.2</span></a>.</p>
</div>
<section id="S4.SS6.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Interleaved data modeling.</h4>

<div id="S4.SS6.SSS0.Px1.p1" class="ltx_para">
<p id="S4.SS6.SSS0.Px1.p1.1" class="ltx_p">The relative performance of the dense model and various MoE configurations is consistent with the trends observed in pre-training loss (§<a href="#S4.SS2" title="4.2 Scaling of Performance with Compute ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.2</span></a>). As shown in Table <a href="#S4.T3" title="Table 3 ‣ 4.6 Inference-time Performance ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>, by using 1 more image expert, the 1.4B MoMa 1t1i model significantly outperforms the dense baseline on most metrics, except for image-to-text conditional perplexity on COCO and Flickr. Similar to the trend shown in §<a href="#S4.SS2" title="4.2 Scaling of Performance with Compute ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.2</span></a>, adding the image expert leads to substantial performance gains on the image modality. Further scaling up the number of experts also improves the performance, with 1.4B MoE 8x achieves the best image-to-text performance. Additionally, the model excels on text-to-text tasks (Table <a href="#S4.T4" title="Table 4 ‣ 4.6 Inference-time Performance ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4</span></a>). 1.4B MoMa 4t4i performs the best on all conditional image perplexity metrics, with the text perplexity closely matches that of 1.4B MoE 8x on most of the benchmarks. Overall, the 1.4B MoMa 4t4i model achieves the best interleaved data modeling results.</p>
</div>
</section>
<section id="S4.SS6.SSS0.Px2" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">MoD performance regression.</h4>

<div id="S4.SS6.SSS0.Px2.p1" class="ltx_para">
<p id="S4.SS6.SSS0.Px2.p1.3" class="ltx_p">According to Table <a href="#S4.T3" title="Table 3 ‣ 4.6 Inference-time Performance ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">3</span></a>, we observe significant performance regression of the 2.3B MoD MoMa 4t4i model across all metrics, leaving it underperforming the 1.4B MoMa 4t4i model despite having better pre-training loss (§<a href="#S4.SS2" title="4.2 Scaling of Performance with Compute ‣ 4 Experiments ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">4.2</span></a>). Our analysis identifies the auxiliary router’s performance as the primary cause. We used a smaller model, 635M MoD MoMa 4t4i, in the analysis. Initially, we validated that using the training router and selecting top <math id="S4.SS6.SSS0.Px2.p1.1.m1.1" class="ltx_Math" alttext="k_{d}" display="inline"><semantics id="S4.SS6.SSS0.Px2.p1.1.m1.1a"><msub id="S4.SS6.SSS0.Px2.p1.1.m1.1.1" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1.cmml"><mi id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.2" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1.2.cmml">k</mi><mi id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.3" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1.3.cmml">d</mi></msub><annotation-xml encoding="MathML-Content" id="S4.SS6.SSS0.Px2.p1.1.m1.1b"><apply id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.cmml" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1"><csymbol cd="ambiguous" id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.1.cmml" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1">subscript</csymbol><ci id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.2.cmml" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1.2">𝑘</ci><ci id="S4.SS6.SSS0.Px2.p1.1.m1.1.1.3.cmml" xref="S4.SS6.SSS0.Px2.p1.1.m1.1.1.3">𝑑</ci></apply></annotation-xml><annotation encoding="application/x-tex" id="S4.SS6.SSS0.Px2.p1.1.m1.1c">k_{d}</annotation></semantics></math> tokens within a batch for inference enables the 635M MoD MoMa 4t4i model to outperform the 435M MoMa 4t4i model. However, when introducing random noise to the training router selection with a noise ratio of <math id="S4.SS6.SSS0.Px2.p1.2.m2.1" class="ltx_Math" alttext="\sigma" display="inline"><semantics id="S4.SS6.SSS0.Px2.p1.2.m2.1a"><mi id="S4.SS6.SSS0.Px2.p1.2.m2.1.1" xref="S4.SS6.SSS0.Px2.p1.2.m2.1.1.cmml">σ</mi><annotation-xml encoding="MathML-Content" id="S4.SS6.SSS0.Px2.p1.2.m2.1b"><ci id="S4.SS6.SSS0.Px2.p1.2.m2.1.1.cmml" xref="S4.SS6.SSS0.Px2.p1.2.m2.1.1">𝜎</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.SS6.SSS0.Px2.p1.2.m2.1c">\sigma</annotation></semantics></math>, the 635M MoD MoMa 4t4i model starts underperforming the 435M MoMa 4t4i model when <math id="S4.SS6.SSS0.Px2.p1.3.m3.1" class="ltx_Math" alttext="\sigma" display="inline"><semantics id="S4.SS6.SSS0.Px2.p1.3.m3.1a"><mi id="S4.SS6.SSS0.Px2.p1.3.m3.1.1" xref="S4.SS6.SSS0.Px2.p1.3.m3.1.1.cmml">σ</mi><annotation-xml encoding="MathML-Content" id="S4.SS6.SSS0.Px2.p1.3.m3.1b"><ci id="S4.SS6.SSS0.Px2.p1.3.m3.1.1.cmml" xref="S4.SS6.SSS0.Px2.p1.3.m3.1.1">𝜎</ci></annotation-xml><annotation encoding="application/x-tex" id="S4.SS6.SSS0.Px2.p1.3.m3.1c">\sigma</annotation></semantics></math> exceeds 0.5%. Achieving an error rate below this threshold is impractical for auxiliary router training. Therefore, further research is necessary to develop a robust approach for practical MoD applications.</p>
</div>
<div id="S4.SS6.SSS0.Px2.p2" class="ltx_para">
<p id="S4.SS6.SSS0.Px2.p2.1" class="ltx_p">On the other hand, MoE models do not exhibit similar regression during inference, suggesting they are less sensitive to router errors. This is likely due to the MoE model’s multi-expert architecture, which allows mis-routed tokens to be processed by other experts, preserving their information. In contrast, the MoD model’s single selection switch per layer means routing mistakes cannot be remedied, leading to greater performance degradation.</p>
</div>
</section>
</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>

<section id="S5.SS0.SSS0.Px1" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Early-fusion vision-language models.</h4>

<div id="S5.SS0.SSS0.Px1.p1" class="ltx_para">
<p id="S5.SS0.SSS0.Px1.p1.1" class="ltx_p">Early fusion techniques have gained traction in multi-modal learning due to their ability to capture cross-modal interactions from the onset of processing. Several notable works have explored this paradigm. PerceiverIO <cite class="ltx_cite ltx_citemacro_citep">(Jaegle et al., <a href="#bib.bib40" title="" class="ltx_ref">2021</a>)</cite> introduced a fully attentional read-process-write architecture that operates over a modality-agnostic latent space for processing diverse inputs, including text and images. NÜWA <cite class="ltx_cite ltx_citemacro_citep">(Wu et al., <a href="#bib.bib41" title="" class="ltx_ref">2021</a>)</cite> presented a 3D-attention transformer capable of both understanding and generating text, image, and video in various combinations. CM3 <cite class="ltx_cite ltx_citemacro_citep">(Aghajanyan et al., <a href="#bib.bib42" title="" class="ltx_ref">2022</a>)</cite> adopted a causally masked transformer to learn from mixed-modal documents on the Internet, showcasing the scalability of this paradigm to large-scale pretraining. We adopt Chameleon as the base transformer architecture and demonstrate that modality-aware sparsity can effectively improve its scaling performance further.</p>
</div>
</section>
<section id="S5.SS0.SSS0.Px2" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Multi-modal representation learning.</h4>

<div id="S5.SS0.SSS0.Px2.p1" class="ltx_para">
<p id="S5.SS0.SSS0.Px2.p1.1" class="ltx_p">BEIT-3 <cite class="ltx_cite ltx_citemacro_citep">(Wang et al., <a href="#bib.bib19" title="" class="ltx_ref">2022b</a>)</cite> introduced a general-purpose multimodal foundation model that achieves state-of-the-art performance on both vision and vision-language tasks. It employs a Multiway Transformer architecture with modality-specific expert modules and uses a unified masked data modeling objective for pretraining on images, texts, and image-text pairs.</p>
</div>
</section>
<section id="S5.SS0.SSS0.Px3" class="ltx_paragraph">
<h4 class="ltx_title ltx_title_paragraph">Sparse multi-modal language models.</h4>

<div id="S5.SS0.SSS0.Px3.p1" class="ltx_para">
<p id="S5.SS0.SSS0.Px3.p1.1" class="ltx_p">Recent advancements in sparse modeling techniques have shown promising results in scaling language models efficiently. VL-MoE <cite class="ltx_cite ltx_citemacro_citep">(Shen et al., <a href="#bib.bib15" title="" class="ltx_ref">2023</a>)</cite> proposed a vision-language model that uses a sparse mixture-of-experts (MoE) to scale efficiently. It applies sparse MoE to the vision and language feed-forward networks in a unified framework, using masked data modeling objectives for pretraining. VL-MoE achieves strong performance on vision-language tasks while using fewer parameters per token than dense models, providing additional evidence for MoE training for multi-modal settings.</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">Our current MoMa implementation relies on matching the token mix ratio in the dataset with the expert mix ratio in the model to maintain load balance across GPUs. Even so, minor imbalance may still occur because there is no hard limit for a batch to deviate from that ratio at per-GPI per-iteration level. We leave further improvements in this area as future work.
</p>
</div>
<div id="S6.p2" class="ltx_para">
<p id="S6.p2.1" class="ltx_p">Expert-choice routing alleviates the expert load balancing issue during training but presents additional challenges for auto-regressive Language Models (LMs) during inference <cite class="ltx_cite ltx_citemacro_citep">(Zhou et al., <a href="#bib.bib16" title="" class="ltx_ref">2022</a>)</cite>. Although auxiliary routers comprise only a small fraction of the network’s parameters, their role is crucial. In our study, we trained the auxiliary router after completing the whole network training and limited this process to a few thousand steps, while previous work have demonstrated the possibility to jointly train such modules with the full network <cite class="ltx_cite ltx_citemacro_citep">(Raposo et al., <a href="#bib.bib10" title="" class="ltx_ref">2024</a>)</cite>. Future research should explore the architecture and training techniques for auxiliary routers to prevent them from becoming a performance bottleneck and ensure generalizability across diverse data distributions. Especially, further investigation is necessary for training mixture-of-depths architectures, including both the auxiliary routers and the original model, to ensure effective performance in causal inference scenarios.</p>
</div>
<div id="S6.p3" class="ltx_para">
<p id="S6.p3.1" class="ltx_p">In our work, we experimented only the vanilla formulation of MoD and its staged integration with MoE. We leave the investigation of other MoD variations, including modality-aware MoD to future work. In addition, batched sequence generation with mixture-of-depths (MoD) is non-trivial as unlike standard sequence generation there are dynamics shapes and dynamic updates to the KV cache for each layer as certain sequences and layers may skip different tokens. There remains further room to optimize the inference implementations for MoD models.</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>Conclusion</h2>

<div id="S7.p1" class="ltx_para">
<p id="S7.p1.1" class="ltx_p">In this work, we have introduced a series of modality-aware sparse architectures for early fusion, mixed-modality foundation models. Our approach leverages domain specificity while preserving cross-modality knowledge sharing and feature interaction. We have developed highly efficient and effective model architectures by incorporating sparsity in both the width dimension (through modality-aware mixture-of-experts) and the depth dimension (via mixture-of-depths).</p>
</div>
<div id="S7.p2" class="ltx_para">
<p id="S7.p2.1" class="ltx_p">Our best architecture, <span id="S7.p2.1.1" class="ltx_text ltx_font_typewriter">Chameleon-MoMa</span>, demonstrate significant improvements over state-of-the-art baselines. In complexity-controlled experiments, we reduce total FLOPs requirements by up to 3.7<math id="S7.p2.1.m1.1" class="ltx_Math" alttext="\times" display="inline"><semantics id="S7.p2.1.m1.1a"><mo id="S7.p2.1.m1.1.1" xref="S7.p2.1.m1.1.1.cmml">×</mo><annotation-xml encoding="MathML-Content" id="S7.p2.1.m1.1b"><times id="S7.p2.1.m1.1.1.cmml" xref="S7.p2.1.m1.1.1"></times></annotation-xml><annotation encoding="application/x-tex" id="S7.p2.1.m1.1c">\times</annotation></semantics></math>. Importantly, our experimental findings reveal that our modality-aware sparse architectures maintain an empirical scaling law. This characteristic suggests that our approach provides immediate performance benefits and a scalable framework for future developments in mixed-modal foundation model training.</p>
</div>
<div id="S7.p3" class="ltx_para">
<p id="S7.p3.1" class="ltx_p">The promising potential of <span id="S7.p3.1.1" class="ltx_text ltx_font_typewriter">Chameleon-MoMa</span> opens up several promising directions for future research. These include exploring more sophisticated routing mechanisms, investigating the impact of different sparsity patterns across modalities, and extending our approach to a broader range of modalities and tasks.</p>
</div>
</section>
<section id="Sx1" class="ltx_section" lang="en">
<h2 class="ltx_title ltx_title_section">Acknowledgments</h2>

<div id="Sx1.p1" class="ltx_para">
<p id="Sx1.p1.1" class="ltx_p">We thank Baptiste Rozière, Alexei Baevski, Naman Goyal, and Horace He for helpful discussions on mixture-of-experts model training. We thank Ram Pasunuru and Andrew Cohen for helpful feedback on model evaluation. We thank the members of the Chameleon Team for insightful discussions and feedback.</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">Gemini et al. (2023)</span>
<span class="ltx_bibblock">
Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al.

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

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

</span>
</li>
<li id="bib.bib2" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Gemini et al. (2024)</span>
<span class="ltx_bibblock">
Gemini, Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, Soroosh Mariooryad, Yifan Ding, Xinyang Geng, Fred Alcober, Roy Frostig, Mark Omernick, Lexi Walker, Cosmin Paduraru, et al.

</span>
<span class="ltx_bibblock">Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024.

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

</span>
</li>
<li id="bib.bib3" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">OpenAI et al. (2024)</span>
<span class="ltx_bibblock">
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, et al.

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

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

</span>
</li>
<li id="bib.bib4" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Lu et al. (2023)</span>
<span class="ltx_bibblock">
Jiasen Lu, Christopher Clark, Sangho Lee, Zichen Zhang, Savya Khosla, Ryan Marten, Derek Hoiem, and Aniruddha Kembhavi.

</span>
<span class="ltx_bibblock">Unified-io 2: Scaling autoregressive multimodal models with vision, language, audio, and action, 2023.

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

</span>
</li>
<li id="bib.bib5" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Chameleon Team (2024)</span>
<span class="ltx_bibblock">
Chameleon Team.

</span>
<span class="ltx_bibblock">Chameleon: Mixed-modal early-fusion foundation models, 2024.

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

</span>
</li>
<li id="bib.bib6" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Lepikhin et al. (2020)</span>
<span class="ltx_bibblock">
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.

</span>
<span class="ltx_bibblock">Gshard: Scaling giant models with conditional computation and automatic sharding, 2020.

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

</span>
</li>
<li id="bib.bib7" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Fedus et al. (2022)</span>
<span class="ltx_bibblock">
William Fedus, Barret Zoph, and Noam Shazeer.

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

</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.bib8" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Clark et al. (2022)</span>
<span class="ltx_bibblock">
Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan.

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

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

</span>
</li>
<li id="bib.bib9" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Jiang et al. (2024)</span>
<span class="ltx_bibblock">
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed.

</span>
<span class="ltx_bibblock">Mixtral of experts, 2024.

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

</span>
</li>
<li id="bib.bib10" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Raposo et al. (2024)</span>
<span class="ltx_bibblock">
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro.

</span>
<span class="ltx_bibblock">Mixture-of-depths: Dynamically allocating compute in transformer-based language models, 2024.

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

</span>
</li>
<li id="bib.bib11" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Mustafa et al. (2022)</span>
<span class="ltx_bibblock">
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby.

</span>
<span class="ltx_bibblock">Multimodal contrastive learning with limoe: the language-image mixture of experts, 2022.

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

</span>
</li>
<li id="bib.bib12" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Liang et al. (2023)</span>
<span class="ltx_bibblock">
Paul Pu Liang, Amir Zadeh, and Louis-Philippe Morency.

</span>
<span class="ltx_bibblock">Foundations and trends in multimodal machine learning: Principles, challenges, and open questions, 2023.

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

</span>
</li>
<li id="bib.bib13" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Bao et al. (2022)</span>
<span class="ltx_bibblock">
Hangbo Bao, Wenhui Wang, Li Dong, Qiang Liu, Owais Khan Mohammed, Kriti Aggarwal, Subhojit Som, and Furu Wei.

</span>
<span class="ltx_bibblock">Vlmo: Unified vision-language pre-training with mixture-of-modality-experts, 2022.

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

</span>
</li>
<li id="bib.bib14" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Wang et al. (2022a)</span>
<span class="ltx_bibblock">
Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorck, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, and Furu Wei.

</span>
<span class="ltx_bibblock">Image as a foreign language: Beit pretraining for all vision and vision-language tasks.

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

</span>
</li>
<li id="bib.bib15" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Shen et al. (2023)</span>
<span class="ltx_bibblock">
Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, and Yuxiong He.

</span>
<span class="ltx_bibblock">Scaling vision-language models with sparse mixture of experts.

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

</span>
</li>
<li id="bib.bib16" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Zhou et al. (2022)</span>
<span class="ltx_bibblock">
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon.

</span>
<span class="ltx_bibblock">Mixture-of-experts with expert choice routing, 2022.

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

</span>
</li>
<li id="bib.bib17" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Paszke et al. (2019)</span>
<span class="ltx_bibblock">
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.

</span>
<span class="ltx_bibblock">PyTorch: An imperative style, high-performance deep learning library.

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

</span>
</li>
<li id="bib.bib18" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Komatsuzaki et al. (2023)</span>
<span class="ltx_bibblock">
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby.

</span>
<span class="ltx_bibblock">Sparse upcycling: Training mixture-of-experts from dense checkpoints, 2023.

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

</span>
</li>
<li id="bib.bib19" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Wang et al. (2022b)</span>
<span class="ltx_bibblock">
Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorck, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, and Furu Wei.

</span>
<span class="ltx_bibblock">Image as a foreign language: Beit pretraining for all vision and vision-language tasks, 2022b.

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

</span>
</li>
<li id="bib.bib20" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Liu et al. (2022a)</span>
<span class="ltx_bibblock">
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo.

</span>
<span class="ltx_bibblock">Swin transformer v2: Scaling up capacity and resolution, 2022a.

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

</span>
</li>
<li id="bib.bib21" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Elhoushi et al. (2024)</span>
<span class="ltx_bibblock">
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, and Carole-Jean Wu.

</span>
<span class="ltx_bibblock">Layerskip: Enabling early exit inference and self-speculative decoding, 2024.

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

</span>
</li>
<li id="bib.bib22" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Xue et al. (2024)</span>
<span class="ltx_bibblock">
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You.

</span>
<span class="ltx_bibblock">Openmoe: An early effort on open mixture-of-experts language models, 2024.

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

</span>
</li>
<li id="bib.bib23" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Liu et al. (2022b)</span>
<span class="ltx_bibblock">
Pengbo Liu, Hailong Cao, and Tiejun Zhao.

</span>
<span class="ltx_bibblock">Gumbel-attention for multi-modal machine translation, 2022b.

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

</span>
</li>
<li id="bib.bib24" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Geng et al. (2020)</span>
<span class="ltx_bibblock">
Xinwei Geng, Longyue Wang, Xing Wang, Bing Qin, Ting Liu, and Zhaopeng Tu.

</span>
<span class="ltx_bibblock">How does selective mechanism improve self-attention networks?, 2020.

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

</span>
</li>
<li id="bib.bib25" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Zhao et al. (2023)</span>
<span class="ltx_bibblock">
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al.

</span>
<span class="ltx_bibblock">Pytorch fsdp: experiences on scaling fully sharded data parallel.

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

</span>
</li>
<li id="bib.bib26" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Gale et al. (2023)</span>
<span class="ltx_bibblock">
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia.

</span>
<span class="ltx_bibblock">Megablocks: Efficient sparse training with mixture-of-experts.

</span>
<span class="ltx_bibblock"><em id="bib.bib26.1.1" class="ltx_emph ltx_font_italic">Proceedings of Machine Learning and Systems</em>, 5:288–304, 2023.

</span>
</li>
<li id="bib.bib27" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Ansel et al. (2024)</span>
<span class="ltx_bibblock">
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, et al.

</span>
<span class="ltx_bibblock">Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation.

</span>
<span class="ltx_bibblock">2024.

</span>
</li>
<li id="bib.bib28" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Artetxe et al. (2021)</span>
<span class="ltx_bibblock">
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O’Horo, Jeff Wang, Luke Zettlemoyer, Mona T. Diab, Zornitsa Kozareva, and Ves Stoyanov.

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

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

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

</span>
</li>
<li id="bib.bib29" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Laurençon et al. (2023)</span>
<span class="ltx_bibblock">
Hugo Laurençon, Lucile Saulnier, Léo Tronchon, Stas Bekman, Amanpreet Singh, Anton Lozhkov, Thomas Wang, Siddharth Karamcheti, Alexander M Rush, Douwe Kiela, et al.

</span>
<span class="ltx_bibblock">Obelisc: An open web-scale filtered dataset of interleaved image-text documents.

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

</span>
</li>
<li id="bib.bib30" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Bisk et al. (2020)</span>
<span class="ltx_bibblock">
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al.

</span>
<span class="ltx_bibblock">Piqa: Reasoning about physical commonsense in natural language.

</span>
<span class="ltx_bibblock">In <em id="bib.bib30.1.1" class="ltx_emph ltx_font_italic">Proceedings of the AAAI conference on artificial intelligence</em>, pages 7432–7439, 2020.

</span>
</li>
<li id="bib.bib31" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Sap et al. (2019)</span>
<span class="ltx_bibblock">
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi.

</span>
<span class="ltx_bibblock">Socialiqa: Commonsense reasoning about social interactions.

</span>
<span class="ltx_bibblock"><em id="bib.bib31.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1904.09728</em>, 2019.

</span>
</li>
<li id="bib.bib32" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Zellers et al. (2019)</span>
<span class="ltx_bibblock">
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.

</span>
<span class="ltx_bibblock">Hellaswag: Can a machine really finish your sentence?

</span>
<span class="ltx_bibblock"><em id="bib.bib32.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1905.07830</em>, 2019.

</span>
</li>
<li id="bib.bib33" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Sakaguchi et al. (2021)</span>
<span class="ltx_bibblock">
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi.

</span>
<span class="ltx_bibblock">Winogrande: An adversarial winograd schema challenge at scale.

</span>
<span class="ltx_bibblock"><em id="bib.bib33.1.1" class="ltx_emph ltx_font_italic">Communications of the ACM</em>, 64(9):99–106, 2021.

</span>
</li>
<li id="bib.bib34" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Clark et al. (2018)</span>
<span class="ltx_bibblock">
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.

</span>
<span class="ltx_bibblock">Think you have solved question answering? try arc, the ai2 reasoning challenge.

</span>
<span class="ltx_bibblock"><em id="bib.bib34.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1803.05457</em>, 2018.

</span>
</li>
<li id="bib.bib35" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Mihaylov et al. (2018)</span>
<span class="ltx_bibblock">
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal.

</span>
<span class="ltx_bibblock">Can a suit of armor conduct electricity? a new dataset for open book question answering.

</span>
<span class="ltx_bibblock"><em id="bib.bib35.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1809.02789</em>, 2018.

</span>
</li>
<li id="bib.bib36" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Clark et al. (2019)</span>
<span class="ltx_bibblock">
Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova.

</span>
<span class="ltx_bibblock">Boolq: Exploring the surprising difficulty of natural yes/no questions.

</span>
<span class="ltx_bibblock"><em id="bib.bib36.1.1" class="ltx_emph ltx_font_italic">arXiv preprint arXiv:1905.10044</em>, 2019.

</span>
</li>
<li id="bib.bib37" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Lin et al. (2014)</span>
<span class="ltx_bibblock">
Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.

</span>
<span class="ltx_bibblock">Microsoft coco: Common objects in context.

</span>
<span class="ltx_bibblock">In <em id="bib.bib37.1.1" class="ltx_emph ltx_font_italic">European conference on computer vision</em>, pages 740–755. Springer, 2014.

</span>
</li>
<li id="bib.bib38" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Plummer et al. (2015)</span>
<span class="ltx_bibblock">
Bryan A Plummer, Liwei Wang, Chris M Cervantes, Juan C Caicedo, Julia Hockenmaier, and Svetlana Lazebnik.

</span>
<span class="ltx_bibblock">Flickr30k entities: Collecting region-to-phrase correspondences for richer image-to-sentence models.

</span>
<span class="ltx_bibblock">In <em id="bib.bib38.1.1" class="ltx_emph ltx_font_italic">Proceedings of the IEEE international conference on computer vision</em>, pages 2641–2649, 2015.

</span>
</li>
<li id="bib.bib39" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Goyal et al. (2017)</span>
<span class="ltx_bibblock">
Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh.

</span>
<span class="ltx_bibblock">Making the V in VQA matter: Elevating the role of image understanding in visual question answering.

</span>
<span class="ltx_bibblock">In <em id="bib.bib39.1.1" class="ltx_emph ltx_font_italic">Proceedings of the IEEE conference on computer vision and pattern recognition</em>, pages 6904–6913, 2017.

</span>
</li>
<li id="bib.bib40" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Jaegle et al. (2021)</span>
<span class="ltx_bibblock">
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al.

</span>
<span class="ltx_bibblock">Perceiver io: A general architecture for structured inputs &amp; outputs.

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

</span>
</li>
<li id="bib.bib41" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Wu et al. (2021)</span>
<span class="ltx_bibblock">
Chenfei Wu, Jian Liang, Lei Ji, Fan Yang, Yuejian Fang, Daxin Jiang, and Nan Duan.

</span>
<span class="ltx_bibblock">Nüwa: Visual synthesis pre-training for neural visual world creation.

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

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

</span>
</li>
<li id="bib.bib42" class="ltx_bibitem">
<span class="ltx_tag ltx_role_refnum ltx_tag_bibitem">Aghajanyan et al. (2022)</span>
<span class="ltx_bibblock">
Armen Aghajanyan, Bernie Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, et al.

</span>
<span class="ltx_bibblock">Cm3: A causal masked multimodal model of the internet.

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

</span>
</li>
</ul>
</section>
<div class="ltx_pagination ltx_role_newpage"></div>
<section id="A1" class="ltx_appendix" lang="en">
<h2 class="ltx_title ltx_title_appendix">
<span class="ltx_tag ltx_tag_appendix">Appendix A </span>Additional Implementation Details</h2>

<section id="A1.SS1" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">A.1 </span>Pre-training</h3>

<div id="A1.SS1.p1" class="ltx_para">
<p id="A1.SS1.p1.1" class="ltx_p">Table <a href="#A1.T5" title="Table 5 ‣ A.1 Pre-training ‣ Appendix A Additional Implementation Details ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">5</span></a> reports the hyperparameters of our pre-training aproach.</p>
</div>
<figure id="A1.T5" class="ltx_table">
<div id="A1.T5.2" class="ltx_inline-block ltx_align_center ltx_transformed_outer" style="width:424.1pt;height:66.6pt;vertical-align:-0.9pt;"><span class="ltx_transformed_inner" style="transform:translate(-23.6pt,3.6pt) scale(0.9,0.9) ;">
<table id="A1.T5.2.1" class="ltx_tabular ltx_align_middle">
<tr id="A1.T5.2.1.1" class="ltx_tr">
<td id="A1.T5.2.1.1.1" class="ltx_td ltx_align_left ltx_border_tt">
<span id="A1.T5.2.1.1.1.1" class="ltx_inline-block ltx_parbox ltx_align_middle" style="width:34.7pt;">
<span id="A1.T5.2.1.1.1.1.1" class="ltx_p">model size</span>
</span>
</td>
<td id="A1.T5.2.1.1.2" class="ltx_td ltx_align_center ltx_border_tt">peak lr</td>
<td id="A1.T5.2.1.1.3" class="ltx_td ltx_align_center ltx_border_tt">end lr</td>
<td id="A1.T5.2.1.1.4" class="ltx_td ltx_align_center ltx_border_tt">lr scheduler</td>
<td id="A1.T5.2.1.1.5" class="ltx_td ltx_align_center ltx_border_tt">warm-up</td>
<td id="A1.T5.2.1.1.6" class="ltx_td ltx_align_center ltx_border_tt"># steps</td>
<td id="A1.T5.2.1.1.7" class="ltx_td ltx_align_center ltx_border_tt">batch size (tokens)</td>
<td id="A1.T5.2.1.1.8" class="ltx_td ltx_align_center ltx_border_tt">
<span id="A1.T5.2.1.1.8.1" class="ltx_inline-block ltx_parbox ltx_align_middle" style="width:34.7pt;">
<span id="A1.T5.2.1.1.8.1.1" class="ltx_p">model parallel</span>
</span>
</td>
<td id="A1.T5.2.1.1.9" class="ltx_td ltx_align_center ltx_border_tt">seq len</td>
</tr>
<tr id="A1.T5.2.1.2" class="ltx_tr">
<td id="A1.T5.2.1.2.1" class="ltx_td ltx_align_left ltx_border_t">60M</td>
<td id="A1.T5.2.1.2.2" class="ltx_td ltx_align_center ltx_border_t">1e-4</td>
<td id="A1.T5.2.1.2.3" class="ltx_td ltx_align_center ltx_border_t">1e-6</td>
<td id="A1.T5.2.1.2.4" class="ltx_td ltx_align_center ltx_border_t">linear</td>
<td id="A1.T5.2.1.2.5" class="ltx_td ltx_align_center ltx_border_t">4000</td>
<td id="A1.T5.2.1.2.6" class="ltx_td ltx_align_center ltx_border_t">160000</td>
<td id="A1.T5.2.1.2.7" class="ltx_td ltx_align_center ltx_border_t">6291456</td>
<td id="A1.T5.2.1.2.8" class="ltx_td ltx_align_center ltx_border_t">1</td>
<td id="A1.T5.2.1.2.9" class="ltx_td ltx_align_center ltx_border_t">4096</td>
</tr>
<tr id="A1.T5.2.1.3" class="ltx_tr">
<td id="A1.T5.2.1.3.1" class="ltx_td ltx_align_left">435M</td>
<td id="A1.T5.2.1.3.2" class="ltx_td ltx_align_center">1e-4</td>
<td id="A1.T5.2.1.3.3" class="ltx_td ltx_align_center">1e-6</td>
<td id="A1.T5.2.1.3.4" class="ltx_td ltx_align_center">linear</td>
<td id="A1.T5.2.1.3.5" class="ltx_td ltx_align_center">4000</td>
<td id="A1.T5.2.1.3.6" class="ltx_td ltx_align_center">160000</td>
<td id="A1.T5.2.1.3.7" class="ltx_td ltx_align_center">6291456</td>
<td id="A1.T5.2.1.3.8" class="ltx_td ltx_align_center">1</td>
<td id="A1.T5.2.1.3.9" class="ltx_td ltx_align_center">4096</td>
</tr>
<tr id="A1.T5.2.1.4" class="ltx_tr">
<td id="A1.T5.2.1.4.1" class="ltx_td ltx_align_left ltx_border_bb">1.4B</td>
<td id="A1.T5.2.1.4.2" class="ltx_td ltx_align_center ltx_border_bb">1e-4</td>
<td id="A1.T5.2.1.4.3" class="ltx_td ltx_align_center ltx_border_bb">1e-6</td>
<td id="A1.T5.2.1.4.4" class="ltx_td ltx_align_center ltx_border_bb">linear</td>
<td id="A1.T5.2.1.4.5" class="ltx_td ltx_align_center ltx_border_bb">4000</td>
<td id="A1.T5.2.1.4.6" class="ltx_td ltx_align_center ltx_border_bb">120000</td>
<td id="A1.T5.2.1.4.7" class="ltx_td ltx_align_center ltx_border_bb">8388608</td>
<td id="A1.T5.2.1.4.8" class="ltx_td ltx_align_center ltx_border_bb">1</td>
<td id="A1.T5.2.1.4.9" class="ltx_td ltx_align_center ltx_border_bb">4096</td>
</tr>
</table>
</span></div>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="A1.T5.3.1.1" class="ltx_text" style="font-size:90%;">Table 5</span>: </span><span id="A1.T5.4.2" class="ltx_text" style="font-size:90%;">Hyperparameters used for pre-training our models. We used the same pre-training hyperparameters for models compute-matched to the same dense size.</span></figcaption>
</figure>
</section>
<section id="A1.SS2" class="ltx_subsection">
<h3 class="ltx_title ltx_title_subsection">
<span class="ltx_tag ltx_tag_subsection">A.2 </span>Auxiliary Router Training</h3>

<div id="A1.SS2.p1" class="ltx_para">
<p id="A1.SS2.p1.1" class="ltx_p">Table <a href="#A1.T6" title="Table 6 ‣ A.2 Auxiliary Router Training ‣ Appendix A Additional Implementation Details ‣ MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts" class="ltx_ref"><span class="ltx_text ltx_ref_tag">6</span></a> reports the hyperparameters for training the auxillary routers for both Mixture of Depth and Mixture of Experts.</p>
</div>
<figure id="A1.T6" class="ltx_table">
<div id="A1.T6.2" class="ltx_inline-block ltx_align_center ltx_transformed_outer" style="width:435.1pt;height:114.3pt;vertical-align:-0.9pt;"><span class="ltx_transformed_inner" style="transform:translate(-24.2pt,6.3pt) scale(0.9,0.9) ;">
<table id="A1.T6.2.1" class="ltx_tabular ltx_align_middle">
<tr id="A1.T6.2.1.1" class="ltx_tr">
<td id="A1.T6.2.1.1.1" class="ltx_td ltx_align_left ltx_border_tt">model size</td>
<td id="A1.T6.2.1.1.2" class="ltx_td ltx_align_center ltx_border_tt">peak lr</td>
<td id="A1.T6.2.1.1.3" class="ltx_td ltx_align_center ltx_border_tt">end lr</td>
<td id="A1.T6.2.1.1.4" class="ltx_td ltx_align_center ltx_border_tt">lr scheduler</td>
<td id="A1.T6.2.1.1.5" class="ltx_td ltx_align_center ltx_border_tt">warm-up</td>
<td id="A1.T6.2.1.1.6" class="ltx_td ltx_align_center ltx_border_tt"># steps</td>
<td id="A1.T6.2.1.1.7" class="ltx_td ltx_align_center ltx_border_tt">batch size (tokens)</td>
<td id="A1.T6.2.1.1.8" class="ltx_td ltx_align_center ltx_border_tt">seq len</td>
</tr>
<tr id="A1.T6.2.1.2" class="ltx_tr">
<td id="A1.T6.2.1.2.1" class="ltx_td ltx_align_left ltx_border_t">2.3B MoD MoE 4t4i</td>
<td id="A1.T6.2.1.2.2" class="ltx_td ltx_align_center ltx_border_t">1e-4</td>
<td id="A1.T6.2.1.2.3" class="ltx_td ltx_align_center ltx_border_t">1e-6</td>
<td id="A1.T6.2.1.2.4" class="ltx_td ltx_align_center ltx_border_t">linear</td>
<td id="A1.T6.2.1.2.5" class="ltx_td ltx_align_center ltx_border_t">2000</td>
<td id="A1.T6.2.1.2.6" class="ltx_td ltx_align_center ltx_border_t">5000</td>
<td id="A1.T6.2.1.2.7" class="ltx_td ltx_align_center ltx_border_t">4194304</td>
<td id="A1.T6.2.1.2.8" class="ltx_td ltx_align_center ltx_border_t">4096</td>
</tr>
<tr id="A1.T6.2.1.3" class="ltx_tr">
<td id="A1.T6.2.1.3.1" class="ltx_td ltx_align_left">1.4B MoE 4t4i</td>
<td id="A1.T6.2.1.3.2" class="ltx_td ltx_align_center">1e-4</td>
<td id="A1.T6.2.1.3.3" class="ltx_td ltx_align_center">1e-6</td>
<td id="A1.T6.2.1.3.4" class="ltx_td ltx_align_center">linear</td>
<td id="A1.T6.2.1.3.5" class="ltx_td ltx_align_center">2000</td>
<td id="A1.T6.2.1.3.6" class="ltx_td ltx_align_center">10000</td>
<td id="A1.T6.2.1.3.7" class="ltx_td ltx_align_center">4194304</td>
<td id="A1.T6.2.1.3.8" class="ltx_td ltx_align_center">4096</td>
</tr>
<tr id="A1.T6.2.1.4" class="ltx_tr">
<td id="A1.T6.2.1.4.1" class="ltx_td ltx_align_left">1.4B MoE EC8</td>
<td id="A1.T6.2.1.4.2" class="ltx_td ltx_align_center">1e-4</td>
<td id="A1.T6.2.1.4.3" class="ltx_td ltx_align_center">1e-6</td>
<td id="A1.T6.2.1.4.4" class="ltx_td ltx_align_center">linear</td>
<td id="A1.T6.2.1.4.5" class="ltx_td ltx_align_center">2000</td>
<td id="A1.T6.2.1.4.6" class="ltx_td ltx_align_center">10000</td>
<td id="A1.T6.2.1.4.7" class="ltx_td ltx_align_center">4194304</td>
<td id="A1.T6.2.1.4.8" class="ltx_td ltx_align_center">4096</td>
</tr>
<tr id="A1.T6.2.1.5" class="ltx_tr">
<td id="A1.T6.2.1.5.1" class="ltx_td ltx_align_left">635M MoD MoE 4t4i</td>
<td id="A1.T6.2.1.5.2" class="ltx_td ltx_align_center">1e-4</td>
<td id="A1.T6.2.1.5.3" class="ltx_td ltx_align_center">1e-6</td>
<td id="A1.T6.2.1.5.4" class="ltx_td ltx_align_center">linear</td>
<td id="A1.T6.2.1.5.5" class="ltx_td ltx_align_center">1000</td>
<td id="A1.T6.2.1.5.6" class="ltx_td ltx_align_center">5000</td>
<td id="A1.T6.2.1.5.7" class="ltx_td ltx_align_center">4194304</td>
<td id="A1.T6.2.1.5.8" class="ltx_td ltx_align_center">4096</td>
</tr>
<tr id="A1.T6.2.1.6" class="ltx_tr">
<td id="A1.T6.2.1.6.1" class="ltx_td ltx_align_left">435M MoE 4t4i</td>
<td id="A1.T6.2.1.6.2" class="ltx_td ltx_align_center">1e-4</td>
<td id="A1.T6.2.1.6.3" class="ltx_td ltx_align_center">1e-6</td>
<td id="A1.T6.2.1.6.4" class="ltx_td ltx_align_center">linear</td>
<td id="A1.T6.2.1.6.5" class="ltx_td ltx_align_center">1000</td>
<td id="A1.T6.2.1.6.6" class="ltx_td ltx_align_center">5000</td>
<td id="A1.T6.2.1.6.7" class="ltx_td ltx_align_center">4194304</td>
<td id="A1.T6.2.1.6.8" class="ltx_td ltx_align_center">4096</td>
</tr>
<tr id="A1.T6.2.1.7" class="ltx_tr">
<td id="A1.T6.2.1.7.1" class="ltx_td ltx_align_left ltx_border_bb">435M MoE EC8</td>
<td id="A1.T6.2.1.7.2" class="ltx_td ltx_align_center ltx_border_bb">1e-4</td>
<td id="A1.T6.2.1.7.3" class="ltx_td ltx_align_center ltx_border_bb">1e-6</td>
<td id="A1.T6.2.1.7.4" class="ltx_td ltx_align_center ltx_border_bb">linear</td>
<td id="A1.T6.2.1.7.5" class="ltx_td ltx_align_center ltx_border_bb">1000</td>
<td id="A1.T6.2.1.7.6" class="ltx_td ltx_align_center ltx_border_bb">5000</td>
<td id="A1.T6.2.1.7.7" class="ltx_td ltx_align_center ltx_border_bb">4194304</td>
<td id="A1.T6.2.1.7.8" class="ltx_td ltx_align_center ltx_border_bb">4096</td>
</tr>
</table>
</span></div>
<figcaption class="ltx_caption ltx_centering"><span class="ltx_tag ltx_tag_table"><span id="A1.T6.3.1.1" class="ltx_text" style="font-size:90%;">Table 6</span>: </span><span id="A1.T6.4.2" class="ltx_text" style="font-size:90%;">Hyperparameters for aux router training.</span></figcaption>
</figure>
<div class="ltx_pagination ltx_role_newpage"></div>
</section>
</section>
</article>
</div>
<div class="ar5iv-footer"><a href="/html/2407.21769" 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/2407.21770" 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+2407.21770">Report<br>an issue</a>
    <a href="https://arxiv.org/abs/2407.21770" class="ar5iv-text-button arxiv-ui-theme">View&nbsp;original<br>on&nbsp;arXiv</a><a href="/html/2407.21771" 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 Mon Aug  5 15:48:33 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>

2407.21770
We introduce MoMa, a novel modality-aware mixture-of-experts (MoE)  architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7overall, with 2.6for text and 5.2for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3× overall FLOPs savings (3for text, 2.8for image) . Combining MoMa with mixture-of-depths (MoD)  further improves pretraining FLOPs savings to 4.2overall (text: 3.4, image: 5.3) ,
although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa’s potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.
abstract
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
Abstract
2407.21770
1IntroductionAuto-regressive mixed-modal foundation models have shown significant promise in applications requiring the processing of mixed-modal inputs and the generation of mixed-modal outputs. These models have demonstrated remarkable capabilities in tasks ranging from visual question answering to multimodal content generation, pushing the boundaries of AI’s ability to understand and interact with our inherently multimodal world.While a popular architecture design for mixed-modal foundation models involves fusing modality-specific encoders or decoders, this approach can limit the model’s ability to integrate information across modalities and generate content with interleaved modalities. To address this limitation,Chameleonrecently introduced a single transformer architecture with a next-token prediction objective to model mixed-modal sequences consisting of discrete image and text tokens, allowing for seamless reasoning and generation across modalities.Chameleon, pretrained on approximately 10 trillion mixed-modal tokens, has demonstrated broad vision and language capabilities across various downstream tasks. Notably, it surpasses commercial baselines such as Gemini 1.0 Pro111https://ai.google.dev/and GPT-4V222gpt-4-vision-preview:https://openai.com/api/in generating mixed-modal long-form responses. However, scaling such mixed-modal early-fusion foundation models to greater capacities presents significant computational challenges.To address these challenges, we investigate the application of routed sparse architectures. These architectures have previously shown effectiveness in scaling language and vision-specific foundation models, as well as multimodal contrastive learning. However, their application to mixed-modal early-fusion models presents unique opportunities and challenges.The insight driving our approach is the inherent heterogeneity of modalities: text and image tokens possess distinct information densities and redundancy patterns. While we integrate these tokens into a unified early-fusion architecture, we propose further optimizing this framework by incorporating modality-specific modules. This concept, which we termmodality-aware sparsity(MaS) , enables models to better capture features specific to each modality while maintaining strong cross-modality integration through partial parameter sharing and attention mechanisms. Previous work such as VLMo, BEiT-3and VL-MoEhave applied mixture-of-modality-experts (MoME)  to train vision-language encoders and masked language models. We apply this type of sparse component to the mixed-modal, early-fusion language modeling setting.Specifically, we adoptChameleonas the base transformer architecture and introduce sparsity along two dimensions:1.Width: We applymixture-of-experts(MoE) where tokens are routed across a set of feed-forward blocks (experts)  at each layer. Crucially, we divide these experts into modality-specific groups, with each group processing only tokens of its designated modality while learned routing is done within each group for semantics-based adaptivity. We dub this approachmixture of modality-aware experts(MoMa) .2.Depth: We applymixture-of-depths(MoD) where tokens can selectively skip the attention and feed-forward computation at certain layers.For both dimensions, we adopt expert-choice routingto ensure load balancing and maintain a static computation graph, both important for high training throughput.We conduct extensive FLOPs-controlled experiments comparing our proposed architecture with the dense baseline and multiple sparse variations. With a 1-trillion-token training budget,Chameleon-MoMa1.4B employing 4 text experts and 4 image experts achieves a substantial 3.7overall FLOPs savings (text: 2.6, image: 5.2)  compared to the 1.4B compute-matched dense baseline as measured by pre-training loss,
while maintaining a relatively modest -17% throughput reduction.
In contrast, the standard expert-choice MoE with 8 mixed-modal experts achieves 3FLOPs savings (text: 3, image: 2.8)  under the same setting, with -9% throughput reduction.
Further combination with MoD (Chameleon-MoMaD)  boosts the FLOPs savings to 4.2overall (text: 3.4, image: 5.3)  as measured by pre-training loss. However, the auto-regressive inference performance of the MoD model is compromised due to its sensitivity to routing accuracy, resulting in subpar performance compared toChameleon-MoMa(§4.6) .We also demonstrate that it is possible to enhance the performance ofChameleon-MoMausing a modality-untied upcycling technique. In particular, we can initialize the model with a seed sparse architecture consisting of 1 expert per modality, and effectively improve router learning for each modality. By training the seed model for only 10k steps, the resulting model achieves additional FLOP reduction while maintaining or improving performance.Our main contributions can be summarized as follows:1.We introduce mixture of modality-aware experts (MoMa)  for mixed-modal early-fusion models, allowing for more efficient parameter allocation and utilization.2.We investigated sparse scaling along both the width and depth dimensions of transformers, demonstrating substantial speedup of pre-training loss convergence can be realized in each direction. However, effectively combining these two approaches for better performance in an autor-regressive inference setup remains an open research challenge.3.We present a modality-untied upcycling technique that further enhances the efficiency of our sparse architecture.4.We provide extensive empirical results and analysis, offering insights into our approach’s scaling behavior and efficiency gains.
introduction
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
1Introduction
2407.21770
2ModelFigure 1:Overview of our proposed multimodal early-fusion architecture. (a)  The early-fusion architecture processes interleaved text and image data as a unified token sequence using a transformer model. (b)  Detailed architecture of the mixture of modality-aware experts (MoMa)  transformer block, where layer norms and the residual connection of attention layer are omitted for readability. The input sequence of tokens are first routed to different expert groups based on their modality, undergo learned routing within each group, and then have their output restored to a tensor of the original input shape of the FFN module.2.1Early FusionOur model builds upon the early fusion architecture introduced byChameleon, which represents images and text as a series of discrete tokens within a unified transformer.
The core ofChameleonis a transformer-based model that applies self-attention mechanisms over the combined sequence of image and text tokens. This allows the model to capture complex relationships between and within modalities. The model is trained using a next-token prediction objective, learning to generate both text and image tokens autoregressively.InChameleon, images are tokenized using a learned image tokenizer that encodes a 512512 image into 1024 discrete tokens from a codebook of size 8192. Text is tokenized using a BPE tokenizer with a vocabulary size of 65,536, which includes the 8192 image codebook tokens. This unified tokenization scheme enables the model to process arbitrary sequences of interleaved image and text tokens.By adopting this early fusion approach, our model inherits several key advantages:1.Unified representation:The model learns a shared representation space for images and text, facilitating cross-modal reasoning and generation.2.Flexibility:The architecture can handle arbitrary sequences of images and text, enabling diverse multimodal tasks such as image captioning, visual dialogue, and mixed-modal document generation.3.Scalability:The token-based approach allows for uniform processing of both text and image data, enabling efficient scaling to larger model sizes and diverse datasets. This is evidenced byChameleon’s successful training on approximately 10 trillion mixed-modal tokens.4.End-to-end learning:The entire auto-regressive model, is trained end-to-end, allowing for joint optimization of the representation and task-specific performance.Building on this foundation (Figure1a) , our work introduces modality-aware sparsity techniques to further enhance the efficiency and performance of early fusion models, as detailed in the following sections. These techniques aim to address the computational challenges of scaling early fusion models while maintaining their powerful cross-modal reasoning capabilities.2.2Width Scaling: Mixture of Modality-Aware ExpertsWe propose a width scaling approach that incorporates modality-aware block sparsity in the feed-forward module, extending the standard mixture-of-experts (MoE)  architecture. The insight driving this approach is that tokens of different modalities have distinct characteristics and information densities. By creating separate expert groups for each modality, we allow the model to develop specialized processing pathways while maintaining the ability to integrate information across modalities.We describe the key components in the mixture of modality-aware experts (MoMa)  formulation below (Figure1b) .Modality-Specific Expert Groups.We divide the experts in each MoE layer into distinct groups, each specialized in processing tokens from a specific modality: one group for processing text tokens and another for image tokens333While we focus on text and image modalities in this work, our formulation is designed to be easily extensible to accommodate an arbitrary number of modalities.. This separation allows each group to specialize in features relevant to its respective modality.By implementing modality-aware block sparsity, we aim to achieve several benefits:•Improved Efficiency: by routing tokens to modality-specific experts, we reduce the computational overhead of processing tokens with experts not specialized for their modality.•Enhanced Specialization: modality-specific expert groups can develop more refined features relevant to their respective modalities.•Maintained Cross-Modal Integration: despite the separation into modality-specific groups, the model still can integrate information across modalities through the shared self-attention mechanisms in non-MoE layers.Hierarchical Routing.We adopt a token-based routing mechanism. For an input token, our routing mechanism operates in two stages.1.Modality-aware routing: tokens are first routed to their corresponding modality-specific expert group based on their modality (textor image) .2.Intra-modality routing: within each modality-specific expert group, tokens are then routed to specific experts using a learned routing function. Specifically, we use a projection matrix, to compute the token-to-expert affinity scores, following other established MoE formulations.Expert choice.Within each modality group, we implemented expert-choice (EC)  routing, where each expert has a fixed bucket size and processes the top-tokens in a batch. We set,
whereis a capacity factor andis the total number of tokens of modalityin a batch.
As a result, each token can be routed to a variable number of experts.EC routing ensures balanced expert utilization during training and eliminates needing a separate load-balancing loss term. Maintaining a single loss term promotes more stable optimization and a more straightforward training process.
However, EC routing compromises causality in auto-regressive language modeling, as each expert selects the top tokens to process in a batch by comparing token scores across all tokens. We use a combination of two techniques to address this issue and enable expert-choice training for auto-regressive LMs.1.We use Sigmoid as the non-linearity in the router scoring function, enabling independent calculation of token-to-expert affinity scores for each token.2.We introduce auxiliary routers, inspired by, which predict the likelihood of an expert selecting a token solely based on its hidden state representation. These routers are trained after the main model training is completed and employed during inference to ensure causality. We discuss the details of the auxiliary routers in §2.3.In summary, the MoMa module for an input tokencan be formally defined as:(1) We further apply residual connection and the Swin Transformer normalizationpost MoMa computation. Our experiments, detailed in later sections, demonstrate that MoMa significantly improves efficiency and performance compared to dense baselines and standard MoE architectures.2.3Mixture-of-DepthsFigure 2:Architecture of transformer layer consisting of MoMa combined with mixture-of-depths (MoD) .We further investigate introducing sparsity in the depth dimension. Prior work explores sparisty in depth through either stochastic layer dropor through learnable routers. We focus on learnable routers, and incorporate the recently proposed mixture-of-depths (MoD)  technique.
Specifically, in each MoD layer, we integrate MoD prior to any mixture-of-experts (MoE)  routing, ensuring it is applied to the full batch before modality split (Figure2) .444This corresponding to the staged MoDE implementation presented in. We chose this implementation because it also introduces sparsity in attention, whereas the integrated MoDE implementation proposed in the same paper always computes full attention.Following, for each MoD layer, we use a projection matrixto compute the token-to-layer affinity score, followed by a Sigmoid non-linearity(2) Analogous to expert-choice routing in MoE, we set a fixed capacity oftokens, selected from the top-scoring tokens in a batch. We set, whereis a capacity factor chosen empirically. In practice, we first fixed the MoD layer interval and capacity factor, then adjust the total number of transformer layers to ensure the resulting architecture has a comparable FLOPs per token to the base architecture.2.4InferenceWe cannot directly apply the expert-choice routing for MoE and layer-choice routing for MoD during inference time, as the top-token selection within a batch breaks causality. Inspired by, to ensure causality during inference, we introduce auxiliary routers, to predict the likelihood of a token being selected by an expert or a layer solely based on its hidden representation.Formally, for each MoE layer, we introduce an auxiliary router(3) whereand.For each MoD layer, we introduce an auxiliary router(4) whereand.We employ a two-stage training approach, where the main model and auxiliary routers are trained separately. First, we train the main model to convergence. Then, we train the auxiliary routers using binary cross-entropy loss, supervised by the ground-truth top-routing assignments computed over an entire batch. At inference time, the main routers are only used to generate the weight values, and tokens are selectively routed to an expert or layer based on thresholding on the auxiliary router score (> 0.5) .2.5UpcyclingTraining MoE architectures with learnable routers from scratch presents unique challenges in optimizing both the representation space and routing mechanism.
The insight we identified is that MoE routers are responsible for partitioning the representation space for each expert. However, this representation space is sub-optimal in the early stages of model training, leading to a sub-optimally trained routing function.To address this limitation of router training, we propose anupcyclingapproach, inspired by.
Specifically, we begin by training an architecture consisting of 1 FFN expert per modality. After a predetermined number of steps, we upcycle this model by converting each modality-specific FFN into an expert-choice MoE module, initializing each expert with the expert trained from the first stage.
We reset the learning rate scheduler while preserving the data loader state from the previous stage, thereby ensuring the second stage training is exposed to refreshed data.To promote expert specialization, we augment the MoE routing function with Gumbel noise, allowing our router to differentiably sample experts. This is expressed in Equation5:(5) whereandare independent Gumbel noise samples.The upcycling approach, combined with the Gumbel-Sigmoid technique, allows us to overcome the limitations of learned routers and achieve improved performance in our modality-aware sparse architecture.
In practice, we found that a short period of 10k-20k steps in the first-stage training significantly improves model training efficiency and stability, aligning with the findings of.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
2Model
2407.21770
3Efficiency OptimizationTo facilitate the distributed training of mixture of modality-aware experts (MoMa) , we employ Fully Sharded Data Parallel (FSDP) . However, this approach presents unique efficiency challenges compared to vanilla MoEs. This section discusses these challenges and our strategies to address them.3.1Load BalancingWithout constraints, load imbalance can occur in our system because the ratio of text to image tokens can vary significantly across different GPUs and iterations. Imbalances can create a cascading straggler effect, delaying weight prefetching for subsequent layers and gradient releases from previous layers. This effectively bounds the training latency by the maximum time required to process text and image experts across all GPUs within a batch.To mitigate these issues, we developed a balanced data mix that aligns the text-to-image data ratio with the expert ratio on each GPU. This approach ensures load balancing in expectation. While alternative rebalancing algorithms, such as token redistribution at each FFN layer, are possible, they may introduce additional communication overhead.3.2Efficient Expert ExecutionWe explored several strategies to efficiently execute experts for different modalities. The first approach restricts to homogeneous experts across modalities and prohibits routing of text tokens to image experts and vice versa. This method allows processing of all tokens and modalities simultaneously, provided all experts share the same token count. Alternatively, we could enhance execution efficiency by employing block sparsity, which offers similar benefits to the first approach without requiring perfect expert token balance. Additionally, we considered running experts from different modalities sequentially when the number of modalities is limited. This approach allows better overlap of computation from previous modality experts with weight prefetching for subsequent modality experts, alleviating memory pressure. It also removes assumptions about expert load balance.Given that our experiments involve a sufficiently large token count per GPU, hardware utilization is not a major concern even with multiple batched matrix multiplications. Thus, we find the sequential approach to be a clean and flexible choice for our experimental environment at the current scale.3.3Additional OptimizationsWe implemented several optimizations to further enhance throughput. These include generic optimizations such as gradient communication quantization and automatic GPU kernel fusion, as well as graph optimizations viatorch.compile. Additionally, we developed MoMa-specific optimizations, including the reuse of modality token indices across different layers to minimize device synchronization between CPU and GPU555This optimization is incompatible with MoD in its current form and was not used in our final experiments to ensure fair cross-comparison across model variations. However, it can be modified to eliminate device synchronization with permutation while tracking live token counts in each MoD layer.. We also consolidated per-layer stats communication and moved these operations off the critical path of training.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
3Efficiency Optimization
2407.21770
4Experiments4.1SetupWe use the same pre-training dataset and preprocessing as. To assess scaling performance, we train all models with over 1 trillion tokens. Unless specified otherwise, we employ a sequence length of 4096 tokens and a model parallel size of 1. Our training regimen includes a peak learning rate of, a 4000-step warm-up period, and linear annealing of the learning rate to 1% of its peak value.
For all MoE architectures, we implement MoE in every layer, setting each expert’s training capacitytoto maintain FLOPs per token comparable to the base dense model. In MoD architectures, we implement MoD in alternating layers, beginning with layer 0, using a layer capacity factorof 25%. To achieve FLOPs parity with the base dense model, we increase the total layer count while maintaining a constant hidden dimension. Table1provides detailed configurations for our dense and sparse models. Additional pre-training details are available in AppendixA.1.For model comparison, we report training losses. Given that our 1-trillion-token training budget covers less than one epoch of our extensive pre-training data, we use training loss as a proxy for validation performance. Our use of expert-choice routing in both MoE and MoD modules introduces a caveat: the training loss calculation compromises causality, as token selection considers the top proportion from a batch, including future tokens. We address this in §4.6by reporting test-time performances using auxiliary routers, demonstrating result generalization to validation settings and causal scenarios.DenseMoEMoDMoDEparamsparamsparamsparams90M851220488210M810.125110M14512820.25317M435M2410244096161.9B810.125635M4010241620.253B1.4B2420488192167.1B810.1252.3B3223041820.2512BTable 1:Specifications of dense and sparse architectures used in our experiments. Architectures listed in the same row are compute-matched with the same amount of active FLOPs per token.params: total number of parameters;: number of layers;: transformer hidden dimension;: feed-forward module hidden dimension;: number of attention heads;: number of experts;: MoE/D layer interval;: MoE expert capacity/MoD layer capacity.Figure 3:Scaling of performance with compute. We consider three dense model sizes – 90M, 435M and 1.4B parameters – and their compute-matched sparse variations. We highlight three configurations of mixture-of-experts (MoE) : 1)  using 1 text expert and 1 image expert (moe_1t1i) , 2)  using 4 text experts and 4 image experts (moe_4t4i) , and 3)  using 8 mixed-modal experts (moe_8x)  in addition to the dense baseline. All models were pre-trained and annealed to 1 trillion tokens.
When controlling for FLOPs per token, all sparse architectures demonstrate better scaling properties compared to the dense baseline. In particular, architectures with modality-specific experts consistently outperforms architectures with only mixed-modal parameters, with particular large benefits shown on the image modality.4.2Scaling of Performance with ComputeWe present the scaling performance of various models across multiple compute levels, with FLOPs matched to three dense model sizes: 90M, 435M, and 1.4B parameters. We report two key metrics: (1)  training loss and (2)  pre-training speed-up factor, which indicates that a sparse model can match the pre-training loss of the iso-FLOPs dense baseline using onlyof total FLOPs666Our definition ofis analogous to the speed-up factor proposed by, but is defined in terms of pre-training loss whereas the original definition uses validation perplexity..Figure 4:Effect of combining with mixture-of-depths (MoD) . We consider the sparse model variations compute-matched to three dense model sizes – 90M, 435M and 1.4B parameters. Combining MoD with MoE in general lead to better pre-training loss convergence across model sizes and configurations.Modality untying.Introducing modality-specific expert groups improves pre-training efficiency across different scales, with particularly large benefits for the image modality. As shown in Figure3, themoe_1t1iconfiguration, using 1 image expert and 1 text expert, significantly outperforms the dense baseline. The image loss ofmoe_1t1inearly matches that of the regular expert-choice MoE model using 8 mix-modal experts (moe_8x) , whereas the text loss remains substantially higher, indicating parameter untying has a disproportionate impact on the image modality. Scaling the number of experts within each modality group further improve the model performance. Themoe_4t4iconfiguration, using 4 text experts and 4 image experts, consistently outperformsmoe_8xacross different scales, with a substantial margin in image loss. However, this comparison reveals a slight regression in the text modality, suggesting that processing text modality across more experts may be beneficial.Mixture-of-Depths-and-Experts.We observe accelerated convergence of training loss when adopting both MoE and MoD, as well as their combinations. As shown in Figure4, adding MoD to themoe_1t1iarchitecture (mod_moe_1t1i)  significantly improves model performance across various model sizes. Moreover,mod_moe_1t1iperforms on par with or better thanmoe_4t4iacross model sizes and modalities, indicating that introducing sparsity over the depth dimension can also effectively improve training efficiency. On the other hand, we observe diminishing returns when stacking MoD and MoE.
Adding MoD to themoe_4t4iarchitecture results in only a mild performance boost compared tomod_moe_1t1iandmoe_4t4i. The improvement is also more visible for the text modality, while the image improvement is less significant. These findings suggest that future research may explore the combination of width and depth scaling to further enhance text modality performance. In contrast, improving image modality performance will require additional exploration of alternative approaches.4.3Scaling Number of Experts(a) The same number of experts are allocated to each modality(b) Uneven number of experts are allocated to each modalityFigure 5:Training curves of 435M compute-matched MoE architectures with increasing number experts with various expert allocations.We conducted further ablation experiments to investigate the impact of scaling the number of experts. We explored two scenarios: allocating an equal number of experts per modality (balanced expert allocation)  and allocating different numbers of experts per modality. We included the dense model and expert-choice MoE with 8 mixed-modal experts (moe_8x)  as baselines.Balanced expert allocation.Figure5(a) demonstrates that training loss consistently improves as the number of experts increases. Text and image losses exhibit distinct scaling patterns. While text loss improves steadily with each doubling of expert count, image loss shows diminishing returns as the number of experts increases from 2 to 4. This suggests that the intrinsic characteristics of each modality lead to different behaviors in sparse modeling. This finding aligns with our previous observation in §4.2, where untying modalities over 8 experts significantly enhanced image performance but not text performance, indicating that the text modality may benefit more from a larger number of experts.Imbalanced expert allocation.Given the diminishing returns observed when allocating more than 2 experts to the image modality, we experimented with configurations allocating fewer experts to the image modality and more to the text modality. Figure5(b) compares three configurations with the same total number of experts (8) , varying the allocation between text and image modalities (moe_7t1i,moe_6t2iandmoe_4t4i) . We observed that allocating more experts to a modality generally enhances its performance. However, assigning more than 4 experts to the text modality also yields diminishing returns. The overall training losses of the three configurations converge to similar levels. Considering the better load balancing of balanced expert assignment with our pre-training data mixture, we selectedmoe_4t4ias our optimal MoE configuration, despitemoe_6t2ihaving a slightly lower average loss. We leave the design of sparse architectures that can effectively leverage the intrinsic characteristics of different modalities to future work.4.4UpcyclingWe investigate further the influence of upcycling as described in §2.5. Specifically, we examine the 2.3B MoD model, comparing the training dynamics ofmod_moe_4t4iwhen trained from scratch versus when initialized from anmod_moe_1t1icheckpoint. To ensure a fair comparison, we adjust the data loader and training steps to account for the number of training steps already completed bymod_moe_1t1i, thereby maintaining equivalent training FLOPs. We ablate initializing the model from 10k and 20k steps.Figure6presents a comparison of the training curves for the various model variations. The training curves for the upcycled models have been adjusted to account for the computational cost of the first stage. We experimented with two seed checkpoints:mod_moe_1t1itrained for 10k and 20k steps, respectively. Our results show that upcycling further enhances model training, yielding a 1.2FLOPs gain with 10k steps in the first stage and a 1.16FLOPs gain with 20k steps. We observe the performance gap between upcycled models and the from-scratch model widens throughout training.Optimal Upcycling Period.As shown in Figure6, both initializations outperform the base MoE model, with 10k steps providing a 1.2speedup and 20k steps providing a 1.16speedup compared to the model trained from scratch. This suggests that a sweet spot likely exists for upcycling, where undertraining the seed model offers minimal benefits for router convergence, while overtraining it impedes future specialization. Based on these findings, we recommend upcycling from 10k steps. However, we hypothesize that the optimal number of upcycle steps may change when training beyond 1T tokens, and we leave this exploration for future research on in-depth upcycling of MoMa.Figure 6:Upcycling experiments. We compare the training curves of 2.3Bmod_moe_4t4istarting from three different initialization points: (1)  scratch, (2)  2.3Bmod_moe_1t1itrained for 10k steps, and (3)  2.3Bmod_moe_1t1itrained for 20k steps. All curves are FLOP adjusted to be equivalent, by shifting the upcycled models by the number of steps the dense model has been trained on.4.5Throughput AnalysisSparse models often cannot immediately lead to performance gains due to their added dynamism and associated data balancing issues (§3) . To quantify the impact of our proposals on training efficiency, we conducted a controlled experiment comparing the training throughput of various architectures, including a combination of MoMa, mixture-of-experts (MoE) , and mixture-of-depths (MoD) .
We compare architectures FLOPs-matched to the 435M dense baseline. We performed our experiments on 256 A100 GPUs with a sequence length of 4096 and a batch size of 6 per GPU and summarized the model throughput in Table2.ConfigurationWPS% diffDense31,970–MoMa 8X28,990-9%MoMa 1t1i30,196-6%MoMa 4t4i26,424-17%MoD MoMa 1t1i25,367-21%MoD MoMa 4t4i22,114-31%Table 2:Training throughput (wps/GPU)  of various architectures compared with the same flops as a 435M dense model on 256 A100 GPUs. Modality-conditioned sparsity offers a good quality-throughput tradeoff compared to the dense model, and exhibited reasonable scalability with more experts. On the other hand, while MoD variants achieves the best absolute loss, they generally incurs larger overheads due to additional dynamism and imbalance.Comparing expert-choice MoE (moe_8x)  and dense models, introducing sparsity does incur overheads (-9%) . This loss in throughput likely comes from the need to compute routing decisions and synchronize gradients for all experts despite being FLOPs-equivalent. On the other hand, as discussed in §3, running experts sequentially by modality (moe_1t1i)  does not suffer from large execution overheads, and we can attribute most of the loss (-6%)  to computing token indices for each modality, which can be amortized by precomputing the indices and share with each transformer layer.When combining modality conditioned feed-forward with learned routing (moe_4t4i) , we observe a smooth degradation in throughput as the number of experts increases, incurring additional 11% overheads with 8 experts, which is on par with the throughput loss of 9% when transitioning from the dense model to MoE with 8 experts.While combining MoD and MoE achieves the best training loss in our experiments, introducing MoD triggers an estimated throughput loss of 15% . This is because MoD architectures require an additional router on the depth dimension, introducing complexity and potential bottlenecks. Furthermore, when combined with MoMa, MoD may exacerbate system imbalance due to the varying active tokens for each modality at each layer, deviating from the predefined token mix ratio in the dataset. This deviation undermines the assumption of efficient MoMa execution, particularly with larger world sizes.
To mitigate this issue, we can force the MoD routers to admit tokens based on the predefined token mix ratio. However, the impact on model quality requires further investigation, which we leave for future work.4.6Inference-time PerformanceWe evaluate our models (1.4B dense, 1.4Bmoe_1t1i, 1.4Bmoe_8x, 1.4Bmoe_4t4i, 2.3Bmod_moe_4t4i)  on held-out language modeling data and downstream tasks. We report perplexity on subsets of OBELICSand Shutterstock777https://www.shuitterstock.com/, which are held-out from our pre-training dataset. In addition, we report 0-shot performance on a set of commonsense reasoning of tasks commonly used for benchmarking pre-trained language models. We also selected several vision-language task datasets and report the perplexity of the ground truth output in these datasets for cross model comparison, where we format the examples in a zero-shot manner.888We do not report the generation evaluation metrics for these tasks, as our experiment checkpoints are under-trained and have not achieved sufficient image understanding and generation capabilities to produce meaningful results on downstream tasks.•Commonsense Reasoning:We evaluated on a set of text-only tasks commonly used for benchmarking the commonsense reasoning capabilities of language models:PIQA,SIQA,HellaSwag,WinoGrande,ARC-Easy,ARC-Challenge,OpenBookQA, andBoolQ. We score the prompt with each candidate’s answer and compute accuracy using the candidate with the highest score.•Image Captioning:We take the Karpathy test split of MS-COCOand the Karpathy test split of Flickr30k, and report text-to-image and image-to-text conditional perplexity using these two datasets.•Visual Question Answering:We report the perplexity of ground truth answers on the test-dev split of VQA-v2.Image to text (PPL) Text to image (PPL) Interleaved (PPL) ModelCOCOFlickrVQAV2COCOFlickrObelicsOverallObelicsTextObelicsImageSSTKOverallSSTKTextSSTKImage1.4B Dense21.728.620.0458.3559.444.614.364.7238.548.7245.81.4B MoMa 1t1i21.929.319.8416.7508.641.513.859.3208.17.6221.51.4B MoE 8x18.924.218.9426.5518.540.412.559.2218.833.7226.61.4B MoMa 4t4i20.926.619.8392.5479.639.013.055.8194.710.5205.7\hdashline2.3B MoD MoMa 4t4i22.928.020.5428.1518.041.413.260.0215.115.0226.1Table 3:Interleaved data modeling performance.ModelBoolQPiQASiQAWinogrande1.1OBQAHellaswagArc-EArc-CAvg1.4B Dense61.872.144.354.242.051.059.834.352.41.4B MoMa 1t1i61.571.344.454.746.052.861.134.953.31.4B MoE 8x62.273.544.458.342.258.262.636.954.81.4B MoMa 4t4i62.372.544.655.947.055.260.735.754.2Table 4:Text-to-text commonsense reasoning task performance.For all sparse models, we perform the second-stage training of auxiliary routers and use the auxiliary router for causal inference. The details of auxiliary router training can be found in AppendixA.2.Interleaved data modeling.The relative performance of the dense model and various MoE configurations is consistent with the trends observed in pre-training loss (§4.2) . As shown in Table3, by using 1 more image expert, the 1.4B MoMa 1t1i model significantly outperforms the dense baseline on most metrics, except for image-to-text conditional perplexity on COCO and Flickr. Similar to the trend shown in §4.2, adding the image expert leads to substantial performance gains on the image modality. Further scaling up the number of experts also improves the performance, with 1.4B MoE 8x achieves the best image-to-text performance. Additionally, the model excels on text-to-text tasks (Table4) . 1.4B MoMa 4t4i performs the best on all conditional image perplexity metrics, with the text perplexity closely matches that of 1.4B MoE 8x on most of the benchmarks. Overall, the 1.4B MoMa 4t4i model achieves the best interleaved data modeling results.MoD performance regression.According to Table3, we observe significant performance regression of the 2.3B MoD MoMa 4t4i model across all metrics, leaving it underperforming the 1.4B MoMa 4t4i model despite having better pre-training loss (§4.2) . Our analysis identifies the auxiliary router’s performance as the primary cause. We used a smaller model, 635M MoD MoMa 4t4i, in the analysis. Initially, we validated that using the training router and selecting toptokens within a batch for inference enables the 635M MoD MoMa 4t4i model to outperform the 435M MoMa 4t4i model. However, when introducing random noise to the training router selection with a noise ratio of, the 635M MoD MoMa 4t4i model starts underperforming the 435M MoMa 4t4i model whenexceeds 0.5%. Achieving an error rate below this threshold is impractical for auxiliary router training. Therefore, further research is necessary to develop a robust approach for practical MoD applications.On the other hand, MoE models do not exhibit similar regression during inference, suggesting they are less sensitive to router errors. This is likely due to the MoE model’s multi-expert architecture, which allows mis-routed tokens to be processed by other experts, preserving their information. In contrast, the MoD model’s single selection switch per layer means routing mistakes cannot be remedied, leading to greater performance degradation.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
4Experiments
2407.21770
5Related WorkEarly-fusion vision-language models.Early fusion techniques have gained traction in multi-modal learning due to their ability to capture cross-modal interactions from the onset of processing. Several notable works have explored this paradigm. PerceiverIOintroduced a fully attentional read-process-write architecture that operates over a modality-agnostic latent space for processing diverse inputs, including text and images. NÜWApresented a 3D-attention transformer capable of both understanding and generating text, image, and video in various combinations. CM3adopted a causally masked transformer to learn from mixed-modal documents on the Internet, showcasing the scalability of this paradigm to large-scale pretraining. We adopt Chameleon as the base transformer architecture and demonstrate that modality-aware sparsity can effectively improve its scaling performance further.Multi-modal representation learning.BEIT-3introduced a general-purpose multimodal foundation model that achieves state-of-the-art performance on both vision and vision-language tasks. It employs a Multiway Transformer architecture with modality-specific expert modules and uses a unified masked data modeling objective for pretraining on images, texts, and image-text pairs.Sparse multi-modal language models.Recent advancements in sparse modeling techniques have shown promising results in scaling language models efficiently. VL-MoEproposed a vision-language model that uses a sparse mixture-of-experts (MoE)  to scale efficiently. It applies sparse MoE to the vision and language feed-forward networks in a unified framework, using masked data modeling objectives for pretraining. VL-MoE achieves strong performance on vision-language tasks while using fewer parameters per token than dense models, providing additional evidence for MoE training for multi-modal settings.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
5Related Work
2407.21770
6LimitationsOur current MoMa implementation relies on matching the token mix ratio in the dataset with the expert mix ratio in the model to maintain load balance across GPUs. Even so, minor imbalance may still occur because there is no hard limit for a batch to deviate from that ratio at per-GPI per-iteration level. We leave further improvements in this area as future work.Expert-choice routing alleviates the expert load balancing issue during training but presents additional challenges for auto-regressive Language Models (LMs)  during inference. Although auxiliary routers comprise only a small fraction of the network’s parameters, their role is crucial. In our study, we trained the auxiliary router after completing the whole network training and limited this process to a few thousand steps, while previous work have demonstrated the possibility to jointly train such modules with the full network. Future research should explore the architecture and training techniques for auxiliary routers to prevent them from becoming a performance bottleneck and ensure generalizability across diverse data distributions. Especially, further investigation is necessary for training mixture-of-depths architectures, including both the auxiliary routers and the original model, to ensure effective performance in causal inference scenarios.In our work, we experimented only the vanilla formulation of MoD and its staged integration with MoE. We leave the investigation of other MoD variations, including modality-aware MoD to future work. In addition, batched sequence generation with mixture-of-depths (MoD)  is non-trivial as unlike standard sequence generation there are dynamics shapes and dynamic updates to the KV cache for each layer as certain sequences and layers may skip different tokens. There remains further room to optimize the inference implementations for MoD models.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
6Limitations
2407.21770
7ConclusionIn this work, we have introduced a series of modality-aware sparse architectures for early fusion, mixed-modality foundation models. Our approach leverages domain specificity while preserving cross-modality knowledge sharing and feature interaction. We have developed highly efficient and effective model architectures by incorporating sparsity in both the width dimension (through modality-aware mixture-of-experts)  and the depth dimension (via mixture-of-depths) .Our best architecture,Chameleon-MoMa, demonstrate significant improvements over state-of-the-art baselines. In complexity-controlled experiments, we reduce total FLOPs requirements by up to 3.7. Importantly, our experimental findings reveal that our modality-aware sparse architectures maintain an empirical scaling law. This characteristic suggests that our approach provides immediate performance benefits and a scalable framework for future developments in mixed-modal foundation model training.The promising potential ofChameleon-MoMaopens up several promising directions for future research. These include exploring more sophisticated routing mechanisms, investigating the impact of different sparsity patterns across modalities, and extending our approach to a broader range of modalities and tasks.
body
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
7Conclusion
2407.21770
AcknowledgmentsWe thank Baptiste Rozière, Alexei Baevski, Naman Goyal, and Horace He for helpful discussions on mixture-of-experts model training. We thank Ram Pasunuru and Andrew Cohen for helpful feedback on model evaluation. We thank the members of the Chameleon Team for insightful discussions and feedback.
conclusion
MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
Acknowledgments