PaperlayerView exact revision on arXiv

Source-linked AI summary

Mellum2 Technical Report

Marko Kojic, Ivan Bondyrev, Aral de Moor, Joseph Shtok, Petr Borovlev, Kseniia Lysaniuk, Madeeswaran Kannan, Ivan Dolgov, Nikita Pavlichenko

arXiv:2605.31268v1cs.CL

TL;DR

Developers need coding models that cover increasingly broad software-engineering tasks while remaining practical to deploy, balancing the limitations of dense models against the cost of large MoEs. Mellum 2 addresses this with a sparse, deployment-oriented coding assistant and is competitive with open-weight 4–14B baselines at 2.5B-dense-model per-token compute, while exposing safety trade-offs after RL.

  • Problem

    Coding models must support generation, editing, debugging, reasoning, tool use, agentic workflows, and conversational assistance, while dense models plateau on harder workloads and large MoEs impose high deployment costs.

  • Method

    Mellum 2 is an open-weight 12B-parameter MoE coding assistant with 2.5B active parameters per token, combining deployment-oriented attention changes and a Multi-Token Prediction head.

  • Results

    Mellum 2 is competitive with open-weight baselines in the 4–14B range while using the per-token compute of a 2.5B dense model and matching or exceeding Qwen2.5-7B inference throughput on one H100.

  • Takeaways & Limitations

    The model targets practical, latency-sensitive IDE deployment while supporting code generation, editing, reasoning, tool calls, and agentic workflows.

  • Takeaways & Limitations

    RL increases HarmBench from 8.4% for Mellum 2-SFT to 23.1%, reflecting a regression in refusal behavior and an alignment trade-off in the recipe.

Abstract

from arXiv · show

We present Mellum 2, an open-weight 12B-parameter Mixture-of-Experts (MoE) language model with 2.5B active parameters per token. Mellum 2 is a general-purpose language model specialized in software engineering, spanning code generation and editing, debugging, multi-step reasoning, tool use and function calling, agentic coding, and conversational programming assistance, and it is the successor to the completion-focused 4B dense Mellum model. The architecture builds on the Mixture-of-Experts (64 experts, 8 active) and combines Grouped-Query Attention with 4 KV heads, Sliding Window Attention on three of every four layers, and a single Multi-Token Prediction head that doubles as both an auxiliary pre-training objective and a built-in draft model for speculative decoding; each choice was validated by ablation with inference efficiency on commodity GPUs as a design constraint. Pre-training spans approximately 10.6 trillion tokens through a three-phase curriculum that progressively shifts the mixture from diverse web data toward curated code and mathematical content, optimized with Muon under FP8 hybrid precision and a Warmup-Hold-Decay schedule with linear decay to zero. The pre-trained base is extended to a 128K context window via a layer-selective YaRN and then post-trained in two stages (supervised fine-tuning followed by RLVR), yielding two released variants: an Instruct model that answers directly and a Thinking model that emits an explicit reasoning trace before its final answer. Across code generation, math and reasoning, tool use, knowledge, and safety benchmarks, Mellum 2 is competitive with open-weight baselines in the 4B-14B range while running at the per-token compute of a 2.5B dense model. We release the base, instruct, and thinking checkpoints, together with this report on the architecture decisions, data pipeline, and training recipe behind them, under the Apache 2.0 license.

1 Introduction

