NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

Merging RLVR-Trained Experts via Policy-Shift-Guided Spectral Alignment

Geeho Kim, MINSIK CHOI, Kyle Min, Young Geun Kim, Bohyung Han

Expand abstract

We propose Policy-Shift-Guided Spectral Alignment (PSA), a retraining-free method for merging RLVR post-trained language-model experts by guiding spectral subspace selection with token-level expert-base policy shifts. We first show that existing SFT-oriented merging techniques under-preserve the sparse, directional changes in next-token probabilities that distinguish RLVR experts from the pretrained base. To address this challenge, PSA scores calibration tokens by expert--base probability difference under the same prefix, and converts probability-shift-weighted input activations into column weights. These weights define a weighted SVD objective that prioritizes update directions active on high-shift tokens. After low-rank truncation, PSA applies cross-expert polar alignment to the truncated task-vector bases and restores each expert's per-layer Frobenius scale before aggregation. Experiments on Qwen2.5-7B and Qwen3-1.7B across three RLVR tasks show that PSA consistently outperforms strong merging baselines while preserving expert-level capabilities in a unified model.

Poster
Mask-Conditioned Gradient Masking for Fine-Tuning Mixture-of-Experts Diffusion Language Models

Yiru Tang, Kun Zhou, Xin Zhao, Jing Sha, Zhichao Sheng, Shijin Wang

Expand abstract

While autoregressive models dominate the LLM landscape, Discrete Diffusion Language Models have emerged as a compelling alternative due to their advantages in parallel decoding and bidirectional contextual modeling. To enhance scalability and efficiency, recent works have integrated the Mixture-of-Experts architecture into the DLM framework. However, MoE-DLMs remain underexplored. In this work, we present an empirical study revealing that MoE-DLMs exhibit not only task-level expert specialization, but also exhibit specialization across different mask-rate diffusion regimes. We further show that during fine-tuning, gradients induced by mismatched mask rates can interfere with the updates of regime-specialized experts, leading to suboptimal adaptation. Motivated by this, we propose MaGM, an adaptive mask-conditioned gradient masking method for MoE-DLM fine-tuning, which dynamically masks expert parameter updates based on the mask rate of training instance. Experiments demonstrate that MaGM consistently outperforms standard full-parameter fine-tuning for diffusion language models, validating the benefits of regime-aware expert adaptation. Our code will be publicly released.

Poster
Managing Self-Learning Experts under Per-Round Budget Constraints

Ilgam Latypov, Alexandra Suvorikova, Alexey Kroshnin, Alexander Gasnikov, Iurii Dorn

Expand abstract

This paper addresses the problem of sequential decision-making under learning budget constraints. Such settings naturally arise in applications like managing a portfolio of bandit or reinforcement learning (RL) algorithms. We propose a novel UCB-type algorithm, M-LCB, designed to manage a pool of $K$ self-learning experts in a stochastic environment while accounting for a limited per-round learning budget $M$. At each round, M-LCB selects one expert to make a decision and at most $M \le K$ experts to learn. For selection, M-LCB uses confidence bounds constructed from limited prior knowledge about the experts (i.e., mild assumptions) and their observed training losses. We derive anytime regret bounds for M-LCB that scale with the individual regrets of the experts. In particular, if each expert has regret $\tilde O(T^\alpha)$ by round $T$, then M-LCB guarantees an overall regret of $\tilde O\left(\sqrt{KT/M} + (K/M)^{1-\alpha}T^\alpha\right)$ relative to the best expert in hindsight. Finally, we demonstrate the applicability of M-LCB using self-learning experts instantiated as (i) parametric models and (ii) bandit algorithms.

Poster
MEME: Lightweight Hierarchical Mixture-of-Experts for Unified Affective Computing

Yinan Zhang, Haoyu Zhang, Tianshu Yu

Expand abstract

Multimodal sentiment analysis (MSA) and emotion recognition (ER) are closely related affective understanding tasks, yet they are usually studied separately due to differences in label space, task formulation, dataset distribution, and modality dependence. In this paper, we propose MEME ($\textbf{M}$ixture-of-$\textbf{E}$xperts for $\textbf{M}$ultimodal sentiment analysis and $\textbf{E}$motion recognition), a lightweight hierarchical MoE framework for unified multimodal affective computing across MSA, conversational emotion recognition, and dynamic facial expression recognition. MEME operates on frozen text, visual, and audio features, compresses variable-length modality sequences with learned-query attention pooling, and processes them with a shared hierarchical MoE backbone. Each block first applies hard modality experts for modality-specific refinement and then task-conditioned cross experts for task-aware multimodal interaction. A $\texttt{[TASK]}$ token provides an explicit task anchor for routing, while modality dropping regularizes unified training. We jointly train MEME on nine affective benchmarks and evaluate all datasets using a single composite-best checkpoint without dataset-wise adaptation. Extensive experiments show that MEME outperforms strong task-specific baselines and recent large-model-based unified methods on most benchmarks, while maintaining favorable efficiency, robustness and generalization.

