PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Embeddings Outperforms Scaling Experts in Language Models

Hong Liu, Jiaqi Zhang, Chao Wang, Xing Hu, Linkun Lyu, Jiaqi Sun, Xurui Yang, Bo Wang, Fengcun Li, Yulei Qian, Lingtong Si, Yerui Sun, Rumei Li, Peng Pei, Yuchen Xie, Xunliang Cai

arXiv:2601.21204v2cs.CLcs.AIcs.LG

TL;DR

MoE scaling encounters diminishing returns and system bottlenecks, motivating the study of embedding scaling as an alternative sparse dimension. The paper characterizes when embedding scaling is more effective, develops system optimizations, and reports LongCat-Flash-Lite’s competitiveness against MoE baselines and similar-scale models.

  • Problem

    The paper addresses limited understanding of how embedding scaling compares with expert scaling and how architectural choices and inference I/O affect its effectiveness.

  • Method

    The study systematically compares embedding and expert scaling, analyzes architectural factors, and applies N-gram Embedding with cache, synchronized-kernel, and decoding optimizations.

  • Results

    LongCat-Flash-Lite, with 68.5B total and 2.9B∼4.5B activated parameters, surpasses a parameter-equivalent MoE baseline and is competitive in agentic and coding tasks.

  • Takeaways & Limitations

    Embedding scaling yields a superior Pareto frontier to expert scaling in specific regimes, supporting it as an orthogonal and efficient sparsity dimension.

  • Takeaways & Limitations

    N-gram Embedding adds I/O, computation, and communication overhead, while dynamic scheduling complicates lookup optimization.

Abstract

from arXiv · show

While Mixture-of-Experts (MoE) architectures have become the standard for sparsity scaling in large language models, they increasingly face diminishing returns and system-level bottlenecks. In this work, we explore embedding scaling as a potent, orthogonal dimension for scaling sparsity. Through a comprehensive analysis and experiments, we identify specific regimes where embedding scaling achieves a superior Pareto frontier compared to expert scaling. We systematically characterize the critical architectural factors governing this efficacy -- ranging from parameter budgeting to the interplay with model width and depth. Moreover, by integrating tailored system optimizations and speculative decoding, we effectively convert this sparsity into tangible inference speedups. Guided by these insights, we introduce LongCat-Flash-Lite, a 68.5B parameter model with ~3B activated trained from scratch. Despite allocating over 30B parameters to embeddings, LongCat-Flash-Lite not only surpasses parameter-equivalent MoE baselines but also exhibits exceptional competitiveness against existing models of comparable scale, particularly in agentic and coding domains.

1 Introduction

The paper studies embedding scaling as an alternative sparse dimension to expert scaling, addressing unclear allocation, architectural, and inference-efficiency trade-offs. It reports regimes where embedding scaling improves the Pareto frontier and presents LongCat-Flash-Lite as validation.

  • MoE scaling faces diminishing marginal gains and system bottlenecks as model size and sparsity increase.
  • Embedding layers provide O(1) lookup complexity, enabling large parameter expansion without routing overheads or proportional computation growth.
  • The study examines how parameter budgets, vocabulary size, initialization, hash collisions, width, depth, and integration choices govern embedding-scaling efficacy.
  • N-gram Embedding achieves a superior Pareto frontier to increasing expert numbers in specific regimes and offers robust scalability among evaluated embedding methods.
  • N-gram Cache and synchronized kernels target embedding overhead and convert reduced active parameters into lower latency and higher throughput.
  • LongCat-Flash-Lite uses 68.5B total parameters and 2.9B∼4.5B activated parameters, surpassing a parameter-equivalent MoE baseline and competing particularly in agentic and coding tasks.

2 N-gram Embedding Layer

The paper scales embedding capacity with vocabulary-free N-gram Embedding, which augments token representations using hashed n-gram tables. A decomposed multi-table design and projections improve expressiveness while preserving parameter-count invariance across n-gram order and table count.

  • N-gram Embedding augments the embedding module with a vocabulary-free n-gram embedding table for each sequence token.
  • The method combines the original embedding table with expanded tables indexed by hash mappings over token n-grams.
  • N-gram tables can be decomposed into K sub-tables with different vocabulary sizes to improve expressiveness and reduce hash collisions.
  • Additional linear projections map sub-table outputs back to the original embedding space.
  • Sub-table hidden size is inversely proportional to the number of sub-tables, keeping N-gram Embedding parameter count invariant with respect to N and K.

3 Comparative Analysis of Expert and Embedding Scaling

The comparative experiments evaluate N-gram Embedding scaling against expert scaling through from-scratch pre-training under matched activated-parameter budgets. The framework varies base sparsity and incrementally adds embedding capacity.

  • Experiments compare N-gram Embedding scaling with expert scaling at 280M, 790M, and 1.3B activated-parameter budgets.
  • The study trains MoE models with base sparsity levels from 35% to 98% before incrementally incorporating N-gram Embedding.
  • Each scaling strategy is evaluated through a framework designed to contrast embedding expansion with increasing expert capacity.

3.1 Optimal Timing for N-gram Embedding Integration