Mellum 2 targets the broad demands of modern software engineering while balancing model capability with practical serving cost. It combines an efficiency-aware MoE architecture, large-scale curriculum training, long-context extension, and two post-trained variants, and performs competitively against open-weight baselines.

  • Modern coding assistants must support code generation, editing, debugging, reasoning, tool use, repository navigation, and conversational collaboration at deployable serving cost.
  • Mellum 2 is a 12B-parameter MoE model with 2.5B active parameters per token, designed as a general-purpose successor to the 4B dense Mellum completion model.Its software-engineering scope includes code generation and editing, tools, multi-step workflows, long conversations, and explicit reasoning traces in the Thinking variant.
  • The architecture uses 64 experts with 8 active, four KV heads, sliding-window attention on three of every four layers, and one MTP head for auxiliary training and speculative decoding.These choices were evaluated through ablations with inference efficiency as a design constraint.
  • Pre-training uses a three-phase curriculum over approximately 10.6T tokens, shifting from diverse web data toward curated code and mathematics while using Muon, FP8 precision, and linear learning-rate decay to zero.The code mixture rises from 23% to 42% to 59% across the phases.
  • The pre-trained base is extended to 128K context with layer-selective YaRN, then post-trained into Instruct and Thinking variants using supervised fine-tuning and RLVR.The Instruct model answers directly, while the Thinking model emits an explicit reasoning trace before its final answer.
  • Mellum 2 is released with base, instruct, and thinking checkpoints under the Apache 2.0 license.The release includes the report documenting the architecture, data pipeline, and training recipe.
  • Across code, math, reasoning, tool-use, knowledge, and safety benchmarks, Mellum 2 is competitive with 4–14B open-weight baselines at the per-token compute of a 2.5B dense model.It also matches or exceeds Qwen2.5-7B inference throughput on a single H100.

2 Model Architecture

Mellum 2 uses an efficiency-oriented Qwen3-MoE-style decoder architecture, selected through ablations against a Qwen2.5-7B latency budget. Its 12B total parameters with 2.5B active per token combine sparse experts, GQA, SWA, and MTP to balance quality, latency, and throughput.

  • Architecture search: Dense configurations did not consistently outperform Qwen2.5-7B within the latency budget, motivating the adoption of MoE scaling.MLA reached approximately 5.5B parameters at equivalent speed, but quality gains did not justify its added complexity.
  • Expert configuration: Higher expert sparsity improved inference speed, while more active experts improved benchmark quality; the selected trade-off uses 8 of 64 experts.Two active experts achieved approximately 1.5× lower latency than eight active experts, but the final design favored quality–latency balance.
  • Grouped-Query Attention: Four KV heads were selected because eight degraded throughput and two reduced evaluation quality, with KV-cache effects strongest under high concurrency.In synchronous mode, latency is more strongly dominated by model depth than by KV-cache utilization.
  • Sliding Window Attention: Sliding Window Attention on three of every four layers uses a 1,024-token window to reduce attention cost while retaining full attention periodically.The 1,024-token window outperformed a 512-token window on quality benchmarks, and the remaining full-attention layers preserve long-range context capability.
  • Multi-Token Prediction: The single MTP head predicts one additional future token, adds only 7% training time in ablation, and serves as a speculative-decoding draft model.It is removed during inference and improved HumanEval by +10.4, MMLU by +3.6, MMLU-Pro by +3.3, and GSM8K by +3.0 in the reported ablation.
  • Final architecture: The final architecture is a 28-layer, 2,304-hidden-dimension decoder with 32 query heads, 4 KV heads, 64 experts, and 8 active experts.It totals approximately 12B parameters with approximately 2.5B active per token and uses a 3:1 SWA pattern with a 1,024-token window.

3 Pre-Training