Poster
MDK-MoE: Multi-view Decomposed Kalman Mixture of Experts Framework for Non-stationary Time Series Forecasting

Rui Hou, Yao Liu, Ruilin Jiang, Mengyao Lu, Jingbo Wang, Lanyi Zhang, Qiao Liu

Expand abstract

Non-stationary time series forecasting (NTSF) requires modeling evolving trends and seasonal patterns. Although linear decomposition operators (e.g., FFT) help isolate these patterns, they often introduce operator-dependent biases and idiosyncratic noise. To address this issue, we propose a novel Multi-view Decomposed Kalman Mixture of Experts (MDK-MoE) framework, which \textit{improves forecasting by strategically integrating multiple complementary, yet biased, linear decompositions}. Specifically, a Multi-view Decomposition (MVD) extracts frequency components from three complementary views. To mitigate view-specific noise, a Maximizing Mutual Information (MMI) enhances cross-view consistency, maximizes task-relevant information across diverse views, and suppresses view-dependent task-irrelevant information. To address the conflicts inherent in multi-source observations, the efficient Kalman MoE adaptively fuses distinct multi-view features. Instead of adopting traditional $O(d^3)$ covariance inversion, our learnable Kalman gain design enhances numerical stability and enables reliable state estimation under non-stationary scenarios. Extensive experiments on 13 real-world datasets demonstrate that MDK-MoE consistently achieves state-of-the-art performance.

Poster
Load Balancing Mixture of Experts with Similarity Preserving Routers

Nabil Omi, Siddhartha Sen, Ali Farhadi

Expand abstract

Sparse Mixture of Experts (MoE) models offer a scalable and efficient architecture for training large neural networks by activating only a subset of parameters (“experts”) for each input. A learned router computes a distribution over these experts, and assigns input tokens to a small subset. However, without auxiliary balancing mechanisms, routers often converge to using only a few experts, severely limiting model capacity and degrading performance. Most current load balancing mechanisms encourage a uniform routing probability across experts. Early in pretraining, this can result in inconsistent routing behavior, resulting in the model spending its capacity learning redundant knowledge. We address this by introducing a novel load balancing loss that helps preserve relational structure, encouraging consistent expert choices for similar inputs during training. Our experimental results show that applying our loss to the router results in 36% faster convergence and lower redundancy compared to a popular load balancing loss.

Poster
Knowledge Localization in Mixture-of-Experts LLMs Using Cross-Lingual Inconsistency

Lucas Bandarkar, Alan Ansell, Trevor Cohn

Expand abstract

Modern LLMs continue to exhibit significant variance in behavior across languages, such as being able to recall factual information in some languages but not others. While typically studied as a problem to be mitigated, in this work, we propose leveraging this cross-lingual inconsistency as a tool for interpretability in mixture-of-experts (MoE) LLMs. Our knowledge localization framework contrasts routing for sets of languages where the model correctly recalls information from languages where it fails. This allows us to isolate model components that play a functional role in answering about a piece of knowledge. Our method proceeds in two stages: (1) querying the model with difficult factual questions across a diverse set of languages to generate "success" and "failure" activation buckets and then (2) applying a statistical contrastive analysis to the MoE router logits to identify experts important for knowledge. To validate the necessity of this small number of experts for answering a knowledge question, we deactivate them and re-ask the question. We find that despite only deactivating about 20 out of 6000 experts, the model no longer answers correctly in over 40% of cases. Generally, this method provides a realistic and scalable localization approach to increasingly complex LLMs, but also suggests redundant and highly dispersed knowledge parameterization in MoEs.

Poster
Improving Continual Video Instance Segmentation via Spatial-Temporal Balanced Mixture-of-Experts Adapters

Hao Fu, Sihao Liu, Hanbin Zhao, Jiahua Dong, zhang chao, Hui Qian

Expand abstract

