ERNIE 5.5 Flash Base

TL;DR

Baidu releases ERNIE-5.5-Flash-Base as the end-of-pretraining checkpoint of a 280B-total / 14B-active sparse MoE trained from scratch on 33.5T tokens with Muon in FP8, with 128K training context extended to 1M for fine-tuning / inference. The architecture packs four distinct design-space moves — Multi-Router latent MoE (several independently-initialized sigmoid routers jointly scoring experts), a 1-in-7 MLA + 6-HCA hybrid-sparse-attention pattern where most layers compress every 128 tokens into one cached entry, manifold-constrained Hyper-Connections (mHC) on the residual stream, and a MultiMax piecewise output-head temperature — then ships weights (FP8 + BF16, Hugging Face), an FP8-native distributed Muon recipe, and the FlashMask-4 CuTe-DSL attention kernel for large head dimensions. Headline: strong across MMLU-Pro / HumanEval+ / LiveCodeBench / long-context-RAG, matching or beating DeepSeek-V4-Flash, Hy3-preview and MiMo-V2.5 at comparable active params while compressing the KV cache 128× in 37 of 43 layers.

Key claims

Standard single-router MoE has a silent self-reinforcement bug: because experts left out of top-k receive no gradient from that token, the router gets no signal that it chose wrongly, so wrong selections reinforce themselves — and single-router selection scores correlate only weakly with how useful each expert actually is, especially in middle and deep layers [§2.2].
Multi-Router replaces the single router with several independently-initialized sigmoid routers that score the experts jointly; the routers develop clearly distinct preferences and routing scores correlate noticeably better with real expert usefulness, with the largest gains in middle and deep layers — at identical activated-expert count, FLOPs, and communication volume [§2.2].
Hybrid MLA + HCA attention: only 1 layer in 7 is MLA (token-level access to the full context); the other 36 are HCA, each compressing every 128 tokens into a single cached entry, so an HCA layer still sees the whole sequence while caching a small fraction of what full-context layers need [§2.1].
Layer-specific head dimensions are set by per-layer attention cost rather than kept uniform: HCA layers get wide 512-dim heads at little extra compute, MLA layers use narrower 256-dim query/key/value because they attend at token level and cost more [§2.1].
With HCA in most layers, concentrating the remaining KV-cache budget on a few full-token-retention MLA layers beats concentrating it on many Compressed Sparse Attention (CSA) layers at similar KV-cache budget in pre-training comparisons [§2.1].
Full-context MLA layers use DeepSeek-style DSA sparse attention — a lightweight indexer scores past positions and core attention is only over top-k — matching dense attention in pre-training loss, substantially faster, with per-query core-attention cost constant as context grows [§2.1].
Low-Rank Gated Attention computes the output gate from the low-rank query latent each attention layer already produces rather than from the full hidden state, yielding 75% fewer gate-projection weights and better ablation numbers; Virtual Head Attention adds a lightweight low-rank cross-head mixing step with an identity skip, fused into the attention kernel at inference [§2.1].
Latent MoE (project 4096 → 2048 before experts, shared up-projection back to 4096) roughly halves dispatch/combine traffic and expert FLOPs vs full-width experts, while Stable LatentMoE design (pre-up-projection RMSNorm, SiTU-GLU) + per-expert learnable routing scalars keep expert outputs well-scaled across all 43 layers [§2.2].
Manifold-constrained Hyper-Connections (mHC) replace the standard residual with 4 parallel residual streams mixed through a constrained matrix — gives layers a wider channel for passing information with negligible extra compute [§2.3].
MultiMax output head: a single softmax temperature trades sparsity against keeping plausible next tokens; a learnable piecewise temperature rescales logits element-wise inside the fused cross-entropy kernel, suppressing the long tail of implausible tokens while keeping several plausible ones, with negligible training overhead and no extra memory, and significantly lower pre-training loss [§2.4].
ECHO replaces MTP: future-token prediction reuses the backbone’s own final layer rather than attaching extra MTP modules, so the loss on future tokens directly trains the same parameters that produce the next token; the same layer serves as the draft model for speculative decoding and gets higher acceptance drafting several tokens ahead than an auxiliary MTP module trained the same way [§2.5].
Right-to-left BPE tokenization in groups of up to three digits lines up place-value digits across numbers of different lengths, targeting arithmetic and number-heavy text [§2.6].
Distributed Muon: dense parameter synchronization is moved onto the same expert-parallel path MoE parameters use and combined into one collective — messages become larger, bandwidth is used more fully, inter-node latency and jitter drop, with optimizer updates unchanged; inter-node balanced sharding spreads optimizer states across more GPUs so memory is evenly balanced without shrinking message sizes [§3.2].
FlashMask-4 (CuTe-DSL, SM90 + SM100) supports head dimensions up to qk_head_dim=576 / v_head_dim=512 — beyond FlashAttention-4 — via 2-CTA tiled parallel computation and (on SM100) double-buffer tensor-memory rotation, up to 175% faster forward / 74% faster backward than FA-4 MaskMod on SM100 [§3.2].
Despite 37 of 43 layers attending over 128×-compressed sequences, the 6 MLA-DSA layers are enough to keep performance strong on long-context benchmarks: HELMET-RAG 71.5, HELMET-ICL 85.9, RepoQA 91.7, LongBench-v2 48.7, LeadIng on all four vs V4-Flash / Hy3-preview / MiMo-V2.5 [§4.2, Table].
Benchmark leadership at 14B active params on MMLU 93.3 / MMLU-Pro 78.4 / GPQA 60.2 / BBH-direct 77.1 / HumanEval+ 82.3 / MBPP+ 81.5 / LiveCodeBench 40.1 / RepoBench 50.3, beating DeepSeek V4 Flash (13B active), Hy3-preview (21B active) and MiMo V2.5 (15B active) [§4.3, Table].