N-gram Embedding is most advantageous when introduced after expert scaling reaches high sparsity or exceeds its effective sweet spot. At excessive embedding allocation, however, its advantage can reverse against MoE baselines.

  • 3.1 Optimal Timing for N-gram Embedding Integration: At low parameter ratios, increasing experts reduces loss more effectively than introducing N-gram Embedding.
  • 3.1 Optimal Timing for N-gram Embedding Integration: At high sparsity levels, N-gram Embedding benefits become more pronounced than expert-scaling gains.
  • 3.1 Optimal Timing for N-gram Embedding Integration: N-gram Embedding should be introduced when the number of experts exceeds its sweet spot.
  • 3.1 Optimal Timing for N-gram Embedding Integration: Embedding scaling is a promising dimension orthogonal to expert scaling.
  • 3.1 Optimal Timing for N-gram Embedding Integration: Figure 2 compares MoE and N-gram Embedding scaling by total-to-activated parameter ratio, with logarithmic right-panel axes and dashed base-MoE connections.

3.2 Integration Strategy

N-gram Embedding requires careful parameter allocation, vocabulary sizing, hyperparameter selection, initialization, and residual-signal preservation to outperform MoE baselines.

  • Parameter Budgeting: Embedding scaling eventually loses its advantage when N-gram Embedding receives an excessive parameter share, with the intersection slightly above a ratio of 20.At that point, N-gram Embedding parameters constitute approximately 50% of total parameters.
  • Vocabulary Sizing: Hash collisions increase ambiguity by forcing one embedding vector to represent multiple distinct n-grams, degrading learning efficiency and performance.Collision counts spike when vocabulary size approaches integer multiples of the base vocabulary size.
  • Hyperparameters: N-gram Embedding is robust across configurations with N ≥3 and K ≥2, whereas N = 2 and K = 1 produces notably inferior performance.Increasing N captures richer context but increases sparsity, while increasing K reduces collisions with diminishing returns beyond a threshold.
  • Diagnostics: The scaling curves compare training and validation loss across combinations of N and K, while the layer-wise analysis compares module-output norms with identity-branch norms.The figures use these comparisons to assess hyperparameter sensitivity and residual-signal balance.
  • Initialization: The first attention module output reaches approximately 10× the corresponding identity-branch norm, drowning out the embedding signal during residual summation.This disparity emerges during training despite standard initialization matching baseline output norms initially.
  • Embedding Amplification: Embedding Amplification improves performance over the vanilla baseline, reducing training loss and both validation losses by 0.02.The approach includes normalization or another strategy that increases the embedding contribution to the forward pass.

3.3 Scaling Properties across Model Width and Depth

N-gram Embedding becomes relatively more effective as model width increases, while increasing depth beyond 20 layers contracts its advantage over MoE baselines.

  • Width: The width experiments vary hidden size and module dimensions at fixed depth, using 790M and 1.3B activation scales.Depth is held constant at 10 shortcut layers in these experiments.
  • Width: At 1.3B activation size, N-gram Embedding retains a clear advantage over MoE at total-to-activated parameter ratios as high as 50.At 280M activation size it underperforms beyond 30, while at 790M it underperforms at that ratio only on English validation.
  • Width: Increasing model width expands the regime in which N-gram Embedding outperforms a parameter-equivalent MoE baseline.The intersection between the N-gram Embedding and MoE curves shifts toward higher total-to-activated parameter ratios as activation size increases.
  • Depth: The depth experiments use 20- and 40-layer architectures while maintaining N-gram Embedding at 50% of total parameters.They build on the 1.3B activated-parameter configuration.
  • Depth: Beyond 20 layers, the performance advantage of N-gram Embedding over the MoE baseline contracts sharply.This depth trend contrasts with width scaling, where the performance gap widens.
  • Practical Scope: Within the common range below 40 shortcut layers, the paper reports robust N-gram Embedding performance and expects width scaling to amplify its gains.The stated architectural boundary is below 40 shortcut layers, equivalent to 80 conventional layers.

4 Efficient Inference

N-gram Embedding reduces MoE-layer activation and memory I/O, but its additional lookup overhead requires caching, scheduling-aware kernels, and speculative decoding to realize inference speedups.

  • Activation Reduction: N-gram Embedding reallocates parameters from MoE layers to embeddings, reducing activated MoE parameters while preserving total model size.The embedding lookup cost scales with input tokens rather than the expanded parameter count.
  • Speculative Decoding: Large batch sizes are needed to maximize hardware utilization and convert reduced active parameters into practical efficiency.Speculative decoding expands the effective batch size through multi-step drafting.
  • Overhead: N-gram Embedding introduces additional I/O, computation, and communication overhead relative to a standard embedding layer.Dynamic inference scheduling also complicates precomputing the token sequences needed for optimized lookups.
  • Caching: The N-gram Cache and custom CUDA kernels manage N-gram IDs on-device and reduce synchronization overhead under complex inference scheduling.The design is inspired by KV-cache principles.
  • Draft Optimization: For speculative decoding, conventional draft-model embeddings bypass expensive n-gram lookup, while caching draft-phase n-gram embeddings removes redundant computation.These are complementary strategies for reducing the relative overhead of N-gram Embedding in low-latency draft models.
  • Future Directions: N-gram Embedding may support faster speculative decoding through local-context drafting or early rejection of low-probability draft tokens.These directions are presented as potential ways to reduce target-model verification work.