Effective Continual Video Instance Segmentation (CVIS) requires balancing catastrophic forgetting with the plasticity to learn new tasks. Existing methods typically adapt pre-trained VIS models using trainable visual prompts, but suffer from two key limitations: they overlook the distinct degradation of frame-level spatial knowledge (e.g., misclassification) and video-level temporal knowledge (e.g., trajectory disruption), and rely on rigid prompt retrieval that struggles under task ambiguity and distribution shifts. To address these issues, we propose STEAM, a CVIS framework with Spatial-Temporal Balanced Mixture-of-Experts Adapters. Instead of explicit retrieval, we introduce a Distribution-Aware spatial-temporal Scaling and sHifting (DASH) mechanism to align the pre-trained feature space with sequential tasks and reduce the domain gap. Furthermore, our method adaptively learns Task-Adaptive Spatial-Temporal mixture-of-Experts (TASTE) at both the frame and video levels, while updating them with Subspace-Constrained Orthogonal Adaptation (SCOA) to mitigate interference. Extensive experiments across diverse CVIS settings show that STEAM preserves temporal-spatial coherence and achieves state-of-the-art performance.

Poster
HypMoE-ReID: Hyperspherical Mixture-of-Experts for Large Scale Person Re-Identification

Kunlun Xu, Haotong Cheng, Yufei Guo, Jiahuan Zhou

Expand abstract

With the rapid growth of surveillance data, Person Re-Identification (ReID) is evolving toward large scale scenarios, where datasets exhibit substantial increases in both size and diversity. Under this trend, the widely adopted dense backbone architectures, such as Transformers, suffer from an inherent limitation: Shared parameters couple all input tokens, forcing the model to learn generic representations that are suboptimal for heterogeneous data sources. Although Mixture-of-Experts (MoE) architectures offer a promising solution for handling diverse data, they encounter a fundamental challenge in ReID. Specifically, person images exhibit high inter-instance similarity at the coarse level, while discriminative cues lie in subtle fine-grained details that are easily overlooked. As a result, existing MoE methods often suffer from severe routing collapse, leading to insufficient expert specialization and limited diverse knowledge acquisition capacity. To address these problems, we propose a novel framework, termed Hyperspherical Mixture-of-Experts for Large Scale Person Re-Identification (HypMoE-ReID). Specifically, motivated by theoretical analysis in hyperspherical space, HypMoE-ReID incorporates two key components: (1) a Token Purification mechanism that suppresses noisy tokens detrimental to expert specialization, and (2) a Routing Diversity Learning strategy that explicitly separates experts on a unit hypersphere to enforce specialization and prevent collapse. Furthermore, to facilitate the investigation in large scale ReID, we construct a large-scale ReID benchmark comprising 809,701 images by unifying 13 public datasets with significant distribution gaps. Furthermore, extensive experiments demonstrate that HypMoE-ReID outperforms state-of-the-art MoE-based methods by at least 11.3\%/10.3\% in Average mAP/R@1, and surpasses strong dense backbone models by 2.2\%/2.0\% in Average mAP/R@1. Our code will be released.

Poster
How to Scale Mixture-of-Experts: From muP to the Maximally Scale-Stable Parameterization

Leena Chennuru Vankadara, Moritz Haas, Luke Hayward, Sebastian Bordt, Alessandro Breccia

Expand abstract

Recent frontier large language models predominantly rely on Mixture-of-Experts (MoE) architectures. Despite empirical progress, there is still no principled understanding of how hyperparameters should scale with network width $N$, expert width $N_e$, number of experts $M$, sparsity $K$, and depth $L$ to ensure both stability and optimal performance at scale. We take a principled step toward resolving this gap by analyzing three different scaling regimes: (I) co-scaling $N\asymp N_e$, (II) co-scaling $N\asymp M\asymp K$, and (III) full proportional scaling of $N, N_e, M$, and $K$. For each regime, we develop a novel Dynamical Mean Field Theory (DMFT) description of the limiting training dynamics of MoEs that provides a formal foundation for our analysis. Within this framework, we derive the unique parameterization for SGD and Adam satisfying all maximal-update ($\mu$) desiderata. We then show that the resulting $\mu$P prescription induces neither monotonic improvement with scale nor robust learning-rate transfer. We trace these pathologies to scale-dependent observables in the aggregation dynamics, which motivates a refined set of desiderata that we term *maximal scale stability*. Guided by this principle, we derive a *Maximally Scale-Stable Parameterization* (MSSP) for both SGD and Adam in all three scaling regimes, and characterize the corresponding limiting dynamics - qualitatively distinct from the $\mu$P limit - through a separate DMFT analysis. Experiments verify that MSSP robustly recovers learning rate transfer and monotonic improvement with scale across regimes. Combined with existing depth-scaling theory, these results provide a complete scaling prescription for MoE architectures as a function of width, depth, expert width, and number of experts.