Mellum 2’s pre-training combines a diverse 10.6T-token corpus with a three-phase curriculum that increasingly emphasizes curated code and mathematics. Muon-based optimization and stability analyses support training, while evaluations show competitiveness with larger dense models despite 2.5B active parameters.

  • Data and curriculum: Approximately 10.6T tokens span web and general knowledge, source code, and mathematical content.
  • Data and curriculum: The three phases shift code from 23% to 42% to 59% while moving from broad web data toward curated code and mathematics.Phase 1, 2, and 3 process approximately 6.18T, 2.79T, and 1.69T tokens, respectively.
  • Data and curriculum: Repeated high-quality data supports expert specialization, with raw code seen for three epochs and no dataset repeated more than four times.The raw code corpus contributes roughly 958B tokens.
  • Optimization and stability: Muon configurations converged successfully on the MoE ablation, while the selected Moonlight setup remained stable across dense and MoE architectures.On the dense 7B run, Moonlight reduced validation loss by 0.028, approximately 2.5%, relative to AdamW.
  • Optimization and stability: Adam epsilon 10^-8 provided the best trade-off between training stability and optimization effectiveness.Values as large as 10^-5 caused disproportionate update dampening.
  • Optimization and stability: Training loss decreased steadily across phases, while the router’s load-balancing loss stabilized as expert assignment became effective.Periodic duplicate-induced loss spikes were modest, isolated, and had no measurable effect on training dynamics.
  • Evaluation: Despite activating 2.5B parameters per token, Mellum 2 is competitive with 7B dense models and exceeds them on several reasoning and code tasks.It achieves 59.3% on MMLU-Pro, 74.9% on BBH, 81.7% on GSM8K, 62.4% on MBPP, 61.4% on MBPP+, and 35.0% on GPQA Main.

4 Long Context Extension

Mellum 2 extends its 8,192-token training context to 128K using layer-selective YaRN applied only to global-attention layers. This recipe outperforms uniform and unchanged-parameter alternatives, while long-context gains largely saturate early in extension training.

  • Extension design: The effective context length is extended from 8,192 to 131,072 tokens through a dedicated long-context training stage.
  • Extension design: YaRN is applied only to global full-attention layers, leaving sliding-window layers’ original RoPE parameters unchanged.
  • Recipe ablations: At 64K evaluation context, layer-selective YaRN scores 0.64 on RULER versus 0.52 for uniform θ-bump and 0.33 for unchanged θ.The unchanged-θ run collapses past 32K, while uniform scaling perturbs sliding-window layers unnecessarily.
  • Training data: The extension mix combines rebalanced Phase 3 data with agentic SFT data and subsamples long reasoning traces to avoid skewing general long-context behavior.
  • Recipe ablations: Adding the Longmino mix lowered RULER by roughly 2–3 percentage points at every measured context length under otherwise matched conditions.
  • Training data: Repository-level FIM examples preserve in-IDE completion by placing related files before the prefix, middle, and suffix target.
  • Training dynamics: Within approximately 30B extension tokens, RULER at every measured context length came within approximately 1 percentage point of its final value at 117B tokens.After that point, the MoE router’s load-balancing loss continued to change substantially.

5 Post-Training

Mellum 2 is post-trained from a shared long-context checkpoint through supervised fine-tuning and reinforcement learning, producing direct-answering Instruct and reasoning-trace Thinking variants. The resulting model is strongest on code and developer-oriented interaction, with reinforcement learning improving tool use and reasoning while broad knowledge and some safety behaviors remain weaker.

  • Post-training pipeline: Mellum 2 post-training begins from a long-context YaRN checkpoint and proceeds through supervised fine-tuning followed by reinforcement learning.
  • Post-training pipeline: The Instruct variant answers directly, while the Thinking variant emits a reasoning trace before its final answer.Both variants are trained from the same long-context base checkpoint, with different handling of reasoning traces and loss masking.
  • Supervised fine-tuning data: The SFT corpus covers general instruction following, single-turn coding, repository-level agentic coding, math, knowledge, and reasoning traces.Math examples include pure solving, calculator use, and Python code execution; knowledge data spans multiple academic and technical domains.
  • Reinforcement learning: RL training accounts for inference–training policy discrepancies caused principally by MoE routing differences and removes token contributions whose IcePop ratio falls outside a configured band.Reward shaping also applies a soft overlong penalty and drops rollouts exceeding the response-length cap.
  • Evaluation: 78.4% on EvalPlus is achieved by Mellum 2-RL, leading every listed baseline for robust function-level code synthesis.Mellum 2-SFT-Thinking reaches 75.1 on LiveCodeBench v6, the top score in that panel, while the instruct variant scores 37.2.
  • Evaluation: BFCL v3 improves from 43.1 to 66.3 for Instruct and from 60.5 to 69.4 for Thinking after RL, while RL-Thinking leads BFCL v4 at 45.6.Math performance also improves in the Thinking variant, with GSM-Plus reaching 87.0.
  • Evaluation: MMLU-Redux and GPQA Diamond remain major weaknesses, with Mellum 2 below the Qwen3.5 9B reference on both benchmarks.The authors attribute this profile to training emphasis on code and developer documentation rather than broad encyclopedic coverage, limiting off-domain use.
  • Evaluation: HarmBench improves safety for Mellum 2-SFT at 8.4% but regresses to 23.1% after RL, while XSTest indicates roughly ten points of over-refusal relative to larger baselines.The report identifies both behaviors as targets for future joint optimization.