5 Integration with Per-Layer Embedding

The section compares per-layer embedding strategies and introduces PLNE, which applies N-gram Embedding at each layer. Although PLNE slightly improves over N-gram Embedding in some settings, its advantage is inconsistent when model width or depth increases.

  • 5.1 Per-Layer Embedding: PLE replaces the SwiGLU up-projection output with embedding output, providing an efficient way to inject embedding information into each layer.PLE allocates embedding parameters across layers and injects them into the dense sub-layer’s MLP.
  • 5.2 Per-Layer N-gram Embedding: PLNE replaces base embedding outputs with N-gram Embedding outputs at every layer for more targeted parameter scaling.Each layer uses a layer-specific embedding table and projection matrix.
  • 5.3 Empirical Comparison: PLE underperforms N-gram Embedding, whereas PLNE yields marginal improvements over N-gram Embedding in the comparison.The authors attribute PLE’s weaker result to N-gram Embedding’s superior learning efficiency.
  • 5.3 Empirical Comparison: PLNE performs on par with N-gram Embedding in most increased-width or increased-depth scenarios, without a consistent advantage.Because PLNE also increases activated parameters through per-layer projection matrices, the authors exclude it from larger-scale experiments.
  • 5.3 Empirical Comparison: The optimal allocation of PLNE parameters across layers remains open, including whether to concentrate them in selected layers or distribute them uniformly.This unresolved allocation question motivates further investigation.

6 LongCat-Flash-Lite

LongCat-Flash-Lite integrates N-gram Embedding into a sparsely activated MoE model and evaluates it against an equivalent expert-scaled baseline and comparable models. It achieves lower training loss and broad benchmark strength, especially in agentic tool use and coding, while system optimizations address inference bottlenecks.

  • Architecture: LongCat-Flash-Lite has 68.5B total parameters and dynamically activates 2.9B–4.5B parameters per token.Its embedding module contains 31.4B N-gram Embedding parameters, while each token selects 12 experts from the MoE layers.
  • Training Data: The model is pretrained on 11T tokens, mid-trained on 1.5T tokens with extended context, and then supervised-finetuned.The sequence length increases from 8k during pretraining to 128k during mid-training, with support for sequences up to 256k tokens.
  • Baseline: The parameter-matched LongCat-Flash-Lite-Vanilla baseline converts N-gram Embedding parameters into additional experts and uses the same training strategy and data recipe.This isolates the comparison between embedding scaling and expert scaling.
  • Base Model Evaluation: LongCat-Flash-Lite consistently achieves lower training loss and substantially improves over LongCat-Flash-Lite-Vanilla across most benchmarks in general, reasoning, and coding domains.The evaluation covers MMLU-family, BBH/GPQA/DROP/GSM8K, HumanEval+, MultiPL-E, and BigCodeBench tasks.
  • Agentic Tool Use: LongCat-Flash-Lite leads the comparison models in τ2-Bench tool-use scenarios, scoring 72.8 on Telecom, 73.1 on Retail, and 58.0 on Airline.It also scores 7.00 on VitaBench, above Qwen3-Next-80B-A3B-Instruct’s 5.80 and Gemini 2.5 Flash-Lite’s 4.50.
  • Agentic Coding: 54.4 accuracy on SWE-Bench exceeds Qwen3-Next-80B-A3B-Instruct (37.6), Gemini 2.5 Flash-Lite (41.3), and Kimi-Linear-48B-A3B (32.8).The model also scores 33.75 on TerminalBench, 38.10 on SWE-Bench Multilingual, and 39.63 on PRDBench.
  • Inference Optimization: System optimizations combine three-step speculative decoding, wide expert parallelism, single-batch overlap, and kernel-level changes to address bandwidth and kernel-launch bottlenecks.The authors report exceptional inference performance after applying these optimizations.

7 Conclusions

The report concludes that embedding scaling can outperform expert scaling in specific sparsity regimes and presents LongCat-Flash-Lite as a validation of this framework. Its system optimizations address associated I/O bottlenecks while the model remains competitive in agentic and coding tasks.

  • 7 Conclusions: Embedding scaling achieves a superior Pareto frontier to increasing expert numbers in specific regimes identified through systematic analysis.The study examines architectural constraints and comparative scaling laws for embedding scaling in LLMs.
  • 7 Conclusions: LongCat-Flash-Lite combines over 30B N-gram Embedding parameters with a 68.5B-parameter MoE design and outperforms parameter-equivalent MoE baselines.The model also demonstrates competitive performance in agentic and coding tasks.
  • 7 Conclusions: The N-gram Cache and synchronized kernels resolve associated I/O bottlenecks in the proposed inference pipeline.These optimizations are presented as part of the efficiency framework for future model scaling.
Loading 2601.21204v2…