Poster
Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization

Weilin Wan, Jingtao Han, Debing Zhang, WEIZHONG ZHANG, Cheng Jin

Expand abstract

Scaling laws govern macroscopic resource allocation for LLMs, yet precise architectural configurations for Mixture-of-Experts (MoE) models remain guided by heuristics. Existing MoE scaling studies either incorporate MoE-specific variables into scaling formulas, causing fitted coefficients to grow rapidly without proportional experimental support, or fix all non-MoE factors, implicitly assuming global architecture does not influence local MoE scaling. We propose a reusable framework for holistic MoE architectural optimization. We first reveal that relying solely on FLOPs per token ($M$) biases MoE evaluation, as heterogeneous Attention/FFN computational densities enable parameter inflation without effective compute gains. We therefore establish a joint constraint triad of $M$, active parameters ($N_a$), and total parameters ($N$) for rigorous MoE characterization. To tame the resulting $\mathcal{O}(n^{16})$ search space, we employ mathematical decoupling: structural constraints and a rank-preserving property of the hidden dimension factorize the optimization into an $\mathcal{O}(n^3) + \mathcal{O}(n^2)$ two-phase search. Through extensive validation across 670+ MoE models spanning $10^{18}$ to $3 \times 10^{20}$ FLOPs, we derive globally applicable scaling laws mapping any compute budget to its optimal architecture. A key finding is that the near-optimal configuration band widens with scale, providing a quantitative basis for trade-offs between scaling law recommendations and engineering constraints. Our work delivers actionable blueprints for optimal MoE design under arbitrary compute budgets.

Poster
Geometry is an Operator: Lie-Algebraic Space Routing for View-Robust 3D MLLMs

Jingjun Yi, Chen Hu, Qi Bi, Hao ZHENG, Huimin Huang, Yongkang Li, Wei Ji, Yawen Huang, Xian Wu, Yefeng Zheng

Expand abstract

Multimodal Large Language Models (MLLMs) have shown increasing potential for 3D scene understanding and spatial reasoning from videos. However, even with strong visual geometry priors, the hidden semantics of existing 3D-enhanced MLLMs can drift across views, leading to degraded grounding, captioning, and spatial reasoning performance. The underlying reason can be their geometry injection as additional tokens or additive features, which enriches visual content but fails to define how view-dependent semantic states should transform. To address this issue, we propose Lie Routing Vision-Language Model (LieVLM), a geometry-conditioned Lie routing framework that treats geometry as an operator over hidden semantics. LieVLM first elicits local Lie states from 3D geometry tokens by predicting latent rotation states and routing coefficients. It then converts these rotation states into triplet-wise Lie group actions to route projected 2D visual semantics. Finally, a geometry-aware routed fusion module combines the original semantic path, the direct 3D context path, and the Lie-routed transformation path for language generation. This design preserves pretrained visual-language semantics while allowing geometry to actively correct view-induced semantic drift. Extensive experiments on 3D scene understanding and spatial reasoning benchmarks demonstrate that LieVLM improves view robustness, especially in severe-pose-change regimes, and consistently outperforms additive geometry fusion methods. Our results suggest a shift in 3D MLLM design, i.e., geometry should not merely be injected as context, but should operate on semantic representations.

Poster
Gaussian Mixture Models in Hilbert Spaces via Kernel Methods

Daniel López Montero, Antonio Álvarez-López, Marcos Matabuena

Expand abstract

Modern datasets across many disciplines increasingly consist of time-evolving, potentially infinite-dimensional random objects, such as dynamic functional data, which are naturally modeled in Hilbert spaces. In these settings, characterizing probability measures, for example, through densities, can be ill-defined or technically challenging. Motivated by clustering applications, we propose a Gaussian mixture framework for Hilbert-space-valued data based on kernel mean embeddings and develop efficient optimization algorithms for estimation. We establish theoretical guarantees showing that the proposed algorithm is well defined and that the model yields a dense class of approximations in infinite-dimensional spaces. We evaluate the framework through extensive experiments on diverse structures and data geometries, including $L^2$-functional data and random graphs in Laplacian spaces arising in modern medical applications.

Poster
From Facts to Personas: Interpretable Role Unlearning in LLMs via Mixture-of-Experts