6 Efficiency and Deployment

Mellum 2 was designed for practical deployment in latency-sensitive IDEs, targeting Qwen2.5-7B’s single-H100 performance while improving sustained throughput.

  • The deployment target is inference speed comparable to Qwen2.5-7B for latency-sensitive IDE environments.
  • In sync mode, Mellum 2 reaches 192 tokens/s versus Qwen2.5-7B’s 193 tokens/s, while concurrent throughput is 21% higher.
  • Mellum 2 matches Qwen2.5-7B’s sync latency and delivers 21% higher sustained throughput on a single H100.The benchmark uses vLLM FP8 serving at IS-L/OSL = 2,304/256.

7 Conclusion

Mellum 2 concludes as an open, inference-aware small-MoE coding model with matched Instruct and Thinking variants, broad software-engineering capabilities, and competitive performance at low active-parameter cost. The authors identify further work in repository-level SWE reinforcement learning, RL infrastructure, and long-context training.

  • Mellum 2 is released as matched Instruct and Thinking variants, with 2.5B active parameters and Apache 2.0 licensing.
  • The model supports code generation and editing, engineering reasoning, tool calls, and agentic workflows inside IDEs at practical per-token cost.
  • Architecture, training, context extension, and post-training combine into a model competitive with 4–14B open-weight baselines at 2.5B dense-model-equivalent per-token compute.
  • Future directions include repository-level SWE reinforcement learning, broader RL infrastructure and environment coverage, and revisiting the long-context mid-training mixture.
  • The inference-aware fixed-budget design is presented as a path toward scaling Mellum to larger models.
  • The released checkpoints and report provide an open recipe and inference-aware design point for small-MoE coding models.

A.1 Dense Architecture Exploration

Dense architecture exploration found that deeper or wider variants did not consistently surpass Qwen2.5-7B under the latency constraint, while MLA offered insufficient quality gains at the tested scale.

  • The dense exploration evaluated Qwen3-based variations across two architectural axes.
  • Deeper Qwen3-based dense variants failed to consistently outperform Qwen2.5-7B under the latency constraint because additional sequential operations degraded inference performance.
  • Wider and shallower dense variants improved inference performance but still did not consistently exceed the Qwen2.5-7B quality baseline.
  • MLA scaled the dense model to approximately 5.5B parameters at Qwen2.5-7B latency, but quality improvements were insufficient and KV-cache savings were limited.

A.2 MoE Architecture Exploration

MoE exploration selected 64 experts with 8 active per token as the best quality–latency trade-off, while shared experts and dense interleaving harmed efficiency without matching quality gains.

  • The exploration fixed the expert count at 64 because larger counts exceeded GPU memory constraints.
  • 8 active experts out of 64 provided the best quality–latency trade-off, whereas 2 active experts reduced latency by approximately 1.5× but substantially worsened quality.
  • Up to approximately 15B total parameters were feasible while matching Qwen2.5-7B latency with 8 active experts.
  • Adding a shared expert produced no measurable quality gain and consistently hurt inference performance because of extra always-on FFN computation.
  • Replacing some MoE layers with dense FFN layers hurt inference performance without a matching quality improvement, so all FFN layers remained MoE.
  • The final design retained auxiliary-loss load balancing to fit the Qwen3-MoE module layout despite short-run alternatives with similar expert utilisation.