Method

The model is a 43-layer, 4096-hidden, 64-head sparse MoE. The attention stack is a hybrid of sparsity patterns: 6 of 43 layers are MLA with DSA sparse attention (full-context, token-level, top-k over learned indexer scores), and 37 are HCA (Heavily Compressed Attention) — each HCA layer compresses every 128 tokens into a single cached entry and attends to the compressed sequence plus a short uncompressed recent window, caching a small fraction of what a full-context layer needs. Per-layer head dimension is set by attention cost: HCA layers get 512-dim heads (cheap), MLA layers get 256-dim heads (expensive). Every attention layer uses Low-Rank Gated Attention (gate computed from the low-rank query latent the attention layer already produces, 75% fewer gate weights than hidden-state-sourced gating) and Virtual Head Attention (identity-skipped low-rank cross-head mixing, fused into the attention kernel at inference).

The MoE block is a Multi-Router Latent MoE: 512 routed experts + 1 shared expert (top-10 routed) in all layers except the first, inside a 2048-d latent space (4096 → 2048 projection, shared up-projection back). Rather than one router, several independently-initialized sigmoid routers jointly score the experts — the routers develop distinct preferences so an expert can be chosen because one router strongly favors it or several do together. The claim is that this directly fixes the “experts outside top-k get no gradient, so wrong selections reinforce themselves” failure mode of single-router top-k: multiple scores give the system a signal to correct against. Load balancing is auxiliary-loss-free with bias correction; each expert has its own learnable scalar on its routing weight; Stable LatentMoE (RMSNorm before up-projection, SiTU-GLU expert activation) keeps expert output scale stable across all 43 layers.

Beyond MoE and attention, three more primitives: mHC replaces the residual with 4 parallel residual streams mixed by a constrained matrix (wider inter-layer channel, negligible extra compute); MultiMax replaces the fixed output-head softmax temperature with a learnable piecewise temperature that rescales logits inside the fused CE kernel (suppresses implausible-token long tail while preserving several plausible next tokens); ECHO replaces MTP by reusing the backbone’s own final layer for future-token prediction, so future-token loss trains the same parameters that produce the next token — the same layer doubles as the speculative-decoding draft model with higher acceptance than auxiliary MTP.