Ruihong Zeng, Puning Yang, Jinghui Zhang, Shen Gao, Xiangliang Zhang, Preslav Nakov, Xiuying Chen

Expand abstract

LLMs are widely adopted in multi-role dialogue and persona simulation. However, their strong ability to imitate role-specific linguistic and behavioral patterns can introduce privacy risks, reinforce social biases, and lead to unsafe or uncontrollable outputs. To address these issues, we introduce the Role-playing Unlearning task, which aims to selectively forget the style and knowledge associated with target roles while preserving the language ability and behaviors of the remaining roles. Unlike prior unlearning tasks that focus solely on removing factual information, our task additionally requires forgetting persona-specific characteristics, making the control over model generation more fine-grained. We propose ROLEMOE, a Mixture-of-Experts framework that enhances the separation of role-specific patterns by mapping different roles to more independent expert subspaces, providing structural interpretability. To reinforce this separation, ROLEMOE incorporates two complementary losses: a specialization loss that drives different roles to adopt different expert distributions, and a disentanglement loss that encourages different experts to develop distinct capability subspaces. This decomposed structure enables role unlearning to focus selectively on target-role experts. We finally propose RoleBench-Unlearn for evaluation, and experiments across multiple LLM architectures show that ROLEMOE achieves significantly more effective and controllable role forgetting, while better preserving the linguistic quality, knowledge, and persona consistency of non-target roles. All code and datasets are publicly available at: https://anonymous.4open.science/r/RoleMoE-6328.

Poster
From Experts to Sub-experts: Fine-grained Parameter-Efficient Fine-Tuning for MoE LLMs

Zhentao Tan, Chang Liu, Yao Liu, Yue Wu, Jieping Ye

Expand abstract

As large language models (LLMs) scale rapidly, dense full-parameter adaptation becomes increasingly expensive, motivating sparse and modular architectures such as Mixture-of-Experts (MoE) models. This shift raises a key question for parameter-efficient fine-tuning (PEFT): at what granularity should parameters be selected and updated? Existing PEFT methods such as LoRA operate on predefined weight matrices, while expert-level sparse tuning methods update entire selected experts. However, we observe that activated experts are internally sparse, with only a small fraction of intermediate channels strongly responding to downstream tasks, indicating that expert-level adaptation is still too coarse. We propose NSFT (Neural Sub-expert Fine-Tuning), a fine-grained PEFT framework that refines MoE adaptation from experts to sub-experts. NSFT decomposes each expert along the intermediate dimension into structured channel groups and selects task-relevant sub-experts by combining routing importance with intra-expert activation saliency. To optimize sparse partial updates, NSFT further introduces learning-rate scaling and dynamic gradient scaling to compensate for the reduced effective update magnitude. Experiments on OLMoE and Ling-mini-2.0 across challenging domain-specific tasks and general benchmarks show that NSFT consistently outperforms representative PEFT and expert-level sparse tuning baselines, while using substantially fewer trainable parameters and preserving competitive general capability. These results suggest that sub-expert-level adaptation is a more precise and efficient PEFT paradigm for MoE LLMs.

Poster
FineRMoE: Dimension Expansion for Finer-Grained Expert with an Upcycling Approach

Ning Liao, Xiaoxing Wang, Xiaohan Qin, Junchi Yan

Expand abstract

As revealed by the scaling law of Mixture-of-Experts models, model performance ceases improving once the granularity of the intermediate dimension exceeds the optimal threshold, limiting further gains from single-dimension fine-grained design. To address this bottleneck, we propose FineRMoE (FineR-Grained MoE), an architecture that extends fine-grained expert design to both intermediate and output dimensions, aiming to enhance expert specialization beyond the single-dimension limit. We further introduce a bi-level sparse forward computation paradigm and a specialized routing mechanism to govern the activation. In addition, to obviate the prohibitive cost of pre-training from scratch, we devise a generalized upcycling method to build FineRMoE in a cost-effective manner. Extensive experiments performed under the upcycling paradigm demonstrate the superior performance achieved by FineRMoE across ten standard benchmarks. Compared with the strongest upcycling baseline method, FineRMoE achieves significant parameter and inference efficiency gains.

Poster
FLUX: Geometry-Aware Longitudinal Flow Matching with Mixture of Experts

Josue Ortega Caro, Yongxu Zhang, Hannah M Batchelor, Sizhuang He, Jessica Cardin, Shreya Saxena

Expand abstract