A.3 Hybrid Architecture Exploration

The authors evaluated hybrid attention designs that replace most full-attention layers with Gated DeltaNet, but found them unsuitable for Mellum 2’s short-context, single-batch deployment target at the time of study.

  • Hybrid design: Hybrid variants retained full attention in every fourth layer while replacing the other layers with Gated DeltaNet token mixers.This design provides a fixed-size recurrent state instead of a linearly growing KV cache.
  • Expected efficiency: On long-context, large-batch workloads, these hybrids offer near-constant per-token decode cost by eliminating the linearly growing KV cache.
  • Deployment constraint: At the target of short-context, single-batch in-IDE inference, every tested hybrid exhibited a substantial latency regression.The authors therefore did not select the hybrid architecture for Mellum 2.
  • Future direction: The authors attribute the gap to kernel and framework maturity rather than a fundamental limitation, and plan to revisit hybrids as inference optimization improves.

A.4 MoE Training Hyperparameters

Preliminary MoE experiments selected training settings that prioritize flexible load balancing and unconstrained expert utilization, while adopting dropless routing for throughput benefits during training.

  • Balancing strategy: Global-batch balancing was selected over slightly better per-sequence test loss because it handles variable batch sizes more flexibly.
  • Auxiliary loss: The full pretraining run used an auxiliary loss coefficient of 10^-3 rather than the better-performing 10^-2 short-run setting to avoid over-constraining expert utilization.
  • Token dropping: Dropless routing was adopted because capacity factors from 1.0–1.5 produced no meaningful quality difference, while routing throughput improved as load balancing developed.The residual overhead was approximately 15% at the time of writing.

B Training Hyperparameters (Full)

The full pretraining configuration specifies numerical precision, MoE routing, parallelism, and multi-token prediction settings for Mellum 2.

  • Precision: Mellum 2 uses BF16 base precision with a hybrid FP8 recipe, tensorwise FP8 scaling, recent-parameter gathering, and FP32 gradient reduction.
  • MoE configuration: MoE routing uses global-batch auxiliary loss with coefficient 10^-3, FP32 router precision, disabled token dropping, and fused grouped GEMM, router, and permutation operations.
  • Parallelism: The training setup uses expert parallelism of 8, tensor parallelism of 1, and pipeline parallelism of 1.
  • Multi-token prediction: The model has one additional prediction layer with an MTP loss scaling factor of 0.1.

C Evaluation Notes and Lessons Learned

The evaluation notes identify formatting and termination behaviors that can distort benchmark results, leading the authors to use prompt corrections and bounded reasoning budgets when appropriate.

  • Long-context evaluation: The authors observed near-zero RULER QA scores because the model continued the question instead of answering it, despite normal retrieval and aggregation behavior.
  • Long-context evaluation: They attribute the RULER degradation to prompt formatting rather than capability, and deliberately avoid adding benchmark-style QA prompts to the long-context training mix.
  • Reasoning termination: Small Qwen3 and Qwen3.5 thinking variants sometimes omitted the closing </think> tag and continued reasoning indefinitely.
  • Reasoning termination: Uncapped reasoning is expensive and can yield near-zero benchmark scores when models fill the context window with traces instead of answering.
  • Reasoning budgets: A 32K-token reasoning budget was used for every evaluated thinking model to permit long bounded traces while preventing non-terminating cases from dominating averages.
  • Reasoning budgets: The authors suspect insufficient small-scale on-policy reinforcement learning contributes to unreliable termination, because larger models in the same families terminate more reliably.
Loading 2605.31268v1…