Pre-training used Muon + FP8 mixed precision over 33.5T tokens, with context gradually extended to 128K during pretraining. The infrastructure contribution is a distributed Muon that moves dense-parameter all-reduce onto the same expert-parallel ring MoE uses and combines dense + MoE sync into one collective (larger messages, lower inter-node jitter), plus inter-node balanced sharding that keeps optimizer-state memory flat across nodes. FlashMask-4 (CuTe-DSL) handles the large head dimensions (qk_head_dim up to 576, v_head_dim up to 512) via 2-CTA tiling + SM100 tensor-memory double buffering.

Results

At 14B activated / 280B total params, ERNIE-5.5-Flash-Base matches or beats DeepSeek-V4-Flash-Base (13B active / 284B total), Hy3-preview-Base (21B active / 295B total) and MiMo-V2.5-Base (15B active / 310B total) — all evaluated under the same harness, prompts, few-shot examples and LLM judge. Exam: MMLU 93.3 vs 89.4 / 87.5 / 87.6; MMLU-Pro 78.4 vs 69.2 / 66.9 / 65.7; MMLU-Redux 94.1 vs 90.9 / 88.7 / 89.4 [Table]. Code: HumanEval+ 82.3 vs 69.5 / 78.0 / 74.4; MBPP+ 81.5 vs 67.2 / 78.3 / 74.3; LiveCodeBench 40.1 vs 31.9 / 33.5 / 32.4; RepoBench 50.3 vs 49.1 / 49.1 / 48.9 [Table]. STEM: GPQA 60.2 vs 55.9 / 50.2 / 51.4; MGSM-CoT 87.0 vs 84.6 / 76.9 / 83.8 [Table]. Reasoning: BBH-direct 77.1 vs 69.0 / 67.5 / 68.0 [Table]. Long-context (where 37 of 43 layers attend over 128× compressed sequences): HELMET-RAG 71.5, HELMET-ICL 85.9, RepoQA 91.7, LongBench-v2 48.7 — leading on 4 of 6 long-context benchmarks including RepoQA where it beats MiMo-V2.5 91.7 vs 56.5 [Table]. Knowledge is more mixed — SimpleQA 34.1 leads but ChineseSimpleQA 73.0 trails Hy3-preview-Base 75.3 [Table]. The infrastructure numbers: FlashMask-4 is up to 175% faster forward and 74% faster backward than FlashAttention-4 MaskMod at the model’s head dimensions on SM100 [§3.2].

Why it’s interesting

The Multi-Router story is a direct architectural answer to the silent-expert-death / routing-collapse thread that Mitigate Silent Expert Death In Ultra-Sparse MoE, Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models, and The Myth of Expert Specialization in MoEs: Why Routing Reflects Geometry, Not Necessarily Domain Expertise have been building out: rather than diagnosing why single-router training fails (shallow-expert second-moment collapse, micro-batch-scope LBL forbidding specialization, geometry-not-domain routing), it changes the router topology so that “wrong choice at top-k gets no gradient” stops being a silent trap. Complements Routing-Free Mixture-of-Experts‘s “eliminate the router entirely” and Path-Constrained Mixture-of-Experts‘s “share router params across layers” as a third position on the routing-design board — more routers, independently initialized, voting. Meanwhile the 6-MLA-DSA + 37-HCA hybrid is the most aggressive compressed-attention ratio on the wiki: compare GLM-5: from Vibe Coding to Agentic Engineering (GLM-5 uses DSA on all attention layers) and DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention (V3.2-Exp introduces DSA as a drop-in sparse attention), and ERNIE goes a step further by making most layers lossy (128× compression of history) while keeping just enough token-level layers for long-context recall — and the long-context benchmark numbers (HELMET-RAG, RepoQA) validate that 6 MLA layers out of 43 is enough. Pairs the attention and MoE engineering with mHC (also appearing in mHC: Manifold-Constrained Hyper-Connections and DeepSeek V4.1-Flash per DeepSeek-V4.1-Flash — Pushing the Limits of KV Cache Compression), ECHO reusing the final backbone layer as its own MTP module, and a distributed-Muon-over-EP-ring recipe that extends the Muon-scaling thread from How Muon Lost Its Geometry and Dion3: Full-Stack Orthogonal Updates to the 33.5T-token 280B-MoE regime.