Many biological systems evolve through continuous local dynamics while switching between latent regimes defined by learning, stimulus context, internal state, or developmental stage. These processes are often observed only as unpaired longitudinal snapshots: the same cells, neurons, or animals are not tracked as matched trajectories, even though population states are sampled across successive stages. This creates two coupled challenges. First, trajectories must respect curved low-dimensional manifolds embedded in high-dimensional biological measurements. Second, the model must identify when the transport mechanism itself changes. We introduce FLUX (FLow matching for Unpaired longitudinal data with miXture-of-experts), a geometry-aware longitudinal flow-matching framework for joint transport modeling and unsupervised regime discovery. FLUX learns a data-dependent metric from pooled labeled and unlabeled observations, uses that metric to construct geometry-aware conditional paths between adjacent marginals, and decomposes the resulting velocity field into sparse expert vector fields selected by a Straight-Through Gumbel-Softmax router. Across manifold controls, a regime-switching Lorenz system, widefield cortical calcium imaging during associative learning, and embryoid body single-cell differentiation, FLUX reconstructs longitudinal transport while recovering interpretable regime structure. In neural data, the router separates early/intermediate training from late training, coinciding with the behavioral divergence of CS+ and CS- lick indices. In embryoid body differentiation, FLUX separates expression-evolution regimes associated with pluripotent and more differentiated cell populations. Ablations show that mixture-of-experts routing alone is insufficient: FLUX without geometric learning can fit local transport but fails or weakens regime discovery when regimes are encoded in local dynamics. These results suggest that geometry-aware velocity decomposition provides a general strategy for discovering latent biological state transitions from unpaired longitudinal snapshots.

Poster
Evidence Guided Dual Expert Memory with Joint Routing for VLLMs Online Correction

ZHANG LINTONG, Yearang Lee, Seong-Whan Lee

Expand abstract

Vision-large language models (VLLMs) editing aims to efficiently update model responses while retaining generalization and unrelated knowledge. Recent online editors store previous edits as reusable experts, forming a mixture-of-experts editing memory for scalable VLLM correction. By retrieving relevant experts at inference time, this paradigm avoids repeatedly overwriting model parameters and enables continual reuse of accumulated editing experience. However, existing sample-level residual experts usually encode each edit instance as a whole, entangling target-relevant visual evidence with irrelevant context. This coarse expert construction weakens expert reusability, leads to inaccurate routing, and increases unintended side effects during continual editing. To address this limitation, we reformulate online VLLM editing as evidence guided expert construction and routing. A fine mask identifies edit-relevant visual evidence, from which support-suppress experts are constructed to strengthen target consistent cues and weaken edit conflicting cues. Each expert is stored with routing signals and reliability scores, and later retrieved through joint sparse routing over visual evidence, textual context, and expert quality. Experiments demonstrate consistent improvements in editing reliability, generalization, and interpretability over competitive online editing baselines.

Poster
Efficient Pre-Training with Token Superposition

Bowen Peng, Théo Gigant, Jeffrey Quesnelle

Expand abstract

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.

Poster
EMO: Pretraining Mixture of Experts for Emergent Modularity

Ryan Wang, Akshita Bhagia, Sewon Min

Expand abstract

Large language models are typically deployed as monolithic systems, requiring the full model even when applications need only a narrow subset of capabilities, e.g., code, math, or domain-specific knowledge. Mixture-of-Experts (MoEs) seemingly offer a potential alternative by activating only a subset of experts per input, but in practice, restricting inference to a subset of experts for a given domain leads to severe performance degradation. This limits their practicality in memory-constrained settings, especially as models grow larger and sparser. We introduce EMO, an MoE designed for modularity—the independent use and composition of expert subsets—without requiring human-defined priors. Our key idea is to encourage tokens from similar domains to rely on similar experts. Since tokens within a document often share a domain, EMO restricts them to select experts from a shared pool, while allowing different documents to use different pools. This simple constraint enables coherent expert groupings to emerge during pretraining using document boundaries alone. We pretrain a 1B-active, 14B-total EMO on 1T tokens. As a full model, it matches standard MoE performance. Crucially, it enables selective expert use: retaining only 25\% (12.5\%) of experts incurs just a 1\% (3\%) absolute drop, whereas standard MoEs break under the same setting. We further find that expert subsets in EMO specialize at semantic levels (e.g., domains such as math or code), in contrast to the low-level syntactic specialization observed in standard MoEs. Altogether, our results demonstrate a path toward modular, memory-efficient deployment of large, sparse models and open new opportunities for composable architectures.

Poster