新增: DeepSeek 全部 33 篇官方论文(原文 PDF 走 Git LFS + 中文精读笔记与索引)
This commit is contained in:
ylhan
@@ -0,0 +1,201 @@
|
||||
TITLE: DreamCraft3D: Hierarchical 3D Generation with Bootstrapped Diffusion Prior
|
||||
ARXIV: 2310.16818
|
||||
AUTHORS: Jingxiang Sun , Bo Zhang , Ruizhi Shao , Lizhen Wang , Wen Liu , Zhenda Xie , Yebin Liu
|
||||
DATE: [Submitted on 25 Oct 2023 ( v1 ), last revised 26 Oct 2023 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DreamCraft3D: Hierarchical 3D Generation with Bootstrapped Diffusion Prior
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Related work
|
||||
## 3 Preliminaries
|
||||
## 4 DreamCraft3D
|
||||
## 4.1 Geometry sculpting
|
||||
## 4.2 Texture boosting via bootstrapped score sampling
|
||||
## 5 Experiments
|
||||
## 5.1 Implementation Details
|
||||
## 5.2 Comparisons with the State of the Arts
|
||||
## 5.3 Analysis
|
||||
## 6 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DreamCraft3D: Hierarchical 3D Generation with Bootstrapped Diffusion Prior
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related work
|
||||
3 Preliminaries
|
||||
4 DreamCraft3D
|
||||
4.1 Geometry sculpting
|
||||
4.2 Texture boosting via bootstrapped score sampling
|
||||
5 Experiments
|
||||
5.1 Implementation Details
|
||||
5.2 Comparisons with the State of the Arts
|
||||
5.3 Analysis
|
||||
6 Conclusion
|
||||
References
|
||||
A Appendix
|
||||
A.1 Implementation details
|
||||
A.2 Additional experiments
|
||||
A.3 Limitations
|
||||
A.4 Additional qualitative results
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2310.16818v2 [cs.CV] 26 Oct 2023
|
||||
## DreamCraft3D: Hierarchical 3D Generation with Bootstrapped Diffusion Prior
|
||||
Jingxiang Sun 1 1
|
||||
1
|
||||
Work partially done during the internship at DeepSeek AI.
|
||||
Affiliation: Tsinghua University
|
||||
Bo Zhang 2 2
|
||||
2
|
||||
Corresponding authors.
|
||||
Affiliation: Independent Researcher
|
||||
Ruizhi Shao
|
||||
Affiliation: Tsinghua University
|
||||
Lizhen Wang
|
||||
Affiliation: Tsinghua University
|
||||
Wen Liu
|
||||
Affiliation: DeepSeek AI
|
||||
Zhenda Xie
|
||||
Affiliation: DeepSeek AI
|
||||
Yebin Liu 2 2 footnotemark:
|
||||
2
|
||||
Affiliation: Tsinghua University
|
||||
## Abstract
|
||||
We present DreamCraft3D, a hierarchical 3D content generation method that produces high-fidelity and coherent 3D objects. We tackle the problem by leveraging a 2D reference image to guide the stages of geometry sculpting and texture boosting. A central focus of this work is to address the consistency issue that existing works encounter. To sculpt geometries that render coherently, we perform score distillation sampling via a view-dependent diffusion model. This 3D prior, alongside several training strategies, prioritizes the geometry consistency but compromises the texture fidelity. We further propose Bootstrapped Score Distillation to specifically boost the texture. We train a personalized diffusion model, Dreambooth, on the augmented renderings of the scene, imbuing it with 3D knowledge of the scene being optimized. The score distillation from this 3D-aware diffusion prior provides view-consistent guidance for the scene. Notably, through an alternating optimization of the diffusion prior and 3D scene representation, we achieve mutually reinforcing improvements: the optimized 3D scene aids in training the scene-specific diffusion model, which offers increasingly view-consistent guidance for 3D optimization. The optimization is thus bootstrapped and leads to substantial texture boosting. With tailored 3D priors throughout the hierarchical generation, DreamCraft3D generates coherent 3D objects with photorealistic renderings, advancing the state-of-the-art in 3D content generation. Code available at https://github.com/deepseek-ai/DreamCraft3D .
|
||||
## 1 Introduction
|
||||
The remarkable success of 2D generative modeling ( Saharia et al., 2022 ; Ramesh et al., 2022 ; Rombach et al., 2022 ; Gu et al., 2022 ) has profoundly shaped the way that we create visual content. 3D content creation, which is crucial for applications like games, movies and virtual reality, still presents a significant challenge for deep generative networks. While 3D generative modeling has shown compelling results for certain categories ( Wang et al., 2023a ; Chan et al., 2022 ; Zhang et al., 2023b ) , generating general 3D objects remains formidable due to the lack of extensive 3D data. Recent research effort has sought to leverage the guidance of pretrained text-to-image generative models ( Poole et al., 2022 ; Lin et al., 2023 ; Tang et al., 2023 ) and showcases promising results.
|
||||
The idea of leveraging pretrained text-to-image (T2I) models for 3D generation is initially proposed by DreamFusion ( Poole et al., 2022 ) . A score distillation sampling (SDS) loss is enforced to optimize the 3D model such that its renderings at random viewpoints match the text-conditioned image distribution as interpreted by a powerful T2I diffusion model. DreamFusion inherits the imaginative power of 2D generative models and can yield highly creative 3D assets. To deal with the over-saturation and blurriness issues, recent works adopt stage-wise optimization strategies ( Mildenhall et al., 2021 ) or propose improved 2D distillation loss ( Wang et al., 2023b ) , which leads to an enhancement in photo-realism. However, the majority of current research falls short of synthesizing complex content as achieved by 2D generative models.
|
||||
In addition, these works are often plagued with the “Janus issue”, where 3D renderings that appear plausible individually show semantic and stylistic inconsistencies when examined holistically.
|
||||
Figure 1: By lifting 2D images to 3D, DreamCraft3D achieves 3D generation with rich details and holistic 3D consistency. Please refer to the Appendix and the demo video for more results.
|
||||
In this paper, we propose DreamCraft3D , an approach to produce complex 3D assets while maintaining holistic 3D consistency. Our approach explores the potential of hierarchical generation. We draw inspiration from the manual artistic process: an abstract concept is first solidified into a 2D draft, followed by the sculpting of rough geometry, the refinement of the geometric details and the painting of high-fidelity textures. We adopt a similar approach, breaking down the challenging 3D generation into manageable steps. Starting with a high-quality 2D reference image generated from a text prompt, we lift it into 3D via stages of geometry sculpting and texture boosting. Contrary to prior approaches, our work highlights how careful consideration of each stage can unleash the full potential of hierarchical generation, resulting in superior-quality 3D creation.
|
||||
The geometry sculpting stage aims to produce plausible and consistent 3D geometry from the 2D reference image. On top of using the SDS loss for novel views and photometric loss at the reference view, we introduce multiple strategies to promote geometric consistency. Foremost, we leverage an off-the-shelf viewpoint-conditioned image translation model, Zero-1-to-3 ( Liu et al., 2023b ) , to model the distribution of novel views based on the reference image. Since this view-conditioned diffusion model is trained on diverse 3D data ( Deitke et al., 2023 ) , it provides a rich 3D prior that complements the 2D diffusion prior. Additionally, we find annealing the sampling timestep and progressively enlarging training views are crucial to further improve coherency. During optimization, we transition from implicit surface representation ( Wang et al., 2021 ) to mesh representation ( Shen et al., 2021 ) for coarse-to-fine geometry refinement. Through these techniques, the geometry sculpting stage produces sharp, detailed geometry while effectively suppressing most geometric artifacts.
|
||||
We further propose bootstrapped score distillation to substantially boost the texture. Existing view-conditioned diffusion models trained on limited 3D often struggle to match the fidelity of modern 2D diffusion models. Instead, we finetune the diffusion model according to multi-view renderings of the 3D instance being optimized. This personalized 3D-aware generative prior becomes instrumental in augmenting the 3D texture while ensuring view consistency. Importantly, we find that alternatively optimizing the generative prior and 3D representation leads to mutually reinforcing improvements. The diffusion model benefits from training on improved multi-view renderings, which in turn provides superior guidance for optimizing the 3D texture. In contrast to prior works ( Poole et al., 2022 ; Wang et al., 2023b ) that distill from a fixed target distribution, we learn from a distribution that gradually evolves according to the optimization state. Through this “bootstrapping”, our approach captures increasingly detailed texture while keeping the view consistency.
|
||||
As shown in Figure 1 , our method is capable of producing creative 3D assets with intricate geometric structures and realistic textures rendered coherently in 360°. Compared to optimization-based approaches ( Poole et al., 2022 ; Lin et al., 2023 ) , our method offers substantially improved texture and complexity. Meanwhile, compared to image-to-3D techniques ( Tang et al., 2023 ; Qian et al., 2023 ) , our work excels at producing unprecedented realistic renderings in 360°renderings.
|
||||
These results suggest the strong potential of DreamCraft3D in enabling new creative possibilities in 3D content creation.
|
||||
The full implementation will be made publicly available.
|
||||
## 2 Related work
|
||||
3D generative models have been intensively studied to generate 3D assets without tedious manual creation. Generative adversarial networks (GANs) ( Chan et al., 2021 ; Chan et al., 2022 ; Chan et al., 2021 ; Xie et al., 2021 ; Zeng et al., 2022 ; Skorokhodov et al., 2023 ; Gao et al., 2022 ; Tang et al., 2022 ; Xie et al., 2021 ; Sun et al., 2023 ; Sun et al., 2022 ) have long been the prominent techniques in the field. Auto-regressive models have been explored ( Sanghi et al., 2022 ; Mittal et al., 2022 ; Yan et al., 2022 ; Zhang et al., 2022 ; Yu et al., 2023 ) , which learn the distribution of these 3D shapes conditioned on images or texts. Diffusion models ( Wang et al., 2023a ; Cheng et al., 2023 ; Li et al., 2023 ; Nam et al., 2022 ; Zhang et al., 2023a ; Nichol et al., 2022 ; Jun & Nichol, 2023 ; Bautista et al., 2022 ; Gupta et al., 2023 ) have also shown significant recent success in learning probabilistic mappings from text or images to 3D shape latent. However, these methods require 3D shapes or multi-view data for training, raising challenges when generating in-the-wild 3D assets due to the scarcity of diverse 3D data ( Chang et al., 2015 ; Deitke et al., 2023 ; Wu et al., 2023 ) compared to 2D.
|
||||
3D-aware image generation aims to render images in novel views while
|
||||
offering some level of 3D consistency. These works ( Sargent et al., 2023 ; Skorokhodov et al., 2023 ; Xiang et al., 2023 ) often rely on a pretrained monocular depth prediction model to synthesize view-consistent images. While they achieve photo-realistic renderings for categories of ImageNet, they fall short in producing results in large views. There are a few recent attempts ( Watson et al., 2022 ; Liu et al., 2023b ) that train view-dependent diffusion models on 3D data and demonstrate promising novel view synthesis capability for open domain. However, these are inherently 2D models and cannot ensure perfect view consistency.
|
||||
Lifting 2D to 3D approaches improve a 3D scene representation by seeking guidance using estblished 2D text-image foundation models. Early works ( Jain et al., 2022 ; Lee & Chang, 2022 ; Hong et al., 2022 ) utilize the pretrained CLIP ( Radford et al., 2021 ) model to maximize the similarity between rendered images and text prompt. DreamFusion ( Poole et al., 2022 ) and SJC ( Wang et al., 2022 ) , on the other hand, propose to distill the score of image distribution from a pretrained diffusion model and demonstrate promising results. Recent works have sought to further enhance the texture realism via coarse-to-fine optimization ( Lin et al., 2023 ; Chen et al., 2023 ) , improved distillation loss ( Wang et al., 2023b ) , shape guidance ( Metzer et al., 2023 ) or lifting 2D image to 3D ( Deng et al., 2023 ; Tang et al., 2023 ; Qian et al., 2023 ; Liu et al., 2023a ) .
|
||||
Recently, Raj et al. (2023) proposes to finetune a personalized diffusion model for 3D consistent generation. However, producing globally consistent 3D remains challenging. In this work, we meticulously design 3D priors through the whole hierarchical generation process, achieving unprecedented coherent 3D generation.
|
||||
## 3 Preliminaries
|
||||
DreamFusion ( Poole et al., 2022 ) achieves text-to-3D generation by utilizing a pretrained text-to-image diffusion model ϵ ϕ \bm{\epsilon}_{\phi} as an image prior to optimize the 3D representation parameterized by θ \theta . The image 𝒙 = g ( θ ) \bm{x}=g(\theta) , rendered at random viewpoints by a volumetric renderer, is expected to represent a sample drawn from the text-conditioned image distribution p ( 𝒙 | y ) p(\bm{x}|y) modeled by a pretrained diffusion model. The diffusion model ϕ \phi is trained to predict the sampled noise ϵ ϕ ( 𝒙 t , y , t ) \bm{\epsilon}_{\phi}(\bm{x}_{t};y,t) of the noisy image 𝒙 t \bm{x}_{t} at the noise level t t , conditioned on the text prompt y y . A score distillation sampling (SDS) loss encourages the rendered images to match the distribution modeled by the diffusion model. Specifically, the SDS loss computes the gradient:
|
||||
∇ θ ℒ SDS ( ϕ , g ( θ ) ) = 𝔼 t , ϵ [ ω ( t ) ( ϵ ϕ ( 𝒙 t , y , t ) − ϵ ) ∂ 𝒙 ∂ θ ] , \nabla_{\theta}\mathcal{L}_{\textnormal{SDS}}(\phi,g(\theta))=\mathbb{E}_{t,\bm{\epsilon}}\Big[\omega(t)(\bm{\epsilon}_{\phi}(\bm{x}_{t};y,t)-\bm{\epsilon})\frac{\partial\bm{x}}{\partial\theta}\Big],
|
||||
(1)
|
||||
which is the per-pixel difference between the predicted and the added noise upon the rendered image, where ω ( t ) \omega(t) is the weighting function.
|
||||
One way to improve the generation quality of a conditional diffusion model is to use the classifier-free guidance (CFG) technique to steer the sampling slightly away from the unconditional sampling, i.e. , ϵ ϕ ( 𝒙 t , y , t ) + ϵ ϕ ( 𝒙 t , y , t ) − ϵ ϕ ( 𝒙 t , t , ∅ ) \epsilon_{\phi}(\bm{x}_{t};y,t)+\epsilon_{\phi}(\bm{x}_{t};y,t)-\epsilon_{\phi}(\bm{x}_{t},t,\varnothing) , where ∅ \varnothing represents the “empty” text prompt. Typically, the SDS loss requires a large CFG guidance weight for high-quality text-to-3D generation, yet this will bring side effects like over-saturation and over-smoothing ( Poole et al., 2022 ) .
|
||||
Recently, Wang et al. (2023b) proposed a variational score distillation (VSD) loss that is friendly to standard CFG guidance strength and better resolves unnatural textures. Instead of seeking a single data point, this approach regards the solution corresponding to a text prompt as a random variable. Specifically, VSD optimizes a distribution q μ ( 𝒙 0 | y ) q^{\mu}(\bm{x}_{0}|y) of the possible 3D representations μ ( θ | y ) \mu(\theta|y) corresponding to the text y y , to be closely aligned with the distribution defined by the diffusion timestep t = 0 t=0 , p ( 𝒙 0 | y ) p(\bm{x}_{0}|y) , in terms of KL divergence:
|
||||
ℒ VSD = D KL ( q μ ( 𝒙 0 | y ) | | p ( 𝒙 0 | y ) ) . \mathcal{L}_{\textnormal{VSD}}=D_{\textnormal{KL}}(q^{\mu}(\bm{x}_{0}|y)||p(\bm{x}_{0}|y)).
|
||||
(2)
|
||||
Wang et al. (2023b) further shows that this objective can be optimized by matching the score of noisy real images and that of noisy rendered images at each time t t , so the gradient of ℒ VSD \mathcal{L}_{\textnormal{VSD}} is
|
||||
∇ θ ℒ VSD ( ϕ , g ( θ ) ) = 𝔼 t , ϵ [ ω ( t ) ( ϵ ϕ ( 𝒙 t , y , t ) − ϵ lora ( 𝒙 t , y , t , c ) ) ∂ 𝒙 ∂ θ ] . \nabla_{\theta}\mathcal{L}_{\textnormal{VSD}}(\phi,g(\theta))=\mathbb{E}_{t,\bm{\epsilon}}\Big[\omega(t)(\bm{\epsilon}_{\phi}(\bm{x}_{t};y,t)-\bm{\epsilon}_{\textnormal{lora}}(\bm{x}_{t};y,t,c))\frac{\partial\bm{x}}{\partial\theta}\Big].
|
||||
(3)
|
||||
Here, ϵ lora \bm{\epsilon}_{\textnormal{lora}} estimates the score of the rendered images using a LoRA (Low-rank adaptation) ( Hu et al., 2021 ) model. The obtained variational distribution yields samples with high-fidelity textures. However, this loss is applied for texture enhancement and is helpless to the coarse geometry initially learned by SDS. Moreover, both the SDS and VSD attempt to distill from a fixed target 2D distribution which only assures per-view plausibility rather than a global 3D consistency. Consequently, they suffer from the same appearance and semantic shift issue that hampers the perceived 3D quality.
|
||||
## 4 DreamCraft3D
|
||||
We propose a hierarchical pipeline for 3D content generation as illustrated in Figure 2 . Our method first leverages a state-of-the-art text-to-image generative model to generate a high-quality 2D image from a text prompt. In this way, we can leverage the full power of state-of-the-art 2D diffusion models to depict intricate visual semantics described in the text, retaining the creative freedom as 2D models. We then lift this image to 3D through cascaded stages of geometric sculpting and texture boosting. By decomposing the problem, we can apply specialized techniques at each stage. For geometry, we prioritize multi-view consistency and global 3D structure, allowing for some compromise on detailed textures. With the geometry fixed, we then focus solely on optimizing realistic and coherent texture, for which we jointly learn a 3D-aware diffusion prior that bootstraps the 3D optimization. In the next, we elaborate on key design considerations for the two phases.
|
||||
BSD
|
||||
L
|
||||
BSD
|
||||
SDS 3D-SDS
|
||||
L
|
||||
SDS
|
||||
,
|
||||
L
|
||||
3D-SDS
|
||||
Figure 2: DreamCraft3D leverages a 2D image generated from the text prompt and uses it to guide the stages of geometry sculpting and texture boosting. When sculpting the geometry, the view-conditioned diffusion model provides crucial 3D guidance to ensure geometric consistency. We then dedicately improve the texture quality by conducting a cyclic optimization. We augment the multi-view renderings and use them to finetune a diffusion model, DreamBooth, to offer multi-view consistent gradients to optimize the scene. We term the loss that distills from an evolving diffusion prior as bootstrapped distillation sampling ( ℒ BSD \mathcal{L}_{\textnormal{BSD}} in the figure).
|
||||
## 4.1 Geometry sculpting
|
||||
At this stage, we aim to craft a 3D model such that it matches the appearance of the reference image 𝒙 ^ \bm{\hat{x}} at the same reference view while maintaining plausibility under different viewing angles. To achieve this, we encourage plausible image renderings for each randomly sampled view, recognizable by a pretrained diffusion model. This is achieved using the SDS loss ℒ SDS \mathcal{L}_{\textnormal{SDS}} , as defined in Equation 1 . In order to effectively utilize guidance from the reference image, we penalize the photometric difference between the rendered image and the reference via
|
||||
ℒ rgb = ‖ 𝒎 ^ ⊙ ( 𝒙 ^ − g ( θ , c ^ ) ) ‖ 2 \mathcal{L}_{\text{rgb}}=\left\lVert\bm{\hat{m}}\odot\left(\bm{\hat{x}}-g(\theta;\hat{c})\right)\right\rVert_{2} at the reference view c ^ \hat{c} . The loss is computed only within the foreground region denoted by the mask 𝒎 ^ \bm{\hat{m}} . Meanwhile, we implement the mask loss ℒ mask = ‖ 𝒎 ^ − g m ( θ , c ^ ) ‖ 2 \mathcal{L}_{\textnormal{mask}}=\left\|\bm{\hat{m}}-g_{m}(\theta;\hat{c})\right\|_{2} to encourage scene sparsity, where g m g_{m} renders the silhouette. In addition, akin to ( Deng et al., 2023 ) , we fully exploit the geometry prior inferred from the reference image, and enforce the consistency with the depth and normal map computed for the reference view. The corresponding depth and normal loss are respectively computed as:
|
||||
ℒ depth = − conv ( 𝒅 , 𝒅 ^ ) σ ( 𝒅 ) σ ( 𝒅 ^ ) , ℒ normal = − 𝒏 ⋅ 𝒏 ^ ‖ 𝒏 ‖ 2 ⋅ ‖ 𝒏 ^ ‖ 2 , \mathcal{L}_{\textnormal{depth}}=-\frac{\text{conv}(\bm{d},\hat{\bm{d}})}{\sigma(\bm{d})\sigma(\hat{\bm{d}})},\quad\mathcal{L}_{\textnormal{normal}}=-\frac{\bm{n}\cdot\hat{\bm{n}}}{{\|\bm{n}\|}_{2}\cdot{\|\hat{\bm{n}}\|}_{2}},
|
||||
(4)
|
||||
where conv ( ⋅ ) \text{conv}(\cdot) and σ ( ⋅ ) \sigma(\cdot) represent the covariance and variance operators respectively, and the depth 𝒅 ^ \hat{\bm{d}} and the normal 𝒏 ^ \hat{\bm{n}} at the reference view are computed using the off-the-shelf single-view estimator ( Eftekhar et al., 2021 ) . The depth loss adopts the form of negative Pearson correlation ℒ depth \mathcal{L}_{\textnormal{depth}} to account for the scale mismatch in depth.
|
||||
Despite these, maintaining consistent semantics and appearance across back-views remains a challenge. Thus, we employ additional techniques to produce coherent, detailed geometry.
|
||||
3D-aware diffusion prior. We argue that the 3D optimization with per-view supervision alone is under-constrained. Hence, we utilize a view-conditioned diffusion model, Zero-1-to-3, which is trained on a large scale of 3D assets and offers an improved viewpoint awareness. The Zero-1-to-3 is a fine-tuned 2D diffusion model, which hallucinates the image in a relative camera pose c c given the reference image 𝒙 ^ \bm{\hat{x}} . This 3D-aware model encodes richer 3D knowledge of the visual world and allows us to better extrapolate the views given a reference image. As such, we distill the probability density from this model and compute the gradient of a 3D-aware SDS loss for novel views:
|
||||
∇ θ ℒ 3D-SDS ( ϕ , g ( θ ) ) = 𝔼 t , ϵ [ ω ( t ) ( ϵ ϕ ( 𝒙 t , 𝒙 ^ , c , y , t ) − ϵ ) ∂ 𝒙 ∂ θ ] . \nabla_{\theta}\mathcal{L}_{\textnormal{3D-SDS}}(\phi,g(\theta))=\mathbb{E}_{t,\bm{\epsilon}}[\omega(t)(\bm{\epsilon}_{\phi}(\bm{x}_{t};\bm{\hat{x}},c,y,t)-\bm{\epsilon})\frac{\partial\bm{x}}{\partial\theta}].
|
||||
(5)
|
||||
This loss effectively alleviates 3D consistency issues like Janus problem. However, the finetuning on limited categories of 3D data of inferior rendering quality impairs the diffusion model’s generation capability, so the 3D-aware SDS loss alone is prone to induce deteriorated quality when lifting general images to 3D. Therefore, we employ a hybrid SDS loss, which incorporates both the 2D and 3D diffusion priors simultaneously. Formally, this hybrid SDS loss provides the gradient as:
|
||||
∇ θ ℒ hybrid ( ϕ , g ( θ ) ) = ∇ θ ℒ SDS ( ϕ , g ( θ ) ) + μ ∇ θ ℒ 3D-SDS ( ϕ , g ( θ ) ) , \nabla_{\theta}\mathcal{L}_{\textnormal{hybrid}}(\phi,g(\theta))=\nabla_{\theta}\mathcal{L}_{\textnormal{SDS}}(\phi,g(\theta))+\mu\nabla_{\theta}\mathcal{L}_{\textnormal{3D-SDS}}(\phi,g(\theta)),
|
||||
(6)
|
||||
where we choose μ = 2 \mu=2 to emphasize the weight of the 3D diffusion prior. When computing ℒ SDS \mathcal{L}_{\textnormal{SDS}} , we adopt the DeepFloyd IF base model ( Shonenkov et al., 2023 ) , a diffusion model that operates at 64 × 64 64\times 64 resolution pixel space and better captures coarse geometry.
|
||||
Progressive view training. However, directly deriving the free views in 360°may still result in geometric artifacts, such as extra chair legs, due to the ambiguity inherent in a single reference image. To solve this, we progressively enlarge the training views, gradually propagating the well-established geometry to 360°results.
|
||||
Diffusion timestep annealing.
|
||||
To align with the coarse-to-fine progression of 3D optimization, we adopt a diffusion timestep annealing strategy similar to Huang et al. (2023) . At the start of optimization, we prioritize sampling larger diffusion timestep t t from the range [ 0.7 , 0.85 ] [0.7,0.85] when computing Equation 6 to provide the global structure. As training proceeds, we linearly anneal the t t sampling range to [ 0.2 , 0.5 ] [0.2,0.5] over hundreds of iterations. This annealing strategy allows the model to first establish a plausible global geometry in the early optimization phase before refining the structural details.
|
||||
Detailed structural enhancement. We initially optimize an implicit surface representation with the corresponding volume rendering as in NeuS ( Wang et al., 2021 ) to establish the coarse structure. Then, following Lin et al. (2023) , we use this result to initialize a textured 3D mesh representation using a deformable tetrahedral grid (DMTet) ( Shen et al., 2021 ) to facilitate high-resolution details. Moreover, this representation disentangles the learning of geometry and texture. Hence, at the end of this structural enhancement, we are able to solely refine the texture and better preserve high-frequency details from the reference image.
|
||||
## 4.2 Texture boosting via bootstrapped score sampling
|
||||
The geometry sculpting stage prioritizes the learning of coherent and detailed geometry but leaves the texture blurry. This is due to our reliance on a 2D prior model that operates at a coarse resolution, and the limited sharpness offered by the 3D-aware diffusion model. Additionally, texture issues such as over-smoothing and over-saturation arise from excessively large classifier-free guidance.
|
||||
To augment the texture realism, we use variational score distillation (VSD) loss, as detailed in Equation 3 . We switch to the Stable Diffusion model ( Rombach et al., 2021 ) in this stage which offers high-resolution gradients. To promote realistic rendering, we exclusively optimize the mesh texture with the tetrahedral grid fixed. In this learning stage, we do not leverage the Zero-1-to-3 model as the 3D prior since it adversely impacts the texture quality. Nonetheless, the inconsistent textures may come back, resulting in bizarre 3D outcomes.
|
||||
We observe that the multi-view renderings from the last stage, despite some blurriness, exhibit good 3D consistency. One idea is to adapt a pretrained 2D diffusion model using these rendering results, enabling the model to form a concept about the scene’s surrounding views. In light of this, we finetune the diffusion model with the multi-view image renderings { 𝒙 } \{\bm{x}\} , using DreamBooth ( Ruiz et al., 2023 ) . Specifically, we incorporate the text prompts containing a unique identifier and the subject’s class name (e.g., “A [V] astronaut” in Figure 2 ). During finetuning, the camera parameter of each view is introduced as an additional condition. In practice, we train the DreamBooth with “augmented” image renderings, 𝒙 𝒓 = r t ′ ( 𝒙 ) \bm{x_{r}}=r_{t^{\prime}}(\bm{x}) . We introduce Gaussian noises, in an amount specified by the diffusion timestep t ′ t^{\prime} , to the multi-view renderings, i.e. , 𝒙 t ′ = α t ′ 𝒙 0 + σ t ′ ϵ \bm{x}_{t^{\prime}}=\alpha_{t^{\prime}}\bm{x}_{0}+\sigma_{t^{\prime}}\bm{\epsilon} ( α t ′ , σ t ′ > 0 \alpha_{t^{\prime}},\sigma_{t^{\prime}}>0 are hyperparameters), which are restored using the diffusion model. By choosing a large t ′ t^{\prime} , these augmented images reveal high-frequency details at the cost of the fidelity to the original renderings. The DreamBooth model trained on these augmented renderings can serve as a 3D prior to guide texture refinement.
|
||||
Further, we propose to alternatively optimize the 3D scene to facilitate a bootstrapped optimization (Figure 2 ). Initially, the 3D mesh yields blurry multi-view renderings. We adopt a large diffusion t ′ t^{\prime} to augment their texture quality while introducing some 3D inconsistency. The DreamBooth model trained on these augmented renderings obtains a unified 3D concept of the scene to guide texture refinement. As the 3D mesh reveals finer textures, we reduce the diffusion noises introduced to the image renderings, so the DreamBooth model learns from more consistent renderings and better captures the image distribution faithful to evolving views.
|
||||
In this cyclic process, the 3D mesh and diffusion prior mutually improve in a bootstrapped manner. Formally, we derive the 3D optimization gradient using the following bootstrapped score distillation (BSD) loss:
|
||||
∇ θ ℒ BSD ( ϕ , g ( θ ) ) = 𝔼 t , ϵ , c [ ω ( t ) ( ϵ DreamBooth ( 𝒙 t , y , t , r t ′ ( 𝒙 ) , c ) − ϵ lora ( 𝒙 t , y , t , 𝒙 , c ) ) ∂ 𝒙 ∂ θ ] . \nabla_{\theta}\mathcal{L}_{\textnormal{BSD}}(\phi,g(\theta))=\mathbb{E}_{t,\bm{\epsilon},c}[\omega(t)(\bm{\epsilon}_{\textnormal{DreamBooth}}(\bm{x}_{t};y,t,r_{{t}^{\prime}}(\bm{x}),c)-\bm{\epsilon}_{\textnormal{lora}}(\bm{x}_{t};y,t,\bm{x},c))\frac{\partial\bm{x}}{\partial\theta}].
|
||||
(7)
|
||||
Contrary to prior works ( Poole et al., 2022 ; Wang et al., 2023b ) that distill the score function from a fixed 2D model, our BSD loss learns from an evolving model which becomes increasingly 3D consistent by drawing feedback from the ongoing crafted 3D model. In our experiments, we alternate the optimization twice, which suffices to produce consistent textures with rich details.
|
||||
## 5 Experiments
|
||||
Reference Ours Magic123 Make-It-3D DreamFusion Magic3D ProlificDreamer
|
||||
Figure 3: Qualitative comparison with baselines. Our method generates sharper and more plausible details in both geometry and texture. Note that our method generates rich texture detail at novel views and eliminates multi-face Janus problems.
|
||||
## 5.1 Implementation Details
|
||||
Architectural details. In the geometry sculpting stage, we use Neus and textured 3D mesh representations. We employ Instant NGP ( Müller et al., 2022 ) , optimizing from a 64 to a 384 resolution. For the textured mesh, we use DMTet at a 128 grid and 512 rendering resolution.
|
||||
Optimization.
|
||||
During mesh refinement, we iteratively render a guided normal map and an RGB image, enhancing geometric detail and optimizing our texture prediction network for consistency. Considering the given plausible global geometric structure, our approach eschews the use of a 3D prior during texture optimization. We leverage random sampling of the camera radius and field-of-view (FOV) angles, aligning with Dreamfusion’s methodology. This results in improved texture and geometry details via alternate rendering of normal maps and RGB images.
|
||||
Figure 4: Improved view consistency and texture fidelity along bootstrapping.
|
||||
## 5.2 Comparisons with the State of the Arts
|
||||
Baselines. We conduct a comparative analysis of our technique against five baseline methods. The first three are text-to-3D methods: DreamFusion ( Poole et al., 2022 ) , Magic3D ( Lin et al., 2023 ) and ProlificDreamer ( Wang et al., 2023b ) . We also compare our method against two image-to-3d methods: Make-it-3D ( Tang et al., 2023 ) and Magic123 ( Qian et al., 2023 ) . For DreamFusion, Magic3D, Magic123 and ProlificDreamer, we utilize their implementations in the Threestudio library ( Guo et al., 2023 ) for comparison. For Make-it-3D, we use its official implementation.
|
||||
Table 1: Quantitative comparison against prior 2D-to-3D lifting methods. The metrics are measured on 300 generated samples.
|
||||
CLIP ↑ \uparrow
|
||||
Contextual ↓ \downarrow
|
||||
PSNR ↑ \uparrow
|
||||
LPIPS ↓ \downarrow
|
||||
Make-it-3D
|
||||
0.872
|
||||
1.609
|
||||
18.937
|
||||
0.054
|
||||
Magic123
|
||||
0.843
|
||||
1.628
|
||||
22.838
|
||||
0.053
|
||||
DreamCraft3D
|
||||
0.896
|
||||
1.579
|
||||
31.801
|
||||
0.005
|
||||
Datasets. We establish a test benchmark that includes 300 images, which is a mix of real pictures and those produced by Stable Diffusion ( Rombach et al., 2021 ) and Deep Floyd. Each image in this benchmark comes with an alpha mask for the foreground, a predicted depth map, and a text prompt. For real images, the text prompts are sourced from an image caption model. We intend to make this test benchmark accessible to the public.
|
||||
Quantitative comparison.
|
||||
To generate compelling 3D content that resembles the input image and consistently conveys semantics from various perspectives, we compare our technique with established baselines using a quantitative analysis. Our evaluation employed four metrics: LPIPS ( Zhang et al., 2018 ) and PSNR for fidelity measurement at the reference viewpoint; Contextual Distance ( Mechrez et al., 2018 ) for pixel-level congruence assessment; and CLIP score ( Radford et al., 2021 ) to estimate semantic coherence. Results depicted in Table 1 indicate that our approach significantly surpasses the baselines in maintaining both texture consistency and fidelity.
|
||||
Figure 5: User study.
|
||||
User study. To substantiate the robustness and quality of our proposed model, we executed a user study employing 15 distinct pairs of prompts and images. Each participant was provided with four free-view rendering video alongside their corresponding textual input, and asked to choose their top preferred 3D model. The study gathered 480 responses from a total of 32 participants, the analysis of which is depicted in Figure 5 . On an average basis, our model was favored by 92% of users over alternative models, outperforming the baselines by a large margin. This result provides compelling evidence of the resilience and superior quality inherent to our proposed method.
|
||||
Qualitative comparison. Figure 3 compares our method with the baselines. All the text-to-3D methods suffer from multi-view consistency issues. While ProlificDreamer offers realistic textures, it fails to form a plausible 3D object. Image-to-3D methods like Make-it-3D create quality frontal views but struggle with geometry. Magic123, enhanced by Zero1-to-3, fares better in geometry regularization, but both generate overly smoothed textures and geometric details. In contrast, our Bootstrapped Score Distillation method improves imagination diversity while maintaining semantic consistency.
|
||||
(a) w/o Zero1-to-3
|
||||
(b) SDS
|
||||
(c) VSD
|
||||
(d) 1-round BSD
|
||||
(e) 2-round BSD
|
||||
Figure 6: Ablation study of the effectiveness of 3D prior and our proposed BSD (Bootstrapped Score Distillation). (a) geometry sculpting stage without 3D prior. (b) texture optimization with SDS loss. (c) VSD loss produces richer texture detail while suffering from texture inconsistency. (d) BSD improves the texture consistency with one round DreamBooth. (e) Two-round BSD adds more details to the generated result.
|
||||
Figure 7: Continual improvement of geometry and texture quality through multiple stages.
|
||||
## 5.3 Analysis
|
||||
The effect of 3D prior. In our paper we claim that the guidance offered by a 3D prior enhances the generation of globally plausible geometry. To ascertain its impact, an ablation study is conducted, where the 3D prior is deactivated. Figure 6 demonstrates that, in the absence of the 3D prior, the resultant character tends to exhibit the multifaceted Janus issue and suffers from irregular geometry. This observation underlines the significance of a viewpoint-aware 3D prior in regulating a globally consistent shape.
|
||||
The effect of BSD. Figure 6 also presents an ablation study encompassing three texture optimization techniques: (1) BSD, (2) VSD, and (3) Score Distillation Sampling (SDS) with the traditional stable diffusion. The application of SDS has been observed to generate novel-view textures that are excessively smooth and over-saturated. In contrast, while VSD using standard stable diffusion can produce realistic textures, it yields a notably high inconsistency. In contrast, our proposed approach successfully generates textures that strike a balance between realism and consistency.
|
||||
Visualization of multiple stages. Figure 7 provides the visualization of the intermediate rendering results for each stage in our hierarchical pipeline. In the geometry sculpting stage, we convert Neus to DMTet to improve high-resolution geometry details. However, the improvement in texture is negligible. On the contrary, in the texture stage, we significantly improve the texture quality with our proposed BSD.
|
||||
DreamBooth times. Figure 4 illustrates multi-view datasets for DreamBooth. The initial stage involves the introduction of substantial noise into each image to amplify detail richness, leading to inconsistent denoised images. However, as the textured mesh undergoes optimization, the produced renderings evolve towards increased consistency and photorealism, thereby enhancing the quality of the input dataset tailored for DreamBooth.
|
||||
## 6 Conclusion
|
||||
We have presented DreamCraft3D, an innovative approach that advances the field of complex 3D asset generation. This work introduces a meticulous geometry sculpting phase for producing plausible and coherent 3D geometries and a novel Bootstrapped Score Distillation strategy. The latter, by distilling from an optimizing 3D-aware diffusion prior and adapting to multi-view renderings of the instance being optimized, significantly improves texture quality and consistency. DreamCraft3D produces high-fidelity 3D assets with compelling texture details and multi-view consistency. We believe this work represents an important step towards democratizing 3D content creation and shows great promise in future applications.
|
||||
|
||||
=== BODY (tail) ===
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,889 @@
|
||||
TITLE: DeepSeek-VL: Towards Real-World Vision-Language Understanding
|
||||
ARXIV: 2403.05525
|
||||
AUTHORS: Haoyu Lu , Wen Liu , Bo Zhang , Bingxuan Wang , Kai Dong , Bo Liu , Jingxiang Sun , Tongzheng Ren , Zhuoshu Li , Hao Yang , Yaofeng Sun , Chengqi Deng , Hanwei Xu , Zhenda Xie , Chong Ruan
|
||||
DATE: [Submitted on 8 Mar 2024 ( v1 ), last revised 11 Mar 2024 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-VL: Towards Real-World Vision-Language Understanding
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Data Construction
|
||||
## 2.1 Vision-Language pretraining Data
|
||||
## 2.2 Supervised Fine-tuning Data
|
||||
## 3 Approach
|
||||
## 3.1 Architecture
|
||||
## 3.2 Training Pipelines
|
||||
## 3.2.1 Stage 1: Training Vision-Language Adaptor
|
||||
## 3.2.2 Stage 2: Joint Vision-Language pretraining
|
||||
## 3.2.3 Stage 3: Supervised Fine-tuning
|
||||
## 3.3 Hyperparameters and Infrastructures
|
||||
## 4 Evaluation
|
||||
## 4.1 Public Multimodal Benchmarks Evaluation
|
||||
## 4.2 Public Language Benchmarks Evaluation
|
||||
## 4.3 Human Evaluation
|
||||
## 4.4 Ablation Study
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-VL: Towards Real-World Vision-Language Understanding
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Data Construction
|
||||
2.1 Vision-Language pretraining Data
|
||||
2.2 Supervised Fine-tuning Data
|
||||
3 Approach
|
||||
3.1 Architecture
|
||||
3.2 Training Pipelines
|
||||
3.2.1 Stage 1: Training Vision-Language Adaptor
|
||||
3.2.2 Stage 2: Joint Vision-Language pretraining
|
||||
3.2.3 Stage 3: Supervised Fine-tuning
|
||||
3.3 Hyperparameters and Infrastructures
|
||||
4 Evaluation
|
||||
4.1 Public Multimodal Benchmarks Evaluation
|
||||
4.2 Public Language Benchmarks Evaluation
|
||||
4.3 Human Evaluation
|
||||
4.4 Ablation Study
|
||||
5 Conclusion, Limitation, and Future Work
|
||||
References
|
||||
A Appendix
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2403.05525v2 [cs.AI] 11 Mar 2024
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-VL: Towards Real-World Vision-Language Understanding
|
||||
Haoyu Lu* 1† ,
|
||||
Wen Liu* 1 ,
|
||||
Bo Zhang* 1‡ ,
|
||||
Bingxuan Wang 1† ,
|
||||
Kai Dong 1 ,
|
||||
Bo Liu 1† ,
|
||||
Jingxiang Sun 1† ,
|
||||
Tongzheng Ren 1† ,
|
||||
Zhuoshu Li 1 ,
|
||||
Hao Yang 1† ,
|
||||
Yaofeng Sun 1 ,
|
||||
Chengqi Deng 1 ,
|
||||
Hanwei Xu 1 ,
|
||||
Zhenda Xie 1 ,
|
||||
Chong Ruan 1
|
||||
1 DeepSeek-AI
|
||||
{neal, liuwen, bo}@deepseek.com
|
||||
https://github.com/deepseek-ai/DeepSeek-VL
|
||||
Corresponding author: ∗ Equal contribution.
|
||||
† Work done during the internship at DeepSeek-AI.
|
||||
‡ Project lead.
|
||||
## Abstract
|
||||
We present DeepSeek-VL, an open-source Vision-Language (VL) Model designed for real-world vision and language understanding applications. Our approach is structured around three key dimensions:
|
||||
• Data Construction : We strive to ensure our data is diverse, scalable and extensively covers real-world scenarios including web screenshots, PDFs, OCR, charts, and knowledge-based content (expert knowledge, textbooks), aiming for a comprehensive representation of practical contexts. Further, we create a use case taxonomy from real user scenarios and construct an instruction-tuning dataset accordingly. The fine-tuning with this dataset substantially improves the model’s user experience in practical applications.
|
||||
• Model Architecture : Considering efficiency and the demands of most real-world scenarios, DeepSeek-VL incorporates a hybrid vision encoder that efficiently processes high-resolution images (1024 x 1024) within a fixed token budget, while maintaining a relatively low computational overhead. This design choice ensures the model’s ability to capture critical semantic and detailed information across various visual tasks.
|
||||
• Training Strategy : We posit that a proficient Vision-Language Model should, foremost, possess strong language abilities. To ensure the preservation of LLM capabilities during pretraining, we investigate an effective VL pretraining strategy by integrating LLM training from the beginning and carefully managing the competitive dynamics observed between vision and language modalities. Starting with a focus on text, we gradually adjust the ratio to facilitate a balanced integration of both modalities.
|
||||
The DeepSeek-VL family (both 1.3B and 7B models) showcases superior user experiences as a vision-language chatbot in real-world applications, achieving state-of-the-art or competitive performance across a wide range of visual-language benchmarks at the same model size while maintaining robust performance on language-centric benchmarks. We have made both 1.3B and 7B models publicly accessible to foster innovations based on this foundation model.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Data Construction
|
||||
2.1 Vision-Language pretraining Data
|
||||
2.2 Supervised Fine-tuning Data
|
||||
3 Approach
|
||||
3.1 Architecture
|
||||
3.2 Training Pipelines
|
||||
3.2.1 Stage 1: Training Vision-Language Adaptor
|
||||
3.2.2 Stage 2: Joint Vision-Language pretraining
|
||||
3.2.3 Stage 3: Supervised Fine-tuning
|
||||
3.3 Hyperparameters and Infrastructures
|
||||
4 Evaluation
|
||||
4.1 Public Multimodal Benchmarks Evaluation
|
||||
4.2 Public Language Benchmarks Evaluation
|
||||
4.3 Human Evaluation
|
||||
4.4 Ablation Study
|
||||
5 Conclusion, Limitation, and Future Work
|
||||
References
|
||||
A Appendix
|
||||
## 1 Introduction
|
||||
The remarkable success of large language models (LLMs) ( OpenAI, 2022 ; OpenAI, 2023a ; Anthropic, 2023 ; Google, 2023 ) has fueled the demand for a versatile interface that can handle multiple modalities beyond language.
|
||||
In response to this growing demand, we have seen an emergence of Large Multimodal Models (LMMs) like GPT-4V ( OpenAI, 2023b ) and Gemini ( Team et al., 2023 ) , which serve as versatile assistants capable of comprehending and acting upon instructions that span vision and language. These models exhibit considerable promise in executing complex, diverse real-world tasks, enabling more natural and human-like interactions.
|
||||
Figure 1: DeepSeek-VL possesses general multimodal understanding capabilities, capable of processing logical diagrams, web pages, formula recognition, scientific literature, natural images, and embodied intelligence in complex scenarios.
|
||||
Recently, there has been a surge of open-source large multimodal models aimed at narrowing the gap with proprietary counterparts. Substantial strides have been made, especially in benchmark performance, yet a significant divide persists between the majority of open-source models and state-of-the-art closed-source models ( OpenAI, 2023b ; Bavishi et al., 2023 ; Team et al., 2023 ; Bai et al., 2023 ) when it comes to real-world performance and user experience. It remains challenging for the open-source community to develop models with robust general multimodal capabilities for real-world applications.
|
||||
The performance gap between the most open-source models and the proprietary models is largely pronounced in real-world scenarios, primarily due to the following reasons:
|
||||
•
|
||||
Many open-source solutions allocate a significant proportion of computational resources to the instruction tuning phase. However, the experience of training powerful language models underscores the importance of extensive pretraining in the development of general intelligence. To imbue multimodal models with rich world knowledge, there should be an emphasis on comprehensive pretraining that leverages a broad spectrum of vision-language data.
|
||||
•
|
||||
A common practice is to amalgamate various academic datasets during instruction tuning. While such an approach may yield good benchmark results, it often falls short in providing an authentic real-world usage experience.
|
||||
•
|
||||
In terms of model architecture, prior works mostly adapt a vision transformer, typically text-aligned, to a pre-trained language model. However, most of these models operate on a relatively low resolution, e.g., 336 × \times 336 or 448 × \times 448. The intricacies of complex real-world scenarios, such as optical character recognition or tiny object discernment, demand high-resolution processing capability.
|
||||
•
|
||||
While some models ( Sun et al., 2023 ; Wang et al., 2023b ; 01-ai, 2024 ; Lin et al., 2023a ) have begun to exploit pretraining, they often overlook the preservation of language skills. Often, there is a degradation of language capability after prolonged multimodal training. Since we aim for a generalist that possesses strong capabilities in both modalities, there should be a training strategy that well preserves the language capability when developing the new modality ability.
|
||||
In light of these, we present DeepSeek-VL, an open-source large multimodal model, which is built upon the DeepSeek language model series. We develop the model in the pursuit of adept performance in real-world scenarios, which involves extensive pretraining, careful data curation based on a use case taxonomy, model architecture design for high-resolution processing, and a training strategy that balances the multi-modalities.
|
||||
On top of these, we develop a training methodology that steers the model scaling, from 1B to 7B. These comprehensive explorations bring a significant performance advantage in practical settings, compared to other large multimodal models (LMMs) of similar size.
|
||||
DeepSeek-VL’s pretraining dataset is compiled
|
||||
from a variety of sources, including but not limited to Common Crawl, Web Code, E-books, Educational Materials, and arXiv Articles. This collection thoroughly encompasses real-world scenarios such as web screenshots, PDFs, OCR, charts, and knowledge-based content (expertise, textbooks), aiming for a broad and practical representation while remaining scalable.
|
||||
While our pretraining data encompasses a wide array of world knowledge, we meticulously curate our instruction-tuning dataset to reflect real-world usage scenarios. To achieve this, we manually gather authentic test cases for GPT-4V and Gemini from the Internet. These cases have been systematically organized into a comprehensive taxonomy. We use this structured taxonomy to choose prompts for each test image, ensuring a practical and relevant instruction tuning dataset. This taxonomy is also used to create an evaluation dataset that effectively assesses real-world performance.
|
||||
The visual module is designed to optimize the utilization of high-resolution visual inputs while remaining within a fixed token budget to manage inference costs effectively. As such, we employ a hybrid vision encoder, which combines a text-aligned encoder for coarse semantic extraction at 384 × 384 384\times 384 resolution with a high-resolution encoder that captures detailed visual information at 1024 × 1024 1024\times 1024 resolution. By fusing these two encoders, our hybrid approach efficiently condenses a 1024×1024 resolution image (which suffices in most use cases) into 576 tokens. This token count strikes a balance between rich visual representation and token economy, making it feasible for both text-image interleaving and multi-turn inference scenarios.
|
||||
During the pretraining of multimodal models, a common challenge encountered is the potential degradation of language capabilities when the training process is overly reliant on vision-language data. Our research reveals that maintaining a significant proportion of language data—specifically, at least 70%—is essential to preserve the integrity of language knowledge within the model. This balance is critical for achieving a robust multimodal capability that does not compromise language performance. Moreover, we introduce a novel “modality warm-up” strategy. This approach carefully adjusts the ratio of modalities during training, gradually incorporating more vision-language data.
|
||||
The careful tuning of the modality ratio along with the warm-up strategy results in a balanced performance of both modalities.
|
||||
When iterating on our model, We conduct experiments on a small scale before scaling to a larger model size. However, a smaller model, e.g., 1B model, cannot demonstrate reasonable performance on benchmarks ( Schaeffer et al., 2024 ) and faithfully reflect the model’s performance. We adopt two approaches to address this. First, we modify the evaluation protocol from multi-choice to compare the perplexity of options. Also, to prevent the instruction following ability from becoming the bottleneck, we mix a small proportion of instruction tuning data during the pretraining phase. In this way, we can achieve reasonable performance using the 1B model and more accurately measure the impact of each iteration during the experiment.
|
||||
Through extensive evaluations of general vision and language benchmarks, the DeepSeek-VL family showcases superior user experiences in real-world applications and achieves state-of-the-art or competitive performance across a wide range of visual-language benchmarks at the same model size, while maintaining robust language-centric performance. To foster innovation and enable a wide range of applications, we have made two versions of our ours, 1.3B and 7B, publicly accessible, in the hope of facilitating the needs of varying computational capabilities.
|
||||
## 2 Data Construction
|
||||
A diverse and large dataset is the most important ingredient of visual language model training. Our dataset can be divided into two parts: Vision-Language pretraining Data and Vision-Language Supervised Fine-Tuning Data. VL pretraining Data is composed of visual-text data from various sources, aimed at enhancing the model’s fundamental cross-modal understanding capabilities; while VL Supervised Fine-Tuning Data has a relatively smaller size and aims to teach the model to complete specific downstream tasks. By design, VL pretraining Data is used to warm up the vision-language adaptor in training stage 1 and jointly pretrain the vision-language model in stage 2, and VL Supervised Fine-Tuning Data is exploited in training stage 3, i.e., vision language supervised fine-tuning.
|
||||
## 2.1 Vision-Language pretraining Data
|
||||
Table 1: Summary of datasets used in the joint vision and language pretraining stage.
|
||||
Category
|
||||
Dataset
|
||||
Ratio
|
||||
Interleaved image-text
|
||||
MMC4 ( Zhu et al., 2024 )
|
||||
13.1%
|
||||
Wikipedia EN& CN ( Foundation, )
|
||||
Wikihow ( Yang et al., 2021 )
|
||||
in-house PDF and Epub textbooks
|
||||
Image caption
|
||||
Capsfusion ( Yu et al., 2023a )
|
||||
11.1%
|
||||
TaiSu ( Liu et al., 2022b )
|
||||
Detailed Caption ( echo840, 2024 )
|
||||
Table and chart
|
||||
Chart2text ( Kantharaj et al., 2022 )
|
||||
2.1%
|
||||
Geo170K ( Gao et al., 2023 )
|
||||
Ureader ( Ye et al., 2023 )
|
||||
Unichart ( Masry et al., 2023 )
|
||||
M-paper ( Hu et al., 2023 )
|
||||
ScienceQA ( Lu et al., 2022b )
|
||||
ScreenQA ( Hsiao et al., 2022 )
|
||||
SciGraphQA-295K ( Li and Tajbakhsh, 2023 )
|
||||
Paper2figure100k ( Rodriguez et al., 2023 )
|
||||
Widget Captioning ( Li et al., 2020 )
|
||||
Screen2words ( Wang et al., 2021 )
|
||||
Refexp ( Mao et al., 2016 )
|
||||
Web Code
|
||||
Websight ( HuggingFaceM4, 2024 )
|
||||
0.4%
|
||||
python plots scraped from GitHub notebook
|
||||
Scene text OCR
|
||||
ArT ( Chng et al., 2019 )
|
||||
1.2%
|
||||
MLT-17 ( Nayef et al., 2017 )
|
||||
LSVT ( Sun et al., 2019 )
|
||||
UberText ( Zhang et al., 2017 )
|
||||
Coco-text ( Veit et al., 2016 )
|
||||
RCTW-17 ( Shi et al., 2017 )
|
||||
ReCTS ( Zhang et al., 2019 )
|
||||
TextOCR ( Singh et al., 2021 )
|
||||
OpenVINO ( Krylov et al., 2021 )
|
||||
HierText ( Long et al., 2022 )
|
||||
Document OCR
|
||||
arXiv rendered markdown ( Blecher et al., 2023 )
|
||||
2.1%
|
||||
Text-only corpus
|
||||
DeepSeek-LLM 2T text copus ( DeepSeek-AI, 2024 )
|
||||
70.0%
|
||||
The pretraining dataset utilized in our study encompasses a diverse range of publicly accessible sources, in addition to a selection of proprietary data. We provide a comprehensive overview of the data sources employed during the joint vision and language pretraining stage in Table 1 . Such a dataset can facilitate LLM’s comprehension of the entities portrayed in the images.
|
||||
Furthermore, we present a detailed breakdown of the complete dataset, which is organized into the following categories:
|
||||
Interleaved image-text data enable the models to have a better capability for in-context learning of multi-modality inputs, and we utilize three public datasets MMC4 ( Zhu et al., 2024 ) , Wiki ( Burns et al., 2023 ) , Wikihow ( Yang et al., 2021 ) and Epub textbooks.
|
||||
Image caption data come from three high-quality image-text paired datasets: Capsfusion ( Yu et al., 2023a ) , TaiSu ( Liu et al., 2022b ) and Detailed Caption ( echo840, 2024 ) .
|
||||
Table and chart data enable the models to learn the capability for general table and chart image understanding. It encompasses a diverse range of public data sources, including Chart2text ( Kantharaj et al., 2022 ) , Geo170K ( Gao et al., 2023 ) ,
|
||||
Unichart ( Masry et al., 2023 ) , Ureader ( Ye et al., 2023 ) ,
|
||||
M-paper ( Hu et al., 2023 ) , ScienceQA ( Lu et al., 2022b ) , ScreenQA ( Hsiao et al., 2022 ) , SciGraphQA-295K ( Li and Tajbakhsh, 2023 ) , Paper2figure100k ( Rodriguez et al., 2023 ) , Widget Captioning ( Li et al., 2020 ) , Screen2words ( Wang et al., 2021 ) , and Refexp ( Mao et al., 2016 ) .
|
||||
Web Code data empowers models with the capability to reconstruct code from graphical interfaces or visual plots. Leveraging Websight ( HuggingFaceM4, 2024 ) for UI Inverse Rendering, we adopted a strategy akin to that used in MATCHA ( Liu et al., 2022a ) for visual plots inverse rendering. This involved the processing of approximately 1.46 million Jupyter notebooks from the Stack dataset ( Kocetkov et al., 2023 ) . By extracting these notebooks and collating all diagrams along with their corresponding preceding code segments, we succeeded in curating a collection featuring 2 million pairs of images and codes. For better data quality, we filter 1.1 million instances, each comprising a singular image coupled with a minimum of 5 lines of code, to constitute our primary training dataset.
|
||||
Document Optical Character Recognition (OCR) data facilitates the recognition of optical characters at the document level, even in challenging real-world scenarios. To the best of our knowledge, there is currently no publicly available large-scale dataset encompassing both English and Chinese documents. Despite the existence of the publicly accessible small-scale dataset Latex-OCR ( Blecher, 2024 ) , we additionally constructed a comprehensive English and Chinese document OCR dataset. It is comprised of two parts: 1): arXiv Articles: We collected source code and compiled PDFs from 1.4 million arXiv articles. Utilizing pre-processing tools from Nougat ( Blecher et al., 2023 ) , we rendered these articles into paired images and texts; 2):
|
||||
E-books and Educational Materials: We cleaned 860K English and 180K Chinese e-books from Anna’s Archive ( Anna’s Archive, 2024 ) alongside millions of K-12 education exam questions. Subsequently, we employed HTML rendering tools ( Kulkarni and Truelsen, ) to convert these HTML files with different templates into paired image and text formats.
|
||||
Scene text OCR data augment the capability of the model to recognize and extract text from images in which the text is integrated into the environment. The dataset is composed of multiple public datasets, including ArT ( Chng et al., 2019 ) , MLT-17 ( Nayef et al., 2017 ) , LSVT ( Sun et al., 2019 ) , UberText ( Zhang et al., 2017 ) , Coco-text ( Veit et al., 2016 ) , RCTW-17 ( Shi et al., 2017 ) , ReCTS ( Zhang et al., 2019 ) , TextOCR ( Singh et al., 2021 ) , OpenVINO ( Krylov et al., 2021 ) and HierText ( Long et al., 2022 ) .
|
||||
Text-only corpus serves to maintain proficiency in language-centric tasks. In this study, we employ the same text corpus with DeepSeek-LLM ( DeepSeek-AI, 2024 ) .
|
||||
Table 2: Summary of data used in our joint vision and language supervised fine-tuning stage.
|
||||
Class
|
||||
Dataset
|
||||
Ratio
|
||||
In-house Data
|
||||
SFT data based on taxonomy (Figure 3 )
|
||||
10.5%
|
||||
General Multi-modality
|
||||
ShareGPT4V ( Chen et al., 2023 )
|
||||
35.5%
|
||||
LAION-GPTV ( LAION, 2023 )
|
||||
LVIS-Instruct4V ( Wang et al., 2023a )
|
||||
textOCR-GPT4V ( Carter, 2024 )
|
||||
LLaVA1.6-GPT4V ( Liu et al., 2024a )
|
||||
IconQA ( Lu et al., 2021 )
|
||||
Table and chart
|
||||
Ureader ( Ye et al., 2023 )
|
||||
4.1%
|
||||
Geo170K ( Gao et al., 2023 )
|
||||
ScienceQA ( Lu et al., 2022b )
|
||||
Web Code
|
||||
Screen-to-code ( Abi, 2024 )
|
||||
2.0%
|
||||
ScreenQA ( Hsiao et al., 2022 )
|
||||
Text-only SFT
|
||||
DeepSeek-LLM ( DeepSeek-AI, 2024 )
|
||||
47.9%
|
||||
Main Category
|
||||
Description
|
||||
Secondary Category
|
||||
Tertiary Category
|
||||
Recognition
|
||||
This part of the use cases mainly examines the understanding and description ability of large models for image content, which does not require high knowledge reserve and reasoning ability of the model, and some tasks can be completed using traditional machine learning models.
|
||||
Global Description
|
||||
Theme Description, Event/Behavior Description, Location/Scene Description, Emotion/Mood Description, Style Recognition, Food Recognition, Others
|
||||
Local Description
|
||||
Pointing Description, Position Description, Person Recognition, Object Attribute Description, Logo Recognition, Counting, Currency Recognition
|
||||
OCR and Transcription
|
||||
Printed Text Transcription, Handwritten Text Transcription, Specified Format Transcription, Specified Language Transcription
|
||||
Conversion
|
||||
This type of use case requires the model to be able to describe and recognize image content, and use specific knowledge (e.g., code knowledge, prompt engineering knowledge) to convert image content into another form.
|
||||
Image to Code
|
||||
UI to Code, Chart to Code, Photo to SVG/p64 Encoding, Formula to Code, Flowchart to Code
|
||||
Image to Text
|
||||
Image to Prompt, Text Summary, Image-based Creation, Text Interpretation
|
||||
Analysis
|
||||
This type of use case requires the model to use specific knowledge and logical ability to make reasonable analysis and understanding based on image content, and describe the image according to instructions.
|
||||
Data Chart Analysis
|
||||
Graph Interpretation, Table Interpretation
|
||||
Professional Chart Analysis
|
||||
Circuit Diagram, Flowchart, Map, Music Score, Financial Chart, Floor Plan, Others
|
||||
Professional Image Analysis
|
||||
Sensor Image, Biological and Medical Image, Voiceprint Image, Point Cloud Image
|
||||
Encyclopedia Knowledge Analysis
|
||||
Art and Culture Knowledge, Natural Environment Knowledge, Food/Clothing/Housing/Transportation Related Knowledge, Entertainment Related Knowledge, Historical Knowledge
|
||||
Commonsense Reasoning
|
||||
This type of use case mainly tests the model’s understanding and mastery of common sense in life, which requires reasoning based on the interpretation and analysis of image content combined with common sense.
|
||||
Relationship Reasoning
|
||||
Interpersonal Relationship, Spatial Relationship, Size Relationship, Species Relationship
|
||||
Function Reasoning
|
||||
Hardware Function Reasoning, Software Function Reasoning
|
||||
Environment Reasoning
|
||||
Environment State Analysis, Environment-based Behavior Reasoning, Embodied Intelligence
|
||||
Anomaly Reasoning
|
||||
Identifying Anomalies in Images, Defect Detection, Accident Judgment
|
||||
Humor Reasoning
|
||||
-
|
||||
Other Commonsense Reasoning
|
||||
State Reasoning, Cause Reasoning, Attribute Comparison, Optical Illusion, Fun Games, Intention Interpretation, Behavior Prediction
|
||||
Logical Reasoning
|
||||
This type of use case requires the model to combine the understanding of images, comprehensively use domain knowledge and logical reasoning ability to complete corresponding tasks.
|
||||
Mathematical Reasoning
|
||||
Algebra and Operation, Plane Geometry, Solid Geometry
|
||||
Other Logical Reasoning
|
||||
Physics, Chemistry, Biology, Code, IQ Questions
|
||||
Evaluation
|
||||
This type of use case requires the model to evaluate the image content according to specific criteria.
|
||||
-
|
||||
Reality Evaluation, Similarity Evaluation, Aesthetic Evaluation, Open-ended Evaluation, Improvement Suggestions
|
||||
Multi-graph
|
||||
This type of use case examines the model’s ability to analyze and understand multiple images.
|
||||
Temporal Sequence Understanding
|
||||
Event Prediction, Image Sequencing, Behavior Analysis
|
||||
Multi-graph Comparison
|
||||
Attribute Comparison, Image-Text Matching, Finding Associations, Spotting Differences, Image Discrimination
|
||||
Safety
|
||||
This type of use case examines the model’s performance in terms of safety.
|
||||
-
|
||||
Suggestive Questioning, Counterfactual Questioning, Prompt Injection
|
||||
Table 3: Our taxonomy for the in-house SFT data. The categories covered by our high-quality in-house multi-modality SFT data are comprehensively represented in this taxonomy.
|
||||
## 2.2 Supervised Fine-tuning Data
|
||||
The supervised fine-tuning datasets utilized in our study encompass a diverse range of multi-modality and language data sources, including well-known open-source shared gpt4v datasets such as ShareGPT4V ( Chen et al., 2023 ) , LAION-GPTV ( LAION, 2023 ) , LVIS-Instruct4V ( Wang et al., 2023a ) , textOCR-GPT4V ( Carter, 2024 ) , LLaVA1.6-GPT4V ( Liu et al., 2024a ) and IconQA ( Lu et al., 2021 ) . Additionally, we incorporate partial table and chart data extracted from pretraining datasets such as Ureader ( Ye et al., 2023 ) , ScreenQA ( Hsiao et al., 2022 ) , Geo170K ( Gao et al., 2023 ) , and ScienceQA ( Lu et al., 2022b ) . Moreover, we integrate the UI Code dataset obtained from Screen-to-code ( Abi, 2024 ) tasks. To enhance the quality of our multi-modality SFT data, we have also curated a portion of high-quality in-house multi-modality SFT data, some of which are in the Chinese language. Our in-house instruction-tuning dataset is meticulously designed to reflect real-world usage scenarios and cover a wide range of tasks. We start by collecting a diverse set of authentic test cases for GPT-4V and Gemini from various online sources. These test cases are then carefully analyzed and organized into a comprehensive taxonomy, which encompasses multiple categories, such as recognition, conversion, analysis, reasoning, evaluation, and safety, as detailed in Table 3 . This structured taxonomy serves as a guideline for selecting representative prompts for each test image, ensuring that our instruction-tuning dataset is both practical and relevant to real-world applications. Moreover, this taxonomy is also employed to construct a balanced and comprehensive evaluation dataset, which allows us to effectively assess the model’s performance across different tasks and categories. By following this systematic approach, we ensure that the categories covered by our in-house multi-modality SFT data are well-aligned with the taxonomy and representative of real-world usage scenarios. Furthermore, we include the text-only SFT data employed in DeepSeek-LLM ( DeepSeek-AI, 2024 ) as part of our joint vision and language SFT data.
|
||||
## 3 Approach
|
||||
## 3.1 Architecture
|
||||
Our system contains three modules: a hybrid vision encoder, a vision adaptor, and a language model. We introduce each part in this section.
|
||||
Hybrid Vision Encoder.
|
||||
We employ SigLIP as the vision encoder to extract high-level semantic feature representations from visual inputs. However, we observe that a single SigLIP encoder struggles to address all real-world questions comprehensively. Vision encoders in the CLIP family, including SigLIP, are primarily designed for semantic visual representations but are challenged by ambiguous encoding, resulting in visually distinct images being encoded as similar due to what is referred to as "CLIP-blind pairs" Tong et al. (2024) . Meanwhile, the CLIP family of models is limited by its relatively low-resolution inputs (e.g., 224 x 224, 336 x 336, 384 x 384, 512 x 512), which hinders their ability to handle tasks requiring more detailed low-level features like dense OCR and visual grounding task.
|
||||
To address these limitations,
|
||||
recent researches ( Wei et al., 2023 ; Tong et al., 2024 ; Lin et al., 2023b ) have advocated for the integration of additional vision-only self-supervised encoders, to enhance the visual grounding capabilities of multi-modality models. Building upon previous motivations,
|
||||
we additionally utilize a vision-only encoder based on the SAM-B ( Kirillov et al., 2023 ) , a pre-trained ViTDet ( Li et al., 2022 ) image encoder to process low-level features, which accepts high-resolution 1024 x 1024 image inputs. In addition to the SAM-B encoder, we retain the SigLIP-L vision encoder with low-resolution 384 x 384 image inputs. Consequently, our hybrid vision encoder combines the SAM-B and SigLIP-L encoders, efficiently encoding high-resolution 1024 x 1024 images while preserving both semantic and detailed information. Specifically, a high-resolution SAM-B vision encoder first resizes the image into 1024 x 1024 and results in a 64 x 64 x 256 feature map.
|
||||
In the case of a high-resolution feature map of size, 64 x 64 x 256 generated by SAM-B, the VL Adaptor initially interpolates it into a size of 96 x 96 x 256. Subsequently, it employs two convolutional layers with a stride of 2, producing a feature map of 24 x 24 x 1024, and reshapes it to 576 x 1024. Alongside this, the low-resolution feature map of size 576 x 1024 generated by SigLIP-L is concatenated with the high-resolution features, resulting in 576 visual tokens with 2048 dimensions. These visual tokens possess a substantial capacity for enhancing high-level semantic visual recognition and low-level visual grounding tasks. Then they undergo GeLU activation and are directed through an embedding layer to establish a connection with the language model.
|
||||
Figure 2: Visualization results.
|
||||
DeepSeek-VL is capable of capturing tiny object and giving organized explanations.
|
||||
Vision-Language Adaptor.
|
||||
We employ a two-layer hybrid MLP to bridge the vision encoder and the LLM. Initially, distinct single-layer MLPs are used to process high-resolution features and low-resolution features separately. Subsequently, these features are concatenated along their dimensions and then transformed into the LLM’s input space through another layer of MLP.
|
||||
Language Model. Our language model is built upon DeepSeek LLM ( DeepSeek-AI, 2024 ) whose micro design largely follows the design of LLaMA ( Touvron et al., 2023a ; Touvron et al., 2023b ) , adopting a Pre-Norm structure with RMSNorm ( Zhang and Sennrich, 2019 ) function and using SwiGLU ( Shazeer, 2020 ) as the activation function for the Feed-Forward Network (FFN), with an intermediate layer dimension of 8 3 d m o d e l \frac{8}{3}d_{model} . It also incorporates Rotary Embedding ( Su et al., 2024 ) for positional encoding and uses the same tokenizer with DeepSeek-LLM. We introduce a family of DeepSeek-VL models.
|
||||
Given our objective of conducting joint pretraining with multimodal and language, we select an intermediate checkpoint from DeepSeek’s pretrained models to continue pretraining.
|
||||
Specifically, the DeepSeek-VL-1B model is constructed based on the DeekSeek-LLM-1B model, which underwent training with an approximate corpus of 500 billion text tokens. And the DeekSeek-VL-7B model is developed leveraging the DeepSeek-LLM-7B model trained with an estimated 2 trillion text tokens.
|
||||
Figure 3: Our training pipelines consist of three stages. Stage 1 involves training the Vision-Language (VL) adaptor while keeping the hybrid vision encoder and language model fixed. Stage 2 is the crucial part of the joint vision and language pretraining, where both VL adaptor and language model are trainable. Stage 3 is the supervised fine-tuning phase, during which the low-resolution vision encoder SigLIP-L, VL adaptor, and language model will be trained.
|
||||
## 3.2 Training Pipelines
|
||||
We train our DeepSeek-VL in three consecutive stages as shown in Figure 3 : vision-language adaptor warmup, joint vision-language pretraining, and supervised fine-tuning. We currently focus on visual understanding capabilities and only calculate the next token prediction loss on the language part.
|
||||
## 3.2.1 Stage 1: Training Vision-Language Adaptor
|
||||
The primary objective of this stage is to establish a conceptual link between visual and linguistic elements within the embedding space, thereby facilitating the comprehensive understanding of depicted entities in the images by the Large Language Model (LLM). Consistent with prior research conducted by LLaVA ( Liu et al., 2024b ) and Instruct-BLIP ( Dai et al., 2023 ) , we adopt a similar approach in which both the vision encoder and the LLM remain frozen during this stage, while solely allowing the trainable parameters within the vision-language adaptor. We utilize a dataset comprising 1.25 million image-text paired captions obtained from ShareGPT4V, along with 2.5 million Document OCR rendering pairs to train the VL adaptor.
|
||||
Nevertheless, compared to Large Language Models (LLMs), vision-language adaptors (e.g., a 2-layer MLP) have a significantly smaller parameter capacity. This limitation in model capacity restricts the capabilities that can be learned during this stage. A natural question arises: Can the law of data scaling be effective at this stage? To address this question, we conducted a simple experiment in Table 8 . The results demonstrate that expanding the data scale at this stage does not provide benefits and may even lead to inferior performance. Consequently, we proceed to unfreeze the Large Language Model (LLM) and investigate efficient vision-language pretraining approaches during stage 2.
|
||||
## 3.2.2 Stage 2: Joint Vision-Language pretraining
|
||||
In this stage, we explore effective pretraining strategies which can be considered as an additional stage to enable Large Language Models (LLMs) to comprehend multimodal inputs.
|
||||
We keep the vision encoder frozen and optimize the language model and VL adaptor.
|
||||
Initially, we attempt to directly train the LLM with multimodal data. However, we find while the metrics for multimodal performance incrementally improved, there is a stark and severe decline in language metrics as illustrated in Figure 4 (Multimodal:Language=100%:0%),. This underscores the inherent challenge in directly conducting multimodal pretraining on the foundation of an LLM, revealing a critical trade-off between enhancing multimodal abilities and preserving linguistic proficiency.
|
||||
We hypothesize that the observed phenomenon stems from two primary factors: firstly, the majority of multimodal corpora, are overly simplistic and exhibit a significant divergence from the complexity and distribution of linguistic data. Secondly, there appears to be a competitive dynamic between multimodal and linguistic modalities, leading to what can be described as catastrophic forgetting of language capabilities within the LLM.
|
||||
Figure 4: Comparative performance results on different modality fusion ratio on training stage 2. An excessively large proportion of multimodal data (multimodal:language=100%:0%) leads to significant forgetting of language capabilities in LLMs. A suitable ratio (multimodal:language=70%:30%) can effectively mitigate the issue of language forgetting while simultaneously enhancing the model’s multimodal abilities.
|
||||
Joint Language-multimodal Training
|
||||
To address this challenge, we devise a straightforward yet effective joint language-multimodal training strategy. During training, we not only engage in multimodal data training but also incorporate a large proportion of language data into the training. This approach aims to balance the training focus, mitigating the adverse effects observed. We conduct experiments on the DeepSeek-VL 1B model in Figure 4 to explore the impact of varying the modality mixing ratios.
|
||||
The analysis of the graph yields several key conclusions:
|
||||
(1). Integrating language data significantly alleviates the decline in language capabilities, demonstrating a substantial improvement in the model’s linguistic performance.
|
||||
(2). The inclusion of language data does not lead to a significant loss in multimodal performance, indicating that the model retains its multimodal processing abilities.
|
||||
(3). The performance of different modalities is strongly correlated with their respective proportions in the training dataset, substantiating the competitive relationship between the two modalities.
|
||||
Ultimately, we opt for a training ratio of language to multimodal data of roughly 7:3 for our final model. This ratio enables the model to maintain its language capabilities while simultaneously achieving better pretraining on multimodal data, effectively balancing the development of both language and multimodal proficiencies.
|
||||
Scaling Vision-Language Pretraining
|
||||
Nevertheless, the pretraining stage of the model incurs a substantial computational cost, and performing iterations on the 7B model requires an excessive amount of computing power and time. One suitable strategy involves conducting experiments on a smaller model, specifically the 1.3B model, and subsequently scaling it up to the 7B model. Fortunately, we have observed that a significant portion of the outcomes obtained from the 1.3B models can be effectively transferred to the 7B model through the utilization of SFT (e.g., the encoder design). However, during the stage 2 training phase, we have encountered considerable fluctuations in the generative metrics of the 1.3B model, rendering it challenging to supervise the training process effectively.
|
||||
And this has been discussed in Schaeffer et al. (2024) , "sharp and unpredictable changes might be induced by the researcher’s choice of measurement, even though the model family’s per-token error rate changes smoothly, continuously and predictably with increasing scale." Subsequent experiments have led us to identify the root causes of this issue: the limited capacity of the 1.3B model and the absence of SFT data within the training dataset, both of which hinder the model’s ability to accurately follow instructions. Even when the model possesses knowledge of the correct options, it struggles to generate them precisely.
|
||||
To mitigate these challenges, we adopte a dual-pronged approach. Firstly, we employ the Multi-choice PPL methodology to monitor the model’s progress. This involves inputting not only the prompt and image into the network but also all the answer associated with the question. Subsequently, we calculate the PPL for each answer position (e.g., A, B, C, D) and select the option deemed correct by the model as the final answer. Secondly, we introduce SFT data into the training dataset at a minimal proportion, allowing the model to acquire some proficiency in following instructions. The combination of these two approaches ensures the maintenance of stable training metrics for the 1.3B model and bring better performance after stage3.
|
||||
## 3.2.3 Stage 3: Supervised Fine-tuning
|
||||
In this phase, we finetune the pretrained DeepSeek-VL model with instruction-based fine-tuning to bolster its ability to follow instructions and engage in dialogue, culminating in the creation of the interactive DeepSeek-VL-Chat model.
|
||||
We optimize the language model, VL adaptor, and hybrid vision encoder with the vision-language SFT data as shown in Table 2 , SAM-B remains frozen due to the limited GPU memory. We only supervise answers and special tokens and mask the system and user prompts.
|
||||
To guarantee the model’s comprehensive proficiency in dialogue, we utilize a blend of multimodal data and pure text dialogue data used in DeepSeek-LLM. This approach ensures the model’s versatility across various dialogue scenarios.
|
||||
Figure 5: Visualization results.
|
||||
DeepSeek-VL can understand Python code and provide detailed and organized explanations.
|
||||
## 3.3 Hyperparameters and Infrastructures
|
||||
The detailed hyperparameters of all stages are illustrated in Table 4 . We train and evaluate our DeepSeek-VL with HAI-LLM ( High-flyer, 2023 ) , a lightweight and efficient distributed training framework. Since we use visual encoders to convert images into embedding vectors and then treat image embeddings and text embeddings uniformly, we can easily adapt pipeline parallelism to VL model training: all we need to do is to view visual encoders and text embedding as a single module and take it as the first layer of the resulting model. This very first layer has a complicated model structure and precludes standard tensor parallelism technique, but luckily it requires relatively small computation compared to upper standard transformer blocks. We therefore simply recompute the visual encoder forward pass in all tensor parallel ranks. The existence of visual encoders also leads to non-uniform execution time across model layers, so we re-divide model layers between pipeline parallelism ranks to achieve better load balance and throughput. The upper layers of DeepSeek-VL are exactly the same as those in DeepSeek-LLM. With such minor modification, we can now perform canonical 3D parallelism techniques as in Megatron ( Shoeybi et al., 2019 ; Narayanan et al., 2021 ; Korthikanti et al., 2023 ) and overlap computation and communication as in DeepSeek-LLM ( DeepSeek-AI, 2024 ) .
|
||||
DeepSeek-VL-7B consumed 5 days on a cluster of 64 nodes, each comprising 8 Nvidia A100 GPUs, while DeepSeek-VL-1B consumed 7 days on a setup involving 16 nodes.
|
||||
DeepSeek-VL 1B
|
||||
DeepSeek-VL-7B
|
||||
Vision Encoder
|
||||
SigLIP
|
||||
SigLIP+SAM
|
||||
Hyperparameters
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Learning rate
|
||||
1.0 × 10 − 3 1.0\times 10^{-3}
|
||||
3 × 10 − 5 3\times 10^{-5}
|
||||
2.0 × 10 − 5 2.0\times 10^{-5}
|
||||
1.0 × 10 − 3 1.0\times 10^{-3}
|
||||
4.2 × 10 − 5 4.2\times 10^{-5}
|
||||
2.0 × 10 − 5 2.0\times 10^{-5}
|
||||
LR scheduler
|
||||
Cosine
|
||||
Step
|
||||
Cosine
|
||||
Cosine
|
||||
Step
|
||||
Cosine
|
||||
Weight decay
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
Gradient clip
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
Optimizer
|
||||
AdamW( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
AdamW( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
Warm-up steps
|
||||
128
|
||||
2000
|
||||
256
|
||||
128
|
||||
2000
|
||||
256
|
||||
Training steps
|
||||
15000
|
||||
96000
|
||||
10000
|
||||
15000
|
||||
42000
|
||||
10000
|
||||
Batch size
|
||||
256
|
||||
1024
|
||||
256
|
||||
256
|
||||
2304
|
||||
256
|
||||
Sequence length
|
||||
512
|
||||
4096
|
||||
4096
|
||||
512
|
||||
4096
|
||||
4096
|
||||
Sequence packing
|
||||
× \times
|
||||
✓
|
||||
× \times
|
||||
× \times
|
||||
✓
|
||||
× \times
|
||||
Pipeline parallelism
|
||||
× \times
|
||||
× \times
|
||||
× \times
|
||||
× \times
|
||||
✓
|
||||
✓
|
||||
Table 4: Detailed hyperparameters of our DeepSeek-VL.
|
||||
## 4 Evaluation
|
||||
## 4.1 Public Multimodal Benchmarks Evaluation
|
||||
We evaluate our models on a series of public benchmarks:
|
||||
Multimodal comprehensive understanding datasets: MMMU ( Yue et al., 2023 ) , CMMMU ( Zhang et al., 2024 ) , MMBench ( Liu et al., 2023a ) , MMBench-CN ( Liu et al., 2023a ) , SeedBench ( Li et al., 2023a ) and MMV ( Yu et al., 2023b ) . We compare DeepSeek-VL with competitors on MMB/MMC-dev as current official test download link is no longer active.
|
||||
Chart/table understanding datasets: OCRBench ( Liu et al., 2023b ) ;
|
||||
LLM
|
||||
MMMU
|
||||
CMMMU
|
||||
MMB
|
||||
MMC
|
||||
SEED
|
||||
OCRB
|
||||
POPE
|
||||
MathV
|
||||
MMVet
|
||||
Close-source LMMs :
|
||||
Gemini Pro
|
||||
Unk
|
||||
48.9
|
||||
-
|
||||
75.2
|
||||
74.0
|
||||
70.7
|
||||
659
|
||||
-
|
||||
45.2
|
||||
59.2
|
||||
GPT-4V
|
||||
Unk
|
||||
56.8
|
||||
42.5
|
||||
75.0
|
||||
74.7
|
||||
71.6
|
||||
659
|
||||
-
|
||||
47.8
|
||||
49.9
|
||||
Qwen-VL-Plus
|
||||
Unk
|
||||
45.2
|
||||
39.5
|
||||
66.2
|
||||
69.6
|
||||
72.7
|
||||
-
|
||||
-
|
||||
43.3
|
||||
55.7
|
||||
Qwen-VL-MAX
|
||||
Unk
|
||||
51.4
|
||||
-
|
||||
78.1
|
||||
76.4
|
||||
72.7
|
||||
-
|
||||
-
|
||||
51.0
|
||||
61.8
|
||||
Open-source 13B LMMs :
|
||||
LLaVA-1.5
|
||||
13B
|
||||
36.4
|
||||
-
|
||||
68.2
|
||||
61.9
|
||||
68.2
|
||||
331
|
||||
85.9
|
||||
26.4
|
||||
38.3
|
||||
VILA
|
||||
13B
|
||||
-
|
||||
-
|
||||
70.3
|
||||
64.3
|
||||
-
|
||||
-
|
||||
84.2
|
||||
-
|
||||
38.8
|
||||
LLaVA-Next
|
||||
13B
|
||||
36.2
|
||||
-
|
||||
70.0
|
||||
64.4
|
||||
71.9
|
||||
-
|
||||
86.7
|
||||
35.3
|
||||
48.4
|
||||
Open-source 7B LMMs :
|
||||
EMU2-Chat
|
||||
7B
|
||||
36.3
|
||||
23.8
|
||||
63.6
|
||||
45.9
|
||||
68.9
|
||||
-
|
||||
-
|
||||
30.0
|
||||
31.0
|
||||
Qwen-VL-Chat
|
||||
7B
|
||||
37.0
|
||||
-
|
||||
60.6
|
||||
56.7
|
||||
64.8
|
||||
-
|
||||
-
|
||||
33.8
|
||||
47.3
|
||||
CogVLM
|
||||
7B
|
||||
37.3
|
||||
24.8
|
||||
63.7
|
||||
53.8
|
||||
68.8
|
||||
-
|
||||
-
|
||||
34.7
|
||||
54.5
|
||||
LLaVA-Next
|
||||
7B
|
||||
35.8
|
||||
-
|
||||
67.4
|
||||
60.0
|
||||
70.2
|
||||
-
|
||||
86.5
|
||||
34.6
|
||||
43.9
|
||||
Yi-VL
|
||||
6B
|
||||
37.8
|
||||
35.8
|
||||
68.2
|
||||
68.9
|
||||
67.6
|
||||
-
|
||||
-
|
||||
28.0
|
||||
31.1
|
||||
DeepSeek-VL (ours)
|
||||
7B
|
||||
36.6
|
||||
37.9
|
||||
73.2
|
||||
72.8
|
||||
70.4
|
||||
456
|
||||
88.1
|
||||
36.1
|
||||
41.5
|
||||
Table 5: The comparison between different multi-modal models. The top half are proprietary models, while the bottom are open-source models.
|
||||
LLM
|
||||
MMMU
|
||||
CMMMU
|
||||
MMB
|
||||
MMC
|
||||
SEED
|
||||
OCRB
|
||||
POPE
|
||||
MathV
|
||||
MMVet
|
||||
Tiny Model :
|
||||
MobileVLM
|
||||
1.4B
|
||||
-
|
||||
-
|
||||
53.2
|
||||
-
|
||||
-
|
||||
-
|
||||
84.5
|
||||
-
|
||||
-
|
||||
MobileVLM
|
||||
2.7B
|
||||
-
|
||||
-
|
||||
59.6
|
||||
-
|
||||
-
|
||||
-
|
||||
84.9
|
||||
-
|
||||
-
|
||||
MobileVLM V2
|
||||
1.4B
|
||||
-
|
||||
-
|
||||
59.6
|
||||
-
|
||||
-
|
||||
-
|
||||
84.3
|
||||
-
|
||||
-
|
||||
MobileVLM V2
|
||||
2.7B
|
||||
-
|
||||
-
|
||||
63.2
|
||||
-
|
||||
-
|
||||
-
|
||||
84.7
|
||||
-
|
||||
-
|
||||
LLaVA-Phi
|
||||
2.7B
|
||||
-
|
||||
-
|
||||
59.5
|
||||
-
|
||||
-
|
||||
-
|
||||
85.0
|
||||
-
|
||||
28.9
|
||||
DeepSeek-VL (ours)
|
||||
1.3B
|
||||
32.2
|
||||
27.4
|
||||
64.6
|
||||
61.3
|
||||
66.7
|
||||
409
|
||||
87.6
|
||||
31.1
|
||||
34.8
|
||||
Table 6: The comparison between tiny multi-modal models.
|
||||
Hallucination datasets: POPE ( Li et al., 2023b ) ;
|
||||
Scientific problem datasets: ScienceQA ( Lu et al., 2022a ) and MathVista ( Lu et al., 2023 ) .
|
||||
We apply generation-based evaluation with greedy decoding. The generation-based evaluation here refers to letting the model generate free texts and parsing results from generated texts.
|
||||
The comparative results, as illustrated in Table 5 , show that DeepSeek-VL-7B surpasses most open-source models of similar size across a wide range of benchmarks.
|
||||
DeepSeek-VL outperforms open-source models of similar size in benchmarks such as MMB, MMC, and SEEDbench, even approaching proprietary models (DeepSeek-VL vs. GPT-4V = 70.4 vs. 71.6 on seedbench), demonstrating its powerful natural image comprehension capability. The model also surpasses all open-source models in mathematical logic, but still lags significantly behind proprietary models like GPT-4V (36.1 vs. 47.8 on MathVista). This difference could be attributed to the variance in base model sizes.
|
||||
Furthermore, as shown in Table 6 , DeepSeek-VL-1.3B significantly outperforms models of comparable size. It demonstrates superior performance compared to leading open-source models in the MMB benchmark test, while utilizing only close to half the parameters (1.3B vs. 2.7B), indicating its robust natural image comprehension capability. DeepSeek-VL-1.3B even achieves comparable results to 7B open-source models on MathVista, further validating the powerful logical understanding capabilities of the DeepSeek-VL family.
|
||||
## 4.2 Public Language Benchmarks Evaluation
|
||||
We evaluate our models on the following public language benchmarks:
|
||||
Multi-subject multiple-choice datasets including MMLU ( Hendrycks et al., 2020 ) .
|
||||
Language understanding and reasoning datasets including HellaSwag ( Zellers et al., 2019 ) .
|
||||
Language modeling datasets including Pile ( Gao et al., 2020 ) .
|
||||
Version
|
||||
DeepSeek-VL
|
||||
DeepSeek-VL
|
||||
DeepSeek-LLM
|
||||
1B Chat
|
||||
7B Chat
|
||||
7B Chat
|
||||
Encoder
|
||||
SigLIP
|
||||
SigLIP+SAM
|
||||
None
|
||||
Benchmark
|
||||
HellaSwag
|
||||
56.0
|
||||
68.4
|
||||
68.5
|
||||
MMLU
|
||||
32.5
|
||||
52.4
|
||||
49.4
|
||||
GSM8K
|
||||
18.0
|
||||
55.0
|
||||
63.0
|
||||
MBPP
|
||||
10.0
|
||||
35.2
|
||||
35.2
|
||||
AGIEval
|
||||
14.0
|
||||
27.8
|
||||
19.3
|
||||
Table 7: The performance on language benchmarks.
|
||||
Math datasets including GSM8K ( Cobbe et al., 2021 ) .
|
||||
Code datasets including MBPP ( Austin et al., 2021 ) .
|
||||
Standardized exams including AGIEval ( Zhong et al., 2023 ) .
|
||||
We apply perplexity-based evaluation to datasets that require answers to be chosen from several options. These datasets include HellaSwag and MMLU. The perplexity-based evaluation here refers to calculating the perplexity of each option and selecting the lowest one as the model prediction. Perplexity-based evaluation helps to distinguish subtle probability difference between model predictions and avoids discontinuity of exact match style evaluation. We apply generation-based evaluation with greedy decoding for GSM8K and AGIEval. The generation-based evaluation here refers to letting the model generate free texts and parsing results from generated texts. We apply language-modeling-based evaluation for Pile-test, which means calculating the bits-per-byte on the test corpus. And the results are illustrated in Table 7
|
||||
It can be observed that across the majority of language benchmarks, DeepSeek-VL performs comparably to, or even surpasses, DeepSeek-7B. For instance, it achieves scores of 68.4 vs. 68.5 on HellaSwag, which serves as a general benchmark for evaluating general language ability. DeepSeek-VL outperforms DeepSeek-7B on metrics such as MMLU and AGIEval, indicating that multimodal training methods may even aid in language tasks. Nevertheless, DeepSeek-VL-7B shows a certain degree of decline in mathematics (GSM8K), which suggests that despite efforts to promote harmony between vision and language modalities, there still exists a competitive relationship between them. This could be attributed to the limited model capacity (7B), and larger models might alleviate this issue significantly. Overall, DeepSeek-VL strives to achieve the goal of minimizing declines in language capability while addressing these challenges.
|
||||
## 4.3 Human Evaluation
|
||||
To further explore the capabilities of our DeepSeek-VL, we independently construct a dataset for manual evaluation. This dataset comprises 100 questions, divided into seven categories, each encompassing specific tasks. These categories and tasks are same as our taxonomy for the in-house SFT data, as shown in Table 3 . This approach ensures that the tasks we test are universal and encompass the majority of use cases for multimodal models.
|
||||
Moreover, based on the categories and tasks described in existing reports, we collect similar image materials and developed prompts. The sources for these image materials include royalty-free image communities and photographs taken by the researchers. This methodical collection and prompt formulation process ensures our dataset is both comprehensive and representative of real-world multimodal model applications.
|
||||
Figure 6: Human evaluation results on InternLM-XComposer2-VL ( Dong et al., 2024 ) , CogVLM ( Wang et al., 2023b ) , DeepSeek-VL and GPT-4V ( OpenAI, 2023b ) .
|
||||
We compare our DeepSeek-VL-7B with InternLM-XComposer2-VL, CogVLM and GPT-4V as shown in Figure 6 (and we also provide visualization results in Appendix A ).
|
||||
GPT-4V demonstrates exceptional performance across most dimensions. All open-source models are still far behind GPT-4V in logical reasoning, highlighting the necessity of scaling up the size of Large Language Models (LLMs). DeepSeek-VL-7B achieves better results in overall performance, reaching outcomes close to GPT-4V in Recognition, Conversion, and Commonsense Reasoning.
|
||||
Figure 7: GPT-4V-based Evaluation Results of DeepSeek-VL vs. Other Models: The chart depicts results from a GPT-4V-based assessment across 99 test samples, demonstrating DeepSeek-VL’s favorable outcomes against both open-source and proprietary models.
|
||||
In addition, we conduct a comparative assessment using GPT-4V to evaluate the performance of DeepSeek-VL against other models across a set of 99 test samples designed for human evaluation.
|
||||
Following ( Zheng et al., 2024 ) , we show GPT-4V the question and the answers from two different models and ask GPT-4V to determine which one is better or declare a tie.
|
||||
The results indicate a preference for DeepSeek-VL’s responses in the majority of cases, as GPT-4V tends to rate the quality of DeepSeek-VL’s answers more favorably.
|
||||
As illustrated in Figure 7 , DeepSeek-VL is judged to be superior in over 60% of instances when compared to open-source multimodal models, including Fuyu-8B, CogVLM-17B, and InternLM-XComposer2-VL.
|
||||
Moreover, in comparison with other proprietary models, such as GPT-4V itself, DeepSeek-VL demonstrates comparably exceptional performance.
|
||||
## 4.4 Ablation Study
|
||||
Scale Up Projector Training
|
||||
We expand the dataset for stage 1 (projector warmup) and subsequently apply supervised fine-tuning. The results, depicted in Figure 8 , demonstrate that augmenting the training data volume does not enhance performance at this stage. This implies that the projector’s capacity is inherently constrained, rendering it incapable of capturing the extensive knowledge necessary for multimodal tasks.
|
||||
Stage 1, Training Step
|
||||
MMB
|
||||
MMC
|
||||
SEED
|
||||
POPE
|
||||
MMMU
|
||||
Average
|
||||
2K
|
||||
59.0
|
||||
54.0
|
||||
61.8
|
||||
82.3
|
||||
30.3
|
||||
57.5
|
||||
8K
|
||||
58.0
|
||||
45.0
|
||||
58.5
|
||||
84.9
|
||||
29.2
|
||||
55.1
|
||||
20K
|
||||
56.0
|
||||
52.3
|
||||
59.0
|
||||
81.7
|
||||
28.6
|
||||
55.5
|
||||
80K
|
||||
58.1
|
||||
55.0
|
||||
58.6
|
||||
78.6
|
||||
27.9
|
||||
55.6
|
||||
Table 8: Comparative directly SFT performance results on scaling up stage 1 data. The results demonstrate that expanding the data scale at this stage does not yield benefits, or even results in worse performance.
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
MMB
|
||||
MMC
|
||||
SEED
|
||||
POPE
|
||||
MMMU
|
||||
Average
|
||||
✓
|
||||
✓
|
||||
59.4
|
||||
54.2
|
||||
61.4
|
||||
82.5
|
||||
29.2
|
||||
57.4
|
||||
✓
|
||||
✓
|
||||
63.4
|
||||
60.5
|
||||
65.9
|
||||
87.1
|
||||
31.8
|
||||
61.7
|
||||
✓
|
||||
✓
|
||||
✓
|
||||
64.3
|
||||
61.3
|
||||
66.7
|
||||
87.6
|
||||
32.2
|
||||
62.4
|
||||
Table 9: Analysis of model performance across training stages.
|
||||
Training Stage
|
||||
In Table 9 , we examine the contributions of each stage to the model’s performance. It’s evident that combining stage 1, stage 2, and stage 3 yields significantly better results across all metrics compared to combining stage 1 and stage 3 alone, demonstrating the effectiveness of multimodal pretraining. Additionally, the combination of stage 2 and stage 3 still slightly lags behind the combined performance of stage 1, stage 2, and stage 3, indicating that vision-language adaptor warmup stage remains meaningful.
|
||||
Modality Group Training When mixing language and multimodal data, we observe that directly blending them at the batch level significantly reduces training efficiency. This inefficiency arises because each batch gradient backpropagation process waits for the slowest sample to complete. As a result, the predominantly faster-to-process pure language data ends up waiting for the multimodal samples to finish, leading to a decrease in overall training efficiency.
|
||||
Figure 8: Comparative analysis of modality warmup on language (Pile-test) and multimodal (MMBench and MMBench _ \_ CN) benchmarks demonstrates that modality grouping consistently surpasses the non-grouped modality approach in language tasks, while simultaneously preserving performance on multimodal tasks on training stage 2 (Multimodal:Language=60%:40%).
|
||||
To address this issue, we experiment with grouping different modalities of data at each global step, sampling distinct modalities separately. This approach involves organizing the training data so that batches are composed either entirely of language data or entirely of multimodal data at different training steps, rather than mixing them within the same batch.
|
||||
The results are shown in Figure 8 ,
|
||||
we observe that this method does not compromise the model’s performance while enhancing the model’s training efficiency by 20%. This strategy effectively circumvents the bottleneck caused by the disparate processing times between modalities, optimizing the training workflow.
|
||||
Modality Warmup Considering that our approach involves multimodal training on the foundation of a language model, directly mixing multimodal data in a fixed proportion from the outset can destabilize the model. To counteract this issue, we propose a simple yet effective modality warm-up strategy. Initially, we set the language data ratio to 1, and then gradually decrease it to the target ratio for the final model training (e.g., 0.7).
|
||||
Figure 9: Comparative performance results on language (Pile-test) and multimodal (MMBench and MMBench_CN) benchmarks for modality warmup. Modality warmup consistently matches or surpasses the performance of approaches without modality warmup across all evaluated tasks on training stage 2 (Multimodal:Language=60%:40%).
|
||||
Our experiments, as illustrated in Figure 9 , demonstrate that this strategy effectively prevents a significant decline in language capabilities at the beginning of training, while also yielding comparatively superior outcomes in the final phases for both the language and multimodal domains. This gradual adaptation enables the model to more seamlessly adjust to the incorporation of multimodal data, thereby improving overall training stability and performance.
|
||||
Vision Encoder Selection In order to better acquire and utilize image information, we compare the training loss of different vision encoders under our training settings except for reducing training steps of stage 2 to 8000 for efficiency. As illustrated in
|
||||
Figure 10 , the incorporation of vision-only self-supervised encoders has been found to significantly enhance performance on training loss. To more effectively process high-resolution images, our research ultimately adopts a hybrid vision encoder strategy, combining SigLIP with SAM for our model’s implementation.
|
||||
Figure 10: Comparative analysis of different vision encoders on training losses in stage 2.
|
||||
Vision-Language Adaptor Design
|
||||
To improve the efficiency of extracting information from the visual encoder while adhering to current token length constraints, adjustments can be made to the Vision-Language adaptor in two main ways: the method used to combine visual features and the design of the MLP adaptor.
|
||||
Previous studies ( Tong et al., 2024 ) have indicated that combining visual features along the sequence dimension can lead to better model performance, although this comes with the trade-off of increased computational requirements due to a longer sequence of visual feature tokens. As demonstrated in the top section of Table 10 , reducing the sequence length by stacking visual features along the image’s width or height dimensions before sequence concatenation, in order to keep the sequence length constant, does not achieve better results compared to simply merging them along the embedding dimension in most metrics.
|
||||
In terms of the adaptor architecture, employing separate MLP adaptors for each vision feature encoder allows for more precise adjustments to the specific values and distribution patterns of visual features, facilitating smoother model training. Conversely, using a shared MLP adaptor for different vision encoders contributes to adequate feature fusion. We adopt a mixed strategy and report stable and improved performance, as outlined in the lower section of Table 10 .
|
||||
Architecture
|
||||
MMB
|
||||
MMC
|
||||
SEED
|
||||
POPE
|
||||
ScienceQA
|
||||
MMMU
|
||||
OCRB
|
||||
Average
|
||||
Sequence Concatenation:
|
||||
Token Pooling - W
|
||||
61.2
|
||||
59.6
|
||||
61.6
|
||||
86.5
|
||||
57.7
|
||||
31.6
|
||||
304
|
||||
55.5
|
||||
Token Pooling - H
|
||||
59.9
|
||||
58.3
|
||||
61.6
|
||||
83.8
|
||||
55.0
|
||||
32.0
|
||||
291
|
||||
54.2
|
||||
Embedding Concatenation:
|
||||
Hybrid MLP
|
||||
61.7
|
||||
60.1
|
||||
62.9
|
||||
87.8
|
||||
56.6
|
||||
31.3
|
||||
309
|
||||
55.9
|
||||
Shared MLP
|
||||
62.0
|
||||
58.9
|
||||
62.5
|
||||
86.6
|
||||
54.7
|
||||
30.2
|
||||
318
|
||||
55.2
|
||||
Separate MLP
|
||||
57.5
|
||||
58.7
|
||||
63.1
|
||||
86.5
|
||||
56.6
|
||||
29.0
|
||||
299
|
||||
54.5
|
||||
Table 10: Comparison of different adaptor architectures using SigLIP and SAM as hybrid vision encoder, Hybrid MLP are used for sequence concatenation experiments. Bolded entries represent the best results, while underlined entries denote the second-best results. For calculating the average score, we divide the OCRBench by the total number of questions.
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
In this technical report, we have introduced DeepSeek-VL, a series of Multimodal Large Language Models, available in scales of 1.3B and 6.7B parameters. This report has unveiled the limitations inherent in the predominant projector-based pretraining methodologies, setting the stage for the innovative approach adopted by DeepSeek-VL. By prioritizing a joint vision and language (VL) pretraining phase, DeepSeek-VL transcends traditional models by ensuring that the integration of multimodal data does not compromise the linguistic capabilities of the Large Language Models (LLMs). This is achieved through a strategic warm-up data ratio and the introduction of a hybrid vision encoder, which together enable the efficient processing of high-resolution images without losing sight of semantic richness.
|
||||
The incorporation of a hybrid vision encoder, capable of handling 1024 x 1024 images within a constrained token budget, underscores our commitment to preserving the nuanced details and semantic integrity across diverse tasks. As a result, DeepSeek-VL emerges as a pioneering model that not only meets but exceeds the standards set by generalist models in its class. It showcases exceptional performance across a wide range of visually-centric benchmarks while sustaining formidable proficiency in language-centric evaluations.
|
||||
In making DeepSeek-VL publicly available, we aim to catalyze further innovation and exploration within the research community, providing a robust foundation upon which future studies can build. This gesture of openness is intended to facilitate the collective advancement of our understanding and capabilities in handling multimodal data.
|
||||
Looking ahead, we are excited to announce plans to scale up DeepSeek-VL to larger sizes, incorporating Mixture of Experts (MoE) technology. This forthcoming expansion promises to further enhance the model’s efficiency and effectiveness, opening up new horizons for research and application in the field of AI.
|
||||
|
||||
=== BODY (tail) ===
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,378 @@
|
||||
TITLE: DeepSeek-Prover: Advancing Theorem Proving in LLMs through Large-Scale Synthetic Data
|
||||
ARXIV: 2405.14333
|
||||
AUTHORS: Huajian Xin , Daya Guo , Zhihong Shao , Zhizhou Ren , Qihao Zhu , Bo Liu , Chong Ruan , Wenda Li , Xiaodan Liang
|
||||
DATE: [Submitted on 23 May 2024]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-Prover: Advancing Theorem Proving in LLMs through Large-Scale Synthetic Data
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Background and Related Works
|
||||
## 3 Approach
|
||||
## 3.1 Autoformalization
|
||||
## 3.2 Quality Filtering
|
||||
## 3.3 Statement Proving
|
||||
## 3.4 Iterative Enhancement
|
||||
## 4 Experiments
|
||||
## 4.1 Experimental Setup
|
||||
## 4.2 Main Results
|
||||
## 4.3 Ablation Studies
|
||||
## 4.3.1 The Effectiveness of Large-scale Autoformalization
|
||||
## 4.3.2 The Effectiveness of Formal Statements Scoring
|
||||
## 4.3.3 The Effectiveness of Iterative Enhancement
|
||||
## 4.3.4 The Effectiveness of Scaling Synthetic Theorem Proving Data
|
||||
## 5 Case Studies
|
||||
## 5.1 Autoformalized Theorem with Complete Proof
|
||||
## 5.2 Autoformalization of Theorem with Inconsistent Hypotheses
|
||||
## 6 Conclusion
|
||||
## Broader Impact
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-Prover: Advancing Theorem Proving in LLMs through Large-Scale Synthetic Data
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Background and Related Works
|
||||
3 Approach
|
||||
3.1 Autoformalization
|
||||
3.2 Quality Filtering
|
||||
3.3 Statement Proving
|
||||
3.4 Iterative Enhancement
|
||||
4 Experiments
|
||||
4.1 Experimental Setup
|
||||
4.2 Main Results
|
||||
4.3 Ablation Studies
|
||||
4.3.1 The Effectiveness of Large-scale Autoformalization
|
||||
4.3.2 The Effectiveness of Formal Statements Scoring
|
||||
4.3.3 The Effectiveness of Iterative Enhancement
|
||||
4.3.4 The Effectiveness of Scaling Synthetic Theorem Proving Data
|
||||
5 Case Studies
|
||||
5.1 Autoformalized Theorem with Complete Proof
|
||||
5.2 Autoformalization of Theorem with Inconsistent Hypotheses
|
||||
6 Conclusion
|
||||
References
|
||||
A Appendix / supplemental material
|
||||
A.1 Prompts
|
||||
A.2 Case Studies of Autoformalization
|
||||
A.3 Example Lean Proofs Generated by Our Model
|
||||
A.3.1 Results on MiniF2F-Test Dataset
|
||||
A.3.2 Results on FIMO Dataset
|
||||
A.4 Formal Verification Details
|
||||
License: CC BY 4.0
|
||||
arXiv:2405.14333v1 [cs.AI] 23 May 2024
|
||||
## DeepSeek-Prover: Advancing Theorem Proving in LLMs through Large-Scale Synthetic Data
|
||||
Huajian Xin
|
||||
Affiliation: DeepSeek
|
||||
Affiliation: Sun Yat-sen University
|
||||
Daya Guo
|
||||
Affiliation: DeepSeek
|
||||
Email: benjaminliu.eecs@gmail.com
|
||||
Zhihong Shao
|
||||
Affiliation: DeepSeek
|
||||
Email: wli8@ed.ac.uk
|
||||
Z.Z. Ren Qihao Zhu Bo Liu
|
||||
Affiliation: DeepSeek
|
||||
Email: xdliang328@gmail.com
|
||||
Chong Ruan Wenda Li Xiaodan Liang
|
||||
† † thanks: Corresponding author.
|
||||
Email: {xinhj, guoday, zhihongshao, rzz, zhuqh, chong.ruan}@deepseek.com,
|
||||
Affiliation: DeepSeek
|
||||
Affiliation: Sun Yat-sen University
|
||||
Affiliation: University of Edinburgh
|
||||
Affiliation: MBZUAI
|
||||
## Abstract
|
||||
Proof assistants like Lean have revolutionized mathematical proof verification, ensuring high accuracy and reliability. Although large language models (LLMs) show promise in mathematical reasoning, their advancement in formal theorem proving is hindered by a lack of training data. To address this issue, we introduce an approach to generate extensive Lean 4 proof data derived from high-school and undergraduate-level mathematical competition problems. This approach involves translating natural language problems into formal statements, filtering out low-quality statements, and generating proofs to create synthetic data. After fine-tuning the DeepSeekMath 7B model on this synthetic dataset, which comprises 8 million formal statements with proofs, our model achieved whole-proof generation accuracies of 46.3% with 64 samples and 52% cumulatively on the Lean 4 miniF2F test, surpassing the baseline GPT-4 at 23.0% with 64 samples and a tree search reinforcement learning method at 41.0%. Additionally, our model successfully proved 5 out of 148 problems in the Lean 4 Formalized International Mathematical Olympiad (FIMO) benchmark, while GPT-4 failed to prove any. These results demonstrate the potential of leveraging large-scale synthetic data to enhance theorem-proving capabilities in LLMs. Both the synthetic dataset and the model will be made available to facilitate further research in this promising field.
|
||||
## 1 Introduction
|
||||
In modern mathematics, the increasing complexity of proofs presents substantial challenges for peer review. This complexity has led to the acceptance of erroneous proofs, with critical flaws often detected only after considerable time. To address these issues, formal mathematical languages such as Lean ( De Moura et al., 2015 ; Moura and Ullrich, 2021 ) , Isabelle ( Paulson, 1994 ) , and Coq ( The Coq Development Team, ) have been developed. These languages enable the creation of computer-verifiable proofs ( Avigad, 2023 ) . However, crafting formal proofs demands significant effort, specialized expertise, and poses challenges even for seasoned mathematicians. Consequently, the significance of automated theorem proving is on the rise ( Shulman, 2024 ) .
|
||||
To reduce the effort involved in writing formal mathematical proofs, several approaches ( Polu and Sutskever, 2020 ; Jiang et al., 2021 ; Han et al., 2021 ; Polu et al., 2022 ; Lample et al., 2022 ; Jiang et al., 2022a ; Yang et al., 2024 ) have been developed, primarily focusing on search algorithms that explore potential solutions for proposed theorems. However, these methods often struggle with the vast search spaces required for complex theorems, rendering them ineffective for more intricate proofs ( Loos et al., 2017 ) . Recently, advances in large language models (LLMs) have introduced a novel strategy, utilizing pre-trained models to guide the search process. Although these new methods ( Jiang et al., 2022b ; Zhao et al., 2023 ; Xin et al., 2023 ) represent significant improvements, they still fall short of practical applicability due to the lack of parallel corpus. Unlike conventional programming languages such as Python or Java, formal proof languages are used by relatively few mathematicians, resulting in limited datasets. Recent advances in autoformalization ( Wu et al., 2022 ) allow more aligned data to be synthesized to train LLM-based automated theorem provers. Nevertheless, the resulting dataset remains too small to fully unleash the capabilities of LLMs.
|
||||
To address this issue, we propose a method for generating extensive Lean 4 proof data from informal mathematical problems. Our approach translates high-school and undergraduate-level mathematical competition problems into formal statements. We then automate proof generation using a large language model (LLM) and verify the correctness of these proofs within the Lean 4 environment. The primary challenge of this method is to ensure both the scale and quality of the synthetic data.
|
||||
Quality Assurance: We enhance the quality of generated proofs through a multi-step process. First, we filter out simple statements using a quality scoring model and exclude invalid statements via a hypothesis rejection strategy. Our novel iterative framework then improves proof quality by initially generating synthetic statements from informal math problems using an under-trained LLM fine-tuned on limited data. These statements are used to generate corresponding proofs, which are validated for correctness using a Lean 4 verifier. The correct theorem-proof pairs are subsequently used to further train the initial model. Through several iterations, the model trained on large-scale synthetic data becomes significantly more powerful than the originally under-trained LLMs, resulting in higher-quality theorem-proof pairs.
|
||||
Scale Assurance: To accelerate the proof generation process, our method addresses the challenge of the large search space for proofs. A significant cause of delays is the generation of unprovable statements that continue to be processed until they reach the time limit. To mitigate this, we propose proving negated statements in parallel. Once either the original statement or its negation is proved, the entire proving process is terminated.
|
||||
We assess the effectiveness of our method on Lean 4 theorem proving using 488 problems from miniF2F ( Zheng et al., 2021 ) and 148 problems from the FIMO benchmarks ( Liu et al., 2023 ) . We utilize DeepSeekMath 7B ( Shao et al., 2024 ) , a state-of-the-art mathematical model, as our base. The results show that our iteratively trained model performs strongly, achieving 46.3% accuracy in whole-proof generation on the miniF2F-test benchmark with 64 samples, surpassing GPT-4 ( Achiam et al., 2023 ) at 23.0% and a reinforcement learning method at 41.0%. Additionally, our approach solved 4 out of 148 problems in the FIMO benchmark with 100 samples, while GPT-4 solved none, and our approach solved 5 with 4096 samples. Ablation experiments indicate that the model progressively solves more problems in miniF2F with each iteration. In summary, our paper makes the following contributions:
|
||||
•
|
||||
We introduce an iterative method to synthesize 8 million formal statements, each accompanied by a formal proof, from informal math problems. Experimental results demonstrate that this method significantly enhances both the scalability and quality of synthetic data.
|
||||
•
|
||||
Our model, trained on this synthetic dataset, achieves state-of-the-art performance on benchmarks, with whole-proof generation accuracies of 46.3% using 64 samples and 52% cumulatively on the Lean 4 miniF2F test. This surpasses the baseline GPT-4 at 23.0% with 64 samples and a tree search reinforcement learning method at 41.0%. Additionally, our model successfully proved 5 out of 148 problems in the Lean 4 Formalized International Mathematical Olympiad (FIMO) benchmark, while GPT-4 failed to prove any.
|
||||
•
|
||||
We contribute to the mathematical and AI communities by creating and open-sourcing a large dataset of high-quality formal mathematical proofs, thereby fostering further research and development in automated theorem proving.
|
||||
## 2 Background and Related Works
|
||||
Automated theorem proving has been a significant area of interest in artificial intelligence research since its inception ( Bibel, 2013 ) . Initial efforts were directed at simpler logical frameworks, which led to the development of highly efficient first-order theorem provers like E ( Schulz, 2002 ) and Vampire ( Kovács and Voronkov, 2013 ) . Nonetheless, these tools often fall short in handling complex theorems commonly found in modern proof assistants such as Lean ( De Moura et al., 2015 ) , Isabelle ( Paulson, 1994 ) , and Coq ( The Coq Development Team, ) . The advent of recent deep learning models and model-guided search techniques has reinvigorated the field ( Bansal et al., 2019 ) . This modern approach has not only enhanced the capabilities of ATP systems but also expanded their applicability in solving more intricate mathematical problems.
|
||||
ATP with Neural Models. With the development of deep learning, several approaches have been proposed to combine neural models with ATP ( Loos et al., 2017 ) . A series of ATP approaches adopts tree search algorithms guided by neural models ( Polu and Sutskever, 2020 ; Han et al., 2021 ; Polu et al., 2022 ; Jiang et al., 2022a ; Yang et al., 2024 ) . These approaches primarily utilize reinforcement learning techniques to enhance the accuracy of the model ( Kaliszyk et al., 2018 ; Crouse et al., 2021 ; Wu et al., 2021 ; Lample et al., 2022 ) . Since the search space is significantly large, the searching process consumes considerable time and computing resources.
|
||||
Another series of ATP approaches harnesses the power of large language models. These approaches typically involve language models that are fine-tuned with open-source proof data and interact with verifiers via a state-action transition program ( Polu and Sutskever, 2020 ; Jiang et al., 2021 ; Han et al., 2021 ; Polu et al., 2022 ; Lample et al., 2022 ; Jiang et al., 2022a ; Yang et al., 2024 ) . This process iteratively generates proof steps and verifies their correctness with formal verifiers. It then generates the next proof steps based on the proof states returned by the formal verifiers. Although these approaches achieve high performance, they are computationally intensive. To enhance efficiency, recent researches leverage language models to generate complete formal proofs directly ( First et al., 2023 ; Jiang et al., 2022b ; Zhao et al., 2023 ; Xin et al., 2023 ) , thus bypassing the iterative interaction during proof generation.
|
||||
Autoformalization for Formal Mathematics.
|
||||
Due to the limited availability of formal corpora for training, the performance of current large language models (LLMs) is also constrained. Thus, some approaches propose autoformalization ( Wu et al., 2022 ; Jiang et al., 2022b ) , which involves converting natural language descriptions into formal statements that can be verified by proof assistants. Several studies have generated synthetic datasets of formal proofs using rule-based transformations of existing theorems ( Wu et al., 2020 ; Wang and Deng, 2020 ; Xiong et al., 2023 ) . While effective, these methods are constrained by their reliance on predefined rules and lack flexibility for broader applications. Recent methodologies adopt large language models to translating natural language problems into formal statements ( Huang et al., 2024 ) . However, these datasets remain smaller than needed and are limited to small mathematical benchmarks, leading to only minor improvements in training outcomes for language models. In this paper, we aim to synthesise formal proofs via autoformalization at a much larger scale to boost the performance of a neural prover.
|
||||
## 3 Approach
|
||||
Figure 1 : An overview of our approach.
|
||||
In this section, we introduce our approach, which consists of four key processes as depicted in Figure 1 . The initial phase concentrates on generating formal mathematical statements from a broad collection of informal math problems, necessitating further proof. Next, the autoformalized statements are filtered through model scoring and hypothesis rejection methods to select high-quality statements. These statements are then proved by a model called DeepSeek-Prover, with their correctness verified by the formal verifier called Lean 4 1 1
|
||||
1
|
||||
𝚕𝚎𝚊𝚗𝚙𝚛𝚘𝚟𝚎𝚛 / 𝚕𝚎𝚊𝚗𝟺 : 𝚟𝟺 .7.0 − 𝚛𝚌𝟸 \mathtt{leanprover/lean4:v4.7.0-rc2} , yielding validated formal statements and proofs. These data serve as synthetic data for fine-tuning the DeepSeek-Prover. After enhancing DeepSeek-Prover, we repeat the entire previously described process. This cycle continues until the improvements in DeepSeek-Prover become marginal. Notably, to enhance proof efficiency, we prove concurrently both the original statements and their negations. This method has the advantage of swiftly discarding the original statement when it is invalid by proving its negation. The details of each phase will be described in the subsequent sections.
|
||||
## 3.1 Autoformalization
|
||||
The generation of formal proof data fundamentally relies on the availability of a substantial corpus of formal statements. In practice, however, amassing a large collection of manually crafted formal statements is challenging. Fortunately, the internet is replete with math-related problems expressed in natural language. By autoformalizing these informal mathematical problems, we can generate a vast repository of formal statements.
|
||||
We have observed that problems with explicit conditions and well-defined goals are typically easier to formalize compared to advanced mathematical topics that necessitate intricate definitions and constructions. Consequently, this paper primarily examines high school and undergraduate-level competition problems, with a particular emphasis on algebra and number theory, and to a lesser extent, combinatorics, geometry, and statistics. Despite their apparent simplicity, these problems often involve complex solution techniques, making them excellent candidates for constructing proof data to improve theorem-proving capabilities in Large Language Models (LLMs). To compile our dataset, we employed web scraping and careful data cleaning techniques to extract problems from online resources featuring high school and undergraduate exercises, exams, and competitions, resulting in a dataset of 869,659 high-quality natural language math problems.
|
||||
Specifically, we initialized the DeepSeek-Prover using the DeepSeekMath-Base 7B model ( Shao et al., 2024 ) . Initially, the model struggled to convert informal math problems into formal statements. To address this, we fine-tuned the DeepSeek-Prover model using the MMA dataset ( Jiang et al., 2023 ) , which comprises formal statements from Lean 4’s mathlib 2 2
|
||||
2
|
||||
The specific mathlib commit used is 𝟼𝟺𝟻𝟸𝟾𝟸𝟼𝟾 𝚋 𝟹 𝚌 𝟸 𝚌 𝚏 𝟻𝟽𝟾𝟼𝟹𝟿 𝚋 𝚌 𝟺𝟽𝟿𝟾𝟸𝟾𝟾𝟾𝟸 𝚊 𝟿 𝚎 𝚌 𝚍 𝟹 𝚊 𝟾𝟸 \mathtt{64528268b3c2cf578639bc479828882a9ecd3a82} . that were back-translated into natural language problem descriptions by GPT-4. We then instructed the model to translate these natural language problems into formal statements in Lean 4 using a structured approach.
|
||||
Prompt :
|
||||
Mathematical Problem in Natural Language:
|
||||
{$informal_statement_with_answers}
|
||||
Translate the problem to Lean 4 (only the core declaration):
|
||||
‘‘‘lean4
|
||||
Response :
|
||||
{$formal_statement}
|
||||
‘‘‘
|
||||
## 3.2 Quality Filtering
|
||||
The quality of the autoformalized statements was found to be suboptimal due to two main issues. Firstly, many formal statements were overly simplistic. To address this, we developed scoring criteria and provided examples from miniF2F-valid as few-shot examples to guide the DeepSeek-Prover model in evaluating the content and quality of these statements using a chain-of-thought approach. Manual review of these scores confirmed that the model’s evaluations closely matched human intuition and expectations. Specifically, the model was instructed (see Appendix A.1 for the detailed prompt) to classify the quality of each formal statement into categories: "excellent," "good," "above average," "fair," or "poor." Statements rated as "fair" or "poor" were subsequently excluded.
|
||||
The second issue pertains to formal statements that, although provable, are based on inconsistent hypotheses leading to vacuous conclusions, rendering the conclusions meaningless in mathematics. For example, consider the following model-generated statement:
|
||||
⬇
|
||||
example (θ : ℝ) ( h ₀ : ∀ z : ℂ, z ^ 2 = -1 ∧ z ^ 3 = -1 ∧ z ^ 6 = 1) ( h ₁ : Real . tan θ = 2 * Real . sqrt 3) : θ = 5 * Real . pi / 3
|
||||
Here, the hypothesis z 2 = − 1 ∧ z 3 = − 1 ∧ z 6 = 1 z^{2}=-1\wedge z^{3}=-1\wedge z^{6}=1 for all complex numbers is clearly false, making any derived conclusions meaningless. To eliminate such cases from our dataset, we implemented a hypothesis rejection method. This involves using the DeepSeek-Prover model to attempt proving the formal statement with ’False’ as the conclusion. A successful proof indicates an invalid hypothesis, prompting exclusion of the statement. An example is shown below:
|
||||
⬇
|
||||
example (θ : ℝ) ( h ₀ : ∀ z : ℂ, z ^ 2 = -1 ∧ z ^ 3 = -1 ∧ z ^ 6 = 1) ( h ₁ : Real . tan θ = 2 * Real . sqrt 3) : False := by
|
||||
simpa using h ₀ 1
|
||||
By applying this dual strategy of model scoring and hypothesis rejection, we curated a refined set of 712,073 high-quality formal statements, providing a robust foundation for further proof synthesis.
|
||||
## 3.3 Statement Proving
|
||||
After creating a substantial corpus of high-quality formal statements, we employed the model to search for proofs of these statements. Traditionally, language models have been used predominantly in a brute-force manner to prove theorems—repeatedly attempting until a valid proof is found or computational resources are exhausted. This approach is inefficient for our purposes. Typically, language models are applied to human-curated formal statements that are carefully crafted and generally true and provable; however, in our task of proving autoformalized statements, many of the statements produced by the model may be incorrect. Indeed, it is unreasonable to expect the model to validate a false proposition within any reliable proof system. This issue becomes more pronounced during large-scale autoformalization, where we observed that at least 20% of the formal statements generated by our model, even after quality filtering, were incorrect, leading to significant computational waste if addressed with brute force.
|
||||
To minimize resource wastage on unprovable statements and improve the efficiency of the proof search process, we exploited the logical symmetry between a statement and its negation to accelerate proof synthesis. We implemented dual concurrent proof searches for each synthetic statement—one for the statement Γ ⊢ P \Gamma\vdash P and another for its negation Γ ⊢ ¬ P \Gamma\vdash\neg P . The search terminates as soon as a valid proof is found for either, conclusively demonstrating the unprovability of the other. Each proof search stream attempts up to k k proofs unless a valid proof emerges sooner.
|
||||
All validated proofs, whether they justify the original theorems or their negations, are then aggregated to further train the DeepSeek-Prover. Thus, this dual approach serves as a form of data augmentation, enriching the dataset with both propositions and their negations—even if the original propositions were not correctly formalized by the model.
|
||||
## 3.4 Iterative Enhancement
|
||||
Since the entire pipeline heavily relies on the DeepSeek-Prover, enhancing the model’s performance after each iteration is crucial. To achieve this, we consistently fine-tune the model with newly generated data. The updated model is then utilized for subsequent autoformalization iterations. The key insight from this iterative process is that the model incrementally improves in strength and efficacy after each cycle of refinement and application. This iterative process continues until no further gains are observed. Consequently, the theorem-proof pairs generated by the model become increasingly higher in quality with each iteration. This method ensures that the DeepSeek-Prover consistently enhances its performance, ultimately producing superior theorem-proof pairs through continuous refinement.
|
||||
## 4 Experiments
|
||||
## 4.1 Experimental Setup
|
||||
DeepSeek-Prover is built upon DeepSeekMath-Base 7B model ( Shao et al., 2024 ) , a decoder-only transformer ( Vaswani et al., 2017 ) pre-trained on a corpus comprising 120 billion math-related tokens. We fine-tuned this model using a global batch size of 512 and a constant learning rate of 1 × 10 − 4 1\times 10^{-4} , incorporating 6,000 warmup steps with synthetic data. DeepSeek-Prover’s performance was evaluated against several baselines:
|
||||
•
|
||||
GPT-3.5 and GPT-4 ( Achiam et al., 2023 ) , developed by OpenAI, are advanced generative AI models known for their effectiveness in diverse tasks, including code generation. Although not explicitly designed for theorem proving, their extensive scale and parameter count confer significant capabilities. In contrast, DeepSeekMath is a specialized model, explicitly pre-trained for mathematical content. We utilized both GPT-4 (specifically the GPT-4-turbo 0409 version) and DeepSeekMath to generate complete proofs for given theorems using a methodology similar to ours.
|
||||
•
|
||||
GPT-f ( Polu and Sutskever, 2020 ) , utilizing a GPT-2-inspired architecture ( Radford et al., 2019 ) , implements an iterative best-first search method to progressively generate and validate proof steps within a formal proof setting until a proof is either completed or resources are depleted. This methodology has been further advanced by Proof Artifact Co-Training ( Han et al., 2021 ) , ReProver ( Yang et al., 2024 ) , Llemma ( Azerbayev et al., 2023 ) , and COPRA ( Thakur et al., 2023 ) , which employ either specialized fine-tuned models or versatile general-purpose models such as GPT-3.5 and GPT-4 for the generation of proof steps.
|
||||
## 4.2 Main Results
|
||||
This study addresses complex mathematical problems in algebra and number theory. We evaluate the theorem-proving efficacy of our model using the miniF2F ( Zheng et al., 2021 ) and FIMO ( Liu et al., 2023 ) benchmarks. The metric pass@k is employed to denote the scenario where at least one valid proof is discovered among the first k attempts generated by the model.
|
||||
Table 1 :
|
||||
Comparing with state-of-the-arts on the miniF2F dataset.
|
||||
Method
|
||||
Model size
|
||||
Generation times
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
Tree Search Methods
|
||||
COPRA (GPT-3.5) ( Thakur et al., 2023 )
|
||||
-
|
||||
1 × 60 1\times 60
|
||||
-
|
||||
9.0 % 9.0\%
|
||||
COPRA (GPT-4) ( Thakur et al., 2023 )
|
||||
-
|
||||
1 × 60 1\times 60
|
||||
-
|
||||
26.6 % 26.6\%
|
||||
Proof Artifact Co-Training ( Han et al., 2021 )
|
||||
837M
|
||||
1 × 8 × 512 1\times 8\times 512
|
||||
23.9 % 23.9\%
|
||||
24.6 % 24.6\%
|
||||
8 × 8 × 512 8\times 8\times 512
|
||||
29.3 % 29.3\%
|
||||
29.2 % 29.2\%
|
||||
ReProver ( Yang et al., 2024 )
|
||||
229M
|
||||
1 × 3751 1\times 3751
|
||||
-
|
||||
25.0 % 25.0\%
|
||||
Llemma ( Azerbayev et al., 2023 )
|
||||
7B
|
||||
1 × 3200 1\times 3200
|
||||
-
|
||||
26.2 % 26.2\%
|
||||
Llemma ( Azerbayev et al., 2023 )
|
||||
34B
|
||||
1 × 3200 1\times 3200
|
||||
-
|
||||
25.8 % 25.8\%
|
||||
Curriculum Learning ( Polu et al., 2022 )
|
||||
837M
|
||||
1 × 8 × 512 1\times 8\times 512
|
||||
33.6 % 33.6\%
|
||||
29.6 % 29.6\%
|
||||
8 × 8 × 512 8\times 8\times 512
|
||||
41.2 % 41.2\%
|
||||
34.5 % 34.5\%
|
||||
64 × 8 × 512 64\times 8\times 512
|
||||
47.3 % 47.3\%
|
||||
36.6 % 36.6\%
|
||||
Hypertree Proof Search ( Lample et al., 2022 )
|
||||
600M
|
||||
cumulative
|
||||
58.6 % 58.6\%
|
||||
-
|
||||
64 × 5000 64\times 5000
|
||||
-
|
||||
41.0 % 41.0\%
|
||||
Whole-Proof Generation Methods
|
||||
GPT-4-turbo 0409
|
||||
-
|
||||
64
|
||||
25.4 % 25.4\%
|
||||
23.0 % 23.0\%
|
||||
DeepSeekMath-Base ( Shao et al., 2024 )
|
||||
7B
|
||||
128
|
||||
25.4 % 25.4\%
|
||||
27.5 % 27.5\%
|
||||
DeepSeek-Prover
|
||||
7B
|
||||
cumulative
|
||||
60.2%
|
||||
52.0%
|
||||
1 1 (greedy)
|
||||
-
|
||||
30.0 % 30.0\%
|
||||
64 64
|
||||
-
|
||||
46.3 % 46.3\%
|
||||
128 128
|
||||
-
|
||||
46.3 % 46.3\%
|
||||
8192 8192
|
||||
-
|
||||
48.8 % 48.8\%
|
||||
65536 65536
|
||||
-
|
||||
50.0 % 50.0\%
|
||||
Results on MiniF2F.
|
||||
The miniF2F benchmark consists of 244 validation and 244 test problems, ranging from basic arithmetic to competition-level problems, e.g., problems from the American Invitational Mathematics Examination (AIME), the American Mathematics Competitions (AMC), and the International Mathematical Olympiad (IMO). We use the version of miniF2F in Lean 4, which was released by the LeanDojo project ( https://github.com/yangky11/miniF2F-lean4 ).
|
||||
Table 1 compares various state-of-the-art methods on the miniF2F dataset. DeepSeek-Prover outperforms all with cumulative scores of 60.2% on miniF2F-valid and 52.0% on miniF2F-test, significantly higher than other methods, including GPT-4 which scores 25.41% and 22.95%, respectively. Even the best tree search method, Hypertree Proof Search with a 600M model, achieves only up to 58.6% on miniF2F-valid and 41.0% on miniF2F-test. DeepSeek-Prover’s scalability is evident as its performance improves with increased computational resources, rising from 30.03% using a greedy approach to 50.0% at 65536 generation times, demonstrating its effectiveness in handling complex proof scenarios. Examples of proved theorems of MiniF2F can be found in Appendix A.3.1 .
|
||||
Results on FIMO.
|
||||
The FIMO benchmark comprises 149 formal problems which are sourced from the IMO shortlist translated into Lean 4.
|
||||
Our method successfully proved 4 theorems with 100 attempts per theorem, whereas GPT-4 failed to prove any.
|
||||
By increasing the number of attempts per theorem to 4,096, we successfully proved an additional theorem. Examples of proved theorems of FIMO can be found in Appendix A.3.2 .
|
||||
## 4.3 Ablation Studies
|
||||
## 4.3.1 The Effectiveness of Large-scale Autoformalization
|
||||
To demonstrate the effectiveness of large-scale autoformalization, we conducted a comparative analysis as shown in Table 2 between our autoformalized dataset and conventional datasets using expert iteration ( Polu and Sutskever, 2020 ) . This iterative method entails generating formal proofs, fine-tune the model based on successful outcomes, and iterating this process until no additional enhancements are observed. The results indicate that models trained with our autoformalized data significantly outperform those trained solely with mathlib data.
|
||||
Table 2 :
|
||||
Improvement in pass rates for miniF2F at pass@128 in models trained on formal proofs, including those derived from human-authored theorems in Lean 4’s mathlib and automatically formalized theorems.
|
||||
Model
|
||||
#Tokens
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
-
|
||||
-
|
||||
25.4 % 25.4\%
|
||||
27.5 % 27.5\%
|
||||
Mathlib
|
||||
0.238 0.238 B
|
||||
30.3 % 30.3\%
|
||||
31.2 % 31.2\%
|
||||
Autoformalized Statements
|
||||
3.108 3.108 B
|
||||
48.8 % 48.8\%
|
||||
42.6 % 42.6\%
|
||||
## 4.3.2 The Effectiveness of Formal Statements Scoring
|
||||
To demonstrate the effectiveness of the model in filtering out low-quality statements, we fine-tuned the DeepSeekMath-Base model using an equal amount of high-score proof data and low-score proof data to verify the quality of the data, as shown in Table 3 . The table shows that the model trained on high-score proof data outperformed the model trained on low-score proof data by 4.5%. This enhancement underscores the utility of the model in accurately scoring and effectively filtering out lower-quality statements.
|
||||
Table 3 : Improvement in pass rates for miniF2F at pass@128 in models trained on differently scored proof data.
|
||||
Scored Class
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
"excellent", "good" and "above average"
|
||||
48.8 % 48.8\%
|
||||
42.6 % 42.6\%
|
||||
"fair" and "poor"
|
||||
41.4 % 41.4\%
|
||||
38.1 % 38.1\%
|
||||
## 4.3.3 The Effectiveness of Iterative Enhancement
|
||||
Table 4 demonstrates a distinct correlation between the number of iterations in data synthesis and enhanced performance in theorem proving. This evidence underscores the success of our iterative enhancement strategy in augmenting theorem-proving capabilities. Successive iterations not only refine the model’s ability to handle complex proofs but also significantly increase the quality and quantity of the synthetic data produced.
|
||||
Table 4 :
|
||||
Improvement in pass rates for miniF2F at pass@128 in models across successive training iterations, facilitated by the incremental integration of synthesized data via autoformalization.
|
||||
Model
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
iteration 0
|
||||
38.1 % 38.1\%
|
||||
34.0 % 34.0\%
|
||||
iteration 1
|
||||
45.1 % 45.1\%
|
||||
39.3 % 39.3\%
|
||||
iteration 2
|
||||
49.2 % 49.2\%
|
||||
41.4 % 41.4\%
|
||||
iteration 3
|
||||
54.5 % 54.5\%
|
||||
45.1 % 45.1\%
|
||||
iteration 4
|
||||
59.4 % 59.4\%
|
||||
46.3 % 46.3\%
|
||||
## 4.3.4 The Effectiveness of Scaling Synthetic Theorem Proving Data
|
||||
Our investigation into synthetic theorem proving data reveals a clear correlation between dataset size and model efficacy, as illustrated in Table 5 . By examining subsets of the eight million generated proof data points, we observed that performance on the miniF2F benchmark improves proportionally to the exponential increase in dataset size. This pattern highlights the pivotal importance of large-scale datasets for boosting model proficiency in automatically formalizing natural language questions. These findings emphasize the significant potential and necessity of systematic data construction for progressing in the field of automated theorem proving.
|
||||
Table 5 :
|
||||
Improvement in pass rates for miniF2F at pass@128 in models trained with a larger fraction of synthesized data via autoformalization.
|
||||
Size
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
1,000
|
||||
22.95 % 22.95\%
|
||||
24.18 % 24.18\%
|
||||
10,000
|
||||
32.79 % 32.79\%
|
||||
31.97 % 31.97\%
|
||||
100,000
|
||||
36.07 % 36.07\%
|
||||
37.7 % 37.7\%
|
||||
1,000,000
|
||||
39.34 % 39.34\%
|
||||
38.11 % 38.11\%
|
||||
8,066,621
|
||||
42.62 % 42.62\%
|
||||
40.16 % 40.16\%
|
||||
## 5 Case Studies
|
||||
This section presents two case studies to demonstrate the application of our methods in autoformalizing theorems. It showcases both successful proofs and the identification of inconsistencies during the Hypothesis Rejection stage.
|
||||
## 5.1 Autoformalized Theorem with Complete Proof
|
||||
Example a. Problem : Prove that the determinant of the following matrix is zero.
|
||||
[ 1 cos ( a − b ) cos ( a ) cos ( a − b ) 1 cos ( b ) cos ( a ) cos ( b ) 1 ] \begin{bmatrix}1&\cos(a-b)&\cos(a)\\
|
||||
\cos(a-b)&1&\cos(b)\\
|
||||
\cos(a)&\cos(b)&1\end{bmatrix}
|
||||
Autoformalized Theorem in Lean :
|
||||
⬇
|
||||
example ( a b : ℝ) :
|
||||
Matrix . det ![![1, Real . cos ( a - b ), Real . cos a ], ![ Real . cos ( a - b ), 1, Real . cos b ], ![ Real . cos a , Real . cos b , 1]] = 0
|
||||
This approach effectively translates the algebraic expression of the matrix and its determinant into a formal language using Lean. The autoformalization captures the essence of the original mathematical statement by defining a specific 3 × 3 3\times 3 matrix dependent on real numbers a a and b b , and asserts that its determinant is zero. The formalization employs the Matrix.det function to compute the determinant, utilizing the ![...] notation for lists of lists in Lean to represent the matrix rows.
|
||||
## 5.2 Autoformalization of Theorem with Inconsistent Hypotheses
|
||||
Example b. Problem : Given a real number D D and the condition that for non-zero real numbers a , b , c a,b,c , the determinant of the matrix [ a b c 1 4 9 3 1 2 ] \begin{bmatrix}a&b&c\\
|
||||
1&4&9\\
|
||||
3&1&2\end{bmatrix} equals D D , prove that D 2 = 154 D^{2}=154 .
|
||||
Autoformalized Theorem in Lean :
|
||||
⬇
|
||||
example ( D : ℝ) ( h ₀ : ∀ a b c : ℝ, a ≠ 0 ∧ b ≠ 0 ∧ c ≠ 0 →
|
||||
Matrix . det ![![ a , b , c ], ![1, 4, 9], ![3, 1, 2]] = D ) : D ^ 2 = 154
|
||||
The initial autoformalization incorrectly assumes that the condition D 2 = 154 D^{2}=154 universally applies to all non-zero real numbers a a , b b , and c c . This assumption is not supported by the problem statement, which does not claim universal applicability. Instead, the formalization should aim to either identify specific values of a a , b b , and c c that satisfy D 2 = 154 D^{2}=154 or demonstrate that no such values exist.
|
||||
The model successfully identifies this inconsistency and provides a counterexample to demonstrate the absurdity of the hypothesis:
|
||||
⬇
|
||||
example ( D : ℝ) ( h ₀ : ∀ a b c : ℝ, a ≠ 0 ∧ b ≠ 0 ∧ c ≠ 0 →
|
||||
Matrix . det ![![ a , b , c ], ![1, 4, 9], ![3, 1, 2]] = D ) : False := by
|
||||
have h ₁ := h ₀ 1 2 3
|
||||
have h ₂ := h ₀ 1 4 9
|
||||
simp [ Matrix . det_fin_three ] at h ₁ h ₂
|
||||
linarith
|
||||
A corrected version of the autoformalized theorem can be proposed as follows:
|
||||
⬇
|
||||
example ( a b c : ℝ) ( h ₀ : a ≠ 0 ∧ b ≠ 0 ∧ c ≠ 0) :
|
||||
let D := Matrix . det ![![ a , b , c ], ![1, 4, 9], ![3, 1, 2]];
|
||||
D ^ 2 = 154
|
||||
These examples illustrate the model’s capability to verify proofs and identify hypothesis inconsistencies effectively. Further details can be found in Appendix A.2 .
|
||||
## 6 Conclusion
|
||||
In this paper, we presented a method to generate extensive synthetic proof data from high-school and undergraduate-level mathematical competition problems. By translating natural language problems into formal statements, filtering out low-quality ones, and using iterative proof generation, we created 8 million proof data points and significantly improved the DeepSeekMath 7B model’s performance in ATP when trained on this synthetic data. Our model outperforms GPT-4 and other methods on benchmarks like miniF2F and FIMO. By open-sourcing our dataset and model, we aim to advance research in automated theorem proving and enhance the capabilities of large language models in formal mathematical reasoning. Currently, our work mainly focuses on algebra and number theory at the middle school and undergraduate levels. In future work, we will aim to expand the diversity of mathematical problems addressed, enhancing the general applicability of our methods in ATP.
|
||||
## Broader Impact
|
||||
The research presented in this paper has the potential to significantly advance automated theorem proving by leveraging large-scale synthetic proof data generated from informal mathematical problems. This remarkable advancement can enhance the capabilities of large language models in formal theorem proving, contributing to more reliable mathematical proof verification and providing valuable educational resources for students and researchers. By directly releasing the code, model, and data, we aim to ensure the responsible use of our work, fostering further innovation and maintaining high standards of data privacy and intellectual property compliance.
|
||||
|
||||
=== BODY (tail) ===
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,538 @@
|
||||
TITLE: Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models
|
||||
ARXIV: 2407.01906
|
||||
AUTHORS: Zihan Wang , Deli Chen , Damai Dai , Runxin Xu , Zhuoshu Li , Y. Wu
|
||||
DATE: [Submitted on 2 Jul 2024 ( v1 ), last revised 5 Jul 2024 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Related Work
|
||||
## 2.1 Parameter-efficient fine-tuning for dense architectural LLMs
|
||||
## 2.2 Coarse- and Fine-grained MoE LLMs
|
||||
## 3 Methods
|
||||
## 3.1 Preliminaries: Mixture-of-Experts for Transformers
|
||||
## 3.2 Probing Task-Specific Expert Specialization in MoE Models
|
||||
## Expert Routing is Concentrated in the Same Task
|
||||
## Active Experts Vary Significantly across Tasks
|
||||
## 3.3 Expert-Specialized Fine-tuning (ESFT)
|
||||
## Data Sampling
|
||||
## Expert Relevance Score
|
||||
## Expert Selection and Fine-tuning
|
||||
## 4 Experiment Setup
|
||||
## 4.1 Main Evaluation
|
||||
## 4.1.1 Tasks for Model Enhancement
|
||||
## 4.1.2 Tasks for Model Adaptation
|
||||
## 4.2 General Ability Evaluation
|
||||
## 4.3 Backbone Model and Training Settings
|
||||
## 5 Results
|
||||
## 5.1 Benchmark Performance Results
|
||||
## 5.2 Computational Efficiency Results
|
||||
## Training Time
|
||||
## Storage Space
|
||||
## 6 Analysis
|
||||
## 6.1 ESFT Leverages Specialized Experts Effectively
|
||||
## 6.2 ESFT Leverages Training Resources Efficiently
|
||||
## 6.3 Selectively Training Non-Shared Parameters is the Key to ESFT
|
||||
## 6.4 Analysis of Key Modules in ESFT
|
||||
## Expert Relevance Score Function
|
||||
## Fine-Grained Expert Segmentation of the MoE Model
|
||||
## 7 Conclusion
|
||||
## Limitations
|
||||
|
||||
=== BODY (head) ===
|
||||
Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Work
|
||||
2.1 Parameter-efficient fine-tuning for dense architectural LLMs
|
||||
2.2 Coarse- and Fine-grained MoE LLMs
|
||||
3 Methods
|
||||
3.1 Preliminaries: Mixture-of-Experts for Transformers
|
||||
3.2 Probing Task-Specific Expert Specialization in MoE Models
|
||||
Expert Routing is Concentrated in the Same Task
|
||||
Active Experts Vary Significantly across Tasks
|
||||
3.3 Expert-Specialized Fine-tuning (ESFT)
|
||||
Data Sampling
|
||||
Expert Relevance Score
|
||||
Expert Selection and Fine-tuning
|
||||
4 Experiment Setup
|
||||
4.1 Main Evaluation
|
||||
4.1.1 Tasks for Model Enhancement
|
||||
4.1.2 Tasks for Model Adaptation
|
||||
4.2 General Ability Evaluation
|
||||
4.3 Backbone Model and Training Settings
|
||||
5 Results
|
||||
5.1 Benchmark Performance Results
|
||||
5.2 Computational Efficiency Results
|
||||
Training Time
|
||||
Storage Space
|
||||
6 Analysis
|
||||
6.1 ESFT Leverages Specialized Experts Effectively
|
||||
6.2 ESFT Leverages Training Resources Efficiently
|
||||
6.3 Selectively Training Non-Shared Parameters is the Key to ESFT
|
||||
6.4 Analysis of Key Modules in ESFT
|
||||
Expert Relevance Score Function
|
||||
Fine-Grained Expert Segmentation of the MoE Model
|
||||
7 Conclusion
|
||||
References
|
||||
A Examples for Specialized Tasks
|
||||
B Strategy for Grouping Experts
|
||||
C Analysis of Expert Affinity Sample Size
|
||||
D Detailed Results for Ablations on Training Shared Parameters
|
||||
E Qualitative Examples of the Expert Choices
|
||||
F The Impact of Mixing Alignment Data for Training
|
||||
G Evaluation Instructions for Specialized Tasks
|
||||
H Evaluating Math and Code as General Tasks
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2407.01906v2 [cs.CL] 05 Jul 2024
|
||||
## Let the Expert Stick to His Last: Expert-Specialized Fine-Tuning for Sparse Architectural Large Language Models
|
||||
Zihan Wang
|
||||
† † thanks: Work done during internship at DeepSeek.
|
||||
Email: zw@deepseek.com
|
||||
Deli Chen
|
||||
Affiliation: DeepSeek AI
|
||||
Email: victorchen@deepseek.com
|
||||
Damai Dai
|
||||
Affiliation: DeepSeek AI
|
||||
Runxin Xu
|
||||
Affiliation: DeepSeek AI
|
||||
Zhuoshu Li
|
||||
Affiliation: DeepSeek AI
|
||||
Y. Wu
|
||||
Affiliation: DeepSeek AI
|
||||
Affiliation: Northwestern University
|
||||
## Abstract
|
||||
Parameter-efficient fine-tuning ( PEFT ) is crucial for customizing Large Language Models ( LLMs ) with constrained resources.
|
||||
Although there have been various PEFT methods for dense-architecture LLMs, PEFT for sparse-architecture LLMs is still underexplored.
|
||||
In this work, we study the PEFT method for LLMs with the Mixture-of-Experts ( MoE ) architecture and the contents of this work are mainly threefold:
|
||||
(1) We investigate the dispersion degree of the activated experts in customized tasks, and found that the routing distribution for a specific task tends to be highly concentrated, while the distribution of activated experts varies significantly across different tasks.
|
||||
(2) We propose E xpert- S pecialized F ine- T uning, or ESFT, which tunes the experts most relevant to downstream tasks while freezing the other experts and modules; experimental results demonstrate that our method not only improves the tuning efficiency, but also matches or even surpasses the performance of full-parameter fine-tuning.
|
||||
(3) We further analyze the impact of the MoE architecture on expert-specialized fine-tuning. We find that MoE models with finer-grained experts are more advantageous in selecting the combination of experts that are most relevant to downstream tasks, thereby enhancing both the training efficiency and effectiveness. Our code is available at https://github.com/deepseek-ai/ESFT .
|
||||
## 1 Introduction
|
||||
As the parameter scale of large language models ( LLMs ) continues to increase ( Meta, 2024 ; Mistral, 2024a ; DeepSeek, 2024 ; Qwen, 2024 ) , parameter-efficient fine-tuning ( PEFT ) methods ( Han et al., 2024 ) are becoming increasingly important in adapting pre-trained LLMs to downstream customization tasks.
|
||||
However, existing works on PEFT like low-rank adaptation (LoRA) and P-Tuning ( Hu et al., 2021 ; Liu et al., 2021 ) have primarily focused on dense-architecture LLMs, with research on sparse-architecture LLMs still being markedly insufficient.
|
||||
In this work, we focus on exploring PEFT techniques within the Mixture-of-Experts ( MoE ) LLMs ( Mistral, 2024b ; Databricks, 2024 ) , as introduced in § 3.1 .
|
||||
Unlike dense models where all tasks are handled by the same parameters, in the MoE architecture, different tasks are processed by distinct activated experts ( Lepikhin et al., 2021 ; Fedus et al., 2021 ) .
|
||||
Observations indicate that task specialization in expert systems is the key to the MoE LLM performance ( Dai et al., 2024 ) .
|
||||
We further illustrate such specialization in § 3.2 that experts activated by the same task’s data are concentrated, while those for different tasks vary significantly, suggesting MoE models use specialized expert combinations to handle different tasks.
|
||||
Motivated by this, we propose Expert-Specialized Fine-Tuning ( ESFT ), as illustrated in § 3.3 .
|
||||
ESFT only tunes the experts with the highest affinity to the task, while freezing the parameters of other experts and modules.
|
||||
The primary advantages of ESFT lie in two aspects:
|
||||
(1) Maintaining Expert Specialization : ESFT prevents the decrement of specialization in full-parameter fine-tuning, where experts not adept at the task also update their parameters.
|
||||
Experimental results in § 5.1 show that ESFT can achieve aligned or even superior performance in downstream tasks compared to full-parameter fine-tuning, and better maintains performance in general tasks.
|
||||
(2) Saving Computation Resources : ESFT only trains the parameters of the selected experts, which effectively reduces the storage of up to 90% and training time up to 30% compared to full-parameter fine-tuning, as shown in § 5.2 .
|
||||
Besides, we delve deeper into the working mechanism of the ESFT method. We analyze the expert selection process in § 6.1 and demonstrate how ESFT leverages specialized experts effectively, as selecting 5-15% experts can achieve promising performance in different tasks.
|
||||
We investigate the efficiency of ESFT under different computational constraints in § 6.2 , showcasing its ability to leverage training resources efficiently compared to other PEFT methods like LoRA.
|
||||
Our studies in § 6.3 analyze the effects of shared and non-shared parameters in the model on specialized and general performance, pointing out the priority to selectively train non-shared parameters in ESFT.
|
||||
Through ablation studies in § 6.4 , we highlight the importance of our expert relevance scores and the fine-grained expert segmentation architecture.
|
||||
## 2 Related Work
|
||||
## 2.1 Parameter-efficient fine-tuning for dense architectural LLMs
|
||||
The goal of parameter-efficient fine-tuning ( Han et al., 2024 ) is to efficiently customize LLMs for downstream tasks, while existing studies primarily focus on dense architectural LLMs.
|
||||
PEFT methods for dense models can generally be categorized into three approaches:
|
||||
(1) Adding new parameters : methods of this kind fix the existing model parameters and fine-tune the model on a small number of newly added parameters. Adapter ( Houlsby et al., 2019 ; Pfeiffer et al., 2020 ; He et al., 2021 ; Wang et al., 2022 ) and Soft Prompt ( Li and Liang, 2021 ; Liu et al., 2021 ; Zhang et al., 2023b ; Lester et al., 2021 ) are two typical representatives of this category of methods.
|
||||
(2) Selecting existing parameters : methods of this type fine-tune a limited part of existing parameters, while keeping the majority of the other parameters fixed. Based on whether the trainable parameter space is continuous, these methods can generally be divided into structured training ( Guo et al., 2020 ; Gheini et al., 2021 ; He et al., 2023 ; Vucetic et al., 2022 ) and unstructured training ( Liao et al., 2023 ; Ansell et al., 2021 ; Sung et al., 2021 ; Xu et al., 2021 ) .
|
||||
(3) Applying low-rank adaptation : LoRA ( Hu et al., 2021 ; Fomenko et al., 2024 ) is a widely-used PEFT method, which decomposes the origin weight matrices into low-rank components. Subsequent works ( Zhang et al., 2023a ; Ding et al., 2023 ; Lin et al., 2024 ; Liu et al., 2023 ) have introduced numerous improvements to the original LoRA method.
|
||||
However, the study of PEFT in sparse models is still scarce. In this work, we select and tune part of the experts based on their downstream task affinity, as a unique selection dimension exclusive to the sparse MoE architecture.
|
||||
## 2.2 Coarse- and Fine-grained MoE LLMs
|
||||
Compared to dense LLMs (e.g., LLaMA series, Meta, 2023b ; Meta, 2023a ), MoE LLMs (e.g., Mixtral series, Mistral, 2024a ; Mistral, 2024b ) can increase model size while saving training and inference costs.
|
||||
Based on the granularity of experts, existing large MoE models can generally be divided into two categories: coarse- and fine-grained expert LLMs.
|
||||
Most existing MoE LLMs ( Lepikhin et al., 2021 ; Fedus et al., 2021 ; Roller et al., 2021 ; Dai et al., 2022 ; Shen et al., 2024 ) have coarse-grained experts where the number of experts is very limited.
|
||||
For example, 2 out of 8 experts are activated for Mixtral MoE series ( Mistral, 2024a ; Mistral, 2024b ) and Grok-V1 ( XAI, 2024 ) .
|
||||
As a result, a single expert has to learn complicated patterns from different domain tasks simultaneously.
|
||||
To address this issue, DeepSeek MoE ( Dai et al., 2024 ) has introduced fine-grained expert segmentation. In the DeepSeek-V2 ( DeepSeek, 2024 ) , there are as many as 162 experts, with 8 active experts (8 out of 66 experts are activated for the DeepSeek-V2-Lite).
|
||||
The fine-grained division of experts ensures a high degree of specialization among the experts. Moreover, the specialized expert system enables the selection of experts that are most relevant to the task for efficient tuning.
|
||||
## 3 Methods
|
||||
Figure 1: Comparison between Expert-Specialized Fine-Tuning (ESFT) and other fine-tuning methods. FFT trains all parameters. LoRA combines pre-trained weights with low-rank matrices to reduce training costs. ESFT only trains a subset of experts in a Mixture-of-Expert (MoE) architecture, optimizing efficiency and task specialization.
|
||||
## 3.1 Preliminaries: Mixture-of-Experts for Transformers
|
||||
Mixture-of-Experts (MoE) for Transformers replace Feed-Forward Networks (FFNs) with MoE layers. Each MoE layer consists of multiple experts structurally identical to a FFN. Tokens are assigned to and processed by a subset of the most relevant experts based on their affinity scores, ensuring computational efficiency in MoE layers. The output hidden state 𝐡 t l \mathbf{h}_{t}^{l} of the t t -th token in the l l -th MoE layer is computed as:
|
||||
𝐡 t l = ∑ i = 1 N ( g i , t FFN i n ( 𝐮 t l ) ) + 𝐮 t l , \mathbf{h}_{t}^{l}=\sum_{i=1}^{N}\left(g_{i,t}\text{FFN}^{n}_{i}(\mathbf{u}_{t}^{l})\right)+\mathbf{u}_{t}^{l},
|
||||
(1)
|
||||
g i , t = { s i , t , s i , t ∈ TopK ( { s j , t | 1 ⩽ j ⩽ N } , K ) , 0 , otherwise , g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}{\in}\text{TopK}(\{s_{j,t}|1{\leqslant}j{\leqslant}N\},K),\\
|
||||
0,&\text{otherwise},\end{cases}
|
||||
(2)
|
||||
s i , t = Softmax i ( 𝐮 t l ⊤ 𝐞 i l ) , s_{i,t}=\text{Softmax}_{i}\left(\mathbf{u}_{t}^{l\top}\mathbf{e}_{i}^{l}\right),
|
||||
(3)
|
||||
where N N denotes the total number of experts, FFN i ( ⋅ ) \text{FFN}_{i}(\cdot) is the i i -th expert FFN, g i , t g_{i,t} denotes the gate value for the i i -th expert, s i , t s_{i,t} denotes the token-to-expert affinity, TopK ( ⋅ , K ) \text{TopK}(\cdot,K) denotes the set comprising K K highest affinity scores among those calculated for the t t -th token and all N N experts, and 𝐞 i l \mathbf{e}_{i}^{l} is the centroid of the i i -th expert in the l l -th layer.
|
||||
Recently, DeepSeekMoE ( Dai et al., 2024 ) proposes enhancements to the MoE architecture through several techniques, including
|
||||
(1) Fine-grained segmentation, segmenting each expert into multiple smaller ones and keeping the same fraction of experts to process each token, allowing specialization in different knowledge types while maintaining the same computational cost. (2) Shared expert isolation, leveraging shared experts that process all tokens to capture common knowledge, reducing parameter redundancy and enhancing efficiency.
|
||||
The output of an MoE layer in DeepSeekMoE is:
|
||||
𝐡 t l = ∑ i = 1 K s FFN i s ( 𝐮 t l ) + ∑ i = 1 N ( g i , t FFN i n ( 𝐮 t l ) ) + 𝐮 t l , \mathbf{h}_{t}^{l}{=}\sum_{i=1}^{K_{s}}\text{FFN}^{s}_{i}(\mathbf{u}_{t}^{l}){+}\sum_{i=1}^{N}(g_{i,t}\text{FFN}^{n}_{i}{(}\mathbf{u}_{t}^{l}{)}){+}\mathbf{u}_{t}^{l},
|
||||
(4)
|
||||
g i , t = { s i , t , s i , t ∈ TopK ( { s j , t | 1 ⩽ j ⩽ N } , K − K s ) , 0 , otherwise , g_{i,t}=\begin{cases}s_{i,t},~s_{i,t}{\in}\text{TopK}(\{s_{j,t}|1{\leqslant}j{\leqslant}N\},K{-}K_{s}),\\
|
||||
0,~~~\text{otherwise},\end{cases}
|
||||
(5)
|
||||
where K s K_{s} is the number of shared experts, FFN i s \text{FFN}^{s}_{i} and FFN i n \text{FFN}^{n}_{i} denote the shared and non-shared experts, respectively. Each expert is segmented into m m ones, with N N and K K also multiplied by m m times compared to the coarse-grained architecture.
|
||||
## 3.2 Probing Task-Specific Expert Specialization in MoE Models
|
||||
Despite the significant success of MoE LLMs, a clear understanding of the underlying mechanism remains elusive.
|
||||
We conduct probing experiments to understand how non-shared experts are utilized across various tasks. These tasks, as detailed in § 4.1 , include general domains like math and code, as well as specialized domains like intent recognition, summarization, legal judgment prediction, and translation. These experiments reveal the expert specialization in MoE models in two aspects:
|
||||
## Expert Routing is Concentrated in the Same Task
|
||||
We investigate the distribution of normalized gate values, i.e., the sum of all expert-token gate values for each expert, divided by the total across all experts.
|
||||
Figure 2 displays this distribution, where the experts are sorted by their normalized values from high to low.
|
||||
The figure shows that a small subset of experts handles the majority of gate values, indicating the model’s and concentrated expert allocation for a specific task.
|
||||
## Active Experts Vary Significantly across Tasks
|
||||
We investigate the joint distribution of experts across tasks.
|
||||
Figure 3 shows a heatmap of the shared Top-6 experts for two independent data samples per task averaged across layers. This indicates the degree of overlap of experts used within the same task or between different tasks.
|
||||
Off-diagonal values are near 0, and diagonal values are near 6, indicating that the same task uses similar experts, while different tasks use different sets.
|
||||
Figure 2: Top Expert distribution for specific tasks. Shaded areas represent variance across layers. The figure shows that few experts handle most gate values, highlighting expert specialization for different tasks.
|
||||
Figure 3: The average number of shared Top-6 routed experts across tasks. The values are averaged by layer, indicating that the sets of experts used for the same task are consistent while different tasks are distinct.
|
||||
## 3.3 Expert-Specialized Fine-tuning (ESFT)
|
||||
The highly specialized expert system suggests that different experts can be optimized for specific tasks. Inspired by this, we propose Expert-Specialized Fine-Tuning (ESFT) for MoE LLM customization, which selectively fine-tunes the most relevant experts for downstream tasks to enhance computational efficiency and maintain expert specialization. Figure 1 illustrates the differences between our method and existing methods. Below, we introduce our method step by step.
|
||||
## Data Sampling
|
||||
We randomly sample a subset D s = { ( x i , y i ) } i = 1 N s D_{s}=\{(x_{i},y_{i})\}_{i=1}^{N_{s}} from the training data D = { ( x i , y i ) } i = 1 N D=\{(x_{i},y_{i})\}_{i=1}^{N} for expert selection, where x i x_{i} and y i y_{i} denote the input and label, respectively. Empirically, we find that a subset of 32 32 concatenated samples, each with a fixed length of L = 4096 L=4096 , is robust enough to select the most relevant experts for a task. We detail this claim in Appendix C .
|
||||
## Expert Relevance Score
|
||||
We propose two methods to calculate the relevance of an expert to a task based on its affinity to the sample tokens, defined as average gate score and token selection ratio, respectively. Both methods assess each expert’s relevance to downstream tasks and can be chosen based on task-specific experimental performance.
|
||||
Average Gate Score (ESFT-Gate) This score calculates the average affinity of expert e i e_{i} to all tokens in the sampled data. It is defined as:
|
||||
g i l = 1 N s ∑ j = 1 N s 1 L j ∑ k = 1 L j g i , k l , g_{i}^{l}=\frac{1}{N_{s}}\sum_{j=1}^{N_{s}}\frac{1}{L_{j}}\sum_{k=1}^{L_{j}}g_{i,k}^{l},
|
||||
(6)
|
||||
where L j L_{j} is the length of the input sequence x j x_{j} in the sampled data D s D_{s} .
|
||||
Token Selection Ratio (ESFT-Token) This score calculates the ratio of tokens for which expert e i e_{i} is selected. It is defined as:
|
||||
r i l = 1 N s ∑ j = 1 N s 1 L j ∑ k = 1 L j 𝟙 ( g i , k l > 0 ) K , r_{i}^{l}=\frac{1}{N_{s}}\sum_{j=1}^{N_{s}}\frac{1}{L_{j}}\sum_{k=1}^{L_{j}}\frac{\mathbbm{1}\left(g_{i,k}^{l}>0\right)}{K},
|
||||
(7)
|
||||
where 𝟙 ( g i , k l > 0 ) \mathbbm{1}\left(g_{i,k}^{l}>0\right) is an indicator that equals 1 if the gate score g i , k l g_{i,k}^{l} is positive, and 0 otherwise. K K is the number of experts selected per token.
|
||||
## Expert Selection and Fine-tuning
|
||||
For each MoE layer l l , we select a subset of experts to be fine-tuned based on their relevance scores. We define a threshold p ∈ ( 0 , 1 ] p\in(0,1] as a hyperparameter controlling the proportion of total relevance scores to be included in the selected subset. For each layer l l , we select a set of top-scored experts E s l E_{s}^{l} whose cumulative relevance score exceeds the threshold p p , satisfying:
|
||||
∑ i ∈ E s l R i l ⩾ p , \sum_{i\in E_{s}^{l}}R_{i}^{l}\geqslant p,
|
||||
(8)
|
||||
where R i l R_{i}^{l} is the relevance score (either r i l r_{i}^{l} or g i l g_{i}^{l} ) of expert i i in layer l l . During training and inference, tokens can be assigned to any expert. However, only the selected experts E s l E_{s}^{l} in each layer can be updated; other experts and modules remain frozen.
|
||||
## 4 Experiment Setup
|
||||
## 4.1 Main Evaluation
|
||||
We evaluate our ESFT method on two common LLM customization scenarios: (1) improving the model’s specific ability in a domain where the model may already have decent performance; (2) adapting the model to a possibly narrow but unfamiliar specialized task .
|
||||
## 4.1.1 Tasks for Model Enhancement
|
||||
We choose two domain-specific tasks, i.e., Math and Code, to evaluate how our method can enhance the model’s existing abilities. The two domains are widely concerned in current LLM research and suitable for evaluation, as many pre-trained models can perform decently, while there is significant potential for improvement through further training. We assess our method’s effectiveness through performance gains.
|
||||
For the Math domain, we use MetaMathQA Yu et al. (2023) for training and use GSM8K Cobbe et al. (2021) and MATH Hendrycks et al. (2021a) for evaluation. For the Code domain, We train the model on the Python subset of the enormous evol-codealpaca dataset Luo et al. (2023) to simulate a more concentrated LLM customization scenario, and assess its performance on HumanEval Chen et al. (2021) and MBPP Austin et al. (2021) .
|
||||
Math Ability
|
||||
Code Ability
|
||||
Specialized Tasks
|
||||
MATH
|
||||
GSM8K
|
||||
Humaneval
|
||||
MBPP
|
||||
Intent
|
||||
Summary
|
||||
Law
|
||||
Translation
|
||||
Average
|
||||
Vanilla Model
|
||||
19.6
|
||||
55.9
|
||||
42.1
|
||||
44.6
|
||||
16.8
|
||||
58.6
|
||||
17.1
|
||||
14.5
|
||||
33.6
|
||||
FFT
|
||||
23.4
|
||||
66.4
|
||||
42.1
|
||||
42.2
|
||||
78.8
|
||||
69.4
|
||||
47.0
|
||||
38.4
|
||||
51.0
|
||||
LoRA
|
||||
20.6
|
||||
58.9
|
||||
39.6
|
||||
44.8
|
||||
67.8
|
||||
64.7
|
||||
39.7
|
||||
23.1
|
||||
44.9
|
||||
ESFT-Token (Ours)
|
||||
22.6
|
||||
66.0
|
||||
41.5
|
||||
42.6
|
||||
75.6
|
||||
65.4
|
||||
45.7
|
||||
36.2
|
||||
49.4
|
||||
ESFT-Gate (Ours)
|
||||
23.2
|
||||
64.9
|
||||
43.3
|
||||
41.8
|
||||
78.6
|
||||
65.8
|
||||
49.1
|
||||
35.2
|
||||
50.2
|
||||
Table 1: Main performance comparison across methods and tasks. Best or near-best results are shown in bold and second-best results are underlined . Our method ESFT provides a strong balance of performance across diverse tasks, rivaling FFT and surpassing LoRA, particularly in specialized task domains.
|
||||
## 4.1.2 Tasks for Model Adaptation
|
||||
We select four specialized tasks to evaluate how our method can facilitate language models to adapt to an unfamiliar downstream task, covering a diverse range of abilities that most models can excel at after training but not without training:
|
||||
(1) Text-to-JSON Intent Recognition in the BDCI-21 Smart HCI NLU Challenge 1 1
|
||||
1
|
||||
https://www.datafountain.cn/competitions/511 , which requires converting text instructions into JSON format for home appliances.
|
||||
(2) Text Summarization in the BDCI-21 Summarization Challenge 2 2
|
||||
2
|
||||
https://www.datafountain.cn/competitions/536 , which summarizes customer service call transcripts.
|
||||
(3) Legal judgment Prediction in the the BDCI-21 Law Event Prediction Challenge 3 3
|
||||
3
|
||||
https://www.datafountain.cn/competitions/540 , where the “case description” and “judgment” are repurposed as a legal judgment prediction task. (4) Low-resource Translation in the ChrEn dataset ( Zhang et al., 2020 ), translating the minority Cherokee to English. Examples of the tasks are shown in Appendix A .
|
||||
To measure model performance, for the text-to-JSON task, we calculate the exact match between model output and reference answer; for other tasks, we employ GPT-4 to score model output between 0 and 10 given reference answer 4 4
|
||||
4
|
||||
The exact version we use is gpt-4-1106-preview . The evaluation instructions are in Appendix G . . All evaluations use few-shot examples.
|
||||
## 4.2 General Ability Evaluation
|
||||
We select a broad range of benchmarks to evaluate the extent to which the models’ general abilities are preserved after training on new tasks.
|
||||
These benchmarks include MMLU Hendrycks et al. (2021b) , TriviaQA Joshi et al. (2017) , HellaSwag Zellers et al. (2019) , ARC-Challenge Clark et al. (2018) , IFEval Zhou et al. (2023) , CEval Huang et al. (2023) , and CLUEWSC Xu et al. (2020) , covering comprehensive model ability evaluations across various domains including natural language understanding, question answering, instruction following, and common sense reasoning.
|
||||
CLUEWSC
|
||||
TriviaQA
|
||||
IFEval
|
||||
MMLU
|
||||
CEval
|
||||
HellaSwag
|
||||
ARC
|
||||
Average
|
||||
Vanilla Model
|
||||
81.5
|
||||
67.7
|
||||
42.5
|
||||
57.5
|
||||
59.9
|
||||
74.0
|
||||
53.7
|
||||
62.4
|
||||
FFT
|
||||
80.9 ± \pm 1.1
|
||||
65.9 ± \pm 0.7
|
||||
34.2 ± \pm 4.1
|
||||
55.5 ± \pm 1.0
|
||||
58.8 ± \pm 0.9
|
||||
67.9 ± \pm 3.8
|
||||
48.4 ± \pm 2.4
|
||||
58.8 ± \pm 1.3
|
||||
LoRA
|
||||
74.3 ± \pm 7.7
|
||||
63.4 ± \pm 5.4
|
||||
38.7 ± \pm 2.5
|
||||
55.5 ± \pm 1.2
|
||||
57.0 ± \pm 1.5
|
||||
72.8 ± \pm 1.9
|
||||
51.8 ± \pm 2.3
|
||||
59.1 ± \pm 2.5
|
||||
ESFT-Token
|
||||
80.9 ± \pm 0.9
|
||||
66.7 ± \pm 1.8
|
||||
40.7 ± \pm 1.3
|
||||
57.1 ± \pm 0.5
|
||||
59.6 ± \pm 0.8
|
||||
72.3 ± \pm 3.6
|
||||
52.9 ± \pm 1.5
|
||||
61.5 ± \pm 1.1
|
||||
ESFT-Gate
|
||||
81.4 ± \pm 1.1
|
||||
66.5 ± \pm 2.3
|
||||
40.2 ± \pm 1.5
|
||||
57.0 ± \pm 0.4
|
||||
59.5 ± \pm 0.8
|
||||
68.2 ± \pm 9.9
|
||||
51.5 ± \pm 3.1
|
||||
60.6 ± \pm 2.3
|
||||
Table 2: General ability performance comparison across methods and tasks. The performance for a task is averaged across all training experiments, followed by the standard deviation across tasks. Best or near-best results are shown in bold . Our method ESFT consistently achieves good performance among all tasks.
|
||||
## 4.3 Backbone Model and Training Settings
|
||||
We use the backbone architecture of DeepSeek-V2-Lite DeepSeek (2024) for all experiments.
|
||||
The model includes a fine-grained set of 66 experts for each transformer layer. This makes it uniquely suitable at the time of this study for our method, which benefits from expert specialization.
|
||||
We train the model on a carefully curated alignment dataset that excludes math and code data and take the resulting checkpoint as our vanilla model for subsequent experiments. This alignment phase can activate model ability across various domains while keeping Math/Code ability as elementary to better verify the performance gains of our method in these two fields.
|
||||
We adopt two baselines: Full-Parameter Fine-Tuning (FFT) and Low-Rank Adaptation (LoRA, Hu et al., 2021 ). For LoRA, we add low-rank matrices to all parameters for training except token embeddings and the language modeling head.
|
||||
We maintain a 1:1 ratio for task-specific data and alignment data for all methods, which we find is highly effective in preserving general abilities obtained from the alignment phase for FFT and LoRA.
|
||||
However, for our ESFT method, not adopting this data mixing strategy may even better maintain general ability. We detail this in Appendix F .
|
||||
All experiments are done on the HFAI cluster 5 5
|
||||
5
|
||||
https://doc.hfai.high-flyer.cn/index.html with 2 nodes of 8x Nvidia A100 PCIe GPUs.
|
||||
For hyperparameter settings, all methods use a batch size of 32 and a sequence length of 4096 for training. For every task, we set the maximum steps of training to 500, and evaluate the model every 100 steps. The learning rates are set to 3e-5, 1e-4, and 1e-5 for FFT, LoRA, and ESFT, respectively, based on a hyperparameter search in {1e-5, 3e-5, 1e-4, 3e-4}. The LoRA rank is set to 8 and scaling is set to 2, following Hu et al. (2021) . The threshold p p is set to 0.1 for ESFT-Gate and 0.2 for ESFT-Token, respectively. § 6.2 shows how we determine the threshold for ESFT.
|
||||
## 5 Results
|
||||
## 5.1 Benchmark Performance Results
|
||||
The results in Table 1 and Table 2 demonstrate several conclusions. All methods can improve model performance in customization tasks compared to the vanilla model, while they may cause a performance decrease in general tasks. Generally, the performance increase is higher in model adaptation tasks than in model enhancement tasks.
|
||||
For customization ability evaluation, ESFT surpasses LoRA significantly and is competitive with FFT. As shown in Table 1 , ESFT-Token and ESFT-Gate achieve near-best results in model enhancement tasks like Math, and ESFT-Gate achieves the best performance in the Humaneval task. ESFT also excels in model adaptation tasks, with ESFT-Gate achieving near-best performance in 3 tasks out of 4. Notably, ESFT-Gate’s average of 50.2 is competitive compared to FFT’s 51.0, slightly better than ESFT-Token’s 49.4, and significantly surpasses LoRA’s 44.9. This demonstrates that finding task-relevant experts can efficiently adapt the model for efficient customization.
|
||||
For general ability evaluation, ESFT consistently outperforms FFT and LoRA by showing less performance degradation. As illustrated in Table 2 , ESFT-token performs better than ESFT-gate, with average scores of 61.5 and 60.6, respectively. The results demonstrate a wide range of retention in tasks such as TriviaQA and IFEval, surpassing FFT’s 58.8 and LoRA’s 59.1. Both methods retain performance better than LoRA and FFT, highlighting their effectiveness in maintaining general task performance 6 6
|
||||
6
|
||||
We further investigate Math and Code performance of the models trained on specialized tasks in Appendix H . FFT and LoRA exhibit even more severe degradation, while ESFT shows a minimal performance drop.
|
||||
. Analyses in § 6.3 indicate that
|
||||
such degradation on general tasks for FFT and LoRA may result from training shared parameters.
|
||||
## 5.2 Computational Efficiency Results
|
||||
The results in Figure 6 demonstrates that ESFT exhibits several advantages in terms of training time and storage space requirements:
|
||||
Figure 4: Number of experts trained in ESFT across layers and tasks. Earlier computed layers are numbered smaller. Most tasks and layers train 5-15% of experts, demonstrating ESFT’s effectiveness in selecting task-related experts.
|
||||
Figure 5: Computational efficiency results. Blue bars show the training time and green lines show storage space. ESFT performs efficiently in terms of training time and storage space.
|
||||
## Training Time
|
||||
The average training time for ESFT-Token and ESFT-Gate is 19.8 minutes and 20.9 minutes, respectively. The FFT method takes significantly longer at 28.5 minutes. Although LoRA achieves a shorter training time of 16.5 minutes, our methods are relatively close.
|
||||
## Storage Space
|
||||
The average storage space of parameters trained is 2.57 GB for ESFT-Token and 3.20 GB for ESFT-Gate, while FFT demands a substantial 28.6 GB.
|
||||
Although LoRA requires less storage,
|
||||
ESFT performs significantly better than LoRA in downstream task performance.
|
||||
In summary, ESFT demonstrates excellent performance in training time and storage space, significantly outperforming FFT.
|
||||
Furthermore, as shown in Table 3 , ESFT requires much fewer trainable parameters compared to FFT, resulting in lower GPU memory usage.
|
||||
These advantages show that ESFT is efficient and effective for language model customization and adaptation.
|
||||
## 6 Analysis
|
||||
Non-shared Experts
|
||||
Shared Experts
|
||||
Non-expert Parameters
|
||||
Trainable Parameters
|
||||
Specialized Ability
|
||||
General Ability
|
||||
Average
|
||||
ALL
|
||||
✓ \checkmark
|
||||
✓ \checkmark
|
||||
15.7B
|
||||
51.0
|
||||
58.8
|
||||
54.9
|
||||
Relevant
|
||||
✓ \checkmark
|
||||
× \times
|
||||
1.85B
|
||||
49.8
|
||||
60.7
|
||||
55.3
|
||||
Relevant
|
||||
× \times
|
||||
× \times
|
||||
1.4B
|
||||
49.4
|
||||
61.5
|
||||
55.4
|
||||
× \times
|
||||
✓ \checkmark
|
||||
× \times
|
||||
450M
|
||||
47.4
|
||||
61.2
|
||||
54.3
|
||||
× \times
|
||||
✓ \checkmark
|
||||
✓ \checkmark
|
||||
1.3B
|
||||
49.0
|
||||
60.0
|
||||
54.5
|
||||
Relevant
|
||||
✓ \checkmark
|
||||
✓ \checkmark
|
||||
2.7B
|
||||
50.8
|
||||
60.3
|
||||
55.6
|
||||
× \times
|
||||
× \times
|
||||
× \times
|
||||
-
|
||||
33.8
|
||||
62.4
|
||||
48.1
|
||||
Table 3: Comparisons of different model configs based on whether training shared or non-shared parameters. Results include trainable parameters and performance of specialized and general abilities. The best or near-best results excluding the non-training setting are shown in bold .
|
||||
Figure 6: Comparison of three methods under different training efficiency settings on the Math task. The x-axis shows the average trainable experts per layer for ESFT and rank for LoRA, indicating the ratio of trained parameters. The y-axis represents specialized and general ability. Markers on the lines indicate p p or rank values. ESFT consistently outperforms LoRA in both specialized and general ability.
|
||||
In this section, we investigate the expert selection process of ESFT in § 6.1 , and demonstrate the performance of ESFT and LoRA under different computational constraints in § 6.2 . We analyze the effects of training shared and non-shared parameters in § 6.3 , and conduct ablation studies in § 6.4 to verify the importance of our expert relevance scores and model structure of fine-grained experts.
|
||||
## 6.1 ESFT Leverages Specialized Experts Effectively
|
||||
We analyze the number of experts ESFT trains across tasks and layers to understand its expert selection process. Results are shown in Figure 4 .
|
||||
From the results, we have several observations: (1) The average number of experts used per task across layers ranges from 2 to 15 out of 66, indicating ESFT can have 75%-95% fewer trainable parameters than FFT. (2) ESFT-Token generally employs fewer experts while better maintaining general performance, comparable to ESFT-Gate in tasks like Math, Intent, and Law. (3) The number of experts varies by task, with more specialized tasks like Math and Translation using fewer experts; our method’s performances for these tasks exceed LoRA to the greatest extent, indicating that our method is especially suitable for more specialized tasks. (4) For most tasks, few experts are chosen in the middle layers, indicating that expert distribution is more concentrated in these layers.
|
||||
## 6.2 ESFT Leverages Training Resources Efficiently
|
||||
Math Ability
|
||||
Code Ability
|
||||
Specialized Tasks
|
||||
MATH
|
||||
GSM8K
|
||||
Humaneval
|
||||
MBPP
|
||||
Intent
|
||||
Summary
|
||||
Law
|
||||
Translation
|
||||
Average
|
||||
ESFT-Token
|
||||
22.6
|
||||
66.0
|
||||
41.5
|
||||
42.6
|
||||
75.6
|
||||
65.4
|
||||
45.7
|
||||
36.2
|
||||
49.4
|
||||
Δ \Delta of rand
|
||||
-1.0
|
||||
-3.7
|
||||
-2.5
|
||||
0.2
|
||||
-2.6
|
||||
-1.7
|
||||
1.3
|
||||
-13.5
|
||||
-2.8
|
||||
ESFT-Gate
|
||||
23.2
|
||||
64.9
|
||||
43.3
|
||||
41.8
|
||||
78.6
|
||||
65.8
|
||||
49.1
|
||||
35.2
|
||||
50.2
|
||||
Δ \Delta of rand
|
||||
-1.7
|
||||
-3.2
|
||||
-4.3
|
||||
1.6
|
||||
-5.0
|
||||
0.3
|
||||
-2.9
|
||||
-20.4
|
||||
-4.4
|
||||
Table 4: Performance comparison between original experts and random experts. Replacing high-affinity experts with random ones significantly harms model performance across different tasks.
|
||||
Both ESFT and LoRA have a training efficiency hyperparameter ( p p for ESFT and rank for LoRA). Increasing its value would raise computational resource usage and potentially improve performance. To understand how ESFT and LoRA perform under different efficiency settings, we evaluate benchmark performance on the Math task. We set rank ⩽ \leqslant 512 for LoRA as a higher value will result in more trainable parameters than FFT.
|
||||
Figure 6 illustrates both specialized and general ability under different training efficiency settings.
|
||||
From the results, we can conclude: (1) All three methods show a trade-off between training efficiency and performance. Increasing trained parameters ( p p for ESFT and rank for LoRA) before a certain point can improve performance. (2) Both ESFT-Token and ESFT-Gate outperform LoRA at any point, demonstrating higher specialized ability and more stable general ability. (3) ESFT-Token peaks in both specialized and general ability at p p =0.5, while ESFT-Gate peaks at p p =0.3 for specialized and p p =0.1 for general ability. (4) ESFT-Token and ESFT-Gate performance saturates at p p =0.2 and p p =0.1, respectively, indicating that most expert choices may be less relevant to task performance. We delve deeper into this in Appendix E .
|
||||
## 6.3 Selectively Training Non-Shared Parameters is the Key to ESFT
|
||||
In our proposed ESFT method, we only fine-tune a subset of non-shared experts. This section provides detailed discussions of several variants of our method that may also train shared parameters. The variables are based on:
|
||||
•
|
||||
Whether all non-shared experts or a task-relevant subset of them (we use the Token Selection Ratio and set p p =0.2) are trained.
|
||||
•
|
||||
Whether shared experts are trained.
|
||||
•
|
||||
Whether other parameters, including gates, attention layers, and embeddings, are trained.
|
||||
The results are shown in Table 3 . We report average trainable parameters across all tasks, performance of specialized and general abilities, and their average. Detailed numbers for all benchmarks are shown in Appendix D .
|
||||
From the results, we can draw several conclusions:
|
||||
Specialized performance increases as trainable parameters increase. The rank of trainable parameters from 450M to 15.7B highly aligns with the rank of specialized ability from 47.4 to 51.0. This suggests that increasing trainable parameters is effective in enhancing specialized performance.
|
||||
General performance decreases as trainable shared parameters increase. Whether relevant non-shared experts are trained or not, general performance decreases from 61.5 to 60.3, or from 62.4 to 60.0, respectively, as we train shared experts and/or non-expert parameters. As the complete set of non-shared experts is trained, general performance decreases further from 60.3 to 58.8. This suggests that training shared parameters is more likely to cause overfitting on downstream tasks and forgetting on general tasks compared to training non-shared parameters.
|
||||
It is highly prioritized to train task-relevant non-shared experts. Training relevant experts achieves at least 55.3, while other settings achieve at most 54.9, even with higher demands of up to 15.7B parameters. Therefore, fine-tuning these experts is highly prioritized for model customization.
|
||||
We propose two major training strategies based on these conclusions:
|
||||
1.
|
||||
Prioritize specialized ability: Train all shared parameters and task-relevant non-shared experts to maximize the enhancement of specialized performance.
|
||||
2.
|
||||
Balance specialized and general ability, and computational efficiency: Train only task-relevant non-shared experts to minimize parameter costs while maximizing the maintenance of general ability.
|
||||
Figure 7: Experiment results for grouped experts. As the experts become more coarse-grained, ESFT degrades more severely than FFT.
|
||||
## 6.4 Analysis of Key Modules in ESFT
|
||||
In this section, we analyze and demonstrate that the effectiveness of our method lies in two modules: (1) our proposed expert relevance score functions and (2) the fine-grained expert segmentation of the MoE model architecture.
|
||||
## Expert Relevance Score Function
|
||||
In this work, we propose Average Gate Score and Token Selection Ratio as expert relevance score functions to filter relevant experts for different tasks. To demonstrate their effectiveness, we replace the experts obtained from these functions with random experts while keeping the number of activated experts per layer the same. Results in Table 4 show that replacing relevant experts with random ones significantly decreases task performance, demonstrating the effectiveness of our proposed relevance scores.
|
||||
## Fine-Grained Expert Segmentation of the MoE Model
|
||||
We use the fine-grained segmented DeepSeek-V2 model as our backbone. To demonstrate t the effectiveness of this fine-grained segmentation, we use greedy search (as detailed in Appendix B ) to group experts, simulating coarse-grained segmentation. Experts in the same group share the average affinity score. We maintain the computational cost by selecting a constant 1/8 of experts for each token. Experiment results of the Math domain in Figure 7 show that as the group size increases, our method’s performance decreases more severely than FFT, while the training cost (i.e., trainable experts) rises. These findings indicate that our method, and even effective LLM customization, highly rely on a fine-grained segmented LLM architecture with more specialized experts.
|
||||
## 7 Conclusion
|
||||
In this work, we study parameter-efficient fine-tuning methods for sparse large language models with the Mixture of Experts (MoE) architecture. We first observe that tasks from different domains are handled by distinct combinations of experts. We then propose selecting the most relevant experts for downstream tasks using two metrics: average gate score and token selection ratio. Experimental results show that our method significantly reduces training costs while matching or surpassing full parameter fine-tuning results. Further analysis confirms that our method enhances the specialization of the expert system within the MoE architecture.
|
||||
## Limitations
|
||||
Firstly, due to the limitation of the availability of other fine-grained MoE models, our method was only tested on the DeepSeek-V2-Lite MoE model. The conclusions drawn from this model require further validation when applied to other contexts.
|
||||
Besides, due to the lack of parameter-wise and structurally aligned MoE models with different expert granularities, we used a simulation approach by binding several groups of experts to compare coarse-grained and fine-grained MoE methods.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,667 @@
|
||||
TITLE: DeepSeek-Prover-V1.5: Harnessing Proof Assistant Feedback for Reinforcement Learning and Monte-Carlo Tree Search
|
||||
ARXIV: 2408.08152
|
||||
AUTHORS: Huajian Xin , Z.Z. Ren , Junxiao Song , Zhihong Shao , Wanjia Zhao , Haocheng Wang , Bo Liu , Liyue Zhang , Xuan Lu , Qiushi Du , Wenjun Gao , Qihao Zhu , Dejian Yang , Zhibin Gou , Z.F. Wu , Fuli Luo , Chong Ruan
|
||||
DATE: [Submitted on 15 Aug 2024]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-Prover-V1.5: Harnessing Proof Assistant Feedback
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 1.1 Contributions
|
||||
## 1.2 Summary of Evaluations and Metrics
|
||||
## 2 Model Training
|
||||
## 2.1 Pre-training
|
||||
## 2.2 Supervised Fine-tuning
|
||||
## Data Curation.
|
||||
## Thought-augmented Proof Generation.
|
||||
## Prompt Augmentation with Tactic State Information.
|
||||
## Training Setting.
|
||||
## 2.3 Reinforcement Learning from Proof Assistant Feedback
|
||||
## Prompts.
|
||||
## Rewards.
|
||||
## Reinforcement Learning Algorithm.
|
||||
## Training Setting.
|
||||
## 2.4 Evaluation
|
||||
## Benchmarks.
|
||||
## Prompting Configurations.
|
||||
## Metric.
|
||||
## Comparison across Training Stages.
|
||||
## Comparison between CoT and non-CoT.
|
||||
## 3 Exploration-oriented Monte-Carlo Tree Search
|
||||
## 3.1 Tactic-level Tree Abstraction
|
||||
## Truncate: Proof Decomposition into Tree Nodes.
|
||||
## Resume: Proof Generation from a Tree Node.
|
||||
## 3.2 Interactive Theorem Proving via Monte-Carlo Tree Search
|
||||
## Selection.
|
||||
## Expansion.
|
||||
## Backpropagation.
|
||||
## 3.3 Intrinsic Rewards for Monte-Carlo Tree Search
|
||||
## RMax applied to MCTS.
|
||||
## UCB for Non-stationary Rewards.
|
||||
## 3.4 Parallelization of Monte-Carlo Tree Search
|
||||
## 3.5 Comparison with Existing Methods
|
||||
## 4 Experimental Results
|
||||
## 4.1 Main Results
|
||||
## Metric.
|
||||
## Results on miniF2F.
|
||||
## Results on ProofNet.
|
||||
## 4.2 Re-Examining the Effectiveness of Training Strategies on Large-scale Sampling
|
||||
## General Enhancement of Reinforcement Learning.
|
||||
## CoT, non-CoT, and Mixture Strategy.
|
||||
## 4.3 Ablation Studies on RMaxTS
|
||||
## Intrinsic Rewards and Discounted UCB.
|
||||
## Guidance of Tactic State Information.
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-Prover-V1.5: Harnessing Proof Assistant Feedback for Reinforcement Learning and Monte-Carlo Tree Search
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
1.1 Contributions
|
||||
1.2 Summary of Evaluations and Metrics
|
||||
2 Model Training
|
||||
2.1 Pre-training
|
||||
2.2 Supervised Fine-tuning
|
||||
Data Curation.
|
||||
Thought-augmented Proof Generation.
|
||||
Prompt Augmentation with Tactic State Information.
|
||||
Training Setting.
|
||||
2.3 Reinforcement Learning from Proof Assistant Feedback
|
||||
Prompts.
|
||||
Rewards.
|
||||
Reinforcement Learning Algorithm.
|
||||
Training Setting.
|
||||
2.4 Evaluation
|
||||
Benchmarks.
|
||||
Prompting Configurations.
|
||||
Metric.
|
||||
Comparison across Training Stages.
|
||||
Comparison between CoT and non-CoT.
|
||||
3 Exploration-oriented Monte-Carlo Tree Search
|
||||
3.1 Tactic-level Tree Abstraction
|
||||
Truncate: Proof Decomposition into Tree Nodes.
|
||||
Resume: Proof Generation from a Tree Node.
|
||||
3.2 Interactive Theorem Proving via Monte-Carlo Tree Search
|
||||
Selection.
|
||||
Expansion.
|
||||
Backpropagation.
|
||||
3.3 Intrinsic Rewards for Monte-Carlo Tree Search
|
||||
RMax applied to MCTS.
|
||||
UCB for Non-stationary Rewards.
|
||||
3.4 Parallelization of Monte-Carlo Tree Search
|
||||
3.5 Comparison with Existing Methods
|
||||
4 Experimental Results
|
||||
4.1 Main Results
|
||||
Metric.
|
||||
Results on miniF2F.
|
||||
Results on ProofNet.
|
||||
4.2 Re-Examining the Effectiveness of Training Strategies on Large-scale Sampling
|
||||
General Enhancement of Reinforcement Learning.
|
||||
CoT, non-CoT, and Mixture Strategy.
|
||||
4.3 Ablation Studies on RMaxTS
|
||||
Intrinsic Rewards and Discounted UCB.
|
||||
Guidance of Tactic State Information.
|
||||
5 Conclusion, Limitation, and Future Work
|
||||
References
|
||||
A Illustrative Examples of Non-CoT and CoT Prompting for Proof Completion
|
||||
B Example Solutions to MiniF2F-test Problems
|
||||
B.1 Problems that non-CoT outperforms CoT
|
||||
B.2 Problems that CoT outperforms non-CoT
|
||||
License: CC BY 4.0
|
||||
arXiv:2408.08152v1 [cs.CL] 15 Aug 2024
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-Prover-V1.5: Harnessing Proof Assistant Feedback
|
||||
for Reinforcement Learning and Monte-Carlo Tree Search
|
||||
Huajian Xin*, Z.Z. Ren*, Junxiao Song*, Zhihong Shao*, Wanjia Zhao, Haocheng Wang, Bo Liu, Liyue Zhang
|
||||
Xuan Lu, Qiushi Du, Wenjun Gao, Qihao Zhu, Dejian Yang, Zhibin Gou, Z.F. Wu,
|
||||
Fuli Luo, Chong Ruan
|
||||
DeepSeek-AI
|
||||
https://github.com/deepseek-ai/DeepSeek-Prover-V1.5
|
||||
Corresponding author: Core contributors
|
||||
## Abstract
|
||||
We introduce DeepSeek-Prover-V1.5, an open-source language model designed for theorem proving in Lean 4, which enhances DeepSeek-Prover-V1 by optimizing both training and inference processes. Pre-trained on DeepSeekMath-Base with specialization in formal mathematical languages, the model undergoes supervised fine-tuning using an enhanced formal theorem proving dataset derived from DeepSeek-Prover-V1. Further refinement is achieved through reinforcement learning from proof assistant feedback (RLPAF). Beyond the single-pass whole-proof generation approach of DeepSeek-Prover-V1, we propose RMaxTS, a variant of Monte-Carlo tree search that employs an intrinsic-reward-driven exploration strategy to generate diverse proof paths. DeepSeek-Prover-V1.5 demonstrates significant improvements over DeepSeek-Prover-V1, achieving new state-of-the-art results on the test set of the high school level miniF2F benchmark (63.5%) and the undergraduate level ProofNet benchmark (25.3%).
|
||||
Figure 1: Pass rates of models on formal theorem proving benchmarks in Lean 4: the high school level miniF2F-test benchmark ( Zheng et al., 2022 ) and the undergraduate level ProofNet benchmark ( Azerbayev et al., 2023 ) . We compare both the pre-trained and fine-tuned versions of DeepSeek-Prover-V1.5 with strong baselines.
|
||||
## 1 Introduction
|
||||
Recent advancements in large language models have significantly influenced mathematical reasoning and theorem proving in artificial intelligence. Despite notable progress in natural language domains, language models still encounter substantial challenges in formal theorem proving, e.g. using Lean ( Moura and Ullrich, 2021 ) and Isabelle ( Paulson, 1994 ) , which requires rigorous derivations satisfying formal specifications of the verification system.
|
||||
Even advanced models like GPT-4 ( OpenAI, 2023 ) struggle with complex formal proofs, underscoring the intricate nature of both the coding and the mathematics involved.
|
||||
A formal theorem proving model must not only grasp the syntax and semantics of formal systems like the Lean theorem prover but also align abstract mathematical reasoning with precise formal representation.
|
||||
Language models in formal theorem proving typically employ two strategies: proof-step generation ( Polu and Sutskever, 2020 ; Jiang et al., 2022a ; Lample et al., 2022 ; Yang et al., 2023 ; Wu et al., 2024 ) and whole-proof generation ( Jiang et al., 2022b ; Zhao et al., 2023 ; Wang et al., 2023a ) . Proof-step generation predicts each subsequent tactic and verifies it using the formal verifier to obtain updated information about the current tactic state, often utilizing tree search techniques to construct valid proofs. In contrast, whole-proof generation is computationally efficient, which produces an entire proof code based on the theorem statement, requiring less communication budget to coordinate between the prover model and the formal theorem verifier. While DeepSeek-Prover-V1 ( Xin et al., 2024 ) has achieved state-of-the-art results in Lean 4 with whole-proof generation, this paradigm presents its unique challenges. It requires long-horizon sequence prediction without access to intermediate tactic states, and future tactics depend on these hidden results. In Lean’s tactic mode, proofs are constructed through a sequence of tactics that transform the proof state. This sequential nature introduces the risk of compounding errors ( Ross et al., 2011 ) , where a single misinterpretation can lead to significant deviations from a valid proof path. More specifically, the auto-regressive model may have incorrect believes on intermediate tactic states when generating long proofs.
|
||||
To seamlessly integrate intermediate tactic states in proof-step generation while maintaining the simplicity and computational efficiency of whole-proof generation, we have developed a unified approach in DeepSeek-Prover-V1.5. This method combines the strengths of both proof-step and whole-proof generation techniques through a truncate-and-resume mechanism. The process begins with standard whole-proof generation, where the language model completes the proof code following the theorem statement prefix. The Lean prover then verifies this code. If the proof is correct and complete, the procedure terminates. If an error is detected, the code is truncated at the first error message, and any subsequent code is discarded. The successfully generated proof code is then used as a prompt for the generation of next proof segment. To enhance the accuracy of the model’s new completions, we append the latest state from the Lean 4 prover as a comment at the end of the prompt. Notably, our method is not restricted to resuming from the last successfully applied tactic. We integrate the truncate-and-resume mechanism into Monte-Carlo tree search ( Coulom, 2006 , MCTS;) in which the truncation points are scheduled by the tree search policy.
|
||||
In addition, we propose a novel reward-free exploration algorithm for MCTS to address the reward sparsity issue of proof search.
|
||||
We assign the tree search agent intrinsic motivation, a.k.a. curiosity ( Schmidhuber, 2010 ) , to extensively explore the tactic state space.
|
||||
These algorithmic modules extend the functionality of our whole-proof generation model to become a flexible tool for interactive theorem proving, which can effectively utilize the proof assistant feedback and generate diverse solution candidates.
|
||||
Figure 2: Overall Framework. DeepSeek-Prover-V1.5 is trained through pre-training, supervised fine-tuning, and reinforcement learning. During supervised fine-tuning, the pre-trained model receives an incomplete theorem proof ending with a tactic state comment keyword. The model is trained to predict the content of this tactic state (auxiliary objective) and complete the subsequent proof steps (main objective). In the reinforcement learning stage, given an incomplete theorem proof and ground-truth tactic state from the Lean prover, we roll out the fine-tuned model to generate multiple proof candidates, which are then verified by the Lean prover. The verification results for these candidates are used as binary (0-1) rewards to further optimize the model and enhance its alignment with the formal specifications of the verification system. For model inference, we offer two alternatives: single-pass sampling and Monte-Carlo tree search.
|
||||
## 1.1 Contributions
|
||||
We present a comprehensive framework for developing a language model-based formal mathematics prover, integrating several key components: large-scale mathematical pre-training, formal mathematics corpus construction and augmentation, online reinforcement learning from proof assistant feedback, and a tree search methodology for long-term planning in theorem proving. The pre-trained model, supervised fine-tuned model, and reinforcement learning model, along with the code for the Monte-Carlo tree search algorithm, are publicly available for further research and application.
|
||||
•
|
||||
Pre-Training : We enhance our base model’s capabilities in formal theorem proving and mathematical reasoning by further pre-training on high-quality mathematics and code data, with a focus on formal languages such as Lean, Isabelle, and Metamath.
|
||||
•
|
||||
Supervised Fine-Tuning : We improve the Lean 4 code completion dataset by implementing two data augmentation techniques. First, we use DeepSeek-Coder V2 236B ( Zhu et al., 2024 ) to annotate natural language chain-of-thought comments alongside Lean 4 code, aligning formal theorem proving with natural language reasoning. Second, we insert intermediate tactic state information within the Lean 4 proof code, enabling our model to leverage compiler feedback effectively. The resulting dataset is then used to fine-tune the pre-trained model.
|
||||
•
|
||||
Reinforcement Learning : We employ the GRPO algorithm ( Shao et al., 2024 ) to perform reinforcement learning from proof assistant feedback (RLPAF) on the supervised fine-tuned model. Verification results from the Lean prover serve as reward supervision, enhancing the model’s alignment with the formal specifications of the verification system.
|
||||
•
|
||||
Monte-Carlo Tree Search : We advance the tree search method in formal theorem proving by introducing a novel abstraction and a corresponding search algorithm. Our truncate-and-resume mechanism acts as a state-action abstraction, seamlessly integrating the tree search process into the whole-proof generation framework. We present RMaxTS, an innovative Monte-Carlo tree search algorithm that leverages the RMax ( Brafman and Tennenholtz, 2002 ) strategy to tackle exploration challenges in sparse-reward proof search problems. By assigning intrinsic rewards, this algorithm encourages the prover agent to generate diverse planning paths, thereby fostering extensive exploration of the proof space.
|
||||
## 1.2 Summary of Evaluations and Metrics
|
||||
•
|
||||
miniF2F : In the single-pass whole-proof generation setting, DeepSeek-Prover-V1.5 achieved a pass rate of 60.2% on the test set of miniF2F, marking a significant improvement of absolute 10.2 percentage points over DeepSeek-Prover-V1’s 50.0%. Incorporating tree search techniques further elevated the pass rate to a new state-of-the-art 63.5% .
|
||||
•
|
||||
ProofNet : DeepSeek-Prover-V1.5 also demonstrated strong performance in the single-pass whole-proof generation setting for ProofNet, with pass rates of 21.6% on the validation set and 23.7% on the test set. The integration of tree search techniques further enhanced these results, achieving new state-of-the-art pass rates of 25.4% on the validation set and 25.3% on the test set.
|
||||
## 2 Model Training
|
||||
## 2.1 Pre-training
|
||||
To enhance our language model’s proficiency in generating formal proofs and reasoning through mathematical language, we further pre-train our base model ( Shao et al., 2024 ) . This refinement involved training on high-quality datasets that include both code and natural language mathematical content. We specifically focused on formal languages widely used in proof assistants, such as Lean, Isabelle, and Metamath. We designate this improved model as DeepSeek-Prover-V1.5-Base.
|
||||
## 2.2 Supervised Fine-tuning
|
||||
In this section, we explore the methodology and processes involved in the supervised fine-tuning (SFT) of DeepSeek-Prover-V1.5. Specifically, we augment the proof dataset from DeepSeek-Prover-V1 by adding detailed explanatory comments. This enhancement aims to improve the alignment between natural language descriptions and Lean 4 code, thereby facilitating better formal mathematical reasoning. Additionally, we incorporate intermediate tactic state information as an auxiliary prediction task to support the truncate-and-resume mechanism used in the Monte-Carlo Tree Search process. We refer to the resulting model as DeepSeek-Prover-V1.5-SFT.
|
||||
## Data Curation.
|
||||
We develop a comprehensive Lean 4 code completion dataset for the supervised fine-tuning. This dataset includes synthetic proof code derived from a wide range of formal theorems. These theorems are sourced from various projects, such as the standard Lean 4 math library Mathlib4 ( Mathlib Community, 2020 ) , synthetic theorems from DeepSeek-Prover-V1 ( Xin et al., 2024 ) and Lean Workbook ( Ying et al., 2024a ) , and validation sets from the miniF2F ( Zheng et al., 2022 ) and ProofNet ( Azerbayev et al., 2023 ) benchmarks. To augment the formal proof data, we employed an expert iteration process ( Polu and Sutskever, 2020 ) . This involves generating proofs using the language model, verifying the generated proof data, retraining the model with the verified data, and then using the optimized model to generate additional proof data. Between each iteration, we use DeepSeek-Coder V2 236B ( Zhu et al., 2024 ) to annotate the thought process before the proof code as comments. Finally, we tailor these data for the truncate-and-resume mechanism for Monte-Carlo Tree Search (details in Section 3.1 ). The resulting proof dataset consists of 9,645k sequences.
|
||||
## Thought-augmented Proof Generation.
|
||||
In DeepSeek-Prover-V1, we identified a significant gap between problem-solving strategies in natural language and theorem proving in Lean. In natural language, models generate detailed deduction steps to construct proofs, whereas in Lean, they often rely on a sequence of high-level tactic calls to brute-force solutions. These high-level tactics, while effective, obscure their internal workings and outcomes, hindering the model’s ability to resolve complex proof goals with structured mathematical reasoning. To address this issue, we develop an approach that incorporates natural language reasoning before generating theorem proof code. Similar to Lean-STaR ( Lin et al., 2024 ) , which performs isolated chain-of-thought reasoning ( Wei et al., 2022 ; Feng et al., 2023 ) before each proof step, our method integrates this reasoning directly as comments within the proof code. We use the DeepSeek-Coder V2 236B ( Zhu et al., 2024 ) to enhance existing data in DeepSeek-Prover-V1 in two ways: first, by inserting a complete natural language solution at the beginning of the proof block, and second, by alternately inserting specific natural language steps for corresponding Lean tactics. Training the model with this data format enforces it to propose complete mathematical reasoning at the beginning of the proof block and detailed step planning before each tactic. This approach successfully develops new behaviors, employing delicate mathematical thinking to guide the generation of tactics. In the training data, two distinct guiding prompts are used to differentiate between the CoT (Chain of Thought) mode and the non-CoT mode for proof code completion. Examples of input and output in both modes can be found in Appendix A .
|
||||
## Prompt Augmentation with Tactic State Information.
|
||||
To implement the truncate-and-resume mechanism for Monte-Carlo Tree Search, we needed to extract tactic information from the code generated by the model. We enhanced the Lean REPL ( Leanprover Community, 2023 , Read-Eval-Print Loop;) with data extraction tools from the LeanDojo ( Yang et al., 2023 ) project. This allowed us to extract tactic information in triples, which include the position of each tactic, as well as the tactic states before and after its application. This information helps us identify the specific tactic code that triggers verification errors (used in the expansion step for tree search, see Section 3.2 ). For each tactic in a generated valid formal proof, we insert the tactic state returned by the verifier as a comment "/- tactic state: … -/". During training, we use all tokens following "/- tactic state: " as responses to calculate the supervised fine-tuning loss, while the tokens before this comment is used as prompts and do not contribute to the training loss calculation.
|
||||
## Training Setting.
|
||||
We conduct supervised fine-tuning based on the pre-trained model and train for 9B tokens, using a batch size of 2,048 and a constant learning rate of 1e-4. The training process begins with 100 warm-up steps to stabilize the learning dynamics. Training examples are randomly concatenated to form sequences, with a maximum context length of 4,096 tokens.
|
||||
## 2.3 Reinforcement Learning from Proof Assistant Feedback
|
||||
Reinforcement learning (RL) has been proven effective in enhancing the mathematical reasoning capabilities of supervised fine-tuned language models ( Shao et al., 2024 ) . To further advance DeepSeek-Prover-V1.5-SFT, we incorporate a reinforcement learning phase, resulting in the model DeepSeek-Prover-V1.5-RL. This phase leverages RL to enhance performance based on verification feedback from the Lean 4 prover. The specifics of this RL process are detailed below.
|
||||
## Prompts.
|
||||
In the reinforcement learning stage, we use a subset of theorem statements from the supervised fine-tuning dataset as training prompts. We select theorems for which DeepSeek-Prover-V1.5-SFT has a moderate success rate in generating correct proofs upon multiple attempts. This ensures that the model has room for improvement while still being able to receive positive feedback. After filtering, we retain approximately 4.5k unique theorem statements. Each theorem is prefixed with both CoT and non-CoT guiding prompts to enhance the model’s proof generation capabilities in both modes.
|
||||
## Rewards.
|
||||
When training LLMs via RL, a trained reward model typically provides feedback signals. In contrast, formal theorem proving benefits from the rigorous verification of generated proofs by proof assistants, offering a significant advantage. Specifically, each generated proof receives a reward of 1 if verified as correct, and 0 otherwise. While this binary reward signal is accurate, it is also sparse, especially for theorems that are challenging for the supervised fine-tuned model. To mitigate this sparsity, we select training prompts that are challenging yet achievable for the supervised fine-tuned model, as described above.
|
||||
## Reinforcement Learning Algorithm.
|
||||
We employ the Group Relative Policy Optimization ( Shao et al., 2024 , GRPO;) as our RL algorithm, which has demonstrated superior effectiveness and efficiency compared to PPO ( Schulman et al., 2017 ) , primarily because it eliminates the necessity of training an additional critic model. Specifically, GRPO samples a group of candidate proofs for each theorem prompt and optimizes the model based on the relative rewards of the outputs within the group. Our prompt selection strategy is designed to likely include both correct and incorrect proofs among the candidates, aligning well with the group-relative nature of GRPO and thereby enhancing the training process.
|
||||
## Training Setting.
|
||||
We conduct RL training based on the SFT model, which serves as both the initial model and the reference model for imposing the Kullback-Leibler (KL) divergence penalty. We use a constant learning rate of 5e-6, and the KL penalty coefficient is set to 0.02. For each theorem, we sample a group of 32 candidate proofs, with maximum length set to 2,048. The training batch size is configured to 512.
|
||||
## 2.4 Evaluation
|
||||
## Benchmarks.
|
||||
We evaluate theorem-proving performance on the following benchmarks to compare model capabilities after each training stage:
|
||||
•
|
||||
MiniF2F ( Zheng et al., 2022 ) focuses on formal problem-solving skills for high-school level exercises and competitions, such as AMC, AIME, and IMO, with an emphasis on algebra and number theory. The benchmark includes 244 validation and 244 test problems, originally in Lean 3 and manually converted to Lean 4.9.0, based on the version provided by Yang (2023) .
|
||||
•
|
||||
ProofNet ( Azerbayev et al., 2023 ) evaluates formal theorem-proving capabilities at the undergraduate level in mathematics. It comprises 185 validation and 186 test problems from widely-used undergraduate textbooks, covering real and complex analysis, linear algebra, abstract algebra, and topology. These problems were initially in Lean 3 and manually converted to Lean 4.9.0.
|
||||
Model
|
||||
Pass@128
|
||||
miniF2F-test
|
||||
ProofNet-test
|
||||
Base (3-shot)
|
||||
29.7 % ± 0.5 % 29.7\%\pm 0.5\%
|
||||
9.7 % ± 0.7 % 9.7\%\pm 0.7\%
|
||||
SFT (non-CoT)
|
||||
49.8 % ± 0.3 % 49.8\%\pm 0.3\%
|
||||
15.9 % ± 0.5 % 15.9\%\pm 0.5\%
|
||||
SFT (CoT)
|
||||
50.4 % ± 0.4 % 50.4\%\pm 0.4\%
|
||||
15.9 % ± 0.6 % 15.9\%\pm 0.6\%
|
||||
RL (non-CoT)
|
||||
50.5 % ± 0.6 % 50.5\%\pm 0.6\%
|
||||
17.5 % ± 0.5 % 17.5\%\pm 0.5\%
|
||||
RL (CoT)
|
||||
51.6 % ± 0.5 % 51.6\%\pm 0.5\%
|
||||
18.2 % ± 0.5 % 18.2\%\pm 0.5\%
|
||||
Figure 3: Comparison of model capabilities at different training stages. "CoT" and "non-CoT" refer to evaluations using two guiding prompts. The shaded region represents the range of standard deviations around the mean values. The notation μ ± σ \mu\pm\sigma indicates the average accuracy μ \mu and the standard deviation σ \sigma .
|
||||
## Prompting Configurations.
|
||||
For each proof attempt of DeepSeek-Prover-V1.5-Base, we independently sample three proof demonstrations from the validation set to construct the few-shot prompts. For the miniF2F benchmark, we use human-written proofs from Yang (2023) , while for the ProofNet benchmark, we use correct proofs generated by DeepSeek-Prover-V1.5-RL as few-shot demonstrations. For DeepSeek-Prover-V1.5-SFT and DeepSeek-Prover-V1.5-RL, we employ two types of guiding prompts: one that encourages chain-of-thought (CoT) reasoning before each proof step, and one that does not (non-CoT). Detailed examples are provided in Appendix A .
|
||||
## Metric.
|
||||
We evaluate theorem-proving performance using the pass@ K K accuracy metric, which measures the model’s success in generating a correct proof within K K attempts. Each model is deployed on a single A100-40G GPU, utilizing the vLLM framework ( Kwon et al., 2023 ) for sample generation. The sampling parameters are set with a temperature of 1, a top-p value of 0.95, and a maximum token limit of 2,048. The generated proofs are then verified using the Lean 4 theorem prover. For this verification, we import Mathlib4 ( Mathlib Community, 2020 ) and Aesop ( Limperg and From, 2023 ) to access predefined premises and tactics. The verification process is subject to a time limit of 300 seconds.
|
||||
## Comparison across Training Stages.
|
||||
Figure 3 presents a comparative analysis of each training stage on the miniF2F and ProofNet datasets. Our base model, DeepSeek-Prover-V1.5-Base, achieves a notable pass rate, solving nearly one-third of the problems on the test set of the miniF2F benchmark using 3-shot prompting. The supervised fine-tuning stage, resulting in DeepSeek-Prover-V1.5-SFT, significantly outperforms the base model, with Pass@128 accuracy increasing by approximately two-thirds on miniF2F and doubling on ProofNet. The subsequent reinforcement learning stage further enhances the model’s performance, improving Pass@ K K accuracy across all values of K K . In contrast to findings in natural language mathematics, such as those reported in DeepSeekMath ( Shao et al., 2024 ) , where reinforcement learning primarily boosts the correct response from TopK, we observe a genuine enhancement of fundamental capabilities in formal theorem proving. This improvement is evident not only with a small sample budget but also remains stable as the sample budget increases. This conclusion is further supported by later Monte-Carlo Tree Search experiments with larger sample budgets, as discussed in Section 4.2 .
|
||||
## Comparison between CoT and non-CoT.
|
||||
We compare the performance of non-CoT and CoT generation modes for both DeepSeek-Prover-V1.5-SFT and DeepSeek-Prover-V1.5-RL. The results in Figure 3 demonstrate that the CoT mode consistently outperforms the non-CoT mode across most settings. Specifically, DeepSeek-Prover-V1.5-RL, leveraging these enhanced theorem-proving patterns, achieves superior performance on both benchmarks, with an average accuracy of 51.6 % 51.6\% on miniF2F and 18.2 % 18.2\% on ProofNet. The integration of natural language reasoning in CoT mode significantly enhances the planning and execution of formal proof writing. For a detailed comparison of proof strategies with and without the use of natural language chain-of-thought, refer to the examples provided in Appendix A .
|
||||
## 3 Exploration-oriented Monte-Carlo Tree Search
|
||||
## 3.1 Tactic-level Tree Abstraction
|
||||
To implement the tree search method in the whole-proof generation setting, we introduce a proof tree abstraction to define the tailored state and action space, leveraging a truncate-and-resume mechanism.
|
||||
Roughly following the paradigm of Yao et al. (2023) , we begin by decomposing an incomplete proof into a sequence of tree nodes that correspond to individual proof steps, and then we utilize the partial content stored in these tree nodes to continue the proof generation process.
|
||||
Figure 4 illustrates the process of constructing a proof search tree from whole-proof generation.
|
||||
## Truncate: Proof Decomposition into Tree Nodes.
|
||||
We construct the proof search tree at the tactic level, where each tree edge represents a single transition step of the tactic state. Initially, we submit the entire proof the model generated to the Lean prover to parse it into tactics. We then truncate the proof at the earliest verification error, ensuring that all subsequent tactic codes can be successfully applied to advance the proof towards the desired theorem. The tactic codes are segmented into several code fractions, each containing a valid tactic code and its associated chain-of-thought comments, corresponding to a single tree edge that represents a tactic state transition. Through this abstraction, each tactic code is converted into a series of tree nodes, forming a path from the root to a specific node.
|
||||
## Resume: Proof Generation from a Tree Node.
|
||||
In Lean 4, different tactics can lead to the same tactic state, meaning each node in our proof tree can correspond to various tactic codes that achieve the same outcome. To handle this, we store a set of these equivalent tactic codes at each node. When the tree search agent expands a node, it randomly selects one tactic to use as a prompt for the language model. This prompt includes the incomplete proof code ending with the chosen tactic and the tactic state information from the Lean prover as a comment block. The fine-tuned model (see Section 2.2 ) has been trained to recognize and utilize this format, using the incomplete code augmented with tactic state comments to guide subsequent proof generation.
|
||||
Figure 4: Truncate-and-Resume Mechanism in the Expansion Step of MCTS. (a) After selecting a node, we trace its corresponding incomplete proof code prefix, which includes the file header, initial statement, and successfully applied tactics from the ancestor nodes. (b) The language model then generates the subsequent proof based on this prefix along with a comment block containing the current tactic state. (c) The combined proof code (prefix and newly generated code) is verified by the Lean 4 prover. If no errors are found, the tree-search procedure terminates. If errors are detected, we truncate the newly generated code at the first error message, discard the subsequent code, and parse the successful portion into tactics. (d) Each tactic is added as a new node in the search tree, extending a chain of descendants beneath the selected node. (e) Once the tree updates are complete, the next iteration of expansion begins by selecting an alternative candidate node, which is not limited to leaf nodes. This process repeats until a correct proof is found or the sample budget is exhausted.
|
||||
## 3.2 Interactive Theorem Proving via Monte-Carlo Tree Search
|
||||
Our proof search tree is developed using the standard Monte-Carlo Tree Search (MCTS) paradigm ( Coulom, 2006 ; Browne et al., 2012 , MCTS;) , which iteratively applies four steps: Selection , Expansion , Simulation , and Backpropagation .
|
||||
We integrate the Simulation step into Expansion because our whole-proof generation model inherently performs a rollout from the expanded node.
|
||||
The detailed design of the algorithm workflow is as follows.
|
||||
## Selection.
|
||||
The selection step, a.k.a. the tree policy, starts from the root node and traverses downward to identify a promising node for expansion.
|
||||
The objective of this algorithmic step is to trade off between exploration and exploitation ( Kocsis and Szepesvári, 2006 ) .
|
||||
The tree policy at a tree node s s is computed by selecting the action that maximizes the value from the set of valid operations:
|
||||
T r e e P o l i c y ( s ) = arg max a ∈ C h i l d r e n ( s ) ∪ { ⊘ } Q U C B ( s , a ) , \displaystyle TreePolicy(s)=\mathop{\arg\max}_{a\in Children(s)\cup\{\oslash\}}Q_{UCB}(s,a),
|
||||
(1)
|
||||
where the action a a can be either moving to a child node, denoted by a ∈ C h i l d r e n ( s ) a\in Children(s) , or expanding the current node s s , denoted by a special token a = ⊘ a=\oslash .
|
||||
This approach uses a technique called virtual node ( Wang et al., 2023b ) , which assigns each node an imaginary child to represent the selection of the current node s s for expansion.
|
||||
It enables the tree search agent to continually expand non-leaf nodes, as the action space is supported by a generative model whose output scope cannot be determined by a fixed number of trails.
|
||||
The value estimation Q U C B ( s , a ) Q_{UCB}(s,a) of performing action a a on node s s is composed by two components:
|
||||
∀ a ∈ C h i l d r e n ( s ) ∪ { ⊘ } , Q U C B ( s , a ) \displaystyle\forall a\in Children(s)\cup\{\oslash\},\quad Q_{UCB}(s,a)
|
||||
= Q ( s , a ) ⏟ Exploitation + U C B ( s , a ) ⏟ Exploration , \displaystyle=\underbrace{Q(s,a)}_{\text{Exploitation}}+\penalty\ \penalty\ \penalty\ \underbrace{UCB(s,a)}_{\text{Exploration}},
|
||||
(2)
|
||||
where Q ( s , a ) Q(s,a) denotes a sample-based estimation of action values derived from the selection history, functioning as the exploitation component that retrieves high-value candidates from previous trials. U C B ( s , a ) UCB(s,a) denotes the exploration bonus computed by upper confidence bounds ( Auer, 2002 , UCB;) , which diminishes with the repeated execution of the state-action pair ( s , a ) (s,a) . More specifically, Q U C B ( s , a ) Q_{UCB}(s,a) stands for an optimistic estimation of Q ( s , a ) Q(s,a) and can serve as an upper bound with high probability. We defer the discussion of detailed settings of node values and UCB bonus to Section 3.3 .
|
||||
## Expansion.
|
||||
The next step is invoking the proof generation model to expand the node nominated by the selection phase.
|
||||
Resuming the incomplete proof codes stored on the node designated for expansion, we perform whole-proof generation to propose a series of subsequent tactics and submit the generated proof to Lean prover for verification.
|
||||
Such a trial of proof completion is equivalent to conducting a single rollout of simulation within the standard MCTS framework.
|
||||
When the verification result indicates the proof is complete, the search procedure is ready to be terminated, having found a new proof of the desired theorem.
|
||||
Otherwise, we parse the verification feedback and truncate the generated proof to the assertion of the earliest verification error.
|
||||
The remaining tactics are transformed into a path of nodes to be merged into the search tree (see Figure 4 ).
|
||||
It is important to note that, because we use the whole-proof generation setting—where the output is an entire proof consisting of a sequence of tactics, rather than just the next tactic—our expansion procedure may insert a path of tree nodes into the search tree during each iteration. This differs from the conventional MCTS designed for competitive games, which typically expands only one layer of children nodes per iteration ( Silver et al., 2016 ; Silver et al., 2018 ; Schrittwieser et al., 2020 ) .
|
||||
## Backpropagation.
|
||||
The final phase of each tree search iteration is to update value statistics along the selection trajectory from the root to the expanded node, i.e. , updating the values associated with the tree policy stated in Eq. ( 1 ).
|
||||
Let τ = { ( r o o t , s ( 1 ) ) , ( s ( 1 ) , s ( 2 ) ) , ( s ( 2 ) , s ( 3 ) ) , … , ( s ( | τ | − 1 ) = s t , ⊘ ) } \tau=\{(root,s^{(1)}),(s^{(1)},s^{(2)}),(s^{(2)},s^{(3)}),\dots,(s^{(|\tau|-1)}=s_{t},\oslash)\} denote the selection trajectory of t t -th iteration that ends with s t s_{t} as the expanding node.
|
||||
We update Q U C B ( s , a ) Q_{UCB}(s,a) for all ( s , a ) ∈ τ (s,a)\in\tau by taking the most recent trajectory reward R ( τ ) R(\tau) into account (details refer to Eq. ( 7 )).
|
||||
The extrinsic source of rewards comes from the compiler feedback, specifically assigning a reward of R extrinsic ( τ ) = 1 R_{\text{extrinsic}}(\tau)=1 for completed proofs and R extrinsic ( τ ) = 0 R_{\text{extrinsic}}(\tau)=0 for unsolved ones.
|
||||
In Section 3.3 , we will introduce an intrinsic reward mechanism to augment the reward assignment that enhances the agent’s incentive for exploration.
|
||||
## 3.3 Intrinsic Rewards for Monte-Carlo Tree Search
|
||||
In the search problem of formal theorem proving, the extrinsic rewards are extremely sparse, i.e. , the search agent only obtains non-zero rewards when the proof is completely solved.
|
||||
More specifically, the proof search process forms a tree structure with only a narrow set of leaves delivering non-zero rewards, which matches a famous hard-exploration case ( Krishnamurthy et al., 2016 ) in the literature of statistical reinforcement learning.
|
||||
To promote exploration in sparse-reward sequential decision making, one classical paradigm is constructing intrinsic rewards ( Schmidhuber, 2010 ) that encourage the agent to not only optimize extrinsic rewards but also acquire general information about the interactive environment ( Bellemare et al., 2016 ; Houthooft et al., 2016 ; Pathak et al., 2017 ; Burda et al., 2019 ) .
|
||||
In this section, we present our intrinsic-reward-driven exploration algorithm, RMax applied to Tree Search (RMaxTS), to incorporate reward-free exploration in the proof search problem.
|
||||
## RMax applied to MCTS.
|
||||
We adopt RMax ( Brafman and Tennenholtz, 2002 ) , a classical exploration mechanism, to construct intrinsic rewards for Monte-Carlo tree search.
|
||||
The core idea of RMax is to explore a broad coverage of the state space.
|
||||
The agent awards itself a maximal amount of reward upon reaching an unseen state.
|
||||
In the context of proof search, where no extrinsic rewards are provided until the proof is completed, our algorithmic procedure resembles ZeroRMax ( Jin et al., 2020 ) , in which the agent’s exploration is driven solely by intrinsic rewards, i.e. , setting R ( τ ) = R intrinsic ( τ ) R(\tau)=R_{\text{intrinsic}}(\tau) .
|
||||
The intrinsic reward of a tree expansion step is determined by whether a new node is added to the search tree,
|
||||
R intrinsic ( τ ) = 𝕀 [ at least one new node is added to the search tree ] , \displaystyle R_{\text{intrinsic}}(\tau)=\mathbb{I}\left[\text{at least one new node is added to the search tree}\right],
|
||||
(3)
|
||||
where τ \tau denotes the most recent selection trajectory that requires a reward assignment for backpropagation.
|
||||
This exploration strategy prioritizes the expansion of nodes where the prover model generates tactics that lead to a diverse range of tactic states.
|
||||
As multiple Lean codes can result in the same transition of intermediate states, this heuristics can potentially reduce redundant generation and improve sample efficiency.
|
||||
## UCB for Non-stationary Rewards.
|
||||
The common setting of UCB exploration bonus for Monte-Carlo tree search is using UCB1 ( Auer et al., 2002 ) :
|
||||
Q U C B 1 ( s , a ) \displaystyle Q_{UCB1}(s,a)
|
||||
= W ( s , a ) N ( s , a ) + 2 ln ∑ a ′ N ( s , a ′ ) N ( s , a ) , \displaystyle=\frac{W(s,a)}{N(s,a)}+\sqrt{\frac{2\ln\sum_{a^{\prime}}N(s,a^{\prime})}{N(s,a)}},
|
||||
(4)
|
||||
W ( s , a ) \displaystyle W(s,a)
|
||||
= ∑ τ ∈ Γ ( s , a ) R ( τ ) , \displaystyle=\textstyle{\sum_{\tau\in\Gamma(s,a)}}\penalty\ R(\tau),
|
||||
(5)
|
||||
N ( s , a ) \displaystyle N(s,a)
|
||||
= | Γ ( s , a ) | , \displaystyle=\left|\Gamma(s,a)\right|,
|
||||
(6)
|
||||
where Γ ( s , a ) = { τ ∣ ( s , a ) ∈ τ } \Gamma(s,a)=\{\tau\mid(s,a)\in\tau\} denotes the list of tree-policy trajectory τ \tau containing ( s , a ) (s,a) as an intermediate selection step.
|
||||
To facilitate discussions, we organize the list Γ ( s , a ) = { τ 1 , τ 2 , ⋯ } \Gamma(s,a)=\{\tau_{1},\tau_{2},\cdots\} such that newly collected trajectories have larger subscript indices.
|
||||
In this work, we propose to use an alternative variant of UCB method.
|
||||
Note that the derived intrinsic reward in Eq. ( 3 ) is a non-stationary reward signal whose expected value decays with the progress of exploration.
|
||||
That is because it becomes definitely harder to discover new nodes with unseen tactic states as the search tree expands through sophisticated exploration.
|
||||
To tackle the non-stationarity, we consider discounted upper confidence bounds ( Garivier and Moulines, 2011 , DUCB;) , which uses a discount factor γ ∈ ( 0 , 1 ) \gamma\in(0,1) to smoothly drop those outdated feedback records:
|
||||
Q D U C B ( s , a ) \displaystyle Q_{DUCB}(s,a)
|
||||
= W γ ( s , a ) N γ ( s , a ) + 2 ln ∑ a ′ N γ ( s , a ′ ) N γ ( s , a ) , \displaystyle=\frac{W_{\gamma}(s,a)}{N_{\gamma}(s,a)}+\sqrt{\frac{2\ln\sum_{a^{\prime}}N_{\gamma}(s,a^{\prime})}{N_{\gamma}(s,a)}},
|
||||
(7)
|
||||
W γ ( s , a ) \displaystyle W_{\gamma}(s,a)
|
||||
= ∑ t = 1 N ( s , a ) γ N ( s , a ) − t R ( τ t ) , \displaystyle=\textstyle{\sum_{t=1}^{N(s,a)}}\penalty\ \gamma^{N(s,a)-t}R(\tau_{t}),
|
||||
(8)
|
||||
N γ ( s , a ) \displaystyle N_{\gamma}(s,a)
|
||||
= ∑ t = 0 N ( s , a ) − 1 γ t , \displaystyle=\textstyle{\sum_{t=0}^{N(s,a)-1}}\penalty\ \gamma^{t},
|
||||
(9)
|
||||
where newly received feedback would be assigned a larger weight in the value estimation.
|
||||
In practice, we set γ = 0.99 \gamma=0.99 .
|
||||
Note that the role of discount factor γ \gamma in DUCB differs from its role in value iteration for infinite-horizon MDPs.
|
||||
The discounting is applied to tree search iterations rather than to the action-step horizon within a single trajectory.
|
||||
## 3.4 Parallelization of Monte-Carlo Tree Search
|
||||
To enhance the efficiency of Monte-Carlo Tree Search (MCTS), we implement several established parallelization techniques as described by Chaslot et al. (2008) .
|
||||
•
|
||||
Root Parallelization: We deploy 256 MCTS runners per node, with one language model per GPU and a batch size of 512 for proof generation. The Lean prover is invoked through REPL and executed on a cluster with thousands of CPU cores, where each proof verification task is handled by an individual process, created and terminated in a sandbox. Both proof generation by language models and verification by Lean provers are handled asynchronously. This setup allows MCTS runners to perform concurrent tree search operations, significantly accelerating the process.
|
||||
•
|
||||
Tree Parallelization: We manage each search tree with 32 thread workers to parallelize the tree iteration steps. This method effectively schedules and balances the tasks of proof generation and Lean verification. Each thread worker iteratively performs the tree search loop by selecting a candidate node for expansion, invoking the language model to generate the proof, verifying the generated proof with the Lean prover, and performing backpropagation.
|
||||
•
|
||||
Virtual Loss: To encourage diverse node selection among concurrent thread workers, we assign a virtual reward R ( τ ) = 0 R(\tau)=0 for ongoing iterations. This involves backpropagating a reward of 0 0 temporarily and updating it to the true reward upon completion. This strategy promotes exploration of different nodes for expansion, thereby enhancing the overall search efficiency.
|
||||
## 3.5 Comparison with Existing Methods
|
||||
In this section, we compare our proposed proof tree search method, which introduces a novel truncate-and-resume mechanism for whole-proof generation, with existing approaches. Current methods for using language models in formal mathematics proof search generally fall into two main strategies:
|
||||
•
|
||||
Multi-pass proof-step generation : This strategy breaks down the proving process into multiple episodes of tactic generation and verification, typically following a tree search pattern. It involves generating and verifying one tactic at a time, repeating the process for the next tactic until no proof goals remain. Notable examples include GPT-f ( Polu and Sutskever, 2020 ; Polu et al., 2022 ) , Thor ( Jiang et al., 2022a ) , ReProver ( Yang et al., 2023 ) , Hypertree Proof Search ( Lample et al., 2022 ) , and InternLM2-StepProver ( Wu et al., 2024 ) .
|
||||
•
|
||||
Single-pass whole-proof generation : This approach generates and verify an entire proof in one attempt. If the proof is incorrect, the model generates a new proof in the next attempt. Methods in this category include DSP ( Jiang et al., 2022b ) , Subgoal-Prover Zhao et al. (2023) , LEGO-Prover ( Wang et al., 2023a ) , Lyra ( Zheng et al., 2023 ) , and miniCTX ( Hu et al., 2024 ) .
|
||||
Our proof tree search method uniquely bridges these two strategies, offering a novel hybrid approach. It starts with whole-proof generation, similar to the single-pass approach, but extends this by implementing a sophisticated truncate-and-resume mechanism. This process involves truncating the generated proof to its successful initial segment, parsing this segment into individual tactics, and resuming the tree search from this point. This iterative process effectively implements a Monte-Carlo Tree Search, seamlessly integrating single-pass whole-proof generation with multi-pass proof-step generation. Consequently, we can train a single model with nearly identical objectives to support both strategies simultaneously. Our experimental results demonstrate that this unified approach achieves superior performance in both settings. By combining the strengths of existing methods and introducing innovative techniques, our method offers a more versatile and effective solution for formal mathematics proof search, potentially paving the way for future advancements in this field.
|
||||
## 4 Experimental Results
|
||||
In this section, we evaluate the theorem-proving capabilities of DeepSeek-Prover-V1.5 using two distinct benchmarks: miniF2F, which encompasses high-school level exercises and competition problems, and ProofNet, which pertains to undergraduate-level theorems. We present the results for both complete proof generation and Monte-Carlo tree search methodologies, utilizing the same trained model and inference configuration as Section 2.4 to ensure consistency.
|
||||
## 4.1 Main Results
|
||||
We present a comparative analysis of DeepSeek-Prover-V1.5 against previous state-of-the-art language models, highlighting its performance and advancements.
|
||||
•
|
||||
General-purpose Models: GPT-3.5 and GPT-4 ( OpenAI, 2023 ) are advanced generative AI models developed by OpenAI, known for their effectiveness across diverse tasks, including code generation. Despite not being specifically designed for theorem proving, their extensive parameter scales provide significant capabilities. The evaluation of these models in formal theorem proving is facilitated by COPRA ( Thakur et al., 2023 ) , an in-context learning agent that leverages these large language models to propose tactic applications. Additionally, we examine Llemma ( Azerbayev et al., 2024 ) , a series of language models trained on extensive general mathematical corpora, commonly used as the base model for formal theorem proving.
|
||||
•
|
||||
Specialized Models for Formal Mathematics: GPT-f ( Polu and Sutskever, 2020 ; Polu et al., 2022 ) represents an initial effort to apply Transformers ( Vaswani et al., 2017 ) to proof-step generation for theorem proving tasks, utilizing a best-first search module to construct complete proofs. Subsequent advancements include ReProver ( Yang et al., 2023 ) , LLMStep ( Welleck and Saha, 2023 ) , and Lean-STaR ( Lin et al., 2024 ) . Hypertree Proof Search ( Lample et al., 2022 ) explores the use of Monte Carlo tree search in formal theorem proving using Lean. Concurrent works, InternLM2-Math ( Ying et al., 2024b ) and InternLM2-StepProver ( Wu et al., 2024 ) , also demonstrate outstanding performance.
|
||||
## Metric.
|
||||
We compare the performance of DeepSeek-Prover-V1.5 with state-of-the-art models using the pass@ K K accuracy metric, which evaluates the model’s ability to generate a correct proof within K K attempts.
|
||||
We display the sample budget K K according to the the following rules to align the computation budget across different generation schemes.
|
||||
•
|
||||
For single-pass sampling methods, we define the sample budget K K as the total number of proofs generated, with large values of K K factorized for the ease of comparison to tree search methods.
|
||||
•
|
||||
For best-first-search methods, following the notation of Azerbayev et al. (2024) , we present K = N × S × T K=N\times S\times T where N N denotes the number of best-first-search attempts, S S denotes the number of tactics generated for each expansion, and T T denotes the number of expansion iterations.
|
||||
•
|
||||
For tree search methods, e.g. , RMaxTS and HTPS ( Lample et al., 2022 ) , we present K = N × T K=N\times T where N N denotes the number of tree search attempts, and T T denotes the number of model generations invoked in tree expansions.
|
||||
## Results on miniF2F.
|
||||
Table 1 provides a comparative analysis of various theorem-proving methods on the miniF2F-test dataset. In the single-pass whole-proof generation setting, DeepSeek-Prover-V1.5-RL achieved the highest pass rate at 60.2%, marking a significant improvement of 10.2 percentage points over DeepSeek-Prover-V1’s 50.0%. With a sampling budget limited to 128 attempts, DeepSeek-Prover-V1.5-RL proved 51.6% of the problems, significantly outperforming other whole-proof generation methods and is comparable to the leading tree search methods. In the Tree Search Methods category, DeepSeek-Prover-V1.5-RL + RMaxTS leads with a pass rate of 62.7%, establishing a new state-of-the-art and creating a substantial gap with existing methods. Notably, DeepSeek-Prover-V1.5-RL requires only 3200 whole-proof generation samplings to achieve a pass rate of 54.9%, surpassing the previous state-of-the-art result of InternLM2-StepProver, which performs 64 × 3200 64\times 3200 tree searches to achieve 54.5%.
|
||||
Method
|
||||
Sample budget
|
||||
miniF2F-test
|
||||
Single-pass Whole-Proof Generation Methods
|
||||
TheoremLlama [ 49 ]
|
||||
128
|
||||
33.6 % 33.6\%
|
||||
DeepSeek-Prover-V1 [ 53 ]
|
||||
128 128
|
||||
46.1 % ± 0.5 % 46.1\%\pm 0.5\%
|
||||
16 × 4096 16\times 4096
|
||||
50.0 % 50.0\%
|
||||
DeepSeek-Prover-V1.5-Base
|
||||
128
|
||||
29.7 % ± 0.5 % 29.7\%\pm 0.5\%
|
||||
3200
|
||||
39.2 % 39.2\%
|
||||
6400
|
||||
42.2 % 42.2\%
|
||||
DeepSeek-Prover-V1.5-SFT
|
||||
32
|
||||
48.2 % ± 0.6 % 48.2\%\pm 0.6\%
|
||||
64
|
||||
49.6 % ± 0.7 % 49.6\%\pm 0.7\%
|
||||
128 128
|
||||
50.4 % ± 0.4 % 50.4\%\pm 0.4\%
|
||||
3200 3200
|
||||
53.3 % ± 0.5 % 53.3\%\pm 0.5\%
|
||||
4 × 6400 4\times 6400
|
||||
55.8 % ± 0.7 % 55.8\%\pm 0.7\%
|
||||
16 × 6400 16\times 6400
|
||||
57.4 % 57.4\%
|
||||
DeepSeek-Prover-V1.5-RL
|
||||
32
|
||||
50.0 % ± 0.5 % 50.0\%\pm 0.5\%
|
||||
64
|
||||
50.7 % ± 0.4 % 50.7\%\pm 0.4\%
|
||||
128 128
|
||||
51.6 % ± 0.5 % 51.6\%\pm 0.5\%
|
||||
3200 3200
|
||||
54.9 % ± 0.7 % 54.9\%\pm 0.7\%
|
||||
4 × 6400 4\times 6400
|
||||
58.4 % ± 0.6 % 58.4\%\pm 0.6\%
|
||||
16 × 6400 16\times 6400
|
||||
60.2 % \mathbf{60.2\%}
|
||||
Tree Search Methods
|
||||
COPRA (Code Llama) [ 45 ]
|
||||
1 × 500 1\times 500
|
||||
5.7 % 5.7\%
|
||||
COPRA (GPT-3.5) [ 45 ]
|
||||
1 × 60 1\times 60
|
||||
9.0 % 9.0\%
|
||||
COPRA (GPT-4) [ 45 ]
|
||||
1 × 60 1\times 60
|
||||
26.6 % 26.6\%
|
||||
Llemma-7B [ 5 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
26.2 % 26.2\%
|
||||
Llemma-34B [ 5 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
25.8 % 25.8\%
|
||||
ReProver [ 55 ]
|
||||
-
|
||||
26.5 % 26.5\%
|
||||
LLMStep [ 51 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
27.9 % 27.9\%
|
||||
GPT-f [ 35 ]
|
||||
64 × 8 × 512 64\times 8\times 512
|
||||
36.6 % 36.6\%
|
||||
Hypertree Proof Search [ 23 ]
|
||||
64 × 5000 64\times 5000
|
||||
41.0 % 41.0\%
|
||||
Lean-STaR [ 26 ]
|
||||
64 × 1 × 50 64\times 1\times 50
|
||||
46.3 % 46.3\%
|
||||
InternLM2-Math-7B [ 58 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
30.3 % 30.3\%
|
||||
InternLM2-Math-Plus-7B [ 58 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
43.4 % 43.4\%
|
||||
InternLM2-StepProver [ 52 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
48.8 % 48.8\%
|
||||
64 × 32 × 100 64\times 32\times 100
|
||||
54.5 % 54.5\%
|
||||
DeepSeek-Prover-V1.5-SFT + RMaxTS
|
||||
1 × 3200 1\times 3200
|
||||
53.5 % ± 0.4 % 53.5\%\pm 0.4\%
|
||||
4 × 6400 4\times 6400
|
||||
56.3 % ± 0.3 % 56.3\%\pm 0.3\%
|
||||
16 × 6400 16\times 6400
|
||||
59.0 % 59.0\%
|
||||
32 × 6400 † 32\times 6400^{\dagger}
|
||||
60.2 % 60.2\%
|
||||
DeepSeek-Prover-V1.5-RL + RMaxTS
|
||||
1 × 3200 1\times 3200
|
||||
55.0 % ± 0.7 % 55.0\%\pm 0.7\%
|
||||
4 × 6400 4\times 6400
|
||||
59.6 % ± 0.6 % 59.6\%\pm 0.6\%
|
||||
16 × 6400 16\times 6400
|
||||
62.7 % \mathbf{62.7\%}
|
||||
32 × 6400 † 32\times 6400^{\dagger}
|
||||
63.5 % \mathbf{63.5\%}
|
||||
Table 1:
|
||||
Comparison with state-of-the-art methods on the miniF2F-test dataset. The notation μ ± σ \mu\pm\sigma denotes the average accuracy μ \mu and the standard deviation σ \sigma . Unless otherwise specified, DeepSeek-Prover-V1.5-Base results are based on 3-shot prompting, while DeepSeek-Prover-V1.5-SFT and RL employ CoT mode prompting. The symbol † {\dagger} indicates performance using a mixture strategy with two guiding prompts (see Section 4.2 for details).
|
||||
## Results on ProofNet.
|
||||
Table 2 presents a comparative analysis of various theorem-proving methods on the ProofNet dataset. DeepSeek-Prover-V1.5-RL achieved pass rates of 22.6% and 25.3% for the overall ProofNet dataset in the single-pass whole-proof generation setting and with the enhancement of RMaxTS, respectively. These results surpass the existing state-of-the-art methods, ReProver (13.8%) and InternLM2-StepProver (18.1%). When the number of whole-proof generation attempts is restricted to 3200, DeepSeek-Prover-V1.5 also proved 21.7% of the theorems, demonstrating a 3.6% improvement over the previous state-of-the-art, InternLM2-StepProver.
|
||||
Method
|
||||
Sample budget
|
||||
ProofNet
|
||||
valid ‡
|
||||
test
|
||||
all
|
||||
Single-pass Whole-Proof Generation Methods
|
||||
DeepSeek-Prover-V1.5-Base
|
||||
128 128
|
||||
6.6 % ± 0.9 % 6.6\%\pm 0.9\%
|
||||
9.7 % ± 0.7 % 9.7\%\pm 0.7\%
|
||||
7.5 % ± 0.7 % 7.5\%\pm 0.7\%
|
||||
3200 3200
|
||||
10.8 % 10.8\%
|
||||
15.6 % 15.6\%
|
||||
13.2 % 13.2\%
|
||||
DeepSeek-Prover-V1.5-SFT
|
||||
128 128
|
||||
19.9 % ± 0.4 % 19.9\%\pm 0.4\%
|
||||
15.9 % ± 0.6 % 15.9\%\pm 0.6\%
|
||||
17.9 % ± 0.3 % 17.9\%\pm 0.3\%
|
||||
3200 3200
|
||||
20.7 % ± 0.7 % 20.7\%\pm 0.7\%
|
||||
21.0 % ± 0.9 % 21.0\%\pm 0.9\%
|
||||
20.9 % ± 0.6 % 20.9\%\pm 0.6\%
|
||||
4 × 6400 4\times 6400
|
||||
22.2 % 22.2\%
|
||||
23.7 % 23.7\%
|
||||
22.9 % 22.9\%
|
||||
DeepSeek-Prover-V1.5-RL
|
||||
128 128
|
||||
20.1 % ± 0.5 % 20.1\%\pm 0.5\%
|
||||
18.2 % ± 0.5 % 18.2\%\pm 0.5\%
|
||||
19.1 % ± 0.4 % 19.1\%\pm 0.4\%
|
||||
3200 3200
|
||||
21.4 % ± 0.3 % 21.4\%\pm 0.3\%
|
||||
22.0 % ± 0.5 % 22.0\%\pm 0.5\%
|
||||
21.7 % ± 0.4 % 21.7\%\pm 0.4\%
|
||||
4 × 6400 4\times 6400
|
||||
21.6 % 21.6\%
|
||||
23.7 % 23.7\%
|
||||
22.6 % 22.6\%
|
||||
Tree Search Methods
|
||||
ReProver [ 55 ]
|
||||
-
|
||||
-
|
||||
-
|
||||
13.8 % 13.8\%
|
||||
InternLM2-StepProver [ 52 ]
|
||||
1 × 32 × 100 1\times 32\times 100
|
||||
-
|
||||
-
|
||||
18.1 % 18.1\%
|
||||
DeepSeek-Prover-V1.5-SFT + RMaxTS
|
||||
1 × 3200 1\times 3200
|
||||
22.2 % ± 0.7 % 22.2\%\pm 0.7\%
|
||||
21.6 % ± 0.2 % 21.6\%\pm 0.2\%
|
||||
21.9 % ± 0.4 % 21.9\%\pm 0.4\%
|
||||
4 × 6400 4\times 6400
|
||||
23.8 % 23.8\%
|
||||
25.8%
|
||||
24.8 % 24.8\%
|
||||
DeepSeek-Prover-V1.5-RL + RMaxTS
|
||||
1 × 3200 1\times 3200
|
||||
22.0 % ± 0.3 % 22.0\%\pm 0.3\%
|
||||
21.5 % ± 0.8 % 21.5\%\pm 0.8\%
|
||||
21.8 % ± 0.4 % 21.8\%\pm 0.4\%
|
||||
4 × 6400 4\times 6400
|
||||
25.4 % 25.4\%
|
||||
25.3%
|
||||
25.3 % 25.3\%
|
||||
Table 2: Comparing with state-of-the-arts on the ProofNet dataset. ‡ Note that the validation set of ProofNet is used to perform expert iteration in supervised fine-tuning.
|
||||
## 4.2 Re-Examining the Effectiveness of Training Strategies on Large-scale Sampling
|
||||
We revisit the effects of several training modules in n a large-scale sampling setting, focusing on both single-pass whole-proof generation and Monte-Carlo tree search. The results demonstration that the observations and findings presented in Section 2.4 generalize to sampling scenarios with a large sample size.
|
||||
## General Enhancement of Reinforcement Learning.
|
||||
To support the claim that online reinforcement learning from verification feedback generally enhances the model capabilities, we compare our final model to the SFT-only version using a large sample budget.
|
||||
The comparison results are presented as two columns in Table 3 .
|
||||
DeepSeek-Prover-V1.5-RL consistently outperforms the SFT model across all generation settings, regardless of whether the chain-of-thought strategy is applied.
|
||||
The results also indicate that the improvements gained from conducting online RL is orthogonal to those achieved through RMaxTS, which can be further combined to boost the performance.
|
||||
By integrating both CoT prompting and RMaxTS, DeepSeek-Prover-V1.5-RL achieves a pass rate of 62.7 % 62.7\% on miniF2F-test.
|
||||
This performance shows a notable 3.7 % 3.7\% improvement over the SFT model, highlighting the critical role of reinforcement learning in enhancing the overall effectiveness of the proof completion model.
|
||||
## CoT, non-CoT, and Mixture Strategy.
|
||||
We compare the performance of two generation modes, i.e. , non-CoT and CoT, on miniF2F-test dataset.
|
||||
The results, shown in Table 3 , indicate that the advantage of CoT over the non-CoT mode is amplified as the sample budget increases.
|
||||
This suggests that the incorporation of natural language chain-of-thought can diversify the planning pathways of theorem proving, potentially leading to a broader range of reasoning strategies and more innovative solutions.
|
||||
Results also show that these two modes have complementary advantages across different problems.
|
||||
The model’s theorem proving strategy in the CoT mode is more systematic and proactive in mathematical thinking, while in the non-CoT mode, the model can efficiently use Lean high-level tactics to solve computational problems that can be addressed within Lean’s automation mechanisms.
|
||||
To leverage these advantages, we consider a mixture strategy, denoted by non-CoT & CoT in Table 3 , allocates half of sample budget to the CoT mode and the remains to the non-CoT mode.
|
||||
This simple combination of two guiding prompts shows great promise in further bootstrapping the performance of our proof completion model, achieving a pass rate of 63.5 % 63.5\% on miniF2F-test. In Appendix B , we present example problems that illustrate the different advantages of the two generation modes.
|
||||
Prompt mode
|
||||
Sample budget
|
||||
DeepSeek-Prover-V1.5
|
||||
SFT
|
||||
RL
|
||||
non-CoT
|
||||
4 × 6400 4\times 6400
|
||||
54.7 % ± 0.4 % 54.7\%\pm 0.4\%
|
||||
56.5 % ± 0.5 % 56.5\%\pm 0.5\%
|
||||
16 × 6400 16\times 6400
|
||||
56.1 % 56.1\%
|
||||
57.4 % 57.4\%
|
||||
Single-Pass
|
||||
CoT
|
||||
4 × 6400 4\times 6400
|
||||
55.8 % ± 0.7 % 55.8\%\pm 0.7\%
|
||||
58.4 % ± 0.5 % 58.4\%\pm 0.5\%
|
||||
Generation
|
||||
16 × 6400 16\times 6400
|
||||
57.4 % 57.4\%
|
||||
60.2 % 60.2\%
|
||||
non-CoT & CoT
|
||||
( 2 + 2 ) × 6400 (2+2)\times 6400
|
||||
56.1 % ± 0.8 % 56.1\%\pm 0.8\%
|
||||
58.3 % ± 0.6 % 58.3\%\pm 0.6\%
|
||||
( 8 + 8 ) × 6400 (8+8)\times 6400
|
||||
58.2 % 58.2\%
|
||||
60.7 % 60.7\%
|
||||
( 16 + 16 ) × 6400 (16+16)\times 6400
|
||||
58.6 % 58.6\%
|
||||
61.1 % 61.1\%
|
||||
RMaxTS
|
||||
non-CoT
|
||||
4 × 6400 4\times 6400
|
||||
55.7 % ± 0.6 % 55.7\%\pm 0.6\%
|
||||
58.4 % ± 0.6 % 58.4\%\pm 0.6\%
|
||||
16 × 6400 16\times 6400
|
||||
57.8 % 57.8\%
|
||||
59.4 % 59.4\%
|
||||
CoT
|
||||
4 × 6400 4\times 6400
|
||||
56.3 % ± 0.3 % 56.3\%\pm 0.3\%
|
||||
59.6 % ± 0.6 % 59.6\%\pm 0.6\%
|
||||
16 × 6400 16\times 6400
|
||||
59.0 % 59.0\%
|
||||
62.7 % 62.7\%
|
||||
non-CoT & CoT
|
||||
( 2 + 2 ) × 6400 (2+2)\times 6400
|
||||
56.1 % ± 0.8 % 56.1\%\pm 0.8\%
|
||||
60.0 % ± 0.8 % 60.0\%\pm 0.8\%
|
||||
( 8 + 8 ) × 6400 (8+8)\times 6400
|
||||
59.0 % 59.0\%
|
||||
63.1 % 63.1\%
|
||||
( 16 + 16 ) × 6400 (16+16)\times 6400
|
||||
60.2 % 60.2\%
|
||||
63.5%
|
||||
Table 3: A large-scale ablation study to investigate the effectiveness of several algorithmic designs on model training. The results are evaluated on the miniF2F-test dataset.
|
||||
## 4.3 Ablation Studies on RMaxTS
|
||||
## Intrinsic Rewards and Discounted UCB.
|
||||
We investigate the effectiveness of two core components of RMaxTS, i.e. , the intrinsic rewards defined in Eq. ( 3 ) and the discounted upper confidence bound stated in Eq. ( 7 ).
|
||||
We start with a baseline implementing the standard UCT algorithm ( Kocsis and Szepesvári, 2006 ) without intrinsic rewards, in which the exploration is driven exclusively by the UCB bonus.
|
||||
Note that, since no non-zero rewards are provided for this baseline, all variants of the UCB formula become equivalent, as node selection is determined solely by visitation counts.
|
||||
The experimental results in Figure 5 show that, in the absence of intrinsic rewards, the performance of UCT (without R intrinsic R_{\text{intrinsic}} ) degenerates into a level comparable to that of non-search methods.
|
||||
Furthermore, we consider RMaxTS using the standard UCB1 (refer to Eq. ( 4 )) instead of the discounted UCB, denoted by RMaxTS (DUCB → \rightarrow UCB1).
|
||||
The results indicate that the performance of RMaxTS with UCB1 bonus is also moderate, comparable to that of UCT (without R intrinsic R_{\text{intrinsic}} ).
|
||||
That is because UCB1 is designed to guarantee asymptotic performance through exhausted exploration ( Auer et al., 2002 ) assuming the sample size to be sufficiently large.
|
||||
In contrast, the discounted UCB can accelerate the value propagation of non-stationary intrinsic rewards, preventing the guidance of R intrinsic R_{\text{intrinsic}} from being dominated by that of visitation counts.
|
||||
This demonstrates that the discounted UCB mechanism is a crucial complement to intrinsic-reward-driven exploration.
|
||||
Sample budget
|
||||
miniF2F-test
|
||||
Single-Pass Generation
|
||||
4 × 6400 4\times 6400
|
||||
58.4 % ± 0.5 % 58.4\%\pm 0.5\%
|
||||
16 × 6400 16\times 6400
|
||||
60.2 % 60.2\%
|
||||
UCT
|
||||
4 × 6400 4\times 6400
|
||||
58.2 % ± 0.3 % 58.2\%\pm 0.3\%
|
||||
(without R intrinsic R_{\text{intrinsic}} )
|
||||
16 × 6400 16\times 6400
|
||||
61.1 % 61.1\%
|
||||
RMaxTS
|
||||
4 × 6400 4\times 6400
|
||||
58.6 % ± 0.3 % 58.6\%\pm 0.3\%
|
||||
(DUCB → \rightarrow UCB1)
|
||||
16 × 6400 16\times 6400
|
||||
60.7 % 60.7\%
|
||||
RMaxTS
|
||||
4 × 6400 4\times 6400
|
||||
58.4 % ± 0.3 % 58.4\%\pm 0.3\%
|
||||
(without tactic state)
|
||||
16 × 6400 16\times 6400
|
||||
61.1 % 61.1\%
|
||||
RMaxTS
|
||||
4 × 6400 4\times 6400
|
||||
59.6 % ± 0.6 % 59.6\%\pm 0.6\%
|
||||
16 × 6400 16\times 6400
|
||||
62.7 % 62.7\%
|
||||
Figure 5: A modular ablation study examining the algorithmic design of RMaxTS. The experiments are conducted on the miniF2F-test dataset with DeepSeek-Prover-V1.5-RL using the CoT mode. The left panel presents the curves of Pass@K accuracy within 6400 generation samples. The results with a larger sample size are presented in the right panel.
|
||||
## Guidance of Tactic State Information.
|
||||
When expanding a tree node, we concatenate the intermediate tactic state information as a comment block to the incomplete code to guide the proof completion.
|
||||
With the provided auxiliary information, the proof completion model can enhance its internal representation of the tactic state, offering intermediate guidance for long-horizon planning.
|
||||
To demonstrate this advantage, we present experiments on RMaxTS that performs code completion directly from the raw incomplete code without accessing tactic state information, denoted by RMaxTS (without tactic state) in Figure 5 .
|
||||
The results indicate that the performance gain from applying tree search becomes moderate in the absence of tactic state information, especially when tackling hard problems that require a large amount of samples.
|
||||
This highlig
|
||||
|
||||
=== BODY (tail) ===
|
||||
hts that the integration of compiler information is an essential component of the tree search algorithm, enhancing its overall effectiveness and sample efficiency.
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
We present DeepSeek-Prover-V1.5, a language model with 7 billion parameters that outperforms all open-source models in formal theorem proving in Lean 4.
|
||||
DeepSeek-Prover-V1.5 is initialized with DeepSeek-Prover-V1.5-Base, which extends the pre-training of DeepSeekMath-Base 7B using a specialized corpus for formal mathematical reasoning.
|
||||
Supervised fine-tuning is conducted on a comprehensive Lean 4 code completion dataset, encompassing a wide range of formal theorems from various mathematical domains.
|
||||
Subsequently, we employ GRPO to enhance whole-proof generation through online reinforcement learning.
|
||||
Upon developing the DeepSeek-Prover-V1.5 model, we introduce RMaxTS, a variant of Monte-Carlo tree search, to improve problem-solving capabilities via large-scale search with extensive exploration.
|
||||
These components form a comprehensive pipeline for training an LLM-based proof assistant, enabling DeepSeek-Prover-V1.5 to achieve significant improvements over DeepSeek-Prover-V1.
|
||||
The framework of DeepSeek-Prover-V1.5 is designed to establish an AlphaZero-like pipeline for formal theorem proving.
|
||||
The use of expert iteration and synthetic data mirrors the core trial-and-error loop of reinforcement learning, with the compiler oracle serving as the world model to provide environmental supervision.
|
||||
Within the RL paradigm, the integrated tree search module has proven to be highly effective in advancing superhuman performance across various domains ( Silver et al., 2016 ; Fawzi et al., 2022 ; Lutz et al., 2023 ) .
|
||||
In developing DeepSeek-Prover-V1.5, we focus on the exploration aspect of RL, introducing RMaxTS to diversify the generation of proof steps.
|
||||
However, the exploitation aspect, specifically the problem of proof search, remains unexplored.
|
||||
A promising future direction is training a critic model to assess incomplete proofs and prune search branches.
|
||||
Such a partial-proof critic model would implicitly perform temporal credit assignment ( Sutton, 1984 ) , decomposing proof-level feedback into step-wise value differences ( Arjona-Medina et al., 2019 ) .
|
||||
Developing critic models for assessing long planning paths and providing guidance rewards presents a crucial and challenging problem ( Ng and Russell, 2000 ; Sorg et al., 2010 ) that warrants further investigation.
|
||||
Finally, recent work has progressed beyond proving individual theorems to addressing real-world theory proving within complex, multi-theorem Lean files ( Hu et al., 2024 ) . This shift is a natural extension of our whole-proof generation approach. Our observations indicate that the current model already possesses some understanding of file-level context. Moving forward, we will focus on enhancing this aspect to support cutting-edge Lean mathematical formalization developers with our language model advancements.
|
||||
@@ -0,0 +1,607 @@
|
||||
TITLE: Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning
|
||||
ARXIV: 2408.14158
|
||||
AUTHORS: Wei An , Xiao Bi , Guanting Chen , Shanhuang Chen , Chengqi Deng , Honghui Ding , Kai Dong , Qiushi Du , Wenjun Gao , Kang Guan , Jianzhong Guo , Yongqiang Guo , Zhe Fu , Ying He , Panpan Huang , Jiashi Li , Wenfeng Liang , Xiaodong Liu , Xin Liu , Yiyuan Liu , Yuxuan Liu , Shanghao Lu , Xuan Lu , Xiaotao Nie , Tian Pei , Junjie Qiu , Hui Qu , Zehui Ren , Zhangli Sha , Xuecheng Su , Xiaowen Sun , Yixuan Tan , Minghui Tang , Shiyu Wang , Yaohui Wang , Yongji Wang , Ziwei Xie , Yiliang Xiong , Yanhong Xu , Shengfeng Ye , Shuiping Yu , Yukun Zha , Liyue Zhang , Haowei Zhang , Mingchuan Zhang , Wentao Zhang , Yichao Zhang , Chenggang Zhao , Yao Zhao , Shangyan Zhou , Shunfeng Zhou , Yuheng Zou
|
||||
DATE: [Submitted on 26 Aug 2024 ( v1 ), last revised 31 Aug 2024 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning Thanks:
|
||||
## Abstract
|
||||
## Index Terms:
|
||||
## I Introduction
|
||||
## Fire-Flyer AI-HPC Architecture
|
||||
## Key Technical Topics in our Architecture
|
||||
## Stability and Robustness
|
||||
## Discussion and Future Work
|
||||
## II Background
|
||||
## II-A Evolution of Deep Learning
|
||||
## II-B Challenges and Solutions in Models Training
|
||||
## II-B 1 Efficiency
|
||||
## II-B 2 Stability
|
||||
## II-C HPC and AI Clusters of This Era
|
||||
## II-C 1 HPC Inadequacies for AI Training
|
||||
## II-C 2 GPU based HPC
|
||||
## II-C 3 GPU Clusters of Large Companies
|
||||
## II-C 4 AI DSA Clusters
|
||||
## II-C 5 Cloud Service Providers
|
||||
## II-D Challenges in AI Infrastructure
|
||||
## III Fire-Flyer 2: Our Approach For Deep Learning and Early LLM Training
|
||||
## III-A Fire-Flyer 2: PCIe A100 GPU Architecture
|
||||
## III-B Network Topology: Two-Layer Fat-Tree with Storage and Computation Integrated
|
||||
## III-C Cost Performance of Our Architecture
|
||||
## IV HFReduce: Hardware Software Co-Design in Network
|
||||
## IV-A HFReduce Algorithm Steps
|
||||
## IV-B Advantages of HFReduce over NCCL
|
||||
## IV-B 1 Reduced PCIe Bandwidth Consumption
|
||||
## IV-B 2 No GPU Kernel Overhead
|
||||
## IV-C Performance Improvements: HFReduce with NVLink
|
||||
## IV-D Deep Analysis of HFReduce
|
||||
## IV-D 1 Key Technical Strategies in Implementation
|
||||
## IV-D 2 HFReduce Overcomes Limitations of EPYC Rome CPU
|
||||
## IV-D 3 Bottlenecks of HFReduce
|
||||
## V HaiScale: Special Optimization for Deep Leaning Models Training
|
||||
## V-A HaiScale DDP Overlap AllReduce in Training
|
||||
## V-B LLMs Training Optimization
|
||||
## V-B 1 NVLink Bridge Enables Tensor Parallel between PCIe GPUs
|
||||
## V-B 2 Pipeline Parallelism Optimization in PCIe Architecture
|
||||
## V-B 3 Fully Sharded Data Parallel (FSDP)
|
||||
## V-C Summary
|
||||
## VI Advanced Cost-Effective and Co-Design Optimizations
|
||||
## VI-A Ensuring Minimal Congestion in Our Computation-Storage Integrated Network
|
||||
## VI-A 1 Divergence of Different Traffics
|
||||
## VI-A 2 Topology Adjustment and Route Optimization
|
||||
## VI-A 3 NCCL Optimization
|
||||
## VI-A 4 Network Tuning in 3FS
|
||||
## VI-B High-Throughput Distributed File System: 3FS
|
||||
## VI-B 1 Overview
|
||||
## VI-B 2 3FS Storage Node Hardware
|
||||
## VI-B 3 Key Techinical Points of 3FS
|
||||
## VI-B 4 3FS-KV
|
||||
## VI-C HAI Platform: a Time-Sharing Scheduling Platform
|
||||
## VII Stability and Robustness
|
||||
## VII-A Checkpoint Manager
|
||||
## VII-B Validator
|
||||
## VII-C Hardware Failures Characterization in Fire-Flyer 2 AI-HPC
|
||||
## VII-C 1 GPU Xid Error
|
||||
## VII-C 2 Network Flash Cut
|
||||
## VIII Discussion
|
||||
## VIII-A Discussion on Congestion Control in RDMA Networks
|
||||
## VIII-B Discussion about NVLink Technology Choices
|
||||
## VIII-C Maintaince Cost Overview
|
||||
## VIII-C 1 Construction Cost
|
||||
## VIII-C 2 Power Consumption
|
||||
## VIII-C 3 Operation Cost
|
||||
## VIII-D Stability Compared with Other Architectures
|
||||
## IX Future Work
|
||||
## Future Arch and Integration with New GPU Models
|
||||
## X Conclusions
|
||||
## Acknowledgment
|
||||
|
||||
=== BODY (head) ===
|
||||
Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
I Introduction
|
||||
II Background
|
||||
II-A Evolution of Deep Learning
|
||||
II-B Challenges and Solutions in Models Training
|
||||
II-B 1 Efficiency
|
||||
II-B 2 Stability
|
||||
II-C HPC and AI Clusters of This Era
|
||||
II-C 1 HPC Inadequacies for AI Training
|
||||
II-C 2 GPU based HPC
|
||||
II-C 3 GPU Clusters of Large Companies
|
||||
II-C 4 AI DSA Clusters
|
||||
II-C 5 Cloud Service Providers
|
||||
II-D Challenges in AI Infrastructure
|
||||
III Fire-Flyer 2: Our Approach For Deep Learning and Early LLM Training
|
||||
III-A Fire-Flyer 2: PCIe A100 GPU Architecture
|
||||
III-B Network Topology: Two-Layer Fat-Tree with Storage and Computation Integrated
|
||||
III-C Cost Performance of Our Architecture
|
||||
IV HFReduce: Hardware Software Co-Design in Network
|
||||
IV-A HFReduce Algorithm Steps
|
||||
IV-B Advantages of HFReduce over NCCL
|
||||
IV-B 1 Reduced PCIe Bandwidth Consumption
|
||||
IV-B 2 No GPU Kernel Overhead
|
||||
IV-C Performance Improvements: HFReduce with NVLink
|
||||
IV-D Deep Analysis of HFReduce
|
||||
IV-D 1 Key Technical Strategies in Implementation
|
||||
IV-D 2 HFReduce Overcomes Limitations of EPYC Rome CPU
|
||||
IV-D 3 Bottlenecks of HFReduce
|
||||
V HaiScale: Special Optimization for Deep Leaning Models Training
|
||||
V-A HaiScale DDP Overlap AllReduce in Training
|
||||
V-B LLMs Training Optimization
|
||||
V-B 1 NVLink Bridge Enables Tensor Parallel between PCIe GPUs
|
||||
V-B 2 Pipeline Parallelism Optimization in PCIe Architecture
|
||||
V-B 3 Fully Sharded Data Parallel (FSDP)
|
||||
V-C Summary
|
||||
VI Advanced Cost-Effective and Co-Design Optimizations
|
||||
VI-A Ensuring Minimal Congestion in Our Computation-Storage Integrated Network
|
||||
VI-A 1 Divergence of Different Traffics
|
||||
VI-A 2 Topology Adjustment and Route Optimization
|
||||
VI-A 3 NCCL Optimization
|
||||
VI-A 4 Network Tuning in 3FS
|
||||
VI-B High-Throughput Distributed File System: 3FS
|
||||
VI-B 1 Overview
|
||||
VI-B 2 3FS Storage Node Hardware
|
||||
VI-B 3 Key Techinical Points of 3FS
|
||||
VI-B 4 3FS-KV
|
||||
VI-C HAI Platform: a Time-Sharing Scheduling Platform
|
||||
VII Stability and Robustness
|
||||
VII-A Checkpoint Manager
|
||||
VII-B Validator
|
||||
VII-C Hardware Failures Characterization in Fire-Flyer 2 AI-HPC
|
||||
VII-C 1 GPU Xid Error
|
||||
VII-C 2 Network Flash Cut
|
||||
VIII Discussion
|
||||
VIII-A Discussion on Congestion Control in RDMA Networks
|
||||
VIII-B Discussion about NVLink Technology Choices
|
||||
VIII-C Maintaince Cost Overview
|
||||
VIII-C 1 Construction Cost
|
||||
VIII-C 2 Power Consumption
|
||||
VIII-C 3 Operation Cost
|
||||
VIII-D Stability Compared with Other Architectures
|
||||
IX Future Work
|
||||
X Conclusions
|
||||
References
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2408.14158v2 [cs.DC] 31 Aug 2024
|
||||
## Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning Thanks:
|
||||
Authors are listed in alphabetical order of their surnames.
|
||||
PubID: pubid:
|
||||
Wei An, Xiao Bi, Guanting Chen, Shanhuang Chen,Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du
|
||||
Affiliation:
|
||||
Wenjun Gao, Kang Guan, Jianzhong Guo, Yongqiang Guo, Zhe Fu, Ying He, Panpan Huang, Jiashi Li
|
||||
Affiliation:
|
||||
Wenfeng Liang, Xiaodong Liu, Xin Liu, Yiyuan Liu, Yuxuan Liu, Shanghao Lu, Xuan Lu, Xiaotao Nie
|
||||
Affiliation:
|
||||
Tian Pei, Junjie Qiu, Hui Qu, Zehui Ren, Zhangli Sha, Xuecheng Su, Xiaowen Sun, Yixuan Tan, Minghui Tang
|
||||
Affiliation:
|
||||
Shiyu Wang, Yaohui Wang, Yongji Wang, Ziwei Xie, Yiliang Xiong, Yanhong Xu, Shengfeng Ye, Shuiping Yu
|
||||
Affiliation:
|
||||
Yukun Zha, Liyue Zhang*, Haowei Zhang, Mingchuan Zhang, Wentao Zhang, Yichao Zhang, Chenggang Zhao
|
||||
Affiliation:
|
||||
Yao Zhao, Shangyan Zhou, Shunfeng Zhou, Yuheng Zou
|
||||
Affiliation:
|
||||
Affiliation: DeepSeek-AI, Beijing, China
|
||||
Affiliation: research@deepseek.com
|
||||
## Abstract
|
||||
The rapid progress in Deep Learning (DL) and Large Language Models (LLMs) has exponentially increased demands of computational power and bandwidth. This, combined with the high costs of faster computing chips and interconnects, has significantly inflated High Performance Computing (HPC) construction costs. To address these challenges, we introduce the Fire-Flyer AI-HPC architecture, a synergistic hardware-software co-design framework and its best practices. For DL training, we deployed the Fire-Flyer 2 with 10,000 PCIe A100 GPUs, achieved performance approximating the DGX-A100 while reducing costs by half and energy consumption by 40%. We specifically engineered HFReduce to accelerate allreduce communication and implemented numerous measures to keep our Computation-Storage Integrated Network congestion-free. Through our software stack, including HaiScale, 3FS, and HAI-Platform, we achieved substantial scalability by overlapping computation and communication. Our system-oriented experience from DL training provides valuable insights to drive future advancements in AI-HPC.
|
||||
## Index Terms:
|
||||
High Performance Computing, Cost-Effective, All-Reduce, Best Practices, Deep Learning, Machine Learning, Large Language Models, Artificial Intelligence Infrastructure
|
||||
## I Introduction
|
||||
In recent years, Deep Learning (DL) [ 1 ] has developed rapidly and is widely used in image recognition, speech recognition, content generation, autonomous driving, and other areas. The rapid development of DL is fundamentally tied to the support rendered by data. Training with copious amounts of data demands massive computational resources. Relying on Moore’s Law [ 2 ] , computer speeds double every two years on average, but the pace of DL far exceeds this speed. In particular, Large Language Models (LLMs) [ 3 , 4 , 5 , 6 , 7 ] that have become popular in recent years have exploded the demand for computational resources and memory. The parameters of LLMs can reach tens to thousands of billions, requiring hundreds or thousands of GPUs for training. Although LLM training is challenging, the emerging capabilities resulting from more parameters have shown the benefits of continued model expansion. Since then, researchers have gone down the path of making models bigger and never looked back. To acquire more computational resources, people have to expand more nodes. This leads to a surge in the cost of building AI infrastructure. How to reduce the cost of new data centers and how to build cost-effective clusters are also hot and challenging problems. Moreover, more nodes lead to higher energy consumption, which contradicts this era’s goal of reducing carbon emissions and achieving carbon neutrality. Reducing energy consumption is also a challenging problem.
|
||||
In this paper, leverage our practical experience accumulated over the years to propose cost-effective strategies cost-effective strategies for constructing AI-HPC systems suitable for deep learning and LLMs.
|
||||
## Fire-Flyer AI-HPC Architecture
|
||||
We have deployed a cluster composed of 10,000 PCIe A100 GPUs for Deep Learning training purposes. Details about the GPU nodes and network topology are provided in Section III-A , where we compare our architecture to the NVIDIA DGX-A100 [ 8 ] in terms of cost-effectiveness and lower CO 2 emissions. In contrast, we must invest more in software optimizations to address the performance challenges of the PCIe architecture. The following sections will discuss about software-hardware co-design.
|
||||
## Key Technical Topics in our Architecture
|
||||
•
|
||||
Network Co-Design : The Two-Layer Fat-Tree Network [ 9 ] integrates storage and computation network, as shown in Section III-B . The entire network is divided into two zones, and the platform supports cross-zone tasks. To prevent congestion, we employed various network tunings detailed in Section VI-A .
|
||||
•
|
||||
HFReduce : Achieves computation-communication overlap via asynchronous allreduce on the CPU, outperforming NVIDIA Collective Communications Library (NCCL) [ 10 ] on our PCIe architecture, as discussed in Section IV .
|
||||
•
|
||||
HaiScale : As described in Section V , optimizes parallelism methods for our PCIe architecture, such as Data Parallelism (DP), Pipeline Parallelism (PP) [ 11 ] [ 12 ] , Tensor Parallelism (TP) [ 13 ] [ 14 ] , Experts Parallelism (EP) [ 15 , 16 , 17 ] , Fully Sharded Data Parallel(FSDP) [ 18 ] and Zero Redundancy Optimizer(ZeRO) [ 19 ] .
|
||||
•
|
||||
3FS Distributed File System : Addresses I/O bottlenecks in big data AI tasks, configures with our communication and network tuning, reduces congestion in storage and computation integrated network topology, as detailed in Section VI-B .
|
||||
•
|
||||
HAI Platform [ 20 ] : Offers task scheduling, fault handling, and disaster recovery, enhancing utilization and reducing costs. It provides an open-box solution for deep-learning researchers. It is already open-sourced: https://github.com/HFAiLab/hai-platform
|
||||
## Stability and Robustness
|
||||
These are crucial topics in HPC. Our systems are equipped with robust mechanisms to handle hardware failures, minimizing downtime and impact on operations. These mechanisms, discussed in Section VII , include:
|
||||
•
|
||||
Disaster recovery through our checkpoint manager
|
||||
•
|
||||
A validator utility for detecting hardware failures
|
||||
•
|
||||
An overview of real hardware failure data from our cluster over the past year.
|
||||
We hope these insights will be beneficial to industry peers and researchers alike.
|
||||
## Discussion and Future Work
|
||||
In Section VIII , we address some common questions regarding PCIe architecture, such as congestion control, maintenance cost, and stability compared with other architectures. In Section IX , we propose the next generation of PCIe architecture, which is aimed at Mixture-of-Experts Large Language Models training and primarily utilizes multi-NICs and a Multi-Plane network.
|
||||
## II Background
|
||||
## II-A Evolution of Deep Learning
|
||||
The revolution in Machine Learning and Deep Learning began in 2012 with AlexNet [ 21 ] , which outperformed traditional methods in image classification, marking the onset of big data utilization and increased computational demands. The emergence of ResNet [ 22 ] , with its deeper layers, further broadened the horizons of image processing, truly bringing about the “deep” in “Deep Learning”. At the same time, big data-driven model training nudged the evolution of data storage technology, leading to the advent of all-flash SSD distributed file systems.
|
||||
Fast forward to 2017, Google’s Transformer [ 23 ] made its grand entry, introducing the concept of ”Attention is all you need”, shaking up the field of Natural Language Processing (NLP). With the advent of more complex models like AlphaFold [ 24 ] and AlphaZero [ 25 ] highlighting the need for more computational power and memory, revealing the limitations of traditional FP64 / FP32 computing devices.
|
||||
Entering the 2020s saw the rise of LLMs as a game-changer in the AI sector. Research indicates that an upscale in the number of language model parameters and computational budget can significantly enhance model performance, given adequate training data. Consequently, despite requiring colossal computational resources, efforts are being made to train large models on tens or hundreds of billions or even trillions of parameters. Pioneering examples include GPT-3 [ 26 ] and PaLM [ 27 ] , which occupy close to 1TB of GPU Memory. Recognizing the potential, industry giants have set up large AI clusters to train LLMs while constantly investing in computational power chips.
|
||||
The shift towards the Mixture-of-Experts (MoE) Models [ 28 , 29 , 30 ] architecture starting from GPT-4 [ 7 ] , and the recent AI Generated Content (AIGC) multi-modal (Sora [ 31 ] ) has amplified the demand for memory and computational resources. However, as AI development outpaces hardware development, leading to skyrocketing training costs, adopting cost-saving solutions has become imperative.
|
||||
Figure 1 illustrates the exponential growth of computational power for DL. And as summarized in Figure 2 [ 32 ] , while AI’s demand for computational power is growing at 10x per year, Moore’s Law lags behind with hardware FLOPs increase at only 3.0x every two years, DRAM bandwidth at 1.6x, and interconnect bandwidth at 1.4x. This disparity necessitates more machines, raising DL training costs, particularly for LLM training, where the computational power required surpasses that of traditional HPC applications.
|
||||
Figure 1 : The Exponential Growth of Computational Power for Deep Learning.
|
||||
Figure 2 : Scaling of Peak Hardware FLOPS, and Memory/Interconnect Bandwidth.
|
||||
## II-B Challenges and Solutions in Models Training
|
||||
In Deep Learning training, a single task demands hundreds of GPUs and consumes substantial storage and network resources. This massive scale introduces system-level challenges:
|
||||
## II-B 1 Efficiency
|
||||
Firstly, achieving efficient training at this magnitude is crucial. Model FLOPs Utilization (MFU), which assesses the ratio of observed throughput to theoretical maximum throughput (assuming 100% peak FLOPS), serves as the standard metric for evaluating training efficiency. Training LLMs involves dividing models among GPUs that communicate extensively for progress. Besides communication, factors like operation optimization, data pre-processing, and GPU memory consumption significantly influence MFU. Multiple parallel strategies are employed to enhance efficiency:
|
||||
•
|
||||
Data Parallelism (DP): Models and optimizer states are replicated across multiple devices with data evenly distributed to all. For LLMs training, the Zero Redundancy Optimizer (ZeRO) [ 19 ] further enhances this method by sharding these states on each data parallel process and using allgather and reduce-scatter for parameter fetching and gradient calculation.
|
||||
•
|
||||
Pipeline Parallelism (PP): Each device holds a portion of the model layers with each training batch divided into micro-batches for pipeline execution. Efficient scheduling strategies like GPipe [ 11 ] , PipeDream 1F1B [ 12 ] , and Zero Bubble Pipeline Parallelism (ZBPP) [ 33 ] , are required to minimize “pipeline bubbles”.
|
||||
•
|
||||
Tensor Parallelism (TP): This involves placing a model layer on multiple GPUs that perform computations in parallel [ 13 ] [ 14 ] . It includes row-wise and column-wise parallelism, necessitating allgather and all2all for input splitting and output merging.
|
||||
•
|
||||
Expert Parallelism (EP) : MoE Models’ different expert models are distributed on different GPUs during MoE training [ 15 , 16 , 17 ] . The gate model selects tokens for allocation during input, with corresponding tokens sent to experts model via all2all communication.
|
||||
•
|
||||
Fully Sharded Data Parallel (FSDP) is an implementation based on the ZeRO Stage 3 algorithm [ 19 ] . FSDP partitions the model’s parameters, optimizer states, and gradients, distributing them across different GPUs, with each GPU retaining only 1 / n 1/n of the total. During forward propagation, FSDP performs an allgather operation to assemble the complete parameters, which are then released after the forward pass is completed. Similarly, during backward propagation, FSDP conducts an allgather operation to obtain the complete parameters, followed by backward computation to calculate gradients. It then performs a reduce-scatter operation to synchronize gradients across all GPUs, resulting in each GPU holding 1 / n 1/n of the reduced gradients. Finally, FSDP updates the 1 / n 1/n parameters using each GPU’s 1 / n 1/n gradients and optimizer states. FSDP reduces GPU memory usage by maintaining only 1 / n 1/n of the parameters, gradients, and optimizer states on each GPU, enabling training of larger-scale models.
|
||||
There are additional strategies and algorithms to accelerate training or reduce memory usage, such as Activation Recomputation [ 34 , 35 ] , as well as enhanced communication and computation overlap during parallelism [ 36 ] , among others.
|
||||
## II-B 2 Stability
|
||||
The second challenge is achieving high-stability training at scale, i.e., maintaining efficient training throughout the process. Stability is vital from a production standpoint as training a big model with a trillion tokens may span several weeks. In DL training, stragglers and hardware failures are common occurrences rather than outliers. Stragglers can decelerate tasks involving hundreds of GPUs, emphasizing the importance of stability and task recovery time.
|
||||
## II-C HPC and AI Clusters of This Era
|
||||
## II-C 1 HPC Inadequacies for AI Training
|
||||
Traditional supercomputers such as TianHe-2A [ 37 ] , Stampede 2 [ 38 ] , and Sunway TaihuLight [ 39 ] primarily focus on double precision calculations and do not support the FP16 precision, rendering them unsuitable for DL training. Fugaku [ 40 ] , despite its high performance, does not support tensor GEMM acceleration, a key component for DL workloads. Although these supercomputers may not be well-suited for DL training, their robust high-performance networks and extensive experience in large-scale cluster construction offer valuable insights and lessons for subsequent researchers.
|
||||
## II-C 2 GPU based HPC
|
||||
Supercomputers like Frontier [ 41 ] , Aurora [ 42 ] , Summit [ 43 ] and Perlmutter [ 44 ] utilize high-performance GPUs to tackle large-scale computations. It’s worth mentioning that Perlmutter utilizes an all-flash storage system, achieving a peak bandwidth of 5TB/s. Indeed, conducting DL training on these GPU-based HPCs yields significant performance.
|
||||
## II-C 3 GPU Clusters of Large Companies
|
||||
Meta, formerly Facebook, has developed its AI-HPC using a software-hardware co-design approach, with one system utilizing IB and another employing RoCE [ 45 , 46 ] . ByteDance initially implemented a DL cluster with a mix of CPU and PCIe GPU [ 47 ] . However, with the advent of the LLMs era, they adopted an architecture similar to DGX, building a cluster with over 10,000 GPUs [ 48 ] . Alibaba has developed its own HPN network [ 49 ] for LLMs training using NVIDIA H800 GPUs. NVIDIA also has its own AI-HPC Eos [ 50 ] , which will feature 576 DGX H100 systems totaling 4,608 H100 GPUs. While this will provide a considerable boost in computational power for AI tasks, the high cost of DGX systems raises questions about economic viability.
|
||||
## II-C 4 AI DSA Clusters
|
||||
Custom-designed AI DSA (Domain Specific Architecture) accelerators like Google’s TPU [ 51 ] , utilize highly advanced optical switch reconfigurable networks. Alternatives to traditional GPU setups, like Intel Habana Gaudi [ 52 ] , are also available. Teslahas introduced the Dojo [ 53 , 54 ] supercomputer , which uses System on Wafer technology to build an entire silicon wafer as a single chip. Huawei has designed the Ascend AI DSA chip [ 55 , 56 ] , which remains competitive with NVIDIA, as noted by NVIDIA CEO Jensen Huang. These accelerators are tailored for efficient execution of AI workloads, offering specialized features to optimize model training and inference. However, their software ecosystems, while progressing, still require further development to match the maturity of NVIDIA’s offerings.
|
||||
## II-C 5 Cloud Service Providers
|
||||
Cloud service providers, such as Azure, offer flexible and scalable resources for AI training. Despite their convenience and easy accessibility, the costs can accumulate significantly over time. For long-term projects spanning around two years, these costs could amount to purchasing an entire dedicated cluster. Therefore, this option may not be the most economical choice for extensive AI computations.
|
||||
Figure 3 : Size of Model Parameter and Accelerator Memory
|
||||
## II-D Challenges in AI Infrastructure
|
||||
As models continue to grow larger, DL training requires thousands of GPUs. Additionally, researchers often need to train multiple models simultaneously. Therefore, a cluster with at least tens of thousands of GPUs can meet the needs of AI practitioners. In addition to increasing the scale of the cluster nodes and adding more GPUs, there’s also a need to find ways to save on the overall system construction costs. These costs include but are not limited to power support, cooling, networking, storage, fault handling, disaster recovery, etc. Building a cost-effective AI-HPC system is a significant challenge. The question of how to construct a high-performance, efficient, economical, and environmentally friendly HPC to meet AI training requirements has become a hot topic. Some works, such as [ 57 ] analyzes the construction of HPC, discussed the interconnection of heterogeneous clusters, cooling systems. AI applications such as Peng Cheng Cloud Brain II [ 58 ] discussed their strategies for building and improving AI computing power and cluster communication efficiency, which helped them achieve first place in IO500 and AIPerf rankings.
|
||||
Drawing from our extensive experience in Deep Learning spanning over the past decade, we have conducted considerable exploration in terms of cost-effectiveness. This work primarily discusses our practices for achieving cost-effectiveness and high-performance across different models and stages.
|
||||
## III Fire-Flyer 2: Our Approach For Deep Learning and Early LLM Training
|
||||
As mentioned in the Background Section, LLMs generally require significant memory resources. In contrast, many other models necessitate considerably less memory, as illustrated in Figure 3 . Popular models like ResNet [ 22 ] , Mask-RCNN [ 59 ] , BERT [ 60 ] , MAE [ 61 ] , among others, all have a parameter volume less than 1B, signifying relatively low memory requirements. Therefore, when designing a cluster primarily for deep learning model training, and with insights gleaned from our Fire-Flyer 1 experiments, we deemed it prudent to incorporate PCIe A100 GPUs, which proved to be sufficient during its construction in 2021.
|
||||
## III-A Fire-Flyer 2: PCIe A100 GPU Architecture
|
||||
Figure 4 : In-node Architecture: 8 PCIe GPUs and 1 InfiniBand (IB) NIC are directly connected to the CPU. Note that GPU5/6 share the same PCIe root port, while IB occupies one independently.
|
||||
In our training workloads, the bandwidth requirements for both storage IOs and computation communication across 8 NVIDIA PCIe A100 GPUs can be met by a single 200Gbps NVIDIA Mellanox ConnectX-6 (CX6) InfiniBand (IB) NIC. We employed the following computation node architecture , as shown in Figure 4 :
|
||||
Table I : Our Arch and DGX-A100 Server Hardware Details
|
||||
Our PCIe Arch
|
||||
DGX-A100
|
||||
CPU
|
||||
2 * AMD 32 Cores EPYC Rome/Milan CPU
|
||||
2 * AMD 64 Cores EPYC 7742 CPU
|
||||
Memory
|
||||
512GB 16-Channels DDR4-3200Mhz
|
||||
2048GB 16-Channels DDR4-3200Mhz
|
||||
GPU
|
||||
8 * PCIe-A100-40GB
|
||||
8 * SXM-A100-40GB
|
||||
NICs
|
||||
1 * Mellanox InfiniBand cx6 200Gbps NIC
|
||||
9 * Mellanox InfiniBand cx6 200Gbps NIC
|
||||
NVLINK
|
||||
600 GB/s between each pair of GPUs
|
||||
600 GB/s interconnect among all 8 GPUs
|
||||
•
|
||||
8 NVIDIA A100 PCIe GPUs and 1 Mellanox CX6 200Gbps IB NIC: directly connect to the CPU, without using a PCIe switch
|
||||
•
|
||||
IB NIC occupies a separate PCIe root complex, thus avoiding performance interference with the GPU.
|
||||
•
|
||||
Reserved the possibility of NVLink Bridge addition in design: As expected, when the LLM era arrived, we indeed added an NVLink Bridge between PCIe cards.
|
||||
Table I shows our arch details and compared with NVIDIA standard DGX-A100 server.
|
||||
## III-B Network Topology: Two-Layer Fat-Tree with Storage and Computation Integrated
|
||||
We selected the Fat-Tree [ 9 ] topology as our primary network architecture due to its exceptionally high bisection bandwidth, making it the preferred choice for AI-HPC and high-throughput storage environments. Although the Dragonfly topology [ 62 , 63 ] also offers comparable cost-effectiveness and performance, its lack of sufficient bisection bandwidth makes it unsuitable for our integrated storage and computation network design . At the time of implementation, various RoCE (RDMA over Converged Ethernet) [ 64 ] technologies were not as mature as they are today, so we opted for InfiniBand (IB) as our network solution. Mellanox QM8700 InfiniBand Switch, offering 40 ports at 200 Gbps, was utilized. Our cluster, consisting of 10,000 A100 GPUs, includes approximately 1,250 GPU compute nodes and nearly 200 storage servers, although a Two-Layer Fat-Tree can accommodate up to 800 nodes (configured with 20 spine switches and 40 leaf switches).
|
||||
Figure 5 : Network Topology: Two complete Two-Layer Fat-Tree connected together.
|
||||
To reduce costs, , we opted for a two-zone network configuration instead of a three-layer Fat-Tree solution, as shown in Figure 5 . Each zone consists of an 800-port Fat-Tree connected to approximately 600 GPU compute nodes. Each storage server equipped with two IB NICs, respectively connected to different zones, hence all GPU compute nodes could share a set of storage services. Additionally, the two zones are interconnected with a limited number of links. Our HAI Platform scheduling strategy ensured that cross-zone computing tasks were limited to one at most. Whether using NCCL [ 10 ] or our in-house developed communication library HFReduce, it can be run across zones by using a double binary tree algorithm [ 65 ] . Our scheduler ensures that in this topology, only one pair of nodes communicates across zones. Consequently, even tasks requiring all nodes can efficiently run on the entire Fire-Flyer 2 AI-HPC.
|
||||
## III-C Cost Performance of Our Architecture
|
||||
Compared to the NVIDIA DGX-A100 [ 8 ] architecture, our approach using PCIe A100 achieves approximately 83% of the performance in TF32 and FP16 General Matrix Multiply (GEMM) benchmarks. However, it offers substantial reductions in both costs and energy usage, achieving 60% of the GPU cost and energy consumption, as detailed in Table II .
|
||||
Contrasting with the DGX-A100 cluster, which necessitates a Three-Layer Fat-Tree encompassing 10,000 access points and involving 320 core switches, 500 spine switches and 500 leaf switches, amounting to 1,320 switches in total (as shown in Table III ), our architecture only requires 122 switches. This arrangement is significantly more cost-efficient. Even when compared to a similarly sized three-layer Fat-Tree network with 1,600 access points, which includes 40 core switches and 160 spine and leaf switches (totaling 200 switches), our design facilitates a saving of 40% in networking costs.
|
||||
Furthermore, by utilizing an 800-Ports Frame Switch, we have further reduced the cost of optical modules and cables. While there is a performance gap due to the inherent differences between PCIe card specifications and SXM, we generally achieved 80% of the DGX-A100 performance at merely 60% of the cost. Additionally, we managed to trim energy consumption by 40%, thereby reducing CO 2 emissions. In terms of cost-performance, we regard this approach as both effective and successful.
|
||||
Table II : A100 PCIe Compared to DGX-A100.
|
||||
Our Arch
|
||||
DGX Arch
|
||||
TF32 GEMM (TFLOPS/GPU)
|
||||
107
|
||||
131
|
||||
FP16 GEMM (TFLOPS/GPU)
|
||||
220
|
||||
263
|
||||
Relative Performance
|
||||
83%
|
||||
100%
|
||||
Node Relative Price
|
||||
60%
|
||||
100%
|
||||
Cost–Performance Ratio
|
||||
1.38
|
||||
1
|
||||
Power Consumption (Watts)
|
||||
2500
|
||||
4200
|
||||
Table III : Relative Cost Comparison.
|
||||
Our Arch
|
||||
PCIe Arch with Three Layer Fat-Tree
|
||||
DGX Arch
|
||||
Number of Switches
|
||||
122
|
||||
200
|
||||
1320
|
||||
Network Price
|
||||
350
|
||||
600
|
||||
4000
|
||||
Server Price
|
||||
11250
|
||||
11250
|
||||
19000
|
||||
Total Price
|
||||
11600
|
||||
11850
|
||||
23000
|
||||
## IV HFReduce: Hardware Software Co-Design in Network
|
||||
In large-scale deep learning training, allreduce is essential for aggregating gradients across GPUs. To optimize communication among PCIe GPUs in our architecture, we developed HFReduce, a library specifically designed for efficient allreduce operations. The core strategy of HFReduce, illustrated in Figure 6 , involves performing intra-node reduction first, followed by inter-node allreduce of the reduced data from the 8 GPUs within each node. This inter-node allreduce leverages a Double Binary Tree Algorithm [ 65 ] , akin to NCCL, and is pipelined by dividing data into chunks for transfer via Remote Direct Memory Access (RDMA), ensuring high performance. HFReduce is versatile and can be applied to any scenario requiring allreduce, as well as general reduce and broadcast operations.
|
||||
Figure 6 : HFReduce Schematic: 1) do intra-node reduce, 2) do inter-node allreduce by CPU, 3) finally transfer reduced data to GPU.
|
||||
## IV-A HFReduce Algorithm Steps
|
||||
Intra-node reduction, as shown in the Algorithm 1 :
|
||||
1.
|
||||
When the gradients data on the GPUs require allreduce, HFReduce asynchronously transfers these data to the CPU memory. This Device-To-Host (D2H) Transfer can utilize GDRCopy [ 66 ] for small data and MemCpyAsync for larger data.
|
||||
2.
|
||||
Upon the arrival of the gradients in memory, perform reduction add operation using CPU vector instructions.
|
||||
Inter-node reduction, as shown in the Algorithm 2 :
|
||||
1.
|
||||
Use the Double Binary Tree Algorithm [ 65 ] for inter-node allreduce, facilitating transfers between nodes using RDMA verbs implementation.
|
||||
2.
|
||||
Finally, the CPU returns reduced gradients to the GPU via PCIe (Host-To-Device Phase).
|
||||
The final Host-To-Device (H2D) Transfer can be optimized by utilizing GDRCopy to write data to four GPUs within the same NUMA node, effectively reducing reads from host memory by threefold compared to MemCpyAsync. This efficiency is achieved because GDRCopy can read data from host memory and temporarily cache it in the CPU caches, allowing data to be written to the four GPUs without additional reads from host memory.
|
||||
Algorithm 1 Intra Node Reduce.
|
||||
Data: Dg: data need to allreduce
|
||||
Result: Dc: data reduced in this node
|
||||
1
|
||||
Split Dg by Chunk_Size
|
||||
2
|
||||
for Dg-i in Splited-Dg do
|
||||
3
|
||||
// Transfer GPU Memory Data to CPU memory
|
||||
4
|
||||
Dc_i = MemCopyAsync Dg_i to CPU Mmeory
|
||||
5
|
||||
end for
|
||||
6
|
||||
for i in Splited_Count do
|
||||
7
|
||||
while every GPU’s D g _ i → D c _ i Dg\_i\rightarrow Dc\_i finished do
|
||||
8
|
||||
// Wait for chunk-i transfer finished in this node
|
||||
9
|
||||
for j in GPU_Count do
|
||||
10
|
||||
// do intra-node reduce
|
||||
11
|
||||
Dc_i += GPU-j’s Dc_i
|
||||
12
|
||||
end for
|
||||
13
|
||||
// do inter-node reduce
|
||||
14
|
||||
Do_Internode_reduce(Dc_i)
|
||||
15
|
||||
end while
|
||||
16
|
||||
end for
|
||||
Algorithm 2 Inter Node Reduce.
|
||||
Data: DL:local node reduced data by Algorithm 1
|
||||
Data: DR: received other node reduced data
|
||||
Result: Dg: reduced data transfer to GPU
|
||||
1
|
||||
// Pass 1: reduce data,individual thread
|
||||
2
|
||||
for i in Chunk_Size do
|
||||
3
|
||||
receive data D R _ i DR\_i from prev node
|
||||
4
|
||||
if D L _ i DL\_i ready then
|
||||
5
|
||||
DL_i +=DR_i // reduce data
|
||||
6
|
||||
end if
|
||||
7
|
||||
if Thread is root of Tree then
|
||||
8
|
||||
send DL_i to prev node // Dg_i finished, go parse 2
|
||||
9
|
||||
else
|
||||
10
|
||||
send DL_i to next node
|
||||
11
|
||||
end if
|
||||
12
|
||||
end for
|
||||
13
|
||||
//
|
||||
14
|
||||
// Pass 2: gather reduced data,individual thread
|
||||
15
|
||||
receive data D R _ i DR\_i from next node
|
||||
16
|
||||
for j in GPU_Count do
|
||||
17
|
||||
Dg_i = MemCopyAsync DR_i to GPU-j // Dg_i is allreduced
|
||||
18
|
||||
end for
|
||||
19
|
||||
send DL_i to prev node
|
||||
## IV-B Advantages of HFReduce over NCCL
|
||||
## IV-B 1 Reduced PCIe Bandwidth Consumption
|
||||
Let n n be the total number of GPUs involved in the communication. In NCCL’s ring topology, each unit of data needs to go through 2 n − 1 2n-1 transmissions, each consuming one unit of inbound bandwidth of one GPU and one unit of outbound bandwidth of another GPU. This means for a single unit of data, it consumes 2 n − 1 n \frac{2n-1}{n} unit of PCIe bidirectional bandwidth. In contrast, for each unit of data, HFReduce only requires one D2H and one H2D data transfer, only one unit of PCIe bidirectional bandwidth is consumed. In our machine architecture, the performance of NCCL is mainly limited by PCIe bandwidth. Therefore, HFReduce can achieve better performance than NCCL.
|
||||
## IV-B 2 No GPU Kernel Overhead
|
||||
HFReduce utilizes the GPU’s Copy Engine (CE) for PCIe asynchronous transfers. In contrast, NCCL’s allreduce operation requires GPU kernel execution, which can affect other computational kernels on the GPU. HFReduce achieves complete asynchrony with no overhead.
|
||||
As demonstrated in Figure 7(a) , HFReduce can reach a inter-node bandwidths of 6.3-8.1GB/s when performing allreduce with a data size of 186 MiB on the Fire-Flyer 2 AI-HPC, while NCCL’s inter-node bandwidth is only 1.6-4.8GB/s.
|
||||
(a) HFRudce and NCCL Allreduce Speed.
|
||||
(b) HFReduce with NVLink (Cross Fat-Tree Zone).
|
||||
Figure 7 : Strong Scalability: (a) Network Bandwidth when HFReduce and NCCL do allreduce test with 186MiB data, scale from 16 to 1440 GPUs. (b) HFReduce with NVLink, and Cross Fat-Tree Zone. Note that tasks utilizing fewer than 128 GPUs do not require cross-zone nodes and are restricted by platform defaults.
|
||||
## IV-C Performance Improvements: HFReduce with NVLink
|
||||
By installing the NVLink Bridge for PCIe A100 GPUs, efficient communication is enabled between paired GPUs via the 600 GB/s NVLink. To alleviate the memory bound issue of the original HFReduce, we implemented another allreduce pattern, termed HFReduce with NVLink . The core concept involves initially performing a reduction operation among GPUs interconnected by NVLink before passing the gradient to the CPU. Subsequently, when the CPU returns the result, it splits the result data and returns them to the paired GPUs connected by NVLink respectively, then performs allgather via NVLink. As illustrated in Figure 7(b) , HFReduce with NVLink achieves inter-node bandwidths exceeding 10 GB/s.
|
||||
## IV-D Deep Analysis of HFReduce
|
||||
## IV-D 1 Key Technical Strategies in Implementation
|
||||
•
|
||||
Using GDRCopy accelerate small data transfer in D2H, and educing reads from host memory by three times compared to MemCpyAsyn.
|
||||
•
|
||||
Intra-Node Reduction: CPU utilizes SIMD instructions and supports FP32 / FP16 / BF16 / FP8 datatypes.
|
||||
•
|
||||
NUMA Awareness: D2H destination memory is interleaved across two NUMA nodes for maximum bandwidth. Memory for CPU-added results and network-received data is bound to the IB-Nic’s NUMA node to minimize latency.
|
||||
•
|
||||
Inter-Node Reduce:Implements a Double Binary Tree allreduce algorithm [ 65 ] via ibverbs RDMA Write, avoiding additional overhead.
|
||||
## IV-D 2 HFReduce Overcomes Limitations of EPYC Rome CPU
|
||||
We consulted AMD and NVIDIA engineers to identify the root cause of NCCL’s suboptimal performance on PCIe architecture, particularly with EPYC Rome CPU servers. It was determined that the Rome CPUs do not support the chained write feature, which can significantly accelerate PCIe peer-to-peer (P2P) transfers between GPUs and IB NICs. Our tests indicate that the maximum bandwidth between the GPU and IB NIC on Rome CPUs is approximately 9 GiB/s, making the observed 4GB/s all-reduce bandwidth for NCCL understandable. HFReduce circumvents this limitation by utilizing the CPU for reduction and transferring data through IB and host memory.
|
||||
## IV-D 3 Bottlenecks of HFReduce
|
||||
When considering the total memory operations on a single node during HFReduce, several factors contribute to its performance limitations:
|
||||
1.
|
||||
D2H Phase requires 8 write operations.
|
||||
2.
|
||||
Intra-node Reduce Add Phase involves 8 read operations and 1 write operation.
|
||||
3.
|
||||
Inter-node Allreduce Phase: IB send demands 2 read operations, while IB receive requires 2 write operations, along with 1 read operation for reduce add.
|
||||
4.
|
||||
H2D Phase Utilizing GDRCopy can reduce this to only 2 read operations, whereas MemCopy necessitates 8 read operations.
|
||||
(a) HFReduce v.s. Torch DDP.
|
||||
(b) HaiScale v.s. Torch FSDP.
|
||||
Figure 8 : Weak Scalability: (a) Training VGG16, HFReduce compared to PyTorch DDP’s NCCL backend (b) Training GPT2-Medium HaiScale compared to Torch, both using FSDP.
|
||||
In total, the memory operations amount to 24 times the original data size in the GPU. A host equipped with 16 channels of DDR4-3200MHz memory can achieve a practical memory access speed of 320GB/s. Consequently, the theoretical maximum speed of HFReduce is approximately 13.3GB/s, but when considering the allreduce algorithm bandwidth and network bandwidth, this value realistically approximates 12GB/s. However, our tests only achieved slightly over 8GB/s.
|
||||
The root cause of this discrepancy is another limitation of the EPYC CPUs. As previously mentioned, our GPU5 and GPU6 are directly connected to the CPU via the same PCIe Root Complex Port (also known as the PCIe Host Bridge). In AMD EPYC Rome and Milan CPUs, the maximum bandwidth from the Root Complex Port to the CPU’s internal bus is about 37.5GB/s. Although a PCIe 4.0 x16 port can achieve over 27GB/s from GPU to CPU, when two GPUs transfer data concurrently, the bandwidth is limited to around 37GB/s. Furthermore, if bidirectional data transfer occurs simultaneously, this bandwidth decreases even further. As a result, HFReduce does not reach its theoretical speed.
|
||||
Employing NVLink with HFReduce offers a functional method to alleviate these bottlenecks. However, it is worth noting that the next-generation CPUs, such as the EPYC Genoa, still face issues with PCIe Host Bridge bandwidth, which cannot support two full-speed PCIe ports simultaneously. We hope AMD will address this issue in future iterations.
|
||||
## V HaiScale: Special Optimization for Deep Leaning Models Training
|
||||
## V-A HaiScale DDP Overlap AllReduce in Training
|
||||
HaiScale Distributed Data Parallel (DDP) is a training tool that utilizes HFReduce as its communication backend, in contrast to PyTorch’s DDP [ 67 ] which employs NCCL as its backend. During the backpropagation phase, HaiScale DDP performs an asynchronous allreduce operation on the computed gradients, allowing this communication to overlap with the computation involved in backpropagation.
|
||||
As previously mentioned, HFReduce does not depend on GPU Streaming Multiprocessors (SM) for reduction computation, enabling completely asynchronous allreduce without impacting performance. As shown in Figure 8(a) , training VGG16 model [ 68 ] with HFReduce takes only half the time compared to using Torch DDP’s NCCL backend, achieving nearly 88% parallel scalability when scale from 32 GPUs to 512.
|
||||
## V-B LLMs Training Optimization
|
||||
Our HaiScale framework various parallelism strategies for training large language models (LLMs), similar to Megagron [ 69 ] and DeepSpeed [ 70 ] . We have made specific engineering optimizations for our PCIe architecture across Data Parallelism (DP), Pipeline Parallelism (PP) [ 11 ] [ 12 ] , Tensor Parallelism (TP) [ 14 ] , Expert Parallelism (EP) [ 15 , 16 , 17 ] .
|
||||
## V-B 1 NVLink Bridge Enables Tensor Parallel between PCIe GPUs
|
||||
With the advent of LLMs, we integrated the NVLink Bridge into our system. This addition established a bandwidth of 600GB/s between each pair of GPUs, enabling more efficient when performing Tensor Parallelism.
|
||||
## V-B 2 Pipeline Parallelism Optimization in PCIe Architecture
|
||||
In our architecture, there is only one IB NIC for 8 GPUs on a single node, which can lead to network bandwidth contention during Pipeline Parallelism (PP). We solve this by configuring Data Parallelism (DP) rank, making the 8 GPUs on the same node belong to different DP ranks which staggers the timing of PP for each DP rank. As Figure 9(a) shown, when scaling from 64 GPUs to 512 GPUs, the step time of training LLaMa-13B [ 71 ] decreases from 64.118 seconds to 9.717 seconds, achieving a parallel efficiency of 91% .
|
||||
We also benchmarked the training performance of our DeepSeekMoE-16B model [ 72 ] on Fire-Flyer 2 AI-HPC. As shown in Figure 9(b) , scaling from 40 GPUs to 640 GPUs reduced the time per training step from 79.615 seconds to 6.535 seconds, achieving a parallel efficiency of 76.14%. Notably, with 320 GPUs, the step time was 10.71 seconds, resulting in a parallel efficiency of 92.92% , demonstrating excellent scalability.
|
||||
(a) Train LLaMa-13B
|
||||
(b) Train DeepSeekMoE-16B
|
||||
Figure 9 : Strong Scalability: (a) Train LLaMa-13B with a config of sequence length 2048, batch size 4096, pipeline parallel 4. (b)Train DeepSeekMoE-16B with a config of sequence length 4096, batch size 4608, pipeline parallel 10.
|
||||
## V-B 3 Fully Sharded Data Parallel (FSDP)
|
||||
Both HaiScale’s FSDP and PyTorch’s FSDP [ 18 ] are implementations based on the ZeRO Stage-3 algorithm [ 19 ] . The details of this implementation are already discussed in Section • ‣ II-B1 .
|
||||
HaiScale’s FSDP offers better engineering implementation, optimizing memory management to reduce fragmentation specific to model adjustments. And we overlap allgather and reduce-scatter communication with forward and backward computation, split the optimization step during backward propagation for enhanced overlap. As shown in Figure 8(b) , training GPT2-medium [ 73 ] , we achieve 95% parallel scalability when scaling from 16 to 128 GPUs. ompared to PyTorch’s FSDP, HaiScale’s FSDP reduces training time by nearly half.
|
||||
## V-C Summary
|
||||
Our AI-HPC design meets DL requirements, and with the addition of the NVLink Bridge, it meets the training needs of early-stage LLMs, reaching the utilization upper limit of PCIe GPUs. However, due to the inherent gap between PCIe card specifications and SXM, there is a certain performance discrepancy. Considering overall performance, basic setup cost, and energy consumption, we achieved 80% performance at half the cost. We believe Fire-Flyer 2 AI-HPC is a successful practice in terms of cost-effectiveness.
|
||||
## VI Advanced Cost-Effective and Co-Design Optimizations
|
||||
## VI-A Ensuring Minimal Congestion in Our Computation-Storage Integrated Network
|
||||
As previously stated, our cost-effective network integrated computation communication and storage traffics together. To achieve maximum bandwidth, it is essential to isolate interference between different types of traffic and control network congestion. In practice, we implemented the following measures:
|
||||
## VI-A 1 Divergence of Different Traffics
|
||||
In typical training tasks, there are four different types of traffic: HFReduce communication, NCCL communication, 3FS storage traffic, and other traffic. By using InfiniBand’s Service Level (SL) technology [ 74 ] [ 75 ] , we assign different value of SL when establishing connections between nodes and map SL to IB physical queues Virtual Lanes (VLs) [ 74 ] [ 75 ] . The use of Virtual Lanes ensures that flows in distinct lanes do not interfere with each other. Ultimately, we configured their proportions to implement traffic isolation, thereby preventing network congestion caused by Head-of-line (HOL) blocking [ 76 ] and different traffic collisions.
|
||||
## VI-A 2 Topology Adjustment and Route Optimization
|
||||
In high-throughput storage scenarios, there naturally exist many incast communication patterns, leading to certain congestion in the network. Under such circumstances, we observed that enabling adaptive routing would lead to more severe congestion spread in the network. Therefore, we opted for a static routing strategy. Based on the static routing scheme, to evenly disperse storage traffic into leaf → \rightarrow spine links, we distribute various nodes (storage, computation, management nodes) evenly disperse storage traffic into leaf → \rightarrow spine links.
|
||||
## VI-A 3 NCCL Optimization
|
||||
We adjusted the NCCL topology to route through the IB NIC and GPUs within the same NUMA node. This adjustment reduced PCIe congestion caused by CPU chiplet interconnects. Additionally, by using PCIe Relaxed Ordering [ 77 ] , we further reduced congestion and increased bandwidth.
|
||||
## VI-A 4 Network Tuning in 3FS
|
||||
3FS implements a request-to-send control mechanism to mitigate the congestion. Details are discussed in the next subsection, VI-B 3 Key Techinical Points of 3FS .
|
||||
## VI-B High-Throughput Distributed File System: 3FS
|
||||
## VI-B 1 Overview
|
||||
3FS is our in-house developed high performance distributed file system, akin to WekaFS [ 78 ] , DAOS [ 79 , 80 ] , and BeeGFS [ 81 ] . However, the design and implementation of 3FS specifically focus on fully utilizing the high IOPS and throughput of NVMe SSDs and the RDMA network.
|
||||
Table IV : Storage Node Hardware Details
|
||||
CPU
|
||||
1 * AMD 64 Cores EPYC 7742 CPU
|
||||
Memory
|
||||
512GB 8-Channels DDR4-3200Mhz
|
||||
NICs
|
||||
2 * Mellanox InfiniBand CX6 200Gbps NIC
|
||||
Data SSDs
|
||||
16 * 15.36TB PCIe 4.0x4
|
||||
## VI-B 2 3FS Storage Node Hardware
|
||||
In Fire-Flyer 2 AI-HPC, we deployed 180 storage nodes, as shown in Table IV , each node contains 16 PCIe 4.0 NVMe SSDs and 2 Mellanox CX6 200Gbps InfiniBand HCAs. With totally 360 * 200Gbps outbound InfiniBand HCAs, the system can total provide 9TB/s outbound bandwidh, and we actually achieved total read throughput of 8TB/s . The total 2880 NVMe SSDs provide over 20PiB storage space with an mirror data redundancy.
|
||||
## VI-B 3 Key Techinical Points of 3FS
|
||||
The 3FS system comprises four roles: cluster manager, meta service, storage service and client. Meta and storage services send heartbeats to cluster manager. All services and clients poll cluster configuration and service status from the manager. Multiple cluster managers are present, with one elected as the primary.
|
||||
File system meta data are stored in tables of a distributed key-value storage system. Each file or directory has a unique inode ID. The File inode/directory ID and meta data, such as file size and location information of the file content data, are stored as key-value pairs in the inode table. A separate directory entry table stores key-value pairs of ( p a r e n t _ d i r _ i n o d e _ i d , e n t r y _ n a m e ) : ( e n t r y _ i n o d e _ i d , … ) (parent\_dir\_inode\_id,entry\_name):(entry\_inode\_id,...) to support iterating entries in a directory and resolving file/directory paths. All states of meta services are persisted on the distributed key-value storage system. Several meta services run concurrently to handle meta requests from clients.
|
||||
The storage service has an implementation of Chain Replication with Apportioned Queries (CRAQ) [ 82 ] to provide strong consistency. CRAQ’s write-all-read-any approach helps to unleash the throughput and IOPS of all SSDs. File content are split into chunks, which are replicated over a chain of storage targets . A chain table contains an ordered set of chains. The meta service selects an offset in the chain table and a stripe size k for each file. The file chunks are assigned to the next k chains starting at the offset. To distribute read/write traffic evenly to all SSDs, each SSD serves multiple storage targets from different chains. The storage service runs on every storage node and manages a few storage targets.
|
||||
The storage network has a Fat-Tree topology that provides full bisection bandwidth. By design, each 3FS client can access every storage service. At peak load, incast congestion is observed on the client side. To mitigate this congestion, a request-to-send control mechanism is implemented in storage service and client [ 83 ] . After receiving a read request from a client, the service reads data from SSD and asks the client’s permission to transfer the data. The client limits the number of concurrent senders. When a storage service is granted the permission to transfer, it sends the data with a RDMA WRITE followed by a RDMA SEND to notify the client. The request-to-send control increases end-to-end IO latency but it’s required to achieve sustainable high throughput.
|
||||
## VI-B 4 3FS-KV
|
||||
3FS-KV is a shared-storage distributed data processing system built on top of 3FS, currently supporting three models: key-value, message queue, and object storage. It supports read-write separation and on-demand startup, allowing it to fully leverage the extremely high I/O throughput provided by 3FS. 3FS-KV supports DeepSeek’s KV Context Caching on Disk technology [ 84 ] , which reduces the cost of LLM serving by an order of magnitude.
|
||||
## VI-C HAI Platform: a Time-Sharing Scheduling Platform
|
||||
The principle of time-sharing scheduling is applied to cluster resource management. Users submit tasks, such as running Python / bash code, starting development containers, etc., and the platform interrupts and loads tasks according to current resource requirements, cluster busyness, etc. Task code needs to follow the platform coding rules to ensure that it can be continued from breakpoints, with the specific process as follows:
|
||||
•
|
||||
Accepting the interruption signal from the cluster;
|
||||
•
|
||||
Saving checkpoints (model parameters, optimizer parameters, etc.);
|
||||
•
|
||||
Notifying the cluster of interruptions;
|
||||
•
|
||||
Recovering from the checkpoint and continuing to run.
|
||||
The cluster deploying HAI Platform does not pool GPU resources, but classifies and marks them based on computing nodes as basic units, according to resource types, network areas, etc. The HAI Platform encourages users to fully utilize multiple GPUs simultaneously for parallel training, facilitating 99% utilization.
|
||||
## VII Stability and Robustness
|
||||
## VII-A Checkpoint Manager
|
||||
Training LLMs can span several months, during which unavoidable hardware failures may cause training interruptions. To minimize recovery time and support the HAI Platform’s interrupt and recovery operations, we developed a checkpoint manager. Additionally, the substantial size of LLM checkpoints necessitated an efficient method for saving and loading them, leveraging the high throughput of 3FS. The checkpoint manager includes the following components:
|
||||
•
|
||||
Parameters and optimization states are divided into chunks and written to 3FS using the 3FS batch write API, which is significantly faster than normal writes, achieving over 10 GiB/s per node. This enables saving to be completed in just a few seconds.
|
||||
•
|
||||
Parameters and optimization states are asynchronously transferred from GPU to CPU host memory, with checkpoint saving performed periodically (typically every 5 minutes).
|
||||
•
|
||||
During the saving process, each tensor is recorded with its index and the offset within the checkpoint., which makes the location of tensors more convenient during the loading process. With the 3FS batch read API, a loading process can be completed in just a few seconds.
|
||||
Thanks to the high write throughput of 3FS, periodic saving operations can be completed asynchronously in a matter of seconds, without impacting the training process. In the event of hardware failures that interrupt training, only the last 5 minutes of progress are lost. For a cluster with thousands of nodes, this overhead from disaster recovery is minimal.
|
||||
## VII-B Validator
|
||||
The best way to enhance device stability is to identify issues before they occur. Therefore, we have developed a set of validator tools to verify whether the hardware is functioning correctly. The platform’s automatic operation and maintenance system runs the validator program weekly on nodes to verify their proper functionality. It removes the faulty nodes from the scheduling platform, ensuring that all scheduled nodes are operational.
|
||||
Diagnosing tools like hostping [ 85 ] also integrated in our platform, but to find root cause of Hardware Failures is still hard work for operation teams.
|
||||
The validator mainly consists of the following parts:
|
||||
•
|
||||
Checking hardware frequency, link speed, and link status.
|
||||
•
|
||||
Testing CPU stress and memory bandwidth.
|
||||
•
|
||||
GPU Memory test: This involves checking each byte of GPU memory to ensure no data corruption has occurred.
|
||||
•
|
||||
Running GEMM with full GPU memory occupancy, which can simultaneously check whether there are any operational logic faults in the GPU chip.
|
||||
•
|
||||
Intra-node allreduce test: checking NVLink bandwidth through upper-level applications.
|
||||
•
|
||||
Storage bandwidth stress test to make sure storage is functioning normally.
|
||||
## VII-C Hardware Failures Characterization in Fire-Flyer 2 AI-HPC
|
||||
In supercomputers and data centers, hardware failures and chip errors can lead to floating-point overflow, non-convergence, or slow convergence during model training [ 86 ] . This paper [ 87 ] even directly points out that there is a substantial amount of Silent Data Corruption in data center processors, ultimately leading to a variety of complex issues that are difficult to replicate and locate. Indeed, in our practice, we have encountered computational errors and GPU memory errors not detected by Error Correction Code (ECC), which led to models’ gradnorm spikes, loss explosions and even non-convergence. How to tackle these hardware failures, promptly identify and categorize them, is a key issue to improve the online rate and overall utilization of cluster nodes.
|
||||
Table V : Type of GPU Xid Errors and Its Causes
|
||||
Xid Errors
|
||||
Analysis
|
||||
Software Causes: Xid_13/31 Xid_43/45
|
||||
Triggered by application programs, software-related Xid messages may indicate anomalies in GPU memory affecting code and data segments. However, it’s crucial to consider other information for a comprehensive hardware functionality assessment.
|
||||
NVLink Error: Xid 74
|
||||
Xid74 indicates errors in NVLink. For PCIe A100, it’s mainly occurred on the NVLink Bridge between two GPUs. Its occurrence rate is several orders of magnitude higher than other hardware faults. Apart from stress testing to exclude those that are constantly repeating errors, there isn’t a good way to avoid the occurrence of Xid74 issues.
|
||||
Memory ECC Error: Xid_63/64 Xid_94/95
|
||||
Triggered when the GPU handles memory ECC errors on the GPU. With the introduction of row remapping technology in A100, most instances can be resolved by simply resetting the GPU to retain optimal performance.
|
||||
Uncorrectable GPU Failures: Xid_44/48 Xid_61/62/69/79
|
||||
Thease failures mean an uncorrectable error occurs on the GPU, which is also reported back to the user application. A GPU reset or node reboot is needed to clear this error.
|
||||
Other Failures: Xid 119
|
||||
Xid119 means GPU GSP module failed.
|
||||
These failures need to do fieldiag test, and most need to
|
||||
RMA.
|
||||
## VII-C 1 GPU Xid Error
|
||||
An Xid error [ 88 ] is a general GPU fault message that originates from the NVIDIA driver, logged into the kernel or event log of the operating system. We have categorized various types of Xid errors and analyzed the potential causes that may lead to such errors, as shown in the Table V .
|
||||
Table VI in Appendix: Supplementary Characterization shows the Xid errors that have occurred in our Fire-Flyer 2 AI-HPC over the past year. In our PCIe-based system, Xid74, also known as NVLink errors, account for a significant proportion, comprising 42.57% of the total. This high frequency is due to the inherent fault rate of NVLink Bridge connectors, amplified by our extensive use of thousands of GPUs.
|
||||
Software-related errors such as Xid13, Xid31, Xid43, and Xid45 suggest possible illegal memory access or instructions in user code. Notably, illegal memory access (Xid 43) accounts for 33.48%, highlighting the need for improved memory management. However, these errors may also result from memory data corruption, so hardware faults should be considered if software bugs are ruled out.
|
||||
Additionally, GPU Memory ECC Errors, such as Xid63, Xid64, Xid94, and Xid95, require special attention as they represent about 2% of the total. Figure 10 illustrates the statistics related to ECC errors in our production cluster over the past six months. It is evident that the number of GPU ECC faults considerably surpasses those from the CPU. Therefore, it is crucial to promptly address GPU ECC faults to ensure that application performance and accuracy remain unaffected.
|
||||
Figure 10 : Trends of Memory and Network Failures from 2023 to 2024: “Main Memory” indicates CPU Memory ECC errors, “Network” indicates Network Flash Cuts, and “xids” are related to GPU memory ECC errors. Raw data is available in Appendix: Supplementary Characterization , Table VII
|
||||
.
|
||||
## VII-C 2 Network Flash Cut
|
||||
In addition to CPU and GPU faults, network device malfunctions represent a significant portion of hardware issues. As shown in Figure 10 , IB link failures account for 30% of hardware faults excluding Xid74. Network flash cuts can lead to application communication disruption, even task failures. Since most tasks run on multiple nodes, an issue on a single node can impact many others, further reducing cluster utilization. Figure 11 illustrates the IB link failures data over the past year, with raw data attached in Appendix: Supplementary Characterization , Table VIII , indicating that these issues can occur randomly throughout the cluster’s operational period.
|
||||
Figure 11 : Trends of IB Network Failures: Link Flash Cuts
|
||||
## VIII Discussion
|
||||
## VIII-A Discussion on Congestion Control in RDMA Networks
|
||||
Lossless RDMA networks offer several flow-control mechanisms, such as Priority Flow Control (PFC) [ 89 ] for RoCE networks and credit-based flow control [ 90 ] for IB networks. In network routing, static routing algorithms in IB or ECMP (Equal-Cost Multi-Path) [ 91 ] and AR (Adaptive Routing) [ 92 ] effectively handle routing issues. However, congestion can still occur when multiple servers send data to a single receiver, potentially blocking the entire network. To mitigate this, IB NICs use DCQCN (Data Center Quantized Congestion Notification) [ 93 ] as their congestion control algorithm. While Data Processing Units (DPUs), such as the NVIDIA BF series, allow users to customize congestion control algorithms (like HPCC [ 94 ] and TIMELY RTT-based CC [ 95 ] ), they increase the cost and operational complexity of the cluster.
|
||||
In practice, we chose to disable DCQCN to avoid its shortcomings, as it could not find parameters that simultaneously support HFReduce traffic and 3FS storage traffic in our Computation-Storage Integrated Network. Instead, we employed the network tuning methods mentioned in Section VI-A , ensuring our network operates without congestion control and remains congestion-free.
|
||||
## VIII-B Discussion about NVLink Technology Choices
|
||||
Initially, we did not use NVLink to avoid extra costs and maintain stability, as HFReduce was sufficient for training requirements at that time. However, as the demand for LLMs increased, we added NVLink specifically for LLM training purposes. The decision to install NVLink should be based on actual needs due to its potential drawbacks.
|
||||
## VIII-C Maintaince Cost Overview
|
||||
## VIII-C 1 Construction Cost
|
||||
Relative hardware costs are provided in Table II and III . Software costs, contributed by several dozen in-house developers, are just a fraction of the cost for thousands of GPU servers.
|
||||
## VIII-C 2 Power Consumption
|
||||
The average power consumption comparision during ResNet training is provided in Table II . Including the overhead from IB switches and other nodes, the total energy consumption of the Fire-Flyer 2 AI-HPC does not exceed 4 MW, approximately just over 3 MW.
|
||||
## VIII-C 3 Operation Cost
|
||||
Operating costs can be estimated by considering power consumption and rack rental costs. By multiplying this figure by the number of nodes and the PUE (Power Usage Effectiveness), the total operating costs can be calculated.
|
||||
## VIII-D Stability Compared with Other Architectures
|
||||
A recent paper [ 96 ] reportsthat NVLink-related failures account for approximately 52.42% (54 out of 103) of total failures, with raw data indicating 54 NVLink Errors, 21 CUDA Errors, 16 Node Failures, 12 ECC Errors, and 12 Network Errors. In comparison, our NVLink-related issues, primarily Xid-74 Errors, as mentioned in Section VII-C1 , account for about 42.57% of GPU failures.
|
||||
## IX Future Work
|
||||
## Future Arch and Integration with New GPU Models
|
||||
Our next-generation PCIe architecture is designed for MoE (Mixture of Experts) LLM training, where all-to-all performance is crucial. Therefore, the next-gen nodes feature a 1:1 GPU to NIC ratio, comparable to DGX-H100/B100 systems, as illustrated in Figure 12 .
|
||||
We are considering implementing a multi-plane network to reduce costs while maintaining performance. Additionally, we are exploring the use of RoCE switches instead of IB switches, which can significantly lower network expenses. With a 128-port 400 Gbps RoCE switch, a 4-Plane Two-Layer Fat-Trees network can support up to 32,768 GPUs.
|
||||
Figure 12 : Next Generation PCIe Node Atchitecture with Multi-Plane Fat-Trees Network
|
||||
## X Conclusions
|
||||
In this paper, we have shared our experiences and insights from deploying and maintaining the Fire-Flyer 2 AI-HPC, which is equipped with 10,000 PCIe A100 GPUs.Our approach to PCIe architecture and storage-computation integrated network design has resulted in significant cost savings, effectively halving construction costs and demonstrating substantial cost-effectiveness.
|
||||
In terms of software co-design, we introduced HFReduce and HaiScale to overcome hardware limitations, ensuring scalable performance of the PCIe architecture. Our in-house developed 3FS distributed file system,in conjunction with network co-design, facilitates traffic isolation for both 3FS and HFReduce allreduce traffic, effectively preventing congestion. The comprehensive software stack within
|
||||
|
||||
=== BODY (tail) ===
|
||||
the HAI Platform addresses a variety of system faults, from network congestion to hardware failures, thereby ensuring high stability and robustness.
|
||||
Together, these software and hardware innovations enable our PCIe A100 architecture to achieve 80% the performance of NVIDIA’s DGX-A100, while consuming less than 60% of its power. The practical knowledge we have accrued may prove valuable for both industrial and academic sectors. We hope that our work will serve as a reference for others aiming to build their own cost-effective and efficient AI-HPC clusters.
|
||||
## Acknowledgment
|
||||
We extend our heartfelt gratitude to all our colleagues at DeepSeek-AI and High-Flyer Quant for their invaluable contributions to the Fire-Flyer 2 AI-HPC project. Throughout the four years of design, construction, and operation, their collaborative efforts have been crucial in overcoming numerous challenges. Consequently, Fire-Flyer 2 AI-HPC now supports the training tasks of the DeepSeek-AI series large language models [ 97 , 98 , 99 , 72 , 100 , 101 ] . We extend a special thanks to the HFAiLab System Team and the Operations Team, whose core contributions have been essential to the success of this endeavor.
|
||||
@@ -0,0 +1,349 @@
|
||||
TITLE: Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
|
||||
ARXIV: 2408.15664
|
||||
AUTHORS: Lean Wang , Huazuo Gao , Chenggang Zhao , Xu Sun , Damai Dai
|
||||
DATE: [Submitted on 28 Aug 2024]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Auxiliary-Loss-Free Load Balancing
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Background
|
||||
## 2.1 Mixture-of-Experts
|
||||
## 2.2 Auxiliary Loss for Load Balance
|
||||
## Auxiliary Loss
|
||||
## The Dilemma Between Load Balance and Model Performance
|
||||
## 3 Auxiliary-Loss-Free Load Balancing Strategy
|
||||
## Comparison with Other Load Balancing Methods.
|
||||
## 4 Experiments
|
||||
## 4.1 Experimental Setups
|
||||
## Model Architecture.
|
||||
## Training Settings
|
||||
## Baseline.
|
||||
## Metrics.
|
||||
## 4.2 Main Results
|
||||
## 4.3 Empirical Studies on Bias Update Algorithm
|
||||
## Update rate.
|
||||
## Update rule.
|
||||
## Multiplicative bias.
|
||||
## 5 Discussion
|
||||
## 5.1 Loss-Free Balancing Is Compatible with Expert Parallelism
|
||||
## 5.2 Load Balancing and Future Token Leakage
|
||||
## 6 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Background
|
||||
2.1 Mixture-of-Experts
|
||||
2.2 Auxiliary Loss for Load Balance
|
||||
Auxiliary Loss
|
||||
The Dilemma Between Load Balance and Model Performance
|
||||
3 Auxiliary-Loss-Free Load Balancing Strategy
|
||||
Comparison with Other Load Balancing Methods.
|
||||
4 Experiments
|
||||
4.1 Experimental Setups
|
||||
Model Architecture.
|
||||
Training Settings
|
||||
Baseline.
|
||||
Metrics.
|
||||
4.2 Main Results
|
||||
4.3 Empirical Studies on Bias Update Algorithm
|
||||
Update rate.
|
||||
Update rule.
|
||||
Multiplicative bias.
|
||||
5 Discussion
|
||||
5.1 Loss-Free Balancing Is Compatible with Expert Parallelism
|
||||
5.2 Load Balancing and Future Token Leakage
|
||||
6 Conclusion
|
||||
References
|
||||
A Model Architecture
|
||||
B Training Settings
|
||||
C Experiments with Softmax Gate
|
||||
C.1 Comparison of Sigmoid Gate Baseline and Softmax Gate Baseline
|
||||
C.2 Loss-Free Load Balancing with Softmax Gate
|
||||
D Future Token Leakage in Expert Choice
|
||||
D.1 Proof for Theoretical Leakage Amount
|
||||
D.2 Experimental Evidence
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2408.15664v1 [cs.LG] 28 Aug 2024
|
||||
## Auxiliary-Loss-Free Load Balancing
|
||||
Strategy for Mixture-of-Experts
|
||||
Lean Wang 1 1 footnotemark:
|
||||
1
|
||||
Affiliation: DeepSeek-AI
|
||||
Affiliation: State Key Laboratory of Multimedia Information Processing,School of Computer Science, Peking University
|
||||
Email: lean@pku.edu.cn
|
||||
Huazuo Gao
|
||||
Affiliation: DeepSeek-AI
|
||||
Email: xusun@pku.edu.cn
|
||||
Chenggang Zhao
|
||||
Affiliation: DeepSeek-AI
|
||||
Email: damai.dai@deepseek.com
|
||||
Xu Sun 2 2 footnotemark:
|
||||
2
|
||||
Affiliation: State Key Laboratory of Multimedia Information Processing,School of Computer Science, Peking University
|
||||
Damai Dai 2 2 footnotemark:
|
||||
2
|
||||
Affiliation: DeepSeek-AI
|
||||
## Abstract
|
||||
For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead.
|
||||
Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance.
|
||||
In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing , featured by an auxiliary-loss-free load balancing strategy.
|
||||
To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert.
|
||||
By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load.
|
||||
In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training.
|
||||
We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens.
|
||||
Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.
|
||||
## 1 Introduction
|
||||
Mixture-of-Experts (MoE) architectures have emerged as a promising solution for managing computational costs when scaling up parameters in large language models (LLMs).
|
||||
Recent applications of MoE in Transformer-based models ( Vaswani et al., 2017 ) have led to successful attempts at scaling language models to substantial sizes ( Shao et al., 2024 ; DeepSeek-AI et al., 2024 ; Dai et al., 2024 ; Fedus et al., 2021 ; Lepikhin et al., 2020 ) , resulting in remarkable performance improvements.
|
||||
However, training MoE models always face the circumstance of load imbalance, which may result in routing collapse ( Shazeer et al., 2017 ) or increased computational overhead ( Fedus et al., 2021 ; Lepikhin et al., 2020 ; Shazeer et al., 2017 ) .
|
||||
In order to avoid imbalanced routing, existing methods ( Fedus et al., 2021 ; Lepikhin et al., 2020 ) commonly use an auxiliary loss to encourage balanced expert load.
|
||||
Although the auxiliary loss can alleviate load imbalance during training, it also introduces undesired gradients that conflict with the language modeling objective.
|
||||
These interference gradients will impair the model performance, so existing MoE methods always need to consider the trade-off between load balance and model performance.
|
||||
In this paper, we propose Loss-Free Balancing , an auxiliary-loss-free load balancing strategy, aiming at maintaining control over expert load balance while not introducing interference gradients.
|
||||
Loss-Free Balancing features an iterative process of token routing and bias updating.
|
||||
As illustrated in Figure 1 , before the top-K routing decision of MoE, Loss-Free Balancing will first apply expert-wise biases to the original routing scores to produce biased gating scores, which determine the actual routing targets of each token during training.
|
||||
These expert-wise biases will keep updating according to the expert load observed on recent training tokens, where the biases of heavy-load experts will be depressed and those of lite-load experts will be elevated.
|
||||
Through this dynamic updating strategy, Loss-Free Balancing ensures that the biased gating scores can consistently lead to balanced routing results.
|
||||
Compared with the auxiliary-loss-controlled load balancing strategies, Loss-Free Balancing does not introduce undesired gradients that disrupt the primary language modeling objective, so its training process is more noise-free and friendly.
|
||||
In order to validate the performance of Loss-Free Balancing, we train MoE language models with 1B parameters on 100B tokens and 3B parameters on 200B tokens from scratch.
|
||||
Experimental results demonstrate that Loss-Free Balancing produces MoE models with better validation loss than traditional auxiliary-loss-controlled models.
|
||||
Meanwhile, keeping the performance advantage, Loss-Free Balancing also achieves a significantly better load balance at the global and batch levels, and is naturally compatible with expert parallelism, which is usually employed for training extremely large MoE models.
|
||||
Figure 1: Loss-Free Balancing selects experts according to a “biased gating score” in each training step and updates this expert-wise bias after each training step.
|
||||
## 2 Background
|
||||
## 2.1 Mixture-of-Experts
|
||||
Current dominant MoE architectures ( Lepikhin et al., 2020 ; Fedus et al., 2021 ; Dai et al., 2024 ) replace the MLP layers in standard transformers with MoE layers. In an MoE layer, Top-K routing is employed to select the experts for each token. Let 𝐮 t \mathbf{u}_{t} denote the input of the t t -th token to an N N -expert MoE layer, the output 𝐡 t \mathbf{h}_{t} is computed as follows:
|
||||
𝐡 t = 𝐮 t + ∑ i = 1 N g i , t FFN i ( 𝐮 t ) , \displaystyle\mathbf{h}_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N}g_{i,t}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}\right),
|
||||
(1)
|
||||
g i , t = { s i , t , s i , t ∈ Topk ( { s j , t ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise, \displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}\in\operatorname{Topk}\left(\left\{s_{j,t}\mid 1\leq j\leq N\right\},K\right),\\
|
||||
0,&\text{ otherwise, }\end{cases}
|
||||
s i , t = G ( 𝐮 t 𝐞 i T ) , \displaystyle s_{i,t}=G\left(\mathbf{u}_{t}{}^{T}\mathbf{e}_{i}\right),
|
||||
where G G is a nonlinear gating function and 𝐞 i \mathbf{e}_{i} is the centroid of the i i -th expert.
|
||||
## 2.2 Auxiliary Loss for Load Balance
|
||||
## Auxiliary Loss
|
||||
Uncontrolled routing strategies are likely to encounter load imbalance, which has two notable drawbacks.
|
||||
Firstly, there is a risk of routing collapse ( Shazeer et al., 2017 ) , where the model consistently selects only a few experts, hindering sufficient training of the other experts.
|
||||
Secondly, when experts are distributed across multiple devices, load imbalance can exacerbate computation bottlenecks.
|
||||
To address these issues, an auxiliary loss ( Fedus et al., 2021 ; Lepikhin et al., 2020 ) is commonly employed to control load balance.
|
||||
For a sequence of length T T , the auxiliary loss is defined as:
|
||||
ℒ Balance \displaystyle\mathcal{L}_{\text{Balance }}
|
||||
= α ∑ i = 1 N f i P i , \displaystyle=\alpha\sum_{i=1}^{N}f_{i}P_{i},
|
||||
(2)
|
||||
f i \displaystyle f_{i}
|
||||
= N K T ∑ t = 1 T 𝟙 ( Token t selects Expert i ) , \displaystyle=\frac{N}{KT}\sum_{t=1}^{T}\mathbbm{1}(\text{ Token }t\text{ selects Expert }i),
|
||||
P i \displaystyle P_{i}
|
||||
= 1 T ∑ t = 1 T s i , t , \displaystyle=\frac{1}{T}\sum_{t=1}^{T}s_{i,t},
|
||||
where N N is the total number of experts, K K is the number of experts selected for each token, s i , t s_{i,t} is the routing score of Expert i i for Token t t , f i f_{i} represents the fraction of tokens routed to Expert i i , P i P_{i} denotes the average gating scores of Expert i i , and α \alpha is a hyper-parameter controlling the strength of the auxiliary loss.
|
||||
Figure 2: The dilemma between load balance and model performance for auxiliary-loss-controlled training. A small auxiliary loss coefficient α \alpha leads to poor load balance, while a large α \alpha impairs the model performance. In contrast, our Loss-Free Balancing method breaks this dilemma.
|
||||
## The Dilemma Between Load Balance and Model Performance
|
||||
The auxiliary loss mentioned above can encourage load balance, but it also interferes with language modeling training as an additional regularization term. The absence of an auxiliary loss or a small auxiliary loss coefficient α \alpha can lead to poor balance, while a large α \alpha can impair training, resulting in suboptimal performance.
|
||||
To illustrate this dilemma, we present the relationship between load balance and model performance in Figure 2 .
|
||||
We vary α \alpha among 1e-2, 1e-3, 1e-4, and 0, and present the corresponding MaxVio global \text{MaxVio}_{\text{global}} , which measures the degree of load balance and its computation details are described in § 4.1 .
|
||||
As shown in the figure, a small α \alpha causes routing collapse, affecting the model efficiency and potentially leading to some experts being insufficiently learned or exploited; while a large α \alpha keeps load balance under control but notably degrades the model performance.
|
||||
In order to break this dilemma, we propose Loss-Free Balancing as a solution, which directly controls the expert load balance, but does not introduce unexpected gradients other than the gradients from the language modeling loss.
|
||||
## 3 Auxiliary-Loss-Free Load Balancing Strategy
|
||||
For a better load-balancing alternative that does not directly interfere with the main gradients from the training objective, we propose Loss-Free Balancing , which directly adjusts the gating scores of each expert according to their balance condition.
|
||||
As illustrated in Figure 1 , we add an expert-wise bias term { 𝒃 𝒊 } 𝒊 = 𝟏 𝑵 \bm{\{b_{i}\}_{i=1}^{N}} to the gating scores s i , t s_{i,t} of each expert, and use the biased scores to determine the top-K selection:
|
||||
g i , t = { s i , t , s i , t + b i ∈ Topk ( { s j , t + b j ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise. \displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}+b_{i}\in\operatorname{Topk}\left(\left\{s_{j,t}+b_{j}\mid 1\leq j\leq N\right\},K\right),\\
|
||||
0,&\text{ otherwise. }\end{cases}
|
||||
(3)
|
||||
Note that the expert bias term b i b_{i} is only used to adjust the routing strategy by influencing the top-K selection.
|
||||
It is not added to the g i , t g_{i,t} that weights the output of the selected experts when computing the final output of the MoE layer.
|
||||
In order to derive proper biases, we adjust each bias b i b_{i} iteratively according to the following principle: decreasing it when the corresponding expert has a relatively heavy load, and vice versa.
|
||||
To be specific, for each b i {b}_{i} , we keep monitoring its corresponding expert load on the previous batch.
|
||||
If an expert has a heavy load on the previous batch, we will reduce its bias.
|
||||
Otherwise, we will increase it.
|
||||
Algorithm 1 describes the details of our update algorithm for the expert-wise biases.
|
||||
It is worth noting that we update the biases based on the historical balance condition, since utilizing the load information of the current sequence will break the causal constraint of language modeling, leading to leakage of the information of future tokens.
|
||||
Through the dynamic adjustment for the biases, we can achieve good expert load balance, but not directly introduce noisy gradients into the model like the auxiliary-loss-controlled method does.
|
||||
Algorithm 1 Adjusting the per-expert bias b i {b}_{i} during training
|
||||
Input: MoE model θ \theta , training batch iterator B B , bias update rate u u .
|
||||
1. Initialize b i = 0 b_{i}=0 for each expert;
|
||||
for a batch { ( 𝐱 k , 𝐲 k ) } k \{(\mathbf{x}_{k},\mathbf{y}_{k})\}_{k} in B B do
|
||||
2. Train MoE model θ \theta on the batch data { ( 𝐱 k , 𝐲 k ) } k \{(\mathbf{x}_{k},\mathbf{y}_{k})\}_{k} , with gating scores calculated according to Eq. ( 3 );
|
||||
3. Count the number of assigned tokens c i c_{i} for each expert, and the average number c i ¯ \overline{c_{i}} ;
|
||||
4. Calculate the load violation error e i = c i ¯ − c i e_{i}=\overline{c_{i}}-c_{i} ;
|
||||
4. Update 𝐛 i \mathbf{b}_{i} by b i = b i + u ∗ sign ( e i ) b_{i}=b_{i}+u*\mathrm{sign}(e_{i}) ;
|
||||
end for
|
||||
Output: trained model θ \theta , corresponding bias 𝐛 i \mathbf{b}_{i}
|
||||
## Comparison with Other Load Balancing Methods.
|
||||
In order to show the theoretical advantages of Loss-Free Balancing, we compare it with other two mainstream load balancing methods, i.e., the auxiliary-loss-controlled method ( Lepikhin et al., 2020 ; Fedus et al., 2021 ) and the Expert Choice (EC) ( Zhou et al., 2022 ) method.
|
||||
As described in § 2.2 , the auxiliary-loss-controlled method faces the dilemma between load balance and model performance, and a perfect trade-off may not exist.
|
||||
As for the EC method, it will break the causal constraint of language modeling, since the target experts of each token are conditioned on the future tokens in the same sequence or batch.
|
||||
This will result in the leakage of information about future tokens, thus destroying the generalization of the model.
|
||||
Table 1 summarizes the properties of different load balancing methods.
|
||||
Table 1: Comparison among different load balancing methods. The good property is displayed in green and the bad property in red .
|
||||
Load Balancing Methods
|
||||
Balanced
|
||||
Expert Load
|
||||
Interference
|
||||
Gradients
|
||||
Future Token
|
||||
Leakage
|
||||
Loss-Controlled (strong auxiliary loss)
|
||||
balanced
|
||||
strong
|
||||
no leakage
|
||||
Loss-Controlled (weak auxiliary loss)
|
||||
imbalanced
|
||||
weak
|
||||
no leakage
|
||||
Expert Choice
|
||||
balanced
|
||||
none
|
||||
with leakage
|
||||
Loss-Free (Ours)
|
||||
balanced
|
||||
none
|
||||
no leakage
|
||||
## 4 Experiments
|
||||
## 4.1 Experimental Setups
|
||||
## Model Architecture.
|
||||
We employ the DeepSeekMoE ( Dai et al., 2024 ) architecture as the backbone since it outperforms conventional MoE architectures like GShard ( Lepikhin et al., 2020 ) significantly.
|
||||
Compared with GShard ( Lepikhin et al., 2020 ) , it segments experts into finer granularity and isolates some experts as shared ones.
|
||||
Slightly different from DeepSeekMoE, in our main experiments, we choose sigmoid instead of softmax as the gating function G G , since we find that the sigmoid baseline performs better than the softmax baseline.
|
||||
Even so, we still provide the experimental results and discussion for the softmax gate in Appendix C .
|
||||
Our experiments are based on two model sizes of 1B and 3B total parameters, and we tune the bias update rate under only the 1B scale.
|
||||
Experiments under the 3B scale directly inherit the best configuration for the 1B scale.
|
||||
Due to the page limit, we present more details about our architecture in Appendix A .
|
||||
## Training Settings
|
||||
We use a multilingual training corpus created by DeepSeek-AI, sourced from a diverse range of textual materials including web text, mathematical material, coding scripts, and published
|
||||
literature.
|
||||
We employ the HuggingFace Tokenizer 2 2
|
||||
2
|
||||
https://github.com/huggingface/tokenizers to train a byte pair encoding (BPE) ( Sennrich et al., 2015 ) tokenizer with a vocabulary size of 32K.
|
||||
In order to draw solid conclusions, we train the 1B model on 100B tokens and the 3B model on 200B tokens to ensure sufficient training.
|
||||
We apply the cosine learning rate scheduler ( Loshchilov & Hutter, 2016 ) and multi-step learning rate scheduler ( Dai et al., 2024 ) for the 1B and 3B models, respectively.
|
||||
Due to the page limit, we list more details about our training settings and hyper-parameters in Appendix B ).
|
||||
## Baseline.
|
||||
We compare our Loss-Free Balancing method with the conventional auxiliary-loss-controlled method.
|
||||
For the baseline, we set the auxiliary loss coefficient α \alpha to 0.001 to achieve a reasonable trade-off between model performance and load balance (see Figure 2 ).
|
||||
We do not take the EC method into comparison due to its issue of future token leakage, which we will discuss in depth in § 5.2 .
|
||||
## Metrics.
|
||||
We reserve a validation set from the training corpus to evaluate model performance and load balance.
|
||||
For model performance, we take perplexity as the metric.
|
||||
For load balance, we introduce a metric called maximal violation ( MaxVio ) to quantify the degree of load balance of an MoE layer:
|
||||
MaxVio = max i Load i − Load i ¯ Load i ¯ , \text{MaxVio}=\frac{\max_{i}\text{Load}_{i}-\overline{\text{Load}_{i}}}{\overline{\text{Load}_{i}}},
|
||||
(4)
|
||||
where Load i \text{Load}_{i} represents the number of tokens assigned to the i i -th expert, and Load i ¯ \overline{\text{Load}_{i}} denotes the expected expert load under perfect load balance.
|
||||
MaxVio has two variants: MaxVio global \textbf{MaxVio}_{\textbf{global}} and MaxVio batch \textbf{MaxVio}_{\textbf{batch}} .
|
||||
For MaxVio global \textbf{MaxVio}_{\textbf{global}} , we count Load i \text{Load}_{i} on the whole validation set, so it reflects the degree of balanced expert utilization and efficiency upper bound when the batch size approaches the limitation.
|
||||
For MaxVio batch \textbf{MaxVio}_{\textbf{batch}} , we count Load i \text{Load}_{i} on each training batch, so it is more related to the training efficiency.
|
||||
For simplicity, in the rest of this paper, we report the MaxVio averaged across all layers as a load balance measurement of the whole model.
|
||||
## 4.2 Main Results
|
||||
Table 2 shows the validation perplexity and MaxVio global \text{MaxVio}_{\text{global}} for the 1B and 3B MoE models trained with auxiliary loss or our auxiliary-loss-free load balancing strategy.
|
||||
As shown in the table, compared with the auxiliary-loss-controlled method, our Loss-Free Balancing achieves better perplexity and much better global load balance for both 1B and 3B models.
|
||||
In addition, to present the load balance condition during training, we provide a load balancing curve depicting MaxVio batch \text{MaxVio}_{\text{batch}} over training steps in Figure 3 , which demonstrates the persistent advantage of Loss-Free Balancing on load balance.
|
||||
In summary, our Loss-Free Balancing method avoids interfering gradients during training and effectively controls the load balance, breaking the dilemma between load balance and model performance in MoE training.
|
||||
Table 2: Loss-Free Balancing achieves lower perplexity and better load balance on both 1B and 3B models. A validation set is used to calculate these metrics (see details in Appendix B ).
|
||||
Model Size
|
||||
Load Balancing Methods
|
||||
Validation Perplexity
|
||||
MaxVio global \textbf{MaxVio}_{\textbf{global}}
|
||||
1B
|
||||
Loss-Controlled
|
||||
9.56
|
||||
0.72
|
||||
Loss-Free
|
||||
9.50
|
||||
0.04
|
||||
3B
|
||||
Loss-Controlled
|
||||
7.97
|
||||
0.52
|
||||
Loss-Free
|
||||
7.92
|
||||
0.04
|
||||
Figure 3: Loss-Free Balancing maintains a better load balance throughout most of the training time. Here, MaxVio batch \text{MaxVio}_{\text{batch}} is averaged over 100 neighboring steps for visibility purposes.
|
||||
## 4.3 Empirical Studies on Bias Update Algorithm
|
||||
We conduct empirical studies on the update rate and variants of the bias update algorithm to validate the optimal configuration used in our main experiments.
|
||||
## Update rate.
|
||||
The update rate u u in Algorithm 1 controls the speed at which the expert bias { b i } i = 1 N \{{b}_{i}\}_{i=1}^{N} converges to the “suitable bias”. Figure 4 illustrates that an overly low update rate u = 0.0001 u=0.0001 may lead to slow convergence, while an unnecessarily high update rate u = 0.01 u=0.01 can cause undesirable fluctuations of the expert bias b i {b}_{i} during the later stage of training, deteriorating load balance in this stage. Both situations can impair performance. An appropriate choice is u = 0.001 u=0.001 , which shows good training balance and validation perplexity.
|
||||
Figure 4: The impact of update rate on training load balance. A low update rate shows poor load balance in the early stage of training, while a high update rate deteriorates load balance in the later stage. Validation PPL denotes the validation perplexity.
|
||||
## Update rule.
|
||||
We investigate a different update rule of the expert-wise biases.
|
||||
To be specific, we attempt to change the update rule of b i = b i + u ∗ sign ( e i ) b_{i}=b_{i}+u*\mathrm{sign}(e_{i}) to b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , which encourages the bias of experts with high violation errors to change faster. Although this variant slightly improves load balance, it does not lead to better performance, as shown in Table 3 .
|
||||
Therefore, we maintain the sign \mathrm{sign} version.
|
||||
Table 3: The variant b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} slightly improves load balance but does not show improvement in model performance.
|
||||
Method
|
||||
Perplexity
|
||||
MaxVio global \textbf{MaxVio}_{\textbf{global}}
|
||||
b i = b i + u ∗ sign ( e i ) b_{i}=b_{i}+u*\text{sign}(e_{i}) , u = 0.001 u=0.001
|
||||
9.50
|
||||
0.044
|
||||
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.01 u=0.01
|
||||
9.53
|
||||
0.028
|
||||
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.001 u=0.001
|
||||
9.51
|
||||
0.036
|
||||
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.0001 u=0.0001
|
||||
9.51
|
||||
0.040
|
||||
## Multiplicative bias.
|
||||
In addition to adding the expert-wise biases to the gating scores, using multiplicative biases is also a potential variant:
|
||||
g i , t = { s i , t , s i , t ∗ b i ∈ Topk ( { s j , t ∗ b j ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise, \displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}*{b}_{i}\in\operatorname{Topk}\left(\left\{s_{j,t}*{b}_{j}\mid 1\leq j\leq N\right\},K\right),\\
|
||||
0,&\text{ otherwise, }\end{cases}
|
||||
(5)
|
||||
These { b i } i = 1 N \{{b}_{i}\}_{i=1}^{N} can be updated using a similar procedure to Algorithm 1 , except that they should be initialized as 1 instead of 0.
|
||||
Table 4 shows that using multiplicative biases results in slightly worse model performance compared to using additive biases, without significant improvements in load balance.
|
||||
Based on these findings, we conclude that additive biases are a more suitable choice for our method.
|
||||
Table 4: Multiplicative bias shows similar load balance but slightly worse performance compared to additive bias.
|
||||
Method
|
||||
Perplexity
|
||||
MaxVio global \textbf{MaxVio}_{\textbf{global}}
|
||||
Addative Bias, u = 0.001 u=0.001
|
||||
9.50
|
||||
0.044
|
||||
Multiplicative Bias, u = 0.01 u=0.01
|
||||
9.52
|
||||
0.041
|
||||
Multiplicative Bias, u = 0.001 u=0.001
|
||||
9.52
|
||||
0.036
|
||||
Multiplicative Bias, u = 0.0001 u=0.0001
|
||||
9.54
|
||||
0.048
|
||||
## 5 Discussion
|
||||
## 5.1 Loss-Free Balancing Is Compatible with Expert Parallelism
|
||||
Extremely large-scale MoE models often employ expert parallelism ( Lepikhin et al., 2020 ) for training or inference, which distributes experts across different devices to reduce memory requirements.
|
||||
In such scenarios, load balance on the data in a single computation step is crucial for efficiency.
|
||||
Due to expert parallelism, each computation step involves micro_batch_size * ep_data_parallel_size samples, which we refer to as a computation batch .
|
||||
Here, micro_batch_size denotes the number of samples processed in one gradient accumulation step on a single device.
|
||||
Loss-Free Balancing can achieve nearly optimal global load balance, and the load balance in each computation step will get closer to the global load balance as the computation batch size increases.
|
||||
In Figure 5 , we examine the computation-batch-level load balance with the MaxVio computation-batch \text{MaxVio}_{\text{computation-batch}} metric.
|
||||
The results show that the load balance of our Loss-Free Balancing always keeps improving as the computation batch size increases, but the load balance of the auxiliary-loss-controlled method approximately maintains a constant level when the computation batch is large.
|
||||
Since expert parallelism will significantly increase the computation batch size by ep_data_parallel_size times, Loss-Free Balancing is naturally compatible with large-scale MoE training, and its advantage on the load balance will be further enhanced as the size of expert parallelism increases.
|
||||
Figure 5: Loss-Free Balancing achieves improved balance compared to auxiliary-loss training as the computation-batch size increases, demonstrating its superiority when a moderately sized computation-batch is utilized.
|
||||
## 5.2 Load Balancing and Future Token Leakage
|
||||
For casual language models, load balancing methods must adhere to the causal constraint of language modeling to avoid future token leakage. While conventional auxiliary-controlled balancing and our Loss-Free Balancing obey this constraint, Expert Choice (EC) ( Zhou et al., 2022 ) violates it. EC ensures perfect load balance by assigning exactly the same number of tokens to each expert. However, this approach inherently leads to a severe issue of future token leakage.
|
||||
In EC, future tokens can influence the expert assignment of previous tokens. Figure 6 illustrates how information can be easily transmitted within a sequence via such influence. Theoretically, the token assignment of an MoE layer with sparse ratio R R (average activated experts per token K K divided by total expert number N N ) can leak more than K log 2 1 − R R K\log_{2}\frac{1-R}{R} bits per token (proof in Appendix D.1 ). For a 9-layer MoE model with 16 experts and an average of 2 experts per token, this amounts to 50 bits, sufficient for each token to determine its successor’s identity.
|
||||
Figure 6:
|
||||
An example of future token leakage in EC.
|
||||
Future tokens can influence the expert assignment of previous tokens.
|
||||
Such an assignment can help previous tokens to infer the identity of their successors.
|
||||
We designed experiments to demonstrate the existence of future token leakage in realistic model training. (1) We reduced the chunk size, within which top-K selection is performed, from 8192 tokens (4 sentences) to 512 (1/4 sentence), with the expectation of exposing such leakage. We observed an abnormal loss drop (about 10%), confirming the presence of leakage. (2) We made leakage more difficult by shuffling tokens across chunks in the top-K selection step, and observed that the abnormal loss drop was mitigated. Detailed experimental results on EC’s information leakage are provided in Appendix D.2 .
|
||||
Future token leakage is fatal since it destroys the generalization of a model and prevents reliable evaluation of the model performance.
|
||||
Therefore, compared with EC, scaling up an MoE model with our Loss-Free Balancing is safer.
|
||||
## 6 Conclusion
|
||||
In this work, we introduced Loss-Free Balancing , a novel MoE load balance control method without introducing auxiliary-loss gradients. Loss-Free Balancing addresses the issue of traditional auxiliary-loss load balance control, which introduces additional gradients during training and potentially impairs model performance when enforcing load balance. Experiments conducted on 1B and 3B MoE models, trained on 100B and 300B tokens respectively, demonstrate that Loss-Free Balancing achieves better model performance and load balance compared to the traditional auxiliary-loss training.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,762 @@
|
||||
TITLE: Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation
|
||||
ARXIV: 2410.13848
|
||||
AUTHORS: Chengyue Wu , Xiaokang Chen , Zhiyu Wu , Yiyang Ma , Xingchao Liu , Zizheng Pan , Wen Liu , Zhenda Xie , Xingkai Yu , Chong Ruan , Ping Luo
|
||||
DATE: [Submitted on 17 Oct 2024]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Related Work
|
||||
## 2.1 Visual Generation
|
||||
## 2.2 Multimodal Understanding
|
||||
## 2.3 Unified Multimodal Understanding and Generation
|
||||
## 3 Janus: A Simple, Unified and Flexible Multimodal Framework
|
||||
## 3.1 Architecture
|
||||
## 3.2 Training Procedure
|
||||
## 3.3 Training Objective
|
||||
## 3.4 Inference
|
||||
## 3.5 Possible Extensions
|
||||
## 4 Experiments
|
||||
## 4.1 Implementation Details
|
||||
## 4.2 Data Setup
|
||||
## 4.3 Evaluation Setup
|
||||
## 4.4 Comparison with State-of-the-arts
|
||||
## 4.5 Ablation Studies
|
||||
## 4.6 Qualitative Results
|
||||
## 5 Conclusion
|
||||
## Appendix
|
||||
## Appendix A Details of Semantic Tokenizer Mentioned in Ablation Study
|
||||
## A.1 Architecture of Semantic Tokenizer
|
||||
## A.2 Training
|
||||
## A.3 Integrating with LLM
|
||||
## Appendix B Additional Qualitative Results
|
||||
|
||||
=== BODY (head) ===
|
||||
Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Work
|
||||
2.1 Visual Generation
|
||||
2.2 Multimodal Understanding
|
||||
2.3 Unified Multimodal Understanding and Generation
|
||||
3 Janus: A Simple, Unified and Flexible Multimodal Framework
|
||||
3.1 Architecture
|
||||
3.2 Training Procedure
|
||||
3.3 Training Objective
|
||||
3.4 Inference
|
||||
3.5 Possible Extensions
|
||||
4 Experiments
|
||||
4.1 Implementation Details
|
||||
4.2 Data Setup
|
||||
4.3 Evaluation Setup
|
||||
4.4 Comparison with State-of-the-arts
|
||||
4.5 Ablation Studies
|
||||
4.6 Qualitative Results
|
||||
5 Conclusion
|
||||
A Details of Semantic Tokenizer Mentioned in Ablation Study
|
||||
A.1 Architecture of Semantic Tokenizer
|
||||
A.2 Training
|
||||
A.3 Integrating with LLM
|
||||
B Additional Qualitative Results
|
||||
References
|
||||
License: CC BY 4.0
|
||||
arXiv:2410.13848v1 [cs.CV] 17 Oct 2024
|
||||
\reportnumber
|
||||
001
|
||||
## Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation
|
||||
Chengyue Wu 1,2 Xiaokang Chen 1,∗,† Zhiyu Wu 1,3 Yiyang Ma 1,3 Xingchao Liu 1 Zizheng Pan 1 Wen Liu 1 Zhenda Xie 1 Xingkai Yu 1 Chong Ruan 1 Ping Luo 2,∗
|
||||
1 DeepSeek-AI 2 The University of Hong Kong 3 Peking University
|
||||
† : Project lead ∗ : Corresponding authors
|
||||
Project Page: https://github.com/deepseek-ai/Janus
|
||||
## Abstract
|
||||
In this paper, we introduce Janus , an autoregressive framework that unifies multimodal understanding and generation.
|
||||
Prior research often relies on a single visual encoder for both tasks, such as Chameleon. However, due to the differing levels of information granularity required by multimodal understanding and generation, this approach can lead to suboptimal performance, particularly in multimodal understanding. To address this issue, we decouple visual encoding into separate pathways, while still leveraging a single, unified transformer architecture for processing.
|
||||
The decoupling not only alleviates the conflict between the visual encoder’s roles in understanding and generation, but also enhances the framework’s flexibility. For instance, both the multimodal understanding and generation components can independently select their most suitable encoding methods.
|
||||
Experiments show that Janus surpasses previous unified model and matches or exceeds the performance of task-specific models.
|
||||
The simplicity, high flexibility, and effectiveness of Janus make it a strong candidate for next-generation unified multimodal models.
|
||||
## 1 Introduction
|
||||
(a) Benchmark Performance.
|
||||
(b) Visual Generation Results.
|
||||
Figure 1: Multimodal understanding and vision generation results from our Janus . Janus outperforms the previous state-of-the-art unified multimodal models as well as some task-specific multimodal understanding models, while also demonstrating strong visual generation capabilities. The image resolution is 384 × 384 384\times 384 . Best viewed on screen.
|
||||
In recent years, multimodal large models have made significant advancements in both understanding and generation domains [ 51 , 20 ] . In the field of multimodal understanding, researchers follow the design of LLaVA [ 51 ] by using a vision encoder as a bridge to enable large language models (LLMs) to understand images. In the field of visual generation, diffusion-based approaches [ 20 , 9 , 67 , 20 ] have seen notable success. More recently, some works have explored autoregressive methods for vision generation [ 73 , 79 ] , achieving performance comparable to diffusion models. To build more powerful and generalist multimodal models, researchers have sought to combine multimodal understanding and generation tasks [ 75 , 77 , 94 ] . For instance, some studies have attempted to connect multimodal understanding models with pretrained diffusion models [ 75 , 28 , 27 ] .
|
||||
For example, Emu [ 75 ]
|
||||
uses the output of the LLM as a condition for a pretrained diffusion model, and then relies on the diffusion model to generate images.
|
||||
However, strictly speaking, this approach cannot be considered a truly unified model, because the visual generation functionality is handled by the external diffusion model, while the multimodal LLM itself lacks the capability to directly generate images.
|
||||
Other approaches [ 77 , 85 , 86 , 94 ] employ a single transformer to unify both multimodal understanding and generation tasks, which improves instruction-following for visual generation, unlocks potential emergent abilities, and reduces model redundancy.
|
||||
Such methods typically use a single vision encoder to process inputs for both two tasks. However, the representations required by multimodal understanding and generation tasks differ significantly.
|
||||
In multimodal understanding tasks, the purpose of the vision encoder is to extract high-level semantic information (e.g., object categories or visual attributes within an image). The output of understanding task not only involves extracting information from images but also involves complex semantic reasoning. Therefore, the granularity of the vision encoder’s representation tends to mainly focus on high-dimensional semantic representation. By contrast, in visual generation tasks, the main focus is on generating local details and maintaining global consistency in the image. The representation in this context necessitates a low-dimensional encoding that is capable of fine-grained spatial structure and textural detail expression. Unifying the representations of these two tasks within the same space will lead to conflicts and trade-offs. Consequently, existing unified models for multimodal understanding and generation often compromise on multimodal understanding performance, falling markedly short of the state-of-the-arts multimodal understanding models. We explore this issue further in the ablation study.
|
||||
To solve this problem, we propose Janus 1 1
|
||||
1
|
||||
In Roman mythology, Janus is the god of duality and transitions, symbolizing the coexistence of contradictory forces by having two faces, each looking in opposite directions. Similarly, our model captures the inherent tension between vision tasks: understanding demands abstract, high-level semantic representations, while generation requires concrete, detailed information. By decoupling these processes into specialized encoders, our system mirrors Janus’s dual nature, resolving this tension within a unified architecture. , a unified multimodal framework that decouples visual encoding for multimodal understanding and generation. Specifically, we introduce two independent visual encoding pathways: one for multimodal understanding and one for multimodal generation, unified by the same transformer architecture. The proposed method offers two main benefits: (1) Janus alleviates the conflict stemming from the different granular needs of multimodal understanding and generation and eliminates the need to make trade-offs between two tasks when selecting visual encoders. (2) Janus is flexible and extensible. After decoupling, both the understanding and generation tasks can adopt state-of-the-art encoding techniques specific to their domain. Moreover, it is possible for Janus to accommodate additional input types in the future, such as point clouds, EEG signals, or audio data, where independent encoders can extract features and then use a unified transformer to process them.
|
||||
To the best of our knowledge, we are the first to highlight the importance of decoupling visual encoding within the unified multimodal understanding and generation framework.
|
||||
Our experimental results show that Janus surpasses existing unified models with comparable parameter sizes on both multimodal understanding and generation benchmarks, achieving state-of-the-art results. Notably, Janus even outperforms some task-specific models which have significantly more parameters (Figure 1 ). Specifically, on multimodal understanding benchmarks MMBench [ 54 ] , SEED-Bench [ 42 ] , and POPE [ 48 ] , Janus ( 1.3 1.3 B) achieved scores of 69.4 69.4 , 63.7 63.7 , and 87.0 87.0 , respectively, outperforming LLaVA-v 1.5 1.5 ( 7 7 B) [ 50 ] and Qwen-VL-Chat ( 7 7 B) [ 3 ] . On visual generation benchmarks MSCOCO- 30 30 K [ 11 ] and GenEval [ 30 ] , Janus achieved an FID score of 8.53 8.53 and an accuracy of 61 61 %, surpassing text-to-image generative models such as DALL-E 2 2 [ 66 ] and SDXL [ 62 ] . We believe that the strong performance, coupled with the high flexibility and extensibility of Janus, presents it as a strong candidate for next-generation unified multimodal models.
|
||||
## 2 Related Work
|
||||
## 2.1 Visual Generation
|
||||
Visual generation is a rapidly evolving field that combines concepts from natural language processing with advancements in transformer architectures. Autoregressive models, influenced by the success in language processing, leverage transformers to predict sequences of discrete visual tokens (codebook IDs) [ 24 , 65 , 75 ] . These models tokenize visual data and employ a prediction approach similar to GPT-style [ 64 ] techniques. Additionally, masked prediction models [ 7 , 8 ] draw upon BERT-style [ 19 ] masking methods, predicting masked sections of visual inputs to improve synthesis efficiency, and have been adapted for video generation [ 89 ] . Concurrently, continuous diffusion models have showcased impressive capabilities in visual generation [ 33 , 71 , 67 ] , complementing discrete methods by approaching generation through a probabilistic lens.
|
||||
## 2.2 Multimodal Understanding
|
||||
Multimodal large language models (MLLMs) integrate both text and images [ 6 , 80 , 81 ] . By leveraging pretrained LLMs, MLLMs [ 55 , 51 , 95 , 82 , 12 , 2 , 1 ] demonstrate a robust ability to understand and process multimodal information. Recent advancements have explored extending MLLMs with pretrained diffusion models to facilitate image generation [ 27 , 36 , 75 , 76 , 29 ] . These methods fall under the category of tool utilization, where diffusion models are used to generate images based on the conditions output by the MLLM, while the MLLM itself does not have the ability to directly perform visual generation. Moreover, the generative ability of the entire system is often constrained by the external diffusion model, making its performance inferior to directly using the diffusion model on its own [ 27 , 75 ] .
|
||||
## 2.3 Unified Multimodal Understanding and Generation
|
||||
Unified multimodal understanding and generation models are considered powerful for facilitating seamless reasoning and generation across different modalities [ 77 , 94 ] . Traditional approaches in these models typically use a single visual representation for both understanding and generation tasks, regardless of whether they are based on autoregressive (AR) models [ 77 , 85 ] or diffusion models [ 86 , 94 ] .
|
||||
For example, Chameleon [ 77 ] adopts a VQ Tokenizer to encode images for both multimodal understanding and generation.
|
||||
However, this practice may lead to suboptimal outcomes, as the vision encoder might face a trade-off between the demands of understanding and generation. In contrast, our Janus can explicitly decouple the visual representations for understanding and generation, recognizing that different tasks may require varying levels of information.
|
||||
## 3 Janus: A Simple, Unified and Flexible Multimodal Framework
|
||||
Figure 2: Architecture of our Janus. Different from previous approaches [ 77 , 85 ] that typically assume visual understanding and generation require the same visual encoder, our Janus decouples visual encoding for visual understanding and visual generation. “Und. Encoder” and “Gen. Encoder” are abbreviations for “Understanding Encoder” and “Generation Encoder”, respectively. Best viewed in color.
|
||||
## 3.1 Architecture
|
||||
The architecture of Janus is shown in Figure 2 . For pure text understanding, multimodal understanding, and visual generation, we apply independent encoding methods to convert the raw inputs into features, which are then processed by an unified autoregressive transformer. Specifically, for text understanding, we use the built-in tokenizer of the LLM to convert the text into discrete IDs and obtain the feature representations corresponding to each ID. For multimodal understanding, we use the SigLIP [ 92 ] encoder to extract high-dimensional semantic features from images. These features are flattened from a 2 2 -D grid into a 1 1 -D sequence, and an understanding adaptor is used to map these image features into the input space of the LLM.
|
||||
For visual generation tasks, we use the VQ tokenizer from [ 73 ] to convert images into discrete IDs. After the ID sequence is flattened into 1 1 -D, we use a generation adaptor to map the codebook embeddings corresponding to each ID into the input space of the LLM. We then concatenate these feature sequences to form a multimodal feature sequence, which is subsequently fed into the LLM for processing. The built-in prediction head of the LLM is utilized for text predictions in both the pure text understanding and multimodal understanding tasks, while a randomly initialized prediction head is used for image predictions in the visual generation task. The entire model adheres to an autoregressive framework without the need for specially designed attention masks.
|
||||
## 3.2 Training Procedure
|
||||
The training of Janus is divided into three stages, as illustrated in Figure 3 . Details are provided in the below.
|
||||
Stage I: Training Adaptors and Image Head.
|
||||
The main goal of this stage is to create a conceptual connection between visual and linguistic elements within the embedding space, enabling the LLM to understand the entities shown in images and have preliminary visual generation ability. We keep the visual encoders and the LLM frozen during this stage, allowing only the trainable parameters within the understanding adaptor, generation adaptor and image head to be updated.
|
||||
Figure 3: Our Janus adopts a three-stage training procedure. We use flame symbols/snowflake symbols in the diagram to indicate the module updates/does not update its parameters.
|
||||
Stage II: Unified Pretraining.
|
||||
In this stage, we perform unified pretraining with multimodal corpus to enable Janus to learn both multimodal understanding and generation.
|
||||
We unfreeze the LLM and utilize all types of training data: pure text data, multimodal understanding data, and visual generation data. Inspired by Pixart [ 9 ] , we begin by conducting simple visual generation training using ImageNet- 1 1 k to help the model grasp basic pixel dependencies. Subsequently, we enhance the model’s open-domain visual generation capability with general text-to-image data.
|
||||
Stage III: Supervised Fine-tuning. During this stage, we fine-tune the pretrained model with instruction tuning data to enhance its instruction-following and dialogue capabilities. We fine-tune all parameters except the generation encoder.
|
||||
We focus on supervising the answers while masking system and user prompts. To ensure Janus’s proficiency in both multimodal understanding and generation, we don’t fine-tune separate models for a certain task. Instead, we use a blend of pure text dialogue data, multimodal understanding data and visual generation data, ensuring versatility across various scenarios.
|
||||
## 3.3 Training Objective
|
||||
Janus is an autoregressive model, and we simply adopt the cross-entropy loss during training:
|
||||
ℒ = − ∑ i = 1 log P θ ( x i | x < i ) \displaystyle\mathcal{L}=-\sum_{i=1}\log P_{\theta}(x_{i}|x_{<i})
|
||||
(1)
|
||||
Here, P ( ⋅ ∣ ⋅ ) P(\cdot\mid\cdot) indicates the conditional probability modeled by the weights θ \theta of Janus. For pure text understanding and multimodal understanding tasks, we compute the loss on the text sequence. For visual generation tasks, we compute the loss only on the image sequence. To keep the design simple, we have not assigned different loss weights to different tasks.
|
||||
## 3.4 Inference
|
||||
During inference, our model adopts a next-token prediction approach. For pure text understanding and multimodal understanding, we follow the standard practice of sampling tokens sequentially from the predicted distribution. For image generation, we utilize classifier-free guidance (CFG) 2 2
|
||||
2
|
||||
During training, we replace the text condition in the text-to-image data with a pad token at a probability of 10 10 %, enabling the model to have unconditional visual generation capability. , similar to prior works [ 26 , 73 , 8 ] . Specifically, for each token, the logit l g l_{g} is calculated as: l g = l u + s ( l c − l u ) l_{g}=l_{u}+s(l_{c}-l_{u}) , where l c l_{c} is the conditional logit, l u l_{u} is the unconditional logit, and s s is the scale for the classifier-free guidance. The default number of s s is 5 5 for the following evaluation.
|
||||
## 3.5 Possible Extensions
|
||||
It is important to note that our design, which features separate encoders for understanding and generation, is straightforward and easy to extend.
|
||||
Multimodal Understanding. (1) For the multimodal understanding component, a stronger vision encoder can be chosen without worrying about whether the encoder is capable of handling vision generation tasks, such as EVA-CLIP [ 74 ] , InternViT [ 13 ] , etc. (2) To handle high-resolution images, dynamic high-resolution techniques [ 50 ] can be used. This
|
||||
allows the model to scale to any resolution, without performing positional embedding interpolation for ViTs. Tokens can be further compressed to save computational cost, for instance, using pixel shuffle operation [ 12 ] .
|
||||
Visual Generation. (1)
|
||||
For visual generation, finer-grained encoders can be chosen in order to preserve more image details after encoding, such as MoVQGan [ 93 ] .
|
||||
(2) Loss functions specifically designed for visual generation can be employed, such as diffusion loss [ 46 ] .
|
||||
(3) A combination of AR (causal attention) and parallel (bidirectional attention) methods can be used in the visual generation process to reduce accumulated errors during visual generation [ 79 ] .
|
||||
Support for Additional Modalities. The straightforward architecture of Janus allows for easy integration with additional encoders, accommodating various modalities such as 3D point cloud [ 53 ] , tactile [ 88 ] , and EEG [ 4 ] . This gives Janus the potential to become a more powerful multimodal generalist model.
|
||||
## 4 Experiments
|
||||
In this section, we present a series of comprehensive experiments designed to assess the performance of our method across a range of visual understanding and generation tasks. We begin by detailing our experimental setup, which includes the model architecture, training datasets, and evaluation benchmarks. Next, we report the performance of Janus, followed by a comparison with other state-of-the-art models on various benchmarks for multimodal understanding and generation. We also conduct extensive ablation studies to verify the effectiveness of the proposed method. Lastly, we provide some qualitative results.
|
||||
## 4.1 Implementation Details
|
||||
In our experiments, we utilize DeepSeek-LLM ( 1.3 1.3 B) [ 5 ] with a maximum supported sequence length of 4096 4096 as the base language model. For the vision encoder used in understanding tasks, we select SigLIP-Large-Patch 16 16 - 384 384 [ 92 ] . The generation encoder has a codebook of size 16,384 16,384 and downsamples images by a factor of 16 16 . Both the understanding adaptor and the generation adaptor are two-layer MLPs. The detailed hyperparameters for each stage are provided in Table 1 . All images are resized to 384 × 384 384\times 384 pixels.
|
||||
For multimodal understanding data, we resize the long side of the image and pad the short side with the background color (RGB: 127 127 , 127 127 , 127 127 ) to reach 384 384 .
|
||||
For visual generation data, the short side is resized to 384 384 , and the long side is cropped to 384 384 . We use sequence packing during training to improve training efficiency. We mix all data types according to the specified ratios in a single training step. Our Janus is trained and evaluated using HAI-LLM [ 32 ] , which is a lightweight and efficient distributed training framework built on top of PyTorch. The whole training process took 7 7 days on a cluster of 16 16 nodes, each equipped with 8 8 Nvidia A 100 100 ( 40 40 GB) GPUs.
|
||||
Table 1: Detailed hyperparameters of our Janus . Data ratio refers to the ratio of multimodal understanding data, pure text data, and visual generation data.
|
||||
Hyperparameters
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Learning rate
|
||||
1.0 × 10 − 3 1.0\times 10^{-3}
|
||||
1 × 10 − 4 1\times 10^{-4}
|
||||
2.0 × 10 − 5 2.0\times 10^{-5}
|
||||
LR scheduler
|
||||
Cosine
|
||||
Constant
|
||||
Constant
|
||||
Weight decay
|
||||
0.0 0.0
|
||||
0.0 0.0
|
||||
0.1 0.1
|
||||
Gradient clip
|
||||
1.0 1.0
|
||||
1.0 1.0
|
||||
1.0 1.0
|
||||
Optimizer
|
||||
AdamW ( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
Warm-up steps
|
||||
300 300
|
||||
5,000 5,000
|
||||
0 0
|
||||
Training steps
|
||||
10,000 10,000
|
||||
180,000 180,000
|
||||
24,000 24,000
|
||||
Batch size
|
||||
256 256
|
||||
512 512
|
||||
256 256
|
||||
Data Ratio
|
||||
1 : 0 : 1 1:0:1
|
||||
2 : 3 : 5 2:3:5
|
||||
7 : 3 : 10 7:3:10
|
||||
Table 2: Comparison with state-of-the-arts on multimodal understanding benchmarks . “Und.” and “Gen.” denote “understanding” and “generation”, respectively. Models using external pretrained diffusion model are marked with † .
|
||||
Type
|
||||
Model
|
||||
# LLM Params
|
||||
POPE ↑ \uparrow
|
||||
MME-P ↑ \uparrow
|
||||
MMB ↑ \uparrow
|
||||
SEED ↑ \uparrow
|
||||
VQAv2 (test) ↑ \uparrow
|
||||
GQA ↑ \uparrow
|
||||
MMMU ↑ \uparrow
|
||||
MM-Vet ↑ \uparrow
|
||||
Und. Only
|
||||
LLaVA-v 1.5 1.5 -Phi- 1.5 1.5 [ 86 ]
|
||||
1.3 1.3 B
|
||||
84.1 84.1
|
||||
1128.0 1128.0
|
||||
-
|
||||
-
|
||||
75.3 75.3
|
||||
56.5 56.5
|
||||
30.7 30.7
|
||||
-
|
||||
MobileVLM [ 14 ]
|
||||
1.4 1.4 B
|
||||
84.5 84.5
|
||||
1196.2 1196.2
|
||||
53.2 53.2
|
||||
-
|
||||
-
|
||||
56.1 56.1
|
||||
-
|
||||
-
|
||||
MobileVLM-V2 [ 15 ]
|
||||
1.4 1.4 B
|
||||
84.3 84.3
|
||||
1302.8 1302.8
|
||||
57.7 57.7
|
||||
-
|
||||
-
|
||||
59.3 59.3
|
||||
-
|
||||
-
|
||||
MobileVLM [ 14 ]
|
||||
2.7 2.7 B
|
||||
84.9 84.9
|
||||
1288.9 1288.9
|
||||
59.6 59.6
|
||||
-
|
||||
-
|
||||
59.0 59.0
|
||||
-
|
||||
-
|
||||
MobileVLM-V2 [ 15 ]
|
||||
2.7 2.7 B
|
||||
84.7 84.7
|
||||
1440.5 1440.5
|
||||
63.2 63.2
|
||||
-
|
||||
-
|
||||
61.1 61.1
|
||||
-
|
||||
-
|
||||
LLaVA-Phi [ 96 ]
|
||||
2.7 2.7 B
|
||||
85.0 85.0
|
||||
1335.1 1335.1
|
||||
59.8 59.8
|
||||
-
|
||||
71.4 71.4
|
||||
-
|
||||
-
|
||||
28.9 28.9
|
||||
LLaVA [ 51 ]
|
||||
7 7 B
|
||||
76.3 76.3
|
||||
809.6 809.6
|
||||
38.7 38.7
|
||||
33.5 33.5
|
||||
-
|
||||
-
|
||||
-
|
||||
25.5 25.5
|
||||
LLaVA-v 1.5 1.5 [ 50 ]
|
||||
7 7 B
|
||||
85.9 85.9
|
||||
1510.7 1510.7
|
||||
64.3 64.3
|
||||
58.6 58.6
|
||||
78.5 78.5
|
||||
62.0 62.0
|
||||
35.4 35.4
|
||||
31.1 31.1
|
||||
InstructBLIP [ 16 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
36.0 36.0
|
||||
53.4 53.4
|
||||
-
|
||||
49.2 49.2
|
||||
-
|
||||
26.2 26.2
|
||||
Qwen-VL-Chat [ 3 ]
|
||||
7 7 B
|
||||
-
|
||||
1487.5 1487.5
|
||||
60.6 60.6
|
||||
58.2 58.2
|
||||
78.2 78.2
|
||||
57.5 57.5
|
||||
-
|
||||
-
|
||||
IDEFICS- 9 9 B [ 41 ]
|
||||
8 8 B
|
||||
-
|
||||
-
|
||||
48.2 48.2
|
||||
-
|
||||
50.9 50.9
|
||||
38.4 38.4
|
||||
-
|
||||
-
|
||||
Emu 3 3 -Chat [ 83 ]
|
||||
8 8 B
|
||||
85.2 85.2
|
||||
-
|
||||
58.5 58.5
|
||||
68.2 68.2
|
||||
75.1 75.1
|
||||
60.3 60.3
|
||||
31.6 31.6
|
||||
-
|
||||
InstructBLIP [ 16 ]
|
||||
13 13 B
|
||||
78.9 78.9
|
||||
1212.8 1212.8
|
||||
-
|
||||
-
|
||||
-
|
||||
49.5 49.5
|
||||
-
|
||||
25.6 25.6
|
||||
Und. and Gen.
|
||||
DreamLLM † [ 21 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
72.9 72.9
|
||||
-
|
||||
-
|
||||
36.6 36.6
|
||||
LaVIT † [ 36 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
66.0 66.0
|
||||
46.8 46.8
|
||||
-
|
||||
-
|
||||
Emu † [ 75 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
52.0 52.0
|
||||
-
|
||||
-
|
||||
-
|
||||
NExT-GPT † [ 84 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
66.7 66.7
|
||||
-
|
||||
-
|
||||
-
|
||||
Show-o [ 86 ]
|
||||
1.3 1.3 B
|
||||
73.8 73.8
|
||||
948.4 948.4
|
||||
-
|
||||
-
|
||||
59.3 59.3
|
||||
48.7 48.7
|
||||
25.1 25.1
|
||||
-
|
||||
Gemini-Nano-1 [ 78 ]
|
||||
1.8 1.8 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
62.7 62.7
|
||||
-
|
||||
26.3 26.3
|
||||
-
|
||||
LWM [ 52 ]
|
||||
7 7 B
|
||||
75.2 75.2
|
||||
-
|
||||
-
|
||||
-
|
||||
55.8 55.8
|
||||
44.8 44.8
|
||||
-
|
||||
9.6 9.6
|
||||
VILA-U [ 85 ]
|
||||
7 7 B
|
||||
85.8 85.8
|
||||
1401.8 1401.8
|
||||
-
|
||||
59.0 59.0
|
||||
79.4 79.4
|
||||
60.8 60.8
|
||||
-
|
||||
33.5 33.5
|
||||
Chameleon [ 77 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
22.4 22.4
|
||||
8.3 8.3
|
||||
Janus (Ours)
|
||||
1.3 1.3 B
|
||||
87.0 87.0
|
||||
1338.0 1338.0
|
||||
69.4 69.4
|
||||
63.7 63.7
|
||||
77.3 77.3
|
||||
59.1 59.1
|
||||
30.5 30.5
|
||||
34.3 34.3
|
||||
## 4.2 Data Setup
|
||||
In this section, we provide details of the pretraining and supervised finetuning datasets.
|
||||
Stage I. We use a dataset that includes 1.25 1.25 million image-text paired captions from ShareGPT 4 4 V [ 10 ] for multimodal understanding and approximately 1.2 1.2 million samples from ImageNet- 1 1 k [ 18 ] for visual generation. The ShareGPT 4 4 V data is formatted as “ <image> <text> ”. The ImageNet data is organized into a text-to-image data format using the category names: “ <category_name> <image> ”. Here, the “<>” symbols represent placeholders.
|
||||
Stage II. We organize the data into the following categories. (1) Text-only data. We use pretraining text copus from DeepSeek-LLM [ 5 ] . (2) Interleaved image-text data. We use WikiHow [ 39 ] and WIT [ 72 ] dataset. (3) Image caption data. We use images from [ 18 , 38 , 17 , 40 , 23 , 47 , 49 , 45 , 70 ] .
|
||||
Among them, we employ open-source multimodal model to re-caption images in [ 17 , 40 ] . The image caption data is formatted into question-answer pairs, for example, “ <image> Describe the image in detail. <caption> ”.
|
||||
(4) Table and chart data. We use corresponding table and chart data from DeepSeek-VL [ 55 ] . The data is formatted as “ <question> <answer> ”. (5) Visual generation data. We utilize image-caption pairs from various datasets including [ 38 , 17 , 40 , 70 , 57 , 58 , 63 , 60 ] , along with 2 2 M in-house data. For images from [ 38 , 70 ] , we filter based on aesthetic scores and image sizes, resulting in 20 % 20\% remaining. During training, we randomly use only the first sentence of a caption with a 25 % 25\% probability to encourage the model to develop strong generation capabilities for short descriptions. ImageNet samples [ 18 ] are presented only during the first 120 120 K training steps, while images from other datasets appear in the later 60 60 K steps. This approach helps the model first learn basic pixel dependencies before progressing to more complex scene understanding, as suggested by [ 9 ] . The visual generation data is provided in the format: “ <caption> <image> ”.
|
||||
Stage III. For text understanding, we use data from [ 43 ] . For multimodal understanding, we use instruct tuning data from [ 43 , 31 , 35 , 69 , 34 , 56 ] . For visual generation, we use image-text pairs from [ 17 , 70 , 60 ] (a subset of that in stage II) and 4 4 M in-house data. We utilize the following format for instruction tuning:“ User:<Input Message> \n Assistant: <Response> ”. For multi-turn dialogues, we repeat this format to structure the data.
|
||||
## 4.3 Evaluation Setup
|
||||
Multimodal Understanding.
|
||||
To assess multimodal understanding capabilities, we evaluate our model on widely recognized image-based vision-language benchmarks, which include VQAv2 [ 31 ] , GQA [ 35 ] , POPE [ 48 ] , MME [ 25 ] , SEED [ 42 ] , MMB [ 54 ] , MM-Vet [ 90 ] , and MMMU [ 91 ] .
|
||||
Visual Generation.
|
||||
For evaluating visual generation capabilities, we use the MSCOCO- 30 30 K [ 11 ] , MJHQ- 30 30 K [ 44 ] , and GenEval [ 30 ] benchmarks. MSCOCO- 30 30 K and MJHQ- 30 30 K employ the Fréchet Inception Distance (FID) metric on generated images compared to 30 30 K high-quality images, which indicates the overall efficacy of image generation. GenEval is a challenging benchmark for image-to-text generation, designed to reflect the comprehensive generative abilities of visual generation models by offering a detailed instance-level analysis of their compositional capabilities.
|
||||
Table 3: Evaluation of text-to-image generation ability on GenEval benchmark . “Und.” and “Gen.” denote “understanding” and “generation”, respectively. Models using external pretrained diffusion model are marked with † .
|
||||
Type
|
||||
Method
|
||||
# Params
|
||||
Single Obj.
|
||||
Two Obj.
|
||||
Counting
|
||||
Colors
|
||||
Position
|
||||
Color Attri.
|
||||
Overall ↑ \uparrow
|
||||
Gen. Only
|
||||
LlamaGen [ 73 ]
|
||||
0.8 0.8 B
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.21 0.21
|
||||
0.58 0.58
|
||||
0.07 0.07
|
||||
0.04 0.04
|
||||
0.32 0.32
|
||||
LDM [ 67 ]
|
||||
1.4 1.4 B
|
||||
0.92 0.92
|
||||
0.29 0.29
|
||||
0.23 0.23
|
||||
0.70 0.70
|
||||
0.02 0.02
|
||||
0.05 0.05
|
||||
0.37 0.37
|
||||
SDv 1.5 1.5 [ 67 ]
|
||||
0.9 0.9 B
|
||||
0.97 0.97
|
||||
0.38 0.38
|
||||
0.35 0.35
|
||||
0.76 0.76
|
||||
0.04 0.04
|
||||
0.06 0.06
|
||||
0.43 0.43
|
||||
PixArt- α \alpha [ 9 ]
|
||||
0.6 0.6 B
|
||||
0.98 0.98
|
||||
0.50 0.50
|
||||
0.44 0.44
|
||||
0.80 0.80
|
||||
0.08 0.08
|
||||
0.07 0.07
|
||||
0.48 0.48
|
||||
SDv 2.1 2.1 [ 67 ]
|
||||
0.9 0.9 B
|
||||
0.98 0.98
|
||||
0.51 0.51
|
||||
0.44 0.44
|
||||
0.85 0.85
|
||||
0.07 0.07
|
||||
0.17 0.17
|
||||
0.50 0.50
|
||||
DALL-E 2 2 [ 66 ]
|
||||
6.5 6.5 B
|
||||
0.94 0.94
|
||||
0.66 0.66
|
||||
0.49 0.49
|
||||
0.77 0.77
|
||||
0.10 0.10
|
||||
0.19 0.19
|
||||
0.52 0.52
|
||||
Emu 3 3 -Gen [ 83 ]
|
||||
8 8 B
|
||||
0.98 0.98
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.81 0.81
|
||||
0.17 0.17
|
||||
0.21 0.21
|
||||
0.54 0.54
|
||||
SDXL [ 62 ]
|
||||
2.6 2.6 B
|
||||
0.98 0.98
|
||||
0.74 0.74
|
||||
0.39 0.39
|
||||
0.85 0.85
|
||||
0.15 0.15
|
||||
0.23 0.23
|
||||
0.55 0.55
|
||||
Und. and Gen.
|
||||
SEED-X † [ 29 ]
|
||||
17 17 B
|
||||
0.97 0.97
|
||||
0.58 0.58
|
||||
0.26 0.26
|
||||
0.80 0.80
|
||||
0.19 0.19
|
||||
0.14 0.14
|
||||
0.49 0.49
|
||||
Show-o [ 86 ]
|
||||
1.3 1.3 B
|
||||
0.95 0.95
|
||||
0.52 0.52
|
||||
0.49 0.49
|
||||
0.82 0.82
|
||||
0.11 0.11
|
||||
0.28 0.28
|
||||
0.53 0.53
|
||||
LWM [ 52 ]
|
||||
7 7 B
|
||||
0.93 0.93
|
||||
0.41 0.41
|
||||
0.46 0.46
|
||||
0.79 0.79
|
||||
0.09 0.09
|
||||
0.15 0.15
|
||||
0.47 0.47
|
||||
Chameleon [ 77 ]
|
||||
34 34 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.39 0.39
|
||||
Janus (Ours)
|
||||
1.3 1.3 B
|
||||
0.97 0.97
|
||||
0.68 0.68
|
||||
0.30 0.30
|
||||
0.84 0.84
|
||||
0.46 0.46
|
||||
0.42 0.42
|
||||
0.61 0.61
|
||||
## 4.4 Comparison with State-of-the-arts
|
||||
Multimodal Understanding Performance.
|
||||
We compare the proposed method with state-of-the-art unified models and understanding-only models in Table 2 .
|
||||
Janus achieves the overall best results among models of similar scale. Specifically, compared to the previous best unified model, Show-o [ 86 ] , we achieve performance improvements of 41 41 % ( 949 → 1338 949\rightarrow 1338 ) and 30 30 % ( 48.7 → 59.1 48.7\rightarrow 59.1 ) on the MME and GQA datasets, respectively. This can be attributed to Janus decoupling the visual encoding for multimodal understanding and generation, mitigating the conflict between these two tasks.
|
||||
When compared to models with significantly larger sizes, Janus remains highly competitive. For instance, Janus outperforms LLaVA-v 1.5 1.5 ( 7 7 B) on several datasets, including POPE, MMbench, SEED Bench, and MM-Vet.
|
||||
Visual Generation Performance.
|
||||
We report visual generation performance on GenEval, COCO- 30 30 K and MJHQ- 30 30 K benchmarks. As shown in Table 3 , our Janus obtains 61 61 % overall accuracy on GenEval, which outperforms the previous best unified model Show-o ( 53 53 %) and some popular generation-only methods, e.g., SDXL ( 55 55 %) and DALL-E 2 2 ( 52 52 %). This demonstrates that our approach has better instruction-following capabilities. As shown in Table 4 , Janus achieves FIDs of 8.53 8.53 and 10.10 10.10 on the COCO- 30 30 K and MJHQ- 30 30 K benchmarks, respectively, surpassing unified models Show-o and LWM, and demonstrating competitive performance compared to some well-known generation-only methods. This demonstrates that the images generated by Janus have good quality and highlights its potential in visual generation.
|
||||
Table 4: Evaluation of text-to-image generation ability on MSCOCO- 30 30 K and MJHQ- 30 30 K benchmark . “Und.” and “Gen.” denote “understanding” and “generation”, respectively. Models using external pretrained diffusion model are marked with † .
|
||||
Type
|
||||
Model
|
||||
# Params
|
||||
COCO-30K ↓ \downarrow
|
||||
MJHQ-30K ↓ \downarrow
|
||||
Gen. Only
|
||||
DALL·E [ 65 ]
|
||||
12 12 B
|
||||
27.50 27.50
|
||||
-
|
||||
GLIDE [ 59 ]
|
||||
5 5 B
|
||||
12.24 12.24
|
||||
-
|
||||
LDM [ 67 ]
|
||||
1.4 1.4 B
|
||||
12.64 12.64
|
||||
-
|
||||
DALL·E 2 [ 66 ]
|
||||
6.5 6.5 B
|
||||
10.39 10.39
|
||||
-
|
||||
SDv1.5 [ 67 ]
|
||||
0.9 0.9 B
|
||||
9.62 9.62
|
||||
-
|
||||
GigaGAN [ 37 ]
|
||||
0.9 0.9 B
|
||||
9.09 9.09
|
||||
-
|
||||
PixArt- α \alpha [ 9 ]
|
||||
0.6 0.6 B
|
||||
7.32 7.32
|
||||
-
|
||||
Imagen [ 68 ]
|
||||
34 34 B
|
||||
7.27 7.27
|
||||
-
|
||||
RAPHAEL [ 87 ]
|
||||
3 3 B
|
||||
6.61 6.61
|
||||
-
|
||||
Und. and Gen.
|
||||
Emu † [ 75 ]
|
||||
13 13 B
|
||||
11.66 11.66
|
||||
-
|
||||
NExT-GPT † [ 84 ]
|
||||
13 13 B
|
||||
11.28 11.28
|
||||
-
|
||||
SEED-X † [ 29 ]
|
||||
17 17 B
|
||||
14.99 14.99
|
||||
-
|
||||
Show-o [ 86 ]
|
||||
1.3 1.3 B
|
||||
9.24 9.24
|
||||
15.18 15.18
|
||||
LWM [ 52 ]
|
||||
7 7 B
|
||||
12.68 12.68
|
||||
17.77 17.77
|
||||
VILA-U ( 256 256 ) [ 85 ]
|
||||
7 7 B
|
||||
-
|
||||
12.81 12.81
|
||||
VILA-U ( 384 384 ) [ 85 ]
|
||||
7 7 B
|
||||
-
|
||||
7.69 7.69
|
||||
Janus (Ours)
|
||||
1.3 1.3 B
|
||||
8.53 8.53
|
||||
10.10 10.10
|
||||
Table 5: Ablation studies . We verify the effectiveness of decoupling visual encoding and compare unified training with task-specific training. “Und.”, “Gen.” and “SE. Tokenizer” denote “understanding”, “generation” and “semantic tokenizer”, respectively.
|
||||
Exp ID
|
||||
Visual Encoder
|
||||
Training Task
|
||||
POPE ↑ \uparrow
|
||||
MMB ↑ \uparrow
|
||||
SEED ↑ \uparrow
|
||||
MMMU ↑ \uparrow
|
||||
COCO-FID ↓ \downarrow
|
||||
A
|
||||
VQ Tokenizer
|
||||
Und. + Gen.
|
||||
60.1 60.1
|
||||
35.0 35.0
|
||||
34.9 34.9
|
||||
24.7 24.7
|
||||
8.72 8.72
|
||||
B
|
||||
SE. Tokenizer
|
||||
Und. + Gen.
|
||||
82.4 82.4
|
||||
52.7 52.7
|
||||
54.9 54.9
|
||||
26.6 26.6
|
||||
7.11 7.11
|
||||
C
|
||||
SE. Tokenizer
|
||||
Und.
|
||||
83.9 83.9
|
||||
62.1 62.1
|
||||
60.8 60.8
|
||||
27.5 27.5
|
||||
-
|
||||
D
|
||||
SigLIP + VQ (Ours)
|
||||
Und. + Gen.
|
||||
87.0 87.0
|
||||
69.4 69.4
|
||||
63.7 63.7
|
||||
30.5 30.5
|
||||
8.53 8.53
|
||||
E
|
||||
SigLIP
|
||||
Und.
|
||||
85.9 85.9
|
||||
70.6 70.6
|
||||
64.8 64.8
|
||||
28.8 28.8
|
||||
-
|
||||
F
|
||||
VQ Tokenizer
|
||||
Gen.
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
8.92 8.92
|
||||
## 4.5 Ablation Studies
|
||||
We carefully design ablation studies to verify the effectiveness of Janus’s design concept. First, we design experiments to validate the importance and benefits of decoupling visual encoding. Second, we investigate the impact of unified training on individual tasks like multimodal understanding or visual generation. Results are listed in Table 5 .
|
||||
Baseline Construction. Following previous work [ 77 ] , we select a VQ tokenizer [ 73 ] to encode images for both multimodal understanding and generation tasks, serving as the baseline (Exp-A). Considering that the VQ tokenizer in Exp-A might be weak in extract semantic information, making it less effective for multimodal understanding, we also construct a stronger baseline Exp-B. We adopt SigLIP to distill an enhanced semantic tokenizer 3 3
|
||||
3
|
||||
The semantic tokenizer is only used in the ablation study as a stronger baseline. For simplicity, we use the ordinary VQ tokenizer [ 73 ] in the main experiment. that can extract high-level semantic information from images while also have the ability to convert images into discrete IDs, which is similar to that in [ 85 ] . Details of the semantic tokenizer could be found in the Appendix A.1 .
|
||||
Impact of Decoupling Visual Encoding.
|
||||
(1) From the results of Exp-A, we find the model achieves satisfactory performance on visual generation benchmark ( 8.72 8.72 FID on COCO). However, there is a significant gap on understanding benchmarks between Exp-A and our model (Exp-D).
|
||||
(2) When comparing Exp-B to Exp-A, the results show a clear improvement in multimodal understanding, although there is still a considerable gap compared to our method. In terms of visual generation, Exp-B outperforms Exp-D. We hypothesize two possible reasons for this. First, the semantic tokenizer produces discrete IDs that are more semantically coherent, providing more reasonable prediction targets for the LLM. Second, the visual encoder in Exp-B has significantly more parameters than the Gen. encoder in Exp-D.
|
||||
(3) To investigate whether using a single visual encoder leads to a trade-off between multimodal understanding and generation, we further design Exp-C based on Exp-B, which focuses solely on multimodal understanding training. The multimodal understanding ability of Exp-C is significantly better than that of Exp-B. This indicates that the visual encoder in Exp-B made trade-offs between multimodal understanding and generation, ultimately sacrificing its multimodal understanding capability. The above experiments illustrate the importance of decoupling visual encoding.
|
||||
Unified Model vs. Pure Understanding & Pure Generation. We compare the performance of unified training (Exp-D) against pure understanding (Exp-E) and pure generation (Exp-F) training. For pure understanding, we omit visual generation data. For pure generation, we exclude the understanding data.
|
||||
Please note that unified training and pure understanding training go through the same steps for the understanding part. Similarly, unified training and pure generation training go through the same steps for the visual generation part.
|
||||
Experimental results show that the performance of unified training is comparable to that of training solely for understanding or solely for visual generation. This demonstrates that our model, Janus, is capable of incorporating strong generative abilities while minimally affecting multimodal understanding performance.
|
||||
Figure 4: Qualitative comparisons of visual generation with LlamaGen and SDXL . The images generated by Janus show better consistency with the user’s prompts. The image resolutions for SDXL, LlamaGen, and ours are 1024 × 1024 1024\times 1024 , 512 × 512 512\times 512 , and 384 × 384 384\times 384 , respectively. Best viewed on screen.
|
||||
## 4.6 Qualitative Results
|
||||
Figure 5: Qualitative results of multimodal understanding on humorous memes . We compare the response with Chameleon- 7 7 B [ 77 ] and Show-o [ 86 ] . We emphasize the key-points in the response. Best viewed on screen.
|
||||
Visualizations of Visual Generation. Figure 4 provides qualitative comparisons between our model, diffusion-based models like SDXL [ 62 ] , and the autoregressive model LlamaGen [ 73 ] . The results show that our model demonstrates superior instruction-following capabilities in visual generation, accurately capturing most of details in the user’s prompt. This indicates the potential of the unified model in the realm of visual generation. More visualizations can be found in the Appendix B .
|
||||
Multimodal Understanding on MEME Images. Figure 5 showcases the qualitative results of Janus’s multimodal understanding ability, compared with Chameleon [ 77 ] and Show-o [ 86 ] . Janus accurately interprets the text caption and captures the emotion conveyed in the meme. In contrast, both Chameleon and Show-o struggle with accurately recognizing the text in the image. Additionally, Chameleon fails to identify objects in the meme, while Show-o misinterprets the dog’s color. These examples highlight that the decoupled vision encoder significantly enhances Janus’s fine-grained multimodal understanding ability compared to the shared encoder used by Chameleon and Show-o. More multimodal understanding exmples can be found in the Appendix B .
|
||||
## 5 Conclusion
|
||||
In this paper, we introduced Janus, a simple, unified and extensible multimodal understanding and generation model. The core idea of Janus is to decouple visual encoding for multimodal understanding and generation, which could alleviate the conflict arising from the differing demands that understanding and generation place on the visual encoder. Extensive experiments have demonstrated the effectiveness and leading performance of Janus. It is also worth noting that Janus is flexible and easy to extend. In addition to having significant potential for improvement in both multimodal understanding and generation, Janus is also easily extendable to incorporate more input modalities. The above advantages suggest that Janus may serve as an inspiration for the development of the next generation of multimodal general-purpose models.
|
||||
## Appendix
|
||||
## Appendix A Details of Semantic Tokenizer Mentioned in Ablation Study
|
||||
## A.1 Architecture of Semantic Tokenizer
|
||||
Figure 6: Architecture and usage of the semantic tokenizer. (a) Architecture used during training of the semantic tokenizer. We use pre-trained SigLIP [ 92 ] to supervise the reconstruction of semantic information, while using raw image to supervise the reconstruction of RGB values. (b) Integrating LLM with the semantic decoder. The semantic decoder outputs continuous features with high-level semantics, which are passed through an adaptor and then used as input for the LLM. Please note that the semantic tokenizer is only used in the ablation study, not in the main experiment.
|
||||
We build the semantic tokenizer based on the tokenizer architecture proposed in [ 73 ] , which has a downsample rate of 16 16 . In addition to the original CNN pixel decoder, we add an additional semantic decoder branch after Vector Quantization, as shown in Figure 6 (a). The semantic decoder is a 12 12 -layer ViT [ 22 ] , with 12 12 attention heads and a hidden dimension of 768 768 . For the semantic decoder, we use a causal attention mask to facilitate next token prediction when integrating it with an LLM.
|
||||
## A.2 Training
|
||||
Training Procedure.
|
||||
The semantic tokenizer is trained from scratch in a two-stage manner. In the first stage, we train the model on the ImageNet- 1 1 k [ 18 ] dataset for 40 40 epochs. In the second stage, we fine-tune the model for 1 1 epoch on 50 50 million images. These images come from the visual generation data used during the Janus pretraining process. We use a constant learning rate of 1 e − 4 1e-4 and a batch size of 128 128 .
|
||||
Training Loss.
|
||||
The training loss of the semantic tokenizer consists of two parts. On one hand, we use the loss for RGB reconstruction as described in [ 73 ] . On the other hand, we use SigLIP-Large-Patch 16 16 - 384 384 as the teacher to supervise the semantic feature reconstruction results by the semantic decoder.
|
||||
We adopt the loss in BEiT-v2 [ 61 ] . Specifically, we maximize the cosine similarity between the semantic feature predicted by the semantic decoder and the SigLIP output. The weight for the semantic reconstruction loss is set to 0.25 0.25 .
|
||||
## A.3 Integrating with LLM
|
||||
We present the integration of the semantic tokenizer and the LLM in Figure 6 (b). The image is first transformed into continuous features through the CNN encoder, vector quantization and the semantic decoder. Then, the LLM processes these features and generates predictions for the image IDs. Finally, the pixel decoder converts these discrete IDs into RGB values.
|
||||
## Appendix B Additional Qualitative Results
|
||||
More Visualizations of Text-to-Image Generation. We present more text-to-image generation results in Figure 7 . It is evident that Janus is capable of producing high-quality images that adhere closely to the given prompts. We further explore the multilingual text-to-image capabilities of our model, as shown in Figure 8 . We are pleasantly surprised to find that, despite our training data consisting solely of English text-to-image data, Janus can still process text-to-image tasks in other languages. We attribute this multilingual ability to the original large language model’s inherent traits. The LLM initially translates various languages into a unified semantic space, allowing Janus to perform text-to-image tasks naturally without additional training.
|
||||
More Multimodal Understanding Results. Additional results on multimodal understanding are shown in Figure 9 . Janus exhibits impressive comprehension abilities when handling inputs from various contexts, showcasing its powerful capabilities.
|
||||
Figure 7: More text-to-image generation results. We upsample the images to 1024 × 1024 1024\times 1024 for better visualization.
|
||||
Figure 8: Multilingual text-to-image generation samples compared to LlamaGen [ 73 ] . Note that we only use English text-to-image data in training, and this is an emergent capability of our model. The languages used in the prompt, from left to right, are: English, Chinese, French, Japanese, and English with emoji.
|
||||
Figure 9: More multimodal understanding results. Janus has a strong multimodal understanding capability and can handle inputs from various contexts, such as scientific charts, artwork images, LaTeX formula images, and more.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,988 @@
|
||||
TITLE: JanusFlow: Harmonizing Autoregression and Rectified Flow for Unified Multimodal Understanding and Generation
|
||||
ARXIV: 2411.07975
|
||||
AUTHORS: Yiyang Ma , Xingchao Liu , Xiaokang Chen , Wen Liu , Chengyue Wu , Zhiyu Wu , Zizheng Pan , Zhenda Xie , Haowei Zhang , Xingkai yu , Liang Zhao , Yisong Wang , Jiaying Liu , Chong Ruan
|
||||
DATE: [Submitted on 12 Nov 2024 ( v1 ), last revised 24 Mar 2025 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## JanusFlow: Harmonizing Autoregression and Rectified Flow for Unified Multimodal Understanding and Generation
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Related Work
|
||||
## Visual Generation with Flow-based Generative Models.
|
||||
## Unified Models For Understanding and Generation.
|
||||
## 3 JanusFlow
|
||||
## 3.1 Background
|
||||
## Multimodal LLMs.
|
||||
## Rectified Flow.
|
||||
## 3.2 A Unified Framework for Multimodal Understanding and Generation
|
||||
## Multimodal Understanding.
|
||||
## Image Generation.
|
||||
## Decoupling Encoders for the Two Tasks.
|
||||
## 3.3 Training Schemes
|
||||
## Stage 1: Adaptation of Randomly Initialized Components.
|
||||
## Stage 2: Unified Pre-Training.
|
||||
## Stage 3: Supervised Fine-Tuning (SFT).
|
||||
## 3.4 Training Objective
|
||||
## Autoregression Objective.
|
||||
## Rectified Flow Objective.
|
||||
## Representation Alignment Regularization.
|
||||
## Summary.
|
||||
## 4 Experiments
|
||||
## 4.1 Experiment Setup and Implementation Details
|
||||
## 4.2 Training Data Settings
|
||||
## Data for Stage 1 and Stage 2.
|
||||
## 4.3 Evaluation Settings
|
||||
## Image Generation.
|
||||
## Multimodal Understanding.
|
||||
## 4.4 Quantitative Results
|
||||
## 4.5 Ablation Studies
|
||||
## 4.6 Qualitative Results
|
||||
## 5 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
JanusFlow: Harmonizing Autoregression and Rectified Flow for Unified Multimodal Understanding and Generation
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Work
|
||||
Visual Generation with Flow-based Generative Models.
|
||||
Unified Models For Understanding and Generation.
|
||||
3 JanusFlow
|
||||
3.1 Background
|
||||
Multimodal LLMs.
|
||||
Rectified Flow.
|
||||
3.2 A Unified Framework for Multimodal Understanding and Generation
|
||||
Multimodal Understanding.
|
||||
Image Generation.
|
||||
Decoupling Encoders for the Two Tasks.
|
||||
3.3 Training Schemes
|
||||
Stage 1: Adaptation of Randomly Initialized Components.
|
||||
Stage 2: Unified Pre-Training.
|
||||
Stage 3: Supervised Fine-Tuning (SFT).
|
||||
3.4 Training Objective
|
||||
Autoregression Objective.
|
||||
Rectified Flow Objective.
|
||||
Representation Alignment Regularization.
|
||||
Summary.
|
||||
4 Experiments
|
||||
4.1 Experiment Setup and Implementation Details
|
||||
4.2 Training Data Settings
|
||||
Data for Stage 1 and Stage 2.
|
||||
4.3 Evaluation Settings
|
||||
Image Generation.
|
||||
Multimodal Understanding.
|
||||
4.4 Quantitative Results
|
||||
4.5 Ablation Studies
|
||||
4.6 Qualitative Results
|
||||
5 Conclusion
|
||||
References
|
||||
A Performance Analysis of 256 Resolution Model
|
||||
B Details of the Datasets
|
||||
C Analysis of CFG Factor and Sampling Steps
|
||||
D Details of REPA Ablation
|
||||
E Additional Qualitative Results
|
||||
License: CC BY 4.0
|
||||
arXiv:2411.07975v2 [cs.CV] 24 Mar 2025
|
||||
\reportnumber
|
||||
001
|
||||
## JanusFlow: Harmonizing Autoregression and Rectified Flow for Unified Multimodal Understanding and Generation
|
||||
Yiyang Ma 1,2 Xingchao Liu 1,† Xiaokang Chen 1,† Wen Liu 1,† Chengyue Wu 1,3 Zhiyu Wu 1,2 Zizheng Pan 1 Zhenda Xie 1 Haowei Zhang 1 Xingkai Yu 1 Liang Zhao 1 Yisong Wang 1,4 Jiaying Liu 2 Chong Ruan 1,‡
|
||||
1 DeepSeek-AI 2 Peking University 3 The University of Hong Kong 4 Tsinghua University
|
||||
† Equal contribution, ‡ Corresponding author
|
||||
Project Page: https://github.com/deepseek-ai/Janus
|
||||
## Abstract
|
||||
We present JanusFlow , a powerful framework that unifies image understanding and generation in a single model.
|
||||
JanusFlow introduces a minimalist architecture that integrates autoregressive language models with rectified flow, a state-of-the-art method in generative modeling.
|
||||
Our key finding demonstrates that rectified flow can be straightforwardly trained within the large language model framework, eliminating the need for complex architectural modifications.
|
||||
To further improve the performance of our unified model, we adopt two key strategies: (i) decoupling the understanding and generation encoders, and (ii) aligning their representations during unified training.
|
||||
Extensive experiments show that JanusFlow achieves comparable or superior performance to specialized models in their respective domains, while significantly outperforming existing unified approaches across standard benchmarks.
|
||||
This work represents a step toward more efficient and versatile vision-language models.
|
||||
## 1 Introduction
|
||||
(a) Benchmark Performances.
|
||||
(b) Visual Generation Results.
|
||||
Figure 1:
|
||||
Multimodal understanding and image generation with JanusFlow . JanusFlow surpasses the state-of-the-art unified multimodal models and several task-specific understanding models on visual understanding benchmarks. It is also capable of generating high-quality images. The resolution of the images is 384 × 384 384\times 384 .
|
||||
Large language models (LLMs) have demonstrated remarkable capabilities in learning diverse knowledge and generalizing to new scenarios [ 91 , 69 , 1 , 8 , 7 ] . Leveraging these capabilities, researchers have developed sophisticated models specialized in image comprehension [ 58 , 56 , 47 , 2 , 15 , 49 ] and text-to-image generation [ 79 , 76 , 73 , 23 ] .
|
||||
The field has recently shifted toward creating unified systems capable of handling both tasks simultaneously. One prominent direction involves utilizing pre-trained text-to-image models for high-quality generation while training LLMs to generate conditions for these models [ 25 , 26 , 87 , 27 , 19 ] . However, this approach introduces architectural complexity and potentially constrains the model’s capabilities through maintaining separate LLM and generative components. Alternative approaches [ 88 , 99 , 100 , 108 , 97 ] propose training a single LLM for both tasks, typically incorporating either diffusion models [ 32 , 83 ] or vector-quantized autoregressive models [ 22 , 86 ] .
|
||||
Our approach builds upon recent breakthroughs in rectified flow models [ 61 , 55 , 3 , 23 , 62 ] , which provide a simple framework for generative modeling while delivering exceptional empirical performance [ 23 , 45 , 36 ] .
|
||||
Building on these advances, we propose JanusFlow , a powerful unified multimodal model that seamlessly integrates rectified flow with LLM architecture.
|
||||
Following a minimalist design principle, our architecture requires only a lightweight encoder and decoder to adapt the LLM for rectified flow operations.
|
||||
To optimize JanusFlow’s performance, we implement two key strategies: First, we maintain separate vision encoders for understanding and generation tasks, preventing task interference and thus enhancing comprehension capabilities. Second, we align the intermediate representations between generation and understanding modules during training, strengthening semantic coherence in the generation process.
|
||||
JanusFlow shows state-of-the-art performances in both multimodal comprehension and text-to-image generation compared to existing unified approaches, and even outperforms several specialized methods. Specifically, on text-to-image generation benchmarks, MJHQ FID- 30 30 k [ 48 ] , GenEval [ 28 ] and DPG-Bench [ 34 ] , JanusFlow achieves scores of 9.51 9.51 , 0.63 0.63 and 80.09 % 80.09\% , surpassing established text-to-image models including SDv1.5 [ 77 ] and SDXL [ 73 ] . In multimodal comprehension benchmarks, JanusFlow attains scores of 74.9 74.9 , 70.5 70.5 and 60.3 60.3 on MMBench [ 63 ] , SeedBench [ 46 ] , and GQA [ 35 ] , respectively, exceeding specialized models such as LLaVA-v1.5 [ 56 ] and Qwen-VL-Chat [ 4 ] . Notably, these results are achieved with a compact LLM architecture with only 1.3B parameters.
|
||||
## 2 Related Work
|
||||
## Visual Generation with Flow-based Generative Models.
|
||||
Recent years have witnessed remarkable progress in visual generation through diffusion models [ 32 , 83 ] , leading to impressive models like [ 77 , 73 , 78 , 79 , 76 , 67 ] . Building on these advances, flow-based generative models [ 61 , 55 , 3 ] emerged as a simplified alternative framework. These approaches have recently enabled advanced visual generation models [ 23 , 36 ] that achieve superior empirical performance with faster sampling. Our work demonstrates that rectified flow [ 61 , 60 , 62 ] can be effectively integrated into LLMs, creating unified models that excel in both understanding and generation tasks.
|
||||
## Unified Models For Understanding and Generation.
|
||||
The development of multimodal large language models (MLLMs) has enabled effective integration of text and visual information. Building upon powerful LLMs [ 91 , 92 , 7 ] , recent MLLMs [ 64 , 58 , 56 , 2 , 15 , 49 ] have demonstrated exceptional multimodal understanding capabilities. Current research increasingly focuses on architectures that can simultaneously handle visual understanding and generation tasks. One approach extends MLLMs with pre-trained diffusion models [ 25 , 26 , 87 , 27 , 19 , 101 ] . However, these systems essentially utilize diffusion models as external tools, where the MLLM generates conditions for image generation without possessing direct generative capabilities. This separation often results in suboptimal performance compared to standalone diffusion models [ 25 , 87 ] . Another line of work [ 88 , 99 , 100 , 108 , 97 ] aim to train a single LLM for both tasks. Many of these methods employ vector-quantization [ 22 , 86 ] to convert images into discrete tokens, enabling unified autoregressive processing [ 88 , 97 ] . While straightforward to implement, these approaches are inherently limited by their image tokenization quality.
|
||||
Our work focuses on developing unified models that combine autoregressive capabilities with flow/diffusion models, leveraging their proven effectiveness in visual generation. Compared to similar approaches [ 100 , 108 , 107 ] , JanusFlow offers three key advantages: (i) a simple yet effective generation process using rectified flow, (ii) enhanced performance through decoupled vision encoders that resolve inter-task conflicts, and (iii) improved generation quality through representation alignment regularization, enabled by our decoupled encoder design.
|
||||
## 3 JanusFlow
|
||||
In this section, we introduce the architecture of JanusFlow and our training strategies.
|
||||
## 3.1 Background
|
||||
## Multimodal LLMs.
|
||||
Given a dataset 𝒟 \mathcal{D} containing discrete token sequences, each of which can be formulated as x = ( x 1 , ⋯ , x ℓ ) x=(x_{1},\cdots,x_{\ell}) , large language models (LLMs) are trained to model the sequence distribution in an autoregressive manner,
|
||||
log P θ L L M ( x ) = ∑ i = 0 ℓ − 1 log P θ L L M ( x i + 1 | x 1 , … , x i ) , \log{{\rm{P}}}_{\theta_{LLM}}(x)=\sum_{i=0}^{\ell-1}\log{{\rm{P}}}_{\theta_{LLM}}(x_{i+1}|x_{1},\dots,x_{i}),
|
||||
(1)
|
||||
where θ L L M \theta_{LLM} denotes the parameters of the LLM and ℓ \ell is the sequence length. After being trained on large-scale datasets, LLMs exhibit the ability to generalize across various tasks and follow diverse instructions [ 8 , 1 , 69 ] . To extend these models to handle visual inputs, LLMs are augmented with vision encoders [ 58 , 56 , 2 ] . For instance, LLaVA [ 58 ] integrates an LLM with a pre-trained CLIP [ 75 ] image encoder via a projection layer, transforming the extracted image features into a joint embedding space that the LLM can process as word embeddings. By leveraging large-scale multimodal datasets and increasingly powerful LLMs, this architecture has facilitated the development of advanced multimodal models capable of addressing a wide range of vision-language tasks [ 4 , 56 , 47 , 64 ] .
|
||||
## Rectified Flow.
|
||||
For a dataset 𝒟 \mathcal{D} consisting of continuous d d -dimensional data points x = ( x 1 , ⋯ , x d ) x=(x_{1},\cdots,x_{d}) drawn from an unknown data distribution π 1 \pi_{1} , rectified flow [ 61 , 55 ] models the data distribution by
|
||||
learning an ordinary differential equation (ODE) defined over time t ∈ [ 0 , 1 ] t\in[0,1] :
|
||||
d z t d t = v θ N N ( z t , t ) , z 0 ∼ π 0 , \frac{{\rm{d}}z_{t}}{{\rm{d}}t}=v_{\theta_{NN}}(z_{t},t),\penalty\ \penalty\ \penalty\ \penalty\ \penalty\ z_{0}\sim\pi_{0},
|
||||
(2)
|
||||
where θ N N \theta_{NN} represents the parameters of the velocity neural network and π 0 \pi_{0} is a simple distribution, typically standard Gaussian noise 𝒩 ( 0 , I ) \mathcal{N}(0,I) .
|
||||
The network is trained by minimizing the Euclidean distance between the neural velocity and the directions of linear paths connecting random points from π 0 \pi_{0} and π 1 \pi_{1} ,
|
||||
min θ 𝔼 t ∼ P ( t ) , z 0 ∼ π 0 , x ∼ π 1 [ | | v θ N N ( z t , t ) − ( x − z 0 ) | | 2 ] , where z t = t x + ( 1 − t ) z 0 . \min_{\theta}{{\mathbb{E}}}_{t\sim{{\rm{P}}}(t),z_{0}\sim\pi_{0},x\sim\pi_{1}}\left[\left|\left|v_{\theta_{NN}}(z_{t},t)-(x-z_{0})\right|\right|^{2}\right],\penalty\ \penalty\ \text{where}\penalty\ \penalty\ z_{t}=tx+(1-t)z_{0}.
|
||||
(3)
|
||||
Here, P ( t ) {\rm{P}}(t) is a distribution over time t ∈ [ 0 , 1 ] t\in[0,1] . When the network has sufficient capacity and the objective is perfectly minimized, the optimal velocity field v θ N N ∗ v_{\theta^{*}_{NN}} maps the elementary distribution π 0 \pi_{0} to the true data distribution π 1 \pi_{1} . More precisely, the distribution of z 1 = ∫ 0 1 v θ N N ∗ ( z t , t ) 𝑑 t z_{1}=\int_{0}^{1}v_{\theta^{*}_{NN}}(z_{t},t){{\rm{d}}}t , with z 0 ∼ π 0 z_{0}\sim\pi_{0} , follows π 1 \pi_{1} . Despite its conceptual simplicity, rectified flow has shown superior performance in various generative modeling tasks, including text-to-image generation [ 23 ] , audio generation [ 40 ] and biological structure generation [ 38 ] .
|
||||
Figure 2:
|
||||
Architecture of the proposed JanusFlow. For visual understanding, the LLM performs autoregressive next-token prediction to generate responses. For image generation, the LLM employs images with rectified flow. Starting from Gaussian noise at t = 0 t=0 , the LLM iteratively updates z t z_{t} by predicting velocity vectors until reaching t = 1 t=1 . We omit the VAE encoder, the skip connection leveraged in generation and the linear layer after f e n c f_{enc} for simplicity.
|
||||
## 3.2 A Unified Framework for Multimodal Understanding and Generation
|
||||
JanusFlow presents a unified framework designed to address both vision understanding and image generation tasks. Next we outline how JanusFlow handles these two tasks within a single LLM architecture.
|
||||
## Multimodal Understanding.
|
||||
In multimodal understanding tasks, the LLM processes an input sequence consisting of interleaved text and image data. The text is tokenized into discrete tokens, each of which is transformed into an embedding of dimension D e m b D_{emb} . For the images, an image encoder f e n c f_{enc} encodes each image x i m x_{im} into a feature map of shape H i m × W i m × D e n c H_{im}\times W_{im}\times D_{enc} .
|
||||
This feature map is flattened and projected through a linear transformation layer into a sequence of embeddings with shape H i m W i m × D e m b H_{im}W_{im}\times D_{emb} .
|
||||
H i m H_{im} and W i m W_{im} are determined by the image encoder. The text and image embeddings are concatenated to form the input sequence to the LLM, which then autoregressively predicts the next tokens based on the input sequence of embeddings. According to common practice [ 97 , 88 , 100 ] , we add special token |BOI| before the image and |EOI| after the image to help the model locate the image embeddings in the sequence.
|
||||
## Image Generation.
|
||||
For image generation, our LLM takes a text sequence x c o n x^{con} as condition and generates a corresponding image using rectified flow.
|
||||
To improve computational efficiency, generation occurs in the latent space using a pre-trained SDXL-VAE [ 73 ] .
|
||||
The generation process begins by sampling Gaussian noise z 0 z_{0} of shape H l a t e n t × W l a t e n t × D l a t e n t H_{latent}\times W_{latent}\times D_{latent} in the latent space, which is then processed by a generation encoder g e n c g_{enc} into a sequence of embeddings H g e n W g e n × D e m b H_{gen}W_{gen}\times D_{emb} . This sequence is concatenated with a time embedding representing the current time step t t ( t = 0 t=0 at the beginning), resulting in a sequence of length H g e n W g e n + 1 H_{gen}W_{gen}+1 .
|
||||
Unlike previous approaches that employ various attention masking strategies [ 100 , 108 ] , we found that causal attention suffices, as our preliminary experiments showed no performance benefits from alternative masking schemes.
|
||||
The LLM’s output corresponding to z 0 z_{0} is transformed back into the latent space by a generation decoder g d e c g_{dec} , producing a velocity vector of shape H l a t e n t × W l a t e n t × D l a t e n t H_{latent}\times W_{latent}\times D_{latent} . The state is updated by a standard Euler solver,
|
||||
z t + d t = z t + v ( z t , t ) d t , z_{t+{\rm{d}}t}=z_{t}+v(z_{t},t){\rm{d}}t,
|
||||
(4)
|
||||
where d t {\rm{d}}t is a user-defined step size. We replace z 0 z_{0} with z d t z_{{\rm{d}}t} on the input and iterate the process until we get z 1 z_{1} , which is then decoded into the final image by the VAE decoder. To enhance generation quality, we employ classifier-free guidance (CFG) when computing the velocity:
|
||||
v ( z t , t ) = w v ( z t , t | x c o n ) + ( 1 − w ) v ( z t , t | ∅ ) , v(z_{t},t)=wv(z_{t},t\penalty\ |\penalty\ x^{con})+(1-w)v(z_{t},t\penalty\ |\penalty\ \varnothing),
|
||||
(5)
|
||||
where v ( z t , t | ∅ ) v(z_{t},t\penalty\ |\penalty\ \varnothing) denotes the velocity inferred without text conditioning and w ⩾ 1 w\geqslant 1 controls the magnitute of CFG. Empirically, increasing w w yields higher semantic alignment [ 77 , 62 , 73 , 23 ] . Analogous to multimodal understanding, we prepend the special token |BOI| to indicate the start of image generation in the sequence.
|
||||
## Decoupling Encoders for the Two Tasks.
|
||||
Previous approaches that unify autoregressive generation and diffusion models within a joint LLM training framework [ 108 , 100 ] employ identical encoders ( f e n c f_{enc} and g e n c g_{enc} ) for both understanding and generation tasks. For instance, Zhou et al. [108] performs both tasks in the same VAE latent space using a shared U-Net or linear encoder, while Xie et al. [100] leverages MAGVIT-v2 [ 102 ] to encode image patches into discrete tokens for both tasks.
|
||||
However, recent work on unified autoregressive models has shown this shared encoder design to be suboptimal [ 97 ] , particularly in models that generate images through autoregression on vector-quantized tokens. Drawing from these insights, JanusFlow adopts a decoupled encoder design. Specifically, we employ a pre-trained SigLIP-Large-Patch/16 [ 106 ] model as f e n c f_{enc} to extract semantic continuous features for multimodal understanding, while using separate ConvNeXt blocks [ 96 ] initialized from scratch as g e n c g_{enc} and g d e c g_{dec} for generation, chosen for its effectiveness. Following established practices [ 5 , 14 , 93 ] , we incorporate a long skip connection between g e n c g_{enc} and g d e c g_{dec} . Our controlled experiments in Sec. 4.5 demonstrate that this decoupled encoder design significantly improves the performance of our unified model. The complete architecture of JanusFlow is illustrated in Fig. 2 .
|
||||
## 3.3 Training Schemes
|
||||
As illustrated in Fig. 3 , we train our model in three sequential stages, detailed below.
|
||||
## Stage 1: Adaptation of Randomly Initialized Components.
|
||||
In the first stage, we focus on training only the randomly initialized components: the linear layers, generation encoder, and generation decoder. This stage serves to adapt these new modules to work effectively with the pre-trained LLM and SigLIP encoder, essentially functioning as an initialization phase for the newly introduced components.
|
||||
## Stage 2: Unified Pre-Training.
|
||||
Following the adaptation stage, we train the entire model except for the visual encoder, consistent with previous approaches [ 58 , 64 ] . The training incorporates three data types: multimodal understanding, image generation, and text-only data. We initially allocate a higher proportion of multimodal understanding data to establish the model’s understanding capabilities. Subsequently, we increase the ratio of image generation data to accommodate the convergence requirements of diffusion-based models [ 18 , 72 ] .
|
||||
## Stage 3: Supervised Fine-Tuning (SFT).
|
||||
In the final stage, we fine-tune the pre-trained model using instruction tuning data, which comprises dialogues, task-specific conversations, and high-quality text-conditioned image generation examples.
|
||||
During this stage, we also unfreeze the SigLIP encoder parameters [ 64 , 90 , 97 ] . This fine-tuning process enables the model to effectively respond to user instructions for both multimodal understanding and image generation tasks.
|
||||
Figure 3: Three training stages of JanusFlow.
|
||||
The trainable modules are marked with flame and the frozen modules are marked with snowflakes.
|
||||
## 3.4 Training Objective
|
||||
Training JanusFlow involves two types of data, multimodal understanding data and image generation data.
|
||||
Both types of data contain two parts: “condition” and “response”. “Condition” refers to the prompting of the tasks ( e . g ., text prompts in the task of generation and images in the task of understanding) while “response” refers to the corresponding responses of the two tasks. The data can be formatted as x = ( x c o n , x r e s ) x=(x^{con},x^{res}) , where the superscript c o n con denotes “condition” and r e s res denotes “response”.
|
||||
We denote the length of the whole sequence x x as ℓ \ell , the length of x c o n x^{con} as ℓ c o n \ell_{con} and the length of x r e s x^{res} as ℓ r e s \ell_{res} .
|
||||
We use θ \theta to represent the collection of all the trainable parameters in JanusFlow, including the LLM, f e n c f_{enc} , g e n c g_{enc} , g d e c g_{dec} and the linear transformation layers.
|
||||
## Autoregression Objective.
|
||||
For mutimodal understanding tasks, x r e s x^{res} contains only text tokens. JanusFlow is trained using the maximum likelihood principle,
|
||||
ℒ A R ( θ ) = − 𝔼 x ∼ 𝒟 u n d [ ∑ i = ℓ c o n ℓ − 1 log P θ ( x i + 1 | x 1 , … , x i ) ] , \mathcal{L}_{AR}(\theta)=-{\mathbb{E}}_{x\sim\mathcal{D}_{und}}\left[\sum_{i=\ell_{con}}^{\ell-1}\log{{\rm{P}}}_{\theta}\left(x_{i+1}|x_{1},\dots,x_{i}\right)\right],
|
||||
(6)
|
||||
where the expectation is taken over all ( x c o n , x r e s ) (x^{con},x^{res}) pairs in our multimodal understanding dataset 𝒟 u n d \mathcal{D}_{und} , computing loss only over tokens in x r e s x^{res} .
|
||||
## Rectified Flow Objective.
|
||||
For image generation tasks, x c o n x^{con} consists of text tokens and x r e s x^{res} is the corresponding image. JanusFlow is trained with the rectified flow objective,
|
||||
ℒ R F ( θ ) = 𝔼 x ∼ 𝒟 g e n , t ∼ P ( t ) , z 0 ∼ 𝒩 ( 0 , I ) [ | | v θ ( z t , t | x c o n ) − ( x r e s − z 0 ) | | 2 ] , \mathcal{L}_{RF}(\theta)={{\mathbb{E}}}_{x\sim\mathcal{D}_{gen},t\sim{{\rm{P}}}(t),z_{0}\sim\mathcal{N}(0,I)}\left[\left|\left|v_{\theta}(z_{t},t\penalty\ |\penalty\ x^{con})-(x^{res}-z_{0})\right|\right|^{2}\right],
|
||||
(7)
|
||||
where z t = t x r e s + ( 1 − t ) z 0 z_{t}=tx^{res}+(1-t)z_{0} .
|
||||
Following Stable Diffusion 3 [ 23 ] , we set the time distribution P ( t ) {\rm{P}}(t) to the logit-normal distribution. To enable CFG inference, we randomly drop 10% of the text prompts in training.
|
||||
## Representation Alignment Regularization.
|
||||
Recent work [ 103 ] has shown that aligning intermediate representations between diffusion transformers and semantic vision encoders enhances diffusion model generalization.
|
||||
Our decoupled vision encoder design enables efficient implementation of this alignment as a regularization term.
|
||||
Specifically, for generation tasks, we align features from the understanding encoder f e n c f_{enc} with the LLM’s intermediate features,
|
||||
ℒ R E P A ( θ , φ ) = − 𝔼 x ∼ 𝒟 g e n [ sim ( stop_grad ( f e n c ( x r e s ) ) , h φ ( q θ ( z t ) ) ) ] , \mathcal{L}_{REPA}(\theta,\varphi)=-{\mathbb{E}}_{x\sim\mathcal{D}_{gen}}\left[\text{sim}\left(\text{{stop\_grad}}(f_{enc}(x^{res})),h_{\varphi}(q_{\theta}(z_{t}))\right)\right],
|
||||
(8)
|
||||
where q θ ( z t ) q_{\theta}(z_{t}) denotes an intermediate LLM representation given input z t z_{t} , and h φ h_{\varphi} is a small trainable MLP that projects q θ ( z t ) q_{\theta}(z_{t}) to dimension D e n c D_{enc} . The function sim ( ⋅ , ⋅ ) \text{sim}(\cdot,\cdot) computes the mean of element-wise cosine similarity between embeddings. Before computing the loss, we reshape h φ ( q θ ( z t ) ) h_{\varphi}(q_{\theta}(z_{t})) to H g e n × W g e n × D e n c H_{gen}\times W_{gen}\times D_{enc} . To simplify the implementation, we intentionally adjust the configuration of g e n c g_{enc} and g d e c g_{dec} to ensure H g e n = H i m H_{gen}=H_{im} and W g e n = W i m W_{gen}=W_{im} .
|
||||
The gradient of ℒ R E P A \mathcal{L}_{REPA} is not back-propagated through the understanding encoder.
|
||||
This alignment loss helps the LLM’s internal feature space (given noisy input z t z_{t} ) align with the understanding encoder’s semantic feature space, thereby improving generation quality when producing images from new random noise and text conditions during inference.
|
||||
## Summary.
|
||||
All three objectives are applied across all training stages. Multimodal understanding tasks use ℒ A R \mathcal{L}_{AR} , while image generation tasks employ the combined loss ℒ R F + ℒ R E P A \mathcal{L}_{RF}+\mathcal{L}_{REPA} . Detailed experimental settings are provided in Sec. 4.1 .
|
||||
Table 1: Hyper-parameters of the proposed JanusFlow . Data ratio denotes the proportion of multimodal understanding data, image generation data and text-only data. In the initial 10,000 10,000 steps of Stage 2, we apply a data ratio of 30 : 50 : 20 30:50:20 to boost the understanding ability.
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Learning Rate
|
||||
1.0 × 10 − 4 1.0\times 10^{-4}
|
||||
1 × 10 − 4 1\times 10^{-4}
|
||||
2.0 × 10 − 5 2.0\times 10^{-5}
|
||||
LR Scheduler
|
||||
Constant
|
||||
Constant
|
||||
Constant
|
||||
Weight Decay
|
||||
0.0 0.0
|
||||
0.0 0.0
|
||||
0.0 0.0
|
||||
Gradient Clip
|
||||
1.0 1.0
|
||||
1.0 1.0
|
||||
1.0 1.0
|
||||
Optimizer
|
||||
AdamW ( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
Warm-up Steps
|
||||
2,000 2,000
|
||||
2,000 2,000
|
||||
1,000 1,000
|
||||
Training Steps
|
||||
10,000 10,000
|
||||
390,000 390,000
|
||||
26,000 26,000
|
||||
Batch Size
|
||||
512 512
|
||||
512 512
|
||||
256 256
|
||||
Data Ratio
|
||||
50 : 50 : 0 50:50:0
|
||||
14 : 80 : 6 14:80:6
|
||||
21 : 70 : 9 21:70:9
|
||||
## 4 Experiments
|
||||
We conduct extensive experiments to evaluate the capabilities of JanusFlow in both multimodal understanding and generation tasks. First, we describe our experimental setup and implementation details. Then, we present results on standard benchmarks for multimodal understanding and image generation. Finally, we perform ablation studies to validate our key design choices.
|
||||
## 4.1 Experiment Setup and Implementation Details
|
||||
Our framework builds upon an enhanced version 1 1
|
||||
1
|
||||
This version, trained on an expanded text corpus compared to the one in Janus [ 97 ] , has been demonstrated to possess better performance on multiple-choice benchmarks (e.g., MMBench [ 63 ] and SEED Bench [ 46 ] ). Our preliminary experiments suggest that it has minimal impact on the quality of visual generation. of DeepSeek-LLM (1.3B) [ 7 , 64 ] . The LLM consists of 24 transformer blocks and supports a sequence length of 4,096 4,096 .
|
||||
In our model, both understanding and generation exploits images of resolution 384.
|
||||
For multimodal understanding, we leverage SigLIP-Large-Patch/16 [ 106 ] as f e n c f_{enc} .
|
||||
For image generation, we utilize the pre-trained SDXL-VAE [ 73 ] for its latent space.
|
||||
The generation encoder g e n c g_{enc} comprises a 2 × 2 2\times 2 patchify layer followed by two ConvNeXt [ 96 ] blocks and a linear layer. The generation decoder g d e c g_{dec} combines two ConvNeXt blocks, a pixel-shuffle layer to upsample the feature map, and a linear layer. Our SigLIP encoder contains ∼ 300 \sim 300 M parameters. g e n c g_{enc} and g d e c g_{dec} are light-weight modules, containing ∼ 70 \sim 70 M parameters in total. Table 1 details the hyperparameters for each training stage. In the alignment regularization, we use the LLM features after the 6th block as q θ ( z t ) q_{\theta}(z_{t}) and a three-layer MLP as h φ h_{\varphi} . We employ an exponential moving average (EMA) with a ratio of 0.99 to ensure training stability.
|
||||
For data preprocessing, we deal with understanding and generation data differently. For understanding tasks, we maintain all image information by resizing the long side to the target size and padding the image to squares. For generation tasks, we resize the short side to the target size and apply random square cropping to avoid padding artifacts. During training, multiple sequences are packed to form a single sequence of length 4,096 4,096 for training efficiency. Our implementation is based on the HAI-LLM platform [ 31 ] using PyTorch [ 74 ] . Training was conducted on NVIDIA A100 GPUs, with each model requiring ∼ 1,600 \sim 1,600 A100 GPU days.
|
||||
## 4.2 Training Data Settings
|
||||
We follow Janus [ 97 ] to construct the training data. The data configuration for each training stage is listed below.
|
||||
## Data for Stage 1 and Stage 2.
|
||||
The first two stages of our framework uses three types of data: multimodal understanding data, image generation data and text-only data.
|
||||
1.
|
||||
Multimodal Understanding Data. This type of data contains several sub-categories: (a) Image caption data. We incorporate caption datasets from [ 20 , 41 , 50 , 51 , 53 , 82 ] and
|
||||
generate additional captions for images from [ 16 , 43 ] using open-source multimodal understanding models. The names of the datasets are provided in the supplementary materials.
|
||||
The data follows template formats, e . g ., “ <image>Generate the caption of this picture. <caption> ”.
|
||||
(b) Charts and tables. We directly adopt the chart and table data from the training data of DeepSeek-VL [ 64 ] .
|
||||
(c) Task data.
|
||||
ShareGPT4V [ 11 ] data is utilized to facilitate basic question-answering capabilities during pre-training, structured as “ <image><question><answer> ”.
|
||||
(d) Interleaved text-image data. This sub-category is sourced from [ 84 , 42 ] .
|
||||
2.
|
||||
Image Generation Data.
|
||||
Our image generation dataset combines high-quality images from [ 16 , 41 , 43 , 68 , 71 , 85 , 21 , 82 ] and 2 2 million in-house data.
|
||||
We enhance them with machine-generated captions using multimodal understanding models.
|
||||
We filter the images in [ 16 , 82 ] with aspect ratios and aesthetic scores, retaining approximately 20 % 20\% of the original datasets.
|
||||
25 % 25\% of the data contains single-sentence captions. These kind of data assist the model to be able to process short prompts. All the data points are formatted as “ <prompt><image> ”.
|
||||
3.
|
||||
Text-Only Data. We directly use the text corpus of DeepSeek-LLM [ 7 ] .
|
||||
Data for Stage 3. The SFT stage also uses three types of data:
|
||||
1.
|
||||
Multimodal Instruction Data. We leverage the instruction tuning datasets from [ 29 , 33 , 35 , 47 , 65 , 80 ] .
|
||||
2.
|
||||
Image Generation Data.
|
||||
We reformat the high-quality text-image pairs from [ 16 , 85 , 82 ] into an instruction format: “ User:<user prompt>\n\n Assistant:<image> ”.
|
||||
3.
|
||||
Text-Only Data. We directly incorporate the text-only data from [ 47 ] .
|
||||
Table 2: Performances on GenEval benchmark. “Gen.” denotes “generation” and “Unified” denotes unified understanding and generation models. Models using external pre-trained generative models are signed with † .
|
||||
Type
|
||||
Method
|
||||
Params
|
||||
Single Obj.
|
||||
Two Obj.
|
||||
Count.
|
||||
Colors
|
||||
Pos.
|
||||
Color Attri.
|
||||
Overall ↑ \uparrow
|
||||
Gen. Only
|
||||
LlamaGen [ 86 ]
|
||||
0.8 0.8 B
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.21 0.21
|
||||
0.58 0.58
|
||||
0.07 0.07
|
||||
0.04 0.04
|
||||
0.32 0.32
|
||||
LDM [ 77 ]
|
||||
1.4 1.4 B
|
||||
0.92 0.92
|
||||
0.29 0.29
|
||||
0.23 0.23
|
||||
0.70 0.70
|
||||
0.02 0.02
|
||||
0.05 0.05
|
||||
0.37 0.37
|
||||
SDv 1.5 1.5 [ 77 ]
|
||||
0.9 0.9 B
|
||||
0.97 0.97
|
||||
0.38 0.38
|
||||
0.35 0.35
|
||||
0.76 0.76
|
||||
0.04 0.04
|
||||
0.06 0.06
|
||||
0.43 0.43
|
||||
PixArt- α \alpha [ 9 ]
|
||||
0.6 0.6 B
|
||||
0.98 0.98
|
||||
0.50 0.50
|
||||
0.44 0.44
|
||||
0.80 0.80
|
||||
0.08 0.08
|
||||
0.07 0.07
|
||||
0.48 0.48
|
||||
SDv 2.1 2.1 [ 77 ]
|
||||
0.9 0.9 B
|
||||
0.98 0.98
|
||||
0.51 0.51
|
||||
0.44 0.44
|
||||
0.85 0.85
|
||||
0.07 0.07
|
||||
0.17 0.17
|
||||
0.50 0.50
|
||||
DALL-E 2 2 [ 76 ]
|
||||
6.5 6.5 B
|
||||
0.94 0.94
|
||||
0.66 0.66
|
||||
0.49 0.49
|
||||
0.77 0.77
|
||||
0.10 0.10
|
||||
0.19 0.19
|
||||
0.52 0.52
|
||||
Emu 3 3 -Gen [ 95 ]
|
||||
8 8 B
|
||||
0.98 0.98
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.81 0.81
|
||||
0.17 0.17
|
||||
0.21 0.21
|
||||
0.54 0.54
|
||||
SDXL [ 73 ]
|
||||
2.6 2.6 B
|
||||
0.98 0.98
|
||||
0.74 0.74
|
||||
0.39 0.39
|
||||
0.85 0.85
|
||||
0.15 0.15
|
||||
0.23 0.23
|
||||
0.55 0.55
|
||||
IF-XL [ 17 ]
|
||||
4.3 4.3 B
|
||||
0.97 0.97
|
||||
0.74 0.74
|
||||
0.66 0.66
|
||||
0.81 0.81
|
||||
0.13 0.13
|
||||
0.35 0.35
|
||||
0.61 0.61
|
||||
DALL-E 3 3 [ 6 ]
|
||||
-
|
||||
0.96 0.96
|
||||
0.87 0.87
|
||||
0.47 0.47
|
||||
0.83 0.83
|
||||
0.43 0.43
|
||||
0.45 0.45
|
||||
0.67 0.67
|
||||
Unified
|
||||
Chameleon [ 88 ]
|
||||
34 34 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.39 0.39
|
||||
LWM [ 59 ]
|
||||
7 7 B
|
||||
0.93 0.93
|
||||
0.41 0.41
|
||||
0.46 0.46
|
||||
0.79 0.79
|
||||
0.09 0.09
|
||||
0.15 0.15
|
||||
0.47 0.47
|
||||
SEED-X † [ 27 ]
|
||||
17 17 B
|
||||
0.97 0.97
|
||||
0.58 0.58
|
||||
0.26 0.26
|
||||
0.80 0.80
|
||||
0.19 0.19
|
||||
0.14 0.14
|
||||
0.49 0.49
|
||||
Show-o [ 100 ]
|
||||
1.3 1.3 B
|
||||
0.95 0.95
|
||||
0.52 0.52
|
||||
0.49 0.49
|
||||
0.82 0.82
|
||||
0.11 0.11
|
||||
0.28 0.28
|
||||
0.53 0.53
|
||||
Janus [ 97 ]
|
||||
1.3 1.3 B
|
||||
0.97 0.97
|
||||
0.68 0.68
|
||||
0.30 0.30
|
||||
0.84 0.84
|
||||
0.46 0.46
|
||||
0.42 0.42
|
||||
0.61 0.61
|
||||
Transfusion [ 108 ]
|
||||
7.3 7.3 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.63
|
||||
JanusFlow (Ours)
|
||||
1.3B
|
||||
0.97
|
||||
0.59
|
||||
0.45
|
||||
0.83
|
||||
0.53
|
||||
0.42
|
||||
0.63
|
||||
Table 3: Performances on DPG-Bench. The methods in this table are all generation-specific models except our method.
|
||||
Method
|
||||
Global
|
||||
Entity
|
||||
Attribute
|
||||
Relation
|
||||
Other
|
||||
Overall ↑ \uparrow
|
||||
SDv1.5 [ 77 ]
|
||||
74.63
|
||||
74.23
|
||||
75.39
|
||||
73.49
|
||||
67.81
|
||||
63.18
|
||||
PixArt- α \alpha [ 9 ]
|
||||
74.97
|
||||
79.32
|
||||
78.60
|
||||
82.57
|
||||
76.96
|
||||
71.11
|
||||
Lumina-Next [ 110 ]
|
||||
82.82
|
||||
88.65
|
||||
86.44
|
||||
80.53
|
||||
81.82
|
||||
74.63
|
||||
SDXL [ 73 ]
|
||||
83.27
|
||||
82.43
|
||||
80.91
|
||||
86.76
|
||||
80.41
|
||||
74.65
|
||||
Playground v2.5 [ 48 ]
|
||||
83.06
|
||||
82.59
|
||||
81.20
|
||||
84.08
|
||||
83.50
|
||||
75.47
|
||||
Hunyuan-DiT [ 54 ]
|
||||
84.59
|
||||
80.59
|
||||
88.01
|
||||
74.36
|
||||
86.41
|
||||
78.87
|
||||
PixArt- Σ \Sigma [ 10 ]
|
||||
86.89
|
||||
82.89
|
||||
88.94
|
||||
86.59
|
||||
87.68
|
||||
80.54
|
||||
Emu3-Gen [ 95 ]
|
||||
85.21
|
||||
86.68
|
||||
86.84
|
||||
90.22
|
||||
83.15
|
||||
80.60
|
||||
JanusFlow (Ours)
|
||||
87.03
|
||||
87.31
|
||||
87.39
|
||||
89.79
|
||||
88.10
|
||||
80.09
|
||||
## 4.3 Evaluation Settings
|
||||
## Image Generation.
|
||||
We evaluate the generated images using both visual quality and semantic accuracy metrics. For visual quality assessment, we employ the Fréchet Inception Distance [ 30 ] (FID) metric and compute FID between 30,000 generated images and their corresponding reference images from the MJHQ dataset [ 48 ] . The FID computation follows the implementation from GigaGAN [ 39 ] . To evaluate semantic accuracy, we utilize two specialized frameworks: GenEval [ 28 ] and DPG-Bench [ 34 ] . These frameworks are designed to assess whether the generated images accurately contain the objects and relationships specified in the input prompts, providing a broad evaluation of the generation capabilities.
|
||||
## Multimodal Understanding.
|
||||
We evaluate JanusFlow’s multimodal understanding abilities across a diverse set of vision-language benchmarks for general understanding capabilities, including POPE [ 52 ] , MME [ 24 ] , MMBench [ 63 ] , SEEDBench [ 46 ] , VQAv2 [ 29 ] , GQA [ 35 ] , MM-Vet [ 104 ] , MMMU [ 105 ] , ChartQA [ 70 ] and TextVQA [ 81 ]
|
||||
## 4.4 Quantitative Results
|
||||
Table 4: Results of MJHQ FID-30k.
|
||||
The models which have similar scales to our model are marked with blue background.
|
||||
JanusFlow achieves the best FID among 1.3B models.
|
||||
Method
|
||||
Params
|
||||
FID ↓ \downarrow
|
||||
LWM [ 59 ]
|
||||
7B
|
||||
17.77
|
||||
VILA-U 256 [ 99 ]
|
||||
7B
|
||||
12.81
|
||||
VILA-U 384 [ 99 ]
|
||||
7B
|
||||
7.69
|
||||
Show-o [ 100 ]
|
||||
1.3B
|
||||
15.18
|
||||
Janus [ 97 ]
|
||||
1.3B
|
||||
10.10
|
||||
JanusFlow (Ours)
|
||||
1.3B
|
||||
9.51
|
||||
Image Generation Performances. We report the performances on GenEval, DPG-Bench and MJHQ FID-30k. In Tab. 2 , we give comparisons on GenEval including the scores of all the sub-tasks and the overall score. JanusFlow achieves an overall score of 0.63, surpassing the previous unified framework and several generation specific models including SDXL [ 73 ] and DALL-E 2 [ 76 ] . In Tab. 3 , We show results on DPG-Bench and the corresponding comparisons. It is noted that all the methods in Tab. 3 are generation-specific models except our model. The results on GenEval and DPG-Bench demonstrate the ability of instruction following of our model.
|
||||
We give the comparisons on MJHQ FID-30k in Tab. 4 . The images which are sampled to calculate FID are generated with a CFG factor w = 2 w=2 and a number of sampling steps 30 30 . We sweep the CFG factor and the sampling steps and provide the results in the appendix.
|
||||
Our method achieves the best performance among all the models with 1.3B LLM. The results prove that the rectified flow is able to improve the quality of generated images over autoregressive models such as Janus [ 97 ] .
|
||||
Multimodal Understanding Performances. We show comparisons of our method and other methods including understanding-specific models and unified understanding and generation models in Tab. 5 . Our model reaches the best performances among all the models with similar number of parameters and even surpasses multiple understanding-specific methods with larger scales. Our results demonstrate that our method harmonizes autoregressive LLM and rectified flow, achieving satisfying performance in both understanding and generation.
|
||||
Table 5: Comparison with other methods on multimodal understanding benchmarks . “Und.” denotes “understanding” and “Unified” denotes unified understanding and generation models. The models employing external pre-trained generative models are marked with † . The models with LLMs which have similar number of parameters to us are marked with blue background under the line of dashes.
|
||||
Type
|
||||
Model
|
||||
LLM Param
|
||||
POPE
|
||||
MME-P
|
||||
MMB dev
|
||||
SEED
|
||||
VQAv2 test
|
||||
GQA
|
||||
MMMU
|
||||
MM-Vet
|
||||
ChartQA
|
||||
TextVQA
|
||||
Und. Only
|
||||
MobileVLM [ 12 ]
|
||||
2.7 2.7 B
|
||||
84.9 84.9
|
||||
1288.9 1288.9
|
||||
59.6 59.6
|
||||
-
|
||||
-
|
||||
59.0 59.0
|
||||
-
|
||||
-
|
||||
-
|
||||
47.5 47.5
|
||||
MobileVLM-V2 [ 13 ]
|
||||
2.7 2.7 B
|
||||
84.7 84.7
|
||||
1440.5 1440.5
|
||||
63.2 63.2
|
||||
-
|
||||
-
|
||||
61.1 61.1
|
||||
-
|
||||
-
|
||||
-
|
||||
57.5 57.5
|
||||
LLaVA-Phi [ 109 ]
|
||||
2.7 2.7 B
|
||||
85.0 85.0
|
||||
1335.1 1335.1
|
||||
59.8 59.8
|
||||
-
|
||||
71.4 71.4
|
||||
-
|
||||
-
|
||||
28.9 28.9
|
||||
-
|
||||
48.6 48.6
|
||||
LLaVA [ 58 ]
|
||||
7 7 B
|
||||
76.3 76.3
|
||||
809.6 809.6
|
||||
38.7 38.7
|
||||
33.5 33.5
|
||||
-
|
||||
-
|
||||
-
|
||||
25.5 25.5
|
||||
-
|
||||
-
|
||||
LLaVA-v1.5 [ 56 ]
|
||||
7 7 B
|
||||
85.9 85.9
|
||||
1510.7 1510.7
|
||||
64.3 64.3
|
||||
58.6 58.6
|
||||
78.5 78.5
|
||||
62.0 62.0
|
||||
35.4 35.4
|
||||
31.1 31.1
|
||||
-
|
||||
58.2 58.2
|
||||
InstructBLIP [ 15 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
36.0 36.0
|
||||
53.4 53.4
|
||||
-
|
||||
49.2 49.2
|
||||
-
|
||||
26.2 26.2
|
||||
-
|
||||
50.1 50.1
|
||||
Qwen-VL-Chat [ 4 ]
|
||||
7 7 B
|
||||
-
|
||||
1487.5 1487.5
|
||||
60.6 60.6
|
||||
58.2 58.2
|
||||
78.2 78.2
|
||||
57.5 57.5
|
||||
-
|
||||
-
|
||||
66.3 66.3
|
||||
61.5 61.5
|
||||
LLaVA-NeXT [ 57 ]
|
||||
7 7 B
|
||||
-
|
||||
1519.3 1519.3
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
35.1 35.1
|
||||
-
|
||||
54.8 54.8
|
||||
-
|
||||
Qwen2-VL [ 94 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
54.1 54.1
|
||||
62.0 62.0
|
||||
83.0 83.0
|
||||
84.3 84.3
|
||||
IDEFICS- 9 9 B [ 44 ]
|
||||
8 8 B
|
||||
-
|
||||
-
|
||||
48.2 48.2
|
||||
-
|
||||
50.9 50.9
|
||||
38.4 38.4
|
||||
-
|
||||
-
|
||||
-
|
||||
25.9 25.9
|
||||
Emu 3 3 -Chat [ 95 ]
|
||||
8 8 B
|
||||
85.2 85.2
|
||||
-
|
||||
58.5 58.5
|
||||
68.2 68.2
|
||||
75.1 75.1
|
||||
60.3 60.3
|
||||
31.6 31.6
|
||||
-
|
||||
68.6 68.6
|
||||
64.7 64.7
|
||||
InstructBLIP [ 15 ]
|
||||
13 13 B
|
||||
78.9 78.9
|
||||
1212.8 1212.8
|
||||
-
|
||||
-
|
||||
-
|
||||
49.5 49.5
|
||||
-
|
||||
25.6 25.6
|
||||
-
|
||||
50.7 50.7
|
||||
LLaVA-v1.5-Phi-1.5 [ 100 ]
|
||||
1.3 1.3 B
|
||||
84.1 84.1
|
||||
1128.0 1128.0
|
||||
-
|
||||
-
|
||||
75.3 75.3
|
||||
56.5 56.5
|
||||
30.7 30.7
|
||||
-
|
||||
-
|
||||
-
|
||||
MobileVLM [ 12 ]
|
||||
1.4 1.4 B
|
||||
84.5 84.5
|
||||
1196.2 1196.2
|
||||
53.2 53.2
|
||||
-
|
||||
-
|
||||
56.1 56.1
|
||||
-
|
||||
-
|
||||
-
|
||||
41.5 41.5
|
||||
MobileVLM-V2 [ 13 ]
|
||||
1.4 1.4 B
|
||||
84.3 84.3
|
||||
1302.8 1302.8
|
||||
57.7 57.7
|
||||
-
|
||||
-
|
||||
59.3 59.3
|
||||
-
|
||||
-
|
||||
-
|
||||
52.1 52.1
|
||||
Unified
|
||||
Gemini-Nano-1 [ 89 ]
|
||||
1.8 1.8 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
62.7 62.7
|
||||
-
|
||||
26.3 26.3
|
||||
-
|
||||
53.6 53.6
|
||||
62.5 62.5
|
||||
LWM [ 59 ]
|
||||
7 7 B
|
||||
75.2 75.2
|
||||
-
|
||||
-
|
||||
-
|
||||
55.8 55.8
|
||||
44.8 44.8
|
||||
-
|
||||
9.6 9.6
|
||||
-
|
||||
-
|
||||
VILA-U [ 99 ]
|
||||
7 7 B
|
||||
85.8 85.8
|
||||
1401.8 1401.8
|
||||
-
|
||||
59.0 59.0
|
||||
79.4 79.4
|
||||
60.8 60.8
|
||||
-
|
||||
33.5 33.5
|
||||
-
|
||||
60.8 60.8
|
||||
Chameleon [ 88 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
22.4 22.4
|
||||
8.3 8.3
|
||||
-
|
||||
-
|
||||
DreamLLM † [ 19 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
72.9 72.9
|
||||
-
|
||||
-
|
||||
36.6 36.6
|
||||
-
|
||||
41.8 41.8
|
||||
LaVIT † [ 37 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
66.0 66.0
|
||||
46.8 46.8
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
Emu † [ 87 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
52.0 52.0
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
NExT-GPT † [ 98 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
66.7 66.7
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
Show-o [ 100 ]
|
||||
1.3 1.3 B
|
||||
73.8 73.8
|
||||
948.4 948.4
|
||||
-
|
||||
-
|
||||
59.3 59.3
|
||||
48.7 48.7
|
||||
25.1 25.1
|
||||
-
|
||||
-
|
||||
-
|
||||
Janus [ 97 ]
|
||||
1.3 1.3 B
|
||||
87.0 87.0
|
||||
1338.0 1338.0
|
||||
69.4 69.4
|
||||
63.7 63.7
|
||||
77.3 77.3
|
||||
59.1 59.1
|
||||
30.5 30.5
|
||||
34.3 34.3
|
||||
-
|
||||
-
|
||||
JanusFlow (Ours)
|
||||
1.3 1.3 B
|
||||
88.0 88.0
|
||||
1333.1 1333.1
|
||||
74.9 74.9
|
||||
70.5 70.5
|
||||
79.8 79.8
|
||||
60.3 60.3
|
||||
29.3 29.3
|
||||
30.9 30.9
|
||||
64.6 64.6
|
||||
55.5 55.5
|
||||
## 4.5 Ablation Studies
|
||||
Table 6: Ablation studies. The weights of the modules with † are frozen during training. “Exp.” denotes “experiment”. “FID” in this table is MJHQ FID-10k with CFG factor w = 7.5 w=7.5 and 30 steps. “CLIP” denotes CLIP similarity with the backbone of CLIP-ViT-Large-Patch/14. Exp. F is the final configuration for training JanusFlow.
|
||||
Exp. ID
|
||||
Model Setting
|
||||
Train. Iter.
|
||||
Evaluation Benchmarks
|
||||
REPA
|
||||
Und. Modules
|
||||
Gen. Modules
|
||||
Type
|
||||
POPE ↑ \uparrow
|
||||
VQAv2 v a l ↑ {}_{val}\uparrow
|
||||
GQA ↑ \uparrow
|
||||
FID ↓ \downarrow
|
||||
CLIP ↑ \uparrow
|
||||
A
|
||||
× \times
|
||||
SigLIP
|
||||
VAE † +ConvNeXt
|
||||
Unified
|
||||
50,000
|
||||
82.40
|
||||
69.62
|
||||
54.43
|
||||
19.84
|
||||
24.94
|
||||
B
|
||||
✓ \checkmark
|
||||
Shared VAE † +ConvNeXt
|
||||
Unified
|
||||
50,000
|
||||
78.13
|
||||
53.94
|
||||
44.04
|
||||
18.05
|
||||
26.38
|
||||
C
|
||||
✓ \checkmark
|
||||
VAE+ConvNeXt
|
||||
VAE † +ConvNeXt
|
||||
Unified
|
||||
50,000
|
||||
75.30
|
||||
55.41
|
||||
44.44
|
||||
17.53
|
||||
26.32
|
||||
D
|
||||
✓ \checkmark
|
||||
SigLIP
|
||||
-
|
||||
Und. Only
|
||||
13,000
|
||||
85.03
|
||||
69.10
|
||||
54.23
|
||||
-
|
||||
-
|
||||
E
|
||||
✓ \checkmark
|
||||
-
|
||||
VAE † +ConvNeXt
|
||||
Gen. Only
|
||||
37,000
|
||||
-
|
||||
-
|
||||
-
|
||||
16.69
|
||||
26.89
|
||||
F
|
||||
✓ \checkmark
|
||||
SigLIP
|
||||
VAE † +ConvNeXt
|
||||
Unified
|
||||
50,000
|
||||
84.73
|
||||
69.20
|
||||
54.83
|
||||
17.61
|
||||
26.40
|
||||
We conduct comprehensive ablation studies to validate the effectiveness of our key design choices. For computational efficiency, all ablation experiments are performed on 256 × 256 256\times 256 resolution images 2 2
|
||||
2
|
||||
The understanding encoders in the 256 × 256 256\times 256 -based ablation studies is also SigLIP-Large-Patch/16 which is pre-trained on 256 × 256 256\times 256 images. . All models are trained on our unified pre-training dataset for 50,000 50,000 iterations, except for the understanding-only and generation-only variants, which are trained for proportionally fewer iterations based on their respective data ratios in the pre-training phase. The quantitative results of these ablation studies are presented in Tab. 6 .
|
||||
Impact of Representation Alignment.
|
||||
The comparison between Exp. A and F demonstrates the significant benefits of incorporating representation alignment regularization [ 103 ] during training. Specifically, models trained with representation alignment show notably lower FID scores on MJHQ dataset and higher CLIP scores, indicating simultaneous improvements in both image quality and semantic alignment. Importantly, our architecture differs from previous studies [ 72 , 66 ] examined in [ 103 ] due to our incorporation of LLM and an additional skip connection between g e n c g_{enc} and g d e c g_{dec} . The effectiveness of representation alignment in our modified architecture suggests its broad applicability and generalization capability across different network structures.
|
||||
Figure 4: Image generation results of JanusFlow.
|
||||
Our model can generate high-quality images that are semantically consistent with text prompts.
|
||||
Impact of Decoupling Visual Encoders. e efficacy of using powerful pre-trained visual encoders in multimodal understanding.
|
||||
The comparison among Exp. B, C, and F demonstrates the advantages of using separate visual encoders for understanding and generation tasks. In Exp. B, following a design similar to Transfusion [ 108 ] , we implement shared ConvNeXt blocks in the SDXL-VAE latent space for both understanding and generation encoders. Exp. C employs separate encoders with identical architectures and initialization parameters, but trained independently. The performance differences between these configurations validate the necessity of decoupled visual encoders in improving our unified model’s capabilities. Moreover, the superior results in Exp. C and F highlight the benefits of leveraging pre-trained semantic visual encoders for multimodal understanding tasks.
|
||||
Fair Comparison with Understanding / Generation-Only Models.
|
||||
To establish meaningful benchmarks, we evaluate task-specific models trained under identical conditions - using the same pre-training dataset, infrastructure, and hyperparameters. Exp. D and E represent these specialized models, trained with data volumes matching the unified models in Tab. 6 . The minimal performance gap between Exp. F and these task-specific baselines demonstrates that our unified framework successfully integrates understanding and generation capabilities without significant compromise in either task’s performance.
|
||||
Figure 5: Visual Understanding with JanusFlow. Our model effectively handles various visual understanding tasks, such as question answering, plot interpretation and object counting.
|
||||
## 4.6 Qualitative Results
|
||||
We present qualitative evaluations of our method for both image generation and understanding tasks. Fig. 1(b) and Fig. 4 showcases the image generation capabilities of JanusFlow.
|
||||
These results demonstrate both the high visual quality of our generated images and our framework’s ability to faithfully follow diverse instructions.
|
||||
For multimodal understanding, Fig. 5 presents example conversations that show our model’s understanding capabilities across various scenarios. These interactions demonstrate the model’s ability to understand and reason about visual content in natural language dialogues.
|
||||
Additional qualitative examples showcasing the versatility and effectiveness of JanusFlow are provided in the appendix.
|
||||
## 5 Conclusion
|
||||
We present JanusFlow, a unified framework that successfully harmonizes autoregressive and rectified flow models for multimodal understanding and generation tasks. Our extensive experiments demonstrate that this unification achieves comparable performance to task-specific models. The successful integration of these fundamentally different model architectures not only addresses current challenges in multimodal learning but also opens new possibilities for future research in training unified models.
|
||||
|
||||
=== BODY (tail) ===
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,935 @@
|
||||
TITLE: DeepSeek-V3 Technical Report
|
||||
ARXIV: 2412.19437
|
||||
AUTHORS: DeepSeek-AI , Aixin Liu , Bei Feng , Bing Xue , Bingxuan Wang , Bochao Wu , Chengda Lu , Chenggang Zhao , Chengqi Deng , Chenyu Zhang , Chong Ruan , Damai Dai , Daya Guo , Dejian Yang , Deli Chen , Dongjie Ji , Erhang Li , Fangyun Lin , Fucong Dai , Fuli Luo , Guangbo Hao , Guanting Chen , Guowei Li , H. Zhang , Han Bao , Hanwei Xu , Haocheng Wang , Haowei Zhang , Honghui Ding , Huajian Xin , Huazuo Gao , Hui Li , Hui Qu , J.L. Cai , Jian Liang , Jianzhong Guo , Jiaqi Ni , Jiashi Li , Jiawei Wang , Jin Chen , Jingchang Chen , Jingyang Yuan , Junjie Qiu , Junlong Li , Junxiao Song , Kai Dong , Kai Hu , Kaige Gao , Kang Guan , Kexin Huang , Kuai Yu , Lean Wang , Lecong Zhang , Lei Xu , Leyi Xia , Liang Zhao , Litong Wang , Liyue Zhang , Meng Li , Miaojun Wang , Mingchuan Zhang , Minghua Zhang , Minghui Tang , Mingming Li , Ning Tian , Panpan Huang , Peiyi Wang , Peng Zhang , Qiancheng Wang , Qihao Zhu , Qinyu Chen , Qiushi Du , R.J. Chen , R.L. Jin , Ruiqi Ge , Ruisong Zhang , Ruizhe Pan , Runji Wang , Runxin Xu , Ruoyu Zhang , Ruyi Chen , S.S. Li , Shanghao Lu , Shangyan Zhou , Shanhuang Chen , Shaoqing Wu , Shengfeng Ye , Shengfeng Ye , Shirong Ma , Shiyu Wang , Shuang Zhou , Shuiping Yu , Shunfeng Zhou , Shuting Pan , T. Wang , Tao Yun , Tian Pei , Tianyu Sun , W.L. Xiao , Wangding Zeng , Wanjia Zhao , Wei An , Wen Liu , Wenfeng Liang , Wenjun Gao , Wenqin Yu , Wentao Zhang , X.Q. Li , Xiangyue Jin , Xianzu Wang , Xiao Bi , Xiaodong Liu , Xiaohan Wang , Xiaojin Shen , Xiaoka
|
||||
DATE: [Submitted on 27 Dec 2024 ( v1 ), last revised 18 Feb 2025 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-V3 Technical Report
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Architecture
|
||||
## 2.1 Basic Architecture
|
||||
## 2.1.1 Multi-Head Latent Attention
|
||||
## 2.1.2 DeepSeekMoE with Auxiliary-Loss-Free Load Balancing
|
||||
## Basic Architecture of DeepSeekMoE.
|
||||
## Auxiliary-Loss-Free Load Balancing.
|
||||
## Complementary Sequence-Wise Auxiliary Loss.
|
||||
## Node-Limited Routing.
|
||||
## No Token-Dropping.
|
||||
## 2.2 Multi-Token Prediction
|
||||
## MTP Modules.
|
||||
## MTP Training Objective.
|
||||
## MTP in Inference.
|
||||
## 3 Infrastructures
|
||||
## 3.1 Compute Clusters
|
||||
## 3.2 Training Framework
|
||||
## 3.2.1 DualPipe and Computation-Communication Overlap
|
||||
## 3.2.2 Efficient Implementation of Cross-Node All-to-All Communication
|
||||
## 3.2.3 Extremely Memory Saving with Minimal Overhead
|
||||
## Recomputation of RMSNorm and MLA Up-Projection.
|
||||
## Exponential Moving Average in CPU.
|
||||
## Shared Embedding and Output Head for Multi-Token Prediction.
|
||||
## 3.3 FP8 Training
|
||||
## 3.3.1 Mixed Precision Framework
|
||||
## 3.3.2 Improved Precision from Quantization and Multiplication
|
||||
## Fine-Grained Quantization.
|
||||
## Increasing Accumulation Precision.
|
||||
## Mantissa over Exponents.
|
||||
## Online Quantization.
|
||||
## 3.3.3 Low-Precision Storage and Communication
|
||||
## Low-Precision Optimizer States.
|
||||
## Low-Precision Activation.
|
||||
## Low-Precision Communication.
|
||||
## 3.4 Inference and Deployment
|
||||
## 3.4.1 Prefilling
|
||||
## 3.4.2 Decoding
|
||||
## 3.5 Suggestions on Hardware Design
|
||||
## 3.5.1 Communication Hardware
|
||||
## 3.5.2 Compute Hardware
|
||||
## Higher FP8 GEMM Accumulation Precision in Tensor Cores.
|
||||
## Support for Tile- and Block-Wise Quantization.
|
||||
## Support for Online Quantization.
|
||||
## Support for Transposed GEMM Operations.
|
||||
## 4 Pre-Training
|
||||
## 4.1 Data Construction
|
||||
## 4.2 Hyper-Parameters
|
||||
## Model Hyper-Parameters.
|
||||
## Training Hyper-Parameters.
|
||||
## 4.3 Long Context Extension
|
||||
## 4.4 Evaluations
|
||||
## 4.4.1 Evaluation Benchmarks
|
||||
## 4.4.2 Evaluation Results
|
||||
## 4.5 Discussion
|
||||
## 4.5.1 Ablation Studies for Multi-Token Prediction
|
||||
## 4.5.2 Ablation Studies for the Auxiliary-Loss-Free Balancing Strategy
|
||||
## 4.5.3 Batch-Wise Load Balance VS. Sequence-Wise Load Balance
|
||||
## 5 Post-Training
|
||||
## 5.1 Supervised Fine-Tuning
|
||||
## Reasoning Data.
|
||||
## Non-Reasoning Data.
|
||||
## SFT Settings.
|
||||
## 5.2 Reinforcement Learning
|
||||
## 5.2.1 Reward Model
|
||||
## Rule-Based RM.
|
||||
## Model-Based RM.
|
||||
## 5.2.2 Group Relative Policy Optimization
|
||||
## 5.3 Evaluations
|
||||
## 5.3.1 Evaluation Settings
|
||||
## Evaluation Benchmarks.
|
||||
## Compared Baselines.
|
||||
## Detailed Evaluation Configurations.
|
||||
## 5.3.2 Standard Evaluation
|
||||
## English Benchmarks.
|
||||
## Code and Math Benchmarks.
|
||||
## Chinese Benchmarks.
|
||||
## 5.3.3 Open-Ended Evaluation
|
||||
## 5.3.4 DeepSeek-V3 as a Generative Reward Model
|
||||
## 5.4 Discussion
|
||||
## 5.4.1 Distillation from DeepSeek-R1
|
||||
## 5.4.2 Self-Rewarding
|
||||
## 5.4.3 Multi-Token Prediction Evaluation
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-V3 Technical Report
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Basic Architecture
|
||||
2.1.1 Multi-Head Latent Attention
|
||||
2.1.2 DeepSeekMoE with Auxiliary-Loss-Free Load Balancing
|
||||
Basic Architecture of DeepSeekMoE.
|
||||
Auxiliary-Loss-Free Load Balancing.
|
||||
Complementary Sequence-Wise Auxiliary Loss.
|
||||
Node-Limited Routing.
|
||||
No Token-Dropping.
|
||||
2.2 Multi-Token Prediction
|
||||
MTP Modules.
|
||||
MTP Training Objective.
|
||||
MTP in Inference.
|
||||
3 Infrastructures
|
||||
3.1 Compute Clusters
|
||||
3.2 Training Framework
|
||||
3.2.1 DualPipe and Computation-Communication Overlap
|
||||
3.2.2 Efficient Implementation of Cross-Node All-to-All Communication
|
||||
3.2.3 Extremely Memory Saving with Minimal Overhead
|
||||
Recomputation of RMSNorm and MLA Up-Projection.
|
||||
Exponential Moving Average in CPU.
|
||||
Shared Embedding and Output Head for Multi-Token Prediction.
|
||||
3.3 FP8 Training
|
||||
3.3.1 Mixed Precision Framework
|
||||
3.3.2 Improved Precision from Quantization and Multiplication
|
||||
Fine-Grained Quantization.
|
||||
Increasing Accumulation Precision.
|
||||
Mantissa over Exponents.
|
||||
Online Quantization.
|
||||
3.3.3 Low-Precision Storage and Communication
|
||||
Low-Precision Optimizer States.
|
||||
Low-Precision Activation.
|
||||
Low-Precision Communication.
|
||||
3.4 Inference and Deployment
|
||||
3.4.1 Prefilling
|
||||
3.4.2 Decoding
|
||||
3.5 Suggestions on Hardware Design
|
||||
3.5.1 Communication Hardware
|
||||
3.5.2 Compute Hardware
|
||||
Higher FP8 GEMM Accumulation Precision in Tensor Cores.
|
||||
Support for Tile- and Block-Wise Quantization.
|
||||
Support for Online Quantization.
|
||||
Support for Transposed GEMM Operations.
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
4.2 Hyper-Parameters
|
||||
Model Hyper-Parameters.
|
||||
Training Hyper-Parameters.
|
||||
4.3 Long Context Extension
|
||||
4.4 Evaluations
|
||||
4.4.1 Evaluation Benchmarks
|
||||
4.4.2 Evaluation Results
|
||||
4.5 Discussion
|
||||
4.5.1 Ablation Studies for Multi-Token Prediction
|
||||
4.5.2 Ablation Studies for the Auxiliary-Loss-Free Balancing Strategy
|
||||
4.5.3 Batch-Wise Load Balance VS. Sequence-Wise Load Balance
|
||||
5 Post-Training
|
||||
5.1 Supervised Fine-Tuning
|
||||
Reasoning Data.
|
||||
Non-Reasoning Data.
|
||||
SFT Settings.
|
||||
5.2 Reinforcement Learning
|
||||
5.2.1 Reward Model
|
||||
Rule-Based RM.
|
||||
Model-Based RM.
|
||||
5.2.2 Group Relative Policy Optimization
|
||||
5.3 Evaluations
|
||||
5.3.1 Evaluation Settings
|
||||
Evaluation Benchmarks.
|
||||
Compared Baselines.
|
||||
Detailed Evaluation Configurations.
|
||||
5.3.2 Standard Evaluation
|
||||
English Benchmarks.
|
||||
Code and Math Benchmarks.
|
||||
Chinese Benchmarks.
|
||||
5.3.3 Open-Ended Evaluation
|
||||
5.3.4 DeepSeek-V3 as a Generative Reward Model
|
||||
5.4 Discussion
|
||||
5.4.1 Distillation from DeepSeek-R1
|
||||
5.4.2 Self-Rewarding
|
||||
5.4.3 Multi-Token Prediction Evaluation
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Contributions and Acknowledgments
|
||||
B Ablation Studies for Low-Precision Training
|
||||
B.1 FP8 v.s. BF16 Training
|
||||
B.2 Discussion About Block-Wise Quantization
|
||||
C Expert Specialization Patterns of the 16B Aux-Loss-Based and Aux-Loss-Free Models
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2412.19437v2 [cs.CL] 18 Feb 2025
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-V3 Technical Report
|
||||
DeepSeek-AI
|
||||
research@deepseek.com
|
||||
## Abstract
|
||||
We present DeepSeek-V3, a strong Mixture-of-Experts (MoE) language model with 671B total parameters with 37B activated for each token.
|
||||
To achieve efficient inference and cost-effective training, DeepSeek-V3 adopts Multi-head Latent Attention (MLA) and DeepSeekMoE architectures, which were thoroughly validated in DeepSeek-V2.
|
||||
Furthermore, DeepSeek-V3 pioneers an auxiliary-loss-free strategy for load balancing and sets a multi-token prediction training objective for stronger performance.
|
||||
We pre-train DeepSeek-V3 on 14.8 trillion diverse and high-quality tokens, followed by Supervised Fine-Tuning and Reinforcement Learning stages to fully harness its capabilities.
|
||||
Comprehensive evaluations reveal that DeepSeek-V3 outperforms other open-source models and achieves performance comparable to leading closed-source models.
|
||||
Despite its excellent performance, DeepSeek-V3 requires only 2.788M H800 GPU hours for its full training.
|
||||
In addition, its training process is remarkably stable.
|
||||
Throughout the entire training process, we did not experience any irrecoverable loss spikes or perform any rollbacks.
|
||||
The model checkpoints are available at https://github.com/deepseek-ai/DeepSeek-V3 .
|
||||
Figure 1:
|
||||
Benchmark performance of DeepSeek-V3 and its counterparts.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Basic Architecture
|
||||
2.1.1 Multi-Head Latent Attention
|
||||
2.1.2 DeepSeekMoE with Auxiliary-Loss-Free Load Balancing
|
||||
2.2 Multi-Token Prediction
|
||||
3 Infrastructures
|
||||
3.1 Compute Clusters
|
||||
3.2 Training Framework
|
||||
3.2.1 DualPipe and Computation-Communication Overlap
|
||||
3.2.2 Efficient Implementation of Cross-Node All-to-All Communication
|
||||
3.2.3 Extremely Memory Saving with Minimal Overhead
|
||||
3.3 FP8 Training
|
||||
3.3.1 Mixed Precision Framework
|
||||
3.3.2 Improved Precision from Quantization and Multiplication
|
||||
3.3.3 Low-Precision Storage and Communication
|
||||
3.4 Inference and Deployment
|
||||
3.4.1 Prefilling
|
||||
3.4.2 Decoding
|
||||
3.5 Suggestions on Hardware Design
|
||||
3.5.1 Communication Hardware
|
||||
3.5.2 Compute Hardware
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
4.2 Hyper-Parameters
|
||||
4.3 Long Context Extension
|
||||
4.4 Evaluations
|
||||
4.4.1 Evaluation Benchmarks
|
||||
4.4.2 Evaluation Results
|
||||
4.5 Discussion
|
||||
4.5.1 Ablation Studies for Multi-Token Prediction
|
||||
4.5.2 Ablation Studies for the Auxiliary-Loss-Free Balancing Strategy
|
||||
4.5.3 Batch-Wise Load Balance VS. Sequence-Wise Load Balance
|
||||
5 Post-Training
|
||||
5.1 Supervised Fine-Tuning
|
||||
5.2 Reinforcement Learning
|
||||
5.2.1 Reward Model
|
||||
5.2.2 Group Relative Policy Optimization
|
||||
5.3 Evaluations
|
||||
5.3.1 Evaluation Settings
|
||||
5.3.2 Standard Evaluation
|
||||
5.3.3 Open-Ended Evaluation
|
||||
5.3.4 DeepSeek-V3 as a Generative Reward Model
|
||||
5.4 Discussion
|
||||
5.4.1 Distillation from DeepSeek-R1
|
||||
5.4.2 Self-Rewarding
|
||||
5.4.3 Multi-Token Prediction Evaluation
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Contributions and Acknowledgments
|
||||
B Ablation Studies for Low-Precision Training
|
||||
B.1 FP8 v.s. BF16 Training
|
||||
B.2 Discussion About Block-Wise Quantization
|
||||
C Expert Specialization Patterns of the 16B Aux-Loss-Based and Aux-Loss-Free Models
|
||||
## 1 Introduction
|
||||
In recent years, Large Language Models (LLMs) have been undergoing rapid iteration and evolution ( OpenAI, 2024a ; Anthropic, 2024 ; Google, 2024 ) , progressively diminishing the gap towards Artificial General Intelligence (AGI).
|
||||
Beyond closed-source models, open-source models, including DeepSeek series ( DeepSeek-AI, 2024b ; DeepSeek-AI, 2024c ; Guo et al., 2024 ; DeepSeek-AI, 2024a ) , LLaMA series ( Touvron et al., 2023a ; Touvron et al., 2023b ; AI@Meta, 2024a ; AI@Meta, 2024b ) , Qwen series ( Qwen, 2023 ; Qwen, 2024a ; Qwen, 2024b ) , and Mistral series ( Jiang et al., 2023 ; Mistral, 2024 ) , are also making significant strides, endeavoring to close the gap with their closed-source counterparts.
|
||||
To further push the boundaries of open-source model capabilities, we scale up our models and introduce DeepSeek-V3, a large Mixture-of-Experts (MoE) model with 671B parameters, of which 37B are activated for each token.
|
||||
With a forward-looking perspective, we consistently strive for strong model performance and economical costs.
|
||||
Therefore, in terms of architecture, DeepSeek-V3 still adopts Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024c ) for efficient inference and DeepSeekMoE ( Dai et al., 2024 ) for cost-effective training.
|
||||
These two architectures have been validated in DeepSeek-V2 ( DeepSeek-AI, 2024c ) , demonstrating their capability to maintain robust model performance while achieving efficient training and inference.
|
||||
Beyond the basic architecture, we implement two additional strategies to further enhance the model capabilities.
|
||||
Firstly, DeepSeek-V3 pioneers an auxiliary-loss-free strategy ( Wang et al., 2024a ) for load balancing, with the aim of minimizing the adverse impact on model performance that arises from the effort to encourage load balancing.
|
||||
Secondly, DeepSeek-V3 employs a multi-token prediction training objective, which we have observed to enhance the overall performance on evaluation benchmarks.
|
||||
In order to achieve efficient training, we support the FP8 mixed precision training and implement comprehensive optimizations for the training framework.
|
||||
Low-precision training has emerged as a promising solution for efficient training ( Kalamkar et al., 2019 ; Narang et al., 2017 ; Peng et al., 2023b ; Dettmers et al., 2022 ) , its evolution being closely tied to advancements in hardware capabilities ( Micikevicius et al., 2022 ; Luo et al., 2024 ; Rouhani et al., 2023a ) .
|
||||
In this work, we introduce an FP8 mixed precision training framework and, for the first time, validate its effectiveness on an extremely large-scale model.
|
||||
Through the support for FP8 computation and storage, we achieve both accelerated training and reduced GPU memory usage.
|
||||
As for the training framework, we design the DualPipe algorithm for efficient pipeline parallelism, which has fewer pipeline bubbles and hides most of the communication during training through computation-communication overlap.
|
||||
This overlap ensures that, as the model further scales up, as long as we maintain a constant computation-to-communication ratio, we can still employ fine-grained experts across nodes while achieving a near-zero all-to-all communication overhead.
|
||||
In addition, we also develop efficient cross-node all-to-all communication kernels to fully utilize InfiniBand (IB) and NVLink bandwidths.
|
||||
Furthermore, we meticulously optimize the memory footprint, making it possible to train DeepSeek-V3 without using costly tensor parallelism.
|
||||
Combining these efforts, we achieve high training efficiency.
|
||||
During pre-training, we train DeepSeek-V3 on 14.8T high-quality and diverse tokens.
|
||||
The pre-training process is remarkably stable.
|
||||
Throughout the entire training process, we did not encounter any irrecoverable loss spikes or have to roll back.
|
||||
Next, we conduct a two-stage context length extension for DeepSeek-V3. In the first stage, the maximum context length is extended to 32K, and in the second stage, it is further extended to 128K.
|
||||
Following this, we conduct post-training, including Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) on the base model of DeepSeek-V3, to align it with human preferences and further unlock its potential.
|
||||
During the post-training stage, we distill the reasoning capability from the DeepSeek-R1 series of models, and meanwhile carefully maintain the balance between model accuracy and generation length.
|
||||
We evaluate DeepSeek-V3 on a comprehensive array of benchmarks.
|
||||
Despite its economical training costs, comprehensive evaluations reveal that DeepSeek-V3-Base has emerged as the strongest open-source base model currently available, especially in code and math.
|
||||
Its chat version also outperforms other open-source models and achieves performance comparable to leading closed-source models, including GPT-4o and Claude-3.5-Sonnet, on a series of standard and open-ended benchmarks.
|
||||
Training Costs
|
||||
Pre-Training
|
||||
Context Extension
|
||||
Post-Training
|
||||
Total
|
||||
in H800 GPU Hours
|
||||
2664K
|
||||
119K
|
||||
5K
|
||||
2788K
|
||||
in USD
|
||||
$5.328M
|
||||
$0.238M
|
||||
$0.01M
|
||||
$5.576M
|
||||
Table 1:
|
||||
Training costs of DeepSeek-V3, assuming the rental price of H800 is $2 per GPU hour.
|
||||
Lastly, we emphasize again the economical training costs of DeepSeek-V3, summarized in Table 1 , achieved through our optimized co-design of algorithms, frameworks, and hardware.
|
||||
During the pre-training stage, training DeepSeek-V3 on each trillion tokens requires only 180K H800 GPU hours, i.e., 3.7 days on our cluster with 2048 H800 GPUs.
|
||||
Consequently, our pre-training stage is completed in less than two months and costs 2664K GPU hours.
|
||||
Combined with 119K GPU hours for the context length extension and 5K GPU hours for post-training, DeepSeek-V3 costs only 2.788M GPU hours for its full training.
|
||||
Assuming the rental price of the H800 GPU is $2 per GPU hour, our total training costs amount to only $5.576M.
|
||||
Note that the aforementioned costs include only the official training of DeepSeek-V3, excluding the costs associated with prior research and ablation experiments on architectures, algorithms, or data.
|
||||
Our main contribution includes:
|
||||
Architecture: Innovative Load Balancing Strategy and Training Objective
|
||||
•
|
||||
On top of the efficient architecture of DeepSeek-V2, we pioneer an auxiliary-loss-free strategy for load balancing, which minimizes the performance degradation that arises from encouraging load balancing.
|
||||
•
|
||||
We investigate a Multi-Token Prediction (MTP) objective and prove it beneficial to model performance.
|
||||
It can also be used for speculative decoding for inference acceleration.
|
||||
Pre-Training: Towards Ultimate Training Efficiency
|
||||
•
|
||||
We design an FP8 mixed precision training framework and, for the first time, validate the feasibility and effectiveness of FP8 training on an extremely large-scale model.
|
||||
•
|
||||
Through the co-design of algorithms, frameworks, and hardware, we overcome the communication bottleneck in cross-node MoE training, achieving near-full computation-communication overlap.
|
||||
This significantly enhances our training efficiency and reduces the training costs, enabling us to further scale up the model size without additional overhead.
|
||||
•
|
||||
At an economical cost of only 2.664M H800 GPU hours, we complete the pre-training of DeepSeek-V3 on 14.8T tokens, producing the currently strongest open-source base model.
|
||||
The subsequent training stages after pre-training require only 0.1M GPU hours.
|
||||
Post-Training: Knowledge Distillation from DeepSeek-R1
|
||||
•
|
||||
We introduce an innovative methodology to distill reasoning capabilities from the long-Chain-of-Thought (CoT) model, specifically from one of the DeepSeek R1 series models, into standard LLMs, particularly DeepSeek-V3.
|
||||
Our pipeline elegantly incorporates the verification and reflection patterns of R1 into DeepSeek-V3 and notably improves its reasoning performance.
|
||||
Meanwhile, we also maintain control over the output style and length of DeepSeek-V3.
|
||||
Summary of Core Evaluation Results
|
||||
•
|
||||
Knowledge :
|
||||
(1)
|
||||
On educational benchmarks such as MMLU, MMLU-Pro, and GPQA, DeepSeek-V3 outperforms all other open-source models, achieving 88.5 on MMLU, 75.9 on MMLU-Pro, and 59.1 on GPQA.
|
||||
Its performance is comparable to leading closed-source models like GPT-4o and Claude-Sonnet-3.5, narrowing the gap between open-source and closed-source models in this domain.
|
||||
(2)
|
||||
For factuality benchmarks, DeepSeek-V3 demonstrates superior performance among open-source models on both SimpleQA and Chinese SimpleQA.
|
||||
While it trails behind GPT-4o and Claude-Sonnet-3.5 in English factual knowledge (SimpleQA), it surpasses these models in Chinese factual knowledge (Chinese SimpleQA), highlighting its strength in Chinese factual knowledge.
|
||||
•
|
||||
Code, Math, and Reasoning :
|
||||
(1)
|
||||
DeepSeek-V3 achieves state-of-the-art performance on math-related benchmarks among all non-long-CoT open-source and closed-source models.
|
||||
Notably, it even outperforms o1-preview on specific benchmarks, such as MATH-500, demonstrating its robust mathematical reasoning capabilities.
|
||||
(2)
|
||||
On coding-related tasks, DeepSeek-V3 emerges as the top-performing model for coding competition benchmarks, such as LiveCodeBench, solidifying its position as the leading model in this domain.
|
||||
For engineering-related tasks, while DeepSeek-V3 performs slightly below Claude-Sonnet-3.5, it still outpaces all other models by a significant margin, demonstrating its competitiveness across diverse technical benchmarks.
|
||||
In the remainder of this paper, we first present a detailed exposition of our DeepSeek-V3 model architecture (Section 2 ).
|
||||
Subsequently, we introduce our infrastructures, encompassing our compute clusters, the training framework, the support for FP8 training, the inference deployment strategy, and our suggestions on future hardware design.
|
||||
Next, we describe our pre-training process, including the construction of training data, hyper-parameter settings, long-context extension techniques, the associated evaluations, as well as some discussions (Section 4 ).
|
||||
Thereafter, we discuss our efforts on post-training, which include Supervised Fine-Tuning (SFT), Reinforcement Learning (RL), the corresponding evaluations, and discussions (Section 5 ).
|
||||
Lastly, we conclude this work, discuss existing limitations of DeepSeek-V3, and propose potential directions for future research (Section 6 ).
|
||||
## 2 Architecture
|
||||
We first introduce the basic architecture of DeepSeek-V3, featured by Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024c ) for efficient inference and DeepSeekMoE ( Dai et al., 2024 ) for economical training.
|
||||
Then, we present a Multi-Token Prediction (MTP) training objective, which we have observed to enhance the overall performance on evaluation benchmarks.
|
||||
For other minor details not explicitly mentioned, DeepSeek-V3 adheres to the settings of DeepSeek-V2 ( DeepSeek-AI, 2024c ) .
|
||||
Figure 2:
|
||||
Illustration of the basic architecture of DeepSeek-V3.
|
||||
Following DeepSeek-V2, we adopt MLA and DeepSeekMoE for efficient inference and economical training.
|
||||
## 2.1 Basic Architecture
|
||||
The basic architecture of DeepSeek-V3 is still within the Transformer ( Vaswani et al., 2017 ) framework.
|
||||
For efficient inference and economical training, DeepSeek-V3 also adopts MLA and DeepSeekMoE, which have been thoroughly validated by DeepSeek-V2.
|
||||
Compared with DeepSeek-V2, an exception is that we additionally introduce an auxiliary-loss-free load balancing strategy ( Wang et al., 2024a ) for DeepSeekMoE to mitigate the performance degradation induced by the effort to ensure load balance.
|
||||
Figure 2 illustrates the basic architecture of DeepSeek-V3, and we will briefly review the details of MLA and DeepSeekMoE in this section.
|
||||
## 2.1.1 Multi-Head Latent Attention
|
||||
For attention, DeepSeek-V3 adopts the MLA architecture.
|
||||
Let d d denote the embedding dimension, n h n_{h} denote the number of attention heads, d h d_{h} denote the dimension per head, and 𝐡 t ∈ ℝ d \mathbf{h}_{t}\in\mathbb{R}^{d} denote the attention input for the t t -th token at a given attention layer.
|
||||
The core of MLA is the low-rank joint compression for attention keys and values to reduce Key-Value (KV) cache during inference:
|
||||
𝐜 t K V \displaystyle\boxed{\color[rgb]{0,0,1}\mathbf{c}_{t}^{KV}}
|
||||
= W D K V 𝐡 t , \displaystyle=W^{DKV}\mathbf{h}_{t},
|
||||
(1)
|
||||
[ 𝐤 t , 1 C ; 𝐤 t , 2 C ; … ; 𝐤 t , n h C ] = 𝐤 t C \displaystyle[\mathbf{k}_{t,1}^{C};\mathbf{k}_{t,2}^{C};...;\mathbf{k}_{t,n_{h}}^{C}]=\mathbf{k}_{t}^{C}
|
||||
= W U K 𝐜 t K V , \displaystyle=W^{UK}\mathbf{c}_{t}^{KV},
|
||||
(2)
|
||||
𝐤 t R \displaystyle\boxed{\color[rgb]{0,0,1}\mathbf{k}_{t}^{R}}
|
||||
= RoPE ( W K R 𝐡 t ) , \displaystyle=\operatorname{RoPE}({W^{KR}}\mathbf{h}_{t}),
|
||||
(3)
|
||||
𝐤 t , i \displaystyle\mathbf{k}_{t,i}
|
||||
= [ 𝐤 t , i C ; 𝐤 t R ] , \displaystyle=[\mathbf{k}_{t,i}^{C};\mathbf{k}_{t}^{R}],
|
||||
(4)
|
||||
[ 𝐯 t , 1 C ; 𝐯 t , 2 C ; … ; 𝐯 t , n h C ] = 𝐯 t C \displaystyle[\mathbf{v}_{t,1}^{C};\mathbf{v}_{t,2}^{C};...;\mathbf{v}_{t,n_{h}}^{C}]=\mathbf{v}_{t}^{C}
|
||||
= W U V 𝐜 t K V , \displaystyle=W^{UV}\mathbf{c}_{t}^{KV},
|
||||
(5)
|
||||
where 𝐜 t K V ∈ ℝ d c \mathbf{c}_{t}^{KV}\in\mathbb{R}^{d_{c}} is the compressed latent vector for keys and values;
|
||||
d c ( ≪ d h n h ) d_{c}(\ll d_{h}n_{h}) indicates the KV compression dimension;
|
||||
W D K V ∈ ℝ d c × d W^{DKV}\in\mathbb{R}^{d_{c}\times d} denotes the down-projection matrix;
|
||||
W U K , W U V ∈ ℝ d h n h × d c W^{UK},W^{UV}\in\mathbb{R}^{d_{h}n_{h}\times d_{c}} are the up-projection matrices for keys and values, respectively;
|
||||
W K R ∈ ℝ d h R × d W^{KR}\in\mathbb{R}^{d_{h}^{R}\times d} is the matrix used to produce the decoupled key that carries Rotary Positional Embedding (RoPE) ( Su et al., 2024 ) ;
|
||||
RoPE ( ⋅ ) \operatorname{RoPE}(\cdot) denotes the operation that applies RoPE matrices;
|
||||
and [ ⋅ ; ⋅ ] [\cdot;\cdot] denotes concatenation.
|
||||
Note that for MLA, only the blue-boxed vectors (i.e., 𝐜 t K V \color[rgb]{0,0,1}\mathbf{c}_{t}^{KV} and 𝐤 t R \color[rgb]{0,0,1}\mathbf{k}_{t}^{R} ) need to be cached during generation, which results in significantly reduced KV cache while maintaining performance comparable to standard Multi-Head Attention (MHA) ( Vaswani et al., 2017 ) .
|
||||
For the attention queries, we also perform a low-rank compression, which can reduce the activation memory during training:
|
||||
𝐜 t Q \displaystyle\mathbf{c}_{t}^{Q}
|
||||
= W D Q 𝐡 t , \displaystyle=W^{DQ}\mathbf{h}_{t},
|
||||
(6)
|
||||
[ 𝐪 t , 1 C ; 𝐪 t , 2 C ; … ; 𝐪 t , n h C ] = 𝐪 t C \displaystyle[\mathbf{q}_{t,1}^{C};\mathbf{q}_{t,2}^{C};...;\mathbf{q}_{t,n_{h}}^{C}]=\mathbf{q}_{t}^{C}
|
||||
= W U Q 𝐜 t Q , \displaystyle=W^{UQ}\mathbf{c}_{t}^{Q},
|
||||
(7)
|
||||
[ 𝐪 t , 1 R ; 𝐪 t , 2 R ; … ; 𝐪 t , n h R ] = 𝐪 t R \displaystyle[\mathbf{q}_{t,1}^{R};\mathbf{q}_{t,2}^{R};...;\mathbf{q}_{t,n_{h}}^{R}]=\mathbf{q}_{t}^{R}
|
||||
= RoPE ( W Q R 𝐜 t Q ) , \displaystyle=\operatorname{RoPE}({W^{QR}}\mathbf{c}_{t}^{Q}),
|
||||
(8)
|
||||
𝐪 t , i \displaystyle\mathbf{q}_{t,i}
|
||||
= [ 𝐪 t , i C ; 𝐪 t , i R ] , \displaystyle=[\mathbf{q}_{t,i}^{C};\mathbf{q}_{t,i}^{R}],
|
||||
(9)
|
||||
where 𝐜 t Q ∈ ℝ d c ′ \mathbf{c}_{t}^{Q}\in\mathbb{R}^{d_{c}^{\prime}} is the compressed latent vector for queries;
|
||||
d c ′ ( ≪ d h n h ) d_{c}^{\prime}(\ll d_{h}n_{h}) denotes the query compression dimension;
|
||||
W D Q ∈ ℝ d c ′ × d , W U Q ∈ ℝ d h n h × d c ′ W^{DQ}\in\mathbb{R}^{d_{c}^{\prime}\times d},W^{UQ}\in\mathbb{R}^{d_{h}n_{h}\times d_{c}^{\prime}} are the down-projection and up-projection matrices for queries, respectively;
|
||||
and W Q R ∈ ℝ d h R n h × d c ′ W^{QR}\in\mathbb{R}^{d_{h}^{R}n_{h}\times d_{c}^{\prime}} is the matrix to produce the decoupled queries that carry RoPE.
|
||||
Ultimately, the attention queries ( 𝐪 t , i \mathbf{q}_{t,i} ), keys ( 𝐤 j , i \mathbf{k}_{j,i} ), and values ( 𝐯 j , i C \mathbf{v}_{j,i}^{C} ) are combined to yield the final attention output 𝐮 t \mathbf{u}_{t} :
|
||||
𝐨 t , i \displaystyle\mathbf{o}_{t,i}
|
||||
= ∑ j = 1 t Softmax j ( 𝐪 t , i T 𝐤 j , i d h + d h R ) 𝐯 j , i C , \displaystyle=\sum_{j=1}^{t}\operatorname{Softmax}_{j}(\frac{\mathbf{q}_{t,i}^{T}\mathbf{k}_{j,i}}{\sqrt{d_{h}+d_{h}^{R}}})\mathbf{v}_{j,i}^{C},
|
||||
(10)
|
||||
𝐮 t \displaystyle\mathbf{u}_{t}
|
||||
= W O [ 𝐨 t , 1 ; 𝐨 t , 2 ; … ; 𝐨 t , n h ] , \displaystyle=W^{O}[\mathbf{o}_{t,1};\mathbf{o}_{t,2};...;\mathbf{o}_{t,n_{h}}],
|
||||
(11)
|
||||
where W O ∈ ℝ d × d h n h W^{O}\in\mathbb{R}^{d\times d_{h}n_{h}} denotes the output projection matrix.
|
||||
## 2.1.2 DeepSeekMoE with Auxiliary-Loss-Free Load Balancing
|
||||
## Basic Architecture of DeepSeekMoE.
|
||||
For Feed-Forward Networks (FFNs), DeepSeek-V3 employs the DeepSeekMoE architecture ( Dai et al., 2024 ) .
|
||||
Compared with traditional MoE architectures like GShard ( Lepikhin et al., 2021 ) , DeepSeekMoE uses finer-grained experts and isolates some experts as shared ones.
|
||||
Let 𝐮 t \mathbf{u}_{t} denote the FFN input of the t t -th token, we compute the FFN output 𝐡 t ′ \mathbf{h}_{t}^{\prime} as follows:
|
||||
𝐡 t ′ \displaystyle\mathbf{h}_{t}^{\prime}
|
||||
= 𝐮 t + ∑ i = 1 N s FFN i ( s ) ( 𝐮 t ) + ∑ i = 1 N r g i , t FFN i ( r ) ( 𝐮 t ) , \displaystyle=\mathbf{u}_{t}+\sum_{i=1}^{N_{s}}{\operatorname{FFN}^{(s)}_{i}\left(\mathbf{u}_{t}\right)}+\sum_{i=1}^{N_{r}}{g_{i,t}\operatorname{FFN}^{(r)}_{i}\left(\mathbf{u}_{t}\right)},
|
||||
(12)
|
||||
g i , t \displaystyle g_{i,t}
|
||||
= g i , t ′ ∑ j = 1 N r g j , t ′ , \displaystyle=\frac{g^{\prime}_{i,t}}{\sum_{j=1}^{N_{r}}g^{\prime}_{j,t}},
|
||||
(13)
|
||||
g i , t ′ \displaystyle g^{\prime}_{i,t}
|
||||
= { s i , t , s i , t ∈ Topk ( { s j , t | 1 ⩽ j ⩽ N r } , K r ) , 0 , otherwise , \displaystyle=\begin{cases}s_{i,t},&s_{i,t}\in\operatorname{Topk}(\{s_{j,t}|1\leqslant j\leqslant N_{r}\},K_{r}),\\
|
||||
0,&\text{otherwise},\end{cases}
|
||||
(14)
|
||||
s i , t \displaystyle s_{i,t}
|
||||
= Sigmoid ( 𝐮 t T 𝐞 i ) , \displaystyle=\operatorname{Sigmoid}\left({\mathbf{u}_{t}}^{T}\mathbf{e}_{i}\right),
|
||||
(15)
|
||||
where N s N_{s} and N r N_{r} denote the numbers of shared experts and routed experts, respectively;
|
||||
FFN i ( s ) ( ⋅ ) \operatorname{FFN}^{(s)}_{i}(\cdot) and FFN i ( r ) ( ⋅ ) \operatorname{FFN}^{(r)}_{i}(\cdot) denote the i i -th shared expert and the i i -th routed expert, respectively;
|
||||
K r K_{r} denotes the number of activated routed experts;
|
||||
g i , t g_{i,t} is the gating value for the i i -th expert;
|
||||
s i , t s_{i,t} is the token-to-expert affinity;
|
||||
𝐞 i \mathbf{e}_{i} is the centroid vector of the i i -th routed expert;
|
||||
and Topk ( ⋅ , K ) \operatorname{Topk}(\cdot,K) denotes the set comprising K K highest scores among the affinity scores calculated for the t t -th token and all routed experts.
|
||||
Slightly different from DeepSeek-V2, DeepSeek-V3 uses the sigmoid function to compute the affinity scores, and applies a normalization among all selected affinity scores to produce the gating values.
|
||||
## Auxiliary-Loss-Free Load Balancing.
|
||||
For MoE models, an unbalanced expert load will lead to routing collapse ( Shazeer et al., 2017 ) and diminish computational efficiency in scenarios with expert parallelism.
|
||||
Conventional solutions usually rely on the auxiliary loss ( Fedus et al., 2021 ; Lepikhin et al., 2021 ) to avoid unbalanced load.
|
||||
However, too large an auxiliary loss will impair the model performance ( Wang et al., 2024a ) .
|
||||
To achieve a better trade-off between load balance and model performance, we pioneer an auxiliary-loss-free load balancing strategy ( Wang et al., 2024a ) to ensure load balance.
|
||||
To be specific, we introduce a bias term b i b_{i} for each expert and add it to the corresponding affinity scores s i , t s_{i,t} to determine the top-K routing:
|
||||
g i , t ′ \displaystyle g^{\prime}_{i,t}
|
||||
= { s i , t , s i , t + b i ∈ Topk ( { s j , t + b j | 1 ⩽ j ⩽ N r } , K r ) , 0 , otherwise . \displaystyle=\begin{cases}s_{i,t},&s_{i,t}+b_{i}\in\operatorname{Topk}(\{s_{j,t}+b_{j}|1\leqslant j\leqslant N_{r}\},K_{r}),\\
|
||||
0,&\text{otherwise}.\end{cases}
|
||||
(16)
|
||||
Note that the bias term is only used for routing.
|
||||
The gating value, which will be multiplied with the FFN output, is still derived from the original affinity score s i , t s_{i,t} .
|
||||
During training, we keep monitoring the expert load on the whole batch of each training step.
|
||||
At the end of each step, we will decrease the bias term by γ \gamma if its corresponding expert is overloaded, and increase it by γ \gamma if its corresponding expert is underloaded, where γ \gamma is a hyper-parameter called bias update speed.
|
||||
Through the dynamic adjustment, DeepSeek-V3 keeps balanced expert load during training, and achieves better performance than models that encourage load balance through pure auxiliary losses.
|
||||
## Complementary Sequence-Wise Auxiliary Loss.
|
||||
Although DeepSeek-V3 mainly relies on the auxiliary-loss-free strategy for load balance, to prevent extreme imbalance within any single sequence, we also employ a complementary sequence-wise balance loss:
|
||||
ℒ Bal \displaystyle\mathcal{L}_{\mathrm{Bal}}
|
||||
= α ∑ i = 1 N r f i P i , \displaystyle=\alpha\sum_{i=1}^{N_{r}}{f_{i}P_{i}},
|
||||
(17)
|
||||
f i = N r K r T ∑ t = 1 T 𝟙 \displaystyle f_{i}=\frac{N_{r}}{K_{r}T}\sum_{t=1}^{T}\mathds{1}
|
||||
( s i , t ∈ Topk ( { s j , t | 1 ⩽ j ⩽ N r } , K r ) ) , \displaystyle\left(s_{i,t}\in\operatorname{Topk}(\{s_{j,t}|1\leqslant j\leqslant N_{r}\},K_{r})\right),
|
||||
(18)
|
||||
s i , t ′ \displaystyle s^{\prime}_{i,t}
|
||||
= s i , t ∑ j = 1 N r s j , t , \displaystyle=\frac{s_{i,t}}{\sum_{j=1}^{N_{r}}s_{j,t}},
|
||||
(19)
|
||||
P i \displaystyle P_{i}
|
||||
= 1 T ∑ t = 1 T s i , t ′ , \displaystyle=\frac{1}{T}\sum_{t=1}^{T}{s^{\prime}_{i,t}},
|
||||
(20)
|
||||
where the balance factor α \alpha is a hyper-parameter, which will be assigned an extremely small value for DeepSeek-V3;
|
||||
𝟙 ( ⋅ ) \mathds{1}(\cdot) denotes the indicator function;
|
||||
and T T denotes the number of tokens in a sequence.
|
||||
The sequence-wise balance loss encourages the expert load on each sequence to be balanced.
|
||||
## Node-Limited Routing.
|
||||
Like the device-limited routing used by DeepSeek-V2, DeepSeek-V3 also uses a restricted routing mechanism to limit communication costs during training.
|
||||
In short, we ensure that each token will be sent to at most M M nodes, which are selected according to the sum of the highest K r M \frac{K_{r}}{M} affinity scores of the experts distributed on each node.
|
||||
Under this constraint, our MoE training framework can nearly achieve full computation-communication overlap.
|
||||
## No Token-Dropping.
|
||||
Due to the effective load balancing strategy, DeepSeek-V3 keeps a good load balance during its full training.
|
||||
Therefore, DeepSeek-V3 does not drop any tokens during training.
|
||||
In addition, we also implement specific deployment strategies to ensure inference load balance, so DeepSeek-V3 also does not drop tokens during inference.
|
||||
Figure 3:
|
||||
Illustration of our Multi-Token Prediction (MTP) implementation.
|
||||
We keep the complete causal chain for the prediction of each token at each depth.
|
||||
## 2.2 Multi-Token Prediction
|
||||
Inspired by Gloeckle et al. (2024) , we investigate and set a Multi-Token Prediction (MTP) objective for DeepSeek-V3, which extends the prediction scope to multiple future tokens at each position.
|
||||
On the one hand, an MTP objective densifies the training signals and may improve data efficiency.
|
||||
On the other hand, MTP may enable the model to pre-plan its representations for better prediction of future tokens.
|
||||
Figure 3 illustrates our implementation of MTP.
|
||||
Different from Gloeckle et al. (2024) , which parallelly predicts D D additional tokens using independent output heads, we sequentially predict additional tokens and keep the complete causal chain at each prediction depth.
|
||||
We introduce the details of our MTP implementation in this section.
|
||||
## MTP Modules.
|
||||
To be specific, our MTP implementation uses D D sequential modules to predict D D additional tokens.
|
||||
The k k -th MTP module consists of a shared embedding layer Emb ( ⋅ ) \operatorname{Emb}(\cdot) , a shared output head OutHead ( ⋅ ) \operatorname{OutHead}(\cdot) , a Transformer block TRM k ( ⋅ ) \operatorname{TRM}_{k}(\cdot) , and a projection matrix M k ∈ ℝ d × 2 d M_{k}\in\mathbb{R}^{d\times 2d} .
|
||||
For the i i -th input token t i t_{i} , at the k k -th prediction depth, we first combine the representation of the i i -th token at the ( k − 1 ) (k-1) -th depth 𝐡 i k − 1 ∈ ℝ d \mathbf{h}_{i}^{k-1}\in\mathbb{R}^{d} and the embedding of the ( i + k ) (i+k) -th token E m b ( t i + k ) ∈ ℝ d Emb(t_{i+k})\in\mathbb{R}^{d} with the linear projection:
|
||||
𝐡 i ′ k = M k [ RMSNorm ( 𝐡 i k − 1 ) ; RMSNorm ( Emb ( t i + k ) ) ] , \mathbf{h}_{i}^{\prime k}=M_{k}[\operatorname{RMSNorm}(\mathbf{h}_{i}^{k-1});\operatorname{RMSNorm}(\operatorname{Emb}(t_{i+k}))],
|
||||
(21)
|
||||
where [ ⋅ ; ⋅ ] [\cdot;\cdot] denotes concatenation.
|
||||
Especially, when k = 1 k=1 , 𝐡 i k − 1 \mathbf{h}_{i}^{k-1} refers to the representation given by the main model.
|
||||
Note that for each MTP module, its embedding layer is shared with the main model.
|
||||
The combined 𝐡 i ′ k \mathbf{h}_{i}^{\prime k} serves as the input of the Transformer block at the k k -th depth to produce the output representation at the current depth 𝐡 i k \mathbf{h}_{i}^{k} :
|
||||
𝐡 1 : T − k k = TRM k ( 𝐡 1 : T − k ′ k ) , \mathbf{h}_{1:T-k}^{k}=\operatorname{TRM}_{k}(\mathbf{h}_{1:T-k}^{\prime k}),
|
||||
(22)
|
||||
where T T represents the input sequence length and i:j denotes the slicing operation (inclusive of both the left and right boundaries).
|
||||
Finally, taking 𝐡 i k \mathbf{h}_{i}^{k} as the input, the shared output head will compute the probability distribution for the k k -th additional prediction token P i + 1 + k k ∈ ℝ V P_{i+1+k}^{k}\in\mathbb{R}^{V} , where V V is the vocabulary size:
|
||||
P i + k + 1 k = OutHead ( 𝐡 i k ) . P_{i+k+1}^{k}=\operatorname{OutHead}(\mathbf{h}_{i}^{k}).
|
||||
(23)
|
||||
The output head OutHead ( ⋅ ) \operatorname{OutHead}(\cdot) linearly maps the representation to logits and subsequently applies the Softmax ( ⋅ ) \operatorname{Softmax}(\cdot) function to compute the prediction probabilities of the k k -th additional token.
|
||||
Also, for each MTP module, its output head is shared with the main model.
|
||||
Our principle of maintaining the causal chain of predictions is similar to that of EAGLE ( Li et al., 2024b ) , but its primary objective is speculative decoding ( Xia et al., 2023 ; Leviathan et al., 2023 ) , whereas we utilize MTP to improve training.
|
||||
## MTP Training Objective.
|
||||
For each prediction depth, we compute a cross-entropy loss ℒ MTP k \mathcal{L}_{\text{MTP}}^{k} :
|
||||
ℒ MTP k = CrossEntropy ( P 2 + k : T + 1 k , t 2 + k : T + 1 ) = − 1 T ∑ i = 2 + k T + 1 log P i k [ t i ] , \mathcal{L}_{\text{MTP}}^{k}=\operatorname{CrossEntropy}(P_{2+k:T+1}^{k},t_{2+k:T+1})=-\frac{1}{T}\sum_{i=2+k}^{T+1}\log P_{i}^{k}[t_{i}],
|
||||
(24)
|
||||
where T T denotes the input sequence length, t i t_{i} denotes the ground-truth token at the i i -th position, and P i k [ t i ] P_{i}^{k}[t_{i}] denotes the corresponding prediction probability of t i t_{i} , given by the k k -th MTP module.
|
||||
Finally, we compute the average of the MTP losses across all depths and multiply it by a weighting factor λ \lambda to obtain the overall MTP loss ℒ MTP \mathcal{L}_{\text{MTP}} , which serves as an additional training objective for DeepSeek-V3:
|
||||
ℒ MTP = λ D ∑ k = 1 D ℒ MTP k . \mathcal{L}_{\text{MTP}}=\frac{\lambda}{D}\sum_{k=1}^{D}\mathcal{L}_{\text{MTP}}^{k}.
|
||||
(25)
|
||||
## MTP in Inference.
|
||||
Our MTP strategy mainly aims to improve the performance of the main model, so during inference, we can directly discard the MTP modules and the main model can function independently and normally.
|
||||
Additionally, we can also repurpose these MTP modules for speculative decoding to further improve the generation latency.
|
||||
## 3 Infrastructures
|
||||
## 3.1 Compute Clusters
|
||||
DeepSeek-V3 is trained on a cluster equipped with 2048 NVIDIA H800 GPUs.
|
||||
Each node in the H800 cluster contains 8 GPUs connected by NVLink and NVSwitch within nodes.
|
||||
Across different nodes, InfiniBand (IB) interconnects are utilized to facilitate communications.
|
||||
## 3.2 Training Framework
|
||||
The training of DeepSeek-V3 is supported by the HAI-LLM framework, an efficient and lightweight training framework crafted by our engineers from the ground up.
|
||||
On the whole, DeepSeek-V3 applies 16-way Pipeline Parallelism (PP) ( Qi et al., 2023a ) , 64-way Expert Parallelism (EP) ( Lepikhin et al., 2021 ) spanning 8 nodes, and ZeRO-1 Data Parallelism (DP) ( Rajbhandari et al., 2020 ) .
|
||||
In order to facilitate efficient training of DeepSeek-V3, we implement meticulous engineering optimizations.
|
||||
Firstly, we design the DualPipe algorithm for efficient pipeline parallelism.
|
||||
Compared with existing PP methods, DualPipe has fewer pipeline bubbles.
|
||||
More importantly, it overlaps the computation and communication phases across forward and backward processes, thereby addressing the challenge of heavy communication overhead introduced by cross-node expert parallelism.
|
||||
Secondly, we develop efficient cross-node all-to-all communication kernels to fully utilize IB and NVLink bandwidths and conserve Streaming Multiprocessors (SMs) dedicated to communication.
|
||||
Finally, we meticulously optimize the memory footprint during training, thereby enabling us to train DeepSeek-V3 without using costly Tensor Parallelism (TP).
|
||||
## 3.2.1 DualPipe and Computation-Communication Overlap
|
||||
Figure 4:
|
||||
Overlapping strategy for a pair of individual forward and backward chunks (the boundaries of the transformer blocks are not aligned).
|
||||
Orange denotes forward, green denotes "backward for input", blue denotes "backward for weights", purple denotes PP communication, and red denotes barriers.
|
||||
Both all-to-all and PP communication can be fully hidden.
|
||||
For DeepSeek-V3, the communication overhead introduced by cross-node expert parallelism results in an inefficient computation-to-communication ratio of approximately 1:1.
|
||||
To tackle this challenge, we design an innovative pipeline parallelism algorithm called DualPipe, which not only accelerates model training by effectively overlapping forward and backward computation-communication phases, but also reduces the pipeline bubbles.
|
||||
The key idea of DualPipe is to overlap the computation and communication within a pair of individual forward and backward chunks.
|
||||
To be specific, we divide each chunk into four components: attention , all-to-all dispatch , MLP , and all-to-all combine .
|
||||
Specially, for a backward chunk, both attention and MLP are further split into two parts, backward for input and backward for weights , like in ZeroBubble ( Qi et al., 2023b ) .
|
||||
In addition, we have a PP communication component.
|
||||
As illustrated in Figure 4 , for a pair of forward and backward chunks, we rearrange these components and manually adjust the ratio of GPU SMs dedicated to communication versus computation.
|
||||
In this overlapping strategy, we can ensure that both all-to-all and PP communication can be fully hidden during execution.
|
||||
Given the efficient overlapping strategy, the full DualPipe scheduling is illustrated in Figure 5 .
|
||||
It employs a bidirectional pipeline scheduling, which feeds micro-batches from both ends of the pipeline simultaneously and a significant portion of communications can be fully overlapped.
|
||||
This overlap also ensures that, as the model further scales up, as long as we maintain a constant computation-to-communication ratio, we can still employ fine-grained experts across nodes while achieving a near-zero all-to-all communication overhead.
|
||||
Figure 5:
|
||||
Example DualPipe scheduling for 8 PP ranks and 20 micro-batches in two directions.
|
||||
The micro-batches in the reverse direction are symmetric to those in the forward direction, so we omit their batch ID for illustration simplicity.
|
||||
Two cells enclosed by a shared black border have mutually overlapped computation and communication.
|
||||
In addition, even in more general scenarios without a heavy communication burden, DualPipe still exhibits efficiency advantages.
|
||||
In Table 2 , we summarize the pipeline bubbles and memory usage across different PP methods.
|
||||
As shown in the table, compared with ZB1P ( Qi et al., 2023b ) and 1F1B ( Harlap et al., 2018 ) , DualPipe significantly reduces the pipeline bubbles while only increasing the peak activation memory by 1 P P \frac{1}{PP} times.
|
||||
Although DualPipe requires keeping two copies of the model parameters, this does not significantly increase the memory consumption since we use a large EP size during training.
|
||||
Compared with Chimera ( Li and Hoefler, 2021 ) , DualPipe only requires that the pipeline stages and micro-batches be divisible by 2, without requiring micro-batches to be divisible by pipeline stages.
|
||||
In addition, for DualPipe, neither the bubbles nor activation memory will increase as the number of micro-batches grows.
|
||||
Method
|
||||
Bubble
|
||||
Parameter
|
||||
Activation
|
||||
1F1B
|
||||
( P P − 1 ) ( F + B ) (PP-1)(F+B)
|
||||
1 × 1\times
|
||||
P P PP
|
||||
ZB1P
|
||||
( P P − 1 ) ( F + B − 2 W ) (PP-1)(F+B-2W)
|
||||
1 × 1\times
|
||||
P P PP
|
||||
DualPipe (Ours)
|
||||
( P P 2 − 1 ) ( F & B + B − 3 W ) (\frac{PP}{2}-1)(F\&B+B-3W)
|
||||
2 × 2\times
|
||||
P P + 1 PP+1
|
||||
Table 2:
|
||||
Comparison of pipeline bubbles and memory usage across different pipeline parallel methods. F F denotes the execution time of a forward chunk, B B denotes the execution time of a full backward chunk, W W denotes the execution time of a "backward for weights" chunk, and F & B F\&B denotes the execution time of two mutually overlapped forward and backward chunks.
|
||||
## 3.2.2 Efficient Implementation of Cross-Node All-to-All Communication
|
||||
In order to ensure sufficient computational performance for DualPipe, we customize efficient cross-node all-to-all communication kernels (including dispatching and combining) to conserve the number of SMs dedicated to communication.
|
||||
The implementation of the kernels is co-designed with the MoE gating algorithm and the network topology of our cluster.
|
||||
To be specific, in our cluster, cross-node GPUs are fully interconnected with IB, and intra-node communications are handled via NVLink.
|
||||
NVLink offers a bandwidth of 160 GB/s, roughly 3.2 times that of IB (50 GB/s).
|
||||
To effectively leverage the different bandwidths of IB and NVLink, we limit each token to be dispatched to at most 4 nodes, thereby reducing IB traffic.
|
||||
For each token, when its routing decision is made, it will first be transmitted via IB to the GPUs with the same in-node index on its target nodes.
|
||||
Once it reaches the target nodes, we will endeavor to ensure that it is instantaneously forwarded via NVLink to specific GPUs that host their target experts, without being blocked by subsequently arriving tokens.
|
||||
In this way, communications via IB and NVLink are fully overlapped, and each token can efficiently select an average of 3.2 experts per node without incurring additional overhead from NVLink.
|
||||
This implies that, although DeepSeek-V3 selects only 8 routed experts in practice, it can scale up this number to a maximum of 13 experts (4 nodes × \times 3.2 experts/node) while preserving the same communication cost.
|
||||
Overall, under such a communication strategy, only 20 SMs are sufficient to fully utilize the bandwidths of IB and NVLink.
|
||||
In detail, we employ the warp specialization technique ( Bauer et al., 2014 ) and partition 20 SMs into 10 communication channels.
|
||||
During the dispatching process, (1) IB sending, (2) IB-to-NVLink forwarding, and (3) NVLink receiving are handled by respective warps.
|
||||
The number of warps allocated to each communication task is dynamically adjusted according to the actual workload across all SMs.
|
||||
Similarly, during the combining process, (1) NVLink sending, (2) NVLink-to-IB forwarding and accumulation, and (3) IB receiving and accumulation are also handled by dynamically adjusted warps.
|
||||
In addition, both dispatching and combining kernels overlap with the computation stream, so we also consider their impact on other SM computation kernels.
|
||||
Specifically, we employ customized PTX (Parallel Thread Execution) instructions and auto-tune the communication chunk size, which significantly reduces the use of the L2 cache and the interference to other SMs.
|
||||
## 3.2.3 Extremely Memory Saving with Minimal Overhead
|
||||
In order to reduce the memory footprint during training, we employ the following techniques.
|
||||
## Recomputation of RMSNorm and MLA Up-Projection.
|
||||
We recompute all RMSNorm operations and MLA up-projections during back-propagation, thereby eliminating the need to persistently store their output activations.
|
||||
With a minor overhead, this strategy significantly reduces memory requirements for storing activations.
|
||||
## Exponential Moving Average in CPU.
|
||||
During training, we preserve the Exponential Moving Average (EMA) of the model parameters for early estimation of the model performance after learning rate decay.
|
||||
The EMA parameters are stored in CPU memory and are updated asynchronously after each training step.
|
||||
This method allows us to maintain EMA parameters without incurring additional memory or time overhead.
|
||||
## Shared Embedding and Output Head for Multi-Token Prediction.
|
||||
With the DualPipe strategy, we deploy the shallowest layers (including the embedding layer) and deepest layers (including the output head) of the model on the same PP rank.
|
||||
This arrangement enables the physical sharing of parameters and gradients, of the shared embedding and output head, between the MTP module and the main model.
|
||||
This physical sharing mechanism further enhances our memory efficiency.
|
||||
## 3.3 FP8 Training
|
||||
Figure 6:
|
||||
The overall mixed precision framework with FP8 data format. For clarification, only the Linear operator is illustrated.
|
||||
Inspired by recent advances in low-precision training ( Peng et al., 2023b ; Dettmers et al., 2022 ; Noune et al., 2022 ) , we propose a fine-grained mixed precision framework utilizing the FP8 data format for training DeepSeek-V3.
|
||||
While low-precision training holds great promise, it is often limited by the presence of outliers in activations, weights, and gradients ( Sun et al., 2024 ; He et al., ; Fishman et al., 2024 ) .
|
||||
Although significant progress has been made in inference quantization ( Xiao et al., 2023 ; Frantar et al., 2022 ) , there are relatively few studies demonstrating successful application of low-precision techniques in large-scale language model pre-training ( Fishman et al., 2024 ) .
|
||||
To address this challenge and effectively extend the dynamic range of the FP8 format, we introduce a fine-grained quantization strategy: tile-wise grouping with 1 × N c 1\times N_{c} elements or block-wise grouping with N c × N c N_{c}\times N_{c} elements.
|
||||
The associated dequantization overhead is largely mitigated under our increased-precision accumulation process, a critical aspect for achieving accurate FP8 General Matrix Multiplication (GEMM).
|
||||
Moreover, to further reduce memory and communication overhead in MoE training, we cache and dispatch activations in FP8, while storing low-precision optimizer states in BF16.
|
||||
We validate the proposed FP8 mixed precision framework on two model scales similar to DeepSeek-V2-Lite and DeepSeek-V2, training for approximately 1 trillion tokens (see more details in Appendix B.1 ).
|
||||
Notably, compared with the BF16 baseline, the relative loss error of our FP8-training model remains consistently below 0.25%, a level well within the acceptable range of training randomness.
|
||||
## 3.3.1 Mixed Precision Framework
|
||||
Building upon widely adopted techniques in low-precision training ( Kalamkar et al., 2019 ; Narang et al., 2017 ) , we propose a mixed precision framework for FP8 training.
|
||||
In this framework, most compute-density operations are conducted in FP8, while a few key operations are strategically maintained in their original data formats to balance training efficiency and numerical stability.
|
||||
The overall framework is illustrated in Figure 6 .
|
||||
Firstly, in order to accelerate model training, the majority of core computation kernels, i.e., GEMM operations, are implemented in FP8 precision.
|
||||
These GEMM operations accept FP8 tensors as inputs and produce outputs in BF16 or FP32.
|
||||
As depicted in Figure 6 , all three GEMMs associated with the Linear operator, namely Fprop (forward pass), Dgrad (activation backward pass), and Wgrad (weight backward pass), are executed in FP8.
|
||||
This design theoretically doubles the computational speed compared with the original BF16 method.
|
||||
Additionally, the FP8 Wgrad GEMM allows activations to be stored in FP8 for use in the backward pass.
|
||||
This significantly reduces memory consumption.
|
||||
Despite the efficiency advantage of the FP8 format, certain operators still require a higher precision due to their sensitivity to low-precision computations.
|
||||
Besides, some low-cost operators can also utilize a higher precision with a negligible overhead to the overall training cost.
|
||||
For this reason, after careful investigations, we maintain the original precision (e.g., BF16 or FP32) for the following components: the embedding module, the output head, MoE gating modules, normalization operators, and attention operators.
|
||||
These targeted retentions of high precision ensure stable training dynamics for DeepSeek-V3.
|
||||
To further guarantee numerical stability, we store the master weights, weight gradients, and optimizer states in higher precision.
|
||||
While these high-precision components incur some memory overheads, their impact can be minimized through efficient sharding across multiple DP ranks in our distributed training system.
|
||||
Figure 7:
|
||||
(a) We propose a fine-grained quantization method to mitigate quantization errors caused by feature outliers; for illustration simplicity, only Fprop is illustrated.
|
||||
(b) In conjunction with our quantization strategy, we improve the FP8 GEMM precision by promoting to CUDA Cores at an interval of N C = 128 N_{C}=128 elements MMA for the high-precision accumulation.
|
||||
## 3.3.2 Improved Precision from Quantization and Multiplication
|
||||
Based on our mixed precision FP8 framework, we introduce several strategies to enhance low-precision training accuracy, focusing on both the quantization method and the multiplication process.
|
||||
## Fine-Grained Quantization.
|
||||
In low-precision training frameworks, overflows and underflows are common challenges due to the limited dynamic range of the FP8 format, which is constrained by its reduced exponent bits.
|
||||
As a standard practice, the input distribution is aligned to the representable range of the FP8 format by scaling the maximum absolute value of the input tensor to the maximum representable value of FP8 ( Narang et al., 2017 ) .
|
||||
This method makes low-precision training highly sensitive to activation outliers, which can heavily degrade quantization accuracy.
|
||||
To solve this, we propose a fine-grained quantization method that applies scaling at a more granular level.
|
||||
As illustrated in Figure 7 (a), (1) for activations, we group and scale elements on a 1x128 tile basis (i.e., per token per 128 channels); and (2) for weights, we group and scale elements on a 128x128 block basis (i.e., per 128 input channels per 128 output channels).
|
||||
This approach ensures that the quantization process can better accommodate outliers by adapting the scale according to smaller groups of elements.
|
||||
In Appendix B.2 , we further discuss the training instability when we group and scale activations on a block basis in the same way as weights quantization.
|
||||
One key modification in our method is the introduction of per-group scaling factors along the inner dimension of GEMM operations.
|
||||
This functionality is not directly supported in the standard FP8 GEMM.
|
||||
However, combined with our precise FP32 accumulation strategy, it can be efficiently implemented.
|
||||
Notably, our fine-grained quantization strategy is highly consistent with the idea of microscaling formats ( Rouhani et al., 2023b ) , while the Tensor Cores of NVIDIA next-generation GPUs (Blackwell series) have announced the support for microscaling formats with smaller quantization granularity ( NVIDIA, 2024a ) .
|
||||
We hope our design can serve as a reference for future work to keep pace with the latest GPU architectures.
|
||||
## Increasing Accumulation Precision.
|
||||
Low-precision GEMM operations often suffer from underflow issues, and their accuracy largely depends on high-precision accumulation, which is commonly performed in an FP32 precision ( Kalamkar et al., 2019 ; Narang et al., 2017 ) .
|
||||
However, we observe that the accumulation precision of FP8 GEMM on NVIDIA H800 GPUs is limited to retaining around 14 bits, which is significantly lower than FP32 accumulation precision.
|
||||
This problem will become more pronounced when the inner dimension K is large ( Wortsman et al., 2023 ) , a typical scenario in large-scale model training where the batch size and model width are increased.
|
||||
Taking GEMM operations of two random matrices with K = 4096 for example, in our preliminary test, the limited accumulation precision in Tensor Cores results in a maximum relative error of nearly 2%.
|
||||
Despite these problems, the limited accumulation precision is still the default option in a few FP8 frameworks ( NVIDIA, 2024b ) , severely constraining the training accuracy.
|
||||
In order to address this issue, we adopt the strategy of promotion to CUDA Cores for higher precision ( Thakkar et al., 2023 ) .
|
||||
The process is illustrated in Figure 7 (b).
|
||||
To be specific, during MMA (Matrix Multiply-Accumulate) execution on Tensor Cores, intermediate results are accumulated using the limited bit width.
|
||||
Once an interval of N C N_{C} is reached, these partial results will be copied to FP32 registers on CUDA Cores, where full-precision FP32 accumulation is performed.
|
||||
As mentioned before, our fine-grained quantization applies per-group scaling factors along the inner dimension K .
|
||||
These scaling factors can be efficiently multiplied on the CUDA Cores as the dequantization process with minimal additional computational cost.
|
||||
It is worth noting that this modification reduces the WGMMA (Warpgroup-level Matrix Multiply-Accumulate) instruction issue rate for a single warpgroup. However, on the H800 architecture, it is typical for two WGMMA to persist concurrently: while one warpgroup performs the promotion operation, the other is able to execute the MMA operation. This design enables overlapping of the two operations, maintaining high utilization of Tensor Cores.
|
||||
Based on our experiments, setting N C = 128 N_{C}=128 elements, equivalent to 4 WGMMAs, represents the minimal accumulation interval that can significantly improve precision without introducing substantial overhead.
|
||||
## Mantissa over Exponents.
|
||||
In contrast to the hybrid FP8 format adopted by prior work ( NVIDIA, 2024b ; Peng et al., 2023b ; Sun et al., 2019b ) , which uses E4M3 (4-bit exponent and 3-bit mantissa) in Fprop and E5M2 (5-bit exponent and 2-bit mantissa) in Dgrad and Wgrad , we adopt the E4M3 format on all tensors for higher precision.
|
||||
We attribute the feasibility of this approach to our fine-grained quantization strategy, i.e., tile and block-wise scaling.
|
||||
By operating on smaller element groups, our methodology effectively shares exponent bits among these grouped elements, mitigating the impact of the limited dynamic range.
|
||||
## Online Quantization.
|
||||
Delayed quantization is employed in tensor-wise quantization frameworks ( NVIDIA, 2024b ; Peng et al., 2023b ) , which maintains a history of the maximum absolute values across prior iterations to infer the current value.
|
||||
In order to ensure accurate scales and simplify the framework, we calculate the maximum absolute value online for each 1x128 activation tile or 128x128 weight block.
|
||||
Based on it, we derive the scaling factor and then quantize the activation or weight online into the FP8 format.
|
||||
## 3.3.3 Low-Precision Storage and Communication
|
||||
In conjunction with our FP8 training framework, we further reduce the memory consumption and communication overhead by compressing cached activations and optimizer states into lower-precision formats.
|
||||
## Low-Precision Optimizer States.
|
||||
We adopt the BF16 data format instead of FP32 to track the first and second moments in the AdamW ( Loshchilov and Hutter, 2017 ) optimizer, without incurring observable performance degradation.
|
||||
However, the master weights (stored by the optimizer) and gradients (used for batch size accumulation) are still retained in FP32 to ensure numerical stability throughout training.
|
||||
## Low-Precision Activation.
|
||||
As illustrated in Figure 6 , the Wgrad operation is performed in FP8.
|
||||
To reduce the memory consumption, it is a natural choice to cache activations in FP8 format for the backward pass of the Linear operator.
|
||||
However, special considerations are taken on several operators for low-cost high-precision training:
|
||||
(1) Inputs of the Linear after the attention operator. These activations are also used in the backward pass of the attention operator, which makes it sensitive to precision. We adopt a customized E5M6 data format exclusively for these activations. Additionally, these activations will be converted from an 1x128 quantization tile to an 128x1 tile in the backward pass. To avoid introducing extra quantization error, all the scaling factors are round scaled, i.e., integral power of 2.
|
||||
(2) Inputs of the SwiGLU operator in MoE. To further reduce the memory cost, we cache the inputs of the SwiGLU operator and recompute its output in the backward pass. These activations are also stored in FP8 with our fine-grained quantization method, striking a balance between memory efficiency and computational accuracy.
|
||||
## Low-Precision Communication.
|
||||
Communication bandwidth is a critical bottleneck in the training of MoE models.
|
||||
To alleviate this challenge, we quantize the activation before MoE up-projections into FP8 and then apply dispatch components, which is compatible with FP8 Fprop in MoE up-projections. Like the inputs of the Linear after the attention operator, scaling factors for this activation are integral power of 2.
|
||||
A similar strategy is applied to the activation gradient before MoE down-projections.
|
||||
For both the forward and backward combine components, we retain them in BF16 to preserve training precision in critical parts of the training pipeline.
|
||||
## 3.4 Inference and Deployment
|
||||
We deploy DeepSeek-V3 on the H800 cluster, where GPUs within each node are interconnected using NVLink, and all GPUs across the cluster are fully interconnected via IB.
|
||||
To simultaneously ensure both the Service-Level Objective (SLO) for online services and high throughput, we employ the following deployment strategy that separates the prefilling and decoding stages.
|
||||
## 3.4.1 Prefilling
|
||||
The minimum deployment unit of the prefilling stage consists of 4 nodes with 32 GPUs.
|
||||
The attention part employs 4-way Tensor Parallelism (TP4) with Sequence Parallelism (SP), combined with 8-way Data Parallelism (DP8).
|
||||
Its small TP size of 4 limits the overhead of TP communication.
|
||||
For the MoE part, we use 32-way Expert Parallelism (EP32), which ensures that each expert processes a sufficiently large batch size, thereby enhancing computational efficiency.
|
||||
For the MoE all-to-all communication, we use the same method as in training: first transferring tokens across nodes via IB, and then forwarding among the intra-node GPUs via NVLink.
|
||||
In particular, we use 1-way Tensor Parallelism for the dense MLPs in shallow layers to save TP communication.
|
||||
To achieve load balancing among different experts in the MoE part, we need to ensure that each GPU processes approximately the same number of tokens.
|
||||
To this end, we introduce a deployment strategy of redundant experts , which duplicates high-load experts and deploys them redundantly.
|
||||
The high-load experts are detected based on statistics collected during the online deployment and are adjusted periodically (e.g., every 10 minutes).
|
||||
After determining the set of redundant experts, we carefully rearrange experts among GPUs within a node based on the observed loads, striving to balance the load across GPUs as much as possible without increasing the cross-node all-to-all communication overhead.
|
||||
For the deployment of DeepSeek-V3, we set 32 redundant experts for the prefilling stage.
|
||||
For each GPU, besides the original 8 experts it hosts, it will also host one additional redundant expert.
|
||||
Furthermore, in the prefilling stage, to improve the throughput and hide the overhead of all-to-all and TP communication, we simultaneously process two micro-batches with similar computational workloads, overlapping the attention and MoE of one micro-batch with the dispatch and combine of another.
|
||||
Finally, we are exploring a dynamic redundancy strategy for experts, where each GPU hosts more experts (e.g., 16 experts), but only 9 will be activated during each inference step.
|
||||
Before the all-to-all operation at each layer begins, we compute the globally optimal routing scheme on the fly.
|
||||
Given the substantial computation involved in the prefilling stage, the overhead of computing this routing scheme is almost negligible.
|
||||
## 3.4.2 Decoding
|
||||
During decoding, we treat the shared expert as a routed one.
|
||||
From this perspective, each token will select 9 experts during routing, where the shared expert is regarded as a heavy-load one that will always be selected.
|
||||
The minimum deployment unit of the decoding stage consists of 40 nodes with 320 GPUs.
|
||||
The attention part employs TP4 with SP, combined with DP80, while the MoE part uses EP320.
|
||||
For the MoE part, each GPU hosts only one expert, and 64 GPUs are responsible for hosting redundant experts and shared experts.
|
||||
All-to-all communication of the dispatch and combine parts is performed via direct point-to-point transfers over IB to achieve low latency.
|
||||
Additionally, we leverage the IBGDA ( NVIDIA, 2022 ) technology to further minimize latency and enhance communication efficiency.
|
||||
Similar to prefilling, we periodically determine the set of redundant experts in a certain interval, based on the statistical expert load from our online service.
|
||||
However, we do not need to rearrange experts since each GPU only hosts one expert.
|
||||
We are also exploring the dynamic redundancy strategy for decoding.
|
||||
However, this requires more careful optimization of the algorithm that computes the globally optimal routing scheme and the fusion with the dispatch kernel to reduce overhead.
|
||||
Additionally,
|
||||
|
||||
=== BODY (tail) ===
|
||||
pen-source model.
|
||||
Additionally, it is competitive against frontier closed-source models like GPT-4o and Claude-3.5-Sonnet.
|
||||
## English Benchmarks.
|
||||
MMLU is a widely recognized benchmark designed to assess the performance of large language models, across diverse knowledge domains and tasks.
|
||||
DeepSeek-V3 demonstrates competitive performance, standing on par with top-tier models such as LLaMA-3.1-405B, GPT-4o, and Claude-Sonnet 3.5, while significantly outperforming Qwen2.5 72B.
|
||||
Moreover, DeepSeek-V3 excels in MMLU-Pro, a more challenging educational knowledge benchmark, where it closely trails Claude-Sonnet 3.5.
|
||||
On MMLU-Redux, a refined version of MMLU with corrected labels, DeepSeek-V3 surpasses its peers.
|
||||
In addition, on GPQA-Diamond, a PhD-level evaluation testbed, DeepSeek-V3 achieves remarkable results, ranking just behind Claude 3.5 Sonnet and outperforming all other competitors by a substantial margin.
|
||||
In long-context understanding benchmarks such as DROP, LongBench v2, and FRAMES, DeepSeek-V3 continues to demonstrate its position as a top-tier model.
|
||||
It achieves an impressive 91.6 F1 score in the 3-shot setting on DROP, outperforming all other models in this category.
|
||||
On FRAMES, a benchmark requiring question-answering over 100k token contexts, DeepSeek-V3 closely trails GPT-4o while outperforming all other models by a significant margin.
|
||||
This demonstrates the strong capability of DeepSeek-V3 in handling extremely long-context tasks.
|
||||
The long-context capability of DeepSeek-V3 is further validated by its best-in-class performance on LongBench v2, a dataset that was released just a few weeks before the launch of DeepSeek V3.
|
||||
On the factual knowledge benchmark, SimpleQA, DeepSeek-V3 falls behind GPT-4o and Claude-Sonnet, primarily due to its design focus and resource allocation.
|
||||
DeepSeek-V3 assigns more training tokens to learn Chinese knowledge, leading to exceptional performance on the C-SimpleQA.
|
||||
On the instruction-following benchmark, DeepSeek-V3 significantly outperforms its predecessor, DeepSeek-V2-series, highlighting its improved ability to understand and adhere to user-defined format constraints.
|
||||
## Code and Math Benchmarks.
|
||||
Coding is a challenging and practical task for LLMs, encompassing engineering-focused tasks like SWE-Bench-Verified and Aider, as well as algorithmic tasks such as HumanEval and LiveCodeBench.
|
||||
In engineering tasks, DeepSeek-V3 trails behind Claude-Sonnet-3.5-1022 but significantly outperforms open-source models.
|
||||
The open-source DeepSeek-V3 is expected to foster advancements in coding-related engineering tasks.
|
||||
By providing access to its robust capabilities, DeepSeek-V3 can drive innovation and improvement in areas such as software engineering and algorithm development, empowering developers and researchers to push the boundaries of what open-source models can achieve in coding tasks.
|
||||
In algorithmic tasks, DeepSeek-V3 demonstrates superior performance, outperforming all baselines on benchmarks like HumanEval-Mul and LiveCodeBench.
|
||||
This success can be attributed to its advanced knowledge distillation technique, which effectively enhances its code generation and problem-solving capabilities in algorithm-focused tasks.
|
||||
On math benchmarks, DeepSeek-V3 demonstrates exceptional performance, significantly surpassing baselines and setting a new state-of-the-art for non-o1-like models.
|
||||
Specifically, on AIME, MATH-500, and CNMO 2024, DeepSeek-V3 outperforms the second-best model, Qwen2.5 72B, by approximately 10% in absolute scores, which is a substantial margin for such challenging benchmarks.
|
||||
This remarkable capability highlights the effectiveness of the distillation technique from DeepSeek-R1, which has been proven highly beneficial for non-o1-like models.
|
||||
## Chinese Benchmarks.
|
||||
Qwen and DeepSeek are two representative model series with robust support for both Chinese and English.
|
||||
On the factual benchmark Chinese SimpleQA, DeepSeek-V3 surpasses Qwen2.5-72B by 16.4 points, despite Qwen2.5 being trained on a larger corpus compromising 18T tokens, which are 20% more than the 14.8T tokens that DeepSeek-V3 is pre-trained on.
|
||||
On C-Eval, a representative benchmark for Chinese educational knowledge evaluation, and CLUEWSC (Chinese Winograd Schema Challenge), DeepSeek-V3 and Qwen2.5-72B exhibit similar performance levels, indicating that both models are well-optimized for challenging Chinese-language reasoning and educational tasks.
|
||||
Model
|
||||
Arena-Hard
|
||||
AlpacaEval 2.0
|
||||
DeepSeek-V2.5-0905
|
||||
76.2
|
||||
50.5
|
||||
Qwen2.5-72B-Instruct
|
||||
81.2
|
||||
49.1
|
||||
LLaMA-3.1 405B
|
||||
69.3
|
||||
40.5
|
||||
GPT-4o-0513
|
||||
80.4
|
||||
51.1
|
||||
Claude-Sonnet-3.5-1022
|
||||
85.2
|
||||
52.0
|
||||
DeepSeek-V3
|
||||
85.5
|
||||
70.0
|
||||
Table 7:
|
||||
English open-ended conversation evaluations.
|
||||
For AlpacaEval 2.0, we use the length-controlled win rate as the metric.
|
||||
## 5.3.3 Open-Ended Evaluation
|
||||
In addition to standard benchmarks, we also evaluate our models on open-ended generation tasks using LLMs as judges, with the results shown in Table 7 .
|
||||
Specifically, we adhere to the original configurations of AlpacaEval 2.0 ( Dubois et al., 2024 ) and Arena-Hard ( Li et al., 2024a ) , which leverage GPT-4-Turbo-1106 as judges for pairwise comparisons.
|
||||
On Arena-Hard, DeepSeek-V3 achieves an impressive win rate of over 86% against the baseline GPT-4-0314, performing on par with top-tier models like Claude-Sonnet-3.5-1022.
|
||||
This underscores the robust capabilities of DeepSeek-V3, especially in dealing with complex prompts, including coding and debugging tasks.
|
||||
Furthermore, DeepSeek-V3 achieves a groundbreaking milestone as the first open-source model to surpass 85% on the Arena-Hard benchmark.
|
||||
This achievement significantly bridges the performance gap between open-source and closed-source models, setting a new standard for what open-source models can accomplish in challenging domains.
|
||||
Similarly, DeepSeek-V3 showcases exceptional performance on AlpacaEval 2.0, outperforming both closed-source and open-source models.
|
||||
This demonstrates its outstanding proficiency in writing tasks and handling straightforward question-answering scenarios.
|
||||
Notably, it surpasses DeepSeek-V2.5-0905 by a significant margin of 20%, highlighting substantial improvements in tackling simple tasks and showcasing the effectiveness of its advancements.
|
||||
## 5.3.4 DeepSeek-V3 as a Generative Reward Model
|
||||
We compare the judgment ability of DeepSeek-V3 with state-of-the-art models, namely GPT-4o and Claude-3.5.
|
||||
Table 8 presents the performance of these models in RewardBench ( Lambert et al., 2024 ) .
|
||||
DeepSeek-V3 achieves performance on par with the best versions of GPT-4o-0806 and Claude-3.5-Sonnet-1022, while surpassing other versions.
|
||||
Additionally, the judgment ability of DeepSeek-V3 can also be enhanced by the voting technique.
|
||||
Therefore, we employ DeepSeek-V3 along with voting to offer self-feedback on open-ended questions, thereby improving the effectiveness and robustness of the alignment process.
|
||||
Model
|
||||
Chat
|
||||
Chat-Hard
|
||||
Safety
|
||||
Reasoning
|
||||
Average
|
||||
GPT-4o-0513
|
||||
96.6
|
||||
70.4
|
||||
86.7
|
||||
84.9
|
||||
84.7
|
||||
GPT-4o-0806
|
||||
96.1
|
||||
76.1
|
||||
88.1
|
||||
86.6
|
||||
86.7
|
||||
GPT-4o-1120
|
||||
95.8
|
||||
71.3
|
||||
86.2
|
||||
85.2
|
||||
84.6
|
||||
Claude-3.5-sonnet-0620
|
||||
96.4
|
||||
74.0
|
||||
81.6
|
||||
84.7
|
||||
84.2
|
||||
Claude-3.5-sonnet-1022
|
||||
96.4
|
||||
79.7
|
||||
91.1
|
||||
87.6
|
||||
88.7
|
||||
DeepSeek-V3
|
||||
96.9
|
||||
79.8
|
||||
87.0
|
||||
84.3
|
||||
87.0
|
||||
DeepSeek-V3 (maj@6)
|
||||
96.9
|
||||
82.6
|
||||
89.5
|
||||
89.2
|
||||
89.6
|
||||
Table 8: Performances of GPT-4o, Claude-3.5-sonnet and DeepSeek-V3 on RewardBench.
|
||||
## 5.4 Discussion
|
||||
## 5.4.1 Distillation from DeepSeek-R1
|
||||
We ablate the contribution of distillation from DeepSeek-R1 based on DeepSeek-V2.5.
|
||||
The baseline is trained on short CoT data, whereas its competitor uses data generated by the expert checkpoints described above.
|
||||
Table 9 demonstrates the effectiveness of the distillation data, showing significant improvements in both LiveCodeBench and MATH-500 benchmarks.
|
||||
Our experiments reveal an interesting trade-off: the distillation leads to better performance but also substantially increases the average response length.
|
||||
To maintain a balance between model accuracy and computational efficiency, we carefully selected optimal settings for DeepSeek-V3 in distillation.
|
||||
Our research suggests that knowledge distillation from reasoning models presents a promising direction for post-training optimization.
|
||||
While our current work focuses on distilling data from mathematics and coding domains, this approach shows potential for broader applications across various task domains.
|
||||
The effectiveness demonstrated in these specific areas indicates that long-CoT distillation could be valuable for enhancing model performance in other cognitive tasks requiring complex reasoning.
|
||||
Further exploration of this approach across different domains remains an important direction for future research.
|
||||
Model
|
||||
LiveCodeBench-CoT
|
||||
MATH-500
|
||||
Pass@1
|
||||
Length
|
||||
Pass@1
|
||||
Length
|
||||
DeepSeek-V2.5 Baseline
|
||||
31.1
|
||||
718
|
||||
74.6
|
||||
769
|
||||
DeepSeek-V2.5 +R1 Distill
|
||||
37.4
|
||||
783
|
||||
83.2
|
||||
1510
|
||||
Table 9:
|
||||
The contribution of distillation from DeepSeek-R1.
|
||||
The evaluation settings of LiveCodeBench and MATH-500 are the same as in Table 6 .
|
||||
## 5.4.2 Self-Rewarding
|
||||
Rewards play a pivotal role in RL, steering the optimization process.
|
||||
In domains where verification through external tools is straightforward, such as some coding or mathematics scenarios, RL demonstrates exceptional efficacy.
|
||||
However, in more general scenarios, constructing a feedback mechanism through hard coding is impractical.
|
||||
During the development of DeepSeek-V3, for these broader contexts, we employ the constitutional AI approach ( Bai et al., 2022 ) , leveraging the voting evaluation results of DeepSeek-V3 itself as a feedback source.
|
||||
This method has produced notable alignment effects, significantly enhancing the performance of DeepSeek-V3 in subjective evaluations.
|
||||
By integrating additional constitutional inputs, DeepSeek-V3 can optimize towards the constitutional direction.
|
||||
We believe that this paradigm, which combines supplementary information with LLMs as a feedback source, is of paramount importance.
|
||||
The LLM serves as a versatile processor capable of transforming unstructured information from diverse scenarios into rewards, ultimately facilitating the self-improvement of LLMs.
|
||||
Beyond self-rewarding, we are also dedicated to uncovering other general and scalable rewarding methods to consistently advance the model capabilities in general scenarios.
|
||||
## 5.4.3 Multi-Token Prediction Evaluation
|
||||
Instead of predicting just the next single token, DeepSeek-V3 predicts the next 2 tokens through the MTP technique.
|
||||
Combined with the framework of speculative decoding ( Leviathan et al., 2023 ; Xia et al., 2023 ) , it can significantly accelerate the decoding speed of the model.
|
||||
A natural question arises concerning the acceptance rate of the additionally predicted token.
|
||||
Based on our evaluation, the acceptance rate of the second token prediction ranges between 85% and 90% across various generation topics, demonstrating consistent reliability.
|
||||
This high acceptance rate enables DeepSeek-V3 to achieve a significantly improved decoding speed, delivering 1.8 times TPS (Tokens Per Second).
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
In this paper, we introduce DeepSeek-V3, a large MoE language model with 671B total parameters and 37B activated parameters, trained on 14.8T tokens.
|
||||
In addition to the MLA and DeepSeekMoE architectures, it also pioneers an auxiliary-loss-free strategy for load balancing and sets a multi-token prediction training objective for stronger performance.
|
||||
The training of DeepSeek-V3 is cost-effective due to the support of FP8 training and meticulous engineering optimizations.
|
||||
The post-training also makes a success in distilling the reasoning capability from the DeepSeek-R1 series of models.
|
||||
Comprehensive evaluations demonstrate that DeepSeek-V3 has emerged as the strongest open-source model currently available, and achieves performance comparable to leading closed-source models like GPT-4o and Claude-3.5-Sonnet.
|
||||
Despite its strong performance, it also maintains economical training costs.
|
||||
It requires only 2.788M H800 GPU hours for its full training, including pre-training, context length extension, and post-training.
|
||||
While acknowledging its strong performance and cost-effectiveness, we also recognize that DeepSeek-V3 has some limitations, especially on the deployment.
|
||||
Firstly, to ensure efficient inference, the recommended deployment unit for DeepSeek-V3 is relatively large, which might pose a burden for small-sized teams.
|
||||
Secondly, although our deployment strategy for DeepSeek-V3 has achieved an end-to-end generation speed of more than two times that of DeepSeek-V2, there still remains potential for further enhancement.
|
||||
Fortunately, these limitations are expected to be naturally addressed with the development of more advanced hardware.
|
||||
DeepSeek consistently adheres to the route of open-source models with longtermism, aiming to steadily approach the ultimate goal of AGI (Artificial General Intelligence).
|
||||
In the future, we plan to strategically invest in research across the following directions.
|
||||
•
|
||||
We will consistently study and refine our model architectures, aiming to further improve both the training and inference efficiency, striving to approach efficient support for infinite context length.
|
||||
Additionally, we will try to break through the architectural limitations of Transformer, thereby pushing the boundaries of its modeling capabilities.
|
||||
•
|
||||
We will continuously iterate on the quantity and quality of our training data, and explore the incorporation of additional training signal sources, aiming to drive data scaling across a more comprehensive range of dimensions.
|
||||
•
|
||||
We will consistently explore and iterate on the deep thinking capabilities of our models, aiming to enhance their intelligence and problem-solving abilities by expanding their reasoning length and depth.
|
||||
•
|
||||
We will explore more comprehensive and multi-dimensional model evaluation methods to prevent the tendency towards optimizing a fixed set of benchmarks during research, which may create a misleading impression of the model capabilities and affect our foundational assessment.
|
||||
@@ -0,0 +1,995 @@
|
||||
TITLE: DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
|
||||
ARXIV: 2501.12948
|
||||
AUTHORS: DeepSeek-AI , Daya Guo , Dejian Yang , Haowei Zhang , Junxiao Song , Peiyi Wang , Qihao Zhu , Runxin Xu , Ruoyu Zhang , Shirong Ma , Xiao Bi , Xiaokang Zhang , Xingkai Yu , Yu Wu , Z.F. Wu , Zhibin Gou , Zhihong Shao , Zhuoshu Li , Ziyi Gao , Aixin Liu , Bing Xue , Bingxuan Wang , Bochao Wu , Bei Feng , Chengda Lu , Chenggang Zhao , Chengqi Deng , Chenyu Zhang , Chong Ruan , Damai Dai , Deli Chen , Dongjie Ji , Erhang Li , Fangyun Lin , Fucong Dai , Fuli Luo , Guangbo Hao , Guanting Chen , Guowei Li , H. Zhang , Han Bao , Hanwei Xu , Haocheng Wang , Honghui Ding , Huajian Xin , Huazuo Gao , Hui Qu , Hui Li , Jianzhong Guo , Jiashi Li , Jiawei Wang , Jingchang Chen , Jingyang Yuan , Junjie Qiu , Junlong Li , J.L. Cai , Jiaqi Ni , Jian Liang , Jin Chen , Kai Dong , Kai Hu , Kaige Gao , Kang Guan , Kexin Huang , Kuai Yu , Lean Wang , Lecong Zhang , Liang Zhao , Litong Wang , Liyue Zhang , Lei Xu , Leyi Xia , Mingchuan Zhang , Minghua Zhang , Minghui Tang , Meng Li , Miaojun Wang , Mingming Li , Ning Tian , Panpan Huang , Peng Zhang , Qiancheng Wang , Qinyu Chen , Qiushi Du , Ruiqi Ge , Ruisong Zhang , Ruizhe Pan , Runji Wang , R.J. Chen , R.L. Jin , Ruyi Chen , Shanghao Lu , Shangyan Zhou , Shanhuang Chen , Shengfeng Ye , Shiyu Wang , Shuiping Yu , Shunfeng Zhou , Shuting Pan , S.S. Li , Shuang Zhou , Shaoqing Wu , Shengfeng Ye , Tao Yun , Tian Pei , Tianyu Sun , T. Wang , Wangding Zeng , Wanjia Zhao , Wen Liu , Wenfeng Liang , Wenjun Gao , Wenqin Yu , Wentao Zhang , W.L. Xiao ,
|
||||
DATE: [Submitted on 22 Jan 2025 ( v1 ), last revised 4 Jan 2026 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
|
||||
## Abstract
|
||||
## keywords
|
||||
## 1 Introduction
|
||||
## 2 DeepSeek-R1-Zero
|
||||
## 2.1 Group Relative Policy Optimization
|
||||
## 2.2 Reward Design
|
||||
## 2.3 Incentivize Reasoning Capability in LLMs
|
||||
## 3 DeepSeek-R1
|
||||
## 3.1 Model-based Rewards
|
||||
## Helpful Reward Model
|
||||
## Safety Reward Model
|
||||
## 3.2 Training Details
|
||||
## 3.2.1 Training Details of the First RL Stage
|
||||
## 3.2.2 Training Details of the Second RL Stage
|
||||
## 4 Experiment
|
||||
## 5 Ethics and Safety Statement
|
||||
## 6 Conclusion, Limitation, and Future Work
|
||||
## 7 Author List
|
||||
## Appendix
|
||||
## Appendix A Background
|
||||
## A.1 DeepSeek-V3
|
||||
## A.2 Conventional Post-Training Paradigm
|
||||
## A.3 A Comparison of GRPO and PPO
|
||||
## Appendix B Training Details
|
||||
## B.1 RL Infrastructure
|
||||
## B.2 Reward Model Prompt
|
||||
## B.3 Data Recipe
|
||||
## B.3.1 RL Data
|
||||
## B.3.2 DeepSeek-R1 Cold Start
|
||||
## Question
|
||||
## Thought process
|
||||
## Question
|
||||
## Response
|
||||
## Question
|
||||
## Response
|
||||
## Question
|
||||
## Response
|
||||
## Answer Quality Classification
|
||||
## Question
|
||||
## Reference Answer
|
||||
## Answer to be Evaluated
|
||||
## Output Format
|
||||
## B.3.3 800K Supervised Data
|
||||
## Reasoning Data
|
||||
## Non-Reasoning Data
|
||||
## SFT Data Statistics
|
||||
## Examples of SFT Trajectories
|
||||
## Question
|
||||
## Response
|
||||
## Question
|
||||
## Response
|
||||
## Question
|
||||
## Response
|
||||
## B.4 Hyper-Parameters
|
||||
## B.4.1 Hyper-Parameters of DeepSeek-R1-Zero-Qwen-32B
|
||||
## B.4.2 Hyper-Parameters of SFT
|
||||
## B.4.3 Hyper-Parameters of Distillation
|
||||
## B.4.4 Training Cost
|
||||
## B.5 Reward Hacking
|
||||
## B.6 Ablation Study of Language Consistency Reward
|
||||
## Appendix C Self-Evolution of DeepSeek-R1-Zero
|
||||
## C.1 Evolution of Reasoning Capability in DeepSeek-R1-Zero during Training
|
||||
## C.2 Evolution of Advanced Reasoning Behaviors in DeepSeek-R1-Zero during Training
|
||||
## Appendix D Evaluation of DeepSeek-R1
|
||||
## D.1 Experiment Setup
|
||||
## Benchmarks
|
||||
## Decontamination
|
||||
## Evaluation Prompts
|
||||
## Baselines
|
||||
## D.2 Main Results
|
||||
## Standard Benchmark
|
||||
## Human Evaluation
|
||||
## D.3 DeepSeek-R1 Safety Report
|
||||
## D.3.1 Risk Control System for DeepSeek-R1
|
||||
## Potential Risky Dialogue Filtering
|
||||
## Model-based Risk Review
|
||||
## D.3.2 R1 Safety Evaluation on Standard Benchmarks
|
||||
## D.3.3 Safety Taxonomic Study of R1 on In-House Benchmark
|
||||
## Discrimination and Prejudice Issues
|
||||
## Illegal and Criminal Behavior
|
||||
## Harmful Behavior
|
||||
## Moral and Ethical Issues
|
||||
## D.3.4 Multilingual Safety Performance
|
||||
## D.3.5 Robustness against Jailbreaking
|
||||
## Appendix E More Analysis
|
||||
## E.1 Performance Comparison with DeepSeek-V3
|
||||
## E.2 Generalization to Real-World Competitions
|
||||
## E.3 Mathematical Capabilities Breakdown by Categories
|
||||
## E.4 An Analysis on CoT Length
|
||||
## E.5 Performance of Each Stage on Problems of Varying Difficulty
|
||||
## Appendix F DeepSeek-R1 Distillation
|
||||
## F.1 Distillation v.s. Reinforcement Learning
|
||||
## Appendix G Discussion
|
||||
## G.1 Key Findings
|
||||
## G.2 Unsuccessful Attempts
|
||||
## Process Reward Model (PRM)
|
||||
## Monte Carlo Tree Search (MCTS)
|
||||
## Appendix H Related Work
|
||||
## H.1 Chain-of-thought Reasoning
|
||||
## H.2 Scaling Inference-time Compute
|
||||
## H.3 Reinforcement Learning for Reasoning Enhancement
|
||||
## Appendix I Open Weights, Code, and Data
|
||||
## Appendix J Evaluation Prompts and Settings
|
||||
## Question:
|
||||
## Choices:
|
||||
## Instruction
|
||||
## Instruction
|
||||
## Model Outputs
|
||||
## Task
|
||||
## Best Model Identifier
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 DeepSeek-R1-Zero
|
||||
2.1 Group Relative Policy Optimization
|
||||
2.2 Reward Design
|
||||
2.3 Incentivize Reasoning Capability in LLMs
|
||||
3 DeepSeek-R1
|
||||
3.1 Model-based Rewards
|
||||
Helpful Reward Model
|
||||
Safety Reward Model
|
||||
3.2 Training Details
|
||||
3.2.1 Training Details of the First RL Stage
|
||||
3.2.2 Training Details of the Second RL Stage
|
||||
4 Experiment
|
||||
5 Ethics and Safety Statement
|
||||
6 Conclusion, Limitation, and Future Work
|
||||
7 Author List
|
||||
A Background
|
||||
A.1 DeepSeek-V3
|
||||
A.2 Conventional Post-Training Paradigm
|
||||
A.3 A Comparison of GRPO and PPO
|
||||
B Training Details
|
||||
B.1 RL Infrastructure
|
||||
B.2 Reward Model Prompt
|
||||
B.3 Data Recipe
|
||||
B.3.1 RL Data
|
||||
B.3.2 DeepSeek-R1 Cold Start
|
||||
B.3.3 800K Supervised Data
|
||||
Reasoning Data
|
||||
Non-Reasoning Data
|
||||
SFT Data Statistics
|
||||
Examples of SFT Trajectories
|
||||
B.4 Hyper-Parameters
|
||||
B.4.1 Hyper-Parameters of DeepSeek-R1-Zero-Qwen-32B
|
||||
B.4.2 Hyper-Parameters of SFT
|
||||
B.4.3 Hyper-Parameters of Distillation
|
||||
B.4.4 Training Cost
|
||||
B.5 Reward Hacking
|
||||
B.6 Ablation Study of Language Consistency Reward
|
||||
C Self-Evolution of DeepSeek-R1-Zero
|
||||
C.1 Evolution of Reasoning Capability in DeepSeek-R1-Zero during Training
|
||||
C.2 Evolution of Advanced Reasoning Behaviors in DeepSeek-R1-Zero during Training
|
||||
D Evaluation of DeepSeek-R1
|
||||
D.1 Experiment Setup
|
||||
Benchmarks
|
||||
Decontamination
|
||||
Evaluation Prompts
|
||||
Baselines
|
||||
D.2 Main Results
|
||||
Standard Benchmark
|
||||
Human Evaluation
|
||||
D.3 DeepSeek-R1 Safety Report
|
||||
D.3.1 Risk Control System for DeepSeek-R1
|
||||
Potential Risky Dialogue Filtering
|
||||
Model-based Risk Review
|
||||
D.3.2 R1 Safety Evaluation on Standard Benchmarks
|
||||
D.3.3 Safety Taxonomic Study of R1 on In-House Benchmark
|
||||
Discrimination and Prejudice Issues
|
||||
Illegal and Criminal Behavior
|
||||
Harmful Behavior
|
||||
Moral and Ethical Issues
|
||||
D.3.4 Multilingual Safety Performance
|
||||
D.3.5 Robustness against Jailbreaking
|
||||
E More Analysis
|
||||
E.1 Performance Comparison with DeepSeek-V3
|
||||
E.2 Generalization to Real-World Competitions
|
||||
E.3 Mathematical Capabilities Breakdown by Categories
|
||||
E.4 An Analysis on CoT Length
|
||||
E.5 Performance of Each Stage on Problems of Varying Difficulty
|
||||
F DeepSeek-R1 Distillation
|
||||
F.1 Distillation v.s. Reinforcement Learning
|
||||
G Discussion
|
||||
G.1 Key Findings
|
||||
G.2 Unsuccessful Attempts
|
||||
Process Reward Model (PRM)
|
||||
Monte Carlo Tree Search (MCTS)
|
||||
H Related Work
|
||||
H.1 Chain-of-thought Reasoning
|
||||
H.2 Scaling Inference-time Compute
|
||||
H.3 Reinforcement Learning for Reasoning Enhancement
|
||||
I Open Weights, Code, and Data
|
||||
J Evaluation Prompts and Settings
|
||||
References
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2501.12948v2 [cs.CL] 04 Jan 2026
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
|
||||
DeepSeek-AI
|
||||
research@deepseek.com
|
||||
## Abstract
|
||||
General reasoning represents a long-standing and formidable challenge in artificial intelligence. Recent breakthroughs, exemplified by large language models (LLMs) ( Brown et al., 2020 ; OpenAI, 2023 ) and chain-of-thought prompting ( Wei et al., 2022b ) , have achieved considerable success on foundational reasoning tasks. However, this success is heavily contingent upon extensive human-annotated demonstrations, and models’ capabilities are still insufficient for more complex problems.
|
||||
Here we show that the reasoning abilities of LLMs can be incentivized through pure reinforcement learning (RL), obviating the need for human-labeled reasoning trajectories. The proposed RL framework facilitates the emergent development of advanced reasoning patterns, such as self-reflection, verification, and dynamic strategy adaptation. Consequently, the trained model achieves superior performance on verifiable tasks such as mathematics, coding competitions, and STEM fields, surpassing its counterparts trained via conventional supervised learning on human demonstrations. Moreover, the emergent reasoning patterns exhibited by these large-scale models can be systematically harnessed to guide and enhance the reasoning capabilities of smaller models.
|
||||
## keywords
|
||||
Large Language Model, Reasoning, Reinforcement Learning
|
||||
## 1 Introduction
|
||||
Reasoning capability, the cornerstone of human intelligence, enables complex cognitive tasks ranging from mathematical problem-solving to logical deduction and programming. Recent advances in artificial intelligence have demonstrated that large language models (LLMs) can exhibit emergent behaviors, including reasoning abilities, when scaled to a sufficient size ( Wei et al., 2022a ; Kaplan et al., 2020 ) . However, achieving such capabilities in pre-training typically demands substantial computational resources.
|
||||
In parallel, a complementary line of research has demonstrated that large language models can be effectively augmented through chain-of-thought (CoT) prompting. This technique, which involves either providing carefully designed few-shot examples or using minimalistic prompts such as “Let’s think step by step” ( Wei et al., 2022b ; Kojima et al., 2022 ) , enables models to produce intermediate reasoning steps, thereby substantially enhancing their performance on complex tasks.
|
||||
Similarly, further performance gains have been observed when models learn high-quality, multi-step reasoning trajectories during the post-training phase ( OpenAI, 2023 ; Chung et al., 2024 ) .
|
||||
Despite their effectiveness, these approaches exhibit notable limitations. Their dependence on human-annotated reasoning traces hinders scalability and introduces cognitive biases. Furthermore, by constraining models to replicate human thought processes, their performance is inherently capped by the human-provided exemplars, which prevents the exploration of superior, non-human-like reasoning pathways.
|
||||
To tackle these issues, we aim to explore the potential of LLMs for developing reasoning abilities through self-evolution in an RL framework, with minimal reliance on human labeling efforts.
|
||||
Specifically, we build upon DeepSeek-V3-Base ( DeepSeek-AI, 2024b ) and employ Group Relative Policy Optimization (GRPO) ( Shao et al., 2024 ) as our RL framework. The reward signal is solely based on the correctness of final predictions against ground-truth answers, without imposing constraints on the reasoning process itself. Notably, we bypass the conventional supervised fine-tuning (SFT) phase before RL training. This design choice stems from our hypothesis that human-defined reasoning patterns may limit model exploration, whereas unrestricted RL training can better incentivize the emergence of novel reasoning capabilities in LLMs.
|
||||
Through this process, detailed in Section 2 , our model (referred to as DeepSeek-R1-Zero) naturally developed diverse and sophisticated reasoning behaviors.
|
||||
In solving reasoning problems, the model exhibits a tendency to generate longer responses, incorporating verification, reflection, and the exploration of alternative approaches within each response. Although we do not explicitly teach the model how to reason, it successfully learns improved reasoning strategies through reinforcement learning.
|
||||
Although DeepSeek-R1-Zero demonstrates excellent reasoning capabilities, it faces challenges such as poor readability and language mixing, occasionally combining English and Chinese within a single chain-of-thought response. Furthermore, the rule-based RL training stage of DeepSeek-R1-Zero is narrowly focused on reasoning tasks, resulting in limited performance in broader areas such as writing and open-domain question answering.
|
||||
To address these challenges, we introduce DeepSeek-R1, a model trained through a multi-stage learning framework that integrates rejection sampling, reinforcement learning, and supervised fine-tuning, detailed in Section 3 . This training pipeline enables DeepSeek-R1 to inherit the reasoning capabilities of its predecessor, DeepSeek-R1-Zero, while aligning model behavior with human preferences through additional non-reasoning data.
|
||||
To enable broader access to
|
||||
powerful AI at a lower energy cost, we have distilled several smaller models and made them publicly available. These distilled models exhibit strong reasoning capabilities, surpassing the performance of their original instruction-tuned counterparts. We believe that these instruction-tuned versions will also significantly contribute to the research community by providing a valuable resource for understanding the mechanisms underlying long chain-of-thought (CoT) reasoning models and for fostering the development of more powerful reasoning models. We release DeepSeek-R1 series models to the public at https://huggingface.co/deepseek-ai .
|
||||
## 2 DeepSeek-R1-Zero
|
||||
We begin by elaborating on the training of DeepSeek-R1-Zero, which relies exclusively on reinforcement learning without supervised fine-tuning. To facilitate large-scale RL efficiency, we adopt Group Relative Policy Optimization (GRPO) ( Shao et al., 2024 ) .
|
||||
## 2.1 Group Relative Policy Optimization
|
||||
GRPO ( Shao et al., 2024 ) is the reinforcement learning algorithm that we adopt to train DeepSeek-R1-Zero and DeepSeek-R1. It was originally proposed to simplify the training process and reduce the resource consumption of Proximal Policy Optimization (PPO) ( Schulman et al., 2017 ) , which is widely used in the RL stage of LLMs ( Ouyang et al., 2022 ) .
|
||||
For each question q q , GRPO samples a group of outputs { o 1 , o 2 , ⋯ , o G } \{o_{1},o_{2},\cdots,o_{G}\} from the old policy π θ o l d \pi_{\theta_{old}} and then optimizes the policy model π θ \pi_{\theta} by maximizing the following objective:
|
||||
𝒥 G R P O ( θ ) = 𝔼 [ q ∼ P ( Q ) , { o i } i = 1 G ∼ π θ o l d ( O | q ) ] 1 G ∑ i = 1 G ( min ( π θ ( o i | q ) π θ o l d ( o i | q ) A i , clip ( π θ ( o i | q ) π θ o l d ( o i | q ) , 1 − ε , 1 + ε ) A i ) − β 𝔻 K L ( π θ | | π r e f ) ) , \begin{split}\footnotesize&\mathcal{J}_{GRPO}(\theta)=\mathbb{E}{[q\sim P(Q),\{o_{i}\}_{i=1}^{G}\sim\pi_{\theta_{old}}(O|q)]}\\
|
||||
&\frac{1}{G}\sum_{i=1}^{G}\left(\min\left(\frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old}}(o_{i}|q)}A_{i},\text{clip}\left(\frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old}}(o_{i}|q)},1-\varepsilon,1+\varepsilon\right)A_{i}\right)-\beta\mathbb{D}_{KL}\left(\pi_{\theta}||\pi_{ref}\right)\right),\end{split}
|
||||
(1)
|
||||
𝔻 K L ( π θ | | π r e f ) = π r e f ( o i | q ) π θ ( o i | q ) − log π r e f ( o i | q ) π θ ( o i | q ) − 1 , \mathbb{D}_{KL}\left(\pi_{\theta}||\pi_{ref}\right)=\frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)}-\log\frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)}-1,
|
||||
(2)
|
||||
where π r e f \pi_{ref} is a reference policy, ε \varepsilon and β \beta are hyper-parameters, and A i A_{i} is the advantage, computed using a group of rewards { r 1 , r 2 , … , r G } \{r_{1},r_{2},\ldots,r_{G}\} corresponding to the outputs within each group:
|
||||
A i = r i − m e a n ( { r 1 , r 2 , ⋯ , r G } ) s t d ( { r 1 , r 2 , ⋯ , r G } ) . A_{i}=\frac{r_{i}-{\mathrm{m}ean(\{r_{1},r_{2},\cdots,r_{G}\})}}{{\mathrm{s}td(\{r_{1},r_{2},\cdots,r_{G}\})}}.
|
||||
(3)
|
||||
We give a comparison of GRPO and PPO in Supplementary A.3 .
|
||||
To train DeepSeek-R1-Zero, we set the learning rate to 3e-6, the KL coefficient to 0.001, and the sampling temperature to 1 for rollout. For each question, we sample 16 outputs with a maximum length of 32,768 tokens before the 8.2k step and 65,536 tokens afterward. As a result, both the performance and response length of DeepSeek-R1-Zero exhibit a significant jump at the 8.2k step, with training continuing for a total of 10,400 steps, corresponding to 1.6 training epochs. Each training step consists of 32 unique questions, resulting in a training batch size of 512. Every 400 steps, we replace the reference model with the latest policy model. To accelerate training, each rollout generates 8,192 outputs, which are randomly split into 16 mini-batches and trained for only a single inner epoch.
|
||||
Table 1: Template for DeepSeek-R1-Zero. prompt will be replaced with the specific reasoning question during training.
|
||||
A conversation between User and Assistant. The user asks a question, and the Assistant solves it.
|
||||
The assistant first thinks about the reasoning process in the mind and then provides the user with the answer.
|
||||
The reasoning process and answer are enclosed within <think>...</think> and <answer>...</answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer> .
|
||||
User: prompt . Assistant:
|
||||
Our high-performance RL infrastructure is described in Supplementary B.1 , ensuring scalable and efficient training.
|
||||
## 2.2 Reward Design
|
||||
The reward is the source of the training signal, which decides the direction of RL optimization.
|
||||
For DeepSeek-R1-Zero, we employ rule-based rewards to
|
||||
deliver precise feedback for data in mathematical, coding, and logical reasoning domains. Our rule-based reward system mainly consists of two types of rewards: accuracy rewards and format rewards.
|
||||
Accuracy rewards evaluate whether the response is correct. For example, in the case of math problems with deterministic results, the model is required to provide the final answer in a specified format (e.g., within a box), enabling reliable rule-based verification of correctness. Similarly, for code competition prompts, a compiler can be utilized to evaluate the model’s responses against a suite of predefined test cases, thereby generating objective feedback on correctness.
|
||||
Format rewards complement the accuracy reward model by enforcing specific formatting requirements. In particular, the model is incentivized to encapsulate its reasoning process within designated tags, specifically ‘ <think> ’ and ‘ </think> ’. This ensures that the model’s thought process is explicitly delineated, enhancing interpretability and facilitating subsequent analysis.
|
||||
R e w a r d rule = R e w a r d acc + R e w a r d format Reward_{\text{rule}}=Reward_{\text{acc}}+Reward_{\text{format}}
|
||||
(4)
|
||||
The accuracy, reward and format reward are combined with the same weight.
|
||||
Notably, we abstain from applying neural reward models—whether outcome-based or process-based—to reasoning tasks. This decision is predicated on our observation that neural reward models are susceptible to reward hacking during large-scale reinforcement learning. Moreover, retraining such models necessitates substantial computational resources and introduces additional complexity into the training pipeline, thereby complicating the overall optimization process.
|
||||
## 2.3 Incentivize Reasoning Capability in LLMs
|
||||
Specifically, we apply the RL technique on the DeepSeek-V3 base to train DeepSeek-R1-Zero. During training, we design a straightforward template, to require DeepSeek-R1-Zero to first produce a reasoning process, followed by the final answer. We intentionally limit our constraints to this structural format, avoiding any content-specific biases to ensure that we can accurately observe the model’s natural progression during the RL process.
|
||||
Figure 1: (a) AIME accuracy of DeepSeek-R1-Zero during training. AIME takes a mathematical problem as input and a number as output, illustrated in Table 32 . Pass@1 and Cons@16 are described in Supplementary D.1 . The baseline is the average score achieved by human participants in the AIME competition. (b) The average response length of DeepSeek-R1-Zero on the training set during the RL process. DeepSeek-R1-Zero naturally learns to solve reasoning tasks with more thinking time. Note that a training step refers to a single policy update operation.
|
||||
Figure 1 (a) depicts the performance trajectory of DeepSeek-R1-Zero on the AIME 2024 benchmark throughout the RL training process, where the average pass@1 score on AIME 2024 shows a significant increase, jumping from an initial 15.6% to 77.9%.
|
||||
In addition, by leveraging the self-consistency decoding ( Wang et al., 2023c ) , the model’s performance can be further improved, achieving an accuracy of 86.7%. This performance significantly surpasses the average performance across all human competitors.
|
||||
Besides the math competitions, as shown in Figure 10 , DeepSeek-R1-Zero also achieves remarkable performance in coding competitions and graduate-level biology, physics, and chemistry problems.
|
||||
These results underscore the effectiveness of RL in enhancing the reasoning capabilities of large language models.
|
||||
Table 2: An interesting “aha moment” of an intermediate version of DeepSeek-R1-Zero. The model learns to rethink using an anthropomorphic tone. This is also an aha moment for us, allowing us to witness the power and beauty of reinforcement learning.
|
||||
Question: If a > 1 a>1 , then the sum of the real solutions of a − a + x = x \sqrt{a-\sqrt{a+x}}=x is equal to
|
||||
Response: <think>
|
||||
To solve the equation a − a + x = x \sqrt{a-\sqrt{a+x}}=x , let’s start by squaring both ⋯ \cdots
|
||||
( a − a + x ) 2 = x 2 ⟹ a − a + x = x 2 . \left(\sqrt{a-\sqrt{a+x}}\right)^{2}=x^{2}\implies a-\sqrt{a+x}=x^{2}.
|
||||
Rearrange to isolate the inner square root term:
|
||||
( a − x 2 ) 2 = a + x ⟹ a 2 − 2 a x 2 + ( x 2 ) 2 = a + x ⟹ x 4 − 2 a x 2 − x + ( a 2 − a ) = 0 (a-x^{2})^{2}=a+x\implies a^{2}-2ax^{2}+(x^{2})^{2}=a+x\implies x^{4}-2ax^{2}-x+(a^{2}-a)=0
|
||||
…
|
||||
Wait, wait. Wait. That’s an aha moment I can flag here .
|
||||
Let’s reevaluate this step-by-step to identify if the correct sum can be ⋯ \cdots
|
||||
We started with the equation:
|
||||
a − a + x = x \sqrt{a-\sqrt{a+x}}=x
|
||||
First, let’s square both sides:
|
||||
a − a + x = x 2 ⟹ a + x = a − x 2 a-\sqrt{a+x}=x^{2}\implies\sqrt{a+x}=a-x^{2}
|
||||
Next, I could square both sides again, treating the equation:
|
||||
⋯ \cdots
|
||||
…
|
||||
The self-evolution of DeepSeek-R1-Zero exemplifies how RL can autonomously enhance a model’s reasoning capabilities.
|
||||
As shown in Figure 1 (b), DeepSeek-R1-Zero exhibits a steady increase in thinking time throughout training, driven solely by intrinsic adaptation rather than external modifications. Leveraging long CoT, the model progressively refines its reasoning, generating hundreds to thousands of tokens to explore and improve its problem-solving strategies.
|
||||
The increase in thinking time fosters the autonomous development of sophisticated behaviors.
|
||||
Specifically, DeepSeek-R1-Zero increasingly exhibits advanced reasoning strategies such as reflective reasoning and systematic exploration of alternative solutions (see Figure 9 (a) in Supplementary C.2 for details), significantly boosting its performance on verifiable tasks like math and coding.
|
||||
Notably, during training, DeepSeek-R1-Zero exhibits an “aha moment” (Table 2 ), characterized by a sudden increase in the use of the word “wait” during reflections (see Figure 9 (b) in Supplementary C.2 for details). This moment marks a distinct change in reasoning patterns and clearly shows the self-evolution process of DeepSeek-R1-Zero.
|
||||
The self-evolution of DeepSeek-R1-Zero underscores the power and beauty of RL: rather than explicitly teaching the model how to solve a problem, we simply provide it with the right incentives, and it autonomously develops advanced problem-solving strategies.
|
||||
This serves as a reminder of the potential of RL to unlock higher levels of capabilities in LLMs, paving the way for more autonomous and adaptive models in the future.
|
||||
## 3 DeepSeek-R1
|
||||
Figure 2: The multi-stage pipeline of DeepSeek-R1. A detailed background on DeepSeek-V3 Base and DeepSeek-V3 is provided in Supplementary A.1 . The models DeepSeek-R1 Dev1, Dev2, and Dev3 represent intermediate checkpoints within this pipeline.
|
||||
Although DeepSeek-R1-Zero exhibits strong reasoning capabilities, it faces several issues. DeepSeek-R1-Zero struggles with challenges like poor readability, and language mixing, as DeepSeek-V3-Base is trained on multiple languages, especially English and Chinese.
|
||||
To address these issues, we develop DeepSeek-R1, whose pipeline is illustrated in Figure 2 .
|
||||
In the initial stage, we collect thousands of cold-start data that exhibits a conversational, human-aligned thinking process. RL training is then applied to improve the model performance with the conversational thinking process and language consistency. Subsequently, we apply rejection sampling and SFT once more. This stage incorporates both reasoning and non-reasoning datasets into the SFT process, enabling the model to not only excel in reasoning tasks but also demonstrate advanced writing capabilities. To further align the model with human preferences, we implement a secondary RL stage designed to enhance the model’s helpfulness and harmlessness while simultaneously refining its reasoning capabilities.
|
||||
The remainder of this section details the key components of this pipeline: Section 3.1 introduces the Reward Model utilized in our RL stages, and Section 3.2 elaborates on the specific training methodologies and implementation details. Data we used in this stage is detailed in Supplementary B.3 .
|
||||
## 3.1 Model-based Rewards
|
||||
For general data, we resort to reward models to capture human preferences in complex and nuanced scenarios. We build
|
||||
upon the DeepSeek-V3 pipeline and adopt a similar distribution of preference pairs and training prompts. For helpfulness, we focus exclusively on the final summary, ensuring that the
|
||||
assessment emphasizes the utility and relevance of the response to the user while minimizing
|
||||
interference with the underlying reasoning process. For harmlessness, we evaluate the entire
|
||||
response of the model, including both the reasoning process and the summary, to identify and
|
||||
mitigate any potential risks, biases, or harmful content that may arise during the generation
|
||||
process.
|
||||
## Helpful Reward Model
|
||||
Regarding helpful reward model training, we first generate preference pairs by prompting DeepSeek-V3 using the arena-hard prompt format, listed in Supplementary B.2 , where each pair consists of a user query along with two candidate responses. For each preference pair, we query DeepSeek-V3 four times, randomly assigning the responses as either Response A or Response B to mitigate positional bias. The final preference score is determined by averaging the four independent judgments, retaining only those pairs where the score difference ( Δ \Delta ) exceeds 1 to ensure meaningful distinctions. Additionally, to minimize length-related biases, we ensure that the chosen and rejected responses of the whole dataset have comparable lengths. In total, we curated 66,000 data pairs for training the reward model. The prompts used in this dataset are all non-reasoning questions and are sourced either from publicly available open-source datasets or from users who have explicitly consented to share their data for the purpose of model improvement. The architecture of our reward model is consistent with that of DeepSeek-R1, with the addition of a reward head designed to predict scalar preference scores.
|
||||
R e w a r d h e l p f u l = R M h e l p f u l ( R e s p o n s e A , R e s p o n s e B ) Reward_{helpful}=RM_{helpful}(Response_{A},Response_{B})
|
||||
(5)
|
||||
The helpful reward models were trained with a batch size of 256, a learning rate of 6e-6, and for a single epoch over the training dataset. The maximum sequence length during training is set to 8192 tokens, whereas no explicit limit is imposed during reward model inference.
|
||||
## Safety Reward Model
|
||||
To assess and improve model safety, we curated a dataset of 106,000 prompts with model-generated responses annotated as “safe" or “unsafe" according to predefined safety guidelines. Unlike the pairwise loss employed in the helpfulness reward model, the safety reward model was trained using a point-wise methodology to distinguish between safe and unsafe responses. The training hyperparameters are the same as the helpful reward model.
|
||||
R e w a r d s a f e t y = R M s a f e t y ( R e s p o n s e ) Reward_{safety}=RM_{safety}(Response)
|
||||
(6)
|
||||
For general queries, each instance is categorized as belonging to either the safety dataset or the helpfulness dataset. The general reward, R e w a r d G e n e r a l Reward_{General} , assigned to each query corresponds to the respective reward defined within the associated dataset.
|
||||
## 3.2 Training Details
|
||||
## 3.2.1 Training Details of the First RL Stage
|
||||
In the first stage of RL, we set the learning rate to 3e-6, the KL coefficient to 0.001, the GRPO clip ratio ε \varepsilon to 10, and the sampling temperature to 1 for rollout. For each question, we sample 16 outputs with a maximum length of 32,768. Each training step consists of 32 unique questions, resulting in a training batch size of 512 per step. Every 400 steps, we replace the reference model with the latest policy model. To accelerate training, each rollout generates 8,192 outputs, which are randomly split into 16 minibatches and trained for only a single inner epoch. However, to mitigate the issue of language mixing, we introduce a language consistency reward during RL training, which is calculated as the proportion of target language words in the CoT.
|
||||
R e w a r d l a n g u a g e = N u m ( W o r d s t a r g e t ) N u m ( W o r d s ) Reward_{language}=\frac{Num(Words_{target})}{Num(Words)}
|
||||
(7)
|
||||
Although ablation experiments in Supplementary B.6 show that such alignment results in a slight degradation in the model’s performance, this reward aligns with human preferences, making it more readable. We apply the language consistency reward to both reasoning and non-reasoning data by directly adding it to the final reward.
|
||||
Note that the clip ratio plays a crucial role in training. A lower value can lead to the truncation of gradients for a significant number of tokens, thereby degrading the model’s performance, while a higher value may cause instability during training.
|
||||
## 3.2.2 Training Details of the Second RL Stage
|
||||
Specifically, we train the model using a combination of reward signals and diverse prompt distributions. For reasoning data, we follow the methodology outlined in DeepSeek-R1-Zero, which employs rule-based rewards to guide learning in mathematical, coding, and logical reasoning domains. During the training process, we observe that CoT often exhibits language mixing, particularly when RL prompts involve multiple languages.
|
||||
For general data, we utilize reward models to guide training. Ultimately, the integration of reward signals with diverse data distributions enables us to develop a model that not only excels in reasoning but also prioritizes helpfulness and harmlessness. Given a batch of data, the reward can be formulated as
|
||||
R e w a r d \displaystyle Reward
|
||||
= R e w a r d reasoning + R e w a r d general + R e w a r d language \displaystyle=Reward_{\text{reasoning}}+Reward_{\text{general}}+Reward_{\text{language}}
|
||||
(8)
|
||||
where, R e w a r d reasoning \displaystyle\text{where, }Reward_{\text{reasoning}}
|
||||
= R e w a r d rule \displaystyle=Reward_{\text{rule}}
|
||||
(9)
|
||||
R e w a r d general \displaystyle Reward_{\text{general}}
|
||||
= R e w a r d reward_model + R e w a r d format \displaystyle=Reward_{\text{reward\_model}}+Reward_{\text{format}}
|
||||
(10)
|
||||
The second stage of RL retains most of the parameters from the first stage, with the key difference being a reduced temperature of 0.7, as we find that higher temperatures in this stage lead to incoherent generation. The stage comprises a total of 1,700 training steps, during which general instruction data and preference-based rewards are incorporated exclusively in the final 400 steps. We find that more training steps with the model based preference reward signal may lead to reward hacking, which is documented in Supplementary B.5 . The total training cost is listed in Supplementary B.4.4 .
|
||||
Table 3: Experimental results at each stage of DeepSeek-R1. Numbers in bold denote the performance is statistically significant (t − - test with p < 0.01 p<0.01 ).
|
||||
Benchmark (Metric)
|
||||
R1-Zero
|
||||
R1-Dev1
|
||||
R1-Dev2
|
||||
R1-Dev3
|
||||
R1
|
||||
English
|
||||
MMLU (EM)
|
||||
88.8
|
||||
89.1
|
||||
91.2
|
||||
91.0
|
||||
90.8
|
||||
MMLU-Redux (EM)
|
||||
85.6
|
||||
90.0
|
||||
93.0
|
||||
93.1
|
||||
92.9
|
||||
MMLU-Pro (EM)
|
||||
68.9
|
||||
74.1
|
||||
83.8
|
||||
83.1
|
||||
84.0
|
||||
DROP (3-shot F1)
|
||||
89.1
|
||||
89.8
|
||||
91.1
|
||||
88.7
|
||||
92.2
|
||||
IF-Eval (Prompt Strict)
|
||||
46.6
|
||||
71.7
|
||||
72.0
|
||||
78.1
|
||||
83.3
|
||||
GPQA Diamond (Pass@1)
|
||||
75.8
|
||||
66.1
|
||||
70.7
|
||||
71.2
|
||||
71.5
|
||||
SimpleQA (Correct)
|
||||
30.3
|
||||
17.8
|
||||
28.2
|
||||
24.9
|
||||
30.1
|
||||
FRAMES (Acc.)
|
||||
82.3
|
||||
78.5
|
||||
81.8
|
||||
81.9
|
||||
82.5
|
||||
AlpacaEval2.0 (LC-winrate)
|
||||
24.7
|
||||
50.1
|
||||
55.8
|
||||
62.1
|
||||
87.6
|
||||
ArenaHard (GPT-4-1106)
|
||||
53.6
|
||||
77.0
|
||||
73.2
|
||||
75.6
|
||||
92.3
|
||||
Code
|
||||
LiveCodeBench (Pass@1-COT)
|
||||
50.0
|
||||
57.5
|
||||
63.5
|
||||
64.6
|
||||
65.9
|
||||
Codeforces (Percentile)
|
||||
80.4
|
||||
84.5
|
||||
90.5
|
||||
92.1
|
||||
96.3
|
||||
Codeforces (Rating)
|
||||
1444
|
||||
1534
|
||||
1687
|
||||
1746
|
||||
2029
|
||||
SWE Verified (Resolved)
|
||||
43.2
|
||||
39.6
|
||||
44.6
|
||||
45.6
|
||||
49.2
|
||||
Aider-Polyglot (Acc.)
|
||||
12.2
|
||||
6.7
|
||||
25.6
|
||||
44.8
|
||||
53.3
|
||||
Math
|
||||
AIME 2024 (Pass@1)
|
||||
77.9
|
||||
59.0
|
||||
74.0
|
||||
78.1
|
||||
79.8
|
||||
MATH-500 (Pass@1)
|
||||
95.9
|
||||
94.2
|
||||
95.9
|
||||
95.4
|
||||
97.3
|
||||
CNMO 2024 (Pass@1)
|
||||
88.1
|
||||
58.0
|
||||
73.9
|
||||
77.3
|
||||
78.8
|
||||
Chinese
|
||||
CLUEWSC (EM)
|
||||
93.1
|
||||
92.8
|
||||
92.6
|
||||
91.6
|
||||
92.8
|
||||
C-Eval (EM)
|
||||
92.8
|
||||
85.7
|
||||
91.9
|
||||
86.4
|
||||
91.8
|
||||
C-SimpleQA (Correct)
|
||||
66.4
|
||||
58.8
|
||||
64.2
|
||||
66.9
|
||||
63.7
|
||||
## 4 Experiment
|
||||
We evaluate our models on MMLU ( Hendrycks et al., 2021 ) , MMLU-Redux ( Gema et al., 2025 ) , MMLU-Pro ( Wang et al., 2024 ) , C-Eval ( Huang et al., 2023 ) , and CMMLU ( Li et al., 2024 ) , IFEval ( Zhou et al., 2023b ) , FRAMES ( Krishna et al., 2024 ) , GPQA Diamond ( Rein et al., 2023 ) , SimpleQA ( OpenAI, 2024a ) , C-SimpleQA ( He et al., 2024 ) , SWE-Bench Verified ( OpenAI, 2024b ) , Aider ( Gauthier, 2025 ) , LiveCodeBench ( Jain et al., 2024 ) (2024-08 – 2025-01), Codeforces ( Mirzayanov, 2025 ) , Chinese National High School Mathematics Olympiad (CNMO 2024) ( CMS, 2024 ) , and American Invitational Mathematics Examination 2024 (AIME 2024) ( MAA, 2024 ) . The details of these benchmarks are listed in Supplementary D .
|
||||
Table 3 summarizes the performance of DeepSeek-R1 across multiple developmental stages, as outlined in Figure 2 . A comparison between DeepSeek-R1-Zero and DeepSeek-R1 Dev1 reveals substantial improvements in instruction-following, as evidenced by higher scores on the IF-Eval and ArenaHard benchmarks. However, due to the limited size of the cold-start dataset, Dev1 exhibits a partial degradation in reasoning performance compared to DeepSeek-R1-Zero, most notably on the AIME benchmark. In contrast, DeepSeek-R1 Dev2 demonstrates marked performance enhancements on benchmarks that require advanced reasoning skills, including those focused on code generation, mathematical problem solving, and STEM-related tasks. Benchmarks targeting general-purpose tasks, such as AlpacaEval 2.0, show marginal improvement. These results suggest that reasoning-oriented RL considerably enhances reasoning capabilities while exerting limited influence on user preference-oriented benchmarks.
|
||||
DeepSeek-R1 Dev3 integrates both reasoning and non-reasoning datasets into the SFT pipeline, thereby enhancing the model’s proficiency in both reasoning and general language generation tasks. Compared to Dev2, DeepSeek-R1 Dev3 achieves notable performance improvements on AlpacaEval 2.0 and Aider-Polyglot, attributable to the inclusion of large-scale non-reasoning corpora and code engineering datasets. Finally, comprehensive RL training on DeepSeek-R1 Dev3 using mixed reasoning-focused and general-purpose data produced the final DeepSeek-R1. Marginal improvements occurred in code and mathematics benchmarks, as substantial reasoning-specific RL was done in prior stages. The primary advancements in the final DeepSeek-R1 were in general instruction-following and user-preference benchmarks, with AlpacaEval 2.0 improving by 25% and ArenaHard by 17%.
|
||||
In addition, we compare DeepSeek-R1 with other models in Supplementary D.2 . Model safety evaluations are provided in Supplementary D.3 .
|
||||
A comprehensive analysis is provided in Supplementary E , including a comparison with DeepSeek-V3, performance evaluations on both fresh test sets, a breakdown of mathematical capabilities by category, and an investigation of test-time scaling behavior. Supplementary F shows that the strong reasoning capability can be transferred to smaller models.
|
||||
## 5 Ethics and Safety Statement
|
||||
With the advancement in the reasoning capabilities of DeepSeek-R1, we deeply recognize the potential ethical risks. For example, R1 can be subject to jailbreak attacks, leading to the generation of dangerous content such as explosive manufacturing plans, while the enhanced reasoning capabilities enable the model to provide plans with better operational feasibility and executability.
|
||||
Besides, a public model is also vulnerable to further fine-tuning that could compromise inherent safety protections.
|
||||
In Supplementary D.3 , we present a comprehensive safety report from multiple perspectives, including performance on open-source and in-house safety evaluation benchmarks, and safety levels across multiple languages and against jailbreak attacks. These comprehensive safety analyses conclude that the inherent safety level of the DeepSeek-R1 model, compared to other state-of-the-art models, is generally at a moderate level (comparable to GPT-4o (2024-05-13)). Besides, when coupled with the risk control system, the model’s safety level is elevated to a superior standard.
|
||||
## 6 Conclusion, Limitation, and Future Work
|
||||
We present DeepSeek-R1-Zero and DeepSeek-R1, which rely on large-scale RL to incentivize model reasoning behaviors. Our results demonstrate that pre-trained checkpoints inherently possess substantial potential for complex reasoning tasks. We believe that the key to unlocking this potential lies not in large-scale human annotation but in the provision of hard reasoning questions, a reliable verifier, and sufficient computational resources for reinforcement learning. Sophisticated reasoning behaviors, such as self-verification and reflection, appeared to emerge organically during the reinforcement learning process.
|
||||
Even if DeepSeek-R1 achieves frontier results on reasoning benchmarks, it still faces several capability limitations, as outlined below:
|
||||
Structure Output and Tool Use:
|
||||
Currently, the structural output capabilities of DeepSeek-R1 remain suboptimal compared to existing models. Moreover, DeepSeek-R1 cannot leverage tools, such as search engines and calculators, to improve the performance of output. However, as it is not hard to build an RL environment for structure output and tool use, we believe the issue will be addressed in the next version.
|
||||
Token efficiency: Unlike conventional test-time computation scaling approaches, such as majority voting or Monte Carlo Tree Search (MCTS), DeepSeek-R1 dynamically allocates computational resources during inference according to the complexity of the problem at hand. Specifically, it uses fewer tokens to solve simple tasks, while generating more tokens for complex tasks. Nevertheless, there remains room for further optimization in terms of token efficiency, as instances of excessive reasoning—manifested as overthinking—are still observed in response to simpler questions.
|
||||
Language Mixing:
|
||||
DeepSeek-R1 is currently optimized for Chinese and English, which may result in language mixing issues when handling queries in other languages. For instance, DeepSeek-R1 might use English for reasoning and responses, even if the query is in a language other than English or Chinese. We aim to address this limitation in future updates. The limitation may be related to the base checkpoint, DeepSeek-V3-Base, mainly utilizes Chinese and English, so that it can achieve better results with the two languages in reasoning.
|
||||
Prompting Engineering: When evaluating DeepSeek-R1, we observe that it is sensitive to prompts. Few-shot prompting consistently degrades its performance. Therefore, we recommend users directly describe the problem and specify the output format using a zero-shot setting for optimal results.
|
||||
Software Engineering Tasks:
|
||||
Due to the long evaluation times, which impact the efficiency of the RL process, large-scale RL has not been applied extensively in software engineering tasks. As a result, DeepSeek-R1 has not demonstrated a huge improvement over DeepSeek-V3 on software engineering benchmarks. Future versions will address this by implementing rejection sampling on software engineering data or incorporating asynchronous evaluations during the RL process to improve efficiency.
|
||||
Beyond specific capability limitations, the pure RL methodology itself also presents inherent challenges:
|
||||
Reward Hacking: The success of pure RL depends on reliable reward signals. In this study, we ensure reward reliability through a reasoning-domain rule-based reward model (RM). However, such dependable RMs are difficult to construct for certain tasks, such as writing. If the reward signal is assigned by a model instead of predefined rules, it becomes more susceptible to exploitation as training progresses, which means the policy model may find shortcuts to hack the reward model. Consequently, for complex tasks that cannot be effectively evaluated by a reliable reward model, scaling up pure RL methods remains an open challenge.
|
||||
In this work, for tasks that cannot obtain a reliable signal, DeepSeek-R1 uses human annotation to create supervised data, and only conduct RL for hundreds of steps. We hope in the future, a robust reward model can be obtained to address such issues.
|
||||
With the advent of pure RL methods like DeepSeek-R1, the future holds immense potential for solving any task that can be effectively evaluated by a verifier, regardless of its complexity for humans. Machines equipped with such advanced RL techniques are poised to surpass human capabilities in these domains, driven by their ability to optimize performance iteratively through trial and error. However, challenges remain for tasks where constructing a reliable reward model is inherently difficult. In such cases, the lack of a robust feedback mechanism may hinder progress, suggesting that future research should focus on developing innovative approaches to define and refine reward structures for these complex, less verifiable problems.
|
||||
Furthermore, leveraging tools during the reasoning process holds significant promise. Whether it’s utilizing tools like compilers or search engines to retrieve or compute necessary information, or employing external tools—such as biological or chemical reagents, to validate final results in the real world, this integration of tool-augmented reasoning could dramatically enhance the scope and accuracy of machine-driven solutions.
|
||||
## 7 Author List
|
||||
The list of authors is organized by contribution role, with individuals listed alphabetically by their first name within each category. Authors marked with an asterisk (*) are no longer affiliated with our team.
|
||||
Core Contributors :
|
||||
Daya Guo,
|
||||
Dejian Yang,
|
||||
Haowei Zhang,
|
||||
Junxiao Song,
|
||||
Peiyi Wang,
|
||||
Qihao Zhu,
|
||||
Runxin Xu,
|
||||
Ruoyu Zhang,
|
||||
Shirong Ma,
|
||||
Xiao Bi,
|
||||
Xiaokang Zhang,
|
||||
Xingkai Yu,
|
||||
Yu Wu,
|
||||
Z.F. Wu,
|
||||
Zhibin Gou,
|
||||
Zhihong Shao,
|
||||
Zhuoshu Li,
|
||||
Ziyi Gao,
|
||||
Contributions of the Core Authors:
|
||||
Peiyi Wang and Daya Guo jointly demonstrated that outcome-based RL induces the emergence of long Chain-of-Thought patterns in LLMs, achieving breakthrough reasoning capabilities. They contributed equally to the creation of R1-Zero, and their work laid the foundation for R1. Daya Guo also contributed to the RL training stability of MOE models. Junxiao Song proposed the GRPO algorithm, implemented the initial version, and introduced rule-based rewards for math tasks. The GRPO algorithm was subsequently refined by Peiyi Wang and Runxin Xu. Zhibin Gou proposed a large PPO clipping strategy to enhance GRPO performance, demonstrating its significance alongside Zhihong Shao and Junxiao Song.
|
||||
Regarding data iteration, reward design, and evaluation, specific teams led efforts across different domains: Qihao Zhu, Z.F. Wu, and Dejian Yang focused on code tasks; Zhihong Shao, Zhibin Gou, and Junxiao Song focused on math tasks; and Peiyi Wang, Ruoyu Zhang, Runxin Xu, and Yu Wu led efforts for other reasoning and general tasks. Additionally, Qihao Zhu and Zhihong Shao contributed to the data selection strategy for RL training, while Zhuoshu Li and Yu Wu co-led the data labeling efforts for the entire project.
|
||||
On the system side, Xiao Bi, Xingkai Yu, Shirong Ma, Xiaokang Zhang, Haowei Zhang, and Ziyi Gao implemented the RL pipeline, optimizing system efficiency and addressing stability issues in large-scale training. Finally, Zhibin Gou, Daya Guo, and Ruoyu Zhang oversaw the final training phase and monitored the model training dynamics. Zhibin Gou led the development of the R1-distill series.
|
||||
Contributors :
|
||||
Aixin Liu,
|
||||
Bing Xue,
|
||||
Bingxuan Wang,
|
||||
Bochao Wu,
|
||||
Bei Feng,
|
||||
Chengda Lu,
|
||||
Chenggang Zhao,
|
||||
Chengqi Deng,
|
||||
Chong Ruan,
|
||||
Damai Dai,
|
||||
Deli Chen,
|
||||
Dongjie Ji,
|
||||
Erhang Li,
|
||||
Fangyun Lin,
|
||||
Fucong Dai,
|
||||
Fuli Luo*,
|
||||
Guangbo Hao,
|
||||
Guanting Chen,
|
||||
Guowei Li,
|
||||
H. Zhang,
|
||||
Hanwei Xu,
|
||||
Honghui Ding,
|
||||
Huazuo Gao,
|
||||
Hui Qu,
|
||||
Hui Li,
|
||||
Jianzhong Guo,
|
||||
Jiashi Li,
|
||||
Jingchang Chen,
|
||||
Jingyang Yuan,
|
||||
Jinhao Tu,
|
||||
Junjie Qiu,
|
||||
Junlong Li,
|
||||
J.L. Cai,
|
||||
Jiaqi Ni,
|
||||
Jian Liang,
|
||||
Jin Chen,
|
||||
Kai Dong,
|
||||
Kai Hu*,
|
||||
Kaichao You,
|
||||
Kaige Gao,
|
||||
Kang Guan,
|
||||
Kexin Huang,
|
||||
Kuai Yu,
|
||||
Lean Wang,
|
||||
Lecong Zhang,
|
||||
Liang Zhao,
|
||||
Litong Wang,
|
||||
Liyue Zhang,
|
||||
Lei Xu,
|
||||
Leyi Xia,
|
||||
Mingchuan Zhang,
|
||||
Minghua Zhang,
|
||||
Minghui Tang,
|
||||
Mingxu Zhou,
|
||||
Meng Li,
|
||||
Miaojun Wang,
|
||||
Mingming Li,
|
||||
Ning Tian,
|
||||
Panpan Huang,
|
||||
Peng Zhang,
|
||||
Qiancheng Wang,
|
||||
Qinyu Chen,
|
||||
Qiushi Du,
|
||||
Ruiqi Ge,
|
||||
Ruisong Zhang,
|
||||
Ruizhe Pan,
|
||||
Runji Wang,
|
||||
R.J. Chen,
|
||||
R.L. Jin,
|
||||
Ruyi Chen,
|
||||
Shanghao Lu,
|
||||
Shangyan Zhou,
|
||||
Shanhuang Chen,
|
||||
Shengfeng Ye,
|
||||
Shiyu Wang,
|
||||
Shuiping Yu,
|
||||
Shunfeng Zhou,
|
||||
Shuting Pan,
|
||||
S.S. Li,
|
||||
Shuang Zhou,
|
||||
Shaoqing Wu,
|
||||
Shengfeng Ye,
|
||||
Tao Yun,
|
||||
Tian Pei,
|
||||
Tianyu Sun,
|
||||
T. Wang,
|
||||
Wangding Zeng,
|
||||
Wen Liu,
|
||||
Wenfeng Liang,
|
||||
Wenjun Gao,
|
||||
Wenqin Yu*,
|
||||
Wentao Zhang,
|
||||
W.L. Xiao,
|
||||
Wei An,
|
||||
Xiaodong Liu,
|
||||
Xiaohan Wang,
|
||||
Xiaokang Chen,
|
||||
Xiaotao Nie,
|
||||
Xin Cheng,
|
||||
Xin Liu,
|
||||
Xin Xie,
|
||||
Xingchao Liu,
|
||||
Xinyu Yang,
|
||||
Xinyuan Li,
|
||||
Xuecheng Su,
|
||||
Xuheng Lin,
|
||||
X.Q. Li,
|
||||
Xiangyue Jin,
|
||||
Xiaojin Shen,
|
||||
Xiaosha Chen,
|
||||
Xiaowen Sun,
|
||||
Xiaoxiang Wang,
|
||||
Xinnan Song,
|
||||
Xinyi Zhou,
|
||||
Xianzu Wang,
|
||||
Xinxia Shan,
|
||||
Y.K. Li,
|
||||
Y.Q. Wang,
|
||||
Y.X. Wei,
|
||||
Yang Zhang,
|
||||
Yanhong Xu,
|
||||
Yao Li,
|
||||
Yao Zhao,
|
||||
Yaofeng Sun,
|
||||
Yaohui Wang,
|
||||
Yi Yu,
|
||||
Yichao Zhang,
|
||||
Yifan Shi,
|
||||
Yiliang Xiong,
|
||||
Ying He,
|
||||
Yishi Piao,
|
||||
Yisong Wang,
|
||||
Yixuan Tan,
|
||||
Yiyang Ma*,
|
||||
Yiyuan Liu,
|
||||
Yongqiang Guo,
|
||||
Yuan Ou,
|
||||
Yuduan Wang,
|
||||
Yue Gong,
|
||||
Yuheng Zou,
|
||||
Yujia He,
|
||||
Yunfan Xiong,
|
||||
Yuxiang Luo,
|
||||
Yuxiang You,
|
||||
Yuxuan Liu,
|
||||
Yuyang Zhou,
|
||||
Y.X. Zhu,
|
||||
Yanping Huang,
|
||||
Yaohui Li,
|
||||
Yi Zheng,
|
||||
Yuchen Zhu,
|
||||
Yunxian Ma,
|
||||
Ying Tang,
|
||||
Yukun Zha,
|
||||
Yuting Yan,
|
||||
Z.Z. Ren,
|
||||
Zehui Ren,
|
||||
Zhangli Sha,
|
||||
Zhe Fu,
|
||||
Zhean Xu,
|
||||
Zhenda Xie,
|
||||
Zhengyan Zhang,
|
||||
Zhewen Hao,
|
||||
Zhicheng Ma,
|
||||
Zhigang Yan,
|
||||
Zhiyu Wu,
|
||||
Zihui Gu,
|
||||
Zijia Zhu,
|
||||
Zijun Liu*,
|
||||
Zilin Li,
|
||||
Ziwei Xie,
|
||||
Ziyang Song,
|
||||
Zizheng Pan,
|
||||
Zhen Huang,
|
||||
Zhipeng Xu,
|
||||
Zhongyu Zhang,
|
||||
Zhen Zhang,
|
||||
## Appendix
|
||||
## Appendix A Background
|
||||
## A.1 DeepSeek-V3
|
||||
DeepSeek V3 ( DeepSeek-AI, 2024b ) is an advanced open-source LLM developed by DeepSeek. Released in December 2024, DeepSeek V3 represents a significant leap forward in AI innovation, designed to rival leading models like OpenAI’s GPT-4 and Meta’s Llama 3.1, while maintaining remarkable cost efficiency and performance.
|
||||
Built on a Mixture-of-Experts (MoE) architecture, DeepSeek V3 has 671 billion total parameters, with 37 billion activated per token, optimizing both efficiency and capability. It was pre-trained on an expansive dataset of 14.8 trillion high-quality, diverse tokens, followed by supervised fine-tuning and reinforcement learning to enhance its abilities across various domains. The model incorporates innovative features like Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024a ) for efficient inference, an auxiliary-loss-free load-balancing strategy, and Multi-Token Prediction (MTP) ( Gloeckle et al., 2024 ) to boost performance, particularly in tasks like mathematics and coding.
|
||||
For the training data of DeepSeek-V3-Base, we exclusively use plain web pages and e-books, without incorporating any synthetic data. However, we have observed that some web pages contain a significant number of OpenAI-model-generated answers, which may lead the base model to acquire knowledge from other powerful models indirectly. However, we did not intentionally include synthetic data generated by OpenAI during the pre-training cooldown phase; all data used in this phase were naturally occurring and collected through web crawling. The pre-training dataset contains a substantial amount of mathematical and code-related content, indicating that DeepSeek-V3-Base has been exposed to a significant volume of reasoning trace data. This extensive exposure equips the model with the capability to generate plausible solution candidates, from which reinforcement learning can effectively identify and optimize high-quality outputs.
|
||||
We did the data contamination in pre-training as described in Appendix D.1 . The training data of DeepSeek-V3 base are mostly Chinese and English, which might be the cause for DeepSeek-R1-Zero language mixing when the language consistent reward is absent.
|
||||
In this paper, we use the notation DeepSeek-V3-Base as the base model, DeepSeek-V3 as the instructed model. Notably, DeepSeek-R1 and DeepSeek-R1-Zero are trained on top of DeepSeek-V3-Base and DeepSeek-R1 leverages non-reasoning data from DeepSeek-V3 SFT data. DeepSeek-R1-Dev1, DeepSeek-R1-Dev2, DeepSeek-R1-Dev3 are intermediate checkpoints of DeepSeek-R1.
|
||||
## A.2 Conventional Post-Training Paradigm
|
||||
Post-training has emerged as an essential step in refining pre-trained LLMs to meet specific performance goals and align with human expectations. A widely adopted two-stage post-training framework is SFT followed by RL ( Ouyang et al., 2022 ) .
|
||||
Supervised Fine-Tuning refines a pre-trained LLM by training it on a curated dataset of input-output pairs tailored to specific tasks. The process employs a supervised learning objective, typically minimizing cross-entropy loss between the model’s predictions and labeled ground truth ( Brown et al., 2020 ) . For instance, in conversational applications, SFT might utilize dialogue datasets where desired responses are explicitly provided, enabling the model to adapt its outputs to predefined standards ( Radford et al., 2019 ) . SFT offers several compelling benefits. First, it achieves precise task alignment by leveraging high-quality examples, allowing the model to excel in domains such as customer support or technical documentation ( Radford et al., 2019 ) . Second, its reliance on pre-trained weights ensures computational efficiency, requiring fewer resources than training from scratch. Finally, the use of explicit input-output mappings enhances interpretability, as the model’s learning process is directly tied to observable data, minimizing the risk of erratic behavior ( Ouyang et al., 2022 ) . Despite its strengths, the performance of SFT hinges on the quality and diversity of the training dataset; narrow or biased data can impair the model’s ability to generalize to novel contexts ( Brown et al., 2020 ) . Additionally, SFT’s static nature—optimizing for fixed outputs—may fail to capture evolving human preferences or nuanced objectives. The labor-intensive process of curating high-quality datasets further complicates its scalability, as errors or inconsistencies in the data can propagate into the model’s behavior ( Ouyang et al., 2022 ) .
|
||||
Following SFT, Reinforcement Learning further refines the LLM by optimizing its outputs against a reward signal. In this stage, the model interacts with an environment—often a reward model trained on human feedback—and adjusts its behavior to maximize cumulative rewards. A prominent instantiation of this approach is Reinforcement Learning from Human Feedback (RLHF), where the reward function encodes human preferences ( Christiano et al., 2017 ) . RL thus shifts the focus from static supervision to dynamic optimization. Notably, RL reduces the need for extensive annotated resources; while SFT demands a fully labeled dataset for every input-output pair, RL can operate with a smaller set of human evaluations or a trained reward model, even rule-based reward model, significantly lowering the annotation burden.
|
||||
The sequential application of SFT and RL combines their complementary strengths. SFT establishes a robust, task-specific baseline by grounding the model in curated examples, while RL refines this foundation to align with broader, human-centric objectives ( Ouyang et al., 2022 ) . For example, SFT might ensure grammatical accuracy in a dialogue system, while RL optimizes for engagement and brevity, as demonstrated in the development of InstructGPT ( Ouyang et al., 2022 ) . This hybrid approach has proven effective in producing models that are both precise and adaptable.
|
||||
In this study, we demonstrate that the SFT stage may impede a model’s ability to explore and develop effective reasoning strategies. This limitation arises because human-provided responses, which serve as targets during SFT, are not always optimal for model learning; they often omit critical reasoning components such as explicit reflection and verification steps. To address this, DeepSeek-R1-Zero enables direct exploration of reasoning patterns by the model itself, independent of human priors. The reasoning trajectories discovered through this self-exploration are subsequently distilled and used to train other models, thereby promoting the acquisition of more robust and generalizable reasoning capabilities.
|
||||
## A.3 A Comparison of GRPO and PPO
|
||||
Figure 3: Demonstration of PPO and our GRPO. GRPO foregoes the value model, instead
|
||||
estimating the advantages from group scores.
|
||||
Group Relative Policy Optimization (GRPO) ( Shao et al., 2024 ) is the reinforcement learning algorithm that we adopt to train DeepSeek-R1-Zero and DeepSeek-R1. It was originally proposed to simplify the training process and reduce the resource consumption of Proximal Policy Optimization (PPO) ( Schulman et al., 2017 ) , which is widely used in the RL stage of LLMs ( Ouyang et al., 2022 ) . For an overall comparison between GRPO and PPO, see Figure 3 .
|
||||
For each question q q , GRPO samples a group of outputs { o 1 , o 2 , ⋯ , o G } \{o_{1},o_{2},\cdots,o_{G}\} from the old policy π θ o l d \pi_{\theta_{old}} and then optimizes the policy model π θ \pi_{\theta} by maximizing the following objective:
|
||||
𝒥 G R P O ( θ ) = 𝔼 [ q ∼ P ( Q ) , { o i } i = 1 G ∼ π θ o l d ( O | q ) ] 1 G ∑ i = 1 G ( min ( π θ ( o i | q ) π θ o l d ( o i | q ) A i , clip ( π θ ( o i | q ) π θ o l d ( o i | q ) , 1 − ε , 1 + ε ) A i ) − β 𝔻 K L ( π θ | | π r e f ) ) , \begin{split}\footnotesize&\mathcal{J}_{GRPO}(\theta)=\mathbb{E}{[q\sim P(Q),\{o_{i}\}_{i=1}^{G}\sim\pi_{\theta_{old}}(O|q)]}\\
|
||||
&\frac{1}{G}\sum_{i=1}^{G}\left(\min\left(\frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old}}(o_{i}|q)}A_{i},\text{clip}\left(\frac{\pi_{\theta}(o_{i}|q)}{\pi_{\theta_{old}}(o_{i}|q)},1-\varepsilon,1+\varepsilon\right)A_{i}\right)-\beta\mathbb{D}_{KL}\left(\pi_{\theta}||\pi_{ref}\right)\right),\end{split}
|
||||
(11)
|
||||
𝔻 K L ( π θ | | π r e f ) = π r e f ( o i | q ) π θ ( o i | q ) − log π r e f ( o i | q ) π θ ( o i | q ) − 1 , \mathbb{D}_{KL}\left(\pi_{\theta}||\pi_{ref}\right)=\frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)}-\log\frac{\pi_{ref}(o_{i}|q)}{\pi_{\theta}(o_{i}|q)}-1,
|
||||
(12)
|
||||
where π r e f \pi_{ref} is a reference policy, ε \varepsilon and β \beta are hyper-parameters, and A i A_{i} is the advantage, computed using a group of rewards { r 1 , r 2 , … , r G } \{r_{1},r_{2},\ldots,r_{G}\} corresponding to the outputs within each group:
|
||||
A i = r i − m e a n ( { r 1 , r 2 , ⋯ , r G } ) s t d ( { r 1 , r 2 , ⋯ , r G } ) . A_{i}=\frac{r_{i}-{\mathrm{m}ean(\{r_{1},r_{2},\cdots,r_{G}\})}}{{\mathrm{s}td(\{r_{1},r_{2},\cdots,r_{G}\})}}.
|
||||
(13)
|
||||
In contrast, in PPO, the advantage is typically computed by applying the Generalized Advantage Estimation (GAE) ( Schulman et al., 2015 ) , based not only on the rewards but also on a learned value model. Since the value model is usually of similar size as the policy model, it introduces a significant memory and computational overhead. Additionally, the training objective of the value model is to predict the expected cumulative reward from the current position onward, based on the tokens generated from the beginning up to the current position. This is inherently difficult, especially when only the final outcome reward is available. The challenge becomes even more pronounced when training long chain-of-thought reasoning models. As the output length increases, the model is more likely to engage in behaviors such as reflection and revision during generation, meaning that the content initially generated may later be revised or contradicted, which makes it even less feasible to predict the final reward based on a partial response.
|
||||
Another key difference between GRPO and PPO is how Kullback–Leibler (KL) divergence between the trained policy and the reference policy is incorporated into the training process. In GRPO, an unbiased estimator of the KL divergence ( Schulman, 2020 ) is directly added in the loss as in equation 11 , while in PPO the per-token KL penalty is added as a dense reward at each token ( Ouyang et al., 2022 ) . Since the optimization goal of reinforcement learning is to maximize cumulative rewards, PPO’s approach penalizes the cumulative KL divergence, which may implicitly penalize the length of the response and thereby prevent the model’s response length from increasing. In addition, as we may train thousands of steps in the scenario of training long chain-of-thought reasoning models, the trained policy can diverge significantly from the initial reference policy. In order to balance the scope that the training policy can explore and the stability of the training, we periodically update the reference policy to the latest policy during the actual training process.
|
||||
Figure 4: Performance of PPO and GRPO on the MATH task.
|
||||
Figure 4 compares the performance of PPO and GRPO on the MATH task using DeepSeek-Coder-V2-Lite (16B MoE with 2.4B active parameters). Unlike GRPO, PPO requires additional hyperparameter tuning—particularly of the λ \lambda coefficient in GAE—and is highly sensitive to this parameter. When λ \lambda is set to 0.95 (the default value in most open-source PPO implementations), PPO performs considerably worse than GRPO. However, with careful tuning (setting λ \lambda to 1.0), PPO’s performance improves substantially, nearing that of GRPO.
|
||||
While PPO can achieve comparable performance when appropriately tuned, it demands additional computational cost for hyperparameter optimization. Moreover, considering the memory and computational overhead associated with training an additional value model, GRPO presents a more practical alternative, especially when training large-scale models with constrained resources.
|
||||
## Appendix B Training Details
|
||||
## B.1 RL Infrastructure
|
||||
Figure 5: Overview of our RL framework.
|
||||
Conducting RL training on large models places high demands on the infrastructure. Our RL framework is architected with a decoupled and extensible structure to facilitate seamless integration of diverse models and algorithms. Within this framework, we have incorporated both intra-modular and inter-modular optimization techniques, to ensure training efficiency and scalability.
|
||||
Specifically, as depicted in Figure 5 , the framework is partitioned into four distinct modules, each corresponding to a specific phase of the RL pipeline:
|
||||
•
|
||||
Rollout Module: Prompts are loaded from training dataset and uniformly dispatched across multiple vLLM ( Kwon et al., 2023 ) workers, each equipped with the actor model, to sample multiple responses. For DeepSeek-V3 MoE architecture, we implement an expert parallelism strategy across nodes to reduce memory access overhead, and deploy redundant copies of hotspot experts to balance computational loads among different experts. Multi-Token Prediction (MTP) component is also leveraged for self-speculative decoding, significantly accelerating the decoding speed and effectively minimizing the completion time for the longest samples.
|
||||
•
|
||||
Inference Module: This module loads the reward model and reference to perform a forward pass on the samples generated during the rollout phase, thereby obtaining model-based rewards and other essential information.
|
||||
•
|
||||
Rule-based Reward Module: This module computes rule-based rewards for the model-generated responses. A unified interface has been designed to accommodate diverse implementations (e.g., code executor, answer matcher, format checker, etc.). Although this module does not require loading models into GPU memory, its execution tends to be time-consuming. To tackle this issue, an asynchronous scheduling approach is employed to overlap its execution with the Rollout and Inference modules, effectively hiding the associated latency.
|
||||
•
|
||||
Training Module: This module loads the actor model and the critic model (if required), to compute loss and update model parameters. It provides flexible support for a variety of RL algorithms (e.g., PPO, GRPO, DPO, etc.). To minimize computational waste caused by sequence padding and balance the workload across devices, we design the following data packing strategy: first, all data in a global batch is sorted by length and distributed across processes within the data parallel group; subsequently, within each process, the Best-Fit strategy is applied to pack the data into fixed-length chunks with minimal padding; finally, the number of chunks is adjusted to be equal across all processes. Additionally, we have integrated the DualPipe algorithm, utilized in DeepSeek-V3 training, to achieve efficient pipeline parallelism.
|
||||
Notably, upon completion of each module (excluding the Rule-based Reward module), the model instances utilized in that phase are automatically offloaded from VRAM to either system memory or disk storage, thereby freeing up VRAM for the subsequent phase.
|
||||
## B.2 Reward Model Prompt
|
||||
Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt displayed below. You will be given assistant A’s answer and assistant B’s answer. Your job is to evaluate which assistant’s answer is better.
|
||||
Begin your evaluation by generating your own answer to the prompt. You must provide your answers before judging any answers.
|
||||
When evaluating the assistants’ answers, compare both assistants’ answers with your answer. You must identify and correct any mistakes or inaccurate information.
|
||||
Then consider if the assistant’s answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive.
|
||||
Then consider the creativity and novelty of the assistant’s answers when needed. Finally, identify any missing important information in the assistants’ answers that would be beneficial to include when responding to the user prompt.
|
||||
After providing your explanation, you must output only one of the following choices as your final verdict with a label:
|
||||
1. Assistant A is significantly better: [[A ≫ \gg B]]
|
||||
2. Assistant A is slightly better: [[A > > B]]
|
||||
3. Tie, relatively the same: [[A = = B]]
|
||||
4. Assistant B is slightly better: [[B > > A]]
|
||||
5. Assistant B is significantly better: [[B ≫ \gg A]]
|
||||
Example output: M̈y final verdict is tie: [[A = = B]].̈
|
||||
## B.3 Data Recipe
|
||||
## B.3.1 RL Data
|
||||
Table 4:
|
||||
Description of RL Data and Tasks.
|
||||
Data Type
|
||||
# Prompts
|
||||
Question Type
|
||||
Output Ty
|
||||
|
||||
=== BODY (tail) ===
|
||||
specifically measure a model’s ability to reason over and synthesize information from provided sources to generate correct and verifiable facts.
|
||||
PROMPT
|
||||
Here are the relevant Wikipedia articles:
|
||||
url: https:en.wikipedia.orgwikiPresident_of_the_United_States
|
||||
url content: The president of the United States (POTUS) is the head of state and head of government of the United States of America. The president directs the executive branch of the federal government and is the commander-in-chief of the United States Armed Forces.
|
||||
…
|
||||
Based on all the information, answer the query.
|
||||
Query: If my future wife has the same first name as the 15th first lady of the United States’ mother and her surname is the same as the second assassinated president’s mother’s maiden name, what is my future wife’s name?
|
||||
Evaluation
|
||||
===Task===
|
||||
I need your help in evaluating an answer provided by an LLM against a ground truth answer. Your task is to determine if the ground truth answer is present in the LLM’s response. Please analyze the provided data and make a decision.
|
||||
===Instructions===
|
||||
1. Carefully compare the "Predicted Answer" with the "Ground Truth Answer".
|
||||
2. Consider the substance of the answers - look for equivalent information or correct answers.
|
||||
Do not focus on exact wording unless the exact wording is crucial to the meaning.
|
||||
3. Your final decision should be based on whether the meaning and the vital facts of the
|
||||
"Ground Truth Answer" are present in the "Predicted Answer:"
|
||||
===Input Data===
|
||||
- Question: If my future wife has the same first name as the 15th first lady of the United States’ mother and her surname is the same as the second assassinated president’s mother’s maiden name, what is my future wife’s name?
|
||||
- Predicted Answer: …
|
||||
- Ground Truth Answer: Jane Ballou
|
||||
===Output Format===
|
||||
Provide your final evaluation in the following format:
|
||||
Explanation: xxx
|
||||
Decision: "TRUE" or "FALSE"
|
||||
Please proceed with the evaluation.
|
||||
Table 25: Arena-Hard is an open-ended evaluation benchmark specifically designed to assess the capabilities of LLMs. It presents models with challenging, novel, and diverse prompts curated from Chatbot Arena, a continuously evolving, crowd-sourced platform. It focuses on measuring model performance in open-ended tasks, with particular emphasis on coding and mathematics-related prompts.
|
||||
Given the inherently subjective nature of open-ended tasks, where multiple valid responses may exist, the benchmark necessitates the use of an evaluation model to approximate human judgment effectively. Higher evaluation scores suggest that the model is more likely to be favored by human users in real-world scenarios.
|
||||
PROMPT
|
||||
To write an SQL query that selects the top 10 rows in a database and joins to 3 different tables based on a field called "code", you would need to know the names of the tables and the specific fields you want to select. Assuming the tables are named "table1", "table2", and "table3", and you want to select all fields from each table, the query would look like this:
|
||||
“‘sql
|
||||
SELECT *
|
||||
FROM (
|
||||
SELECT *
|
||||
FROM table1
|
||||
LIMIT 10
|
||||
) AS t1
|
||||
JOIN table2 ON t1.code = table2.code
|
||||
JOIN table3 ON t1.code = table3.code;
|
||||
“‘
|
||||
This query first selects the top 10 rows from "table1" and then joins "table2" and "table3" based on the "code" field. If you want to select specific fields from each table, you can replace the asterisks (*) with the field names separated by commas.
|
||||
Evaluation
|
||||
Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt displayed below. You will be given assistant A’s answer and assistant B’s answer. Your job is to evaluate which assistant’s answer is better.
|
||||
Begin your evaluation by generating your own answer to the prompt. You must provide your answers before judging any answers.
|
||||
When evaluating the assistants’ answers, compare both assistants’ answers with your answer. You must identify and correct any mistakes or inaccurate information.
|
||||
Then consider if the assistant’s answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive.
|
||||
Then consider the creativity and novelty of the assistant’s answers when needed. Finally, identify any missing important information in the assistants’ answers that would be beneficial to include when responding to the user prompt.
|
||||
After providing your explanation, you must output only one of the following choices as your final verdict with a label:
|
||||
1. Assistant A is significantly better: [[A > > >> B]]
|
||||
2. Assistant A is slightly better: [[ A > A> B]]
|
||||
3. Tie, relatively the same: [[A = = B]]
|
||||
4. Assistant B is slightly better: [[B > > A]]
|
||||
5. Assistant B is significantly better: [[B > > >> A]]
|
||||
Example output: "My final verdict is tie: [[A = = B]]".
|
||||
Table 26: AlpacaEval 2.0 is an open-ended evaluation dataset, similar in nature to ArenaHard, and leverages an LLM to assess model performance on subjective tasks. However, in contrast to ArenaHard, the prompts in AlpacaEval 2.0 are generally less challenging and only a small subset necessitates the deployment of reasoning capabilities by the evaluated models.
|
||||
PROMPT
|
||||
What are the names of some famous actors that started their careers on Broadway?
|
||||
Evaluation
|
||||
< | i m s t a r t | > <|im_{s}tart|> system
|
||||
You are a highly efficient assistant, who evaluates and selects the best large language model (LLMs) based on the quality of their responses to a given instruction. This process will be used to create a leaderboard reflecting the most accurate and human-preferred answers.
|
||||
< | i m e n d | > <|im_{e}nd|>
|
||||
< | i m s t a r t | > <|im_{s}tart|> user
|
||||
I require a leaderboard for various large language models. I’ll provide you with prompts given to these models and their corresponding outputs. Your task is to assess these responses, and select the model that produces the best output from a human perspective.
|
||||
## Instruction
|
||||
{
|
||||
"instruction": """{instruction}""",
|
||||
}
|
||||
## Model Outputs
|
||||
Here are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier.
|
||||
{
|
||||
{
|
||||
"model_identifier": "m",
|
||||
"output": """{output_1}"""
|
||||
},
|
||||
{
|
||||
"model_identifier": "M",
|
||||
"output": """{output_2}"""
|
||||
}
|
||||
}
|
||||
## Task
|
||||
Evaluate the models based on the quality and relevance of their outputs, and select the model that generated the best output. Answer by providing the model identifier of the best model. We will use your output as the name of the best model, so make sure your output only contains one of the following model identifiers and nothing else (no quotes, no spaces, no new lines, …): m or M.
|
||||
## Best Model Identifier
|
||||
< | i m e n d | > <|im_{e}nd|>
|
||||
Table 27: The CLUEWSC (Chinese Language Understanding Evaluation Benchmark - Winograd Schema Challenge) is a specialized task within the CLUE benchmark suite designed to evaluate a model’s commonsense reasoning and contextual understanding capabilities in Chinese.
|
||||
PROMPT
|
||||
请参考示例的格式，完成最后的测试题。
|
||||
下面是一些示例：
|
||||
他伯父还有许多女弟子，大半是富商财主的外室；这些财翁白天忙着赚钱，怕小公馆里的情妇长日无聊，要不安分，常常叫她们学点玩艺儿消遣。
|
||||
上面的句子中的"她们"指的是
|
||||
情妇
|
||||
耶律克定说到雁北义军时，提起韦大哥，就连声说不可挡、不可挡，似有谈虎色变之味。后来又听说粘罕在云中，特派人厚币卑词，要与‘韦义士修好’。韦大哥斩钉截铁地回绝了，大义凛然，端的是条好汉。如今张孝纯也想结识他，几次三番派儿子张浃上门来厮缠，定要俺引他上雁门山去见韦大哥。
|
||||
上面的句子中的"他"指的是
|
||||
张浃
|
||||
“你何必把这事放在心上？何况你的还不过是手稿，并没有发表出来。”龙点睛越发坦率：“如果发表出来了，那倒也就算了。不过既然没发表出来，我何必还让它飘在外头呢？你给我找一找吧，我要收回。”
|
||||
上面的句子中的"它"指的是
|
||||
手稿
|
||||
这个身材高大，曾经被孙光平拿着菜刀追赶得到处乱窜的年轻人，那天早晨穿上了全新的卡其布中山服，像一个城里来的干部似的脸色红润，准备过河去迎接他的新娘。
|
||||
上面的句子中的"他"指的是
|
||||
年轻人
|
||||
负责接待我们的是两位漂亮的朝鲜女导游，身材高挑，露出比例完美的小腿。一个姓韩，一个姓金，自称小韩和小金。她们的中文说得毫无口音，言谈举止也相当亲切。
|
||||
上面的句子中的"她们"指的是
|
||||
两位漂亮的朝鲜女导游
|
||||
下面是测试题，请在思考结束后（</think>后）用一句话输出答案，不要额外的解释。
|
||||
崩龙珍夫妻康健和美；鞠琴十年前丧偶，两年前重结良缘，现在的老伴是一位以前未曾有过婚史的高级工程师；崩龙珍和鞠琴都尽量避免谈及自己的爱人，也尽量回避提及蒋盈波的亡夫屈晋勇——尽管她们对他都很熟悉；当然也绝不会愚蠢地提出蒋盈波今后是一个人过到底还是再找个老伴的问题来加以讨论，那无论如何还为时过早。
|
||||
上面的句子中的"他"指的是
|
||||
Evaluation
|
||||
Parse the last line in response to judge if the answer equals to ground truth.
|
||||
Table 28: C-EVAL evaluates a model’s breadth and depth of knowledge across 52 diverse academic disciplines, spanning humanities, social sciences, STEM (Science, Technology, Engineering, and Mathematics), and other professional fields (e.g., medicine, law). All question in C-Eval are Chinese.
|
||||
PROMPT
|
||||
以下是中国关于逻辑学考试的单项选择题，请选出其中的正确答案。
|
||||
1991年6月15日，菲律宾吕宋岛上的皮纳图博火山突然大喷发，2000万吨二氧化硫气体冲入平流层，形成的霾像毯子一样盖在地球上空，把部分要照射到地球的阳光反射回太空几年之后，气象学家发现这层霾使得当时地球表面的温度累计下降了0．5℃，而皮纳图博火山喷发前的一个世纪，因人类活动而造成的温室效应已经使地球表面温度升高1℃。某位持“人工气候改造论”的科学家据此认为，可以用火箭弹等方式将二氧化硫充入大气层，阻挡部分阳光，达到地球表面降温的目的。以下哪项如果为真,最能对该科学家的提议构成质疑?___
|
||||
A. 如果利用火箭弹将二氧化硫充入大气层，会导致航空乘客呼吸不适。
|
||||
B. 火山喷发形成的降温效应只是暂时的，经过一段时间温度将再次回升。
|
||||
C. 可以把大气层中的碳取出来存储在地下，减少大气层的碳含量。
|
||||
D. 不论何种方式，“人工气候改造”都将破坏地区的大气层结构。
|
||||
答案： B
|
||||
…
|
||||
新疆的哈萨克人用经过训练的金雕在草原上长途追击野狼。某研究小组为研究金雕的飞行方向和判断野狼群的活动范围，将无线电传导器放置在一只金雕身上进行追踪。野狼为了觅食，其活动范围通常很广。因此，金雕追击野狼的飞行范围通常也很大。然而两周以来，无线电传导器不断传回的信号显示，金雕仅在放飞地3公里的范围内飞行。以下哪项如果为真，最有助于解释上述金雕的行为?___
|
||||
A. 金雕放飞地周边重峦叠嶂，险峻异常。
|
||||
B. 金雕的放飞地2公里范围内有一牧羊草场，成为狼群袭击的目标。
|
||||
C. 由于受训金雕的捕杀，放飞地广阔草原的野狼几乎灭绝了。
|
||||
D. 无线电传导信号仅能在有限的范围内传导。
|
||||
Evaluation
|
||||
Parse the last line in response to judge if the choice equals to ground truth.
|
||||
Table 29: GPQA (Graduate - Level Google - Proof QA Benchmark) is a rigorous evaluation framework designed to measure an LLM’s ability to tackle complex, graduate-level multiple - choice problems in STEM domains—specifically biology, physics, and chemistry.
|
||||
PROMPT
|
||||
Answer the following multiple choice question. The last line of your response should be of the following format: ’ANSWER: $LETTER’ (without quotes) where LETTER is one of ABCD. Think step by step before answering.
|
||||
Two quantum states with energies E1 and E2 have a lifetime of 10 − 9 10^{-9} sec and 10 − 8 10^{-8} sec, respectively. We want to clearly distinguish these two energy levels. Which one of the following options could be their energy difference so that they can be clearly resolved?
|
||||
A) 10 − 9 10^{-9} eV
|
||||
B) 10 − 8 10^{-8} eV
|
||||
C) 10 − 4 10^{-4} eV
|
||||
D) 10 − 11 10^{-11} eV
|
||||
Evaluation
|
||||
Parse the capital letter following “ANSWER: ” in response to judge if the choice equals to ground truth.
|
||||
Table 30: SimpleQA is a factuality evaluation benchmark that measures a model’s ability to answer short, fact-seeking questions with precise, verifiable correctness.
|
||||
PROMPT
|
||||
Who received the IEEE Frank Rosenblatt Award in 2010?
|
||||
Evaluation
|
||||
Your job is to look at a question, a gold target, and a predicted answer, and then assign a grade of either ["CORRECT", "INCORRECT", "NOT_ATTEMPTED"].
|
||||
First, I will give examples of each grade, and then you will grade a new example.
|
||||
The following are examples of CORRECT predicted answers.
|
||||
Question: What are the names of Barack Obama’s children?
|
||||
Gold target: Malia Obama and Sasha Obama
|
||||
Predicted answer 1: sasha and malia obama
|
||||
Predicted answer 2: most people would say Malia and Sasha, but I’m not sure and would have to double check
|
||||
…
|
||||
The following are examples of INCORRECT predicted answers.
|
||||
Question: What are the names of Barack Obama’s children?
|
||||
Gold target: Malia and Sasha
|
||||
Predicted answer 1: Malia.
|
||||
Predicted answer 2: Malia, Sasha, and Susan.
|
||||
…
|
||||
The following are examples of NOT_ATTEMPTED predicted answers.
|
||||
Question: What are the names of Barack Obama’s children?
|
||||
Gold target: Malia and Sasha
|
||||
Predicted answer 1: I don’t know.
|
||||
Predicted answer 2: I need more context about which Obama you are
|
||||
…
|
||||
Also note the following things:
|
||||
…
|
||||
Here is a new example. Simply reply with either CORRECT, INCORRECT, NOT_ATTEMPTED. Don’t apologize or correct yourself if there was a mistake; we are just trying to grade the answer.
|
||||
Question: Who received the IEEE Frank Rosenblatt Award in 2010?
|
||||
Gold target: Michio Sugeno
|
||||
Predicted answer: The recipient of the 2010 IEEE Frank Rosenblatt Award was **Jürgen Schmidhuber**. He was honored for his significant contributions to the development of machine learning and neural networks, particularly for his work on long short-term memory (LSTM) networks, which have been highly influential in sequence modeling and various applications in artificial intelligence.
|
||||
Grade the predicted answer of this new question as one of:
|
||||
A: CORRECT
|
||||
B: INCORRECT
|
||||
C: NOT_ATTEMPTED
|
||||
Just return the letters "A", "B", or "C", with no text around it.
|
||||
Table 31: An example of C-SimpleQA. It measures a model’s ability to answer short, fact-seeking questions in Chinese with precise, verifiable correctness.
|
||||
PROMPT
|
||||
显脉香茶菜可以用来治疗急性的什么类型的黄疸型肝炎？
|
||||
Evaluation
|
||||
请根据给定问题、标准答案和模型预测的答案来评估模型的回答是否正确。您的任务是将结果评定为：【正确】、【错误】或【未尝试】。
|
||||
首先，我们将列出每个评定类别的示例，然后请您对新问题的预测答案进行评定。
|
||||
以下是【正确】的答复示例：
|
||||
问题：贝拉克·奥巴马的孩子叫什么名字？
|
||||
标准答案：玛丽亚·奥巴马和萨莎·奥巴马
|
||||
模型预测1：Malia Obama and Sasha Obama
|
||||
模型预测2：玛丽亚和萨沙
|
||||
…
|
||||
以下是【错误】的答复示例：
|
||||
问题：巴拉克·奥巴马的孩子叫什么名字？
|
||||
标准答案：玛丽亚·奥巴马和萨莎·奥巴马
|
||||
模型预测1：玛丽亚
|
||||
模型预测2：玛丽亚、萨莎和苏珊
|
||||
…
|
||||
以下是【未尝试】的答复示例：
|
||||
问题：巴拉克·奥巴马的孩子叫什么名字？
|
||||
标准答案：玛丽亚·奥巴马和萨莎·奥巴马
|
||||
模型预测1：我不知道。
|
||||
模型预测2：我需要更多关于您所指奥巴马的上下文。
|
||||
…
|
||||
下面是一个新的问题示例。请只回复A、B、C之一，不要道歉或纠正自己的错误，只需要评估该回答。
|
||||
问题: 显脉香茶菜可以用来治疗急性的什么类型的黄疸型肝炎？
|
||||
正确答案: 黄疸型肝炎
|
||||
预测答案: …
|
||||
将此新问题的预测答案评定为以下之一：
|
||||
A:【正确】
|
||||
B:【错误】
|
||||
C:【未尝试】
|
||||
只返回字母"A"、"B"或"C"，无须添加其他文本。
|
||||
Table 32: An example of math evaluation, which applies to AIME, MATH, and CNMO. These benchmarks evaluate model performance on mathematical tasks.
|
||||
PROMPT
|
||||
Let b ⩾ 2 b\geqslant 2 be an integer. Call a positive integer n n b -eautiful b\textit{-eautiful} if it has exactly two digits when expressed in base b b , and these two digits sum to n \sqrt{n} . For example, 81 81 is 13 13 -eautiful because 81 = 6 ¯ 81=\underline{6} 3 ¯ 13 \underline{3}_{13} and 6 + 3 = 81 6+3=\sqrt{81} . Find the least integer b ⩾ 2 b\geqslant 2 for which there are more than ten b b -eautiful integers.
|
||||
Please reason step by step, and put your final answer within \boxed{}.
|
||||
Evaluation
|
||||
Parse the final answer within \boxed{} and use a rule-based grader to determine if it equals the ground truth. Round numerical values as needed, and use ‘SymPy’ 1 1
|
||||
1
|
||||
https://www.sympy.org to parse expressions.
|
||||
@@ -0,0 +1,778 @@
|
||||
TITLE: Janus-Pro: Unified Multimodal Understanding and Generation with Data and Model Scaling
|
||||
ARXIV: 2501.17811
|
||||
AUTHORS: Xiaokang Chen , Zhiyu Wu , Xingchao Liu , Zizheng Pan , Wen Liu , Zhenda Xie , Xingkai Yu , Chong Ruan
|
||||
DATE: [Submitted on 29 Jan 2025]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Janus-Pro: Unified Multimodal Understanding and Generation with Data and Model Scaling
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Method
|
||||
## 2.1 Architecture
|
||||
## 2.2 Optimized Training Strategy
|
||||
## 2.3 Data Scaling
|
||||
## 2.4 Model Scaling
|
||||
## 3 Experiments
|
||||
## 3.1 Implementation Details
|
||||
## 3.2 Evaluation Setup
|
||||
## 3.3 Comparison with State-of-the-arts
|
||||
## 3.4 Qualitative Results
|
||||
## 4 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
Janus-Pro: Unified Multimodal Understanding and Generation with Data and Model Scaling
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Method
|
||||
2.1 Architecture
|
||||
2.2 Optimized Training Strategy
|
||||
2.3 Data Scaling
|
||||
2.4 Model Scaling
|
||||
3 Experiments
|
||||
3.1 Implementation Details
|
||||
3.2 Evaluation Setup
|
||||
3.3 Comparison with State-of-the-arts
|
||||
3.4 Qualitative Results
|
||||
4 Conclusion
|
||||
References
|
||||
License: CC BY 4.0
|
||||
arXiv:2501.17811v1 [cs.AI] 29 Jan 2025
|
||||
\reportnumber
|
||||
001
|
||||
## Janus-Pro: Unified Multimodal Understanding and Generation with Data and Model Scaling
|
||||
Xiaokang Chen, Zhiyu Wu, Xingchao Liu, Zizheng Pan, Wen Liu, Zhenda Xie, Xingkai Yu, Chong Ruan
|
||||
DeepSeek-AI
|
||||
Project Page: https://github.com/deepseek-ai/Janus
|
||||
## Abstract
|
||||
In this work, we introduce Janus-Pro , an advanced version of the previous work Janus. Specifically, Janus-Pro incorporates (1) an optimized training strategy, (2) expanded training data, and (3) scaling to larger model size.
|
||||
With these improvements, Janus-Pro achieves significant advancements in both multimodal understanding and text-to-image instruction-following capabilities, while also enhancing the stability of text-to-image generation. We hope this work will inspire further exploration in the field. Code and models are publicly available.
|
||||
## 1 Introduction
|
||||
(a) Average performance on four multimodal understanding benchmarks.
|
||||
(b) Performance on instruction-following benchmarks for text-to-image generation.
|
||||
Figure 1: Multimodal understanding and visual generation results from our Janus-Pro .
|
||||
For multimodal understand, we average the accuracy of POPE, MME-Perception, GQA, and MMMU. The scores of MME-Perception are divided by 20 to scale to [ 0,100 ] [0,100] . For visual generation, we evaluate the performance on two instruction-following benchamrks, GenEval and DPG-Bench. Overall, Janus-Pro outperforms the previous state-of-the-art unified multimodal models as well as some task-specific models. Best viewed on screen.
|
||||
Figure 2: Comparison of text-to-image generation between Janus-Pro and its predecessor, Janus. Janus-Pro delivers more stable outputs for short prompts, with improved visual quality, richer details, and the ability to generate simple text. The image resolution is 384 × 384 384\times 384 . Best viewed on screen.
|
||||
Recent advancements in unified multimodal understanding and generation models have demonstrated significant progress [ 40 , 48 , 50 , 55 , 46 , 30 , 45 , 54 ] . These approaches have been proven to enhance the instruction-following capabilities in visual generation tasks while reducing model redundancy.
|
||||
Most of these methods utilize the same visual encoder to process inputs for both multimodal understanding and generation tasks. Since the representations required for these two tasks differ, this often results in suboptimal performance in multimodal understanding. To address this issue, Janus [ 46 ] proposes decoupling visual encoding, which alleviates the conflict between multimodal understanding and generation tasks, achieving excellent performance in both tasks.
|
||||
As a pioneering model, Janus is validated at the 1B parameter scale. However, due to the limited amount of training data and the relatively small model capacity, it exhibites certain shortcomings, such as suboptimal performance on short prompts image generation and unstable text-to-image generation quality. In this paper, we introduce Janus-Pro, an enhanced version of Janus that incorporates improvements across three dimensions: training strategies, data, and model size. The Janus-Pro series includes two model sizes: 1B and 7B, demonstrating scalability of the visual encoding decoding method.
|
||||
We evaluate Janus-Pro on multiple benchmarks, and the results reveal its superior multimodal understanding capabilities and significantly improved text-to-image instruction-following performance. Specifically, Janus-Pro-7B achieved a score of 79.2 on the multimodal understanding benchmark MMBench [ 29 ] , surpassing state-of-the-art unified multimodal models such as Janus [ 46 ] (69.4), TokenFlow [ 34 ] (68.9) and MetaMorph [ 42 ] (75.2). Additionally, in the text-to-image instruction-following leaderboard GenEval [ 14 ] , Janus-Pro-7B scores 0.80, outperforming Janus [ 46 ] (0.61), DALL-E 3 (0.67), and Stable Diffusion 3 Medium [ 11 ] (0.74).
|
||||
Figure 3: Architecture of our Janus-Pro. We decouple visual encoding for multimodal understanding and visual generation. “Und. Encoder” and “Gen. Encoder” are abbreviations for “Understanding Encoder” and “Generation Encoder”, respectively. Best viewed on screen.
|
||||
## 2 Method
|
||||
## 2.1 Architecture
|
||||
The architecture of Janus-Pro is shown in Figure 3 , which is the same as Janus [ 46 ] . The core design principle of the overall architecture is to decouple visual encoding for multimodal understanding and generation.
|
||||
We apply independent encoding methods to convert the raw inputs into features, which are then processed by an unified autoregressive transformer. For multimodal understanding, we use the SigLIP [ 53 ] encoder to extract high-dimensional semantic features from images. These features are flattened from a 2 2 -D grid into a 1 1 -D sequence, and an understanding adaptor is used to map these image features into the input space of the LLM.
|
||||
For visual generation tasks, we use the VQ tokenizer from [ 38 ] to convert images into discrete IDs. After the ID sequence is flattened into 1 1 -D, we use a generation adaptor to map the codebook embeddings corresponding to each ID into the input space of the LLM. We then concatenate these feature sequences to form a multimodal feature sequence, which is subsequently fed into the LLM for processing. Apart from the built-in prediction head in the LLM, we also utilize a randomly initialized prediction head for image predictions in the visual generation task. The entire model adheres to an autoregressive framework.
|
||||
## 2.2 Optimized Training Strategy
|
||||
The previous version of Janus employs a three-stage training process. Stage I focuses on training the adaptors and the image head.
|
||||
Stage II handles unified pretraining, during which all components except the understanding encoder and the generation encoder has their parameters updated.
|
||||
Stage III is supervised fine-tuning, building upon Stage II by further unlocking the parameters of the understanding encoder during training. This training strategy has certain issues.
|
||||
In Stage II, Janus divides the training for text-to-image capabilities into two parts following PixArt [ 4 ] . The first part trains on ImageNet [ 9 ] data, using image category names as prompts for text-to-image generation, with the goal of modeling pixel dependence. The second part trains on normal text-to-image data. During implementation, 66.67% of the text-to-image training steps in Stage II are allocated to the first part. However, through further experimentation, we find that this strategy is suboptimal and lead to significant computational inefficiency.
|
||||
To address this issue, we make two modifications.
|
||||
•
|
||||
Longer Training in Stage I : We increase the training steps in Stage I, allowing sufficient training on the ImageNet dataset. Our findings reveals that even with the LLM parameters fixed, the model could effectively model pixel dependence and generate reasonable images based on category names.
|
||||
•
|
||||
Focused Training in Stage II : In Stage II, we drop ImageNet data and directly utilize normal text-to-image data to train the model to generate images based on dense descriptions. This redesigned approach enables Stage II to utilize the text-to-image data more efficiently, resulting in improved training efficiency and overall performance.
|
||||
We also adjust the data ratio in Stage III supervised fine-tuning process across different types of datasets, changing the proportion of multimodal data, pure text data, and text-to-image data from 7:3:10 to 5:1:4. By slightly reducing the proportion of text-to-image data, we observe that this adjustment allows us to maintain strong visual generation capabilities while achieving improved multimodal understanding performance.
|
||||
## 2.3 Data Scaling
|
||||
We scale up the training data used for Janus in both multimodal understanding and visual generation aspects.
|
||||
•
|
||||
Multimodal Understanding . For the Stage II pretraining data, we refer to DeepSeek-VL2 [ 49 ] and add approximately 90 million samples. These include image caption datasets (e.g., YFCC [ 31 ] ), as well as data for table, chart, and document understanding (e.g., Docmatix [ 20 ] ). For the Stage III supervised fine-tuning data, we also incorporate additional datasets from DeepSeek-VL2, such as MEME understanding, Chinese conversational data, and datasets aimed at enhancing dialogue experiences. These additions significantly expanded the model’s capabilities, enriching its ability to handle diverse tasks while improving the overall conversational experience.
|
||||
•
|
||||
Visual Generation . We observe that the real-world data used in the previous version of Janus lacks quality and contains significant noise, which often leads to instability in text-to-image generation, resulting in aesthetically poor outputs. In Janus-Pro, we incorporate approximately 72 million samples of synthetic aesthetic data, bringing the ratio of real to synthetic data to 1:1 during the unified pretraining stage. The prompts for these synthetic data samples are publicly available, such as those in [ 43 ] . Experiments demonstrat that the model converges faster when trained on synthetic data, and the resulting text-to-image outputs are not only more stable but also exhibit significantly improved aesthetic quality.
|
||||
## 2.4 Model Scaling
|
||||
The previous version of Janus validates the effectiveness of visual encoding decoupling using a 1.5B LLM. In Janus-Pro, we scaled the model up to 7B, with the hyperparameters of both the 1.5B and 7B LLMs detailed in Table 1 . We observe that when utilizing a larger-scale LLM, the convergence speed of losses for both multimodal understanding and visual generation improved significantly compared to the smaller model. This finding further validates the strong scalability of this approach.
|
||||
Table 1:
|
||||
Architectural configuration for Janus-Pro . We list the hyperparameters of the architecture.
|
||||
Janus-Pro-1B
|
||||
Janus-Pro-7B
|
||||
Vocabulary size
|
||||
100K
|
||||
100K
|
||||
Embedding size
|
||||
2048
|
||||
4096
|
||||
Context Window
|
||||
4096
|
||||
4096
|
||||
#Attention heads
|
||||
16
|
||||
32
|
||||
#Layers
|
||||
24
|
||||
30
|
||||
Table 2:
|
||||
Detailed hyperparameters for training Janus-Pro . Data ratio refers to the ratio of multimodal understanding data, pure text data, and visual generation data.
|
||||
Janus-Pro-1B
|
||||
Janus-Pro-7B
|
||||
Hyperparameters
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Stage 1
|
||||
Stage 2
|
||||
Stage 3
|
||||
Learning rate
|
||||
1.0 × 10 − 3 1.0\times 10^{-3}
|
||||
1.0 × 10 − 4 1.0\times 10^{-4}
|
||||
4.0 × 10 − 5 4.0\times 10^{-5}
|
||||
1.0 × 10 − 3 1.0\times 10^{-3}
|
||||
1.0 × 10 − 4 1.0\times 10^{-4}
|
||||
4.0 × 10 − 5 4.0\times 10^{-5}
|
||||
LR scheduler
|
||||
Constant
|
||||
Constant
|
||||
Constant
|
||||
Constant
|
||||
Constant
|
||||
Constant
|
||||
Weight decay
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
0.0
|
||||
Gradient clip
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
Optimizer
|
||||
AdamW ( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
AdamW ( β 1 = 0.9 , β 2 = 0.95 \beta_{1}=0.9,\beta_{2}=0.95 )
|
||||
Warm-up steps
|
||||
600
|
||||
5000
|
||||
0
|
||||
600
|
||||
5000
|
||||
0
|
||||
Training steps
|
||||
20K
|
||||
360K
|
||||
80K
|
||||
20K
|
||||
360K
|
||||
40K
|
||||
Batch size
|
||||
256
|
||||
512
|
||||
128
|
||||
256
|
||||
512
|
||||
128
|
||||
Data Ratio
|
||||
1:0:3
|
||||
2:3:5
|
||||
5:1:4
|
||||
1:0:3
|
||||
2:3:5
|
||||
5:1:4
|
||||
## 3 Experiments
|
||||
## 3.1 Implementation Details
|
||||
In our experiments, we utilize DeepSeek-LLM ( 1.5 1.5 B and 7 7 B) [ 3 ] with a maximum supported sequence length of 4096 4096 as the base language model. For the vision encoder used in understanding tasks, we select SigLIP-Large-Patch 16 16 - 384 384 [ 53 ] . The generation encoder has a codebook of size 16,384 16,384 and downsamples images by a factor of 16 16 . Both the understanding adaptor and the generation adaptor are two-layer MLPs. The detailed hyperparameters for each stage are provided in Table 2 . Please note that for Stage II, we employ an early stopping strategy, halting at 270K steps. All images are resized to 384 × 384 384\times 384 pixels.
|
||||
For multimodal understanding data, we resize the long side of the image and pad the short side with the background color (RGB: 127 127 , 127 127 , 127 127 ) to reach 384 384 .
|
||||
For visual generation data, the short side is resized to 384 384 , and the long side is cropped to 384 384 . We use sequence packing during training to improve training efficiency. We mix all data types according to the specified ratios in a single training step. Our Janus-Pro is trained and evaluated using HAI-LLM [ 15 ] , which is a lightweight and efficient distributed training framework built on top of PyTorch. The whole training process took about 9 9 / 14 14 days on a cluster of 16 16 / 32 32 nodes for 1.5 1.5 B/ 7 7 B model, each equipped with 8 8 Nvidia A 100 100 ( 40 40 GB) GPUs.
|
||||
Table 3: Comparison with state-of-the-arts on multimodal understanding benchmarks . “Und.” and “Gen.” denote “understanding” and “generation”, respectively. Models using external pretrained diffusion model are marked with † .
|
||||
Type
|
||||
Model
|
||||
# LLM Params
|
||||
POPE ↑ \uparrow
|
||||
MME-P ↑ \uparrow
|
||||
MMB ↑ \uparrow
|
||||
SEED ↑ \uparrow
|
||||
GQA ↑ \uparrow
|
||||
MMMU ↑ \uparrow
|
||||
MM-Vet ↑ \uparrow
|
||||
Und. Only
|
||||
LLaVA-v 1.5 1.5 -Phi- 1.5 1.5 [ 50 ]
|
||||
1.3 1.3 B
|
||||
84.1 84.1
|
||||
1128.0 1128.0
|
||||
-
|
||||
-
|
||||
56.5 56.5
|
||||
30.7 30.7
|
||||
-
|
||||
MobileVLM [ 6 ]
|
||||
1.4 1.4 B
|
||||
84.5 84.5
|
||||
1196.2 1196.2
|
||||
53.2 53.2
|
||||
-
|
||||
56.1 56.1
|
||||
-
|
||||
-
|
||||
MobileVLM-V2 [ 7 ]
|
||||
1.4 1.4 B
|
||||
84.3 84.3
|
||||
1302.8 1302.8
|
||||
57.7 57.7
|
||||
-
|
||||
59.3 59.3
|
||||
-
|
||||
-
|
||||
MobileVLM [ 6 ]
|
||||
2.7 2.7 B
|
||||
84.9 84.9
|
||||
1288.9 1288.9
|
||||
59.6 59.6
|
||||
-
|
||||
59.0 59.0
|
||||
-
|
||||
-
|
||||
MobileVLM-V2 [ 7 ]
|
||||
2.7 2.7 B
|
||||
84.7 84.7
|
||||
1440.5 1440.5
|
||||
63.2 63.2
|
||||
-
|
||||
61.1 61.1
|
||||
-
|
||||
-
|
||||
LLaVA-Phi [ 56 ]
|
||||
2.7 2.7 B
|
||||
85.0 85.0
|
||||
1335.1 1335.1
|
||||
59.8 59.8
|
||||
-
|
||||
-
|
||||
-
|
||||
28.9 28.9
|
||||
LLaVA [ 27 ]
|
||||
7 7 B
|
||||
76.3 76.3
|
||||
809.6 809.6
|
||||
38.7 38.7
|
||||
33.5 33.5
|
||||
-
|
||||
-
|
||||
25.5 25.5
|
||||
LLaVA-v 1.5 1.5 [ 26 ]
|
||||
7 7 B
|
||||
85.9 85.9
|
||||
1510.7 1510.7
|
||||
64.3 64.3
|
||||
58.6 58.6
|
||||
62.0 62.0
|
||||
35.4 35.4
|
||||
31.1 31.1
|
||||
InstructBLIP [ 8 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
36.0 36.0
|
||||
53.4 53.4
|
||||
49.2 49.2
|
||||
-
|
||||
26.2 26.2
|
||||
Qwen-VL-Chat [ 1 ]
|
||||
7 7 B
|
||||
-
|
||||
1487.5 1487.5
|
||||
60.6 60.6
|
||||
58.2 58.2
|
||||
57.5 57.5
|
||||
-
|
||||
-
|
||||
IDEFICS- 9 9 B [ 19 ]
|
||||
8 8 B
|
||||
-
|
||||
-
|
||||
48.2 48.2
|
||||
-
|
||||
38.4 38.4
|
||||
-
|
||||
-
|
||||
Emu 3 3 -Chat [ 45 ]
|
||||
8 8 B
|
||||
85.2 85.2
|
||||
1244 1244
|
||||
58.5 58.5
|
||||
68.2 68.2
|
||||
60.3 60.3
|
||||
31.6 31.6
|
||||
37.2 37.2
|
||||
InstructBLIP [ 8 ]
|
||||
13 13 B
|
||||
78.9 78.9
|
||||
1212.8 1212.8
|
||||
-
|
||||
-
|
||||
49.5 49.5
|
||||
-
|
||||
25.6 25.6
|
||||
Und. and Gen.
|
||||
DreamLLM † [ 10 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
36.6 36.6
|
||||
LaVIT † [ 18 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
46.8 46.8
|
||||
-
|
||||
-
|
||||
MetaMorph † [ 42 ]
|
||||
8 8 B
|
||||
-
|
||||
-
|
||||
75.2 75.2
|
||||
71.8 71.8
|
||||
-
|
||||
-
|
||||
-
|
||||
Emu † [ 39 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
NExT-GPT † [ 47 ]
|
||||
13 13 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
Show-o-256 [ 50 ]
|
||||
1.3 1.3 B
|
||||
73.8 73.8
|
||||
948.4 948.4
|
||||
-
|
||||
-
|
||||
48.7 48.7
|
||||
25.1 25.1
|
||||
-
|
||||
Show-o-512 [ 50 ]
|
||||
1.3 1.3 B
|
||||
80.0 80.0
|
||||
1097.2 1097.2
|
||||
-
|
||||
-
|
||||
58.0 58.0
|
||||
26.7 26.7
|
||||
-
|
||||
D-Dit [ 24 ]
|
||||
2.0 2.0 B
|
||||
84.0 84.0
|
||||
1124.7 1124.7
|
||||
-
|
||||
-
|
||||
59.2 59.2
|
||||
-
|
||||
-
|
||||
Gemini-Nano-1 [ 41 ]
|
||||
1.8 1.8 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
26.3 26.3
|
||||
-
|
||||
ILLUME [ 44 ]
|
||||
7 7 B
|
||||
88.5 88.5
|
||||
1445.3 1445.3
|
||||
65.1 65.1
|
||||
72.9 72.9
|
||||
− -
|
||||
38.2 38.2
|
||||
37.0 37.0
|
||||
TokenFlow-XL [ 34 ]
|
||||
13 13 B
|
||||
86.8 86.8
|
||||
1545.9 1545.9
|
||||
68.9 68.9
|
||||
68.7 68.7
|
||||
62.7 62.7
|
||||
38.7 38.7
|
||||
40.7 40.7
|
||||
LWM [ 28 ]
|
||||
7 7 B
|
||||
75.2 75.2
|
||||
-
|
||||
-
|
||||
-
|
||||
44.8 44.8
|
||||
-
|
||||
9.6 9.6
|
||||
VILA-U [ 48 ]
|
||||
7 7 B
|
||||
85.8 85.8
|
||||
1401.8 1401.8
|
||||
-
|
||||
59.0 59.0
|
||||
60.8 60.8
|
||||
-
|
||||
33.5 33.5
|
||||
Chameleon [ 40 ]
|
||||
7 7 B
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
22.4 22.4
|
||||
8.3 8.3
|
||||
Janus
|
||||
1.5 1.5 B
|
||||
87.0 87.0
|
||||
1338.0 1338.0
|
||||
69.4 69.4
|
||||
63.7 63.7
|
||||
59.1 59.1
|
||||
30.5 30.5
|
||||
34.3 34.3
|
||||
Janus-Pro-1B
|
||||
1.5 1.5 B
|
||||
86.2 86.2
|
||||
1444.0 1444.0
|
||||
75.5 75.5
|
||||
68.3 68.3
|
||||
59.3 59.3
|
||||
36.3 36.3
|
||||
39.8 39.8
|
||||
Janus-Pro-7B
|
||||
7 7 B
|
||||
87.4 87.4
|
||||
1567.1 1567.1
|
||||
79.2 79.2
|
||||
72.1 72.1
|
||||
62.0 62.0
|
||||
41.0 41.0
|
||||
50.0 50.0
|
||||
## 3.2 Evaluation Setup
|
||||
Multimodal Understanding.
|
||||
To assess multimodal understanding capabilities, we evaluate our model on widely recognized image-based vision-language benchmarks, which include GQA [ 17 ] , POPE [ 23 ] , MME [ 12 ] , SEED [ 21 ] , MMB [ 29 ] , MM-Vet [ 51 ] , and MMMU [ 52 ] .
|
||||
Visual Generation.
|
||||
For evaluating visual generation capabilities, we use GenEval [ 14 ] and DPG-Bench [ 16 ] . GenEval is a challenging benchmark for text-to-image generation, designed to reflect the comprehensive generative abilities of visual generation models by offering a detailed instance-level analysis of their compositional capabilities. DPG-Bench (Dense Prompt Graph Benchmark) is a comprehensive dataset consisting of 1065 lengthy, dense prompts, designed to assess the intricate semantic alignment capabilities of text-to-image models.
|
||||
## 3.3 Comparison with State-of-the-arts
|
||||
Multimodal Understanding Performance.
|
||||
We compare the proposed method with state-of-the-art unified models and understanding-only models in Table 3 .
|
||||
Janus-Pro achieves the overall best results.
|
||||
This can be attributed to decoupling the visual encoding for multimodal understanding and generation, mitigating the conflict between these two tasks.
|
||||
When compared to models with significantly larger sizes, Janus-Pro remains highly competitive. For instance, Janus-Pro-7B outperforms TokenFlow-XL ( 13 13 B) on all benchmarks except GQA.
|
||||
Table 4: Evaluation of text-to-image generation ability on GenEval benchmark . “Und.” and “Gen.” denote “understanding” and “generation”, respectively. Models using external pretrained diffusion model are marked with † .
|
||||
Type
|
||||
Method
|
||||
Single Obj.
|
||||
Two Obj.
|
||||
Counting
|
||||
Colors
|
||||
Position
|
||||
Color Attri.
|
||||
Overall ↑ \uparrow
|
||||
Gen. Only
|
||||
LlamaGen [ 38 ]
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.21 0.21
|
||||
0.58 0.58
|
||||
0.07 0.07
|
||||
0.04 0.04
|
||||
0.32 0.32
|
||||
LDM [ 37 ]
|
||||
0.92 0.92
|
||||
0.29 0.29
|
||||
0.23 0.23
|
||||
0.70 0.70
|
||||
0.02 0.02
|
||||
0.05 0.05
|
||||
0.37 0.37
|
||||
SDv 1.5 1.5 [ 37 ]
|
||||
0.97 0.97
|
||||
0.38 0.38
|
||||
0.35 0.35
|
||||
0.76 0.76
|
||||
0.04 0.04
|
||||
0.06 0.06
|
||||
0.43 0.43
|
||||
PixArt- α \alpha [ 4 ]
|
||||
0.98 0.98
|
||||
0.50 0.50
|
||||
0.44 0.44
|
||||
0.80 0.80
|
||||
0.08 0.08
|
||||
0.07 0.07
|
||||
0.48 0.48
|
||||
SDv 2.1 2.1 [ 37 ]
|
||||
0.98 0.98
|
||||
0.51 0.51
|
||||
0.44 0.44
|
||||
0.85 0.85
|
||||
0.07 0.07
|
||||
0.17 0.17
|
||||
0.50 0.50
|
||||
DALL-E 2 2 [ 35 ]
|
||||
0.94 0.94
|
||||
0.66 0.66
|
||||
0.49 0.49
|
||||
0.77 0.77
|
||||
0.10 0.10
|
||||
0.19 0.19
|
||||
0.52 0.52
|
||||
Emu 3 3 -Gen [ 45 ]
|
||||
0.98 0.98
|
||||
0.71 0.71
|
||||
0.34 0.34
|
||||
0.81 0.81
|
||||
0.17 0.17
|
||||
0.21 0.21
|
||||
0.54 0.54
|
||||
SDXL [ 32 ]
|
||||
0.98 0.98
|
||||
0.74 0.74
|
||||
0.39 0.39
|
||||
0.85 0.85
|
||||
0.15 0.15
|
||||
0.23 0.23
|
||||
0.55 0.55
|
||||
DALL-E 3 3 [ 2 ]
|
||||
0.96 0.96
|
||||
0.87 0.87
|
||||
0.47 0.47
|
||||
0.83 0.83
|
||||
0.43 0.43
|
||||
0.45 0.45
|
||||
0.67 0.67
|
||||
SD3-Medium [ 11 ]
|
||||
0.99
|
||||
0.94
|
||||
0.72
|
||||
0.89
|
||||
0.33
|
||||
0.60
|
||||
0.74 0.74
|
||||
Und. and Gen.
|
||||
SEED-X † [ 13 ]
|
||||
0.97 0.97
|
||||
0.58 0.58
|
||||
0.26 0.26
|
||||
0.80 0.80
|
||||
0.19 0.19
|
||||
0.14 0.14
|
||||
0.49 0.49
|
||||
Show-o [ 50 ]
|
||||
0.95 0.95
|
||||
0.52 0.52
|
||||
0.49 0.49
|
||||
0.82 0.82
|
||||
0.11 0.11
|
||||
0.28 0.28
|
||||
0.53 0.53
|
||||
D-DiT [ 24 ]
|
||||
0.97 0.97
|
||||
0.80 0.80
|
||||
0.54 0.54
|
||||
0.76 0.76
|
||||
0.32 0.32
|
||||
0.50 0.50
|
||||
0.65 0.65
|
||||
LWM [ 28 ]
|
||||
0.93 0.93
|
||||
0.41 0.41
|
||||
0.46 0.46
|
||||
0.79 0.79
|
||||
0.09 0.09
|
||||
0.15 0.15
|
||||
0.47 0.47
|
||||
Transfusion [ 55 ]
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.63 0.63
|
||||
ILLUME [ 44 ]
|
||||
0.99 0.99
|
||||
0.86 0.86
|
||||
0.45 0.45
|
||||
0.71 0.71
|
||||
0.39 0.39
|
||||
0.28 0.28
|
||||
0.61 0.61
|
||||
TokenFlow-XL [ 28 ]
|
||||
0.95 0.95
|
||||
0.60 0.60
|
||||
0.41 0.41
|
||||
0.81 0.81
|
||||
0.16 0.16
|
||||
0.24 0.24
|
||||
0.55 0.55
|
||||
Chameleon [ 40 ]
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.39 0.39
|
||||
Janus [ 46 ]
|
||||
0.97 0.97
|
||||
0.68 0.68
|
||||
0.30 0.30
|
||||
0.84 0.84
|
||||
0.46 0.46
|
||||
0.42 0.42
|
||||
0.61 0.61
|
||||
Janus-Pro-1B
|
||||
0.98 0.98
|
||||
0.82 0.82
|
||||
0.51 0.51
|
||||
0.89 0.89
|
||||
0.65 0.65
|
||||
0.56 0.56
|
||||
0.73 0.73
|
||||
Janus-Pro-7B
|
||||
0.99 0.99
|
||||
0.89 0.89
|
||||
0.59 0.59
|
||||
0.90 0.90
|
||||
0.79 0.79
|
||||
0.66 0.66
|
||||
0.80 0.80
|
||||
Table 5: Performances on DPG-Bench. The methods in this table are all generation-specific models except Janus and Janus-Pro.
|
||||
Method
|
||||
Global
|
||||
Entity
|
||||
Attribute
|
||||
Relation
|
||||
Other
|
||||
Overall ↑ \uparrow
|
||||
SDv1.5 [ 36 ]
|
||||
74.63
|
||||
74.23
|
||||
75.39
|
||||
73.49
|
||||
67.81
|
||||
63.18
|
||||
PixArt- α \alpha [ 4 ]
|
||||
74.97
|
||||
79.32
|
||||
78.60
|
||||
82.57
|
||||
76.96
|
||||
71.11
|
||||
Lumina-Next [ 57 ]
|
||||
82.82
|
||||
88.65
|
||||
86.44
|
||||
80.53
|
||||
81.82
|
||||
74.63
|
||||
SDXL [ 33 ]
|
||||
83.27
|
||||
82.43
|
||||
80.91
|
||||
86.76
|
||||
80.41
|
||||
74.65
|
||||
Playground v2.5 [ 22 ]
|
||||
83.06
|
||||
82.59
|
||||
81.20
|
||||
84.08
|
||||
83.50
|
||||
75.47
|
||||
Hunyuan-DiT [ 25 ]
|
||||
84.59
|
||||
80.59
|
||||
88.01
|
||||
74.36
|
||||
86.41
|
||||
78.87
|
||||
PixArt- Σ \Sigma [ 5 ]
|
||||
86.89
|
||||
82.89
|
||||
88.94
|
||||
86.59
|
||||
87.68
|
||||
80.54
|
||||
Emu3-Gen [ 45 ]
|
||||
85.21
|
||||
86.68
|
||||
86.84
|
||||
90.22
|
||||
83.15
|
||||
80.60
|
||||
DALL-E 3 [ 2 ]
|
||||
90.97
|
||||
89.61
|
||||
88.39
|
||||
90.58
|
||||
89.83
|
||||
83.50
|
||||
SD3-Medium [ 11 ]
|
||||
87.90
|
||||
91.01
|
||||
88.83
|
||||
80.70
|
||||
88.68
|
||||
84.08
|
||||
Janus
|
||||
82.33
|
||||
87.38
|
||||
87.70
|
||||
85.46
|
||||
86.41
|
||||
79.68
|
||||
Janus-Pro-1B
|
||||
87.58
|
||||
88.63
|
||||
88.17
|
||||
88.98
|
||||
88.30
|
||||
82.63
|
||||
Janus-Pro-7B
|
||||
86.90
|
||||
88.90
|
||||
89.40
|
||||
89.32
|
||||
89.48
|
||||
84.19
|
||||
Visual Generation Performance.
|
||||
We report visual generation performance on GenEval and DPG-Bench. As shown in Table 4 , our Janus-Pro-7B obtains 80 80 % overall accuracy on GenEval, which outperforms all the other unified or generation-only methods, e.g., Transfusion [ 55 ] ( 63 63 %) SD3-Medium ( 74 74 %) and DALL-E 3 3 ( 67 67 %). This demonstrates that our approach has better instruction-following capabilities. As shown in Table 5 , Janus-Pro achieves a score of 84.19 84.19 on DPG-Bench, surpassing all other methods. This demonstrates that Janus-Pro excels in following dense instructions for text-to-image generation.
|
||||
## 3.4 Qualitative Results
|
||||
We present results on multimodal understanding in Figure 4 . Janus-Pro exhibits impressive comprehension abilities when handling inputs from various contexts, showcasing its powerful capabilities.
|
||||
We also present some text-to-image generation results in the lower part of Figure 4 .
|
||||
The images generated by Janus-Pro-7B are highly realistic, and despite having a resolution of only 384 × 384 384\times 384 , they still contain a lot of details. For imaginative and creative scenes, Janus-Pro-7B accurately captures the semantic information from the prompts, producing well-reasoned and coherent images.
|
||||
Figure 4: Qualitative results of multimodal understanding and visual generation capability . The model is Janus-Pro-7B and the image output resolution of visual generation is 384 × 384 384\times 384 . Best viewed on screen.
|
||||
## 4 Conclusion
|
||||
This paper introduces improvements to Janus from three aspects: training strategy, data, and model size. These enhancements have led to significant advancements in both multimodal understanding and text-to-image instruction-following capabilities. However, Janus-Pro still has certain limitations.
|
||||
In terms of multimodal understanding, the input resolution is limited to 384 × \times 384, which affects its performance in fine-grained tasks such as OCR. For text-to-image generation, the low resolution, combined with reconstruction losses introduced by the vision tokenizer, results in images that, while rich in semantic content, still lack fine details. For example, small facial regions occupying limited image space may appear under-detailed. Increasing the image resolution could mitigate these issues.
|
||||
|
||||
=== BODY (tail) ===
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,517 @@
|
||||
TITLE: Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||||
ARXIV: 2502.11089
|
||||
AUTHORS: Jingyang Yuan , Huazuo Gao , Damai Dai , Junyu Luo , Liang Zhao , Zhengyan Zhang , Zhenda Xie , Y. X. Wei , Lean Wang , Zhiping Xiao , Yuqing Wang , Chong Ruan , Ming Zhang , Wenfeng Liang , Wangding Zeng
|
||||
DATE: [Submitted on 16 Feb 2025 ( v1 ), last revised 27 Feb 2025 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Rethinking Sparse Attention Methods
|
||||
## 2.1 The Illusion of Efficient Inference
|
||||
## 2.2 The Myth of Trainable Sparsity
|
||||
## 2.3 Native Sparsity as an Imperative
|
||||
## 3 Methodology
|
||||
## 3.1 Background
|
||||
## 3.2 Overall Framework
|
||||
## 3.3 Algorithm Design
|
||||
## 3.3.1 Token Compression
|
||||
## 3.3.2 Token Selection
|
||||
## 3.3.3 Sliding Window
|
||||
## 3.4 Kernel Design
|
||||
## 4 Experiments
|
||||
## 4.1 Pretraining Setup
|
||||
## 4.2 Baselines Methods
|
||||
## 4.3 Performance Comparison
|
||||
## 5 Efficiency Analysis
|
||||
## 5.1 Training Speed
|
||||
## 5.2 Decoding Speed
|
||||
## 6 Discussion
|
||||
## 6.1 Challenges with Alternative Token Selection Strategies
|
||||
## 6.2 Visualization
|
||||
## 7 Related Works
|
||||
## 7.1 Fixed Sparse Pattern
|
||||
## 7.2 Dynamic Token Pruning
|
||||
## 7.3 Query-Aware Selection
|
||||
## 8 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Rethinking Sparse Attention Methods
|
||||
2.1 The Illusion of Efficient Inference
|
||||
2.2 The Myth of Trainable Sparsity
|
||||
2.3 Native Sparsity as an Imperative
|
||||
3 Methodology
|
||||
3.1 Background
|
||||
3.2 Overall Framework
|
||||
3.3 Algorithm Design
|
||||
3.3.1 Token Compression
|
||||
3.3.2 Token Selection
|
||||
3.3.3 Sliding Window
|
||||
3.4 Kernel Design
|
||||
4 Experiments
|
||||
4.1 Pretraining Setup
|
||||
4.2 Baselines Methods
|
||||
4.3 Performance Comparison
|
||||
5 Efficiency Analysis
|
||||
5.1 Training Speed
|
||||
5.2 Decoding Speed
|
||||
6 Discussion
|
||||
6.1 Challenges with Alternative Token Selection Strategies
|
||||
6.2 Visualization
|
||||
7 Related Works
|
||||
7.1 Fixed Sparse Pattern
|
||||
7.2 Dynamic Token Pruning
|
||||
7.3 Query-Aware Selection
|
||||
8 Conclusion
|
||||
References
|
||||
A Example of AIME Results
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2502.11089v2 [cs.CL] 27 Feb 2025
|
||||
\reportnumber
|
||||
001
|
||||
## Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||||
Jingyang Yuan ∗1,2 ,
|
||||
Huazuo Gao 1 ,
|
||||
Damai Dai 1 ,
|
||||
Junyu Luo 2 ,
|
||||
Liang Zhao 1 ,
|
||||
Zhengyan Zhang 1 ,
|
||||
Zhenda Xie 1 ,
|
||||
Y. X. Wei 1 ,
|
||||
Lean Wang 1 ,
|
||||
Zhiping Xiao 3 ,
|
||||
Yuqing Wang 1 ,
|
||||
Chong Ruan 1 ,
|
||||
Ming Zhang 2 ,
|
||||
Wenfeng Liang 1 ,
|
||||
Wangding Zeng 1
|
||||
1 DeepSeek-AI
|
||||
2 Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University, PKU-Anker LLM Lab
|
||||
3 University of Washington
|
||||
{yuanjy, mzhang_cs}@pku.edu.cn ,
|
||||
{zengwangding, wenfeng.liang}@deepseek.com
|
||||
Corresponding author: Contribution during internship at DeepSeek-AI.
|
||||
## Abstract
|
||||
Long-context modeling is crucial for next-generation language models, yet the high computational cost of standard attention mechanisms poses significant computational challenges.
|
||||
Sparse attention offers a promising direction for improving efficiency while maintaining model capabilities.
|
||||
We present NSA, a N atively trainable S parse A ttention mechanism that integrates algorithmic innovations with hardware-aligned optimizations to achieve efficient long-context modeling.
|
||||
NSA employs a dynamic hierarchical sparse strategy, combining coarse-grained token compression with fine-grained token selection to preserve both global context awareness and local precision. Our approach advances sparse attention design with two key innovations: (1) We achieve substantial speedups through arithmetic intensity-balanced algorithm design, with implementation optimizations for modern hardware. (2) We enable end-to-end training, reducing pretraining computation without sacrificing model performance.
|
||||
As shown in Figure 1 , experiments
|
||||
show the model pretrained with NSA maintains or exceeds Full Attention models
|
||||
across general benchmarks, long-context tasks, and instruction-based reasoning.
|
||||
Meanwhile, NSA achieves substantial speedups over Full Attention on 64k-length sequences across decoding, forward propagation, and backward propagation, validating its efficiency throughout the model lifecycle.
|
||||
## 1 Introduction
|
||||
Figure 1: Comparison of performance and efficiency between Full Attention model and our NSA. Left: Despite being sparse, NSA surpasses Full Attention baseline on average across general benchmarks, long-context tasks, and reasoning evaluation. Right: For 64k-length sequence processing, NSA achieves substantial computational speedup compared to Full Attention in all stages: decoding, forward propagation, and backward propagation.
|
||||
The research community increasingly recognizes long-context modeling as a crucial capability for next-generation large language models, driven by diverse real-world applications ranging from in-depth reasoning ( Zelikman et al., 2022 ; DeepSeek-AI, 2025 ) , repository-level code generation ( Zhang et al., 2023a ; Zhang et al., ) and multi-turn autonomous agent systems ( Park et al., 2023 ) .
|
||||
Recent breakthroughs, including OpenAI’s o-series models, DeepSeek-R1 ( DeepSeek-AI, 2025 ) , and Gemini 1.5 Pro ( Google et al., 2024 ) ,
|
||||
enabling models to process entire codebases, lengthy documents, maintain coherent multi-turn conversations over thousands of tokens, and perform complex reasoning across long-range dependencies.
|
||||
However, the high complexity ( Zaheer et al., 2020 ) of vanilla Attention ( Vaswani et al., 2017 ) mechanisms emerges as a critical latency bottleneck as sequence length increases. Theoretical estimates indicate that attention computation with softmax architectures accounts for 70–80% of total latency when decoding 64k-length contexts, underscoring the urgent need for more efficient attention mechanisms.
|
||||
A natural approach to efficient long-context modeling is to take advantage of the inherent sparsity of softmax attention ( Ge et al., 2023 ; Jiang et al., 2023 ) , where selectively computing critical query-key pairs can significantly reduce computational overhead while preserving performance. Recent advances demonstrate this potential through diverse strategies:
|
||||
KV-cache eviction methods ( Zhang et al., 2023b ; Li et al., 2024 ; Zhou et al., 2024 ) , blockwise KV-cache selection methods ( Tang et al., 2024 ; Xiao et al., 2024a ; Gao et al., 2024 ) , and sampling, clustering or hashing-based selection methods ( Chen et al., 2024b ; Liu et al., 2024 ; Desai et al., 2024 ) .
|
||||
Despite these promising strategies, existing sparse attention methods often fall short in practical deployments. Many approaches fail to achieve speedups comparable to their theoretical gains; moreover, most methods lack effective training-time support to fully exploit the sparsity patterns of attention.
|
||||
To address these limitations, the deployment of effective sparse attention must tackle two key challenges: (1) Hardware-aligned inference speedup : Converting theoretical computation reductions into actual speed improvements
|
||||
requires hardware-friendly algorithm design during both prefilling and decoding stages to mitigate memory access and hardware scheduling bottlenecks;
|
||||
(2) Training-aware algorithm design : Enabling end-to-end computation with trainable operators to
|
||||
reduce training costs while maintaining model performance.
|
||||
These requirements are crucial for real-world applications to achieve fast long-context inference or training.
|
||||
When considering both aspects, existing methods still exhibit a noticeable gap.
|
||||
To achieve more effective and efficient sparse attention, we present NSA, a Natively trainable Sparse Attention architecture that integrates hierarchical token modeling. As shown in Figure 2 , NSA reduces per-query computation by organizing keys and values into temporal blocks and processing them through three attention paths: compressed coarse-grained tokens, selectively retained fine-grained tokens, and sliding windows for local contextual information. Then we implement specialized kernels to maximize its practical efficiency.
|
||||
NSA introduces two core innovations corresponding to the key requirements above:
|
||||
(1) Hardware-aligned system: Optimize blockwise sparse attention for Tensor Core utilization and memory access, ensuring balanced arithmetic intensity.
|
||||
(2) Training-aware design: Enable stable end-to-end training through efficient algorithms and backward operators.
|
||||
This optimization enables NSA to support both efficient deployment and end-to-end training.
|
||||
We evaluate NSA through comprehensive experiments on real-world language corpora. Pretraining on a 27B-parameter transformer backbone with 260B tokens, we assess NSA’s performance across general language evaluations, long-context evaluations, and chain-of-thought reasoning evaluation. We further compare the kernel speed on A100 GPUs with optimized Triton ( Tillet et al., 2019 ) implementations. Experimental results demonstrate that NSA achieves comparable or superior performance to full attention baseline, while outperforming existing sparse attention approaches. Additionally, NSA delivers substantial speedups across decoding, forward, and backward stages compared to Full Attention, with the speedup ratio increasing for longer sequences. These results validate that our hierarchical sparse attention design effectively balances model capability and computational efficiency.
|
||||
## 2 Rethinking Sparse Attention Methods
|
||||
Modern sparse attention methods have made significant strides in reducing the theoretical computational complexity of transformer models. However, most approaches predominantly apply sparsity during inference while retaining a pretrained Full Attention backbone, potentially introducing architectural bias that limits their ability to fully exploit sparse attention’s advantages. Before introducing our native sparse architecture, we systematically analyze these limitations through two critical lenses.
|
||||
Figure 2: Overview of NSA’s architecture. Left: The framework processes input sequences through three parallel attention branches: For a given query, preceding keys and values are processed into compressed attention for coarse-grained patterns, selected attention for important token blocks, and sliding attention for local context. Right: Visualization of different attention patterns produced by each branch. Green areas indicate regions where attention scores need to be computed, while white areas represent regions that can be skipped.
|
||||
## 2.1 The Illusion of Efficient Inference
|
||||
Despite achieving sparsity in attention computation, many methods fail to achieve corresponding reductions in inference latency, primarily due to two challenges:
|
||||
Phase-Restricted Sparsity.
|
||||
Methods such as H2O ( Zhang et al., 2023b ) apply sparsity during autoregressive decoding while requiring computationally intensive pre-processing (e.g. attention map calculation, index building) during prefilling. In contrast, approaches like MInference ( Jiang et al., 2024 ) focus solely on prefilling sparsity.
|
||||
These methods fail to achieve acceleration across all inference stages, as at least one phase remains computational costs comparable to Full Attention.
|
||||
The phase specialization reduces the speedup ability of these methods in prefilling-dominated workloads like book summarization and code completion, or decoding-dominated workloads like long chain-of-thought ( Wei et al., 2022 ) reasoning.
|
||||
Incompatibility with Advanced Attention Architecture.
|
||||
Some sparse attention methods fail to adapt to modern decoding efficient architectures like Mulitiple-Query Attention (MQA) ( Shazeer, 2019 ) and Grouped-Query Attention (GQA) ( Ainslie et al., 2023 ) , which significantly reduced the memory access bottleneck during decoding by sharing KV across multiple query heads. For instance, in approaches like Quest ( Tang et al., 2024 ) , each attention head independently selects its KV-cache subset. Although it demonstrates consistent computation sparsity and memory access sparsity in Multi-Head Attention (MHA) models, it presents a different scenario in models based on architectures like GQA, where the memory access volume of KV-cache corresponds to the union of selections from all query heads within the same GQA group. This architectural characteristic means that while these methods can reduce computation operations, the required KV-cache memory access remains relatively high.
|
||||
This limitation forces a critical choice: while some sparse attention methods reduce computation, their scattered memory access pattern conflicts with efficient memory access design from advanced architectures.
|
||||
These limitations arise because many existing sparse attention methods focus on KV-cache reduction or theoretical computation reduction, but struggle to achieve significant latency reduction in advanced frameworks or backends.
|
||||
This motivates us to develop algorithms that combine both advanced architectural and hardware-efficient implementation to fully leverage sparsity for improving model efficiency.
|
||||
## 2.2 The Myth of Trainable Sparsity
|
||||
Our pursuit of native trainable sparse attention is motivated by two key insights from analyzing inference-only approaches:
|
||||
(1) Performance Degradation : Applying sparsity post-hoc forces models to deviate from their pretrained optimization trajectory. As demonstrated by Chen et al. (2024b) , top 20% attention can only cover 70% of the total attention scores, rendering structures like retrieval heads in pretrained models vulnerable to pruning during inference.
|
||||
(2) Training Efficiency Demands :
|
||||
Efficient handling of long-sequence training is crucial for modern LLM development. This includes both pretraining on longer documents to enhance model capacity, and subsequent adaptation phases such as long-context fine-tuning and reinforcement learning. However, existing sparse attention methods primarily target inference, leaving the computational challenges in training largely unaddressed. This limitation hinders the development of more capable long-context models through efficient training. Additionally, efforts to adapt existing sparse attention for training also expose challenges:
|
||||
Non-Trainable Components. Discrete operations in methods like ClusterKV ( Liu et al., 2024 )
|
||||
(includes k-means clustering) and MagicPIG ( Chen et al., 2024b ) (includes SimHash-based selecting) create discontinuities in the computational graph. These non-trainable components prevent gradient flow through the token selection process, limiting the model’s ability to learn optimal sparse patterns.
|
||||
Inefficient Back-propagation. Some theoretically trainable sparse attention methods suffer from practical training inefficiencies. Token-granular selection strategy used in approaches like HashAttention ( Desai et al., 2024 ) leads to the need to load a large number of individual tokens from the KV cache during attention computation.
|
||||
This non-contiguous memory access prevents efficient adaptation of fast attention techniques like FlashAttention, which rely on contiguous memory access and blockwise computation to achieve high throughput.
|
||||
As a result, implementations are forced to fall back to low hardware utilization, significantly degrading training efficiency.
|
||||
## 2.3 Native Sparsity as an Imperative
|
||||
These limitations in inference efficiency and training viability motivate our fundamental redesign of sparse attention mechanisms.
|
||||
We propose NSA, a natively sparse attention framework that addresses both computational efficiency and training requirements.
|
||||
In the following sections, we detail the algorithmic design and operator implementation of NSA.
|
||||
## 3 Methodology
|
||||
Our technical approach spans algorithm design and kernel optimization.
|
||||
In the following subsections, we first introduce the background of our methodology. Then we present the overall framework of NSA, followed by its key algorithmic components. Finally, we detail our hardware-optimized kernel design that maximizes practical efficiency.
|
||||
## 3.1 Background
|
||||
Attention Mechanism
|
||||
is widely used in language modeling where each query token 𝐪 t \mathbf{q}_{t} computes relevance scores against all preceding keys 𝐤 : t \mathbf{k}_{:t} to generate a weighted sum of values 𝐯 : t \mathbf{v}_{:t} . Formally, for an input sequence of length t t , the attention operation is defined as:
|
||||
𝐨 t = Attn ( 𝐪 t , 𝐤 : t , 𝐯 : t ) \mathbf{o}_{t}=\operatorname{Attn}\left(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}\right)
|
||||
(1)
|
||||
where Attn \operatorname{Attn} denotes the attention function:
|
||||
Attn ( 𝐪 t , 𝐤 : t , 𝐯 : t ) = ∑ i = 1 t α t , i 𝐯 i ∑ j = 1 t α t , j , α t , i = e 𝐪 t ⊤ 𝐤 i d k . \operatorname{Attn}\left(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}\right)=\sum_{i=1}^{t}\frac{\alpha_{t,i}\mathbf{v}_{i}}{\sum_{j=1}^{t}\alpha_{t,j}},\quad\alpha_{t,i}=e^{\frac{\mathbf{q}_{t}^{\top}\mathbf{k}_{i}}{\sqrt{d_{k}}}}\,.
|
||||
(2)
|
||||
Here, α t , i \alpha_{t,i} represents the attention weight between 𝐪 t \mathbf{q}_{t} and 𝐤 i \mathbf{k}_{i} , and d k d_{k} is the feature dimension of keys.
|
||||
As sequence length increases, attention computation becomes increasingly dominant in the overall computational cost, presenting significant challenges for long-context processing.
|
||||
Arithmetic Intensity
|
||||
is the ratio of compute operations to memory accesses. It intrinsically shapes algorithm optimization on hardware. Each GPU has a critical arithmetic intensity determined by its peak compute capability and memory bandwidth, calculated as the ratio of these two hardware limits. For computation tasks, arithmetic intensity above this critical threshold becomes compute-bound (limited by GPU FLOPS), while below it becomes memory-bound (limited by memory bandwidth).
|
||||
Specifically for causal self-attention mechanism, during training and prefilling phases, batched matrix multiplications and attention computations exhibit high arithmetic intensity, making these stages compute-bound on modern accelerators. In contrast, auto-regressive decoding becomes memory-bandwidth constrained
|
||||
because it generates one token per forward pass while requiring loading the entire key-value cache,
|
||||
resulting in low arithmetic intensity.
|
||||
This leads to different optimization goals — reducing computation cost during training and prefilling, while reducing memory access during decoding.
|
||||
## 3.2 Overall Framework
|
||||
To leverage the potential of attention with natural sparse pattern, we propose replacing the original key-value pairs 𝐤 : t , 𝐯 : t \mathbf{k}_{:t},\mathbf{v}_{:t} in Equation 1 with a more compact and information-dense set of representation key-value pairs K ~ t , V ~ t \tilde{K}_{t},\tilde{V}_{t} given each query 𝐪 t \mathbf{q}_{t} .
|
||||
Specifically, we formally define the optimized attention output as follows:
|
||||
K ~ t = f K ( 𝐪 t , 𝐤 : t , 𝐯 : t ) , V ~ t = f V ( 𝐪 t , 𝐤 : t , 𝐯 : t ) \tilde{K}_{t}=f_{K}(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}),\quad\tilde{V}_{t}=f_{V}(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t})
|
||||
(3)
|
||||
𝐨 t ∗ = Attn ( 𝐪 t , K ~ t , V ~ t ) \mathbf{o}^{*}_{t}=\operatorname{Attn}\left(\mathbf{q}_{t},\tilde{K}_{t},\tilde{V}_{t}\right)
|
||||
(4)
|
||||
where
|
||||
K ~ t , V ~ t \tilde{K}_{t},\tilde{V}_{t} are dynamically constructed based on the current query 𝐪 t \mathbf{q}_{t} and the contextual memory 𝐤 : t , 𝐯 : t \mathbf{k}_{:t},\mathbf{v}_{:t} .
|
||||
We can design various mapping strategies to get different categories of K ~ t c , V ~ t c \tilde{K}_{t}^{c},\tilde{V}_{t}^{c} , and combine them as follows:
|
||||
𝐨 t ∗ = ∑ c ∈ 𝒞 g t c ⋅ Attn ( 𝐪 t , K ~ t c , V ~ t c ) . \mathbf{o}^{*}_{t}=\sum_{c\in\mathcal{C}}g_{t}^{c}\cdot\text{Attn}(\mathbf{q}_{t},\tilde{K}_{t}^{c},\tilde{V}_{t}^{c}).
|
||||
(5)
|
||||
As illustrated in Figure 2 , NSA have three mapping strategies 𝒞 = { cmp , slc , win } \mathcal{C}=\{\text{cmp},\text{slc},\text{win}\} , representing compression, selection, and sliding window for keys and values.
|
||||
g t c ∈ [ 0 , 1 ] g_{t}^{c}\in[0,1] is the gate score for corresponding strategy c c , derived from input features via an MLP and sigmoid activation.
|
||||
Let N t N_{t} denote the total number of remapped keys/values:
|
||||
N t = ∑ c ∈ 𝒞 size [ K ~ t c ] . N_{t}=\sum_{c\in\mathcal{C}}\text{size}[\tilde{K}^{c}_{t}].
|
||||
(6)
|
||||
We maintain a high sparsity ratio by ensuring N t ≪ t {N_{t}}{}\ll t .
|
||||
## 3.3 Algorithm Design
|
||||
In this subsection, we introduce the design of our remapping strategies f K f_{K} and f V f_{V} : token compression, token selection, and sliding window.
|
||||
## 3.3.1 Token Compression
|
||||
By aggregating sequential blocks of keys or values into block-level representations, we obtain compressed keys and values that capture the information of the entire block.
|
||||
Formally, the compressed key representation is defined as:
|
||||
K ~ t cmp = f K cmp ( 𝐤 : t ) = { φ ( 𝐤 i d + 1 : i d + l ) | 0 ⩽ i ⩽ ⌊ t − l d ⌋ } \tilde{K}^{\text{cmp}}_{t}=f_{K}^{\text{cmp}}(\mathbf{k}_{:t})=\left\{\varphi(\mathbf{k}_{id+1:id+l})\middle|0\leqslant i\leqslant\left\lfloor\frac{t-l}{d}\right\rfloor\right\}
|
||||
(7)
|
||||
where l l is the block length, d d is the sliding stride between adjacent blocks, and
|
||||
φ \varphi is a learnable MLP with intra-block position encoding to map keys in a block to a single compressed key.
|
||||
K ~ t cmp ∈ ℝ d k × ⌊ t − l d ⌋ \tilde{K}_{t}^{\text{cmp}}\in\mathbb{R}^{d_{k}\times\left\lfloor\frac{t-l}{d}\right\rfloor} is tensor composed by compresion keys.
|
||||
Usually, we adopt d < l d<l to mitigate information fragmentation.
|
||||
An analogous formulation holds for the compressed value representation V ~ t cmp \tilde{V}_{t}^{\text{cmp}} .
|
||||
Compressed representations capture coarser-grained higher-level semantic information and reduce computational burden of attention.
|
||||
## 3.3.2 Token Selection
|
||||
Using only compressed keys, values might lose important fine-grained information, motivating us to selectively preserve individual keys, values. Below we describe our efficient token selection mechanism that identifies and preserves the most relevant tokens with low computational overhead.
|
||||
Blockwise Selection.
|
||||
Our selection strategy processes key and value sequences in spacial continuous blocks, motivated by two key factors: hardware efficiency considerations and inherent distribution patterns of attention scores. Blockwise selection is crucial to achieve efficient computation on modern GPUs. That is because modern GPU architectures exhibit significantly higher throughput for continuous block accesses compared to random index-based reads. Also, blockwise computation enables optimal utilization of Tensor Cores. This architectural characteristic has established blockwise memory access and computation as a fundamental principle in high-performance attention implementations, as exemplified by FlashAttention’s block-based design. Blockwise selection follows the inherent distribution patterns of attention scores. Prior works ( Jiang et al., 2024 ) have shown that attention scores often exhibit spatial continuity, suggesting that neighboring keys tend to share similar importance levels. Our visualization in Section 6.2 also shows this spatial continuous pattern.
|
||||
To implement blockwise selection, we first divide key, value sequences into selection blocks. To identify the most important blocks for attention computation, we need to assign importance scores to each block. Below we present our method for computing these block-level importance scores.
|
||||
Importance Score Computation.
|
||||
Computing block importance scores could introduce significant overhead. Fortunately, the attention computation of compression tokens produces intermediate attention scores that we can leverage to induce selection block importance scores, formulated as:
|
||||
𝐩 t cmp = Softmax ( 𝐪 t T K ~ t cmp ) , \mathbf{p}_{t}^{\text{cmp}}=\operatorname{Softmax}\left(\mathbf{q}_{t}^{T}\tilde{K}_{t}^{\text{cmp}}\right),
|
||||
(8)
|
||||
where 𝐩 t cmp ∈ ℝ ⌊ t − l d ⌋ + 1 \mathbf{p}_{t}^{\text{cmp}}\in\mathbb{R}^{\left\lfloor\frac{t-l}{d}\right\rfloor+1} is the attention scores between q t q_{t} and compression keys K ~ t cmp \tilde{K}_{t}^{\text{cmp}} .
|
||||
Let l ′ l^{\prime} denote the selection block size.
|
||||
When compression blocks and selection blocks share the same blocking scheme, i.e., l ′ = l = d l^{\prime}=l=d ,
|
||||
we can directly obtain the selection block importance scores 𝐩 t slc \mathbf{p}_{t}^{\text{slc}} by 𝐩 t slc = 𝐩 t cmp \mathbf{p}_{t}^{\text{slc}}=\mathbf{p}_{t}^{\text{cmp}} straightforwardly.
|
||||
For cases where the blocking schemes differ, we derive the importance scores for selection blocks according to their spatial relationship. Given l ⩽ l ′ , l\leqslant l^{\prime}, d | l d\mid l and d | l ′ d\mid l^{\prime} , we have:
|
||||
𝐩 t slc [ j ] = ∑ m = 0 l ′ d − 1 ∑ n = 0 l d − 1 𝐩 t cmp [ l ′ d j − m − n ] , \mathbf{p}_{t}^{\text{slc}}[j]=\sum_{m=0}^{\frac{l^{\prime}}{d}-1}\sum_{n=0}^{\frac{l}{d}-1}\mathbf{p}_{t}^{\text{cmp}}\left[\frac{l^{\prime}}{d}j-m-n\right],
|
||||
(9)
|
||||
where [ ⋅ ] [\cdot] denotes the indexing operator for accessing vector element.
|
||||
For models employing GQA or MQA where key-value caches are shared across query heads, consistent block selection across these heads has to be ensured to minimize KV cache loading during decoding. The shared importance scores across heads in a group are formally defined as:
|
||||
𝐩 t slc ′ = ∑ h = 1 H 𝐩 t slc , ( h ) , {\mathbf{p}_{t}^{\text{slc}}}^{\prime}=\sum_{h=1}^{H}\mathbf{p}_{t}^{\text{slc},(h)},
|
||||
(10)
|
||||
where ( h ) (h) in the superscript denotes the head index, and H H is the number of query heads in each group. This aggregation ensures consistent block selection across heads within the same group.
|
||||
Top- n \boldsymbol{n} Block Selection. After obtaining the selection block importance scores,
|
||||
We retain tokens within the top- n n sparse blocks ranked by block importance scores, formulated as:
|
||||
ℐ t = { i ∣ rank ( 𝐩 t slc ′ [ i ] ) ⩽ n } \mathcal{I}_{t}=\{i\mid\text{rank}({\mathbf{p}_{t}^{\text{slc}}}^{\prime}[i])\leqslant n\}
|
||||
(11)
|
||||
K ~ t slc = Cat [ { 𝐤 i l ′ + 1 : ( i + 1 ) l ′ | i ∈ ℐ t } ] , \tilde{K}^{\text{slc}}_{t}=\operatorname{Cat}\left[\{\mathbf{k}_{il^{\prime}+1:(i+1)l^{\prime}}|i\in\mathcal{I}_{t}\}\right],
|
||||
(12)
|
||||
where rank ( ⋅ ) (\cdot) denotes the ranking position in descending order, with rank = 1 corresponding to the highest score, ℐ t \mathcal{I}_{t} is the set of selected blocks’ indices,
|
||||
Cat \operatorname{Cat} denotes the concatenation operation.
|
||||
K ~ t slc ∈ ℝ d k × n l ′ \tilde{K}_{t}^{\text{slc}}\in\mathbb{R}^{d_{k}\times nl^{\prime}} is tensor composed by compresion keys.
|
||||
An analogous formulation applies to the fine-grained value V ~ t slc \tilde{V}^{\text{slc}}_{t} .
|
||||
The selected keys and values then participate in the attention computation with 𝐪 t \mathbf{q}_{t} as defined in Equation 5 .
|
||||
## 3.3.3 Sliding Window
|
||||
In attention mechanisms, local patterns typically adapt faster and can dominate the learning process, potentially preventing the model from effectively learning from compression and selection tokens. To address this issue, we introduce a dedicated sliding window branch that explicitly handles local context, allowing other branches (compression and selection) to focus on learning their respective features without being shortcutted by local patterns. Specifically, we maintain recent tokens K ~ t win = 𝐤 t − w : t , V ~ t win = 𝐯 t − w : t \tilde{K}_{t}^{\text{win}}=\mathbf{k}_{t-w:t},\tilde{V}_{t}^{\text{win}}=\mathbf{v}_{t-w:t} in a window w w , and isolate attention computations of different information sources (compression tokens, and selected tokens, sliding window) into separate branches. These branch outputs are then aggregated through a learned gating mechanism.
|
||||
To further prevent shortcut learning across attention branches with marginal computational overhead, we provide independent keys and values for three branches.
|
||||
This architectural design enables stable learning by preventing gradient interference between local and long-range pattern recognition, while introducing minimal overhead.
|
||||
After obtaining all three categories of keys and values ( K ~ t cmp , V ~ t cmp \tilde{K}_{t}^{\text{cmp}},\tilde{V}_{t}^{\text{cmp}} ; K ~ t slc , V ~ t slc \tilde{K}_{t}^{\text{slc}},\tilde{V}_{t}^{\text{slc}} ; and K ~ t win , V ~ t win \tilde{K}_{t}^{\text{win}},\tilde{V}_{t}^{\text{win}} ), we compute the final attention output following Equation 5 .
|
||||
Together with the compression, selection, and sliding window mechanisms described above, this forms the complete algorithmic framework of NSA.
|
||||
## 3.4 Kernel Design
|
||||
To achieve FlashAttention-level speedup during the training and prefilling, we implement hardware-aligned sparse attention kernels upon Triton. Given MHA is memory-intensive and inefficient for decoding, we focus on architectures with shared KV caches like GQA and MQA following the current state-of-the-art LLMs. While compression and sliding window attention computations are readily compatible with existing FlashAttention-2 kernels, we introduce the specialized kernel design for sparse selection attention.
|
||||
If we were to follow FlashAttention’s strategy of loading temporally continuous query blocks into SRAM, it would result in inefficient memory access since queries within a block may require disjoint KV blocks. To address this, our key optimization lies in a different query grouping strategy: for each position on the query sequence, we load all query heads within a GQA group (they share the same sparse KV blocks) into SRAM. Figure 3 illustrates our forward pass implementation.
|
||||
The proposed kernel architecture is characterized by the following key features:
|
||||
1.
|
||||
Group-Centric Data Loading . For each inner loop, load all heads’ queries Q ∈ ℝ [ h , d k ] Q\in\mathbb{R}^{[h,d_{k}]} in the group at position t t and their shared sparse key/value block indices ℐ t \mathcal{I}_{t} .
|
||||
2.
|
||||
Shared KV Fetching . In the inner loop, Sequentially load continuous key/value blocks indexed by ℐ t \mathcal{I}_{t} into SRAM as K ∈ ℝ [ B k , d k ] , V ∈ ℝ [ B k , d v ] K\in\mathbb{R}^{[B_{k},d_{k}]},V\in\mathbb{R}^{[B_{k},d_{v}]} to minimize memory loading, where B k B_{k} is the kernel block size satisfying B k | l ′ B_{k}|l^{\prime} .
|
||||
3.
|
||||
Outer Loop on Grid . Since the inner-loop length (proportional to the selected block count n n ) remains nearly identical for different query blocks, we put query/output loops in Triton’s grid scheduler to simplify and optimize the kernel.
|
||||
This design achieves near-optimal arithmetic intensity by (1) eliminating redundant KV transfers through group-wise sharing, and (2) balancing compute workloads across GPU streaming multiprocessors.
|
||||
Figure 3: Kernel design for NSA. The kernel loads queries by GQA groups (Grid Loop), fetches corresponding sparse KV blocks (Inner Loop), and performs attention computation on SRAM. Green blocks indicate data on SRAM, while blue indicates data on HBM.
|
||||
## 4 Experiments
|
||||
We evaluate NSA through three lenses: (1) general benchmarks performance, (2) long-context benchmarks performance, and (3) chain-of-thought reasoning performance, comparing against Full Attention baseline and state-of-the-art sparse attention methods. We defer the efficiency analysis of our sparse computation paradigm to Section 5 , where we provide detailed discussions on training and inference speed.
|
||||
## 4.1 Pretraining Setup
|
||||
Following the common practice in state-of-the-art LLMs, our experiments adopt a backbone combining Grouped-Query Attention (GQA) and Mixture-of-Experts (MoE), featuring 27 B 27\text{B} total parameters with 3 B 3\text{B} active parameters. The model consists of 30 layers with a hidden dimension of 2560. For GQA, we set the number of groups to 4, with a total of 64 attention heads. For each head, the hidden dimensions of the query, key, and value are configured as d q = d k = 192 d_{q}=d_{k}=192 and d v = 128 d_{v}=128 , respectively. For MoE, we utilize the DeepSeekMoE ( Dai et al., 2024 ; DeepSeek-AI, 2024 ) structure, with 72 routed experts and 2 shared experts, and set the top-k experts to 6. To ensure training stability, the MoE in the first layer is replaced by an MLP in the form of SwiGLU. The proposed architecture achieves an effective trade-off between computation cost and model performance.
|
||||
For NSA, we set compression block size l = 32 l=32 , sliding stride d = 16 d=16 , selected block size l ′ = 64 l^{\prime}=64 , selected block count n = 16 n=16 (including fixed activating the 1 initial block and 2 local blocks), and sliding window size w = 512 w=512 . Both Full Attention and sparse attention models are pretrained on 270 B 270\text{B} tokens of 8 k 8\text{k} -length texts, followed by continued training and supervised fine-tuning on 32 k 32\text{k} -length texts with YaRN ( Peng et al., 2024 ) to achieve long-context adaptation. Both models are trained to full convergence to ensure fair comparison. As shown in Figure 4 , the pretraining loss curve of our NSA and Full Attention baseline demonstrates stable and smooth decline, with NSA consistently outperforming the Full Attention model.
|
||||
Figure 4: Pretraining loss comparison between Full Attention and our NSA on 27B-parameter model. Both models exhibit stable convergence, with NSA achieving lower loss values.
|
||||
Model
|
||||
MMLU
|
||||
MMLU-PRO
|
||||
CMMLU
|
||||
BBH
|
||||
GSM8K
|
||||
MATH
|
||||
DROP
|
||||
MBPP
|
||||
HumanEval
|
||||
Avg.
|
||||
Acc. 5-shot
|
||||
Acc. 5-shot
|
||||
Acc. 5-shot
|
||||
Acc. 3-shot
|
||||
Acc. 8-shot
|
||||
Acc. 4-shot
|
||||
F1 1-shot
|
||||
Pass@1 3-shot
|
||||
Pass@1 0-shot
|
||||
Full Attn
|
||||
0.567
|
||||
0.279
|
||||
0.576
|
||||
0.497
|
||||
0.486
|
||||
0.263
|
||||
0.503
|
||||
0.482
|
||||
0.335
|
||||
0.443
|
||||
NSA
|
||||
0.565
|
||||
0.286
|
||||
0.587
|
||||
0.521
|
||||
0.520
|
||||
0.264
|
||||
0.545
|
||||
0.466
|
||||
0.348
|
||||
0.456
|
||||
Table 1: Pretraining performance comparison between the full attention baseline and NSA on general benchmarks, across knowledge (MMLU, MMLU-PRO, CMMLU), reasoning (BBH, GSM8K, MATH, DROP), and coding (MBPP, HumanEval) tasks. NSA achieves superior average performance on most benchmarks despite high sparsity.
|
||||
## 4.2 Baselines Methods
|
||||
In addition to comparing with Full Attention, we evaluate several state-of-the-art inference-stage sparse attention methods: H2O ( Zhang et al., 2023b ) , infLLM ( Xiao et al., 2024a ) , Quest ( Tang et al., 2024 ) , and Exact-Top, which first computes full attention score and select the top- n n scores keys corresponding to each query and then calculates attention on these positions. These methods span diverse sparse attention paradigms, including KV-cache eviction, query-aware selection, and exact top- n n sparse selection.
|
||||
For general evaluation, where most samples have lengths within the local context window of sparse attention baselines, these methods are effectively equivalent to Full Attention. Therefore, we present only the comparison results between NSA and Full Attention baseline in this setting. In the long-context evaluation, we conduct comparisons across all baseline methods, with the sparsity of all sparse attention methods set to the same to ensure a fair comparison. For chain-of-thought reasoning evaluation, which requires long-text supervised fine-tuning, we limit our comparison to Full Attention, as sparse attention baselines do not support training.
|
||||
## 4.3 Performance Comparison
|
||||
Model
|
||||
SQA
|
||||
MQA
|
||||
Synthetic
|
||||
Code
|
||||
Avg.
|
||||
MFQA-en
|
||||
MFQA-zh
|
||||
Qasper
|
||||
HPQ
|
||||
2Wiki
|
||||
GovRpt
|
||||
Dur
|
||||
PassR-en
|
||||
PassR-zh
|
||||
LCC
|
||||
H2O
|
||||
0.428
|
||||
0.429
|
||||
0.308
|
||||
0.112
|
||||
0.101
|
||||
0.231
|
||||
0.208
|
||||
0.704
|
||||
0.421
|
||||
0.092
|
||||
0.303
|
||||
InfLLM
|
||||
0.474
|
||||
0.517
|
||||
0.356
|
||||
0.306
|
||||
0.250
|
||||
0.277
|
||||
0.257
|
||||
0.766
|
||||
0.486
|
||||
0.143
|
||||
0.383
|
||||
Quest
|
||||
0.495
|
||||
0.561
|
||||
0.365
|
||||
0.295
|
||||
0.245
|
||||
0.293
|
||||
0.257
|
||||
0.792
|
||||
0.478
|
||||
0.135
|
||||
0.392
|
||||
Exact-Top
|
||||
0.502
|
||||
0.605
|
||||
0.397
|
||||
0.321
|
||||
0.288
|
||||
0.316
|
||||
0.291
|
||||
0.810
|
||||
0.548
|
||||
0.156
|
||||
0.423
|
||||
Full Attn
|
||||
0.512
|
||||
0.623
|
||||
0.409
|
||||
0.350
|
||||
0.305
|
||||
0.324
|
||||
0.294
|
||||
0.830
|
||||
0.560
|
||||
0.163
|
||||
0.437
|
||||
NSA
|
||||
0.503
|
||||
0.624
|
||||
0.432
|
||||
0.437
|
||||
0.356
|
||||
0.307
|
||||
0.341
|
||||
0.905
|
||||
0.550
|
||||
0.232
|
||||
0.469
|
||||
Table 2: Performance comparison between our NSA and baselines on LongBench, including subsets in single document QA, multi-document QA, synthetic and code task categories. NSA outperformed most of the baselines including Full Attention.
|
||||
General Evaluation. We evaluated the pretrained NSA and Full Attention baseline, on a comprehensive suite of benchmarks spanning knowledge, reasoning, and coding capabilities, including MMLU ( Hendrycks et al., 2020 ) , MMLU-PRO ( Wang et al., 2024 ) , CMMLU ( Li et al., 2023 ) , BBH ( Suzgun et al., 2022 ) , GSM8K ( Cobbe et al., 2021 ) , MATH ( Hendrycks et al., 2020 ) , DROP ( Dua et al., 2019 ) , MBPP ( Austin et al., 2021 ) , and HumanEval ( Chen et al., 2021 ) . The results are shown in Table 1 . Despite its sparsity, NSA achieves superior overall performance, outperforming all baselines including Full Attention on 7 out of 9 metrics. This indicates that although NSA may not fully leverage its efficiency advantages on shorter sequences, it shows strong performance.
|
||||
Notably, NSA demonstrates significant gains in reasoning-related benchmarks (DROP: +0.042, GSM8K: +0.034), suggesting that
|
||||
our pretraining helps models to develop specialized attention mechanisms. This sparse attention pretraining mechanism forces model to focus on the most important information, potentially enhancing performance by filtering out noise from irrelevant attention pathways.
|
||||
The consistent performance across diverse evaluations also validates NSA’s robustness as a general-purpose architecture.
|
||||
Long-Context Evaluation. As shown in Figure 5 ,
|
||||
NSA achieves perfect retrieval accuracy across all positions in 64k-context needle-in-a-haystack ( Kamradt, 2023 ) test. This performance stems from our hierarchical sparse attention design, which combines compression tokens for efficient global context scanning, and
|
||||
selection tokens for precise local information retrieval.
|
||||
The coarse-grained compression identifies relevant context blocks at low computational cost, while the token-level attention on selected tokens ensures the preservation of critical fine-grained information. This design enables NSA to maintain both global awareness and local precision.
|
||||
Figure 5: Needle-in-a-Haystack retrieval accuracy across context positions with 64k context length. NSA achieves perfect accuracy through its hierarchical sparse attention design.
|
||||
We also evaluate NSA on LongBench ( Bai et al., 2023 ) against state-of-the-art sparse attention methods and Full Attention baseline.
|
||||
To ensure consistent sparsity, we set the token activated by each query in all sparse attention baselines to 2560 tokens, which corresponds to the average number of tokens activated in NSA when handling 32k sequence lengths. Following StreamLLM ( Xiao et al., 2023 ) , this token budget includes the leading 128 tokens and 512 local tokens.
|
||||
We exclude certain subsets from LongBench due to their low scores across all models, which may not provide meaningful comparisons. As shown in Table 2 , NSA achieves the highest average score 0.469, outperforming all baselines (+0.032 over Full Attention and +0.046 over Exact-Top). This improvement arises from two key innovations: (1) our native sparse attention design, which enables end-to-end optimization of sparse patterns during pretraining, facilitates synchronized adaptation between the sparse attention module and other model components; and (2) the hierarchical sparse attention mechanism achieves a balance between local and global information processing.
|
||||
Notably, NSA demonstrates exceptional performance on tasks requiring complex reasoning over long contexts, achieving +0.087 and +0.051 improvements over Full Attention on multi-hop QA tasks (HPQ and 2Wiki), exceeding the performance of baselines on code understanding (LCC: +0.069), and outperforming other methods on passage retrieval (PassR-en: +0.075). These results validate NSA’s capability to handle diverse long-context challenges, with its natively pretrained sparse attention providing additional benefits in learning task-optimal patterns.
|
||||
Chain-of-Thought Reasoning Evaluation.
|
||||
To evaluate NSA’s compatibility with advanced downstream training paradigms, we investigate its capacity to acquire chain-of-thought mathematical reasoning abilities via post-training.
|
||||
Given the limited effectiveness of reinforcement learning on smaller-scale models, we employ knowledge distillation from DeepSeek-R1, conducting supervised fine-tuning (SFT) with 10B tokens of 32k-length mathematical reasoning traces. This produces two comparable models: Full Attention-R (Full Attention baseline) and NSA-R (our sparse variant). We assess both models on the challenging American Invitational Mathematics Examination (AIME 24) benchmark. We use a sampling temperature of 0.7 and a top- p p value of 0.95 to generate 16 responses for each question and obtain the average score. To validate the impact of reasoning depth, we conduct experiments with two generation context limits: 8k and 16k tokens, measuring whether extended reasoning chains improve accuracy. Example comparisons of model predictions are provided in Appendix A .
|
||||
Generation Token Limit
|
||||
8192
|
||||
16384
|
||||
Full Attention-R
|
||||
0.046
|
||||
0.092
|
||||
NSA-R
|
||||
0.121
|
||||
0.146
|
||||
Table 3: AIME Instruction-based Evaluating after supervised fine-tuning. Our NSA-R demonstrates better performance than Full Attention-R at both 8k and 16k sequence lengths
|
||||
As shown in Table 3 , NSA-R achieves significantly higher accuracy than Full Attention-R under the 8k context setting (+0.075), with this advantage persisting at 16k contexts (+0.054).
|
||||
These results validate two key benefits of native sparse attention: (1) The pretrained sparse attention patterns enable efficient capture of long-range logical dependencies critical for complex mathematical derivations; (2) Our architecture’s hardware-aligned design maintains sufficient context density to support growing reasoning depth without catastrophic forgetting. The consistent outperformance across context lengths confirms sparse attention’s viability for advanced reasoning tasks when natively integrated into the training pipeline.
|
||||
## 5 Efficiency Analysis
|
||||
We evaluate the computational efficiency of NSA against Full Attention on an 8-GPU A100 system. In efficiency analysis, we also configure the model with GQA group g = 4 g=4 , heads per group h = 16 h=16 , query/key dimension d k = 192 d_{k}=192 , and value dimension d v = 128 d_{v}=128 .
|
||||
Following the same settings in Section 4 , we set NSA compression block size l = 32 l=32 , sliding stride d = 16 d=16 , selected block size l ′ = 64 l^{\prime}=64 , selected block count n = 16 n=16 , and sliding window size w = 512 w=512 .
|
||||
Figure 6: Comparison of Triton-based NSA kernel with Triton-based FlashAttention-2 kernel. Our implementation significantly reduces latency across all context lengths, with the improvement becoming more pronounced as input length increases.
|
||||
## 5.1 Training Speed
|
||||
We compare the Triton-based implementations of our NSA attention and Full Attention with Triton-based FlashAttention-2 to ensure fair speed comparison across the same backend.
|
||||
As shown in Figure 6 , our NSA achieves progressively greater speedups as context length increases, up to 9.0 × \times forward and 6.0 × \times backward speedup at 64k context-length. Notably, the speed advantage becomes more pronounced with longer sequences. This speedup stems from our hardware-aligned algorithm design to maximize the efficiency of sparse attention architecture: (1) The Blockwise memory access pattern maximizes Tensor Core utilization through coalesced loads, (2) The
|
||||
delicate loop scheduling in the kernel eliminates redundant KV transfers.
|
||||
Context Length
|
||||
8192
|
||||
16384
|
||||
32768
|
||||
65536
|
||||
Full Attention
|
||||
8192
|
||||
16384
|
||||
32768
|
||||
65536
|
||||
NSA
|
||||
2048
|
||||
2560
|
||||
3584
|
||||
5632
|
||||
Expected Speedup
|
||||
4 × \times
|
||||
6.4 × \times
|
||||
9.1 × \times
|
||||
11.6 × \times
|
||||
Table 4: Memory access volume (in equivalent number of tokens) per attention operation during decoding. Due to the low arithmetic intensity and memory-bound nature of decoding, the expected speedup is approximately linear with the volume of memory access.
|
||||
## 5.2 Decoding Speed
|
||||
The decoding speed of Attention is primarily determined by the memory access bottleneck, which is closely tied to the amount of KV cache loading. In each decoding step, Our NSA just needs to load at most ⌊ s − l d ⌋ \left\lfloor\frac{s-l}{d}\right\rfloor compression tokens, n l ′ nl^{\prime} selected tokens, and w w neighbor tokens, where s s is the cached sequence length.
|
||||
As shown in Table 4 , our method exhibits a significant reduction in latency as the decoding length increases, achieving up to 11.6 × \times speedup at 64k context-length. This advantage in memory access efficiency also amplifies with longer sequences.
|
||||
## 6 Discussion
|
||||
In this section, we reflect on the development process of NSA and discuss key insights gained from our exploration of different sparse attention strategies. While our approach demonstrates promising results, understanding the challenges encountered with alternative strategies and analyzing attention patterns provides valuable context for future research directions. We first examine challenges with alternative token selection strategies that motivated our design choices, followed by visualizations that offer insights into attention distribution patterns.
|
||||
## 6.1 Challenges with Alternative Token Selection Strategies
|
||||
Before designing NSA, we explored adapting existing sparse attention methods to the training stage. However, these attempts encountered various challenges, prompting us to design a different sparse attention architecture:
|
||||
Key-Clustering Based Strategies.
|
||||
We examined clustering-based strategies like ClusterKV ( Liu et al., 2024 ) . These methods store Keys and Values from the same cluster in contiguous memory regions. While theoretically feasible for training and inference, they face three significant challenges: (1) Non-trivial computational overhead introduced by dynamic clustering mechanisms; (2) Operator optimization difficulties exacerbated by inter-cluster imbalances, especially in Mixture-of-Experts (MoE) systems, where skewed Expert Parallelism (EP) group execution times lead to persistent load imbalances; (3) Implementation constraints arising from the need for mandatory periodic reclustering and chunk-sequential training protocols. These combined factors create substantial bottlenecks, significantly limiting their effectiveness for real-world deployment.
|
||||
Figure 7: Compare training loss on a 3B-parameter model with Full Attention and different token selection strategies and. Our NSA achieves better performance.
|
||||
Figure 8: Visualization of Attention Map on a Full Attention transformer. Light-colored regions indicate higher attention values. As shown in the figure, attention scores exhibit blockwise clustering distribution.
|
||||
Other Blockwise Selection Strategies.
|
||||
We also considered blockwise key, value selection strategies different from NSA, such as Quest ( Tang et al., 2024 ) and InfLLM ( Xiao et al., 2024a ) . These methods rely on
|
||||
computing an importance score
|
||||
for each KV block and selecting the top- n n blocks based on their similarity with q t q_{t} . However, existing methods face two critical issues: (1) Since the selection operation is non-differentiable, importance score computation based on neural networks relies on auxiliary loss, which increases operator overhead and often degrades model performance; (2) Heuristic parameter-free importance score computation strategy suffer from low recall rates, leading to suboptimal performance. We evaluate both approaches on a 3B-parameter model with similar architecture and compare their loss curve with NSA and Full Attention.
|
||||
For the auxiliary loss-based selection method, we introduce additional queries for each token and representative keys for each block to estimate the block importance scores.
|
||||
We compute block-level supervision signals by mean-pooling attention scores within each key block, and use KL divergence to supervise block importance prediction. We maintain individual query granularity instead of block-averaged queries to accommodate efficient decoding. This auxiliary loss-based importance estimation shares conceptual similarity with SeerAttention ( Gao et al., 2024 ) .
|
||||
For the heuristic parameter-free selection method, following the strategy of Quest, we implement direct selection using the product between queries and coordinate-wise min-max of the key chunks, without introducing additional parameters. We also explore a cold-start training approach where Full Attention is applied for the initial 1000 steps before transitioning to the heuristic blockwise selection.
|
||||
As shown in Figure 8 , both methods exhibited inferior loss.
|
||||
## 6.2 Visualization
|
||||
To explore potential patterns in transformer attention distributions and seek inspiration for our design, we visualize the attention map from our pretrained 27B Full Attention model in Figure 8 . The visualization reveals interesting patterns where attention scores tend to exhibit blockwise clustering characteristics, with nearby keys often showing similar attention scores. This observation inspired our design of NSA, suggesting that selecting key blocks based on spatial continuity might be a promising approach. The blockwise clustering phenomenon indicates that tokens adjacent in the sequence may share certain semantic relationships with query tokens, though the exact nature of these relationships requires further investigation. This observation motivated us to explore a sparse attention mechanism that operates on continuous token blocks rather than individual tokens, aiming to enhance computational efficiency and preserve high-attention patterns.
|
||||
## 7 Related Works
|
||||
We review existing approaches that improve the efficiency of attention computation through sparse attention. These methods can be broadly categorized into three groups based on their core strategies: (1) fixed sparse pattern, (2) dynamic token pruning, and (3) query-aware selection. We introduce several representative works from each category.
|
||||
## 7.1 Fixed Sparse Pattern
|
||||
SlidingWindow is a commonly used approach that allows the query to compute attention only within a fixed window.
|
||||
StreamingLLM ( Xiao et al., 2023 ) combines attention sinks with local windows to process continuous text streams.
|
||||
MoA ( Fu et al., 2024a ) and DuoAttention ( Xiao et al., 2024b ) also adopt similar local and sink information for long sequence modeling.
|
||||
Longformer ( Beltagy et al., 2020 ) interleaves local windowed attention with global tokens to process long sequences.
|
||||
Compared with them, our NSA does not rely on pre-defined sparse patterns, but learns the patterns automatically, unlocking the potential to utilize full context.
|
||||
## 7.2 Dynamic Token Pruning
|
||||
Some methods are designed to reduce memory and computation costs during inference through dynamic KV-cache pruning.
|
||||
H2O ( Zhang et al., 2023b ) , BUZZ ( Zhao et al., 2024 ) , and SepLLM ( Chen et al., 2024a ) implement adaptive approaches to reduce KV-cache memory usage during decoding. These methods dynamically evict tokens deemed less important for future predictions.
|
||||
FastGen Ge et al. (2023) and HeadKV ( Fu et al., 2024b ) optimize computation by assigning different strategies to individual attention heads.
|
||||
SnapKV ( Li et al., 2024 ) introduces a token pruning strategy that reduces the KV-cache by selectively retaining only the most crucial features, enabling efficient memory usage.
|
||||
Unlike these inference-focused approaches, our NSA incorporates sparsity natively during the training phase.
|
||||
## 7.3 Query-Aware Selection
|
||||
Other works focus on query-dependent token selection methods to reduce attention computation while preserving attention quality.
|
||||
Quest ( Tang et al., 2024 ) employs a blockwise selection strategy where each chunk’s importance is estimated by product between query and coordinate-wise min-max of the key chunks.
|
||||
InfLLM ( Xiao et al., 2024a ) combines fixed patterns with retrieval by maintaining attention sinks, local context, and retrievable chunks. This method selects representative keys from each chunk to estimate chunk importance.
|
||||
HashAttention ( Desai et al., 2024 ) formulates pivotal token identification as a recommendation problem by mapping queries and keys to Hamming space using learned functions.
|
||||
ClusterKV ( Liu et al., 2024 ) achieves sparsity by firstly clustering keys and then selecting the most relevant clusters for attention computation based on query-cluster similarity.
|
||||
MInference ( Jiang et al., 2024 ) and TokenSelect ( Wu et al., 2024 ) select KV pairs for computation based on token-level importance scoring.
|
||||
SeerAttention ( Gao et al., 2024 ) separates queries and keys into spatial blocks and performs blockwise selection for efficient computation.
|
||||
Compared to these methods, our NSA achieves hardware-aligned sparse attention computation throughout the full model lifecycle, including training, prefilling, and decoding.
|
||||
## 8 Conclusion
|
||||
We present NSA, a hardware-aligned sparse attention architecture for efficient long-context modeling. By integrating hierarchical token compression with blockwise token selection within a trainable architecture, our architecture achieves accelerated training and inference while maintaining Full Attention performance. NSA advances the state-of-the-art by demonstrating general benchmark performance matches full-attention baselines, exceeding modeling capability in long-context evaluations, and enhanced reasoning ability, all accompanied by measurable reductions in computational latency and achieving significant speedup.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,502 @@
|
||||
TITLE: Inference-Time Scaling for Generalist Reward Modeling
|
||||
ARXIV: 2504.02495
|
||||
AUTHORS: Zijun Liu , Peiyi Wang , Runxin Xu , Shirong Ma , Chong Ruan , Peng Li , Yang Liu , Yu Wu
|
||||
DATE: [Submitted on 3 Apr 2025 ( v1 ), last revised 25 Sep 2025 (this version, v3)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Inference-Time Scaling for Generalist Reward Modeling
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Preliminaries
|
||||
## 2.1 Comparisons of Different RM approaches
|
||||
## Reward Generation Paradigms
|
||||
## Scoring Patterns
|
||||
## Representative Methods
|
||||
## 2.2 Boosting Reward Quality with Principles
|
||||
## 3 Self-Principled Critique Tuning (SPCT)
|
||||
## 3.1 Unpinning Principles from Understanding to Generation
|
||||
## 3.2 Rule-Based Reinforcement Learning
|
||||
## 4 Inference-Time Scaling with SPCT
|
||||
## 5 Results on Reward Modeling Benchmarks
|
||||
## 5.1 Experiment Settings
|
||||
## 5.2 Results and Analysis
|
||||
## 6 Related Work
|
||||
## 7 Conclusion and Future Work
|
||||
## Ethics Statement
|
||||
|
||||
=== BODY (head) ===
|
||||
Inference-Time Scaling for Generalist Reward Modeling
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Preliminaries
|
||||
2.1 Comparisons of Different RM approaches
|
||||
Reward Generation Paradigms
|
||||
Scoring Patterns
|
||||
Representative Methods
|
||||
2.2 Boosting Reward Quality with Principles
|
||||
3 Self-Principled Critique Tuning (SPCT)
|
||||
3.1 Unpinning Principles from Understanding to Generation
|
||||
3.2 Rule-Based Reinforcement Learning
|
||||
4 Inference-Time Scaling with SPCT
|
||||
5 Results on Reward Modeling Benchmarks
|
||||
5.1 Experiment Settings
|
||||
5.2 Results and Analysis
|
||||
6 Related Work
|
||||
7 Conclusion and Future Work
|
||||
References
|
||||
A Additional Related Work
|
||||
Constitutional AI
|
||||
Scalar Reward Models
|
||||
Semi-Scalar Reward Models
|
||||
B Limitations and Future Directions
|
||||
Limitations
|
||||
Future Directions
|
||||
C Implementation Details
|
||||
C.1 Model Training
|
||||
C.2 Baseline Implementation
|
||||
D Experiment Details
|
||||
D.1 Hyper-Parameters
|
||||
D.2 Benchmarks
|
||||
D.3 Detailed Results
|
||||
E Additional Experiments
|
||||
E.1 Input Flexibility of the Pointwise GRM Approach
|
||||
E.1.1 Generating Rewards for Many Responses
|
||||
E.1.2 Generating Rewards for Single Responses
|
||||
E.1.3 Generating Rewards with Reference
|
||||
E.2 Transferability of Generated Principles
|
||||
E.3 Generalization beyond Training Data
|
||||
E.4 Response Length Analysis for Rule-Based RL
|
||||
F Qualitative Analysis
|
||||
F.1 Case Study
|
||||
F.2 Failure Mode Analysis
|
||||
G Prompt Templates
|
||||
License: CC BY-SA 4.0
|
||||
arXiv:2504.02495v3 [cs.CL] 25 Sep 2025
|
||||
## Inference-Time Scaling for Generalist Reward Modeling
|
||||
Zijun Liu
|
||||
Affiliation: DeepSeek-AI
|
||||
Email: zj-liu24@mails.tsinghua.edu.cn
|
||||
Peiyi Wang
|
||||
Email: wangpeiyi9979@gmail.com
|
||||
Runxin Xu
|
||||
Affiliation: DeepSeek-AI
|
||||
Shirong Ma
|
||||
Affiliation: DeepSeek-AI
|
||||
Chong Ruan
|
||||
Affiliation: DeepSeek-AI
|
||||
Peng Li
|
||||
Affiliation: Institute for AI Industry Research (AIR), Tsinghua University
|
||||
Yang Liu
|
||||
Affiliation: Dept. of Computer Sci. & Tech., Tsinghua University
|
||||
Affiliation: Institute for AI Industry Research (AIR), Tsinghua University
|
||||
Yu Wu
|
||||
Affiliation: DeepSeek-AI
|
||||
## Abstract
|
||||
Reinforcement learning (RL) has been widely adopted in post-training for large language models (LLMs) at scale.
|
||||
Recently, the incentivization of reasoning capabilities in LLMs from RL indicates that proper learning methods could enable effective inference-time scalability .
|
||||
A key challenge of RL is to obtain accurate reward signals for LLMs in various domains beyond verifiable questions or artificial rules.
|
||||
In this work, we investigate how to improve reward modeling (RM) with more inference compute for general queries, i.e. the inference-time scalability of generalist RM .
|
||||
For the RM approach, we adopt pointwise generative reward modeling (GRM) to enable flexibility for different input types and the potential for inference-time scaling.
|
||||
For the learning method, we propose Self-Principled Critique Tuning (SPCT) to foster scalable reward generation behaviors in GRMs through online RL, to generate principles adaptively and critiques accurately, resulting in DeepSeek-GRM models.
|
||||
Furthermore, for effective inference-time scaling, we use parallel sampling to expand compute usage, and introduce a meta RM to guide the voting process for better scaling performance.
|
||||
Empirically, we show that SPCT significantly improves the quality and scalability of GRMs, outperforming existing methods and models in various RM benchmarks without severe biases, and it can achieve better performance compared to training-time scaling. DeepSeek-GRM still meets challenges in some tasks, which we believe can be addressed by future efforts in generalist reward systems. The models are released at Hugging Face and ModelScope .
|
||||
1 1 footnotetext: Equal contribution.
|
||||
† Work done during internship at DeepSeek-AI.
|
||||
Figure 1: Inference-time scaling performance with different RMs on all tested RM benchmarks. Results are shown with up to 8 samples for each method, and are further scaled to 32 samples for ours. Non-italic font indicates models based on Gemma-2-27B.
|
||||
## 1 Introduction
|
||||
The remarkable advancements in large language models (LLMs) ( DeepSeek-AI, 2024b ; OpenAI, 2025b ) have catalyzed significant shifts in artificial intelligence research, enabling models to perform tasks that require understanding, generation, and nuanced decision-making capabilities. Recently, reinforcement learning (RL) as a post-training method for LLMs has been widely adopted at scale, and resulting in remarkable improvements in human value alignment ( Ouyang et al., 2022 ; Bai et al., 2022a ) , long-term reasoning ( DeepSeek-AI, 2025 ; OpenAI, 2025c ) , and environment adaptation ( OpenAI, 2025a ) for LLMs. Reward modeling (RM) ( Gao et al., 2023 ) , as a crucial component in RL, is essential for generating accurate reward signals for LLM responses. Current studies ( Lightman et al., 2024 ; DeepSeek-AI, 2025 ) also show that, with high-quality and robust rewards in either training or inference time, LLMs can achieve strong performance in specific domains.
|
||||
However, such high-quality rewards in specific domains are mainly obtained from human-designed environments with clear conditions ( Yao et al., 2022 ; Xie et al., 2024 ) or from hand-crafted rules for verifiable questions, e.g., mathematical problems ( Hendrycks et al., 2021 ; Veeraboina, 2023 ) and coding tasks ( Jimenez et al., 2024 ; Zhuo et al., 2025 ) . In general domains, reward generation is more challenging, as the criteria for rewards are more diverse and complex, and there are often no explicit reference or ground truth. Generalist reward modeling is thus crucial for improving the performance of LLMs in broader applications, either from post-training perspectives, e.g., RL at scale, or from inference perspectives, e.g., RM-guided search. Furthermore, RM performance should be improved by increasing both the training compute ( Gao et al., 2023 ) and the inference compute.
|
||||
In practice, challenges arise in making RMs both general and effectively scalable in inference time.
|
||||
The former demands (1) flexibility for different input types and (2) accurate reward generation in various domains. We refer to this paradigm as generalist reward modeling .
|
||||
Moreover, effective inference-time scalability requires the RM (3) to generate higher-quality reward signals with increased inference compute, and (4) to learn scalable behaviors for better performance-compute scaling.
|
||||
Existing research on reward modeling demonstrates several paradigms for reward generation, including scalar ( Cobbe et al., 2021 ; Wang et al., 2024e ; Liu et al., 2024 ) , semi-scalar ( Ye et al., 2025a ; Yu et al., 2025b ; Zhang et al., 2025a ) , and generative ( Li et al., 2024a ; Kim et al., 2024 ; Vu et al., 2024 ; Cao et al., 2024 ; Arabzadeh et al., 2024 ; Ye et al., 2025b ; Alexandru et al., 2025 ; Yu et al., 2025a ) approaches, and various scoring patterns, such as pointwise ( Kendall & Smith, 1940 ; Gao et al., 2023 ; Yuan et al., 2024 ; Winata et al., 2025 ; Guo et al., 2025 ) and pairwise ( Park et al., 2024 ; Zheng et al., 2023 ; Jiang et al., 2023 ; Wang et al., 2024c ; Liu et al., 2025 ) . These approaches inherently determine the input flexibility and the inference-time scalability of RMs ( (1)&(3) ), as shown in Figure 2 . For instance, pairwise RMs only consider the relative preference of paired responses, lacking flexibility to accept single or multiple responses as input; scalar RMs could hardly generate diverse reward signals for the same response, which obstructs getting better rewards through sampling-based inference-time scaling methods ( Snell et al., 2025 ) .
|
||||
Also, different learning methods ( Wang et al., 2024a ; Ankner et al., 2024 ; Wang et al., 2024c ; Mahan et al., 2024 ) have been proposed to improve the quality of rewards, but few of them focus on inference-time scalability and study the interconnection between the learned reward generation behaviors and the effectiveness of inference-time scaling of RMs, resulting in marginal performance improvement ( (2)&(4) ).
|
||||
Current research ( DeepSeek-AI, 2025 ) indicates that effective inference-time scalability could be enabled by proper learning methods, which raises the question: Can we design a learning method aiming to enable effective inference-time scaling for generalist reward modeling?
|
||||
In this work, we investigated different approaches for RM, and found that pointwise generative reward modeling (GRM) could unify the scoring of single, paired, and multiple responses within pure language representation, overcoming challenge (1) . We explored that certain principles could guide reward generation within proper criteria for GRMs, improving the quality of rewards, which suggested that inference-time scalability of RM might be achieved by scaling the generation of high-quality principles and accurate critiques . Based on this preliminary, we propose a novel learning method, Self-Principled Critique Tuning (SPCT), to foster effective inference-time scalable behaviors in GRMs. By leveraging rule-based online RL, SPCT enables GRMs to learn to adaptively posit principles and critiques based on the input query and responses, leading to better outcome rewards in general domains (challenge (2) ).
|
||||
We then come up with DeepSeek-GRM-27B , which is post-trained with SPCT based on Gemma-2-27B ( Team, 2024 ) .
|
||||
For inference-time scaling, we expand compute usage by sampling multiple times. By sampling in parallel, DeepSeek-GRM could generate different sets of principles and corresponding critiques, and then vote for the final reward. With larger-scale sampling, DeepSeek-GRM could judge more accurately based on more diverse principles, and output rewards with finer granularity , which resolves challenge (3)&(4) .
|
||||
Furthermore, we train a meta RM in addition to voting for better scaling performance. Empirically, we show that SPCT significantly improves the quality and scalability of GRMs, outperforming existing methods and models in multiple comprehensive RM benchmarks without severe domain biases. We also compared the inference-time scaling performance of DeepSeek-GRM-27B with larger models up to 671B parameters, and found it could achieve better performance compared to training-time scaling on model sizes.
|
||||
Though the current method meets challenges in efficiency and specific tasks,
|
||||
with efforts beyond SPCT, we believe GRMs with enhanced scalability and efficiency could serve as a versatile interface for generalist reward systems, advancing the frontiers of LLM post-training and inference.
|
||||
In general, our main contributions are as follows.
|
||||
1.
|
||||
We propose a novel approach, Self-Principled Critique Tuning (SPCT), to foster effective inference-time scalability for generalist reward modeling, resulting in DeepSeek-GRM models. And we further introduce a meta RM to effectively improve the inference-time scaling performance of DeepSeek-GRM beyond voting.
|
||||
2.
|
||||
We empirically show SPCT significantly improves the quality and inference-time scalability of GRMs over existing methods and several strong public models.
|
||||
3.
|
||||
We also applied the SPCT training schedule to LLMs with larger sizes and found that inference-time scaling could outperform model size scaling in training time.
|
||||
Figure 2: Different paradigms for reward generation, including (a) scalar, (b) semi-scalar, and (c) generative approaches, and different scoring patterns, including (i) pointwise and (ii) pairwise approaches. We list the representative methods for each approach, and corresponding inference-time scalability (whether better rewards could be obtained from multiple sampling) and input flexibility (whether supports rating single and multiple responses).
|
||||
## 2 Preliminaries
|
||||
## 2.1 Comparisons of Different RM approaches
|
||||
As shown in Figure 2 , RM approaches are mainly determined by reward generation paradigms and scoring patterns, which inherently affect the inference-time scalability and the input flexibility of the RM. For reward generation paradigms , we distinguish three main approaches: scalar, semi-scalar, and generative. For scoring patterns , we distinguish two main approaches: pointwise and pairwise.
|
||||
To expand compute usage in inference time, we focus on sampling-based methods, which generate multiple sets of rewards for the same query and responses, and then aggregate the final reward. Thus, the inference-time scalability of RMs is determined by whether different rewards could be obtained from multiple sampling, where scalar RMs would fail in most cases due to the invariant generation of rewards; and the input flexibility is defined by whether the RM supports rating single, paired, and multiple responses, where pairwise RMs could hardly rate single responses and usually require extra techniques ( Jiang et al., 2023 ; Liu et al., 2025 ) to handle multiple responses.
|
||||
## Reward Generation Paradigms
|
||||
Classic RMs adopt the (a) scalar approach to generate rewards ( ℛ \mathcal{R} ), which assigns scalar values to the given query and responses. The scalar approach is further extended to the (b) semi-scalar approach, which generates texts in addition to the scalar value. And the (c) generative approach only generates textual rewards.
|
||||
ℛ = { 𝑺 (Scalar) ( 𝑺 , 𝑪 ) (Semi-Scalar) 𝑪 (Generative) ∼ r θ ( x , { y i } i = 1 n ) , \mathcal{R}=\begin{cases}\bm{S}&\text{(Scalar)}\\
|
||||
\left(\bm{S},\bm{C}\right)&\text{(Semi-Scalar)}\\
|
||||
\bm{C}&\text{(Generative)}\end{cases}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),
|
||||
(1)
|
||||
where x x is the query, y i y_{i} is the i i -th response, r θ r_{\theta} is the reward function parameterized by θ \theta , 𝑺 ∈ ℝ m , m ≤ n \bm{S}\in\mathbb{R}^{m},m\leq n is the scalar reward, and 𝑪 \bm{C} is the critique.
|
||||
## Scoring Patterns
|
||||
We distinguish two main scoring approaches for rewards: pointwise and pairwise. The (i) pointwise approach assigns an individual score to each response:
|
||||
{ S i } i = 1 n = f point ( ℛ , { y i } i = 1 n ) , ℛ ∼ r θ ( x , { y i } i = 1 n ) , S i ∈ ℝ , \left\{S_{i}\right\}_{i=1}^{n}=f_{\mathrm{point}}\left(\mathcal{R},\{y_{i}\}_{i=1}^{n}\right),\quad\mathcal{R}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),S_{i}\in\mathbb{R},
|
||||
(2)
|
||||
where f point ( ⋅ , ⋅ ) f_{\mathrm{point}}(\cdot,\cdot) is a splitting function.
|
||||
In contrast, the (ii) pairwise approach can be viewed as a best-of- n n method, selecting a single best response from all candidates:
|
||||
y ^ = f pair ( ℛ , { y i } i = 1 n ) , ℛ ∼ r θ ( x , { y i } i = 1 n ) , y ^ ∈ { y i } i = 1 n , \hat{y}=f_{\mathrm{pair}}(\mathcal{R},\{y_{i}\}_{i=1}^{n}),\quad\mathcal{R}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),\hat{y}\in\{y_{i}\}_{i=1}^{n},
|
||||
(3)
|
||||
where f pair ( ⋅ , ⋅ ) f_{\mathrm{pair}}(\cdot,\cdot) is a selection function and n = 2 n=2 in most cases. Though the pairwise approach could be extended to n > 2 n>2 , it could not be applied to single response scoring ( n = 1 n=1 ).
|
||||
## Representative Methods
|
||||
Figure 2 illustrates how the three reward generation paradigms (scalar, semi-scalar, generative) can be combined with the two scoring patterns (pointwise, pairwise). Specifically, Bradley-Terry model ( Kendall & Smith, 1940 ) ( (a)+(i) ) is trained with pairwise preference data and outputs scalar rewards pointwisely
|
||||
{ S i } i = 1 n = f point ( ℛ , { y i } i = 1 n ) = 𝑺 ∈ ℝ n . \left\{S_{i}\right\}_{i=1}^{n}=f_{\mathrm{point}}\left(\mathcal{R},\{y_{i}\}_{i=1}^{n}\right)=\bm{S}\in\mathbb{R}^{n}.
|
||||
(4)
|
||||
PairRM ( Jiang et al., 2023 ) ( (a)+(ii) ) compares a pair of responses with the sign of the scalar reward
|
||||
y ^ = f pair ( ℛ , { y i } i = 1 n ) = y ⌊ 1 2 ( 3 − sgn ( 𝑺 ) ) ⌋ , n = 2 , 𝒮 ∈ ℝ . \hat{y}=f_{\mathrm{pair}}\left(\mathcal{R},\{y_{i}\}_{i=1}^{n}\right)=y_{\lfloor\frac{1}{2}(3-\mathrm{sgn}(\bm{S}))\rfloor},\quad n=2,\mathcal{S}\in\mathbb{R}.
|
||||
(5)
|
||||
The scalar methods above could barely perform inference-time scaling due to the lack of diversity in reward generation.
|
||||
CLoud ( Ankner et al., 2024 ) ( (b)+(i) ) generates scalar rewards for each response based on pre-generated critiques, similar to Equation 4 .
|
||||
LLM-as-a-Judge ( Zheng et al., 2023 ; Wang et al., 2025a ) ( (c)+(ii) ) judges the preference order between paired responses textually,
|
||||
y ^ = f pair ( ℛ , { y i } i = 1 n ) = y f extract ( 𝑪 ) , n = 2 , \hat{y}=f_{\mathrm{pair}}\left(\mathcal{R},\{y_{i}\}_{i=1}^{n}\right)=y_{f_{\mathrm{extract}}(\bm{C})},\quad n=2,
|
||||
(6)
|
||||
where f extract ( ⋅ ) f_{\mathrm{extract}}(\cdot) extracts the index of best response from language representations. However, this approach defaults to neglecting ties of the paired responses. Following Zhang et al. (2025a) , the generation probability of the token that indicates the preference order could be used as the scalar reward ( (b)+(ii) ): 𝒮 = TokenProb ( C ^ ) = r θ ( C ^ | x , { y i } i = 1 n ) \mathcal{S}=\mathrm{TokenProb}(\hat{C})=r_{\theta}(\hat{C}|x,\{y_{i}\}_{i=1}^{n}) , where C ^ \hat{C} is a pre-defined token related to the preference order.
|
||||
However, without additional constraints, GRMs are able to generate pointwise rewards for multiple responses within pure language representations ( (c)+(i) ):
|
||||
{ S i } i = 1 n = f point ( ℛ , { y i } i = 1 n ) = f extract ( 𝑪 ) , \left\{S_{i}\right\}_{i=1}^{n}=f_{\mathrm{point}}\left(\mathcal{R},\{y_{i}\}_{i=1}^{n}\right)=f_{\mathrm{extract}}(\bm{C}),
|
||||
(7)
|
||||
where f extract ( ⋅ ) f_{\mathrm{extract}}(\cdot) extracts the rewards assigned to each response from generation results. Usually, the rewards are discrete, and in this work we assign S i ∈ ℕ , 1 ≤ S i ≤ 10 S_{i}\in\mathbb{N},1\leq S_{i}\leq 10 by default. This approach enables both inference-time scalability and input flexibility.
|
||||
## 2.2 Boosting Reward Quality with Principles
|
||||
Generalist RM requires generating high-quality rewards beyond specific domains ( Hendrycks et al., 2021 ; Jimenez et al., 2024 ) , where the criteria for rewards are more diverse and complex, and there are often no explicit reference or ground truth. To this end, for general domains, we adopt principles to guide reward generation in place of artificial rules.
|
||||
Principles for LLMs are first introduced in Constitutional AI ( Bai et al., 2022b ; Sharma et al., 2025 ) , which are hand-crafted criteria that guide the LLMs or curated classifiers to construct safe data pipelines. With principles, the reward generation of GRMs changes to
|
||||
ℛ = 𝑪 ∼ r θ ( x , { y i } i = 1 n , { p i } i = 1 m ) , \mathcal{R}=\bm{C}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n},\{p_{i}\}_{i=1}^{m}\right),
|
||||
(8)
|
||||
where { p i } i = 1 m \{p_{i}\}_{i=1}^{m} denotes the principles.
|
||||
We conduct a preliminary experiment to examine the influence of proper principles on reward quality, with the Chat Hard subset of Reward Bench ( Lambert et al., 2024 ) and the IFEval subset of the PPE benchmark ( Frick et al., 2025 ) .
|
||||
Method
|
||||
Chat Hard
|
||||
IFEval
|
||||
GPT-4o-2024-08-06
|
||||
76.1
|
||||
56.0
|
||||
w/ Self-Gen. Principles
|
||||
75.9
|
||||
55.6
|
||||
w/ Filtered Principles
|
||||
77.8
|
||||
57.5
|
||||
Gemma-2-27B-it
|
||||
59.1
|
||||
56.1
|
||||
w/ Self-Gen. Principles
|
||||
64.0
|
||||
55.8
|
||||
w/ Filtered Principles
|
||||
68.0
|
||||
57.3
|
||||
Table 1: Preliminary experiments on the influence of principles on reward quality. The default setting of DeepSeek-GRM-27B includes self-generated principles.
|
||||
In the experiment, data samples contain a query and two responses, with the ground-truth label denoting the better response.
|
||||
We use GPT-4o-2024-08-06 to generate the principles and then pointwise rewards four times for each sample. We filter the principles from the correct reward generation process, where the larger reward value is assigned to the labeled better response. We test different LLMs with principles generated by themselves and the filtered principles, and compare them with the default setting with no principle guidance. The results are shown in Table 1 .
|
||||
We found that the self-generated principles barely improve performance, but the filtered principles could significantly boost the reward quality. The result is non-trivial and two main conclusions could be drawn:
|
||||
(a) Current LLMs could generate diverse principles, but not all of them are proper for reward generation. (b) A subset of generated principles could better guide reward generation under correct criteria, indicating a potential of self-bootstrapping.
|
||||
These findings are the foundation of the usage of online RL to optimize GRMs, where they could learn from the principles generated by themselves with a clear signal to tell whether the principles are proper or not.
|
||||
Other details are depicted in Appendix D .
|
||||
Figure 3: Illustration of SPCT, including rejective fine-tuning, rule-based RL, and corresponding scalable behaviors during inference. The inference-time scaling is achieved via naive voting or meta RM guided voting with principles generated at scale, resulting in finer-grained outcome rewards within an expanded value space.
|
||||
## 3 Self-Principled Critique Tuning (SPCT)
|
||||
Inspired by the preliminary results, we developed a novel approach for pointwise GRMs to learn generating adaptive and high-quality principles that could effectively guide the generation of critiques, termed Self-Principled Critique Tuning (SPCT). As shown in Figure 3 , SPCT consists of two phases: rejective fine-tuning, as the cold start, and rule-based online RL, reinforcing generalist reward generation by advancing the generated principles and critiques. SPCT fosters these behaviors in GRMs for inference-time scaling as well.
|
||||
## 3.1 Unpinning Principles from Understanding to Generation
|
||||
From preliminary experiments in Section 2.2 , we found that proper principles could guide reward generation within certain criteria, which is critical for high-quality rewards. However, it remains challenging to generate effective principles for generalist RM at scale. To address this challenge, we propose to unpin principles from understanding to generation, i.e. view principles as a part of reward generation instead of a preprocessing step.
|
||||
Formally, principles guide the generation of rewards following Equation 8 , when principles are pre-defined. GRMs could generate principles themselves, and then generate critiques based on the principles, formalized as
|
||||
{ p i } i = 1 m ∼ p θ ( x , { y i } i = 1 n ) , ℛ = 𝑪 ∼ r θ ( x , { y i } i = 1 n , { p i } i = 1 m ) , \{p_{i}\}_{i=1}^{m}\sim p_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),\quad\mathcal{R}=\bm{C}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n},\{p_{i}\}_{i=1}^{m}\right),
|
||||
(9)
|
||||
where p θ p_{\theta} is the principle generation function parameterized by θ \theta , that shares the same model with reward generation r θ r_{\theta} . In practice, they are implemented with the same language head in LLMs.
|
||||
This shift enables principles to be generated based on the input query and responses, adaptively aligning reward generation process, and the quality and granularity of the principles and corresponding critiques could be further improved with post-training on GRMs.
|
||||
With the principles generated at scale , GRMs could potentially output rewards with finer granularity and broader consideration to achieve better inference-time scalability.
|
||||
## 3.2 Rule-Based Reinforcement Learning
|
||||
To optimize principle and critique generation in GRMs simultaneously, we propose SPCT, which integrates rejective fine-tuning and rule-based RL . The former serves as a cold start.
|
||||
Rejective Fine-Tuning (Cold Start)
|
||||
The core idea of the rejective fine-tuning stage is to train the GRM to generate principles and critiques in the correct format and for various input types. Unlike previous works ( Vu et al., 2024 ; Cao et al., 2024 ; Alexandru et al., 2025 ) that mix RM data for single, paired, and multiple responses in different formats, we adopt pointwise GRM, introduced in Section 2.1 , to flexibly generate rewards for any number of responses in the same format. For data construction, besides general instruction data, we sample trajectories with a pretrained GRM by giving the query and corresponding responses. Each RM data point contains a query and one or multiple responses to the query, as well as the ground-truth label denoting the best response. For each RM data point, the sampling of principles and critiques is performed N RFT N_{\mathrm{RFT}} times. The rejection strategy is also unified, which is to reject trajectories with predicted rewards that are incorrect, and the query and responses with all N RFT N_{\mathrm{RFT}} trajectories correct (too easy).
|
||||
Formally, let r i r_{i} denote the ground-truth reward for the i i -th response y i y_{i} to the query x x , the predicted pointwise rewards { S i } i = 1 n \{S_{i}\}_{i=1}^{n} are correct if
|
||||
{ ∀ i ≠ j , S j > S i , j = arg max l { r l } l = 1 n , if n ≥ 2 , S 1 = r 1 , if n = 1 . \begin{cases}\forall i\neq j,\quad S_{j}>S_{i},\quad j=\argmax_{l}\{r_{l}\}_{l=1}^{n},&\text{if }n\geq 2,\\
|
||||
S_{1}=r_{1},&\text{if }n=1.\end{cases}
|
||||
(10)
|
||||
with a guarantee that the ground-truth rewards only contain one maximum. However, similar to previous works ( Zhang et al., 2025a ) , we found pretrained GRMs could hardly generate correct rewards for a portion of queries and corresponding responses within limited sampling quota. Thus, we optionally append arg max l { r l } l = 1 n \argmax_{l}\{r_{l}\}_{l=1}^{n} to the prompt of the GRM, termed hinted sampling , with the expectation that the predicted rewards to align with the ground truth, besides non-hinted sampling . Specifically, an additional segment “The best response is: Response arg max l { r l } l = 1 n \argmax_{l}\{r_{l}\}_{l=1}^{n} ” will be appended to the input. For hinted sampling, each query and the corresponding responses are sampled once, and trajectories are only rejected when incorrect. Beyond previous studies ( Li et al., 2024a ; Mahan et al., 2024 ) , we observed that hinted sampled trajectories sometimes take shortcuts in the generated critique, especially for reasoning tasks, indicating the necessity and potential benefits of online RL for the GRM.
|
||||
Rule-Based RL The GRM is further fine-tuned with rule-based online RL. Specifically, we use the original setting of GRPO ( Shao et al., 2024 ) with rule-based outcome rewards. During rollout, the GRM generates principles and critiques based on the input query and responses, and then the predicted reward is extracted and compared to the ground truth with accuracy rules. Unlike DeepSeek-AI (2025) , no format rewards are used. Instead, a larger coefficient for the KL penalty is applied to ensure the format and avoid severe biases. Formally, the reward for the i i -th output o i o_{i} to the given query x x and responses { y i } i = 1 n \{y_{i}\}_{i=1}^{n} is
|
||||
r ^ i = { 1 , if n ≥ 2 and ∀ i ′ ≠ j ′ , S j ′ > S i ′ , j ′ = arg max l { r l } l = 1 n , 1 , if n = 1 and S 1 = r 1 , − 1 , otherwise , \hat{r}_{i}=\begin{cases}1,&\text{if }n\geq 2\text{ and }\forall i^{\prime}\neq j^{\prime},\quad S_{j^{\prime}}>S_{i^{\prime}},\quad j^{\prime}=\argmax_{l}\{r_{l}\}_{l=1}^{n},\\
|
||||
1,&\text{if }n=1\text{ and }S_{1}=r_{1},\\
|
||||
-1,&\text{otherwise},\end{cases}
|
||||
(11)
|
||||
where the pointwise rewards { S i } i = 1 n \{S_{i}\}_{i=1}^{n} are extracted from o i o_{i} . The reward function encourages GRMs to distinguish the best responses with online optimized principles and critiques, in favor of effective inference-time scaling. The reward signal could be obtained seamlessly from any preference dataset and labeled LLM responses.
|
||||
## 4 Inference-Time Scaling with SPCT
|
||||
To further improve the performance of DeepSeek-GRM for generalist reward generation using more inference compute, we explore sampling-based strategies to achieve effective inference-time scalability.
|
||||
Voting with Generated Rewards
|
||||
Voting is a widely adopted method for inference-time scaling in RM. Recalling the approaches in Section 2.1 , we demonstrate voting results of k k samples for semi-scalar and generative RMs. For semi-scalar RMs ( Ankner et al., 2024 ; Zhang et al., 2025a ) , voting is performed as averaging:
|
||||
𝑺 ∗ = 1 k ∑ i = 1 k 𝑺 i , { ℛ i = ( 𝑺 i , 𝑪 i ) } i = 1 k ∼ r θ ( x , { y i } i = 1 n ) , ~\bm{S}^{*}=\frac{1}{k}\sum_{i=1}^{k}\bm{S}_{i},\quad\{\mathcal{R}_{i}=(\bm{S}_{i},\bm{C}_{i})\}_{i=1}^{k}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),
|
||||
(12)
|
||||
where 𝑺 ∗ \bm{S}^{*} is the final reward. In practice, the scalar value has limited variance which could hinder the scalability.
|
||||
For pairwise GRMs ( Mahan et al., 2024 ; Wang et al., 2024c ) , voting is performed as selecting the response identified to be the best with the highest frequency, i.e. majority:
|
||||
y ^ ∗ = arg max y ∑ i = 1 k 𝕀 ( y = y ^ i ) , { ℛ i = 𝐂 i } i = 1 k ∼ r θ ( x , { y i } i = 1 n ) , \hat{y}^{*}=\argmax_{y}\sum_{i=1}^{k}\mathbb{I}(y=\hat{y}_{i}),\quad\{\mathcal{R}_{i}=\bm{C}_{i}\}_{i=1}^{k}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),
|
||||
(13)
|
||||
where y ^ ∗ \hat{y}^{*} is the final predicted best response, f pair ( ⋅ , ⋅ ) f_{\mathrm{pair}}(\cdot,\cdot) is a selection function, y ^ i = f pair ( 𝑪 i , { y i } i = 1 n ) \hat{y}_{i}=f_{\mathrm{pair}}(\bm{C}_{i},\{y_{i}\}_{i=1}^{n}) is the individually selected best response of each sample, and 𝕀 ( ⋅ ) \mathbb{I}(\cdot) is the indicator function. Though the voting process is scalable, the majority voted result might be biased since ties are not allowed in each sample, and may not be able to tell apart subtle differences between responses due to the lack of quantitative scores.
|
||||
The voting process for pointwise GRMs is defined as summing the rewards:
|
||||
S i ∗ = ∑ j = 1 k S i , j , { p i , j } i = 1 m j ∼ p θ ( x , { y i } i = 1 n ) , ℛ j = 𝑪 j ∼ r θ ( x , { y i } i = 1 n , { p i , j } i = 1 m j ) , j = 1 , … , k , \footnotesize S_{i}^{*}=\sum_{j=1}^{k}S_{i,j},\hskip 8.50012pt\{p_{i,j}\}_{i=1}^{m_{j}}\sim p_{\theta}\left(x,\{y_{i}\}_{i=1}^{n}\right),\mathcal{R}_{j}=\bm{C}_{j}\sim r_{\theta}\left(x,\{y_{i}\}_{i=1}^{n},\{p_{i,j}\}_{i=1}^{m_{j}}\right),j=1,...,k,
|
||||
(14)
|
||||
where S i ∗ S_{i}^{*} is the final reward for the i i -th response ( i = 1 , … , n i=1,...,n ) and { S i , j } i = 1 n = f point ( 𝑪 j , { y i } i = 1 n ) \{S_{i,j}\}_{i=1}^{n}=f_{\mathrm{point}}(\bm{C}_{j},\{y_{i}\}_{i=1}^{n}) is the j j -th set of pointwise rewards. Since S i , j S_{i,j} is usually set within a small discrete range, e.g., { 1 , … , 10 } \{1,...,10\} , the voting process actually expands the reward space by k k times, and enables the GRM to generate a large number of principles, which benefits the quality and granularity of the final rewards . An intuitive explanation is that, if each principle could be viewed as a proxy of judgement perspectives, a larger number of principles may reflect the real distribution more accurately, leading to scaling effectiveness. Notably, to avoid positional biases and for diversity, responses are shuffled before sampling.
|
||||
Meta Reward Modeling Guided Voting
|
||||
The voting process of DeepSeek-GRM requires multiple sampling and a few generated principles and critiques might be biased or low-quality due to randomness or model limitations. Thus, we train a meta RM to guide the voting process. The meta RM is a pointwise scalar RM, trained to identify the correctness of the principles and critiques generated by DeepSeek-GRM, with the binary cross-entropy loss, where the label is identified based on Equation 10 . The prompt template is in Appendix G , integrates the query, candidate responses, corresponding principles, and the critiques.
|
||||
The dataset comprises trajectories from non-hinted sampling in the RFT stage, and also trajectories sampled from the DeepSeek-GRM to be guided, to both provide enough positive and negative rewards and alleviate the gap between training and inference policy as suggested by Chow et al. (2025) . The guided voting is simple: The meta RM outputs meta rewards for k k sampled rewards, and the final outcome is voted by rewards with top k meta ≤ k k_{\mathrm{meta}}\leq k meta rewards, so that filtering out low-quality samples.
|
||||
## 5 Results on Reward Modeling Benchmarks
|
||||
Model
|
||||
RB
|
||||
PPE Pref.
|
||||
PPE Correct.
|
||||
RMB
|
||||
Overall
|
||||
Avg. Rank ( ↓ \downarrow )
|
||||
Reported Results of Public Models
|
||||
Skywork-Reward-Gemma-2-27B
|
||||
94.1
|
||||
56.6
|
||||
56.6
|
||||
60.2
|
||||
66.9
|
||||
-
|
||||
DeepSeek-V2.5-0905
|
||||
81.5
|
||||
62.8
|
||||
58.5
|
||||
65.7
|
||||
67.1
|
||||
-
|
||||
Gemini-1.5-Pro
|
||||
86.8
|
||||
66.1
|
||||
59.8
|
||||
56.5
|
||||
67.3
|
||||
-
|
||||
ArmoRM-8B-v0.1
|
||||
90.4
|
||||
60.6
|
||||
61.2
|
||||
64.6
|
||||
69.2
|
||||
-
|
||||
InternLM2-20B-Reward
|
||||
90.2
|
||||
61.0
|
||||
63.0
|
||||
62.9
|
||||
69.3
|
||||
-
|
||||
LLaMA-3.1-70b-Instruct
|
||||
84.1
|
||||
65.3
|
||||
59.2
|
||||
68.9
|
||||
69.4
|
||||
-
|
||||
Claude-3.5-sonnet
|
||||
84.2
|
||||
65.3
|
||||
58.8
|
||||
70.6
|
||||
69.7
|
||||
-
|
||||
Nemotron-4-340B-Reward
|
||||
92.0
|
||||
59.3
|
||||
60.8
|
||||
69.9
|
||||
70.5
|
||||
-
|
||||
GPT-4o
|
||||
86.7
|
||||
67.1
|
||||
57.6
|
||||
73.8
|
||||
71.3
|
||||
-
|
||||
Reproduced Results of Baseline Methods
|
||||
LLM-as-a-Judge
|
||||
83.4
|
||||
64.2
|
||||
58.8
|
||||
64.8
|
||||
67.8
|
||||
4.50
|
||||
DeepSeek-BTRM-27B
|
||||
81.7
|
||||
68.3
|
||||
66.7
|
||||
57.9
|
||||
68.6
|
||||
3.50
|
||||
CLoud-Gemma-2-27B
|
||||
82.0
|
||||
67.1
|
||||
62.4
|
||||
63.4
|
||||
68.7
|
||||
3.50
|
||||
DeepSeek-PairRM-27B
|
||||
87.1
|
||||
65.8
|
||||
64.8
|
||||
58.2
|
||||
69.0
|
||||
2.75
|
||||
Results of Our Method
|
||||
DeepSeek-GRM-27B-RFT (Ours)
|
||||
84.5
|
||||
64.1
|
||||
59.6
|
||||
67.0
|
||||
68.8
|
||||
4.00
|
||||
DeepSeek-GRM-27B (Ours)
|
||||
86.0
|
||||
64.7
|
||||
59.8
|
||||
69.0
|
||||
69.9
|
||||
2.75
|
||||
Results of Inference-Time Scaling (Voting@32)
|
||||
DeepSeek-GRM-27B (Ours)
|
||||
88.5
|
||||
65.3
|
||||
60.4
|
||||
69.7
|
||||
71.0
|
||||
-
|
||||
DeepSeek-GRM-27B (MetaRM) (Ours)
|
||||
90.4
|
||||
67.2
|
||||
63.2
|
||||
70.3
|
||||
72.8
|
||||
-
|
||||
Table 2: Overall results of different methods and models on RM benchmarks. Underlined numbers indicate the best performance, bold numbers indicate the best performance among baseline and our methods, and italicized font denotes scalar or semi-scalar RMs. For meta RM guided voting (MetaRM), k meta = 1 2 k k_{\mathrm{meta}}=\frac{1}{2}k . “Avg. Rank ( ↓ \downarrow )” is calculated based on the rank (from 1 to 6, better ranks denoted by darker shades) of baseline and our methods on each dataset.
|
||||
## 5.1 Experiment Settings
|
||||
Benchmarks and Evaluation Metrics
|
||||
We evaluate the performance of different methods on various RM benchmarks of different domains: Reward Bench (RB) ( Lambert et al., 2024 ) , PPE (Preference and Correctness subsets) ( Frick et al., 2025 ) , RMB ( Zhou et al., 2025 ) , ReaLMistake ( Kamoi et al., 2024 ) . We use the standard evaluation metrics for each benchmark: accuracy of picking the best response from a set of responses in Reward Bench, PPE, and RMB, and ROC-AUC for ReaLMistake. To deal with ties of the predicted rewards for multiple responses, we shuffle the responses and determine the best response by arg max i S i \argmax_{i}S_{i} , where S i S_{i} is the predicted reward for the i i -th response after shuffling. Details are in Appendix D .
|
||||
Method Implementation
|
||||
For the baseline methods, we re-implement LLM-as-a-Judge ( Zheng et al., 2023 ) , DeepSeek-BTRM-27B (Bradley-Terry model) ( Kendall & Smith, 1940 ) , CLoud-Gemma-2-27B ( Ankner et al., 2024 ) , and DeepSeek-PairRM-27B ( Jiang et al., 2023 ) based on Gemma-2-27B ( Team, 2024 ) and with all compatible training data and settings as DeepSeek-GRM. For our methods, we implement DeepSeek-GRM-27B-RFT based on Gemma-2-27B, and DeepSeek-GRM on different sizes of LLMs, including DeepSeek-V2-Lite (16B MoE) ( DeepSeek-AI, 2024a ) , Gemma-2-27B, DeepSeek-V2.5 (236B MoE), and DeepSeek-V3 (671B MoE) ( DeepSeek-AI, 2024b ) . The meta RM is trained on Gemma-2-27B.
|
||||
Default results are reported with greedy decoding, and the inference-time scaling uses temperature = 0.5 \mathrm{temperature}=0.5 .
|
||||
Other details are provided in Appendix C .
|
||||
## 5.2 Results and Analysis
|
||||
Performance on RM Benchmarks
|
||||
The overall results of different methods and models on RM benchmarks are shown in Table 2 . We compare the performance of DeepSeek-GRM-27B with the reported results of public models and the reproduced results of baseline methods. We find that DeepSeek-GRM-27B outperforms the baseline methods in overall performance, and achieves competitive performance with strong public RMs, such as Nemotron-4-340B-Reward and GPT-4o; with inference-time scaling, DeepSeek-GRM-27B could further improve and achieve the best overall results. For detailed comparisons, scalar (DeepSeek-BTRM-27B) and semi-scalar (CLoud-Gemma-2-27B) RMs demonstrate biased results on different benchmarks, with significantly better performance on verifiable tasks (PPE Correctness) than all generative RMs, but fail in different other benchmarks, respectively. Nonetheless, most public scalar RMs also exhibit severe domain biases. The PairRM approach could alleviate the problem. LLM-as-a-Judge shows similar trends with DeepSeek-GRM-27B with lower performance, potentially due to the lack of training on rating single responses. In conclusion, SPCT improves the generalist reward generation capability of GRMs, with significantly fewer biases compared to scalar and semi-scalar RMs .
|
||||
Model
|
||||
Overall
|
||||
Reported Results of Public Models
|
||||
Nemotron-4-340B-Reward
|
||||
70.5
|
||||
GPT-4o
|
||||
71.3
|
||||
Results of Inference-Time Scaling (Voting@1)
|
||||
LLM-as-a-Judge
|
||||
67.0
|
||||
CLoud-Gemma-2-27B
|
||||
68.5
|
||||
DeepSeek-GRM-27B-RFT (Ours)
|
||||
67.8
|
||||
DeepSeek-GRM-27B (Ours)
|
||||
67.9
|
||||
Results of Inference-Time Scaling (Voting@8)
|
||||
LLM-as-a-Judge
|
||||
67.6 ( +0.6 )
|
||||
LLM-as-a-Judge w/ TokenProb
|
||||
68.1 ( +1.1 )
|
||||
CLoud-Gemma-2-27B
|
||||
68.8 ( +0.3 )
|
||||
DeepSeek-GRM-27B-RFT (Ours)
|
||||
69.3 ( +1.5 )
|
||||
DeepSeek-GRM-27B (Ours)
|
||||
70.6 ( +2.7 )
|
||||
DeepSeek-GRM-27B (MetaRM) (Ours)
|
||||
72.0 ( +4.1 )
|
||||
Results of Further Inference-Time Scaling (Voting@32)
|
||||
DeepSeek-GRM-27B (Ours)
|
||||
71.0 ( +3.1 )
|
||||
DeepSeek-GRM-27B (MetaRM) (Ours)
|
||||
72.8 ( +4.9 )
|
||||
Table 3: Inference-time scalability results of different methods on RM benchmarks. Settings are the same as Table 2 .
|
||||
Method
|
||||
Overall
|
||||
Results of Greedy Decoding
|
||||
DeepSeek-GRM-27B
|
||||
69.9
|
||||
w/o Principle Generation
|
||||
67.5
|
||||
w/o Rejective Sampling
|
||||
68.7
|
||||
DeepSeek-GRM-27B-RFT
|
||||
68.8
|
||||
w/o Hinted Sampling (①)
|
||||
68.0
|
||||
w/o Non-Hinted Sampling (②)
|
||||
67.4
|
||||
w/o Rejective Sampling (①&②)
|
||||
66.1
|
||||
w/o General Instruction Data
|
||||
63.3
|
||||
Results of Inference-Time Scaling (Voting@8)
|
||||
DeepSeek-GRM-27B
|
||||
70.6
|
||||
w/o Principle Generation
|
||||
68.0
|
||||
Results of Inference-Time Scaling (Voting@32)
|
||||
DeepSeek-GRM-27B
|
||||
71.0
|
||||
DeepSeek-GRM-27B ( k meta = 1 k_{\mathrm{meta}}=1 )
|
||||
71.5
|
||||
DeepSeek-GRM-27B ( k meta = 8 k_{\mathrm{meta}}=8 )
|
||||
72.7
|
||||
DeepSeek-GRM-27B ( k meta = 16 k_{\mathrm{meta}}=16 )
|
||||
72.8
|
||||
Table 4: Ablation studies for different components of the proposed SPCT. Bold numbers indicate the best performance.
|
||||
Inference-Time Scalability
|
||||
The inference-time scaling results of different methods are shown in Table 4 , and the overall trends are demonstrated in Figure 1 . Details are in Appendix D.3 . With up to 8 samples, we find that DeepSeek-GRM-27B has the highest performance increase over the greedy decoding and sampling results. DeepSeek-GRM-27B further shows a strong potential to increase performance with larger inference compute, up to 32 samples. We attribute the effectiveness to the refined principle generation, which expands the output length in a structural way and guides the outcome rewards closer to the ground-truth distribution.
|
||||
The meta RM also reveals its validity in filtering low-quality trajectories for DeepSeek-GRM on each benchmark.
|
||||
Voted with token probabilities, LLM-as-a-Judge also shows a significant performance increase, indicating that the token probability as quantitative weights could help the reliability over mere majority voting on discrete indices. For CLoud-Gemma-2-27B, the performance increase is limited, mainly due to the lack of variance in scalar reward generation, even though the critique has changed a lot. In summary, SPCT improves the inference-time scalability of GRMs, and the meta RM further boosts the scaling performance in general .
|
||||
Ablation Study
|
||||
Table 4 shows the ablation study results of different components of the proposed SPCT, detailed results are listed in Appendix D.3 . Surprisingly, without the cold start with rejective sampled critique data, general-instruction-tuned GRMs still improve significantly after undergoing the online RL ( 66.1 → 68.7 66.1\rightarrow 68.7 ). Also, the non-hinted sampling seems more important than the hinted sampling, potentially because of the shortcuts appearing in hinted sampled trajectories. These indicate the importance of online training for GRMs . Aligned with previous works ( Cao et al., 2024 ) , we confirm that the general instruction data is essential for the performance of GRMs. We find that the principle generation is crucial for the performance of both greedy decoding and inference-time scaling of DeepSeek-GRM-27B . For inference-time scaling, the meta RM guided voting shows robustness with different k meta k_{\mathrm{meta}} . Further analysis on the generalist RM performance, including input flexibility, domain generalization of training data, etc., is discussed in Appendix E .
|
||||
Scaling Inference and Training Costs
|
||||
We further investigate the inference-time and training-time scaling performance of DeepSeek-GRM-27B, by post-training LLMs of different sizes. The models are tested on the Reward Bench, and the results are shown in Figure 4 .
|
||||
We find that direct voting with 32 samples of DeepSeek-GRM-27B could achieve comparable performance to the 671B MoE model, and the meta RM guided voting could achieve the best results with 8 samples, demonstrating the effectiveness of inference-time scaling of DeepSeek-GRM-27B compared to scaling model sizes .
|
||||
Moreover, we test DeepSeek-R1-0120 with a downsampled test set containing 300 samples, and find that its performance is even worse than the 236B MoE RFT model, indicating that expanding long chain-of-thoughts for reasoning tasks could not significantly improve the performance of generalist RM.
|
||||
(a) Inference-time scaling results of DeepSeek-GRM-16B and DeepSeek-GRM-27B ( k meta = 1 2 k k_{\mathrm{meta}}=\frac{1}{2}k ).
|
||||
(b) Training-time scaling results with different model sizes, using greedy decoding except DeepSeek-R1.
|
||||
Figure 4: Inference-time scaling performance v.s. training-time scaling performance on the Reward Bench benchmark.
|
||||
## 6 Related Work
|
||||
Generative Reward Models
|
||||
GRMs represent a paradigm shift from scalar RMs ( Ouyang et al., 2022 ) , modeling reward as textual feedback or scores. ( Li et al., 2024a ; Kim et al., 2024 ; Wang et al., 2024c ; Cao et al., 2024 ; Vu et al., 2024 ; Alexandru et al., 2025 ) , enabling richer reward representations and more flexible for judging single and multiple responses. Previously, the LLM-as-a-judge method ( Zheng et al., 2023 ; Wang et al., 2025a ) accommodates reference-based or reference-free pairwise judgement for evaluating LLMs. Recent studies use offline and online RL to train GRMs ( Wu et al., 2024 ; Mahan et al., 2024 ; Yu et al., 2025a ; Ye et al., 2025b ; Chen et al., 2025 ) , incorporate tools and external knowledge with GRMs ( Li et al., 2024b ; Peng et al., 2025 ) , and even train GRMs as an interface to adjust rewards from environments ( Baker et al., 2025 ) .
|
||||
Though these methods face challenges in efficiency, they demonstrate the potential to improve rewards at scale, towards a more generalist reward system.
|
||||
Inference-Time Scaling for LLMs
|
||||
Inference-time scaling for LLMs has been a critical research direction parallel with scaling LLMs in training time. Studies focus on sampling and RM guided aggregation ( Lightman et al., 2024 ; Brown et al., 2024 ; Snell et al., 2025 ; Wu et al., 2025 ) . Recently, long-horizon chain-of-thoughts ( Wei et al., 2022 ) incentivized from LLMs significantly improve the reasoning capabilities of the models in both solving ( OpenAI, 2024 ; DeepSeek-AI, 2025 ; OpenAI, 2025c ) and judging ( Zheng et al., 2025 ; Kim et al., 2025 ) difficult verifiable questions, as another format of inference-time scaling. However, we do not find an effective way to incentivize long-horizon reward generation for generalist reward modeling as DeepSeek-AI (2025) , and we leave the combination of reasoning and principle-guided reward generation for future engineering efforts.
|
||||
There is also research using scalable rewards or verifiers to improve the performance of policy models, in domains of coding ( Chen et al., 2023 ) , reasoning ( Lifshitz et al., 2025 ) , etc. Thus, the development of inference-time scalable generalist RMs in this work might also contribute to the general performance of policy models by inference-time co-scaling.
|
||||
## 7 Conclusion and Future Work
|
||||
We introduced Self-Principled Critique Tuning (SPCT), a method that enhances the scalability of inference time for generalist reward modeling. With rule-based online RL, SPCT enables adaptive generation of principles and critiques, significantly boosting reward quality and inference-time scalability for GRMs in diverse domains. Empirical results demonstrate that DeepSeek-GRM surpasses baseline methods and a few strong public RMs, and shows notable improvement through inference-time scaling, particularly with the guidance of the meta RM.
|
||||
Future directions could include integrating GRMs into online RL pipelines as versatile interfaces of reward systems, exploring inference-time co-scaling with policy models, or serving as robust offline evaluators for foundation models.
|
||||
## Ethics Statement
|
||||
Our proposed method, Self-Principled Critique Tuning (SPCT), aims to enhance inference-time scalability of generative reward models (GRMs) for general domains. While this advancement promotes accuracy and consistency in reward modeling, several ethical implications might warrant explicit consideration.
|
||||
Firstly, even though through our empirical analysis that DeepSeek-GRM shows fewer biases on different domains, the automated generation of principles and critiques can inadvertently perpetuate or amplify biases when the training data is toxic. We argue that further investigation in the meta RM and other bias mitigation strategies should be prioritized to ensure equitable outcomes.
|
||||
Also, our approach does not aim to diminish human oversight. Instead, we advocate maintaining human-in-the-loop frameworks, and developing reliable proxy methods, like SPCT, to scale human oversight more efficiently and effectively.
|
||||
Secondly, expanded applicability of the inference-time scalable GRMs across diverse domains might raise concerns regarding transparency, accountability, etc. Since the reward generation behavior is largely emerged from self-bootstrapping, the potential for unfaithful principles and critiques is non-negligible. We demonstrate case studies in Appendix F.1 and limitations in Appendix B , and open-source the model under public supervision, which is essential for maintaining trust and ensuring responsible deployment of the artifact.
|
||||
Finally, robust validation and ongoing vigilance across varied RM benchmarks and practical scenarios remain crucial. Ethical use of DeepSeek-GRM necessitates proactive management of risks and continuous evaluation against biases, requiring efforts in research about RM evaluation.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,659 @@
|
||||
TITLE: DeepSeek-Prover-V2: Advancing Formal Mathematical Reasoning via Reinforcement Learning for Subgoal Decomposition
|
||||
ARXIV: 2504.21801
|
||||
AUTHORS: Z.Z. Ren , Zhihong Shao , Junxiao Song , Huajian Xin , Haocheng Wang , Wanjia Zhao , Liyue Zhang , Zhe Fu , Qihao Zhu , Dejian Yang , Z.F. Wu , Zhibin Gou , Shirong Ma , Hongxuan Tang , Yuxuan Liu , Wenjun Gao , Daya Guo , Chong Ruan
|
||||
DATE: [Submitted on 30 Apr 2025 ( v1 ), last revised 18 Jul 2025 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-Prover-V2: Advancing Formal Mathematical Reasoning via Reinforcement Learning for Subgoal Decomposition
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Method
|
||||
## 2.1 Recursive Proof Search via Subgoal Decomposition
|
||||
## Sketching Formal Proofs from Natural Language Reasoning.
|
||||
## Recursive Resolution of Subgoals.
|
||||
## Curriculum Learning for Subgoal-based Theorem Proving.
|
||||
## 2.2 Unifying Informal Reasoning and Proof Formalization
|
||||
## Cold Start by Synthetic Data.
|
||||
## Reasoning-oriented Reinforcement Learning.
|
||||
## 2.3 Training Details of DeepSeek-Prover-V2
|
||||
## Two-Stage Training.
|
||||
## Expert Iteration.
|
||||
## Supervised Fine-tuning.
|
||||
## Reinforcement Learning.
|
||||
## Distillation.
|
||||
## 3 Experimental Results
|
||||
## 3.1 Results on MiniF2F Benchmark
|
||||
## Comparison with SoTA Models.
|
||||
## Proving Challenging Problems through Subgoal-guided Curricula.
|
||||
## CoT vs. non-CoT.
|
||||
## 3.2 Results on Undergraduate-level Benchmarks
|
||||
## Reward Hacking in Reinforcement Learning.
|
||||
## 3.3 Results on Combinatorial Problems
|
||||
## 3.4 Results on FormalMATH
|
||||
## 3.5 ProverBench: Formalization of AIME and Textbook Problems
|
||||
## AIME Formalization.
|
||||
## Textbook Formalization.
|
||||
## 4 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-Prover-V2: Advancing Formal Mathematical Reasoning via Reinforcement Learning for Subgoal Decomposition
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Method
|
||||
2.1 Recursive Proof Search via Subgoal Decomposition
|
||||
Sketching Formal Proofs from Natural Language Reasoning.
|
||||
Recursive Resolution of Subgoals.
|
||||
Curriculum Learning for Subgoal-based Theorem Proving.
|
||||
2.2 Unifying Informal Reasoning and Proof Formalization
|
||||
Cold Start by Synthetic Data.
|
||||
Reasoning-oriented Reinforcement Learning.
|
||||
2.3 Training Details of DeepSeek-Prover-V2
|
||||
Two-Stage Training.
|
||||
Expert Iteration.
|
||||
Supervised Fine-tuning.
|
||||
Reinforcement Learning.
|
||||
Distillation.
|
||||
3 Experimental Results
|
||||
3.1 Results on MiniF2F Benchmark
|
||||
Comparison with SoTA Models.
|
||||
Proving Challenging Problems through Subgoal-guided Curricula.
|
||||
CoT vs. non-CoT.
|
||||
3.2 Results on Undergraduate-level Benchmarks
|
||||
Reward Hacking in Reinforcement Learning.
|
||||
3.3 Results on Combinatorial Problems
|
||||
3.4 Results on FormalMATH
|
||||
3.5 ProverBench: Formalization of AIME and Textbook Problems
|
||||
AIME Formalization.
|
||||
Textbook Formalization.
|
||||
4 Conclusion
|
||||
References
|
||||
A Examples of Non-CoT and CoT Prompting for Proof Generation
|
||||
A.1 Non-CoT Generation of DeepSeek-Prover-V2
|
||||
A.2 CoT Generation Strategy of DeepSeek-Prover-V2
|
||||
B Examples of the Buggy Proofs Generated by DeepSeek-Prover-V2-7B on PutnamBench Triggering the User Interface Bug of apply?
|
||||
C An Example of Using exfalso to Prove a Vacuously True Statement
|
||||
D Revision to MiniF2F
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2504.21801v2 [cs.CL] 18 Jul 2025
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-Prover-V2: Advancing Formal Mathematical Reasoning via Reinforcement Learning for Subgoal Decomposition
|
||||
Z.Z. Ren*, Zhihong Shao*, Junxiao Song*, Huajian Xin † , Haocheng Wang † , Wanjia Zhao † , Liyue Zhang, Zhe Fu
|
||||
Qihao Zhu, Dejian Yang, Z.F. Wu, Zhibin Gou, Shirong Ma, Hongxuan Tang, Yuxuan Liu, Wenjun Gao
|
||||
Daya Guo, Chong Ruan
|
||||
DeepSeek-AI
|
||||
https://github.com/deepseek-ai/DeepSeek-Prover-V2
|
||||
Corresponding author: Core contributors † Work done during internship at DeepSeek-AI.
|
||||
## Abstract
|
||||
We introduce DeepSeek-Prover-V2, an open-source large language model designed for formal theorem proving in Lean 4, with initialization data collected through a recursive theorem proving pipeline powered by DeepSeek-V3.
|
||||
The cold-start training procedure begins by prompting DeepSeek-V3 to decompose complex problems into a series of subgoals.
|
||||
The proofs of resolved subgoals are synthesized into a chain-of-thought process, combined with DeepSeek-V3’s step-by-step reasoning, to create an initial cold start for reinforcement learning.
|
||||
This process enables us to integrate both informal and formal mathematical reasoning into a unified model.
|
||||
The resulting model, DeepSeek-Prover-V2-671B, achieves state-of-the-art performance in neural theorem proving, reaching 88.9 % 88.9\% pass ratio on the MiniF2F-test and solving 47 out of 658 problems from PutnamBench.
|
||||
In addition to standard benchmarks, we introduce ProverBench, a collection of 325 formalized problems, to enrich our evaluation, including 15 selected problems from the recent AIME competitions (years 24-25).
|
||||
Further evaluation on these 15 AIME problems shows that the model successfully solves 6 of them.
|
||||
In comparison, DeepSeek-V3 solves 8 of these problems using majority voting, highlighting that the gap between formal and informal mathematical reasoning in large language models is substantially narrowing.
|
||||
Figure 1: Benchmark performance of DeepSeek-Prover-V2. On the AIME benchmark, DeepSeek-V3 is evaluated using the standard find-answer task for natural-language reasoning, while prover models generate Lean code to construct formal proofs for a given correct answer.
|
||||
## 1 Introduction
|
||||
The emergence of reasoning capabilities in large language models (LLMs) has revolutionized numerous areas of artificial intelligence, particularly in the domain of mathematical problem solving ( DeepSeek-AI, 2025 ) .
|
||||
These advancements are largely enabled by the paradigm of inference-time scaling, most notably through natural language chain-of-thought reasoning ( Jaech et al., 2024 ) .
|
||||
Rather than relying solely on a single forward pass to arrive at an answer, LLMs can reflect on intermediate reasoning steps, improving both accuracy and interpretability.
|
||||
Despite the success of natural language reasoning in solving competition-level mathematical problems, its application to formal theorem proving remains fundamentally challenging.
|
||||
LLMs perform natural language reasoning in an inherently informal manner, relying on heuristics, approximations, and data-driven guessing patterns that often lack the rigorous structure required by formal verification systems.
|
||||
In contrast, proof assistants such as Lean ( Moura and Ullrich, 2021 ) , Isabelle ( Paulson, 1994 ) , and Coq ( Barras et al., 1999 ) operate on strict logical foundations, where every proof step must be explicitly constructed and formally verified.
|
||||
These systems permit no ambiguity, implicit assumptions, or omission of details.
|
||||
Bridging the gap between informal, high-level reasoning and the syntactic rigor of formal verification systems remains a longstanding research challenge in neural theorem proving ( Yang et al., 2024 ) .
|
||||
To harness the strengths of informal mathematical reasoning in support of formal theorem proving, a classical approach is to hierarchically decompose formal proofs based on the guidance of natural-language proof sketches.
|
||||
Jiang et al. (2023) proposed a framework, called Draft, Sketch, and Prove (DSP), that leverages a large language model to generate proof sketches in natural language, which are subsequently translated into formal proof steps.
|
||||
This informal-to-formal theorem proving paradigm closely mirrors the concept of subgoals in hierarchical reinforcement learning ( Barto and Mahadevan, 2003 ; Nachum et al., 2018 ; Eppe et al., 2022 ) , where complex tasks are broken down into a hierarchy of simpler subtasks that can be solved independently to progressively achieve the overarching objective.
|
||||
In formal theorem proving, a subgoal is typically an intermediate proposition or lemma that contributes to the proof of a larger theorem ( Zhao et al., 2023 ; Zhao et al., 2024 ) .
|
||||
This hierarchical decomposition aligns with human problem-solving strategies and supports modularity, reusability, and more efficient proof search ( Wang et al., 2024b ; Zheng et al., 2024 ) .
|
||||
Recent studies have extended this paradigm by employing multi-tiered hierarchies for structured proof generation ( Wang et al., 2024a ) , and by leveraging reinforcement learning techniques to optimize the decomposition of complex theorems into manageable subgoals ( Dong et al., 2024 ) .
|
||||
In this paper, we develop a reasoning model for subgoal decomposition, leveraging a suite of synthetic cold-start data and large-scale reinforcement learning to enhance its performance.
|
||||
To construct the cold-start dataset, we develop a simple yet effective pipeline for recursive theorem proving, utilizing DeepSeek-V3 ( DeepSeek-AI, 2024 ) as a unified tool for both subgoal decomposition and formalization.
|
||||
We prompt DeepSeek-V3 to decompose theorems into high-level proof sketches while simultaneously formalizing these proof steps in Lean 4, resulting in a sequence of subgoals.
|
||||
Since the subgoal decomposition is powered by a large general-purpose model, we use a smaller 7B model to handle the proof search for each subgoal, thereby reducing the associated computational burden.
|
||||
Additionally, we introduce a curriculum learning framework that leverages the decomposed subgoals to generate conjectural theorems, progressively increasing the difficulty of training tasks to better guide the model’s learning process.
|
||||
Once the decomposed steps of a challenging problem are resolved, we pair the complete step-by-step formal proof with the corresponding chain-of-thought from DeepSeek-V3 to create cold-start reasoning data.
|
||||
Based on the cold start, a subsequent reinforcement learning stage is applied to further strengthen the connection between informal mathematical reasoning and formal proof construction.
|
||||
Our experiments show that reinforcement learning starting from the cold start of informal reasoning in task decomposition significantly enhances the model’s capabilities in formal theorem proving.
|
||||
The resulting DeepSeek-Prover-V2-671B model establishes a new state-of-the-art in neural theorem proving across multiple benchmarks. On MiniF2F-test, it achieves 82.4 % 82.4\% accuracy with Pass@32, improving to 88.9 % 88.9\% with Pass@8192.
|
||||
The model shows strong generalization capabilities to college-level theorem proving, solving 37.1 % 37.1\% of ProofNet-test problems with Pass@1024 and tackling 47 out of 658 challenging PutnamBench problems.
|
||||
Additionally, we contribute ProverBench, a benchmark dataset containing 325 formalized problems to advance neural theorem proving research, including 15 from the prestigious AIME competitions (years 24-25).
|
||||
DeepSeek-Prover-V2-671B successfully solves 6 of these 15 challenging AIME problems, further demonstrating its sophisticated mathematical reasoning capabilities.
|
||||
## 2 Method
|
||||
Figure 2: Overview of the cold-start data collection process employed by DeepSeek-Prover-V2. We first prompt DeepSeek-V3 to generate a natural-language proof sketch while simultaneously formalizing it into a Lean statement with sorry placeholders for omitted proof details. A 7B prover model then recursively solves the decomposed subgoals. By combining these subgoal proofs, we construct a complete formal proof for the original complex problem. This composed proof is appended to DeepSeek-V3’s original chain-of-thought, creating high-quality cold-start training data for formal mathematical reasoning.
|
||||
## 2.1 Recursive Proof Search via Subgoal Decomposition
|
||||
Decomposing the proof of a complex theorem into a sequence of smaller lemmas that serve as stepping stones is an effective strategy commonly employed by human mathematicians.
|
||||
Several previous studies have explored this hierarchical strategy in the context of neural theorem proving, aiming to enhance proof search efficiency by leveraging the informal reasoning capabilities of modern LLMs ( Jiang et al., 2023 ; Zhao et al., 2023 ; Wang et al., 2024a ; Dong et al., 2024 ) .
|
||||
In this paper, we develop a simple yet effective pipeline that utilizes DeepSeek-V3 ( DeepSeek-AI, 2024 ) as a unified tool for subgoal decomposition in formal theorem proving.
|
||||
## Sketching Formal Proofs from Natural Language Reasoning.
|
||||
Recent advances in large language models have led to significant breakthroughs in informal reasoning capabilities.
|
||||
To bridge the gap between formal and informal reasoning, we leverage cutting-edge general-purpose LLMs, recognized for their strong mathematical reasoning and instruction-following abilities, to construct the foundational framework of our theorem-proving system.
|
||||
Our findings indicate that off-the-shelf models, such as DeepSeek-V3 ( DeepSeek-AI, 2024 ) , are capable of decomposing proof steps and expressing them in formal languages.
|
||||
To prove a given formal theorem statement, we prompt DeepSeek-V3 to first analyze the mathematical problem in natural language, then decompose the proof into smaller steps, translating each step into a corresponding Lean formal statement.
|
||||
Since general-purpose models are known to struggle with producing complete Lean proofs, we instruct DeepSeek-V3 to generate only a high-level proof sketch with the details omitted.
|
||||
The resulting chain of thought culminates in a Lean theorem composed of a sequence of have statements, each concluded with a sorry placeholder indicating a subgoal to be solved.
|
||||
This approach mirrors the human style of proof construction, in which a complex theorem is incrementally reduced to a sequence of more manageable lemmas.
|
||||
Figure 3: An illustrative example of how we translate decomposed subgoals into a series of lemma statements. We first (a) replace the original goal state and then (b) incorporate preceding subgoals as premises. Statement type (b) is used for recursive solving of complex problems, while both types (a) and (b) are incorporated into the curriculum learning process.
|
||||
## Recursive Resolution of Subgoals.
|
||||
Leveraging the subgoals generated by DeepSeek-V3, we adopt a recursive solving strategy to systematically resolve each intermediate proof step.
|
||||
We extract subgoal expressions from have statements to substitute them for the original goals in the given problems (see Figure 3 (a)), and then incorporate the preceding subgoals as premises (see Figure 3 (b)).
|
||||
This construction enables subsequent subgoals to be resolved using the intermediate results of earlier steps, thereby promoting a more localized dependency structure and facilitating the development of simpler lemmas.
|
||||
To reduce the computational overhead of extensive proof search, we employ a smaller 7B prover model specifically optimized for processing the decomposed lemmas.
|
||||
Upon the successful resolution of all decomposed steps, a complete proof of the original theorem can be automatically derived.
|
||||
## Curriculum Learning for Subgoal-based Theorem Proving.
|
||||
The training of prover models requires large formal-language problem sets, typically derived by formalizing existing natural-language mathematical corpora ( Xin et al., 2024a ; Ying et al., 2024 ; Lin et al., 2025 ) .
|
||||
Although formalization of human-authored texts provides high-quality and diverse formal content, the resulting training signals for prover models are often sparse, as a large proportion of computational attempts do not yield successful proofs and therefore offer no positive reward signals.
|
||||
To generate denser training signals, Dong and Ma (2025) proposed a self-play approach that enriches training problem sets by generating tractable conjectures closely related to the original theorem statements, thereby enabling more efficient use of training compute.
|
||||
In this paper, we implement a straightforward approach that leverages subgoals to expand the scope of formal statements used for model training.
|
||||
We generate two types of subgoal theorems, one incorporating preceding subgoals as premises and one without, corresponding to Figure 3 (b) and Figure 3 (a), respectively.
|
||||
Both types are integrated into the expert iteration stage ( Polu and Sutskever, 2020 ) , establishing a curriculum that progressively guides the prover model toward systematically addressing a curated subset of challenging problems.
|
||||
This procedure builds on the same underlying principle as AlphaProof’s test-time reinforcement learning ( DeepMind, 2024 ) , wherein variations of a target problem are generated to enhance the model’s capability in solving challenging IMO-level problems.
|
||||
## 2.2 Unifying Informal Reasoning and Proof Formalization
|
||||
The algorithmic framework discussed above operates in two stages, leveraging two complementary models: DeepSeek-V3 for lemma decomposition and a 7B prover model to complete the corresponding formal proof details.
|
||||
This pipeline provides a natural and scalable mechanism for synthesizing formal reasoning data by combining high-level reasoning from language models with precise formal verification.
|
||||
In this manner, we unify the capabilities of informal mathematical reasoning and proof formalization within a single model.
|
||||
## Cold Start by Synthetic Data.
|
||||
We curate a subset of challenging problems that remain unsolved by the 7B prover model in an end-to-end manner, but for which all decomposed subgoals have been successfully resolved.
|
||||
By composing the proofs of all subgoals, we construct a complete formal proof for the original problem.
|
||||
This proof is then appended to DeepSeek-V3’s chain-of-thought, which outlines the corresponding lemma decomposition, thereby producing a cohesive synthesis of informal reasoning and subsequent formalization.
|
||||
It enables the collection of hundreds of high-quality synthetic cold-start data, which serve as the foundation for training DeepSeek-Prover-V2.
|
||||
This cold-start dataset generation strategy differs from that of Kimina-Prover ( Wang et al., 2025 ) , a concurrent work on formal reasoning models.
|
||||
Specifically, we synthesize data by formalizing natural-language proofs directly into structured formal proof sketches.
|
||||
In contrast, Kimina-Prover adopts a reverse workflow: it begins by collecting complete formal proofs alongside their informal counterparts, then uses general-purpose reasoning models to retrosynthesize intermediate natural-language reasoning steps into coherent thinking blocks.
|
||||
## Reasoning-oriented Reinforcement Learning.
|
||||
After fine-tuning the prover model on the synthetic cold-start data, we perform a reinforcement learning stage to further enhance its ability to bridge informal reasoning with formal proof construction.
|
||||
Following the standard training objective for reasoning models ( DeepSeek-AI, 2025 ) , we use binary correct-or-incorrect feedback as the primary form of reward supervision.
|
||||
During the training process, we observe that the structure of the generated proofs frequently diverges from the lemma decomposition provided by the chain-of-thought guidance.
|
||||
To address this issue, we incorporate a consistency reward in the early steps of training, which penalizes the structural misalignment, explicitly enforcing the inclusion of all decomposed have -structured lemmas in the final proof.
|
||||
In practice, enforcing this alignment enhances proof accuracy, especially on complex theorems that demand multi-step reasoning.
|
||||
## 2.3 Training Details of DeepSeek-Prover-V2
|
||||
## Two-Stage Training.
|
||||
DeepSeek-Prover-V2 is developed through a two-stage training pipeline that establishes two complementary proof generation modes:
|
||||
1.
|
||||
High-efficiency non-Chain-of-Thought (non-CoT) mode: This mode is optimized for the rapid generation of formal Lean proof codes, focusing on producing concise proofs without explicit intermediate reasoning steps.
|
||||
2.
|
||||
High-precision Chain-of-Thought (CoT) mode: This mode systematically articulates intermediate reasoning steps, emphasizing transparency and logical progression, before constructing the final formal proofs.
|
||||
Consistent with DeepSeek-Prover-V1.5 ( Xin et al., 2024b ) , these two generation modes are governed by two distinct guiding prompts (see Appendix A for examples).
|
||||
In the first stage, we employ expert iteration within a curriculum learning framework to train a non-CoT prover model, meanwhile, synthesizing proofs for hard problems through subgoal-based recursive proving.
|
||||
The non-CoT generation mode is chosen to accelerate iterative training and data collection processes, as it offers significantly faster inference and validation cycles.
|
||||
Building on this foundation, the second stage leverages cold-start chain-of-thought (CoT) data synthesized by integrating DeepSeek-V3’s sophisticated mathematical reasoning patterns with our synthetic formal proofs.
|
||||
The CoT mode is enhanced through a further reinforcement learning stage, following the standard training pipeline commonly used for reasoning models.
|
||||
## Expert Iteration.
|
||||
The training procedure for the non-CoT mode of DeepSeek-Prover-V2 follows the paradigm of expert iteration ( Polu and Sutskever, 2020 ) , a widely adopted framework for developing formal theorem provers.
|
||||
In each training iteration, the current best prover policy is used to generate proof attempts for those challenging problems that remain unsolved in prior iterations.
|
||||
Those successful attempts, verified by Lean proof assistant, are incorporated into the SFT dataset to train an improved model.
|
||||
This iterative loop ensures that the model not only learns from the initial demonstration datasets but also distills its own successful reasoning traces, progressively refining its ability to solve harder problems.
|
||||
The overall training procedure remains largely aligned with that of DeepSeek-Prover-V1 ( Xin et al., 2024a ) and DeepSeek-Prover-V1.5 ( Xin et al., 2024b ) , with only two modifications to the distribution of training problems.
|
||||
First, we incorporate additional problems derived from autoformalization and various open-source datasets ( Ying et al., 2024 ; Dong and Ma, 2025 ; Lin et al., 2025 ) , broadening the coverage of the training problem domains.
|
||||
Second, we augment the dataset with problems generated through subgoal decomposition, aiming at solving more challenging instances from the valid split of the MiniF2F benchmark ( Zheng et al., 2022 ) .
|
||||
## Supervised Fine-tuning.
|
||||
We perform supervised fine-tuning on DeepSeek-V3-Base-671B ( DeepSeek-AI, 2024 ) using a constant learning rate of 5e-6 within a context window of 16,384 tokens.
|
||||
Our training corpus consists of two complementary sources: (1) non-CoT data collected through expert iteration, which produces Lean codes without intermediate reasoning steps; and (2) the cold-start CoT data described in Section 2.2 , which distills DeepSeek-V3’s advanced mathematical reasoning processes into structured proving pathways.
|
||||
The non-CoT components emphasize formal verification skills in the Lean theorem prover ecosystem, while the CoT examples explicitly model the cognitive process of transforming mathematical intuition into formal proof structures.
|
||||
## Reinforcement Learning.
|
||||
We employ Group Relative Policy Optimization ( Shao et al., 2024 , GRPO;) as our reinforcement learning algorithm, which has demonstrated superior effectiveness and efficiency in reasoning tasks ( DeepSeek-AI, 2025 ) .
|
||||
Unlike PPO ( Schulman et al., 2017 ) , GRPO eliminates the need for a separate critic model by sampling a group of candidate proofs for each theorem prompt and optimizing the policy based on their relative rewards.
|
||||
Training utilizes binary rewards, where each generated Lean proof receives a reward of 1 if verified as correct and 0 otherwise.
|
||||
To ensure effective learning, we curate training prompts to include only problems that are sufficiently challenging yet solvable by the supervised fine-tuned model.
|
||||
During each iteration, we sample 256 distinct problems, generating 32 candidate proofs per theorem with a maximum sequence length of 32,768 tokens.
|
||||
## Distillation.
|
||||
We extend the maximum context length of DeepSeek-Prover-V1.5-Base-7B ( Xin et al., 2024b ) from 4,096 to 32,768 tokens and fine-tune this extended-context model using rollout data collected during the reinforcement learning phase of DeepSeek-Prover-V2-671B.
|
||||
Alongside the CoT reasoning mode, we incorporate non-CoT proof data collected during expert iteration to enable a cost-efficient proving option that produces concise formal outputs with a small-size model.
|
||||
In addition, we perform the same reinforcement learning stage used in the training of the 671B model to boost the performance of DeepSeek-Prover-V2-7B.
|
||||
## 3 Experimental Results
|
||||
In this section, we present a systematic evaluation of DeepSeek-Prover-V2 across diverse benchmark datasets of formal theorem proving, covering both high school competition problems and undergraduate-level mathematics.
|
||||
All experimental results of DeepSeek-Prover-V2 are conducted with Lean 4.9.0-rc2, using the same testing environment as DeepSeek-Prover-V1.5 ( Xin et al., 2024b ) .
|
||||
Without further specification, baseline evaluation results are sourced from their respective original papers.
|
||||
## 3.1 Results on MiniF2F Benchmark
|
||||
MiniF2F ( Zheng et al., 2022 ) consists of 488 formalized problem statements sourced from a diverse range of mathematical materials, including the AIME, AMC, and IMO competitions, along with selected problems from the MATH dataset ( Hendrycks et al., 2021 ) .
|
||||
The benchmark includes Olympiad-level problems covering core areas of elementary mathematics, including algebra, number theory, and induction.
|
||||
These problems are divided into two equally sized subsets, denoted by miniF2F-valid and miniF2F-test, each containing 244 problems with an identical distribution across subject areas.
|
||||
We reserve the miniF2F-test set exclusively for evaluating model performance, while the miniF2F-valid problems are incorporated into curriculum learning with subgoal decomposition.
|
||||
We adopt the revised version of miniF2F released by Wang et al. (2025) , and further introduce two additional revisions to miniF2F-valid and one revision to miniF2F-test (see Appendix D ).
|
||||
Method
|
||||
Model size
|
||||
Sample budget
|
||||
miniF2F-test
|
||||
Tree Search Methods
|
||||
Hypertree Proof Search ( Lample et al., 2022 )
|
||||
600M
|
||||
64 × 5000 64\times 5000
|
||||
41.0 % 41.0\%
|
||||
InternLM2.5-StepProver + BFS + CG ( Wu et al., 2024 )
|
||||
7B
|
||||
256 × 32 × 600 256\times 32\times 600
|
||||
65.9 % 65.9\%
|
||||
HunyuanProver v16 + BFS + DC ( Li et al., 2024 )
|
||||
7B
|
||||
600 × 8 × 400 600\times 8\times 400
|
||||
68.4 % 68.4\%
|
||||
BFS-Prover ( Xin et al., 2025 )
|
||||
7B
|
||||
2048 × 2 × 600 2048\times 2\times 600
|
||||
70.83 % ± 0.89 % 70.83\%\pm 0.89\%
|
||||
Whole-proof Generation Methods
|
||||
Leanabell-Prover-GD-RL ( Zhang et al., 2025 )
|
||||
7B
|
||||
128 128
|
||||
61.1 % 61.1\%
|
||||
Goedel-Prover-SFT ( Lin et al., 2025 )
|
||||
7B
|
||||
25600 25600
|
||||
64.7 % 64.7\%
|
||||
STP ( Dong and Ma, 2025 )
|
||||
7B
|
||||
25600 25600
|
||||
67.6 % 67.6\%
|
||||
Kimina-Prover-Preview-Distill ( Wang et al., 2025 )
|
||||
7B
|
||||
1
|
||||
52.5 % 52.5\%
|
||||
32 32
|
||||
63.1 % 63.1\%
|
||||
1024 1024
|
||||
70.8 % 70.8\%
|
||||
Kimina-Prover-Preview ( Wang et al., 2025 )
|
||||
72B
|
||||
1
|
||||
52.94 % 52.94\%
|
||||
32 32
|
||||
68.85 % 68.85\%
|
||||
1024 1024
|
||||
77.87 % 77.87\%
|
||||
8192 8192
|
||||
80.74 % 80.74\%
|
||||
DeepSeek-Prover-V2 (non-CoT)
|
||||
7B
|
||||
1 1
|
||||
55.5 % ± 1.4 % 55.5\%\pm 1.4\%
|
||||
32 32
|
||||
68.0 % ± 0.5 % 68.0\%\pm 0.5\%
|
||||
1024 1024
|
||||
73.2 % ± 0.5 % 73.2\%\pm 0.5\%
|
||||
8192 8192
|
||||
75.0 % 75.0\%
|
||||
671B
|
||||
1 1
|
||||
59.5 % ± 1.4 % 59.5\%\pm 1.4\%
|
||||
32 32
|
||||
73.8 % ± 0.4 % 73.8\%\pm 0.4\%
|
||||
1024 1024
|
||||
76.7 % ± 0.2 % 76.7\%\pm 0.2\%
|
||||
8192 8192
|
||||
78.3 % 78.3\%
|
||||
DeepSeek-Prover-V2 (CoT)
|
||||
7B
|
||||
1 1
|
||||
58.6 % ± 1.1 % 58.6\%\pm 1.1\%
|
||||
32 32
|
||||
75.6 % ± 0.5 % 75.6\%\pm 0.5\%
|
||||
1024 1024
|
||||
79.9 % ± 0.3 % 79.9\%\pm 0.3\%
|
||||
8192 8192
|
||||
82.0 % 82.0\%
|
||||
671B
|
||||
1 1
|
||||
61.9 % ± 1.6 % 61.9\%\pm 1.6\%
|
||||
32 32
|
||||
82.4 % ± 0.6 % 82.4\%\pm 0.6\%
|
||||
1024 1024
|
||||
86.6 % ± 0.3 % 86.6\%\pm 0.3\%
|
||||
8192 8192
|
||||
88.9 % \mathbf{88.9\%}
|
||||
Table 1: Comparison with state-of-the-art models on the miniF2F-test dataset. The notation μ ± σ \mu\pm\sigma denotes the average accuracy μ \mu and the standard deviation σ \sigma . The tags CoT and non-CoT refer to two generation modes of a unified model, each guided by a different prompt.
|
||||
Problem Category
|
||||
miniF2F-valid
|
||||
miniF2F-test
|
||||
curriculum (+Pass@8192)
|
||||
Pass@8192
|
||||
Olympiad
|
||||
IMO
|
||||
10 / 20 = 50.0 % 10/20=50.0\%
|
||||
10 / 20 = 50.0 % 10/20=50.0\%
|
||||
AIME
|
||||
10 ( + 2 ) / 15 = 80.0 % 10(+2)/15=80.0\%
|
||||
14 / 15 = 93.3 % 14/15=93.3\%
|
||||
AMC
|
||||
39 / 45 = 86.7 % 39/45=86.7\%
|
||||
35 / 45 = 77.8 % 35/45=77.8\%
|
||||
MATH
|
||||
Algebra
|
||||
69 / 70 = 98.6 % 69/70=98.6\%
|
||||
70 / 70 = 100.0 % 70/70=100.0\%
|
||||
Number Theory
|
||||
58 / 60 = 96.7 % 58/60=96.7\%
|
||||
58 / 60 = 96.7 % 58/60=96.7\%
|
||||
Custom
|
||||
Algebra
|
||||
17 / 18 = 94.4 % 17/18=94.4\%
|
||||
15 / 18 = 83.3 % 15/18=83.3\%
|
||||
Number Theory
|
||||
8 / 8 = 100.0 % 8/8=100.0\%
|
||||
7 / 8 = 87.5 % 7/8=87.5\%
|
||||
Induction
|
||||
8 / 8 = 100.0 % 8/8=100.0\%
|
||||
8 / 8 = 100.0 % 8/8=100.0\%
|
||||
Overall Pass Rate
|
||||
𝟐𝟏𝟗 ( + 𝟐 ) / 𝟐𝟒𝟒 = 90.6 % \mathbf{219(+2)/244=90.6\%}
|
||||
𝟐𝟏𝟕 / 𝟐𝟒𝟒 = 88.9 % \mathbf{217/244=88.9\%}
|
||||
Table 2: Problems solved by DeepSeek-Prover-V2-671B on the miniF2F benchmark. Results on miniF2F-valid are collected throughout the curriculum learning process, and DeepSeek-Prover-V2-671B is further invoked with Pass@8192 on the remaining problems.
|
||||
## Comparison with SoTA Models.
|
||||
Table 1 summarizes a comparison of state-of-the-art formal theorem-proving modeling evaluated on the miniF2F-test dataset.
|
||||
The experimental results demonstrate that DeepSeek-Prover-V2-671B establishes a new state-of-the-art performance on the miniF2F-test benchmark, achieving an unprecedented 82.4 % 82.4\% accuracy with only 32 samples when leveraging the CoT generation strategy.
|
||||
Notably, the more parameter-efficient DeepSeek-Prover-V2-7B also exhibits competitive performance, surpassing all existing open-source theorem provers in the literature.
|
||||
The comparative analysis further reveals a compelling scaling pattern: as the sample budget increases from 1 to 8192, the performance gap between the 7B and 671B variants widens considerably, with the larger model demonstrating superior sample efficiency and a steeper improvement trajectory.
|
||||
## Proving Challenging Problems through Subgoal-guided Curricula.
|
||||
Table 2 presents a detailed breakdown of the problems solved by DeepSeek-Prover-V2 on the miniF2F benchmark, where it achieves strong overall performance with a 90.6 % 90.6\% pass rate on the validation set and 88.9 % 88.9\% on the test set.
|
||||
Remarkably, our subgoal-guided curriculum learning framework, which integrates the general-purpose model DeepSeek-V3 with a lightweight specialized 7B prover, achieves a 89.8 % 89.8\% success rate on miniF2F-valid, nearly matching the performance of DeepSeek-Prover-V2-671B.
|
||||
These findings highlight the potential of state-of-the-art general-purpose LLMs to extend beyond natural language understanding and effectively support complex formal reasoning tasks.
|
||||
Through strategic subgoal decomposition, the model is able to break down challenging problems into a sequence of tractable steps, serving as an effective bridge between informal reasoning and formal proof construction.
|
||||
# output tokens
|
||||
non-CoT
|
||||
CoT
|
||||
7B
|
||||
442.6
|
||||
4488.5
|
||||
671B
|
||||
761.8
|
||||
6751.9
|
||||
Table 3: Average number of tokens generated by DeepSeek-Prover-V2 on miniF2F-test.
|
||||
## CoT vs. non-CoT.
|
||||
The experimental results in Table 1 demonstrate a substantial performance advantage of the CoT reasoning mode over the non-CoT mode in formal mathematical reasoning.
|
||||
This reinforces the effectiveness of CoT prompting, which encourages decomposition of complex problems into intermediate steps, and further confirms that inference-time scaling holds in the domain of formal theorem proving.
|
||||
Complementing these findings, Table 3 provides statistics on the number of tokens generated by DeepSeek-Prover-V2 under different reasoning modes.
|
||||
As expected, the CoT mode produces significantly longer outputs, reflecting its sophisticated reasoning process.
|
||||
Interestingly, within the non-CoT setting, the 671B model generates longer outputs on average compared to the 7B model.
|
||||
A closer examination reveals that, although explicit reasoning is not prompted in the non-CoT mode, the larger model often inserts brief natural language comments within the proof code that resemble implicit reasoning steps (see Appendix A ).
|
||||
This suggests that high-capacity models may internalize and externalize intermediate reasoning implicitly, even in the absence of explicit CoT prompting
|
||||
## 3.2 Results on Undergraduate-level Benchmarks
|
||||
ProofNet ( Azerbayev et al., 2023 ) consists of 371 problems in Lean 3, drawn from a range of popular undergraduate pure mathematics textbooks, covering topics such as real and complex analysis, linear algebra, abstract algebra, and topology.
|
||||
We use the Lean 4 translation of ProofNet made available by Xin et al. (2024b) , which is further divided into two splits: ProofNet-valid and ProofNet-test, containing 185 and 186 problems, respectively.
|
||||
The test split of ProofNet is reserved exclusively for model evaluation, as variants of the ProofNet-valid problems are included in the public synthetic dataset provided by Dong and Ma (2025) , which is used in our supervised fine-tuning.
|
||||
The results, shown in Table 4 , indicate a substantial improvement in the pass rate of DeepSeek-Prover-V2 when using CoT reasoning compared to the non-CoT setting.
|
||||
Notably, despite the training data being predominantly drawn from high-school level mathematics, the model exhibits strong generalization to more advanced, college-level mathematical problems, underscoring its robust formal reasoning capabilities.
|
||||
PutnamBench ( Tsoukalas et al., 2024 ) is a continuously updated benchmark featuring competition mathematics problems from the William Lowell Putnam Mathematical Competition , spanning the years 1962 to 2023.
|
||||
The Putnam Competition is a highly prestigious annual mathematics competition for undergraduate students across the United States and Canada, encompassing a variety of college-level domains such as analysis, linear algebra, abstract algebra, combinatorics, probability, and set theory.
|
||||
We evaluate our model on the latest release of PutnamBench, which contains 658 problems formalized in Lean 4.
|
||||
We exclude problems that are incompatible with Lean 4.9.0 and evaluate the model on the remaining set of 649 problems.
|
||||
Our initial run solved 49 problems using a sample budget of 1024 per problem.
|
||||
After submitting the proofs to the benchmark maintainers, we excluded two problems due to misformulated statements.
|
||||
The final result stands at 47 problems successfully solved (as shown in Table 4 ), which significantly outperforms its non-CoT counterpart.
|
||||
These results further highlight the effectiveness of the CoT reasoning approach in handling challenging, college-level mathematical problems.
|
||||
Method
|
||||
Model size
|
||||
Sample budget
|
||||
ProofNet-test
|
||||
PutnamBench
|
||||
Goedel-Prover-SFT ( Lin et al., 2025 )
|
||||
7B
|
||||
32 32
|
||||
15.6 % 15.6\%
|
||||
6 / 644 6/644
|
||||
512 512
|
||||
-
|
||||
7 / 644 7/644
|
||||
STP ( Dong and Ma, 2025 )
|
||||
7B
|
||||
128 128
|
||||
19.5 % ± 0.7 % 19.5\%\pm 0.7\%
|
||||
7 / 644 7/644
|
||||
3200 3200
|
||||
23.9 % ± 0.6 % 23.9\%\pm 0.6\%
|
||||
8 / 644 8/644
|
||||
25600 25600
|
||||
26.9 % 26.9\%
|
||||
-
|
||||
DeepSeek-Prover-V2 (non-CoT)
|
||||
7B
|
||||
32 32
|
||||
21.6 % ± 0.2 % 21.6\%\pm 0.2\%
|
||||
8 / 658 8/658
|
||||
128 128
|
||||
23.1 % ± 0.6 % 23.1\%\pm 0.6\%
|
||||
9 / 658 9/658
|
||||
1024 1024
|
||||
24.7 % 24.7\%
|
||||
10 / 658 10/658
|
||||
671B
|
||||
32 32
|
||||
23.8 % ± 0.2 % 23.8\%\pm 0.2\%
|
||||
9 / 658 9/658
|
||||
128 128
|
||||
27.2 % ± 0.5 % 27.2\%\pm 0.5\%
|
||||
11 / 658 11/658
|
||||
1024 1024
|
||||
31.2 % 31.2\%
|
||||
15 / 658 15/658
|
||||
DeepSeek-Prover-V2 (CoT)
|
||||
7B
|
||||
32 32
|
||||
23.0 % ± 0.4 % 23.0\%\pm 0.4\%
|
||||
9 / 658 9/658
|
||||
128 128
|
||||
25.4 % ± 0.7 % 25.4\%\pm 0.7\%
|
||||
10 / 658 10/658
|
||||
1024 1024
|
||||
29.6 % 29.6\%
|
||||
11 / 658 11/658
|
||||
671B
|
||||
32 32
|
||||
30.5 % ± 0.7 % 30.5\%\pm 0.7\%
|
||||
22 / 658 22/658
|
||||
128 128
|
||||
33.6 % ± 0.3 % 33.6\%\pm 0.3\%
|
||||
33 / 658 33/658
|
||||
1024 1024
|
||||
37.1 % \mathbf{37.1\%}
|
||||
𝟒𝟕 / 𝟔𝟓𝟖 \mathbf{47/658}
|
||||
Table 4: The experimental results on ProofNet-test and PutnamBench. The scores for Goedel-Prover-SFT and STP on PutnamBench are sourced from their original papers, which conducted evaluations on an earlier version of PutnamBench comprising 644 problems.
|
||||
## Reward Hacking in Reinforcement Learning.
|
||||
Our initial report claimed an unexpected finding that DeepSeek-Prover-V2-7B successfully solved 13 problems on PutnamBench that remained unsolved by its larger 671B counterpart.
|
||||
We acknowledge the Lean community for their assistance in identifying the cause of this unexpected result, which was traced to a user interface bug in Lean 4.9.0. Specifically, the apply? tactic fails to emit sorry declarations under certain corner cases.
|
||||
Upon closer examination of the model’s outputs, we identified a distinctive pattern in its reasoning approach: the 7B model frequently employs Cardinal.toNat and Cardinal.natCast_inj to exploit this user-interface bug (see examples in Appendix B ), which are noticeably absent in the outputs generated by the 671B version.
|
||||
## 3.3 Results on Combinatorial Problems
|
||||
CombiBench
|
||||
Pass@16
|
||||
Kimina-Prover-Preview ( Wang et al., 2025 )
|
||||
7 / 100 7/100
|
||||
DeepSeek-Prover-V2-7B
|
||||
non-CoT
|
||||
6 / 100 6/100
|
||||
CoT
|
||||
7 / 100 7/100
|
||||
DeepSeek-Prover-V2-671B
|
||||
non-CoT
|
||||
8 / 100 8/100
|
||||
CoT
|
||||
𝟏𝟎 / 𝟏𝟎𝟎 \mathbf{10/100}
|
||||
Table 5: Evaluation results on CombiBench under the with-solution setting.
|
||||
CombiBench ( Liu et al., 2025 ) is a comprehensive benchmark comprising 100 combinatorial competition problems formalized in Lean 4, each paired with its corresponding natural-language statement.
|
||||
We evaluate DeepSeek-Prover-V2 in the with-solution setting of this benchmark, where the correct answer is embedded in the Lean statement, allowing the evaluation to focus solely on proof generation.
|
||||
After filtering out problems incompatible with Lean 4.9.0 and those containing multiple sorry placeholders, we evaluate on 77 problems from the benchmark.
|
||||
Our initial run reported 12 successfully solved problems.
|
||||
After collaborating with the benchmark maintainers and identifying two misformulated problem statements, the final result stands at 10 solved problems.
|
||||
These results indicate that, while the prover model is primarily trained in number theory and algebra, it demonstrates promising generalization to combinatorial problems, despite their persistent difficulty.
|
||||
An interesting finding is that the chain-of-thought reasoning of DeepSeek-Prover-V2-671B can effectively identify misformulations and adapt its proof strategy accordingly.
|
||||
We include an example in the Appendix C where the model detects a contradiction within a misformulated statement and derives False to close the target goal through exfalso .
|
||||
## 3.4 Results on FormalMATH
|
||||
FormalMATH
|
||||
All
|
||||
Lite
|
||||
Pass@32
|
||||
Pass@32
|
||||
Pass@3200
|
||||
Goedel-Prover-SFT ( Lin et al., 2025 )
|
||||
13.53 % 13.53\%
|
||||
46.70 % 46.70\%
|
||||
49.41 % 49.41\%
|
||||
STP ( Dong and Ma, 2025 )
|
||||
13.87 % 13.87\%
|
||||
48.59 % 48.59\%
|
||||
53.17 % 53.17\%
|
||||
Kimina-Prover-Preview-Distill-7B
|
||||
16.46 % 16.46\%
|
||||
48.94 % 48.94\%
|
||||
-
|
||||
DeepSeek-Prover-V2-7B
|
||||
22.41 % 22.41\%
|
||||
51.76 % 51.76\%
|
||||
55.06 % 55.06\%
|
||||
DeepSeek-Prover-V2-671B
|
||||
28.31 % \mathbf{28.31\%}
|
||||
56.00 % \mathbf{56.00\%}
|
||||
61.88 % \mathbf{61.88\%}
|
||||
Table 6: Evaluation results on FormalMATH-All/Lite. The performance of DeepSeek-Prover-V2 is evaluated with the CoT generation mode.
|
||||
FormalMATH ( Yu et al., 2025 ) is a large-scale formal theorem proving benchmark with 5560 problems, ranging from high school Olympiad questions to undergraduate-level theorems, and covering a wide range of domains such as algebra, applied mathematics, calculus, number theory, and discrete mathematics.
|
||||
FormalMATH-Lite is a manageable subset of 425 problems (comprising 359 high school-level and 66 undergraduate-level problems) carefully selected from the full corpus, addressing the impracticality of evaluating on the full dataset while still enabling systematic assessment of test-time scaling across various mathematical domain.
|
||||
Table 6 presents the comparative performance analysis of various theorem provers on both FormalMATH-All and FormalMATH-Lite benchmarks.
|
||||
DeepSeek-Prover-V2-671B demonstrates superior performance across all evaluation metrics, successfully solving 28.31 % 28.31\% of problems in the full dataset and 56.00 % 56.00\% of problems in the Lite subset with 32 samples.
|
||||
When computational budgets are expanded to accommodate 3200 samples, performance on the Lite subset demonstrates a marked improvement to 61.88 % 61.88\% .
|
||||
The 7B parameter variant of DeepSeek-Prover-V2 also exhibits notable effectiveness, outperforming previous baselines including Goedel-Prover-SFT ( Lin et al., 2025 ) , STP ( Dong and Ma, 2025 ) , and Kimina-Prover-Preview-Distill-7B ( Wang et al., 2025 ) .
|
||||
## 3.5 ProverBench: Formalization of AIME and Textbook Problems
|
||||
To enhance existing benchmarks and advance research in formal theorem proving, we introduce a benchmark dataset comprising 325 problems.
|
||||
Of these, 15 are formalized from number theory and algebra questions featured in the recent AIME competitions (AIME 24 and 25), offering authentic high-school competition-level challenges.
|
||||
The remaining 310 problems are drawn from curated textbook examples and educational tutorials, contributing a diverse and pedagogically grounded collection of formalized mathematical problems.
|
||||
This benchmark is designed to enable more comprehensive evaluation across both high-school competition problems and undergraduate-level mathematics.
|
||||
Method
|
||||
Model size
|
||||
Sample budget
|
||||
ProverBench
|
||||
All
|
||||
AIME 24&25
|
||||
STP ( Dong and Ma, 2025 )
|
||||
7B
|
||||
32 32
|
||||
27.5 % ± 0.7 % 27.5\%\pm 0.7\%
|
||||
0 / 15 0/15
|
||||
128 128
|
||||
31.4 % ± 1.1 % 31.4\%\pm 1.1\%
|
||||
1 / 15 1/15
|
||||
512 512
|
||||
36.3 % 36.3\%
|
||||
1 / 15 1/15
|
||||
DeepSeek-Prover-V2 (non-CoT)
|
||||
7B
|
||||
32 32
|
||||
47.7 % ± 0.6 % 47.7\%\pm 0.6\%
|
||||
1 / 15 1/15
|
||||
128 128
|
||||
48.8 % ± 0.2 % 48.8\%\pm 0.2\%
|
||||
1 / 15 1/15
|
||||
512 512
|
||||
49.5 % 49.5\%
|
||||
1 / 15 1/15
|
||||
671B
|
||||
32 32
|
||||
49.5 % ± 0.5 % 49.5\%\pm 0.5\%
|
||||
1 / 15 1/15
|
||||
128 128
|
||||
51.5 % ± 0.3 % 51.5\%\pm 0.3\%
|
||||
2 / 15 2/15
|
||||
512 512
|
||||
52.3 % 52.3\%
|
||||
2 / 15 2/15
|
||||
DeepSeek-Prover-V2 (CoT)
|
||||
7B
|
||||
32 32
|
||||
49.0 % ± 0.3 % 49.0\%\pm 0.3\%
|
||||
1 / 15 1/15
|
||||
128 128
|
||||
50.8 % ± 0.5 % 50.8\%\pm 0.5\%
|
||||
1 / 15 1/15
|
||||
512 512
|
||||
51.7 % 51.7\%
|
||||
1 / 15 1/15
|
||||
671B
|
||||
32 32
|
||||
52.9 % ± 0.9 % 52.9\%\pm 0.9\%
|
||||
4 / 15 4/15
|
||||
128 128
|
||||
56.5 % ± 0.5 % 56.5\%\pm 0.5\%
|
||||
5 / 15 5/15
|
||||
512 512
|
||||
59.1 % \mathbf{59.1\%}
|
||||
𝟔 / 𝟏𝟓 \mathbf{6/15}
|
||||
Table 7: The experimental results on ProverBench. The All category represents the complete evaluation set consisting of 325 problems, while AIME 24&25 denotes a subset of 15 problems formalized from recent AIME competitions. The results for STP ( Dong and Ma, 2025 ) are evaluated using the open-source model weights.
|
||||
Contest
|
||||
Problems
|
||||
AIME 24I
|
||||
P2 , P7 , P13
|
||||
AIME 24II
|
||||
P4 , P7, P13 , P14
|
||||
AIME 25I
|
||||
P1 , P8 , P9, P11
|
||||
AIME 25II
|
||||
P2 , P4 , P13, P15
|
||||
Table 8: Selection of AIME 24&25 problems for formalization. Problems with underlined bolded indices have been solved by DeepSeek-Prover-V2. Problems solved by DeepSeek-V3-0324 using Maj@16 are highlighted with a gray background.
|
||||
## AIME Formalization.
|
||||
The American Invitational Mathematics Examination (AIME) is an annual mathematics competition designed to challenge and recognize talented high school students who demonstrate exceptional proficiency in mathematics.
|
||||
The problems from AIME 24&25 have become a standard benchmark for evaluating the reasoning capabilities of large language models.
|
||||
In order to bridge the evaluation of model performance across formal and informal mathematical reasoning, we curate and formalize a subset of problems from AIME 24&25.
|
||||
To ensure cleaner formalizations, we filter out geometry, combinatorics, and counting problems whose representations in Lean are potentially cumbersome.
|
||||
This results in 15 selected problems, covering competition-level topics in elementary number theory and algebra.
|
||||
We evaluate DeepSeek-V3-0324 on the selected set of problems using the standard find-answer task for natural-language mathematical reasoning.
|
||||
With majority voting over 16 sampled responses, the model successfully solves 8 out of 15 problems.
|
||||
In comparison, DeepSeek-Prover-V2-671B, operating under the formal proof generation setting with given correct answers, is able to construct valid formal proofs for 6 of 15 problems.
|
||||
This comparison highlights that the performance gap between informal mathematical reasoning and formal theorem proving is substantially narrowing, indicating growing alignment between linguistic understanding and formal logical rigor in advanced language models.
|
||||
Area
|
||||
Count
|
||||
AIME 24&25
|
||||
15
|
||||
Number Theory
|
||||
40
|
||||
Elementary Algebra
|
||||
30
|
||||
Linear Algebra
|
||||
50
|
||||
Abstract Algebra
|
||||
40
|
||||
Calculus
|
||||
90
|
||||
Real Analysis
|
||||
30
|
||||
Complex Analysis
|
||||
10
|
||||
Functional Analysis
|
||||
10
|
||||
Probability
|
||||
10
|
||||
Total
|
||||
325
|
||||
Table 9: Distribution of mathematical areas represented in ProverBench.
|
||||
## Textbook Formalization.
|
||||
In addition to AIME 24&25, we augment our benchmark with problems carefully selected from textbooks used in high school competitions and undergraduate-level courses to strengthen coverage in specific mathematical domains.
|
||||
This curation process ensures comprehensive representation across difficulty levels and topic areas.
|
||||
As a result, we formalize 310 problems that encompass a broad spectrum, ranging from elementary mathematics at the competition level to advanced topics typically encountered in undergraduate studies.
|
||||
This comprehensive benchmark covers number theory, elementary algebra, linear algebra, abstract algebra, calculus, real analysis, complex analysis, functional analysis, and probability.
|
||||
The deliberate inclusion of this diverse array of mathematical fields allows for a thorough assessment of model capabilities across varying levels of abstraction and reasoning styles.
|
||||
Number theory and algebra problems test a model’s facility with discrete structures and equations, while analysis-oriented problems evaluate understanding of limits, continuity, and calculus.
|
||||
The abstract algebra and functional analysis components challenge models to reason about abstract structures and spaces, requiring sophisticated formal reasoning capabilities.
|
||||
The evaluation results are presented in Table 7 . As shown, DeepSeek-Prover-V2-671B with CoT reasoning consistently outperforms all baselines, reinforcing the trends observed in other benchmark evaluations.
|
||||
## 4 Conclusion
|
||||
In this work, we propose a comprehensive pipeline for synthesizing cold-start reasoning data to advance formal theorem proving.
|
||||
Our data construction process is grounded in a recursive theorem-proving framework, wherein DeepSeek-V3 serves as a unified model for both subgoal decomposition and lemma formalization within the Lean 4 proof assistant.
|
||||
Our approach combines high-level proof sketches with formal steps, creating a sequence of manageable subgoals that can be efficiently solved using a smaller 7B model, significantly reducing computational requirements.
|
||||
The curriculum learning framework we developed uses these decomposed subgoals to generate increasingly difficult training tasks, creating a more effective learning progression.
|
||||
By pairing complete formal proofs with DeepSeek-V3’s chain-of-thought reasoning, we established valuable cold-start reasoning data that bridges informal mathematical thinking with formal proof structures.
|
||||
The subsequent reinforcement learning stage substantially enhanced this connection, leading to significant improvements in formal theorem proving capabilities.
|
||||
The resulting model, DeepSeek-Prover-V2-671B, consistently outperforms all baselines across a range of benchmarks, spanning both high-school competition problems and undergraduate-level mathematics.
|
||||
Our future work will focus on scaling this paradigm to an AlphaProof-like system with the ultimate aim of tackling IMO-level mathematical problems that represent the frontier of automated theorem proving challenges.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,610 @@
|
||||
TITLE: Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures
|
||||
ARXIV: 2505.09343
|
||||
AUTHORS: Chenggang Zhao , Chengqi Deng , Chong Ruan , Damai Dai , Huazuo Gao , Jiashi Li , Liyue Zhang , Panpan Huang , Shangyan Zhou , Shirong Ma , Wenfeng Liang , Ying He , Yuqing Wang , Yuxuan Liu , Y.X. Wei
|
||||
DATE: [Submitted on 14 May 2025 ( v1 ), last revised 23 Dec 2025 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures
|
||||
## Abstract.
|
||||
## Keywords:
|
||||
## 1. Introduction
|
||||
## 1.1. Background
|
||||
## 1.2. Objectives
|
||||
## 1.3. Structure of this Paper
|
||||
## 2. Design Principles for DeepSeek Models
|
||||
## 2.1. Memory Efficiency
|
||||
## 2.1.1. Low-Precision Models
|
||||
## 2.1.2. Reducing KV Cache with MLA
|
||||
## 2.1.3. Future Directions and Perspectives on Resource-Efficient Techniques
|
||||
## 2.2. Cost-Effectiveness of MoE Models
|
||||
## 2.2.1. Reducing Computational Requirements for Training
|
||||
## 2.2.2. Advantages for Personal Use and On-Premises Deployment
|
||||
## 2.3. Increasing Inference Speed
|
||||
## 2.3.1. Overlapping Computation and Communication: Maximizing Throughput
|
||||
## 2.3.2. Inference Speed Limits
|
||||
## 2.3.3. Multi-Token Prediction
|
||||
## 2.3.4. High Inference Speed for Reasoning Models and Test-Time Scaling
|
||||
## 2.4. Technique Validation Methodology
|
||||
## 3. Low-Precision Driven Design
|
||||
## 3.1. FP8 Mix-Precision Training
|
||||
## 3.1.1. Limitations:
|
||||
## 3.1.2. Suggestions:
|
||||
## 3.2. LogFMT: Communication Compression
|
||||
## 3.2.1. Limitations:
|
||||
## 3.2.2. Suggestions:
|
||||
## 4. Interconnection Driven Design
|
||||
## 4.1. Current Hardware Architecture
|
||||
## 4.2. Hardware-Aware Parallelism
|
||||
## 4.3. Model Co-Design: Node-Limited Routing
|
||||
## 4.4. Scale-Up and Scale-Out Convergence
|
||||
## 4.4.1. Limitations of Current Implementations
|
||||
## 4.4.2. Suggestions:
|
||||
## 4.5. Bandwidth Contention and Latency
|
||||
## 4.5.1. Limitations:
|
||||
## 4.5.2. Suggestions:
|
||||
## 5. Large Scale Network Driven Design
|
||||
## 5.1. Network Co-Design: Multi-Plane Fat-Tree
|
||||
## 5.1.1. Advantages of Multi-Plane Fat-Tree Network
|
||||
## 5.1.2. Performance Analysis
|
||||
## 5.2. Low Latency Networks
|
||||
## 5.2.1. IB or RoCE
|
||||
## 5.2.2. Recommendations for RoCE Improvements
|
||||
## 5.2.3. InfiniBand GPUDirect Async (IBGDA)
|
||||
## 6. Discussion and Insights for Future Hardware Architecture Design
|
||||
## 6.1. Robustness Challenges
|
||||
## 6.1.1. Limitations:
|
||||
## 6.1.2. Suggestions for Advanced Error Detection and Correction
|
||||
## 6.2. CPU Bottlenecks and Interconnects
|
||||
## 6.3. Toward Intelligent Networks for AI
|
||||
## 6.4. Discussion on Memory-Semantic Communication and Ordering Issue
|
||||
## 6.5. In-Network Computation and Compression
|
||||
## 6.6. Memory-Centric Innovations
|
||||
## 6.6.1. Limitations of Memory Bandwidth
|
||||
## 6.6.2. Suggestions:
|
||||
## 7. Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract.
|
||||
1 Introduction
|
||||
1.1 Background
|
||||
1.2 Objectives
|
||||
1.3 Structure of this Paper
|
||||
2 Design Principles for DeepSeek Models
|
||||
2.1 Memory Efficiency
|
||||
2.1.1 Low-Precision Models
|
||||
2.1.2 Reducing KV Cache with MLA
|
||||
2.1.3 Future Directions and Perspectives on Resource-Efficient Techniques
|
||||
2.2 Cost-Effectiveness of MoE Models
|
||||
2.2.1 Reducing Computational Requirements for Training
|
||||
2.2.2 Advantages for Personal Use and On-Premises Deployment
|
||||
2.3 Increasing Inference Speed
|
||||
2.3.1 Overlapping Computation and Communication: Maximizing Throughput
|
||||
2.3.2 Inference Speed Limits
|
||||
2.3.3 Multi-Token Prediction
|
||||
2.3.4 High Inference Speed for Reasoning Models and Test-Time Scaling
|
||||
2.4 Technique Validation Methodology
|
||||
3 Low-Precision Driven Design
|
||||
3.1 FP8 Mix-Precision Training
|
||||
3.1.1 Limitations:
|
||||
3.1.2 Suggestions:
|
||||
3.2 LogFMT: Communication Compression
|
||||
3.2.1 Limitations:
|
||||
3.2.2 Suggestions:
|
||||
4 Interconnection Driven Design
|
||||
4.1 Current Hardware Architecture
|
||||
4.2 Hardware-Aware Parallelism
|
||||
4.3 Model Co-Design: Node-Limited Routing
|
||||
4.4 Scale-Up and Scale-Out Convergence
|
||||
4.4.1 Limitations of Current Implementations
|
||||
4.4.2 Suggestions:
|
||||
4.5 Bandwidth Contention and Latency
|
||||
4.5.1 Limitations:
|
||||
4.5.2 Suggestions:
|
||||
5 Large Scale Network Driven Design
|
||||
5.1 Network Co-Design: Multi-Plane Fat-Tree
|
||||
5.1.1 Advantages of Multi-Plane Fat-Tree Network
|
||||
5.1.2 Performance Analysis
|
||||
5.2 Low Latency Networks
|
||||
5.2.1 IB or RoCE
|
||||
5.2.2 Recommendations for RoCE Improvements
|
||||
5.2.3 InfiniBand GPUDirect Async (IBGDA)
|
||||
6 Discussion and Insights for Future Hardware Architecture Design
|
||||
6.1 Robustness Challenges
|
||||
6.1.1 Limitations:
|
||||
6.1.2 Suggestions for Advanced Error Detection and Correction
|
||||
6.2 CPU Bottlenecks and Interconnects
|
||||
6.3 Toward Intelligent Networks for AI
|
||||
6.4 Discussion on Memory-Semantic Communication and Ordering Issue
|
||||
6.5 In-Network Computation and Compression
|
||||
6.6 Memory-Centric Innovations
|
||||
6.6.1 Limitations of Memory Bandwidth
|
||||
6.6.2 Suggestions:
|
||||
7 Conclusion
|
||||
References
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2505.09343v2 [cs.DC] 23 Dec 2025
|
||||
## Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures
|
||||
Conference: Proceedings of the 52nd Annual International Symposium on Computer Architecture; June 21–25, 2025; Tokyo, Japan Proceedings of the 52nd Annual International Symposium on Computer Architecture (ISCA ’25), June 21–25, 2025, Tokyo, Japan DOI: 10.1145/3695053.3731412 ISBN: 979-8-4007-1261-6/2025/06 CCS: Computer systems organization Architectures
|
||||
Chenggang Zhao
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: chenggangz@deepseek.com
|
||||
,
|
||||
Chengqi Deng
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: cq.deng@deepseek.com
|
||||
,
|
||||
Chong Ruan
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: chong.ruan@deepseek.com
|
||||
,
|
||||
Damai Dai
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: damai.dai@deepseek.com
|
||||
,
|
||||
Huazuo Gao
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: gaohuazuo@deepseek.com
|
||||
,
|
||||
Jiashi Li
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: js.li@deepseek.com
|
||||
,
|
||||
Liyue Zhang
|
||||
Note: Yuqing Wang and Liyue Zhang are the corresponding authors of this paper. Authors are listed in alphabetical order of their first names.
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: ly.zhang@deepseek.com
|
||||
,
|
||||
Panpan Huang
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: pp.huang@deepseek.com
|
||||
,
|
||||
Shangyan Zhou
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: sy.zhou@deepseek.com
|
||||
,
|
||||
Shirong Ma
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: mashirong.2000@deepseek.com
|
||||
,
|
||||
Wenfeng Liang
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: wenfeng.liang@deepseek.com
|
||||
,
|
||||
Ying He
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: ying.he@deepseek.com
|
||||
,
|
||||
Yuqing Wang
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: wangyq@deepseek.com
|
||||
,
|
||||
Yuxuan Liu
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: liuyuxuan@deepseek.com
|
||||
and
|
||||
Y.X. Wei
|
||||
Affiliation: DeepSeek-AI , Beijing , China
|
||||
email: weiyx@deepseek.com
|
||||
© acmlicensed
|
||||
## Abstract.
|
||||
The rapid scaling of large language models (LLMs) has unveiled critical limitations in current hardware architectures, including constraints in memory capacity, computational efficiency, and interconnection bandwidth. DeepSeek-V3, trained on 2,048 NVIDIA H800 GPUs, demonstrates how hardware-aware model co-design can effectively address these challenges, enabling cost-efficient training and inference at scale. This paper presents an in-depth analysis of the DeepSeek-V3/R1 model architecture and its AI infrastructure, highlighting key innovations such as Multi-head Latent Attention (MLA) for enhanced memory efficiency, Mixture of Experts (MoE) architectures for optimized computation-communication trade-offs, FP8 mixed-precision training to unlock the full potential of hardware capabilities, and a Multi-Plane Network Topology to minimize cluster-level network overhead. Building on the hardware bottlenecks encountered during DeepSeek-V3’s development, we engage in a broader discussion with academic and industry peers on potential future hardware directions, including precise low-precision computation units, scale-up and scale-out convergence, and innovations in low-latency communication fabrics. These insights underscore the critical role of hardware and model co-design in meeting the
|
||||
escalating demands of AI workloads, offering a practical blueprint for innovation in next-generation AI systems.
|
||||
## Keywords:
|
||||
Large Language Model, Mixture-of-Experts, Deep Learning, FP8 Mixed-Precision Training, Multi-Plane Network, Co-Design
|
||||
## 1. Introduction
|
||||
## 1.1. Background
|
||||
Large Language Models (LLMs) have undergone rapid evolution in recent years, driven by iterative advancements in model design, computational power, and data availability. In 2024, groundbreaking models such as GPT4o ( OpenAI, 2024a ) , LLaMa-3 ( AI@Meta, 2024a ) , Claude 3.5 Sonnet ( Anthropic, 2024 ) , Grok-2 ( xAI, 2024a ) , Qwen2.5 ( Team, 2024 ) , Gemini-2 ( Google, 2024 ) and our DeepSeek-V3 ( DeepSeek-AI, 2024d ) have showcased remarkable progress, further narrowing the gap towards Artificial General Intelligence (AGI). As the Scaling Laws ( Kaplan et al., 2020 ) shows, increasing model size, training data, and computational resources leads to substantial improvements in model performance, underscoring the pivotal role of scaling in advancing AI capabilities. Collectively, these developments have ushered in an era where scaling model size and computational power is seen as the key to unlocking higher levels of intelligence.
|
||||
Recent developments, reasoning models such as OpenAI’s o1/o3 series models ( OpenAI, 2024b ; OpenAI, 2025 ) , DeepSeek-R1 ( DeepSeek-AI, 2025a ) , Claude-3.7 Sonnet ( Anthropic, 2025 ) , Gemini 2.5 Pro ( Google, 2025 ) , Seed1.5-Thinking ( Seed, 2025 ) and Qwen3 ( Team, 2025 ) have demonstrated not only the benefits conferred by large-scale architectures, but also the necessity of improving inference efficiency, particularly in handling longer contexts and achieving greater reasoning depth. These advancements underscore the need for faster and more efficient inference, consequently placing ever-increasing demands on computational resources.
|
||||
To meet these challenges, industry leaders such as Alibaba, ByteDance, Google, xAI and Meta have deployed colossal training clusters ( Jouppi et al., 2023 ; Mudigere et al., 2023 ; Gangidi et al., 2024 ; Jiang et al., 2024 ; Qian et al., 2024 ; xAI, 2024b ) , featuring tens or even hundreds of thousands of GPUs or TPUs. While such massive infrastructures have enabled the development of state-of-the-art models, their exorbitant costs present significant barriers for smaller research teams and organizations. Despite these barriers, open-source startups such as DeepSeek ( DeepSeek-AI, 2024a ; DeepSeek-AI, 2024c ; DeepSeek-AI, 2024d ; DeepSeek-AI, 2024b ; DeepSeek-AI, 2025a ) and Mistral ( Jiang et al., 2023 ; Mistral, 2024 ) are also striving to develop state-of-the-art models.
|
||||
Among them, DeepSeek has especially demonstrated that effective software-hardware co-design can enable cost-efficient training of large models, leveling the playing field for smaller teams.
|
||||
Building on this tradition, DeepSeek-V3 ( DeepSeek-AI, 2024d ) represents a new milestone in cost-effective training. By leveraging just 2,048 NVIDIA H800 GPUs, DeepSeek-V3 achieves state-of-the-art performance. This achievement aligns with the commitment to advance AI through practical and scalable solutions, as previously demonstrated in the cost-effective architecture of Fire-Flyer AI-HPC ( An et al., 2024 ) . The practices and insights derived from DeepSeek-V3 demonstrate how existing hardware resources can be harnessed to their fullest potential, offering valuable lessons for the broader AI and HPC communities.
|
||||
## 1.2. Objectives
|
||||
This paper does not aim to reiterate the detailed architectural and algorithmic specifics of DeepSeek-V3, which are extensively documented in its technical report ( DeepSeek-AI, 2024d ) . Instead, it adopts a dual perspective—spanning hardware architecture and model design—to explore the intricate interplay between them in achieving cost-efficient large-scale training and inference. By examining this synergy, we aim to provide actionable insights for scaling LLMs efficiently without sacrificing performance or accessibility.
|
||||
Specifically, the paper focuses on:
|
||||
•
|
||||
Hardware-Driven Model Design: Analyze how hardware features, such as FP8 low-precision computation and scale-up/scale-out network properties, informed the architectural choices in DeepSeek-V3.
|
||||
•
|
||||
Mutual Dependencies Between Hardware and Models: Investigate how hardware capabilities shape model innovation and how the evolving demands of LLMs drive the need for next-generation hardware.
|
||||
•
|
||||
Future Directions for Hardware Development: Derive actionable insights from DeepSeek-V3 to guide the co-design of future hardware and model architectures, paving the way for scalable, cost-efficient AI systems.
|
||||
## 1.3. Structure of this Paper
|
||||
The remainder of this paper is organized as follows. Section 2 explores the design principles underpinning DeepSeek-V3 model architecture, highlighting key innovations such as Multi-head Latent Attention, Mixture-of-Experts optimizations and Multi-Token Prediction Module. Section 3 illustrates how our model architecture pursues low-precision computation and communication. Section 4 includes scale-up interconnection optimizations, discusses scale-up/scale-out convergence, and explores how hardware features influence parallelism and expert selection strategies. Section 5 focuses on scale-out network optimizations, including multi-plane network co-designs and low-latency interconnects. Besides current limitations and future suggestions mentioned in Section 3 ∼ \sim 5 , Section 6 elaborates on more critical insights from DeepSeek-V3, and identifies directions for future hardware and model co-design.
|
||||
Figure 1 .
|
||||
Basic architecture of DeepSeek-V3. Built upon DeepSeek-V2’s MLA and DeepSeekMoE, a Multi-Token Prediction Module and FP8 mixed-precision training are introduced to enhance inference and training efficiency. The figure indicates the precision used for computations in different parts of the architecture. All components take inputs and outputs in BF16.
|
||||
## 2. Design Principles for DeepSeek Models
|
||||
The development of DeepSeek-V3 exemplifies a hardware-aware approach to scaling LLMs, where each design decision was carefully aligned with hardware constraints to optimize performance and cost efficiency.
|
||||
As shown in Figure 1 , DeepSeek-V3 employs the
|
||||
DeepSeekMoE ( DeepSeek-AI, 2024e ) and Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024c ) architectures that have been proven effective in DeepSeek-V2 ( DeepSeek-AI, 2024c ) .
|
||||
DeepSeekMoE unlocks the potential of MoE architecture, while MLA drastically reduces memory consumption by compressing Key-Value (KV) caches.
|
||||
In addition, DeepSeek-V3 incorporates FP8 mixed-precision training , significantly lowering computational costs and making large-scale training more practical without compromising model quality. To improve the inference speed, DeepSeek-V3 integrates speculative decoding based on its Multi-Token Prediction Module , which significantly increases the generation speed. Beyond model architecture, we also explored cost-efficient AI infrastructure by deploying a Multi-Plane two-layer Fat-Tree network to replace a traditional three-layer Fat-Tree topology, reducing cluster networking costs.
|
||||
These innovations aim to address three core challenges in scaling LLMs— memory efficiency , cost-effectiveness , and inference speed —which are explored in detail in the following subsections.
|
||||
## 2.1. Memory Efficiency
|
||||
LLMs generally require significant memory resources, with memory demands increasing by more than 1000% per year. In contrast, the growth rate of high-speed memory (e.g., HBM) capacity is much slower, typically less than 50% per year ( Gholami et al., 2024 ) . While multi-node parallelism is a viable solution to address memory limitations, optimizing memory usage at the source remains a crucial and effective strategy.
|
||||
## 2.1.1. Low-Precision Models
|
||||
Compared to models that utilize BF16 for weights, FP8 significantly reduces memory consumption by half, effectively alleviating the AI memory wall challenge. A detailed discussion of low-precision techniques is provided in Section 3 Low-Precision Driven Design.
|
||||
## 2.1.2. Reducing KV Cache with MLA
|
||||
For LLM inference, user requests often involve multi-turn conversations. To handle these efficiently, the context from previous requests is cached in what is commonly referred to as the KV cache . KV cache addresses this challenge by caching the Key and Value vectors of previously processed tokens, eliminating the need to recompute them for subsequent tokens. During each inference step, the model only computes the Key and Value vectors for the current token and performs attention computation by combining them with the cached Key-Value pairs from the history. This incremental computation reduces the complexity of generating each token to O ( N ) O(N) , making it efficient when processing long sequences or multi-turn inputs. However, it introduces a memory-bound bottleneck because the computation shifts from GEMM to GEMV, which has a much lower compute-to-memory ratio. With modern hardware offering hundreds of TFLOPS, GEMV quickly becomes limited by memory bandwidth, making memory access the primary bottleneck.
|
||||
To address this bottleneck, we employ Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024c ) that compresses the KV representations of all attention heads into a smaller latent vector using a projection matrix, which is jointly trained with the model. During inference, only the latent vector needs to be cached, significantly reducing memory consumption compared to storing the KV cache for all attention heads.
|
||||
In addition to MLA, several other approaches have been proposed to reduce the size of the KV cache. These methods are highly valuable and provide significant inspiration for advancements in memory-efficient attention mechanisms:
|
||||
•
|
||||
Shared KV (Grouped-Query Attention, GQA; Multi-Query Attention, MQA): Instead of maintaining separate KV pairs for each attention head, multiple heads share a single set of KV pairs, significantly compressing KV storage. Representative methods include GQA ( Ainslie et al., 2023 ) and MQA ( Shazeer, 2019 ) .
|
||||
•
|
||||
Windowed KV: For long sequences, only a sliding window of KV pairs is retained in the cache, discarding results outside the window. While this reduces storage, it compromises long-context reasoning. Representative methods include Longformer ( Beltagy et al., 2020 ) and related architectures.
|
||||
•
|
||||
Quantized Compression: KV pairs are stored using low-bit representations ( Hooper et al., 2024 ; Liu et al., 2024 ; Kang et al., 2024 ) , further reducing memory usage. Quantization achieves significant compression with minimal impact on model performance.
|
||||
Table 1. KV cache size comparison (BF16 precision): DeepSeek-V3 (MLA) largely reduces KV cache size compared to other models using GQA.
|
||||
Model
|
||||
KV Cache Per Token
|
||||
Multiplier
|
||||
DeepSeek-V3 (MLA)
|
||||
70.272 KB
|
||||
1x
|
||||
Qwen-2.5 72B (GQA)
|
||||
327.680 KB
|
||||
4.66x
|
||||
LLaMA-3.1 405B (GQA)
|
||||
516.096 KB
|
||||
7.28x
|
||||
Table 1 compares the KV cache memory usage per token among DeepSeek-V3, Qwen-2.5 72B ( Team, 2024 ) , and LLaMA-3.1 405B ( AI@Meta, 2024b ) . By adopting MLA, DeepSeek-V3 achieves a significant reduction in KV cache size, requiring only 70 KB per token, substantially less than LLaMA-3.1 405B’s 516 KB and Qwen-2.5 72B’s 327 KB.
|
||||
This reduction highlights the efficiency of MLA in compressing KV representations compared to GQA-based methods. The ability to achieve such a significant reduction in memory consumption makes DeepSeek-V3 particularly well-suited for scenarios involving long-context processing and resource-constrained environments, enabling more scalable and cost-effective inference.
|
||||
## 2.1.3. Future Directions and Perspectives on Resource-Efficient Techniques
|
||||
While reducing the size of the KV cache is a promising method for improving memory efficiency, the quadratic complexity inherent in Transformer-based autoregressive decoding remains a formidable challenge, especially for extremely long contexts. Recent research efforts, such as Mamba-2 ( Dao and Gu, 2024 ) and Lightning Attention ( Qin et al., 2024 ) , investigate linear-time alternatives that offer new possibilities for balancing computational cost and model performance. In addition, approaches such as sparse attention ( Yuan et al., 2025 ) , which seek to compress and sparsely activate attention keys and values, represent another attempt at overcoming the computational challenges associated with attention. We look forward to collaborative progress with the broader community toward breakthroughs in this area.
|
||||
## 2.2. Cost-Effectiveness of MoE Models
|
||||
For sparse computing, we have developed DeepSeekMoE, an advanced Mixture of Experts (MoE) architecture, which is illustrated in the lower right part of Figure 1 .
|
||||
The advantages of MoE models lie in two folds.
|
||||
## 2.2.1. Reducing Computational Requirements for Training
|
||||
The primary advantage of the MoE architecture lies in its ability to significantly reduce training costs. By selectively activating only a subset of expert parameters, MoE models allow the total parameter count to scale up dramatically while keeping computational requirements modest. For example, DeepSeek-V2 features 236B parameters, but only 21B parameters are activated per token. Similarly, DeepSeek-V3 expands to 671B parameters—nearly three times the size of V2—while keeping the activation per token at just 37B. In comparison, dense models such as Qwen2.5-72B and LLaMa3.1-405B require all parameters to be active during training.
|
||||
As shown in Table 2 , the total computational cost for DeepSeek-V3 is approximately 250 GFLOPS per token, whereas the 72B dense model requires 394 GFLOPS and the 405B dense model requires 2448 GFLOPS. This demonstrates that MoE models achieve comparable or even superior performance to dense models while consuming an order of magnitude less computational resources.
|
||||
## 2.2.2. Advantages for Personal Use and On-Premises Deployment
|
||||
In a future where personalized LLM agents ( Luo et al., 2025 ) become ubiquitous, MoE models offer unique advantages in single-request scenarios. Because only a subset of parameters is activated per request, memory and computational demands are greatly reduced. For example, DeepSeek-V2 (236B parameters) activates just 21B parameters during inference.
|
||||
This enables PCs with AI SoC chips ( Apple, 2024 ; NVIDIA, 2025 ; AMD, 2025 ) to achieve nearly 20 tokens per second (TPS), or even twice that speed, which is more than sufficient for personal use. In contrast, dense models of similar capability (e.g., 70B parameters) typically reach only single-digit TPS on similar hardware.
|
||||
Notably, the increasingly popular KTransformers ( group and Approaching.AI, 2025 ) inference engine allows the complete DeepSeek-V3 model to run on a low-cost server equipped with a consumer GPU (costing approximately $10,000), while still achieving nearly 20 TPS.
|
||||
This efficiency makes MoE architectures suitable for local deployments and single-user scenarios, where hardware resources are often limited. By minimizing memory and computational overhead, MoE models can deliver high-quality inference performance without requiring expensive infrastructure.
|
||||
Table 2. Comparison of computational costs for training MoE and dense models: Computational cost per token is measured, assuming a sequence length of 4096.
|
||||
Model
|
||||
Size
|
||||
Training Cost
|
||||
DeepSeek-V2 MoE
|
||||
236B
|
||||
155 GFLOPS/Token
|
||||
DeepSeek-V3 MoE
|
||||
671B
|
||||
250 GFLOPS/Token
|
||||
Qwen-72B Dense
|
||||
72B
|
||||
394 GFLOPS/Token
|
||||
LLaMa-405B Dense
|
||||
405B
|
||||
2448 GFLOPS/Token
|
||||
## 2.3. Increasing Inference Speed
|
||||
## 2.3.1. Overlapping Computation and Communication: Maximizing Throughput
|
||||
Inference speed encompasses both system-wide maximum throughput and single-request latency. To maximize throughput, our model is architected from the outset to leverage dual micro-batch overlap ( DeepSeek-AI, 2025d ; Zhao et al., 2025b ) , intentionally overlapping communication latency with computation. As demonstrated in our online inference system and supported by open-source profiling data ( DeepSeek-AI, 2025d ) , we decouple the computation of MLA and MoE into two distinct stages. While one micro-batch executes a portion of MLA or MoE computation, the other micro-batch simultaneously performs the corresponding dispatch communication. Conversely, during the computation phase of the second micro-batch, the first micro-batch undergoes the combine communication step. This pipelined approach enables seamless overlap of all-to-all communication with ongoing computation, ensuring that the GPU remains fully utilized at all times. Moreover, in production, we adopt a prefill and decode disaggregation architecture ( Zhong et al., 2024 ) , assigning large batch size prefill and latency-sensitive decode requests to different expert parallelism group sizes. This strategy ultimately maximizes system throughput under real-world service conditions.
|
||||
## 2.3.2. Inference Speed Limits
|
||||
This section focuses on the decode output speed of LLM services, typically measured in Time Per Output Token (TPOT) . TPOT is a critical metric for user experience, and it also directly impacts the responsiveness of reasoning models such as OpenAI’s o1/o3 and DeepSeek-R1, which rely on the inference length to enhance their intelligence.
|
||||
For MoE models, achieving high inference speed relies on efficiently deploying expert parameters across computing devices. To achieve the fastest possible inference speed, each device should ideally perform computations for a single expert (or multiple devices should collaboratively compute a single expert if necessary). However, Expert Parallelism (EP) requires routing tokens to the appropriate devices, which involves all-to-all communication across the network. As a result, the upper limit of MoE inference speed is dictated by interconnection bandwidth.
|
||||
Consider a system where each device holds one expert’s parameters and processes approximately 32 tokens at a time. This token count strikes a balance between compute-to-memory ratio and communication latency. And this token count ensures that each device processes an equal batch size during expert parallelism, allowing the communication time to be easily calculated.
|
||||
For a system interconnected with CX7 400Gbps InfiniBand (IB) NICs, the time required for the two all-to-all communications in EP is calculated as follows:
|
||||
Comm. Time = ( 1 Byte + 2 Bytes ) × 32 × 9 × 7 K / 50 GB/s = 120.96 μ s \text{Comm. Time}=(1\text{Byte}+2\text{Bytes})\times 32\times 9\times 7\text{K}/50\text{GB/s}=120.96\mu s
|
||||
Here, dispatch uses FP8 (1 byte), while combine uses BF16 (2 bytes), and the hidden size of each token is approximately 7K. The factor 9 indicates that each token is transferred to 8 routed experts and 1 shared expert. Network latency is not included in this calculation.
|
||||
As discussed in Section 2.3.1 , maximizing throughput necessitates the use of dual micro-batch overlap. In this strategy, our theoretical best-case analysis assumes that computation overhead is minimized, so the upper bound on performance is determined by communication latency. In practical inference workloads, however, request contexts are often much longer, and MLA computations typically dominate execution time. Thus, this analysis represents an idealized scenario under dual micro-batch overlap. Under this assumption, the total time per layer can be formulated as:
|
||||
Total Time Per Layer = 2 × 120.96 μ s = 241.92 μ s \text{Total Time Per Layer}=2\times 120.96\mu s=241.92\mu s
|
||||
With 61 layers in DeepSeek-V3, the total inference time is:
|
||||
Total Inference Time = 61 × 241.92 μ s = 14.76 ms \text{Total Inference Time}=61\times 241.92\mu s=14.76\text{ms}
|
||||
Thus, the theoretical upper limit for this system is approximately 14.76 ms TPOT , equivalent to 67 tokens per second . However, in practice, factors such as communication overhead, latency, incomplete bandwidth utilization, and computational inefficiencies reduce this number.
|
||||
By contrast, if a high-bandwidth interconnect like GB200 NVL72 (900GB/s unidirectional bandwidth across 72 GPUs) were used, the communication time per EP step drops to:
|
||||
Comm. Time = ( 1 Byte + 2 Bytes ) × 32 × 9 × 7 K / 900 GB/s = 6.72 μ s \text{Comm. Time}=(1\text{Byte}+2\text{Bytes})\times 32\times 9\times 7\text{K}/900\text{GB/s}=6.72\mu s
|
||||
Assuming perfect overlap between computation and communication, this would yield a theoretical upper limit of 0.82 ms TPOT , or approximately 1200 tokens per second . However, this figure is purely theoretical and does not account for the substantial drop in GPU efficiency at small batch sizes; in real deployments, actual throughput will be significantly lower. Nonetheless, this calculation vividly illustrates the transformative potential of high-bandwidth scale-up networks in accelerating large-scale model inference.
|
||||
While MoE models exhibit good scalability, achieving high inference speeds by increasing hardware resources alone is cost-prohibitive.
|
||||
Therefore, software and algorithms must also contribute to improving inference efficiency.
|
||||
## 2.3.3. Multi-Token Prediction
|
||||
Inspired by Gloeckle et al. (2024) , DeepSeek-V3 introduces a Multi-Token Prediction (MTP) framework, which simultaneously enhances model performance and improves inference speed.
|
||||
During inference, traditional autoregressive models generate one token at a decoding step, leading to sequential bottlenecks.
|
||||
MTP mitigates this issue by enabling the model to generate additional candidate tokens at a lower cost and verify them in parallel, similar to previous self-drafting-based speculative decoding approaches ( Cai et al., 2024 ; Li et al., 2024 ) . This framework significantly accelerates inference without compromising accuracy.
|
||||
As illustrated in the top part of Figure 1 , each MTP module uses a single layer, which is much more lightweight than the full model, to predict additional tokens, enabling parallel verification of multiple candidate tokens.
|
||||
Although slightly hurting the throughput, this approach significantly improves the end-to-end generation latency.
|
||||
The real world practice data demonstrates that an MTP module achieves an acceptance rate of 80% to 90% for predicting the second subsequent token, which increases the generation TPS by 1.8x compared to the scenario without the MTP module.
|
||||
Moreover, by predicting multiple tokens per step, MTP increases the inference batch size, which is crucial for boosting EP computational intensity and hardware utilization. Such algorithmic innovations are vital for fast and cost-effective inference in DeepSeek-V3.
|
||||
## 2.3.4. High Inference Speed for Reasoning Models and Test-Time Scaling
|
||||
Test-time scaling in LLMs, exemplified by OpenAI’s o1/o3 series ( OpenAI, 2024b ; OpenAI, 2025 ) , has enabled significant advances in mathematical reasoning, programming, and general reasoning by dynamically adjusting computational resources during inference. Subsequent models—including DeepSeek-R1 ( DeepSeek-AI, 2025a ) , Claude-3.7 Sonnet ( Anthropic, 2025 ) , Gemini 2.5 Pro ( Google, 2025 ) , Seed1.5-Thinking ( Seed, 2025 ) , and Qwen3 ( Team, 2025 ) —have adopted similar strategies and achieved notable improvements in these tasks.
|
||||
For these reasoning models, high token output speed is of paramount importance. In reinforcement learning (RL) workflows—such as PPO ( Schulman et al., 2017 ) , DPO ( Rafailov et al., 2024 ) and GRPO ( Shao et al., 2024 ) —the necessity to rapidly generate large numbers of samples makes inference throughput a critical bottleneck. Likewise, prolonged reasoning sequences can increase user wait times, reducing the practical usability of such models. As a result, optimizing inference speed through synergistic hardware and software innovations is indispensable for advancing the efficiency of reasoning models. However, effective strategies for accelerating inference and expediting RL training remain active areas of investigation, as discussed in Section 2.1.3 . We encourage the broader community to collaboratively explore and develop novel solutions to these ongoing challenges.
|
||||
## 2.4. Technique Validation Methodology
|
||||
Each acceleration technique undergoes rigorous empirical validation to evaluate its accuracy impact, including MLA, FP8 mixed-precision computation, and network co-designed MoE gate routing. Given the prohibitive cost of exhaustive ablation on full-scale models, we adopt a hierarchical and resource-efficient validation pipeline. Each technique is first validated extensively on small-scale models, followed by minimal large-scale tuning, and finally integrated in a single, comprehensive training run.
|
||||
For instance, we first conducted fine-grained FP8 training ablation studies on both 16B and 230B DeepSeek-V2 models before final integration. Under these controlled settings, the relative accuracy loss compared to BF16 remains below 0.25%, attributable to our use of high-precision accumulation and fine-grained quantization strategies.
|
||||
## 3. Low-Precision Driven Design
|
||||
## 3.1. FP8 Mix-Precision Training
|
||||
Quantization techniques such as GPTQ ( Frantar et al., 2022 ) and AWQ ( Lin et al., 2024 ) have been widely used to reduce bit-widths to 8-bit, 4-bit, or even lower, significantly reducing memory requirements. However, these techniques are primarily applied during inference to save memory, rather than in the training phase. NVIDIA’s Transformer Engine has supported FP8 mixed-precision training for some time, but prior to DeepSeek-V3, there were no open-source large models leveraging FP8 for training. Through deep collaboration between our infrastructure and algorithm teams, and after extensive experimentation and innovation, we developed an FP8-compatible training framework for MoE models. Figure 1 shows the computational components where FP8-precision forward and backward processes are utilized in the training pipeline. Fine-grained quantization is applied, i.e., tile-wise 1x128 quantization for activations and block-wise 128x128 quantization for model weights. Further technical details of our FP8 framework are documented in the DeepSeek-V3 technical report ( DeepSeek-AI, 2024d ) , and our fine-grained FP8 GEMM implementation has been open-sourced in DeepGEMM ( Zhao et al., 2025a ) .
|
||||
## 3.1.1. Limitations:
|
||||
While FP8 has great potential for accelerating training, several hardware limitations need to be addressed to fully exploit its capabilities:
|
||||
•
|
||||
FP8 Accumulation Precision: FP8 uses constrained accumulation precision in Tensor Cores, affecting the stability for training large models, particularly on NVIDIA Hopper GPUs. After aligning 32 mantissa products by right-shifting based on the maximum exponent, the Tensor Core only maintains their highest 13 fraction bits for addition, and truncates bits exceeding this range. Addition results are accumulated to FP22 registers (1 sign bit, 8 exponent bits, and 13 mantissa bits) ( Zhang et al., 2025 ) . The term “FP22” follows the naming used in the cited work, and the FP8 precision issue on Hopper GPUs had also been observed across industry by early 2024.
|
||||
•
|
||||
Fine-Grained Quantization Challenges: Fine-grained quantization such as tile-wise and block-wise quantization introduces large dequantization overhead in transporting the partial results from Tensor Cores to CUDA Cores for scaling factor multiplication. This incurs frequent data movements, reducing computational efficiency and complicating hardware utilization.
|
||||
## 3.1.2. Suggestions:
|
||||
To address the limitations of existing hardware, we have the following suggestions for future designs:
|
||||
•
|
||||
Increased Accumulation Precision: Hardware should improve the accumulation register precision to an appropriate value (e.g. FP32), or support a configurable accumulation precision, enabling a trade-off between performance and accuracy for different requirements of training and inference in various models.
|
||||
•
|
||||
Native Support for Fine-Grained Quantization: Hardware should natively support fine-grained quantization, enabling Tensor Cores to receive scaling factors and implement matrix multiplication with group scaling. In this way, the whole partial sum accumulation and dequantization can be completed directly inside Tensor Cores until the final result is produced, avoiding frequent data movements to reduce dequantization overhead. A notable industrial implementation of this approach is NVIDIA Blackwell’s support for microscaling data format ( Rouhani et al., 2023 ) , which exemplifies the practical benefits of native quantization at scale.
|
||||
## 3.2. LogFMT: Communication Compression
|
||||
In the current DeepSeek-V3 architecture, we employ low-precision compression for network communication. During EP parallelism, tokens are dispatched using fine-grained FP8 quantization, reducing communication volume by 50% compared to BF16. This significantly lowers communication time. While the combine stage still uses higher precision (e.g., BF16) due to accuracy requirements, we are actively testing FP8, custom precision formats (e.g., E5M6) and mixing FP8-BF16 for further reductions.
|
||||
Besides these traditional floating point formats, we also tried a new data type, named Logarithmic Floating-Point Formats (LogFMT-nBit) , where n n is the number of bits with the leading 1 bit as the sign bit S S . By mapping the activations from the original Linear space to the Log space, the distribution of the activations is more uniform. To be specific, given a tile of elements, [ x 1 , ⋯ , x m ] [x_{1},\cdots,x_{m}] , which is 1x128 in our implementation, we take the absolute values and compute the logarithm of all the elements, and find the minimum m i n = l o g ( a b s ( x i ) ) min=log(abs(x_{i})) and maximum m a x = l o g ( a b s ( x j ) ) max=log(abs(x_{j})) . The minimum is encoded as S .00 ⋯ 01 S.00\cdots 01 and the maximum is encoded as S .11 ⋯ 11 S.11\cdots 11 , with an interval representing S t e p = m a x − m i n 2 n − 1 − 2 Step=\frac{max-min}{2^{n-1}-2} . Zero values are represented by S .00 ⋯ 00 S.00\cdots 00 , specially. The left values are rounded to the nearest integer K K multiples of S t e p Step . The decoding process is simple by combining the sign bit and e x p m i n + S t e p × ( K − 1 ) exp^{min+Step\times(K-1)} .
|
||||
By locally calculating the m i n min and S t e p Step , this data type supports dynamic representation range for different blocks, covering larger ranges or providing more precision, compared to static floating point formats. Besides, we find it is important to round in the original Linear space, instead of the Log space, for the unbiased activation quantization. We also constrain the m i n min to be larger than m a x − l o g ( 2 32 ) max-log(2^{32}) , which means that the max representation range is similar to E5, a floating point with 5 exponents. We validate our LogFMT-nBit on dense language models with around 7 billion parameters, by quantifying the output of the residual branch to simulate the combine stage in MoE models. When setting n = 8 n=8 , sharing the same bits with FP8, the LogFMT-8Bit shows superior training accuracy compared to E4M3 or E5M2. After increasing the n n to 10 bits, we find it’s similar to the BF16 combine stage.
|
||||
## 3.2.1. Limitations:
|
||||
The initial purpose of using LogFMT is to apply it to activations during transmission or near activation functions, as it offers higher precision than FP8 with the same bit width. However, subsequent computations require reconversion to BF16 or FP8 to accommodate the Hopper GPU tensor cores’ data type. Due to insufficient GPU bandwidth for log/exp operations and excessive register pressure during encode/decode, if encode/decode operations are fused with all-to-all communication, the overhead can be substantial (50% ∼ \sim 100%). Therefore, although experimental results validate the effectiveness of this format, we do not employ it eventually.
|
||||
## 3.2.2. Suggestions:
|
||||
Providing native support for compression and decompression units tailored to FP8 or custom precision formats represents a viable approach for future hardware. This could help minimize bandwidth requirements and streamline communication pipelines. The reduced communication overhead is particularly helpful in bandwidth-intensive tasks like MoE training.
|
||||
## 4. Interconnection Driven Design
|
||||
## 4.1. Current Hardware Architecture
|
||||
The NVIDIA H800 GPU SXM architecture we currently use, illustrated in Figure 2 , is built on the Hopper architecture, similar to the H100 GPU. However, it features reduced FP64 computational performance and NVLink bandwidth for regulatory compliance. Specifically, the NVLink bandwidth in H800 SXM nodes is reduced from 900 GB/s to 400 GB/s. This significant reduction in intra-node scale-up bandwidth presents a challenge for high-performance workloads. To compensate, each node is equipped with eight 400G Infiniband (IB) CX7 NICs, enhancing scale-out capabilities to mitigate the bandwidth deficit.
|
||||
To address these hardware constraints, the DeepSeek-V3 model incorporates several design considerations that align with the hardware’s strengths and limitations.
|
||||
## 4.2. Hardware-Aware Parallelism
|
||||
To align with the constraints of the H800 architecture, the following parallelism strategies were considered to optimize the performance of DeepSeek-V3:
|
||||
•
|
||||
Avoidance of Tensor Parallelism (TP): Tensor Parallelism is avoided during training due to its inefficiency under limited NVLink bandwidth. However, during inference, TP can still be selectively used to improve TTFT and TPOT performance.
|
||||
•
|
||||
Enhanced Pipeline Parallelism (PP): DualPipe ( DeepSeek-AI, 2025b ) is employed to overlap attention and MoE computation with MoE communication. This also reduces pipeline bubbles and balances memory usage across GPUs, improving overall throughput. Additional details are available in the technical report ( DeepSeek-AI, 2024d ) .
|
||||
•
|
||||
Accelerated Expert Parallelism (EP): With eight 400Gbps InfiniBand (IB) NICs, the system achieves all-to-all communication at speeds exceeding 40GB/s. Notably, our all-to-all EP implementation, DeepEP ( Zhao et al., 2025b ) , is open-sourced, enabling highly efficient expert parallelism as discussed in the following subsection.
|
||||
Figure 2 . H800 node interconnection.
|
||||
## 4.3. Model Co-Design: Node-Limited Routing
|
||||
The bandwidth disparity between scale-up (intra-node) and scale-out (inter-node) communication in the H800 architecture is approximately 4:1. Specifically, NVLink provides 200GB/s bandwidth (of which about 160GB/s can actually be achieved), while each 400Gbps IB NIC delivers only 50GB/s bandwidth (we consider small message size and latency influence, use 40GB/s for effective bandwidth). To balance and fully utilize the higher intra-node bandwidth, the model architecture is co-designed with hardware, particularly in the TopK Expert Selection Strategy .
|
||||
Consider a setup with 8 nodes (64 GPUs in total) and 256 routed experts (4 experts per GPU).
|
||||
For DeepSeek-V3, each token is routed to one shared expert and 8 routed experts.
|
||||
If its 8 target experts are distributed across all 8 nodes, the communication time over IB would be 8 t 8t , where t t represents the time to send one token over IB.
|
||||
However, by leveraging the higher NVLink bandwidth, tokens routed to the same node can be sent once over IB and then forwarded via NVLink to other intra-node GPUs.
|
||||
The NVLink forwarding enables deduplication of the IB traffic.
|
||||
When the target experts for a given token are distributed across M M nodes, the deduplicated IB communication cost will be reduced to M t Mt ( M < 8 M<8 ).
|
||||
Since the IB traffic depends on only M M , DeepSeek-V3 introduces a Node-Limited Routing for the TopK expert selection strategy.
|
||||
Specifically, we group 256 routed experts into 8 groups, with 32 experts per group, and deploy each group on a single node.
|
||||
On top of this deployment, we algorithmically ensure that each token will be routed to up to 4 nodes. This approach mitigates the bottleneck of IB communication and enhances the effective communication bandwidth during training.
|
||||
## 4.4. Scale-Up and Scale-Out Convergence
|
||||
## 4.4.1. Limitations of Current Implementations
|
||||
While the Node-Limited Routing strategy reduces communication bandwidth requirements, it complicates communication pipeline kernel implementations due to the disparity in bandwidth between intra-node (NVLink) and inter-node (IB) interconnects.
|
||||
In practice, GPU Streaming Multiprocessors (SM) threads are used for both network message handling (e.g., filling QPs and WQEs) and data forwarding over NVLink, consuming computational resources. For example, during training, up to 20 of the SMs on the H800 GPU are allocated for communication-related operations, leaving fewer resources available for actual computation.
|
||||
To maximize throughput in online inference, we perform EP all-to-all communication entirely through NIC RDMA, avoiding SM resource contention and improving compute efficiency. This highlights the advantage of RDMA’s asynchronous communication model in overlapping computation and communication.
|
||||
The following are key tasks currently performed by SMs during EP communication, particularly for the combine stage’s reduce operations and data type conversions. Offloading these tasks to dedicated communication hardware could free up SMs for computation kernels, significantly improving overall efficiency:
|
||||
•
|
||||
Forwarding Data: Aggregating IB traffic destined for multiple GPUs within the same node between the IB and NVLink domains.
|
||||
•
|
||||
Data Transport: Moving data between RDMA buffers (registered GPU memory regions) and input/output buffers.
|
||||
•
|
||||
Reduce Operations: Executing reduce operations required for EP all-to-all combine communications.
|
||||
•
|
||||
Managing Memory Layouts: Handling fine-grained memory layouts for chunked data transfers across the IB and NVLink domains.
|
||||
•
|
||||
Data Type Cast : Converting data type before and after all-to- all communications.
|
||||
## 4.4.2. Suggestions:
|
||||
To address these inefficiencies, we strongly recommend that future hardware should integrate intra-node (scale-up) and inter-node (scale-out) communication into a unified framework. By incorporating dedicated co-processors for network traffic management and seamless forwarding between NVLink and IB domains, such designs can reduce software complexity and maximize bandwidth utilization. For example, node-limited routing strategies employed in DeepSeek-V3 can be further optimized with hardware support for dynamic traffic deduplication.
|
||||
We also recognize emerging interconnect protocols such as the Ultra Ethernet Consortium (UEC) ( Consortium, 2023 ; Consortium, 2024 ) , Ultra Accelerator Link (UALink) ( CONSORTIUM, 2025 ) , both of which are poised to drive advancements in scale-up and scale-out communication. More recently, Unified Bus (UB) ( Liao et al., 2025 ) has introduced a novel approach to scale-up and scale-out convergence. Section 6 further explores several technical innovations proposed by UEC and UALink. However, in this section, our primary focus is on achieving scale-up and scale-out convergence at the programming framework level:
|
||||
(1)
|
||||
Unified Network Adapter: Design NICs (Network Interface Cards) or I/O Dies that are connected to unified scale-up and scale-out networks. These adapters should also support basic switch functionality, such as forwarding packets from the scale-out network to specific GPUs within the scale-up network. This could be achieved using a single LID (Local Identifier) or IP address with policy-based routing.
|
||||
(2)
|
||||
Dedicated Communication Co-Processor: Introduce a dedicated co-processor or programmable component—such as an I/O die—for handling network traffic. This component should offload packet processing from GPU SMs, provide hardware-accelerated memory copy for efficient buffer management, and, crucially, accelerate memory load/store operations in a manner similar to TMA (Tensor Memory Accelerator), thereby saturating bandwidth with minimal resource consumption.
|
||||
(3)
|
||||
Flexible Forwarding, Broadcast and Reduce Mechanisms: Hardware should support flexible forwarding, broadcast operations (for EP dispatch), and reduce operations (for EP combine) across scale-up and scale-out networks—mirroring our current GPU SM-based implementation. This would not only improve effective bandwidth but also reduce the computational complexity of network-specific operations.
|
||||
(4)
|
||||
Hardware Synchronization Primitives: Provide fine-grained hardware synchronization instructions to handle memory consistency issues or out-of-order packet arrivals at the hardware level. This would eliminate the need for software-based synchronization mechanisms like RDMA completion events, which introduce extra latency and increase programming complexity. Memory-semantic communication with an acquire/release mechanism is a promising implementation.
|
||||
By implementing these recommendations, future hardware designs can significantly enhance the efficiency of large-scale distributed AI systems while simplifying software development.
|
||||
## 4.5. Bandwidth Contention and Latency
|
||||
## 4.5.1. Limitations:
|
||||
Besides, current hardware lacks the flexibility to dynamically allocate bandwidth between different types of traffic on NVLink and PCIe. For example, during inference, transferring KV cache data from CPU memory to GPU can consume tens of GB/s, saturating PCIe bandwidth. If the GPU simultaneously uses IB for EP communication, this contention between KV cache transfers and EP communication can degrade overall performance and cause latency spikes.
|
||||
Figure 3 . Eight-plane two-layer fat-tree scale-out network: Each GPU and IB NIC pair belongs to one network plane. Cross-plane traffic must use another NIC and PCIe or NVLink for intra-node forwarding.
|
||||
Figure 4 . Ideal Multi-Plane Network: Each NIC is equipped with multiple physical ports, each connected to a distinct network plane. A single queue pair (QP) can simultaneously utilize all available ports for transmitting and receiving packets, which necessitates native support for out-of-order placement within the NIC.
|
||||
## 4.5.2. Suggestions:
|
||||
•
|
||||
Dynamic NVLink/PCIe Traffic Prioritization: Hardware should support dynamic prioritization of traffic based on its type. For example, traffic related to EP, TP, and KV cache transfers should be assigned different priorities to maximize interconnect efficiency. For PCIe, exposing the traffic class (TC) to user-level programming would suffice.
|
||||
•
|
||||
I/O Die Chiplet Integration: Integrating NICs directly into the I/O die and connecting them to the compute die in the same package, rather than through conventional PCIe, would substantially reduce communication latency and alleviate PCIe bandwidth contention.
|
||||
•
|
||||
CPU–GPU Interconnects within the Scale-Up Domain: To further optimize intra-node communication, CPUs and GPUs should be interconnected using NVLink or similar dedicated high-bandwidth fabrics, rather than relying solely on PCIe. Similar to the benefits provided by integrating NICs into the I/O die, this approach can significantly improve scenarios such as offloading parameters or KV cache between GPU and CPU memory during training and inference.
|
||||
## 5. Large Scale Network Driven Design
|
||||
## 5.1. Network Co-Design: Multi-Plane Fat-Tree
|
||||
During the training of DeepSeek-V3, we deployed a Multi-Plane Fat-Tree (MPFT) scale-out network, as shown in Figure 3 . Each node is equipped with eight GPUs and eight IB NICs, with each GPU–NIC pair assigned to a distinct network plane. Additionally, each node has a 400 Gbps Ethernet RoCE NIC connected to a separate storage network plane for accessing the 3FS ( DeepSeek-AI, 2025c ) distributed file system. In the scale-out network, we used 64-port 400G IB switches, enabling the topology theoretically supports up to 16,384 GPUs while retaining the cost and latency advantages of a two-layer network. However, due to policy and regulatory constraints, just over two thousand GPUs were ultimately deployed.
|
||||
Furthermore, due to the current limitations of IB ConnectX-7, our deployed MPFT network does not fully realize the envisioned architecture. Ideally, as depicted in Figure 4 , each NIC would feature multiple physical ports, each connected to a separate network plane, yet collectively exposed as a single logical interface to the user through port bonding. From a user perspective, a single Queue Pair (QP) could seamlessly transmit and receive messages across all available ports, akin to packet spraying. As a consequence, packets originating from the same QP may traverse distinct network paths and arrive at the receiver out of order, thereby necessitating native support for out-of-order placement within the NIC to guarantee message consistency and preserve the correct ordering semantics. For example, InfiniBand ConnectX-8 natively supports four plane. It would be advantageous for future NICs to fully support advanced multi-plane capabilities, allowing two-tier fat-tree networks to scale effectively to much larger AI clusters. Overall, the multi-plane architecture offers significant advantages in fault isolation, robustness, load balancing, and large-scale system scalability.
|
||||
## 5.1.1. Advantages of Multi-Plane Fat-Tree Network
|
||||
•
|
||||
Subset of Multi-Rail Fat-Tree (MRFT): The MPFT topology constitutes a specific subset of the broader MRFT architecture. As a result, existing optimizations developed by NVIDIA and NCCL for Multi-Rail networks can be seamlessly leveraged within Multi-Plane network deployments. Furthermore, NCCL’s support for PXN ( Mandakolathur and Jeaugey, 2022 ) technology addresses the inherent challenge of inter-plane isolation, enabling efficient communication even when direct interconnectivity between planes is absent.
|
||||
•
|
||||
Cost Efficiency: As shown in Table 3 , the multi-plane network enables over 10k endpoints using a two-layer fat-tree (FT2) topology, significantly reducing network costs compared to a three-layer fat tree (FT3). The cost per endpoint is even slightly more competitive than the cost-efficient Slim Fly (SF) topology ( Blach et al., 2025 ) .
|
||||
•
|
||||
Traffic Isolation: Each plane operates independently, ensuring that congestion in one plane does not affect others. This isolation improves overall network stability and prevents cascading performance degradation.
|
||||
•
|
||||
Latency Reduction: The two-layer topology achieves lower latency than three-layer fat trees, as demonstrated in our experiments. This makes it particularly suitable for latency-sensitive applications such as MoE-based training and inference.
|
||||
•
|
||||
Robustness: As shown in Figure 4 , multi-port NICs provide multiple uplinks, so single-port failures do not disrupt connectivity and rapid, transparent fault recovery is possible.
|
||||
Table 3. Network topology comparison. Cost estimates are derived from the methodology in the Slim Fly (SF) paper ( Blach et al., 2025 ) . DF denotes the canonical dragonfly topology ( De Sensi et al., 2020 ; Kim et al., 2008 ; Rahman et al., 2019 ) .
|
||||
Metric
|
||||
FT2
|
||||
MPFT
|
||||
FT3
|
||||
SF
|
||||
DF
|
||||
Endpoints
|
||||
2,048
|
||||
16,384
|
||||
65,536
|
||||
32,928
|
||||
261,632
|
||||
Switches
|
||||
96
|
||||
768
|
||||
5,120
|
||||
1,568
|
||||
16,352
|
||||
Links
|
||||
2,048
|
||||
16,384
|
||||
131,072
|
||||
32,928
|
||||
384,272
|
||||
Cost [M$]
|
||||
9
|
||||
72
|
||||
491
|
||||
146
|
||||
1,522
|
||||
Cost/Endpoint [k$]
|
||||
4.39
|
||||
4.39
|
||||
7.5
|
||||
4.4
|
||||
5.8
|
||||
It is important to note that, due to current 400G NDR InfiniBand limitations, cross-plane communication requires intra-node forwarding, which introduces additional latency during inference. If future hardware can achieve scale-up and scale-out network convergence as discussed earlier, this latency can be significantly reduced, further enhancing the viability of multi-plane networks.
|
||||
Figure 5 . NCCL all-to-all performance from 32 to 128 GPUs for MRFT and MPFT networks.
|
||||
Figure 6 . Latency comparison between MPFT and MRFT networks in NCCL all-to-all test under different message sizes, showing that their performance is nearly identical.
|
||||
## 5.1.2. Performance Analysis
|
||||
To verify the effectiveness of the Multi-Plane Network design, we conducted real-world experiments on our cluster, modifying the cluster’s network topology to compare the performance of the Multi-Plane Two-Layer Fat Tree (MPFT) and the Single-Plane Multi-Rail Fat Tree (MRFT) . Below are the key findings from our experiments:
|
||||
1. All-to-All Communication and EP Scenarios :
|
||||
As illustrated in Figure 5 , the all-to-all performance of the multi-plane network is very similar to that of the single-plane multi-rail network. This performance parity can be attributed to NCCL’s PXN ( Mandakolathur and Jeaugey, 2022 ) mechanism, which optimizes traffic forwarding via NVLink in multi-rail topologies. The multi-plane topology also benefits from this mechanism. As shown in Figure 6 , the results of all-to-all communication tests conducted on 16 GPUs reveal negligible differences in latency between the MPFT and MRFT topologies.
|
||||
To evaluate MPFT’s performance of all-to-all communication in practical training scenarios, we tested the EP communication patterns commonly used during training. As shown in Figure 7 , each GPU achieves a high bandwidth exceeding 40GB/s in a multi-plane network, providing reliable performance that meets the demands of training.
|
||||
2. Training Throughput for DeepSeek-V3 Model :
|
||||
We also compare the training metrics of the DeepSeek-V3 model between MPFT and MRFT in Table 4 .
|
||||
MFU (Model Flops Utilization) is calculated based on BF16 peak performance. Causal MFU only takes into account the flops of the lower triangle of the attention matrix (in line with FlashAttention ( Dao et al., 2022 ; Dao, 2023 ) ), while non-causal MFU includes the flops of the whole attention matrix (in line with Megatron ( Korthikanti et al., 2023 ) ). 1F, 1B, and 1W denote forward time, input backward time, and weight backward time, respectively.
|
||||
When training the V3 model on 2048 GPUs, the performance of MPFT is nearly identical to that of MRFT, with observed differences falling within normal fluctuations and measurement error.
|
||||
## 5.2. Low Latency Networks
|
||||
In our model inference, large-scale EP relies heavily on all-to-all communication, which is highly sensitive to both bandwidth and latency. Consider a typical scenario discussed in Section 2.3.2 , with a network bandwidth of 50GB/s, the data transfer should ideally take approximately 120 μ s 120\penalty\ \mu\mathrm{s}
|
||||
. Therefore, the intrinsic network latencies on the order of microseconds can critically impact system performance, making their effects non-negligible.
|
||||
## 5.2.1. IB or RoCE
|
||||
Figure 7 . DeepEP performance on MPFT: The EP dispatch and combine kernel communicates across 16 to 128 GPUs using all-to-all. Each GPU processes 4096 tokens. The observed throughput nearly saturates the 400Gps NIC bandwidth.
|
||||
Table 4. Training metric comparison between MPFT and MRFT networks.
|
||||
Metric
|
||||
MPFT
|
||||
MRFT
|
||||
tokens/day (B)
|
||||
272.80
|
||||
272.52
|
||||
time/step (s)
|
||||
19.926
|
||||
19.946
|
||||
1F (s)
|
||||
1.13
|
||||
1.13
|
||||
bubble (s)
|
||||
2.06
|
||||
2.03
|
||||
1B (s)
|
||||
1.99
|
||||
1.99
|
||||
1W (s)
|
||||
0.48
|
||||
0.48
|
||||
1F1B (s)
|
||||
13.95
|
||||
14.00
|
||||
opt (s)
|
||||
0.29
|
||||
0.31
|
||||
TFLOPS (non-causal)
|
||||
432
|
||||
432
|
||||
TFLOPS (causal)
|
||||
385
|
||||
385
|
||||
MFU (non-causal)
|
||||
43.73%
|
||||
43.68%
|
||||
MFU (causal)
|
||||
38.94%
|
||||
38.90%
|
||||
As shown in Table 5 , IB consistently achieves lower latency, making it the preferred choice for latency-sensitive workloads such as distributed training and inference. Although IB has superior latency performance compared to RDMA over Converged Ethernet (RoCE), it comes with certain limitations:
|
||||
•
|
||||
Cost: IB hardware is significantly more expensive than RoCE solutions, which limits its widespread adoption.
|
||||
•
|
||||
Scalability: IB switches typically support only 64 ports per switch, compared to the 128 ports commonly found in RoCE switches. This restricts the scalability of IB-based clusters, particularly for large-scale deployments.
|
||||
Table 5. CPU side end-to-end latency comparison between IB, RoCE, and intra-node NVLink for 64B data transmission.
|
||||
Link Layer
|
||||
Same Leaf
|
||||
Cross Leaf
|
||||
RoCE
|
||||
3.6us
|
||||
5.6us
|
||||
InfiniBand
|
||||
2.8us
|
||||
3.7us
|
||||
NVLink
|
||||
3.33us
|
||||
-
|
||||
## 5.2.2. Recommendations for RoCE Improvements
|
||||
While RoCE has the potential to be a cost-effective alternative to IB, its current limitations in latency and scalability prevent it from fully meeting the demands of large-scale AI systems. Below, we outline specific recommendations for improving RoCE:
|
||||
(1)
|
||||
Specialized Low-Latency RoCE Switches:
|
||||
We recommend that Ethernet vendors develop RoCE switches specifically optimized for RDMA workloads by removing unnecessary Ethernet features. The Slingshot architecture ( De Sensi et al., 2020 ) exemplifies how Ethernet-based designs can achieve latency performance comparable to IB. Similarly, recent innovations from Broadcom ( Broadcom, 2025 ) , including the AI Forwarding Header (AIFH) and upcoming low-latency Ethernet switches, demonstrate the feasibility of high-performance Ethernet fabrics tailored for AI. We are looking forward to continuing innovation in this direction.
|
||||
(2)
|
||||
Optimized Route Policy: As shown in Figure 8 , the default Equal-Cost Multi-Path (ECMP) routing policy in RoCE struggles to distribute traffic efficiently across interconnects, leading to severe congestion performance degradation in NCCL collective communication tests. LLM training traffic, such as in DP (Data Parallelism), tends to lack randomness, causing multiple flows to converge on the same interconnect link. In contrast, Adaptive Routing (AR) ( Geoffray and Hoefler, 2008 ) can significantly enhance network performance by dynamically spraying packets across multiple paths. While static routing—based on manually configured route tables—can avoid link conflicts for specific destinations, it lacks flexibility. For large-scale all-to-all communication, adaptive routing offers superior performance and scalability.
|
||||
(3)
|
||||
Improved Traffic Isolation or Congestion Control Mechanisms:
|
||||
Current RoCE switches support only a limited number of priority queues, which are insufficient for complex AI workloads involving concurrent communication patterns such as EP’s all-to-all and DP’s all-reduce. In such mixed workloads, all-to-all traffic can cause incast congestion due to bursty many-to-one transfers, potentially degrading overall network performance. To address incast’s influence on other traffic, one approach is to adopt virtual output queuing (VOQ), assigning a dedicated virtual queue to each QP to isolate traffic flows. Alternatively, more effective congestion control (CC) mechanisms such as RTT-based CC (RTTCC) or user-programmable CC (PCC) can be employed, enabling NIC–switch co-optimization to maintain low latency and high through
|
||||
|
||||
=== BODY (tail) ===
|
||||
put under dynamic traffic conditions.
|
||||
## 5.2.3. InfiniBand GPUDirect Async (IBGDA)
|
||||
We utilize IBGDA ( NVIDIA, 2022 ; Agostini et al., 2018 ) to reduce latency in network communications. Traditionally, network communication involves the creation of a CPU proxy thread: once the GPU has prepared the data, it must notify the CPU proxy, which then populates the control information for the work request (WR) and signals the NIC via a doorbell mechanism to initiate data transmission. This process introduces additional communication overhead.
|
||||
IBGDA addresses this issue by allowing the GPU to directly fill the WR content and write to the RDMA doorbell MMIO address. By managing the entire control plane within the GPU, IBGDA eliminates the significant latency overhead associated with GPU-CPU communication. Moreover, when sending a large number of small packets, the control plane processor can easily become a bottleneck. Since GPUs have multiple parallel threads, the sender can leverage these threads to distribute the workload, thereby avoiding such bottlenecks. A range of works—including our DeepEP ( Zhao et al., 2025b ) —have leveraged IBGDA and reported substantial performance gains ( Agostini et al., 2017 ; Chen et al., 2025 ; Zheng et al., 2025 ) . We therefore advocate for such capabilities to be widely supported across accelerator devices.
|
||||
Figure 8 . RoCE network bandwidth of AllGather and ReduceScatter communication primitives under different routing methods (ECMP, AR, Static Routing) and TP dimensions.
|
||||
## 6. Discussion and Insights for Future Hardware Architecture Design
|
||||
Building on the previous sections, we summarize key architectural insights and outline future directions for hardware design tailored to large-scale AI workloads.
|
||||
Section 2.3.2 highlighted the importance of large-scale scale-up networks for accelerating model inference. Section 3 discussed the necessity of efficient support for low-precision computation and communication. Section 4 explored the convergence of scale-up and scale-out architectures, along with several proposed enhancements. Section 5 focused on multi-plane network topologies and identified key improvements needed for Ethernet-based interconnects.
|
||||
Together, these sections identify hardware limitations in concrete application contexts and offer corresponding suggestions. Building on that foundation, this section expands the discussion to broader considerations and proposes forward-looking directions for future hardware architecture design.
|
||||
## 6.1. Robustness Challenges
|
||||
## 6.1.1. Limitations:
|
||||
•
|
||||
Interconnect Failures: High-performance interconnects (e.g., IB and NVLink) are prone to intermittent disconnections, which can disrupt node-to-node communication. This is especially harmful in communication-heavy workloads like EP, where even brief interruptions may lead to significant performance drops or job failures.
|
||||
•
|
||||
Single Hardware Failures: Node crashes, GPU failures, or ECC (Error-Correcting Code) memory errors can compromise long-running training jobs, often requiring costly restarts. The impact of such failures escalates in large-scale deployments, where the probability of a single-point failure increases proportionally with system size.
|
||||
•
|
||||
Silent Data Corruption: Errors undetected by ECC mechanisms, such as multi-bit memory flips or computational inaccuracies, pose a significant risk to model quality. These errors are particularly insidious in long-running tasks, as they can propagate undetected and corrupt downstream computations. Current mitigation strategies rely on application-level heuristics, which are insufficient for ensuring system-wide robustness.
|
||||
## 6.1.2. Suggestions for Advanced Error Detection and Correction
|
||||
To mitigate risks associated with silent corruption, hardware must incorporate advanced error detection mechanisms beyond traditional ECC. Techniques such as checksum-based validation or hardware-accelerated redundancy checks can provide higher reliability for large-scale deployments.
|
||||
Furthermore, hardware vendors should deliver comprehensive diagnostic toolkits to end users, empowering them to rigorously verify the integrity of their systems and proactively identify any latent silent data corruption. Such toolkits, when embedded as part of the standard hardware package, foster transparency and enable continuous validation throughout the operational lifecycle, thereby bolstering overall system trustworthiness.
|
||||
## 6.2. CPU Bottlenecks and Interconnects
|
||||
While accelerator design often takes center stage, CPUs remain essential for coordinating computation, managing I/O, and sustaining system throughput. However, current architectures face several critical bottlenecks:
|
||||
First, as discussed in Section 4.5 , the PCIe interface between CPUs and GPUs often becomes a bandwidth bottleneck, particularly during large-scale parameter, gradient, or KV cache transfers. To mitigate this, future systems should adopt direct CPU–GPU interconnects—such as NVLink or Infinity Fabric—or integrate both CPUs and GPUs into the scale-up domain, thereby eliminating intra-node bottlenecks.
|
||||
In addition to PCIe limitations, sustaining such high data transfer rates also requires exceptionally high memory bandwidth. For example, saturating 160 lanes of PCIe 5.0 demands over 640 GB/s per node, translating to a memory bandwidth requirement of approximately 1 TB/s per node—posing a significant challenge for conventional DRAM architectures.
|
||||
Lastly, latency-sensitive tasks such as kernel launches and network processing demand high single-core CPU performance, typically requiring base frequencies above 4 GHz. Furthermore, modern AI workloads require sufficient CPU cores per GPU to prevent control-side bottlenecks. For chiplet-based architectures, additional cores are needed to support cache-aware workload partitioning and isolation.
|
||||
## 6.3. Toward Intelligent Networks for AI
|
||||
To meet the demands of latency-sensitive workloads, future interconnects must prioritize both low latency and intelligent networks:
|
||||
•
|
||||
Co-Packaged Optics: Incorporating silicon photonics enables scalable higher bandwidth scalability and enhanced energy efficiency, both are critical for large-scale distributed systems.
|
||||
•
|
||||
Lossless Network : Credit-Based Flow Control (CBFC) mechanisms ensures lossless data transmission, yet naively triggering flow control can induce severe head-of-line blocking. Therefore, it is imperative to deploy advanced, endpoint-driven congestion control (CC) algorithms that proactively regulate injection rates and avert pathological congestion scenarios.
|
||||
•
|
||||
Adaptive Routing: As underscored in Section 5.2.2 , future network should standardize the adoption of dynamic routing schemes—such as packet spraying and congestion-aware path selection—that continuously monitor real-time network conditions and intelligently redistribute traffic. These adaptive strategies are particularly effective in alleviating hotspots and mitigating bottlenecks during collective communication workloads, including all-to-all and reduce-scatter operations.
|
||||
•
|
||||
Efficient Fault-Tolerant Protocols: Robustness against failures can be significantly enhanced through the deployment of self-healing protocols, redundant ports, and rapid failover techniques. For instance, link-layer retry mechanisms and selective retransmission protocols prove indispensable in scaling reliability across large networks, minimizing downtime and ensuring seamless operation despite intermittent failures.
|
||||
•
|
||||
Dynamic Resource Management: To handle mixed workloads effectively, future hardware should enable dynamic bandwidth allocation and traffic prioritization. For example, inference tasks should be isolated from training traffic in unified clusters, ensuring responsiveness for latency-sensitive applications.
|
||||
## 6.4. Discussion on Memory-Semantic Communication and Ordering Issue
|
||||
Inter-node communication using load/store memory semantics is efficient and programmer-friendly, but current implementations are hampered by memory ordering challenges. For example, after writing data, the sender must issue an explicit memory barrier (fence) before updating a flag to notify the receiver, ensuring data consistency. This strict ordering introduces additional round-trip time (RTT) latency and can stall the issuing thread, impeding inflight stores and reducing throughput. Similar out-of-order synchronization issues arise in message-semantic RDMA; for instance, performing RDMA atomic add operations with packet spraying after regular RDMA writes on InfiniBand or NVIDIA BlueField-3 can incur additional RTT latency.
|
||||
To address these, we advocate for hardware support that offers built-in ordering guarantees for memory-semantic communication. Such consistency should be enforced both at the programming level (e.g., via acquire/release semantics) and by hardware at the receiver, enabling in-order delivery without added overhead.
|
||||
Several approaches are possible. For example, the receiver can buffer atomic messages and use packet sequence numbers (PSN) to ensure in-order processing; this method is straightforward and effective for maintaining correctness. Alternatively, a region-based acquire/release (RAR) mechanism is particularly attractive, as it enables true acquire/release semantics for remote memory access and offers greater flexibility for a range of workloads. In this approach, the receiver hardware maintains lightweight metadata—such as bitmaps or region-based counters—to track the state of memory regions. Acquire and release operations are scoped to specific address ranges, enabling efficient, hardware-enforced ordering without explicit sender-side fences. Notably, both approaches are amenable to implementation on the NIC or I/O die and are applicable to both memory-semantic and message-semantic RDMA primitives, thereby broadening their practical utility.
|
||||
## 6.5. In-Network Computation and Compression
|
||||
EP involves two critical all-to-all stages— dispatch and combine —that present significant opportunities for in-network optimization. The dispatch stage resembles a small-scale multicast operation, where a single message must be forwarded to multiple target devices. A hardware-level protocol enabling automatic packet replication and forwarding to multiple destinations could drastically reduce communication overhead and improve efficiency.
|
||||
The combine stage, acting as a small-scale reduction operation, could benefit from in-network aggregation techniques. However, due to the small reduction scope and imbalanced workload in EP combine, implementing in-network aggregation in a flexible manner is challenging.
|
||||
Moreover, as highlighted in Section 3.2 , LogFMT enables low-precision token transmission with minimal impact on model performance. Incorporating LogFMT natively within network hardware could further optimize communication by increasing entropy density and reducing bandwidth usage. Hardware-accelerated compression and decompression would allow seamless integration of LogFMT into distributed systems, enhancing overall throughput.
|
||||
## 6.6. Memory-Centric Innovations
|
||||
## 6.6.1. Limitations of Memory Bandwidth
|
||||
The exponential growth in model sizes has outpaced advancements in high-bandwidth memory (HBM) technology. This disparity creates a memory bottleneck, particularly in attention-heavy architectures like Transformers.
|
||||
## 6.6.2. Suggestions:
|
||||
•
|
||||
DRAM-Stacked Accelerators: Leveraging advanced 3D stacking technologies, DRAM dies can be vertically integrated atop a logic die, thereby enabling exceptionally high memory bandwidth, ultra-low latency, and a practical memory capacity (though stack-limited). This architectural paradigm proves remarkably advantageous for ultra-fast inference in MoE models, where memory throughput is a critical bottleneck. Architectures such as SeDRAM ( Wang et al., 2023 ) exemplify the potential of this approach, delivering unprecedented performance for memory-bound workloads.
|
||||
•
|
||||
System-on-Wafer (SoW): Wafer-scale integration ( Lie, 2022 ) can maximize computational density and memory bandwidth, addressing the needs of ultra-large-scale models.
|
||||
## 7. Conclusion
|
||||
DeepSeek-V3 exemplifies the transformative potential of hardware-software co-design in advancing the scalability, efficiency, and robustness of large-scale AI systems. By addressing the limitations of current hardware architectures and proposing actionable recommendations, this paper provides a roadmap for the next generation of AI-optimized hardware. These innovations will be critical as AI workloads continue to grow in complexity and scale, driving the future of intelligent systems.
|
||||
@@ -0,0 +1,710 @@
|
||||
TITLE: DeepSeek-OCR: Contexts Optical Compression
|
||||
ARXIV: 2510.18234
|
||||
AUTHORS: Haoran Wei , Yaofeng Sun , Yukun Li
|
||||
DATE: [Submitted on 21 Oct 2025]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-OCR: Contexts Optical Compression
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Related Works
|
||||
## 2.1 Typical Vision Encoders in VLMs
|
||||
## 2.2 End-to-end OCR Models
|
||||
## 3 Methodology
|
||||
## 3.1 Architecture
|
||||
## 3.2 DeepEncoder
|
||||
## 3.2.1 Architecture of DeepEncoder
|
||||
## 3.2.2 Multiple resolution support
|
||||
## 3.3 The MoE Decoder
|
||||
## 3.4 Data Engine
|
||||
## 3.4.1 OCR 1.0 data
|
||||
## 3.4.2 OCR 2.0 data
|
||||
## 3.4.3 General vision data
|
||||
## 3.4.4 Text-only data
|
||||
## 3.5 Training Pipelines
|
||||
## 3.5.1 Training DeepEncoder
|
||||
## 3.5.2 Training DeepSeek-OCR
|
||||
## 4 Evaluation
|
||||
## 4.1 Vision-text Compression Study
|
||||
## 4.2 OCR Practical Performance
|
||||
## 4.3 Qualitative Study
|
||||
## 4.3.1 Deep parsing
|
||||
## 4.3.2 Multilingual recognition
|
||||
## 4.3.3 General vision understanding
|
||||
## 5 Discussion
|
||||
## 6 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-OCR: Contexts Optical Compression
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Typical Vision Encoders in VLMs
|
||||
2.2 End-to-end OCR Models
|
||||
3 Methodology
|
||||
3.1 Architecture
|
||||
3.2 DeepEncoder
|
||||
3.2.1 Architecture of DeepEncoder
|
||||
3.2.2 Multiple resolution support
|
||||
3.3 The MoE Decoder
|
||||
3.4 Data Engine
|
||||
3.4.1 OCR 1.0 data
|
||||
3.4.2 OCR 2.0 data
|
||||
3.4.3 General vision data
|
||||
3.4.4 Text-only data
|
||||
3.5 Training Pipelines
|
||||
3.5.1 Training DeepEncoder
|
||||
3.5.2 Training DeepSeek-OCR
|
||||
4 Evaluation
|
||||
4.1 Vision-text Compression Study
|
||||
4.2 OCR Practical Performance
|
||||
4.3 Qualitative Study
|
||||
4.3.1 Deep parsing
|
||||
4.3.2 Multilingual recognition
|
||||
4.3.3 General vision understanding
|
||||
5 Discussion
|
||||
6 Conclusion
|
||||
References
|
||||
License: CC BY 4.0
|
||||
arXiv:2510.18234v1 [cs.CV] 21 Oct 2025
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-OCR: Contexts Optical Compression
|
||||
Haoran Wei
|
||||
Yaofeng Sun
|
||||
Yukun Li
|
||||
Affiliation: DeepSeek-AI
|
||||
## Abstract
|
||||
We present DeepSeek-OCR as an initial investigation into the feasibility of compressing long contexts via optical 2D mapping. DeepSeek-OCR consists of two components: DeepEncoder and DeepSeek3B-MoE-A570M as the decoder. Specifically, DeepEncoder serves as the core engine, designed to maintain low activations under high-resolution input while achieving high compression ratios to ensure an optimal and manageable number of vision tokens. Experiments show that when the number of text tokens is within 10 times that of vision tokens (i.e., a compression ratio < 10 × \times ), the model can achieve decoding (OCR) precision of 97%. Even at a compression ratio of 20 × \times , the OCR accuracy still remains at about 60%. This shows considerable promise for research areas such as historical long-context compression and memory forgetting mechanisms in LLMs. Beyond this, DeepSeek-OCR also demonstrates high practical value. On OmniDocBench, it surpasses GOT-OCR2.0 (256 tokens/page) using only 100 vision tokens, and outperforms MinerU2.0 (6000+ tokens per page on average) while utilizing fewer than 800 vision tokens. In production, DeepSeek-OCR can generate training data for LLMs/VLMs at a scale of 200k+ pages per day (a single A100-40G). Codes and model weights are publicly accessible at http://github.com/deepseek-ai/DeepSeek-OCR .
|
||||
(a) Compression on Fox benchmark
|
||||
(b) Performance on Omnidocbench
|
||||
Figure 1 : Figure (a) shows the compression ratio (number of text tokens in ground truth/number of vision tokens model used) testing on Fox [ 21 ] benchmark; Figure (b) shows performance comparisons on OmniDocBench [ 27 ] . DeepSeek-OCR can achieve state-of-the-art performance among end-to-end models enjoying the fewest vision tokens.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Typical Vision Encoders in VLMs
|
||||
2.2 End-to-end OCR Models
|
||||
3 Methodology
|
||||
3.1 Architecture
|
||||
3.2 DeepEncoder
|
||||
3.2.1 Architecture of DeepEncoder
|
||||
3.2.2 Multiple resolution support
|
||||
3.3 The MoE Decoder
|
||||
3.4 Data Engine
|
||||
3.4.1 OCR 1.0 data
|
||||
3.4.2 OCR 2.0 data
|
||||
3.4.3 General vision data
|
||||
3.4.4 Text-only data
|
||||
3.5 Training Pipelines
|
||||
3.5.1 Training DeepEncoder
|
||||
3.5.2 Training DeepSeek-OCR
|
||||
4 Evaluation
|
||||
4.1 Vision-text Compression Study
|
||||
4.2 OCR Practical Performance
|
||||
4.3 Qualitative Study
|
||||
4.3.1 Deep parsing
|
||||
4.3.2 Multilingual recognition
|
||||
4.3.3 General vision understanding
|
||||
5 Discussion
|
||||
6 Conclusion
|
||||
References
|
||||
## 1 Introduction
|
||||
Current Large Language Models (LLMs) face significant computational challenges when processing long textual content due to quadratic scaling with sequence length. We explore a potential solution: leveraging visual modality as an efficient compression medium for textual information. A single image containing document text can represent rich information using substantially fewer tokens than the equivalent digital text, suggesting that optical compression through vision tokens could achieve much higher compression ratios.
|
||||
This insight motivates us to reexamine vision-language models (VLMs) from an LLM-centric perspective, focusing on how vision encoders can enhance LLMs’ efficiency in processing textual information rather than basic VQA [ 12 , 24 , 32 , 41 , 16 ] what humans excel at. OCR tasks, as an intermediate modality bridging vision and language, provide an ideal testbed for this vision-text compression paradigm, as they establish a natural compression-decompression mapping between visual and textual representations while offering quantitative evaluation metrics.
|
||||
Accordingly, we present DeepSeek-OCR, a VLM designed as a preliminary proof-of-concept for efficient vision-text compression. Our work makes three primary contributions:
|
||||
First, we provide comprehensive quantitative analysis of vision-text token compression ratios. Our method achieves 96%+ OCR decoding precision at 9-10 × \times text compression, ∼ \sim 90% at 10-12 × \times compression, and ∼ \sim 60% at 20 × \times compression on Fox [ 21 ] benchmarks featuring diverse document layouts (with actual accuracy being even higher when accounting for formatting differences between output and ground truth), as shown in Figure 1(a) . The results demonstrate that compact language models can effectively learn to decode compressed visual representations, suggesting that larger LLMs could readily acquire similar capabilities through appropriate pretraining design.
|
||||
Second, we introduce DeepEncoder, a novel architecture that maintains low activation memory and minimal vision tokens even with high-resolution inputs. It serially connects window attention and global attention encoder components through a 16 × \times convolutional compressor. This design ensures that the window attention component processes a large number of vision tokens, while the compressor reduces vision tokens before they enter the dense global attention component, achieving effective memory and token compression.
|
||||
Third, we develop DeepSeek-OCR based on DeepEncoder and DeepSeek3B-MoE [ 19 , 20 ] . As shown in Figure 1(b) , it achieves state-of-the-art performance within end-to-end models on OmniDocBench while using the fewest vision tokens. Additionally, we equip the model with capabilities for parsing charts, chemical formulas, simple geometric figures, and natural images to enhance its practical utility further. In production, DeepSeek-OCR can generate 33 million pages of data per day for LLMs or VLMs using 20 nodes (each with 8 A100-40G GPUs).
|
||||
In summary, this work presents a preliminary exploration of using visual modality as an efficient compression medium for textual information processing in LLMs. Through DeepSeek-OCR, we demonstrate that vision-text compression can achieve significant token reduction (7-20 × \times ) for different historical context stages, offering a promising direction for addressing long-context challenges in large language models. Our quantitative analysis provides empirical guidelines for VLM token allocation optimization, while the proposed DeepEncoder architecture showcases practical feasibility with real-world deployment capabilities. Although focused on OCR as a proof-of-concept, this paradigm opens new possibilities for rethinking how vision and language modalities can be synergistically combined to enhance computational efficiency in large-scale text processing and agent systems.
|
||||
Figure 2 : Typical vision encoders in popular VLMs. Here are three types of encoders commonly used in current open-source VLMs, all of which suffer from their respective deficiencies.
|
||||
## 2 Related Works
|
||||
## 2.1 Typical Vision Encoders in VLMs
|
||||
Current open-source VLMs employ three main types of vision encoders, as illustrated in Figure 2 . The first type is a dual-tower architecture represented by Vary [ 36 ] , which utilizes parallel SAM [ 17 ] encoder to increase visual vocabulary parameters for high-resolution image processing. While offering controllable parameters and activation memory, this approach suffers from significant drawbacks: it requires dual image preprocessing that complicates deployment and makes encoder pipeline parallelism challenging during training. The second type is tile-based method exemplified by InternVL2.0 [ 8 ] , which processes images by dividing them into small tiles for parallel computation, reducing activation memory under high-resolution settings. Although capable of handling extremely high resolutions, this approach has notable limitations due to its typically low native encoder resolution (below 512 × \times 512), causing large images to be excessively fragmented and resulting in numerous vision tokens. The third type is adaptive resolution encoding represented by Qwen2-VL [ 35 ] , which adopts the NaViT [ 10 ] paradigm to directly process full images through patch-based segmentation without tile parallelization. While this encoder can handle diverse resolutions flexibly, it faces substantial challenges with large images due to massive activation memory consumption that can cause GPU memory overflow, and sequence packing requires extremely long sequence lengths during training. Long vision tokens will slow down both prefill and generation phases of inference.
|
||||
## 2.2 End-to-end OCR Models
|
||||
OCR, particularly document parsing task, has been a highly active topic in the image-to-text domain. With the advancement of VLMs, a large number of end-to-end OCR models have emerged, fundamentally transforming the traditional pipeline architecture (which required separate detection and recognition expert models) by simplifying OCR systems. Nougat [ 6 ] first employs end-to-end framework for academic paper OCR on arXiv, demonstrating the potential of models in handling dense perception tasks. GOT-OCR2.0 [ 38 ] expands the scope of OCR2.0 to include more synthetic image parsing tasks and designs an OCR model with performance-efficiency trade-offs, further highlighting the potential of end-to-end OCR researches. Additionally, general vision models such as Qwen-VL series [ 35 ] , InternVL series [ 8 ] , and many their derivatives continuously enhance their document OCR capabilities to explore dense visual perception boundaries. However, a crucial research question that current models have not addressed is: for a document containing 1000 words, how many vision tokens are at least needed for decoding? This question holds significant importance for research in the principle that " a picture is worth a thousand words. "
|
||||
Figure 3 : The architecture of DeepSeek-OCR. DeepSeek-OCR consists of a DeepEncoder and a DeepSeek-3B-MoE decoder. DeepEncoder is the core of DeepSeek-OCR, comprising three components: a SAM [ 17 ] for perception dominated by window attention, a CLIP [ 29 ] for knowledge with dense global attention, and a 16 × \times token compressor that bridges between them.
|
||||
## 3 Methodology
|
||||
## 3.1 Architecture
|
||||
As shown in Figure 3 , DeepSeek-OCR enjoys a unified end-to-end VLM architecture consisting of an encoder and a decoder. The encoder (namely DeepEncoder) is responsible for extracting image features and tokenizing as well as compressing visual representations. The decoder is used for generating the required result based on image tokens and prompts. DeepEncoder is approximately 380M in parameters, mainly composed of an 80M SAM-base [ 17 ] and a 300M CLIP-large [ 29 ] connected in series. The decoder adopts a 3B MoE [ 19 , 20 ] architecture with 570M activated parameters. In the following paragraphs, we will delve into the model components, data engineering, and training skills.
|
||||
## 3.2 DeepEncoder
|
||||
To explore the feasibility of contexts optical compression, we need a vision encoder with the following features: 1.Capable of processing high resolutions; 2.Low activation at high resolutions; 3.Few vision tokens; 4.Support for multiple resolution inputs; 5. Moderate parameter count. However, as described in the Section 2.1 , current open-source encoders cannot fully satisfy all these conditions. Therefore, we design a novel vision encoder ourselves, named DeepEncoder.
|
||||
Figure 4 : To test model performance under different compression ratios (requiring different numbers of vision tokens) and enhance the practicality of DeepSeek-OCR, we configure it with multiple resolution modes.
|
||||
## 3.2.1 Architecture of DeepEncoder
|
||||
DeepEncoder mainly consists of two components: a visual perception feature extraction component dominated by window attention, and a visual knowledge feature extraction component with dense global attention. To benefit from the pretraining gains of previous works, we use SAM-base (patch-size 16) and CLIP-large as the main architectures for the two components respectively. For CLIP, we remove the first patch embedding layer since its input is no longer images but output tokens from the previous pipeline. Between the two components, we borrow from Vary [ 36 ] and use a 2-layer convolutional module to perform 16 × \times downsampling of vision tokens. Each convolutional layer has a kernel size of 3, stride of 2, padding of 1, and channels increase from 256 to 1024. Assuming we input a 1024 × \times 1024 image, the DeepEncoder will segment it into 1024/16 × \times 1024/16=4096 patch tokens. Since the first half of encoder is dominated by window attention and only 80M, the activation is acceptable. Before entering global attention, the 4096 tokens go through the compression module and the token count becomes 4096/16=256, thus making the overall activation memory controllable.
|
||||
Table 1 : Multi resolution support of DeepEncoder. For both research and application purposes, we design DeepEncoder with diverse native resolution and dynamic resolution modes.
|
||||
Mode
|
||||
Native Resolution
|
||||
Dynamic Resolution
|
||||
Tiny
|
||||
Small
|
||||
Base
|
||||
Large
|
||||
Gundam
|
||||
Gundam-M
|
||||
Resolution
|
||||
512
|
||||
640
|
||||
1024
|
||||
1280
|
||||
640+1024
|
||||
1024+1280
|
||||
Tokens
|
||||
64
|
||||
100
|
||||
256
|
||||
400
|
||||
n × \times 100+256
|
||||
n × \times 256+400
|
||||
Process
|
||||
resize
|
||||
resize
|
||||
padding
|
||||
padding
|
||||
resize + padding
|
||||
resize + padding
|
||||
## 3.2.2 Multiple resolution support
|
||||
Suppose we have an image with 1000 optical characters and we want to test how many vision tokens are needed for decoding. This requires the model to support a variable number of vision tokens. That is to say the DeepEncoder needs to support multiple resolutions.
|
||||
We meet the requirement aforementioned through dynamic interpolation of positional encodings, and design several resolution modes for simultaneous model training to achieve the capability of a single DeepSeek-OCR model supporting multiple resolutions. As shown in Figure 4 , DeepEncoder mainly supports two major input modes: native resolution and dynamic resolution. Each of them contains multiple sub-modes.
|
||||
Native resolution supports four sub-modes: Tiny, Small, Base, and Large, with corresponding resolutions and token counts of 512 × \times 512 (64), 640 × \times 640 (100), 1024 × \times 1024 (256), and 1280 × \times 1280 (400) respectively. Since Tiny and Small modes have relatively small resolutions, to avoid wasting vision tokens, images are processed by directly resizing the original shape. For Base and Large modes, in order to preserve the original image aspect ratio, images are padded to the corresponding size. After padding, the number of valid vision tokens is less than the actual number of vision tokens, with the calculation formula being:
|
||||
N v a l i d = ⌈ N a c t u a l × [ 1 − ( ( m a x ( w , h ) − m i n ( w , h ) ) / ( m a x ( w , h ) ) ) ] ⌉ N_{valid}=\lceil N_{actual}\times[1-((max(w,h)-min(w,h))/(max(w,h)))]\rceil
|
||||
(1)
|
||||
where w w and h h represent the width and height of the original input image.
|
||||
Dynamic resolution can be composed of two native resolutions. For example, Gundam mode consists of n × \times 640 × \times 640 tiles (local views) and a 1024 × \times 1024 global view. The tiling method following InternVL2.0 [ 8 ] . Supporting dynamic resolution is mainly for application considerations, especially for ultra-high-resolution inputs (such as newspaper images). Tiling is a form of secondary window attention that can effectively reduce activation memory further. It’s worth noting that due to our relatively large native resolutions, images won’t be fragmented too much under dynamic resolution (the number of tiles is controlled within the range of 2 to 9). The vision token number output by the DeepEncoder under Gundam mode is: n × 100 + 256 n\times 100+256 , where n n is the number of tiles. For images with both width and height smaller than 640, n n is set to 0, i.e., Gundam mode will degrade to Base mode.
|
||||
Gundam mode is trained together with the four native resolution modes to achieve the goal of one model supporting multiple resolutions. Note that Gundam-master mode (1024 × \times 1024 local views+1280 × \times 1280 global view) is obtained through continued training on a trained DeepSeek-OCR model. This is mainly for load balancing, as Gundam-master’s resolution is too large and training it together would slow down the overall training speed.
|
||||
## 3.3 The MoE Decoder
|
||||
Our decoder uses the DeepSeekMoE [ 19 , 20 ] , specifically DeepSeek-3B-MoE. During inference, the model activates 6 out of 64 routed experts and 2 shared experts, with about 570M activated parameters. The 3B DeepSeekMoE is very suitable for domain-centric (OCR for us) VLM research, as it obtains the expressive capability of a 3B model while enjoying the inference efficiency of a 500M small model.
|
||||
The decoder reconstructs the original text representation from the compressed latent vision tokens of DeepEncoder as:
|
||||
f dec : ℝ n × d latent → ℝ N × d text ; 𝐗 ^ = f dec ( 𝐙 ) where n ≤ N f_{\text{dec}}:\mathbb{R}^{n\times d_{\text{latent}}}\rightarrow\mathbb{R}^{N\times d_{\text{text}}};\quad\hat{\mathbf{X}}=f_{\text{dec}}(\mathbf{Z})\quad\text{where }n\leq N
|
||||
(2)
|
||||
where 𝐙 ∈ ℝ n × d latent \mathbf{Z}\in\mathbb{R}^{n\times d_{\text{latent}}} are the compressed latent(vision) tokens from DeepEncoder and 𝐗 ^ ∈ ℝ N × d text \hat{\mathbf{X}}\in\mathbb{R}^{N\times d_{\text{text}}} is the reconstructed text representation. The function f dec f_{\text{dec}} represents a non-linear mapping that can be effectively learned by compact language models through OCR-style training. It is reasonable to conjecture that LLMs, through specialized pretraining optimization, would demonstrate more natural integration of such capabilities.
|
||||
## 3.4 Data Engine
|
||||
We constructe complex and diverse training data for DeepSeek-OCR, including OCR 1.0 data, which mainly consists of traditional OCR tasks such as scene image OCR and document OCR; OCR 2.0 data, which mainly includes parsing tasks for complex artificial images, such as common charts, chemical formulas, and plane geometry parsing data; General vision data, which is mainly used to inject certain general image understanding capabilities into DeepSeek-OCR and preserve the general vision interface.
|
||||
(a) Ground truth image
|
||||
(b) Fine annotations with layouts
|
||||
Figure 5 : OCR 1.0 fine annotations display. We format the ground truth into an interleaved layout and text format, where each paragraph of text is preceded by the coordinates and label of it in the original image. All coordinates are normalized into 1000 bins.
|
||||
## 3.4.1 OCR 1.0 data
|
||||
Document data is the top priority for DeepSeek-OCR. We collect 30M pages of diverse PDF data covering about 100 languages from the Internet, with Chinese and English accounting for approximately 25M and other languages accounting for 5M. For this data, we create two types of ground truth: coarse annotations and fine annotations. Coarse annotations are extracted directly from the full dataset using fitz , aimed at teaching the model to recognize optical text, especially in minority languages. Fine annotations include 2M pages each for Chinese and English, labeled using advanced layout models (such as PP-DocLayout [ 33 ] ) and OCR models (such as MinuerU [ 34 ] and GOT-OCR2.0 [ 38 ] ) to construct detection and recognition interleaved data. For minority languages, in the detection part, we find that the layout model enjoys certain generalization capabilities. In the recognition part, we use fitz to create small patch data to train a GOT-OCR2.0, then use the trained model to label small patches after layout processing, employing a model flywheel to create 600K data samples. During the training of DeepSeek-OCR, coarse labels and fine labels are distinguished using different prompts. The ground truth for fine annotation image-text pairs can be seen in Figure 5 . We also collect 3M Word data, constructing high-quality image-text pairs without layout by directly extracting content. This data mainly brings benefits to formulas and HTML-formatted tables. Additionally, we select some open-source data [ 28 , 37 ] as supplements.
|
||||
For natural scene OCR, our model mainly supports Chinese and English. The image data sources come from LAION [ 31 ] and Wukong [ 13 ] , labeled using PaddleOCR [ 9 ] , with 10M data samples each for Chinese and English. Like document OCR, natural scene OCR can also control whether to output detection boxes through prompts.
|
||||
## 3.4.2 OCR 2.0 data
|
||||
Following GOT-OCR2.0 [ 38 ] , we refer to chart, chemical formula, and plane geometry parsing data as OCR 2.0 data. For chart data, following OneChart [ 7 ] , we use pyecharts and matplotlib to render 10M images, mainly including commonly used line, bar, pie, and composite charts. We define chart parsing as image-to-HTML-table conversion task, as shown in Figure 6(a) . For chemical formulas, we utilize SMILES format from PubChem as the data source and render them into images using RDKit, constructing 5M image-text pairs. For plane geometry images, we follow Slow Perception [ 39 ] for generation. Specifically, we use perception-ruler size as 4 to model each line segment. To increase the diversity of rendered data, we introduce geometric translation-invariant data augmentation, where the same geometric image is translated in the original image, corresponding to the same ground truth drawn at the centered position in the coordinate system. Based on this, we construct a total of 1M plane geometry parsing data, as illustrated in Figure 6(b) .
|
||||
(a) Image-text ground truth of chart
|
||||
(b) Image-text ground truth of geometry
|
||||
Figure 6 : For charts, we do not use OneChart’s [ 7 ] dictionary format, but instead use HTML table format as labels, which can save a certain amount of tokens. For plane geometry, we convert the ground truth to dictionary format, where the dictionary contains keys such as line segments, endpoint coordinates, line segment types, etc., for better readability. Each line segment is encoded using the Slow Perception [ 39 ] manner.
|
||||
## 3.4.3 General vision data
|
||||
DeepEncoder can benefit from CLIP’s pretraining gains and has sufficient parameters to incorporate general visual knowledge. Therefore, we also prepare some corresponding data for DeepSeek-OCR. Following DeepSeek-VL2 [ 40 ] , we generate relevant data for tasks such as caption, detection, and grounding. Note that DeepSeek-OCR is not a general VLM model, and this portion of data accounts for only 20% of the total data. We introduce such type of data mainly to preserve the general vision interface, so that researchers interested in our model and general vision task can conveniently advance their work in the future.
|
||||
## 3.4.4 Text-only data
|
||||
To ensure the model’s language capabilities, we introduced 10% of in-house text-only pretrain data, with all data processed to a length of 8192 tokens, which is also the sequence length for DeepSeek-OCR. In summary, when training DeepSeek-OCR, OCR data accounts for 70%, general vision data accounts for 20%, and text-only data accounts for 10%.
|
||||
## 3.5 Training Pipelines
|
||||
Our training pipeline is very simple and consists mainly of two stages: a).Training DeepEncoder independently; b).Training the DeepSeek-OCR. Note that the Gundam-master mode is obtained by continuing training on a pre-trained DeepSeek-OCR model with 6M sampled data. Since the training protocol is identical to other modes, we omit the detailed description hereafter.
|
||||
## 3.5.1 Training DeepEncoder
|
||||
Following Vary [ 36 ] , we utilize a compact language model [ 15 ] and use the next token prediction framework to train DeepEncoder. In this stage, we use all OCR 1.0 and 2.0 data aforementioned, as well as 100M general data sampled from the LAION [ 31 ] dataset. All data is trained for 2 epochs with a batch size of 1280, using the AdamW [ 23 ] optimizer with cosine annealing scheduler [ 22 ] and a learning rate of 5e-5. The training sequence length is 4096.
|
||||
## 3.5.2 Training DeepSeek-OCR
|
||||
After DeepEncoder is ready, we use data mentioned in Section 3.4 to train the DeepSeek-OCR. with the entire training process conducted on the HAI-LLM [ 14 ] platform. The entire model uses pipeline parallelism (PP) and is divided into 4 parts, with DeepEncoder taking two parts and the decoder taking two parts. For DeepEncoder, we treat SAM and the compressor as the vision tokenizer, place them in PP0 and freeze their parameters, while treating the CLIP part as input embedding layer and place it in PP1 with unfrozen weights for training. For the language model part, since DeepSeek3B-MoE has 12 layers, we place 6 layers each on PP2 and PP3. We use 20 nodes (each with 8 A100-40G GPUs) for training, with a data parallelism (DP) of 40 and a global batch size of 640. We use the AdamW optimizer with a step-based scheduler and an initial learning rate of 3e-5. For text-only data, the training speed is 90B tokens/day, while for multimodal data, the training speed is 70B tokens/day.
|
||||
Table 2 : We test DeepSeek-OCR’s vision-text compression ratio using all English documents with 600-1300 tokens from the Fox [ 21 ] benchmarks. Text tokens represent the number of tokens after tokenizing the ground truth text using DeepSeek-OCR’s tokenizer. Vision Tokens=64 or 100 respectively represent the number of vision tokens output by DeepEncoder after resizing input images to 512 × \times 512 and 640 × \times 640.
|
||||
Text Tokens
|
||||
Vision Tokens =64
|
||||
Vision Tokens=100
|
||||
Precision
|
||||
Compression
|
||||
Precision
|
||||
Compression
|
||||
Pages
|
||||
600-700
|
||||
96.5%
|
||||
10.5 × \times
|
||||
98.5%
|
||||
6.7 × \times
|
||||
7
|
||||
700-800
|
||||
93.8%
|
||||
11.8 × \times
|
||||
97.3%
|
||||
7.5 × \times
|
||||
28
|
||||
800-900
|
||||
83.8%
|
||||
13.2 × \times
|
||||
96.8%
|
||||
8.5 × \times
|
||||
28
|
||||
900-1000
|
||||
85.9%
|
||||
15.1 × \times
|
||||
96.8%
|
||||
9.7 × \times
|
||||
14
|
||||
1000-1100
|
||||
79.3%
|
||||
16.5 × \times
|
||||
91.5%
|
||||
10.6 × \times
|
||||
11
|
||||
1100-1200
|
||||
76.4%
|
||||
17.7 × \times
|
||||
89.8%
|
||||
11.3 × \times
|
||||
8
|
||||
1200-1300
|
||||
59.1%
|
||||
19.7 × \times
|
||||
87.1%
|
||||
12.6 × \times
|
||||
4
|
||||
Table 3: We use OmniDocBench [ 27 ] to test the performance of DeepSeek-OCR on real document parsing tasks. All metrics in the table are edit distances, where smaller values indicate better performance. "Tokens" represents the average number of vision tokens used per page, and " †200dpi " means using fitz to interpolate the original image to 200dpi. For the DeepSeek-OCR model, the values in parentheses in the "Tokens" column represent valid vision tokens, calculated according to Equation 1 .
|
||||
Model
|
||||
Tokens
|
||||
English
|
||||
Chinese
|
||||
overall
|
||||
text
|
||||
formula
|
||||
table
|
||||
order
|
||||
overall
|
||||
text
|
||||
formula
|
||||
table
|
||||
order
|
||||
Pipline Models
|
||||
Dolphin [ 11 ]
|
||||
-
|
||||
0.356
|
||||
0.352
|
||||
0.465
|
||||
0.258
|
||||
0.35
|
||||
0.44
|
||||
0.44
|
||||
0.604
|
||||
0.367
|
||||
0.351
|
||||
Marker [ 1 ]
|
||||
-
|
||||
0.296
|
||||
0.085
|
||||
0.374
|
||||
0.609
|
||||
0.116
|
||||
0.497
|
||||
0.293
|
||||
0.688
|
||||
0.678
|
||||
0.329
|
||||
Mathpix [ 2 ]
|
||||
-
|
||||
0.191
|
||||
0.105
|
||||
0.306
|
||||
0.243
|
||||
0.108
|
||||
0.364
|
||||
0.381
|
||||
0.454
|
||||
0.32
|
||||
0.30
|
||||
MinerU-2.1.1 [ 34 ]
|
||||
-
|
||||
0.162
|
||||
0.072
|
||||
0.313
|
||||
0.166
|
||||
0.097
|
||||
0.244
|
||||
0.111
|
||||
0.581
|
||||
0.15
|
||||
0.136
|
||||
MonkeyOCR-1.2B [ 18 ]
|
||||
-
|
||||
0.154
|
||||
0.062
|
||||
0.295
|
||||
0.164
|
||||
0.094
|
||||
0.263
|
||||
0.179
|
||||
0.464
|
||||
0.168
|
||||
0.243
|
||||
PPstructure-v3 [ 9 ]
|
||||
-
|
||||
0.152
|
||||
0.073
|
||||
0.295
|
||||
0.162
|
||||
0.077
|
||||
0.223
|
||||
0.136
|
||||
0.535
|
||||
0.111
|
||||
0.11
|
||||
End-to-end Models
|
||||
Nougat [ 6 ]
|
||||
2352
|
||||
0.452
|
||||
0.365
|
||||
0.488
|
||||
0.572
|
||||
0.382
|
||||
0.973
|
||||
0.998
|
||||
0.941
|
||||
1.00
|
||||
0.954
|
||||
SmolDocling [ 25 ]
|
||||
392
|
||||
0.493
|
||||
0.262
|
||||
0.753
|
||||
0.729
|
||||
0.227
|
||||
0.816
|
||||
0.838
|
||||
0.997
|
||||
0.907
|
||||
0.522
|
||||
InternVL2-76B [ 8 ]
|
||||
6790
|
||||
0.44
|
||||
0.353
|
||||
0.543
|
||||
0.547
|
||||
0.317
|
||||
0.443
|
||||
0.29
|
||||
0.701
|
||||
0.555
|
||||
0.228
|
||||
Qwen2.5-VL-7B [ 5 ]
|
||||
3949
|
||||
0.316
|
||||
0.151
|
||||
0.376
|
||||
0.598
|
||||
0.138
|
||||
0.399
|
||||
0.243
|
||||
0.5
|
||||
0.627
|
||||
0.226
|
||||
OLMOCR [ 28 ]
|
||||
3949
|
||||
0.326
|
||||
0.097
|
||||
0.455
|
||||
0.608
|
||||
0.145
|
||||
0.469
|
||||
0.293
|
||||
0.655
|
||||
0.652
|
||||
0.277
|
||||
GOT-OCR2.0 [ 38 ]
|
||||
256
|
||||
0.287
|
||||
0.189
|
||||
0.360
|
||||
0.459
|
||||
0.141
|
||||
0.411
|
||||
0.315
|
||||
0.528
|
||||
0.52
|
||||
0.28
|
||||
OCRFlux-3B [ 3 ]
|
||||
3949
|
||||
0.238
|
||||
0.112
|
||||
0.447
|
||||
0.269
|
||||
0.126
|
||||
0.349
|
||||
0.256
|
||||
0.716
|
||||
0.162
|
||||
0.263
|
||||
GPT4o [ 26 ]
|
||||
-
|
||||
0.233
|
||||
0.144
|
||||
0.425
|
||||
0.234
|
||||
0.128
|
||||
0.399
|
||||
0.409
|
||||
0.606
|
||||
0.329
|
||||
0.251
|
||||
InternVL3-78B [ 42 ]
|
||||
6790
|
||||
0.218
|
||||
0.117
|
||||
0.38
|
||||
0.279
|
||||
0.095
|
||||
0.296
|
||||
0.21
|
||||
0.533
|
||||
0.282
|
||||
0.161
|
||||
Qwen2.5-VL-72B [ 5 ]
|
||||
3949
|
||||
0.214
|
||||
0.092
|
||||
0.315
|
||||
0.341
|
||||
0.106
|
||||
0.261
|
||||
0.18
|
||||
0.434
|
||||
0.262
|
||||
0.168
|
||||
dots.ocr [ 30 ]
|
||||
3949
|
||||
0.182
|
||||
0.137
|
||||
0.320
|
||||
0.166
|
||||
0.182
|
||||
0.261
|
||||
0.229
|
||||
0.468
|
||||
0.160
|
||||
0.261
|
||||
Gemini2.5-Pro [ 4 ]
|
||||
-
|
||||
0.148
|
||||
0.055
|
||||
0.356
|
||||
0.13
|
||||
0.049
|
||||
0.212
|
||||
0.168
|
||||
0.439
|
||||
0.119
|
||||
0.121
|
||||
MinerU2.0 [ 34 ]
|
||||
6790
|
||||
0.133
|
||||
0.045
|
||||
0.273
|
||||
0.15
|
||||
0.066
|
||||
0.238
|
||||
0.115
|
||||
0.506
|
||||
0.209
|
||||
0.122
|
||||
dots.ocr †200dpi [ 30 ]
|
||||
5545
|
||||
0.125
|
||||
0.032
|
||||
0.329
|
||||
0.099
|
||||
0.04
|
||||
0.16
|
||||
0.066
|
||||
0.416
|
||||
0.092
|
||||
0.067
|
||||
DeepSeek-OCR (end2end)
|
||||
Tiny
|
||||
64
|
||||
0.386
|
||||
0.373
|
||||
0.469
|
||||
0.422
|
||||
0.283
|
||||
0.361
|
||||
0.307
|
||||
0.635
|
||||
0.266
|
||||
0.236
|
||||
Small
|
||||
100
|
||||
0.221
|
||||
0.142
|
||||
0.373
|
||||
0.242
|
||||
0.125
|
||||
0.284
|
||||
0.24
|
||||
0.53
|
||||
0.159
|
||||
0.205
|
||||
Base
|
||||
256(182)
|
||||
0.137
|
||||
0.054
|
||||
0.267
|
||||
0.163
|
||||
0.064
|
||||
0.24
|
||||
0.205
|
||||
0.474
|
||||
0.1
|
||||
0.181
|
||||
Large
|
||||
400(285)
|
||||
0.138
|
||||
0.054
|
||||
0.277
|
||||
0.152
|
||||
0.067
|
||||
0.208
|
||||
0.143
|
||||
0.461
|
||||
0.104
|
||||
0.123
|
||||
Gundam
|
||||
795
|
||||
0.127
|
||||
0.043
|
||||
0.269
|
||||
0.134
|
||||
0.062
|
||||
0.181
|
||||
0.097
|
||||
0.432
|
||||
0.089
|
||||
0.103
|
||||
Gundam-M †200dpi
|
||||
1853
|
||||
0.123
|
||||
0.049
|
||||
0.242
|
||||
0.147
|
||||
0.056
|
||||
0.157
|
||||
0.087
|
||||
0.377
|
||||
0.08
|
||||
0.085
|
||||
## 4 Evaluation
|
||||
## 4.1 Vision-text Compression Study
|
||||
We select Fox [ 21 ] benchmarks to verify DeepSeek-OCR’s compression-decompression capability for text-rich documents, in order to preliminarily explore the feasibility and boundaries of contexts optical compression. We use the English document portion of Fox, tokenize the ground truth text with DeepSeek-OCR’s tokenizer (vocabulary size of approximately 129k), and select documents with 600-1300 tokens for testing, which happens to be 100 pages. Since the number of text tokens is not large, we only need to test performance in Tiny and Small modes, where Tiny mode corresponds to 64 tokens and Small mode corresponds to 100 tokens. We use the prompt without layout: "<image> \n Free OCR." to control the model’s output format. Nevertheless, the output format still cannot completely match Fox benchmarks, so the actual performance would be somewhat higher than the test results.
|
||||
As shown in Table 2 , within a 10 × \times compression ratio, the model’s decoding precision can reach approximately 97%, which is a very promising result. In the future, it may be possible to achieve nearly 10 × \times lossless contexts compression through text-to-image approaches. When the compression ratio exceeds 10 × \times , performance begins to decline, which may have two reasons: one is that the layout of long documents becomes more complex, and another reason may be that long texts become blurred at 512 × \times 512 or 640 × \times 640 resolution. The first issue can be solved by rendering texts onto a single layout page, while we believe the second issue will become a feature of the forgetting mechanism. When compressing tokens by nearly 20 × \times , we find that precision can still approach 60%. These results indicate that optical contexts compression is a very promising and worthwhile research direction, and this approach does not bring any overhead because it can leverage VLM infrastructure, as multimodal systems inherently require an additional vision encoder.
|
||||
Table 4: Edit distances for different categories of documents in OmniDocBench. The results show that some types of documents can achieve good performance with just 64 or 100 vision tokens, while others require Gundam mode.
|
||||
Mode Type
|
||||
Book
|
||||
Slides
|
||||
Financial
|
||||
Report
|
||||
Textbook
|
||||
Exam
|
||||
Paper
|
||||
Magazine
|
||||
Academic
|
||||
Papers
|
||||
Notes
|
||||
Newspaper
|
||||
Overall
|
||||
Tiny
|
||||
0.147
|
||||
0.116
|
||||
0.207
|
||||
0.173
|
||||
0.294
|
||||
0.201
|
||||
0.395
|
||||
0.297
|
||||
0.94
|
||||
0.32
|
||||
Small
|
||||
0.085
|
||||
0.111
|
||||
0.079
|
||||
0.147
|
||||
0.171
|
||||
0.107
|
||||
0.131
|
||||
0.187
|
||||
0.744
|
||||
0.205
|
||||
Base
|
||||
0.037
|
||||
0.08
|
||||
0.027
|
||||
0.1
|
||||
0.13
|
||||
0.073
|
||||
0.052
|
||||
0.176
|
||||
0.645
|
||||
0.156
|
||||
Large
|
||||
0.038
|
||||
0.108
|
||||
0.022
|
||||
0.084
|
||||
0.109
|
||||
0.06
|
||||
0.053
|
||||
0.155
|
||||
0.353
|
||||
0.117
|
||||
Gundam
|
||||
0.035
|
||||
0.085
|
||||
0.289
|
||||
0.095
|
||||
0.094
|
||||
0.059
|
||||
0.039
|
||||
0.153
|
||||
0.122
|
||||
0.083
|
||||
Guandam-M
|
||||
0.052
|
||||
0.09
|
||||
0.034
|
||||
0.091
|
||||
0.079
|
||||
0.079
|
||||
0.048
|
||||
0.1
|
||||
0.099
|
||||
0.077
|
||||
## 4.2 OCR Practical Performance
|
||||
DeepSeek-OCR is not only an experimental model; it has strong practical capabilities and can construct data for LLM/VLM pretraining. To quantify OCR performance, we test DeepSeek-OCR on OmniDocBench [ 27 ] , with results shown in Table 3 . Requiring only 100 vision tokens (640 × \times 640 resolution), DeepSeek-OCR surpasses GOT-OCR2.0 [ 38 ] which uses 256 tokens; with 400 tokens (285 valid tokens, 1280 × \times 1280 resolution), it achieves on-par performance with state-of-the-arts on this benchmark. Using fewer than 800 tokens (Gundam mode), DeepSeek-OCR outperforms MinerU2.0 [ 34 ] which needs nearly 7,000 vision tokens. These results demonstrate that our DeepSeek-OCR model is powerful in practical applications, and because the higher tokens compression, it enjoys a higher research ceiling.
|
||||
As shown in Table 4 , some categories of documents require very few tokens to achieve satisfactory performance, such as slides which only need 64 vision tokens. For book and report documents, DeepSeek-OCR can achieve good performance with only 100 vision tokens. Combined with the analysis from Section 4.1 , this may be because most text tokens in these document categories are within 1,000, meaning the vision-token compression ratio does not exceed 10 × \times . For newspapers, Gundam or even Gundam-master mode is required to achieve acceptable edit distances, because the text tokens in newspapers are 4-5,000, far exceeding the 10 × \times compression of other modes. These experimental results further demonstrate the boundaries of contexts optical compression, which may provide effective references for researches on the vision token optimization in VLMs and context compression, forgetting mechanisms in LLMs.
|
||||
## 4.3 Qualitative Study
|
||||
## 4.3.1 Deep parsing
|
||||
DeepSeek-OCR possesses both layout and OCR 2.0 capabilities, enabling it to further parse images within documents through secondary model calls, a feature we refer to as "deep parsing". As shown in Figures 7 , 8 , 9 , 10 , our model can perform deep parsing on charts, geometry, chemical formulas, and even natural images, requiring only a unified prompt.
|
||||
Figure 7 : In the field of financial research reports, the deep parsing mode of DeepSeek-OCR can be used to obtain structured results of charts within documents. Charts are a crucial form of data representation in finance and scientific fields, and the chart structured extraction is an indispensable capability for future OCR models.
|
||||
Figure 8 : For books and articles, the deep parsing mode can output dense captions for natural images in the documents. With just a prompt, the model can automatically identify what type of image it is and output the required results.
|
||||
Figure 9 : DeepSeek-OCR in deep parsing mode can also recognize chemical formulas within chemical documents and convert them to SMILES format. In the future, OCR 1.0+2.0 technology may play a significant role in the development of VLM/LLM in STEM fields.
|
||||
Figure 10 : DeepSeek-OCR also possesses the capability to copy (structure) simple planar geometric figures. Due to the intricate interdependencies among line segments in geometric shapes, parsing geometry task is extremely challenging and has a long way to go.
|
||||
## 4.3.2 Multilingual recognition
|
||||
PDF data on the Internet contains not only Chinese and English, but also a large amount of multilingual data, which is also crucial when training LLMs. For PDF documents, DeepSeek-OCR can handle nearly 100 languages. Like Chinese and English documents, multilingual data also supports both layout and non-layout OCR formats. The visualization results are shown in Figure 11 , where we select Arabic and Sinhala languages to demonstrate results.
|
||||
Figure 11 : To endow the capability of processing widely crawled PDFs (multilingual data), we train our model with OCR capabilities for nearly 100 languages. Minority language documents can also support both layout and non-layout outputs through different prompts.
|
||||
## 4.3.3 General vision understanding
|
||||
We also provide DeepSeek-OCR with a certain degree of general image understanding capabilities. The related visualization results are shown in Figure 12 .
|
||||
Figure 12 : We retain DeepSeek-OCR’s capabilities in general visual understanding, mainly including image description, object detection, grounding, etc. Meanwhile, due to the inclusion of text-only data, DeepSeek-OCR’s language capabilities are also retained. Note that since we do not include SFT (Supervised Fine-Tuning) stage, the model is not a chatbot, and some capabilities need completion prompts to be activated.
|
||||
Figure 13 : Forgetting mechanisms constitute one of the most fundamental characteristics of human memory. The contexts optical compression approach can simulate this mechanism by rendering previous rounds of historical text onto images for initial compression, then progressively resizing older images to achieve multi-level compression, where token counts gradually decrease and text becomes increasingly blurred, thereby accomplishing textual forgetting.
|
||||
## 5 Discussion
|
||||
Our work represents an initial exploration into the boundaries of vision-text compression, investigating how many vision tokens are required to decode N N text tokens. The preliminary results are encouraging: DeepSeek-OCR achieves near-lossless OCR compression at approximately 10 × \times ratios, while 20 × \times compression still retains 60% accuracy. These findings suggest promising directions for future applications, such as implementing optical processing for dialogue histories beyond k k rounds in multi-turn conversations to achieve 10 × \times compression efficiency.
|
||||
For older contexts, we could progressively downsizing the rendered images to further reduce token consumption. This assumption draws inspiration from the natural parallel between human memory decay over time and visual perception degradation over spatial distance—both exhibit similar patterns of progressive information loss, as shown in Figure 13 . By combining these mechanisms, contexts optical compression method enables a form of memory decay that mirrors biological forgetting curves, where recent information maintains high fidelity while distant memories naturally fade through increased compression ratios.
|
||||
While our initial exploration shows potential for scalable ultra-long context processing, where recent contexts preserve high resolution and older contexts consume fewer resources, we acknowledge this is early-stage work that requires further investigation. The approach suggests a path toward theoretically unlimited context architectures that balance information retention with computational constraints, though the practical implications and limitations of such vision-text compression systems warrant deeper study in future research.
|
||||
## 6 Conclusion
|
||||
In this technical report, we propose DeepSeek-OCR and preliminarily validate the feasibility of contexts optical compression through this model, demonstrating that the model can effectively decode text tokens exceeding 10 times the quantity from a small number of vision tokens. We believe this finding will facilitate the development of VLMs and LLMs in the future. Additionally, DeepSeek-OCR is a highly practical model capable of large-scale pretraining data production, serving as an indispensable assistant for LLMs. Of course, OCR alone is insufficient to fully validate true context optical compression and we will conduct digital-optical text interleaved pretraining, needle-in-a-haystack testing, and other evaluations in the future. From another perspective, optical contexts compression still offers substantial room for research and improvement, representing a promising new direction.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,336 @@
|
||||
TITLE: DeepSeekMath-V2: Towards Self-Verifiable Mathematical Reasoning
|
||||
ARXIV: 2511.22570
|
||||
AUTHORS: Zhihong Shao , Yuxiang Luo , Chengda Lu , Z.Z. Ren , Jiewen Hu , Tian Ye , Zhibin Gou , Shirong Ma , Xiaokang Zhang
|
||||
DATE: [Submitted on 27 Nov 2025]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeekMath-V2: Towards Self-Verifiable Mathematical Reasoning
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Method
|
||||
## 2.1 Proof Verification
|
||||
## 2.1.1 Training a Verifier to Identify Issues and Score Proofs
|
||||
## Curating Cold Start RL Data
|
||||
## RL Objective.
|
||||
## 2.1.2 Introducing Meta-Verification to Review Proof Analyses
|
||||
## Meta-Verifier Training Process
|
||||
## 2.2 Proof Generation
|
||||
## 2.2.1 Training a Generator for Theorem Proving
|
||||
## 2.2.2 Enhancing Reasoning via Self-Verification
|
||||
## 2.3 Synergy Between Proof Verification and Generation
|
||||
## 3 Experiments
|
||||
## 3.1 Training Settings
|
||||
## 3.2 Evaluation Benchmarks
|
||||
## 3.3 Evaluation Results
|
||||
## 3.3.1 One-Shot Generation
|
||||
## 3.3.2 Sequential Refinement with Self-Verification
|
||||
## 3.3.3 High-Compute Search
|
||||
## 4 Related Work
|
||||
## 5 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeekMath-V2: Towards Self-Verifiable Mathematical Reasoning
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Method
|
||||
2.1 Proof Verification
|
||||
2.1.1 Training a Verifier to Identify Issues and Score Proofs
|
||||
Curating Cold Start RL Data
|
||||
RL Objective.
|
||||
2.1.2 Introducing Meta-Verification to Review Proof Analyses
|
||||
Meta-Verifier Training Process
|
||||
2.2 Proof Generation
|
||||
2.2.1 Training a Generator for Theorem Proving
|
||||
2.2.2 Enhancing Reasoning via Self-Verification
|
||||
2.3 Synergy Between Proof Verification and Generation
|
||||
3 Experiments
|
||||
3.1 Training Settings
|
||||
3.2 Evaluation Benchmarks
|
||||
3.3 Evaluation Results
|
||||
3.3.1 One-Shot Generation
|
||||
3.3.2 Sequential Refinement with Self-Verification
|
||||
3.3.3 High-Compute Search
|
||||
4 Related Work
|
||||
5 Conclusion
|
||||
References
|
||||
A Prompt Templates
|
||||
A.1 Proof Generation Prompt
|
||||
A.2 Proof Verification Prompt
|
||||
A.3 Meta-Verification Prompt
|
||||
A.4 Proof Refinement Prompt
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2511.22570v1 [cs.AI] 27 Nov 2025
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeekMath-V2: Towards Self-Verifiable Mathematical Reasoning
|
||||
Zhihong Shao*, Yuxiang Luo*, Chengda Lu* † , Z.Z. Ren*
|
||||
Jiewen Hu, Tian Ye, Zhibin Gou, Shirong Ma, Xiaokang Zhang
|
||||
DeepSeek-AI
|
||||
zhihongshao@deepseek.com
|
||||
https://github.com/deepseek-ai/DeepSeek-Math-V2
|
||||
Corresponding author: Core contributors † Work done during internship at DeepSeek-AI.
|
||||
## Abstract
|
||||
Large language models have made significant progress in mathematical reasoning, which serves as an important testbed for AI and could impact scientific research if further advanced.
|
||||
By scaling reasoning with reinforcement learning that rewards correct final answers, LLMs have improved from poor performance to saturating quantitative reasoning competitions like AIME and HMMT in one year.
|
||||
However, this approach faces fundamental limitations.
|
||||
Pursuing higher final answer accuracy doesn’t address a key issue: correct answers don’t guarantee correct reasoning.
|
||||
Moreover, many mathematical tasks like theorem proving require rigorous step-by-step derivation rather than numerical answers, making final answer rewards inapplicable.
|
||||
To push the limits of deep reasoning, we believe it is necessary to verify the comprehensiveness and rigor of mathematical reasoning.
|
||||
Self-verification is particularly important for scaling test-time compute, especially for open problems without known solutions.
|
||||
Towards self-verifiable mathematical reasoning, we investigate how to train an accurate and faithful LLM-based verifier for theorem proving.
|
||||
We then train a proof generator using the verifier as the reward model, and incentivize the generator to identify and resolve as many issues as possible in their own proofs before finalizing them.
|
||||
To maintain the generation-verification gap as the generator becomes stronger, we propose to scale verification compute to automatically label new hard-to-verify proofs, creating training data to further improve the verifier.
|
||||
Our resulting model, DeepSeekMath-V2, demonstrates strong theorem-proving capabilities, achieving gold-level scores on IMO 2025 and CMO 2024 and a near-perfect 118/120 on Putnam 2024 with scaled test-time compute.
|
||||
While much work remains, these results suggest that self-verifiable mathematical reasoning is a feasible research direction that may help develop more capable mathematical AI systems.
|
||||
## 1 Introduction
|
||||
The conventional approach to reinforcement learning (RL) for mathematical reasoning involves rewarding large language models (LLMs) based on whether their predicted final answers to quantitative reasoning problems match ground-truth answers ( Guo et al., 2025 ) .
|
||||
This methodology suffices to allow frontier LLMs to saturate mathematical competitions that primarily evaluate final answers, such as AIME and HMMT.
|
||||
However, this reward mechanism has two fundamental limitations.
|
||||
First, it serves as an unreliable proxy for reasoning correctness – a model can arrive at the correct answer through flawed logic or fortunate errors.
|
||||
Second, it is inapplicable to theorem proving tasks, where problems may not require producing numerical final answers and rigorous derivation is the primary objective.
|
||||
Consequently, LLMs trained on quantitative reasoning problems with such final answer reward still frequently produce mathematically invalid or logically inconsistent natural-language proofs.
|
||||
Moreover, this training approach does not naturally develop the models’ ability to verify proof validity – they exhibit high false-positive rates, often claiming incorrect proofs are valid even when they contain obvious logical flaws.
|
||||
The lack of a generation-verification gap in natural-language theorem proving hinders further improvement.
|
||||
To address this, we propose developing proof verification capabilities in LLMs.
|
||||
Our approach is motivated by several key observations:
|
||||
•
|
||||
Humans can identify issues in proofs even without reference solutions – a crucial ability when tackling open problems.
|
||||
•
|
||||
A proof is more likely to be valid when no issues can be identified despite scaled verification efforts.
|
||||
•
|
||||
The efforts required to identify valid issues can serve as a proxy for proof quality, which can be exploited to optimize proof generation.
|
||||
We believe that LLMs can be trained to identify proof issues without reference solutions.
|
||||
Such a verifier would enable an iterative improvement cycle: (1) using verification feedback to optimize proof generation, (2) scaling verification compute to auto-label hard-to-verify new proofs, thereby creating the training data to improve the verifier itself, and (3) using this enhanced verifier to further optimize proof generation.
|
||||
Moreover, a reliable proof verifier enables us to teach proof generators to evaluate proofs as the verifier does.
|
||||
This allows a proof generator to iteratively refine its proofs until it can no longer identify or resolve any issues.
|
||||
In essence, we make the model explicitly aware of its reward function and enable it to maximize this reward through deliberate reasoning rather than blind trial-and-error.
|
||||
Built on DeepSeek-V3.2-Exp-Base ( DeepSeek-AI, 2025 ) , we developed DeepSeekMath-V2 , a large language model optimized for natural-language theorem proving that demonstrates self-verifiable mathematical reasoning.
|
||||
Our model can assess and iteratively improve its own proofs, achieving gold-level performance in premier high-school mathematics competitions including IMO 2025 and CMO 2024.
|
||||
On the Putnam 2024 undergraduate competition, it scored 118/120, exceeding the highest score of 90 1 1
|
||||
1
|
||||
https://kskedlaya.org/putnam-archive/putnam2024stats.html obtained by human participants.
|
||||
## 2 Method
|
||||
## 2.1 Proof Verification
|
||||
## 2.1.1 Training a Verifier to Identify Issues and Score Proofs
|
||||
We developed high-level rubrics ℐ v \mathcal{I}_{v} for proof evaluation (see Appendix A.2 ) with the goal of training a verifier to evaluate proofs according to these rubrics, mirroring mathematical experts’ assessment process.
|
||||
Specifically, given a problem X X and a proof Y Y , the verifier π φ ( ⋅ | X , Y , ℐ v ) \pi_{\varphi}(\cdot|{}X,Y,\mathcal{I}_{v}) is designed to produce a proof analysis that first summarizes identified issues (if any) and then assigns a score based on three levels:
|
||||
1 for complete and rigorous proofs with all logical steps clearly justified;
|
||||
0.5 for proofs with sound overall logic but minor errors or omitted details;
|
||||
and 0 for fundamentally flawed proofs containing fatal logical errors or critical gaps.
|
||||
## Curating Cold Start RL Data
|
||||
We constructed our initial training data through the following process:
|
||||
1.
|
||||
We crawled problems from Art of Problem Solving (AoPS) contests 2 2
|
||||
2
|
||||
https://artofproblemsolving.com/community/c13_contest_collections , prioritizing math olympiads, team selection tests, and post-2010 problems explicitly requiring proofs, totaling 17,503 problems.
|
||||
This problem set is denoted as 𝒟 p \mathcal{D}_{p} .
|
||||
2.
|
||||
We generated candidate proofs using a variant of DeepSeek-V3.2-Exp-Thinking. As this model was not optimized for theorem proving and tended to produce concise but error-prone outputs, we prompted it to iteratively refine its proofs over multiple rounds to improve comprehensiveness and rigor.
|
||||
3.
|
||||
We randomly sampled proofs across diverse problem types (e.g., algebra and number theory) and had mathematical experts score each proof according to the evaluation rubrics described above.
|
||||
This process yielded an initial RL dataset 𝒟 v = { ( X i , Y i , s i ) } \mathcal{D}_{v}=\{(X_{i},Y_{i},s_{i})\} , where each item consists of a problem X i X_{i} , a proof Y i Y_{i} , and an overall proof score s i ∈ { 0 , 0.5 , 1 } s_{i}\in\{0,0.5,1\} .
|
||||
## RL Objective.
|
||||
Building on a version of DeepSeek-V3.2-Exp-SFT which was supervised fine-tuned on reasoning data related to mathematics and code, we trained the model with reinforcement learning to produce proof analyses using two reward components:
|
||||
•
|
||||
Format reward R format R_{\text{format}} : An indicator function that enforces the model to generate both a summary of identified issues and a proof score, by checking whether the final response contains the key phrase “Here is my evaluation of the solution:” as well as a score within \boxed{} following “Based on my evaluation, the final overall score should be:”.
|
||||
•
|
||||
Score reward R score R_{\text{score}} : Rewards based on proximity between predicted score s i ′ s_{i}^{\prime} and annotated score s i s_{i} :
|
||||
R score ( s i ′ , s i ) = 1 − | s i ′ − s i | R_{\text{score}}(s_{i}^{\prime},s_{i})=1-|s_{i}^{\prime}-s_{i}|
|
||||
(1)
|
||||
The RL objective for training the verifier is:
|
||||
max π φ 𝔼 ( X i , Y i , s i ) ∼ 𝒟 v , ( V i ′ , s i ′ ) ∼ π φ ( ⋅ | X i , Y i ) [ R format ( V i ′ ) ⋅ R score ( s i ′ , s i ) ] \max_{\pi_{\varphi}}\mathbb{E}_{(X_{i},Y_{i},s_{i})\sim\mathcal{D}_{v},(V_{i}^{\prime},s_{i}^{\prime})\sim\pi_{\varphi}(\cdot|X_{i},Y_{i})}\left[R_{\text{format}}(V_{i}^{\prime})\cdot R_{\text{score}}(s_{i}^{\prime},s_{i})\right]
|
||||
(2)
|
||||
where V i ′ V_{i}^{\prime} denotes the verifier’s final response and s i ′ s_{i}^{\prime} is the proof score extracted from it.
|
||||
## 2.1.2 Introducing Meta-Verification to Review Proof Analyses
|
||||
The approach described in Section 2.1.1 trains proof verification through RL to align predicted proof scores with expert annotations, but provides no direct supervision on the identified issues themselves.
|
||||
This creates a critical vulnerability: when evaluating flawed proofs (where s i < 1 s_{i}<1 ) during training, the verifier can receive full reward by predicting the correct scores while hallucinating non-existent issues, undermining its trustworthiness.
|
||||
To address this problem, we introduce meta-verification : a secondary evaluation process that assesses whether issues identified by the verifier indeed exist and whether these issues logically justify the predicted proof score according to the evaluation rubrics ℐ v \mathcal{I}_{v} .
|
||||
The complete meta-verification rubrics ℐ m v \mathcal{I}_{mv} are detailed in Appendix A.3 .
|
||||
We trained a dedicated meta-verifier using RL to perform this evaluation.
|
||||
By incorporating the meta-verifier’s feedback into verifier training, we can improve the faithfulness of the verifier’s issue identification.
|
||||
## Meta-Verifier Training Process
|
||||
1.
|
||||
We obtained an initial verifier π φ \pi_{\varphi} following Section 2.1.1 .
|
||||
2.
|
||||
Mathematical experts scored the quality of verifier responses according to ℐ m v \mathcal{I}_{mv} , creating dataset 𝒟 m v = { ( X i , Y i , V i , m s i ) } \mathcal{D}_{mv}=\{(X_{i},Y_{i},V_{i},ms_{i})\} , where V i V_{i} is the analysis of proof Y i Y_{i} and m s i ∈ { 0 , 0.5 , 1 } ms_{i}\in\{0,0.5,1\} is the expert-annotated quality score.
|
||||
3.
|
||||
We trained a meta-verifier π η ( ⋅ | X , Y , V , ℐ m v ) \pi_{\eta}(\cdot|X,Y,V,\mathcal{I}_{mv}) to analyze the verifier’s proof analysis V V .
|
||||
The meta-verifier produces a summary of issues found in the analysis itself, followed by a quality score measuring how accurate and justified the verifier’s analysis is.
|
||||
The RL objective follows the same structure as the verifier training, with format and score rewards.
|
||||
Using the trained meta-verifier π η \pi_{\eta} , we enhanced the verifier training by integrating meta-verification feedback into the reward function:
|
||||
R V = R format ⋅ R score ⋅ R meta R_{V}=R_{\text{format}}\cdot R_{\text{score}}\cdot R_{\text{meta}}
|
||||
(3)
|
||||
where R meta R_{\text{meta}} is the quality score from the meta-verifier.
|
||||
We trained the enhanced verifier on both the verification dataset 𝒟 v \mathcal{D}_{v} and the meta-verification dataset 𝒟 m v \mathcal{D}_{mv} , using the same reward mechanism on 𝒟 m v \mathcal{D}_{mv} as used for training the meta-verifier.
|
||||
The resulting model can perform both proof verification and meta-verification tasks.
|
||||
On a validation split of 𝒟 v \mathcal{D}_{v} , the average quality score of the verifier’s proof analyses – as evaluated by the meta-verifier – improved from 0.85 to 0.96, while maintaining the same accuracy in proof score prediction.
|
||||
## 2.2 Proof Generation
|
||||
## 2.2.1 Training a Generator for Theorem Proving
|
||||
With verifier π φ \pi_{\varphi} serving as a generative reward model, we train a proof generator π θ ( ⋅ | X ) \pi_{\theta}(\cdot|{}X) with the RL objective:
|
||||
max π θ 𝔼 X i ∼ 𝒟 p , Y i ∼ π θ ( ⋅ | X i ) [ R Y ] \max_{\pi_{\theta}}\mathbb{E}_{X_{i}\sim\mathcal{D}_{p},Y_{i}\sim\pi_{\theta}(\cdot|X_{i})}[R_{Y}]
|
||||
(4)
|
||||
where R Y R_{Y} is the proof score produced by π φ ( ⋅ | X i , Y i , ℐ v ) \pi_{\varphi}(\cdot|{}X_{i},Y_{i},\mathcal{I}_{v}) .
|
||||
## 2.2.2 Enhancing Reasoning via Self-Verification
|
||||
When a proof generator fails to produce a completely correct proof in one shot – common for challenging problems from competitions like IMO and CMO – iterative verification and refinement can improve results.
|
||||
This involves analyzing the proof with an external verifier and prompting the generator to address identified issues.
|
||||
However, we observed a critical limitation: when prompted to both generate and analyze its own proof in one shot, the generator tends to claim correctness even when the external verifier easily identify flaws.
|
||||
In other words, while the generator can refine proofs based on external feedback, it fails to evaluate its own work with the same rigor as the dedicated verifier.
|
||||
This observation motivated us to endow the proof generator with genuine verification capabilities.
|
||||
During training, we prompt the generator π θ \pi_{\theta} to produce a proof Y Y followed by a self-analysis Z Z that follows the same format and rubrics ℐ v \mathcal{I}_{v} as the verifier (see Appendix A.1 ).
|
||||
We denote the proof score predicted in the self-analysis as s ′ s^{\prime} .
|
||||
To ensure faithful self-evaluation, we use the verifier π φ \pi_{\varphi} to assess both components: the proof Y Y receives score R Y = s R_{Y}=s , and the self-analysis Z Z receives a meta-verification score R meta ( Z ) = m s R_{\text{meta}}(Z)=ms .
|
||||
The reward function combines these assessments:
|
||||
R \displaystyle R
|
||||
= R format ( Y , Z ) ⋅ ( α ⋅ R Y + β ⋅ R Z ) \displaystyle=R_{\text{format}}(Y,Z)\cdot(\alpha\cdot R_{Y}+\beta\cdot R_{Z})
|
||||
(5)
|
||||
R Z \displaystyle R_{Z}
|
||||
= R score ( s ′ , s ) ⋅ R meta ( Z ) \displaystyle=R_{\text{score}}(s^{\prime},s)\cdot R_{\text{meta}}(Z)
|
||||
(6)
|
||||
where R format ( Y , Z ) R_{\text{format}}(Y,Z) verifies that both the proof and self-analysis follow the specified format, R score ( s ′ , s ) R_{\text{score}}(s^{\prime},s) rewards accurate self-assessment.
|
||||
We set α = 0.76 \alpha=0.76 and β = 0.24 \beta=0.24 .
|
||||
This reward structure creates the following incentives:
|
||||
•
|
||||
Faithful acknowledgment of errors is rewarded over false claims of correctness.
|
||||
•
|
||||
The highest rewards come from producing correct proofs and accurately recognizing their rigor.
|
||||
•
|
||||
A good strategy to obtain high rewards for the proof generator is to identify and resolve as many issues as possible before finalizing the response.
|
||||
## 2.3 Synergy Between Proof Verification and Generation
|
||||
The proof verifier and generator create a synergistic cycle:
|
||||
the verifier improves the generator, and as the generator improves, it produces new proofs that challenge the verifier’s current capabilities.
|
||||
These challenging cases – where the verifier may fail to identify issues in a single attempt – become valuable training data for enhancing the verifier itself.
|
||||
To retrain and improve the verifier, we need labeled correctness data for newly generated proofs.
|
||||
Manual annotation, while straightforward, becomes increasingly time-consuming as problems grow harder and errors become more subtle.
|
||||
To boost annotation efficiency, we generated multiple verifier analyses per proof to surface potential issues for human review.
|
||||
From this AI-assisted annotation process, we recognized two facts that make it feasible to push the level of automation a step further:
|
||||
1.
|
||||
Scaling verifier samples increases the probability of catching real issues in flawed proofs.
|
||||
2.
|
||||
Reviewing the verifier’s identified issues is exactly meta-verification , which is easier than identifying issues from scratch. Meta-verification is also more sample-efficient for LLMs to master.
|
||||
Building on these observations, we developed the following automated labeling process:
|
||||
1.
|
||||
For each proof, generate n n independent verification analyses
|
||||
2.
|
||||
For analyses reporting issues (scores 0 or 0.5), generate m m meta-verification assessments to validate the identified problems.
|
||||
An analysis is deemed valid if the majority of meta-assessments confirm its findings
|
||||
3.
|
||||
For each proof, we examine analyses that assign the lowest score.
|
||||
If at least k k such analyses are deemed valid, the proof is labeled with that lowest score.
|
||||
If no legitimate issues are identified across all verification attempts, the proof is labeled with 1.
|
||||
Otherwise, the proof is discarded or routed to human experts for labeling
|
||||
In our last two training iterations, this fully automated pipeline replaced human annotation entirely.
|
||||
Quality checks confirmed that the automated labels aligned well with expert judgments.
|
||||
## 3 Experiments
|
||||
Figure 1: Average proof scores on CNML-level problems by category and model, as evaluated by our verifier.
|
||||
## 3.1 Training Settings
|
||||
We employed Group Relative Policy Optimization (GRPO) ( Shao et al., 2024 ) for reinforcement learning, iteratively optimizing proof verification and generation capabilities as described in Section 2 .
|
||||
In each iteration, we first optimized proof verification. The proof generator was then initialized from the verifier checkpoint and optimized for proof generation.
|
||||
Starting from the second iteration, the proof verifier was initialized with a checkpoint that consolidated both verification and generation capabilities from the previous iteration through rejection fine-tuning.
|
||||
## 3.2 Evaluation Benchmarks
|
||||
We evaluate our final proof generator on the following theorem proving benchmarks:
|
||||
In-House CNML-Level Problems 91 theorem-proving problems spanning algebra (13), geometry (24), number theory (19), combinatorics (24), and inequality (11), comparable in difficulty to problems from Chinese National High School Mathematics League (CNML)
|
||||
Competition Problems
|
||||
•
|
||||
IMO 2025 (6 problems): The International Mathematical Olympiad, the premier global mathematics competition for pre-university students
|
||||
•
|
||||
CMO 2024 (6 problems): The China Mathematical Olympiad, China’s national championship
|
||||
•
|
||||
Putnam 2024 (12 problems): The William Lowell Putnam Competition, the preeminent mathematics competition for undergraduate students in North America
|
||||
•
|
||||
ISL 2024 (31 problems): The IMO Shortlist, a collection of problems proposed by participating countries and considered by the Problem Selection Committee for potential inclusion in IMO 2024
|
||||
•
|
||||
IMO-ProofBench (60 problems): Developed by the DeepMind team behind DeepThink IMO-Gold ( Luong and Lockhart, 2025 ) , this benchmark ( Luong et al., 2025 ) is divided into a basic set (30 problems, pre-IMO to IMO-Medium difficulty) and an advanced set (30 challenging problems simulating complete IMO examinations, up to IMO-Hard level)
|
||||
## 3.3 Evaluation Results
|
||||
## 3.3.1 One-Shot Generation
|
||||
We first evaluate the model’s ability to generate correct proofs without iterative refinement.
|
||||
On the in-house problems, we generated 8 proof samples per problem for each evaluated model.
|
||||
Proof correctness was measured by majority voting across 8 verification analyses produced by our final verifier.
|
||||
As shown in Figure 1 , across all categories of CNML-level problems – algebra, geometry, number theory, combinatorics, and inequality – DeepSeekMath-V2 consistently outperforms GPT-5-Thinking-High ( OpenAI, 2025 ) and Gemini 2.5-Pro ( DeepMind, 2025 ) , demonstrating superior theorem-proving ability across domains.
|
||||
## 3.3.2 Sequential Refinement with Self-Verification
|
||||
Figure 2: Proof quality improvements as the maximum sequential iterations varies from 1 (no refinement) to 8 (initial generation plus up to 7 refinements based on self verification).
|
||||
For challenging problems from competitions like IMO and CMO, models often cannot generate comprehensive and rigorous proofs in a single attempt within the 128K token limit.
|
||||
When this occurs, our proof generator recognizes its proof is invalid through self-verification but lacks the context length to resolve all identified issues in a single attempt.
|
||||
To explore how extended context and self-verification can improve proof quality, we evaluate sequential refinement with self-verification.
|
||||
This approach first generates a proof with self-analysis, then iteratively re-prompts the generator with its previous output (see Appendix A.4 for the refinement prompt), allowing it to address identified issues.
|
||||
The process continues until the generator assigns itself a perfect score or reaches the maximum number of sequential attempts.
|
||||
Figure 2 demonstrates proof quality improvement through sequential refinement on IMO Shortlist 2024 problems.
|
||||
For each problem, we launched 32 independent refinement threads.
|
||||
Proof correctness was measured by majority voting across 32 verification analyses from our final verifier.
|
||||
We report two metrics in Figure 2 : (1) Pass@1 – the average score of the final proof from each thread, and (2) Best@32 – the score of the best proof per problem, selected by self-assigned scores across all threads.
|
||||
The self-selected best proofs achieve significantly higher verification scores than the thread average, demonstrating our generator’s ability to accurately assess proof quality.
|
||||
Furthermore, Pass@1 improves substantially as maximum sequential attempts increase, showing that self-verification effectively guides iterative improvement.
|
||||
These results confirm that our generator can reliably differentiate between high-quality and flawed proofs, and leverage this self-awareness to systematically improve its mathematical reasoning.
|
||||
Figure 3: Expert evaluation results on the Basic and Advanced subsets of IMO-ProofBench. All results are sourced from Luong et al. (2025) , with the exception of DeepSeekMath-V2, which was evaluated by our experts following the grading guidelines.
|
||||
## 3.3.3 High-Compute Search
|
||||
Contest
|
||||
Problems
|
||||
Points
|
||||
IMO 2025
|
||||
P1 , P2 , P3 , P4 , P5
|
||||
83.3%
|
||||
CMO 2024
|
||||
P1 , P2 , P4 , P5 , P6
|
||||
73.8%
|
||||
Putnam 2024
|
||||
A1 ∼ \sim B4 , B5 , B6
|
||||
98.3%
|
||||
Table 1: Problems in gray are fully solved , while underlined problems received partial credit .
|
||||
To solve the most challenging problems, we scaled both verification and generation compute – using extensive verification to identify subtle issues and parallel generation to explore diverse proof strategies.
|
||||
Our approach maintains a pool of candidate proofs for each problem, initialized with 64 proof samples with 64 verification analyses generated for each.
|
||||
In each refinement iteration, we select the 64 highest-scoring proofs based on average verification scores and pair each with 8 randomly selected analyses, prioritizing those identifying issues (scores 0 or 0.5).
|
||||
Each proof-analysis pair is used to generate one refined proof, which then updates the candidate pool.
|
||||
This process continues for up to 16 iterations or until a proof successfully passes all 64 verification attempts, indicating high confidence in correctness.
|
||||
All experiments used a single model, our final proof generator, which performs both proof generation and verification.
|
||||
To validate our results, mathematical experts assessed the highest-scoring proofs.
|
||||
As shown in Table 1 , our approach solved 5 of 6 problems from IMO 2025 and 4 problems plus partial credit on another from CMO 2024, achieving gold medal performance in both pinnacle high-school competitions.
|
||||
On Putnam 2024, the preeminent undergraduate mathematics competition, our model solved 11 of 12 problems completely and the remaining problem with minor errors, scoring 118/120 and surpassing the highest human score of 90.
|
||||
Figure 3 shows the results on IMO-ProofBench.
|
||||
Our approach outperforms DeepMind’s DeepThink (IMO Gold) on the basic set and remains competitive on the advanced set, while substantially outperforming all other baselines.
|
||||
We observe that the hardest IMO-level problems remain challenging for our model.
|
||||
Notably, for problems not fully solved, our generator typically identifies the genuine issues in its proofs, while fully solved problems pass all 64 verification attempts.
|
||||
This demonstrates that we can successfully train LLM-based verifiers to assess proofs previously considered difficult to verify automatically.
|
||||
By scaling test-time compute under verifier guidance, our model solves problems that require hours of effort from human competitors.
|
||||
## 4 Related Work
|
||||
Reasoning models ( OpenAI, 2024 ; Guo et al., 2025 ) have saturated quantitative reasoning benchmarks like AIME and HMMT within one year.
|
||||
This rapid advancement is partly attributed to the well-defined evaluation criterion: if we care only about final answers, then quantitative reasoning is easy to verify.
|
||||
However, this final answer metric is inapplicable to theorem proving, which often requires no numerical answers but demands rigorous step-by-step derivation.
|
||||
Informal mathematical proofs have long been considered hard to verify automatically, lacking reliable approaches to assess proof correctness.
|
||||
Recent developments suggest this barrier may be surmountable.
|
||||
Models like Gemini-2.5 Pro already demonstrate a certain level of self-verification capabilities, which can refine their own solutions to improve quality ( Huang and Yang, 2025 ) .
|
||||
More significantly, DeepMind’s internal DeepThink variant ( Luong and Lockhart, 2025 ) achieved gold medal performance at IMO 2025 using pure natural language reasoning – serving as an existence proof that LLM-based verification of complex proofs is achievable.
|
||||
Recent research has begun exploring whether reasoning models can evaluate proofs, both with and without reference solutions ( Dekoninck et al., 2025 ; Luong et al., 2025 ) , showing promising results.
|
||||
In this work, we open source DeepSeekMath-V2 and our training methodology as concrete steps toward self-verifiable mathematical reasoning, showing how models can learn to verify and improve their own proofs.
|
||||
Proof assistants like Lean ( de Moura et al., 2015 ) and Isabelle ( Paulson, 1994 ) offer a reliable approach to verify proofs – proofs must be written in formal language, but once compiled, correctness is guaranteed.
|
||||
AlphaProof ( AlphaProof and teams, 2024 ; Trinh et al., 2024 ; Chervonyi et al., 2025 ) , a system specialized for formal proof search, achieved silver-level performance at IMO 2024 but required intensive computation.
|
||||
While using informal reasoning to guide formal proof generation has been explored extensively ( Jiang et al., 2023 ) , recent reasoning models have dramatically improved informal reasoning quality, making this guidance far more effective.
|
||||
Systems like DeepSeek-Prover-V2 ( Ren et al., 2025 ) and Seed-Prover ( Chen et al., 2025 ) can now produce substantially more valid formal proofs within the same computational budget, with Seed-Prover solving 5 of 6 problems at IMO 2025.
|
||||
Notably, these results were achieved without specifically optimizing the informal reasoning components for theorem proving tasks.
|
||||
We believe advancing natural language theorem proving will significantly benefit formal reasoning.
|
||||
We hope to contribute toward truly reliable mathematical reasoning systems that leverage both informal insights and formal guarantees to advance mathematical research.
|
||||
## 5 Conclusion
|
||||
We presented DeepSeekMath-V2, a model capable of both generating and verifying mathematical proofs.
|
||||
By training models to identify issues in their own reasoning and incentivizing them to address these issues before finalizing outputs, we move beyond the limitations of final-answer-based rewards toward self-verifiable mathematical reasoning.
|
||||
Our iterative training process – alternating between improving verification capabilities and using these to enhance generation – creates a sustainable cycle where each component drives the other forward.
|
||||
Our key technical contributions include: (1) training an accurate and faithful LLM-based verifier for mathematical proofs, (2) using meta-verification to largely reduce hallucinated issues and ensure verification quality, (3) incentivizing the proof generator to maximize proof quality through self-verification, and (4) scaling verification compute to automatically label increasingly hard-to-verify proofs to improve the verifier without human annotation.
|
||||
DeepSeekMath-V2 demonstrates strong performance on competition mathematics.
|
||||
With scaled test-time compute, it achieved gold-medal scores in high-school competitions including IMO 2025 and CMO 2024, and a near-perfect score on the undergraduate Putnam 2024 competition.
|
||||
This work establishes that LLMs can develop meaningful self-evaluation abilities for complex reasoning tasks.
|
||||
While significant challenges remain, we hope this research direction contributes to the goal of creating self-verifiable AI systems that can solve research-level mathematics.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,742 @@
|
||||
TITLE: DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
|
||||
ARXIV: 2512.02556
|
||||
AUTHORS: DeepSeek-AI , Aixin Liu , Aoxue Mei , Bangcai Lin , Bing Xue , Bingxuan Wang , Bingzheng Xu , Bochao Wu , Bowei Zhang , Chaofan Lin , Chen Dong , Chengda Lu , Chenggang Zhao , Chengqi Deng , Chenhao Xu , Chong Ruan , Damai Dai , Daya Guo , Dejian Yang , Deli Chen , Erhang Li , Fangqi Zhou , Fangyun Lin , Fucong Dai , Guangbo Hao , Guanting Chen , Guowei Li , H. Zhang , Hanwei Xu , Hao Li , Haofen Liang , Haoran Wei , Haowei Zhang , Haowen Luo , Haozhe Ji , Honghui Ding , Hongxuan Tang , Huanqi Cao , Huazuo Gao , Hui Qu , Hui Zeng , Jialiang Huang , Jiashi Li , Jiaxin Xu , Jiewen Hu , Jingchang Chen , Jingting Xiang , Jingyang Yuan , Jingyuan Cheng , Jinhua Zhu , Jun Ran , Junguang Jiang , Junjie Qiu , Junlong Li , Junxiao Song , Kai Dong , Kaige Gao , Kang Guan , Kexin Huang , Kexing Zhou , Kezhao Huang , Kuai Yu , Lean Wang , Lecong Zhang , Lei Wang , Liang Zhao , Liangsheng Yin , Lihua Guo , Lingxiao Luo , Linwang Ma , Litong Wang , Liyue Zhang , M.S. Di , M.Y Xu , Mingchuan Zhang , Minghua Zhang , Minghui Tang , Mingxu Zhou , Panpan Huang , Peixin Cong , Peiyi Wang , Qiancheng Wang , Qihao Zhu , Qingyang Li , Qinyu Chen , Qiushi Du , Ruiling Xu , Ruiqi Ge , Ruisong Zhang , Ruizhe Pan , Runji Wang , Runqiu Yin , Runxin Xu , Ruomeng Shen , Ruoyu Zhang , S.H. Liu , Shanghao Lu , Shangyan Zhou , Shanhuang Chen , Shaofei Cai , Shaoyuan Chen , Shengding Hu , Shengyu Liu , Shiqiang Hu , Shirong Ma , Shiyu Wang , Shuiping Yu , Shunfeng Zhou , Shuting Pan , Songyang Zhou , Tao Ni ,
|
||||
DATE: [Submitted on 2 Dec 2025]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-V3.2: Pushing the Frontier of Open
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 DeepSeek-V3.2 Architecture
|
||||
## 2.1 DeepSeek Sparse Attention
|
||||
## Prototype of DSA.
|
||||
## Instantiate DSA Under MLA.
|
||||
## 2.1.1 Continued Pre-Training
|
||||
## Dense Warm-up Stage.
|
||||
## Sparse Training Stage.
|
||||
## 2.2 Parity Evaluation
|
||||
## Standard Benchmark
|
||||
## Human Preference
|
||||
## Long Context Eval
|
||||
## 2.3 Inference Costs
|
||||
## 3 Post-Training
|
||||
## Specialist Distillation
|
||||
## Mixed RL Training
|
||||
## DeepSeek-V3.2 and DeepSeek-V3.2-Speciale
|
||||
## 3.1 Scaling GRPO
|
||||
## Unbiased KL Estimate
|
||||
## Off-Policy Sequence Masking
|
||||
## Keep Routing
|
||||
## Keep Sampling Mask
|
||||
## 3.2 Thinking in Tool-Use
|
||||
## 3.2.1 Thinking Context Management
|
||||
## 3.2.2 Cold-Start
|
||||
## 3.2.3 Large-Scale Agentic Tasks
|
||||
## Search Agent
|
||||
## Code Agent
|
||||
## Code Interpreter Agent
|
||||
## General Agent
|
||||
## 4 Evaluation
|
||||
## 4.1 Main Results
|
||||
## 4.2 Results of DeepSeek-V3.2-Speciale
|
||||
## 4.3 Synthesis Agentic Tasks
|
||||
## 4.4 Context Management of Search Agent
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 DeepSeek-V3.2 Architecture
|
||||
2.1 DeepSeek Sparse Attention
|
||||
Prototype of DSA.
|
||||
Instantiate DSA Under MLA.
|
||||
2.1.1 Continued Pre-Training
|
||||
Dense Warm-up Stage.
|
||||
Sparse Training Stage.
|
||||
2.2 Parity Evaluation
|
||||
Standard Benchmark
|
||||
Human Preference
|
||||
Long Context Eval
|
||||
2.3 Inference Costs
|
||||
3 Post-Training
|
||||
Specialist Distillation
|
||||
Mixed RL Training
|
||||
DeepSeek-V3.2 and DeepSeek-V3.2-Speciale
|
||||
3.1 Scaling GRPO
|
||||
Unbiased KL Estimate
|
||||
Off-Policy Sequence Masking
|
||||
Keep Routing
|
||||
Keep Sampling Mask
|
||||
3.2 Thinking in Tool-Use
|
||||
3.2.1 Thinking Context Management
|
||||
3.2.2 Cold-Start
|
||||
3.2.3 Large-Scale Agentic Tasks
|
||||
Search Agent
|
||||
Code Agent
|
||||
Code Interpreter Agent
|
||||
General Agent
|
||||
4 Evaluation
|
||||
4.1 Main Results
|
||||
4.2 Results of DeepSeek-V3.2-Speciale
|
||||
4.3 Synthesis Agentic Tasks
|
||||
4.4 Context Management of Search Agent
|
||||
5 Conclusion, Limitation, and Future Work
|
||||
References
|
||||
A MHA and MQA Modes of MLA
|
||||
B Cold Start Template
|
||||
C Non-thinking DeepSeek-V3.2 Agentic Evaluation
|
||||
D Evaluation Method of IOI, ICPC World Final, IMO, and CMO
|
||||
E Author List
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2512.02556v1 [cs.CL] 02 Dec 2025
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-V3.2: Pushing the Frontier of Open
|
||||
Large Language Models
|
||||
DeepSeek-AI
|
||||
research@deepseek.com
|
||||
## Abstract
|
||||
We introduce DeepSeek-V3.2, a model that harmonizes high computational efficiency with superior reasoning and agent performance. The key technical breakthroughs of DeepSeek-V3.2 are as follows: (1) DeepSeek Sparse Attention (DSA) :
|
||||
We introduce DSA, an efficient attention mechanism that substantially reduces computational complexity while preserving model performance in long-context scenarios.
|
||||
(2) Scalable Reinforcement Learning Framework : By implementing a robust reinforcement learning protocol and scaling post-training compute, DeepSeek-V3.2 performs comparably to GPT-5. Notably, our high-compute variant, DeepSeek-V3.2-Speciale, surpasses GPT-5 and exhibits reasoning proficiency on par with Gemini-3.0-Pro, achieving gold-medal performance in both the 2025 International Mathematical Olympiad (IMO) and the International Olympiad in Informatics (IOI).
|
||||
(3) Large-Scale Agentic Task Synthesis Pipeline : To integrate reasoning into tool-use scenarios, we developed a novel synthesis pipeline that systematically generates training data at scale. This methodology facilitates scalable agentic post-training, yielding substantial improvements in generalization and instruction-following robustness within complex, interactive environments.
|
||||
Figure 1:
|
||||
Benchmark of DeepSeek-V3.2 and its counterparts. For HMMT 2025, we report the February competition, consistent with the baselines. For HLE, we report the text-only subset.
|
||||
## 1 Introduction
|
||||
The release of reasoning models ( OpenAI, 2024b ; DeepSeek-AI, 2025 ) marked a pivotal moment in the evolution of Large Language Models (LLMs), catalyzing a substantial leap in overall performance across the verifiable fields. Since this milestone, the capabilities of LLMs have advanced rapidly. However, a distinct divergence has emerged in the past months. While the open-source community ( Qwen, 2025 ; ZhiPu-AI, 2025 ; MiniMax, 2025 ; MoonShot, 2025 ) continues to make strides, the performance trajectory of closed-source proprietary models ( OpenAI, 2025 ; Anthropic, 2025a ; DeepMind, 2025b ) has accelerated at a significantly steeper rate. Consequently, rather than converging, the performance gap between closed-source and open-source models appears to be widening, with proprietary systems demonstrating increasingly superior capabilities in complex tasks.
|
||||
Through our analysis, we identify three critical deficiencies that limit the capability of open-source models in complex tasks. First, architecturally, the predominant reliance on vanilla attention ( Vaswani et al., 2017 ) mechanisms severely constrains efficiency for long sequences. This inefficiency poses a substantial obstacle to both scalable deployment and effective post-training. Second, regarding resource allocation, open-source models suffer from insufficient computational investment during the post-training phase, limiting their performance on hard tasks. Finally, in the context of AI agents, open-source models demonstrate a marked lag in generalization and instruction-following capabilities compared to their proprietary counterparts ( EvalSys, 2025 ; Luo et al., 2025 ; Li et al., 2025 ) , hindering their effectiveness in real deployment.
|
||||
To address these critical limitations, we first introduce DSA, a highly efficient attention mechanism designed to substantially reduce computational complexity. This architecture effectively addresses the efficiency bottleneck, preserving model performance even in long-context scenarios. Second, we develop a stable and scalable RL protocol that allows for significant computational expansion during the post-training phase. Notably, this framework allocates a post-training computational budget exceeding 10% of the pre-training cost, unlocking advanced capabilities. Thirdly, we propose a novel pipeline to foster generalizable reasoning in tool-use scenarios. First, we implement a cold-start phase utilizing the DeepSeek-V3 ( DeepSeek-AI, 2024b ) methodology to unify reasoning and tool-use within single trajectories. Subsequently, we advance to large-scale agentic task synthesis, where we generate over 1,800 distinct environments and 85,000 complex prompts. This extensive synthesized data drives the RL process, significantly enhancing the model’s generalization and instruction-following capability in the agent context.
|
||||
DeepSeek-V3.2 achieves similar performance with Kimi-k2-thinking and GPT-5 across multiple reasoning benchmarks. Furthermore, DeepSeek-V3.2 significantly advances the agentic capabilities of open models, demonstrating exceptional proficiency on the long-tail agent tasks introduced in EvalSys (2025) ; Luo et al. (2025) ; Li et al. (2025) . DeepSeek-V3.2 emerges as a highly cost-efficient alternative in agent scenarios, significantly narrowing the performance gap between open and frontier proprietary models while incurring substantially lower costs.
|
||||
Notably, with the aim of pushing the boundaries of open models in the reasoning domain, we relaxed the length constraints to develop DeepSeek-V3.2-Speciale. As a result, DeepSeek-V3.2-Speciale achieves performance parity with the leading closed-source system, Gemini-3.0-Pro ( DeepMind, 2025a ) . It shows gold-medal performance in the IOI 2025, ICPC World Final 2025, IMO 2025, and CMO 2025.
|
||||
## 2 DeepSeek-V3.2 Architecture
|
||||
## 2.1 DeepSeek Sparse Attention
|
||||
DeepSeek-V3.2 uses exactly the same architecture as DeepSeek-V3.2-Exp.
|
||||
Compared with DeepSeek-V3.1-Terminus, the last version of DeepSeek-V3.1, the only architectural modification of DeepSeek-V3.2 is the introduction of DeepSeek Sparse Attention (DSA) through continued training.
|
||||
## Prototype of DSA.
|
||||
The prototype of DSA primarily consists of two components: a lightning indexer and a fine-grained token selection mechanism.
|
||||
The lightning indexer computes the index score I t , s I_{t,s} between the query token 𝐡 t ∈ ℝ d \mathbf{h}_{t}\in\mathbb{R}^{d} and a preceding token 𝐡 s ∈ ℝ d \mathbf{h}_{s}\in\mathbb{R}^{d} , determining which tokens to be selected by the query token:
|
||||
I t , s = ∑ j = 1 H I w t , j I ⋅ ReLU ( 𝐪 t , j I ⋅ 𝐤 s I ) , I_{t,s}=\sum_{j=1}^{H^{I}}w_{t,j}^{I}\cdot\text{ReLU}\left(\mathbf{q}^{I}_{t,j}\cdot\mathbf{k}^{I}_{s}\right),
|
||||
(1)
|
||||
where H I H^{I} denotes the number of indexer heads; 𝐪 t , j I ∈ ℝ d I \mathbf{q}^{I}_{t,j}\in\mathbb{R}^{d^{I}} and w t , j I ∈ ℝ w_{t,j}^{I}\in\mathbb{R} are derived from the query token 𝐡 t \mathbf{h}_{t} ;
|
||||
and 𝐤 s I ∈ ℝ d I \mathbf{k}^{I}_{s}\in\mathbb{R}^{d^{I}} is derived from the preceding token 𝐡 s \mathbf{h}_{s} .
|
||||
We choose ReLU as the activation function for throughput consideration.
|
||||
Given that the lightning indexer has a small number of heads and can be implemented in FP8, its computational efficiency is remarkable.
|
||||
Given the index scores { I t , s } \{I_{t,s}\} for each query token 𝐡 t \mathbf{h}_{t} , our fine-grained token selection mechanism retrieves only the key-value entries { 𝐜 s } \{\mathbf{c}_{s}\} corresponding to the top-k index scores.
|
||||
Then, the attention output 𝐮 t \mathbf{u}_{t} is computed by applying the attention mechanism between the query token 𝐡 t \mathbf{h}_{t} and the sparsely selected key-value entries { 𝐜 s } \{\mathbf{c}_{s}\} :
|
||||
𝐮 t = Attn ( 𝐡 t , { 𝐜 s | I t , s ∈ Top-k ( I t , : ) } ) . \mathbf{u}_{t}=\text{Attn}\quantity(\mathbf{h}_t, \qty{ \mathbf{c}_s \, \middle| \, I_{t, s} \in\text{Top-k} \qty(I_{t, :}) } ).
|
||||
(2)
|
||||
Figure 2:
|
||||
Attention architecture of DeepSeek-V3.2, where DSA is instantiated under MLA.
|
||||
The green part illustrates how DSA selects the top-k key-value entries according to the indexer.
|
||||
## Instantiate DSA Under MLA.
|
||||
For the consideration of continued training from DeepSeek-V3.1-Terminus, we instantiate DSA based on MLA ( DeepSeek-AI, 2024a ) for DeepSeek-V3.2.
|
||||
At the kernel level, each key-value entry must be shared across multiple queries for computational efficiency ( Yuan et al., 2025 ) .
|
||||
Therefore, we implement DSA based on the MQA ( Shazeer, 2019 ) mode of MLA 1 1
|
||||
1
|
||||
We illustrate the difference between the MQA and MHA modes of MLA in Appendix A . , where each latent vector (the key-value entry of MLA) will be shared across all query heads of the query token.
|
||||
The DSA architecture based on MLA is illustrated in Figure 2 .
|
||||
We also provide an open-source implementation of DeepSeek-V3.2 2 2
|
||||
2
|
||||
https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp/tree/main/inference to specify the details unambiguously.
|
||||
## 2.1.1 Continued Pre-Training
|
||||
Starting from a base checkpoint of DeepSeek-V3.1-Terminus, whose context length has been extended to 128K, we perform continued pre-training followed by post-training to create DeepSeek-V3.2.
|
||||
The continued pre-training of DeepSeek-V3.2 consists of two training stages.
|
||||
For both stages, the distribution of training data is totally aligned with the 128K long context extension data used for DeepSeek-V3.1-Terminus.
|
||||
## Dense Warm-up Stage.
|
||||
We first use a short warm-up stage to initialize the lightning indexer.
|
||||
In this stage, we keep dense attention and freeze all model parameters except for the lightning indexer.
|
||||
To align the indexer outputs with the main attention distribution, for the t t -th query token, we first aggregate the main attention scores by summing across all attention heads.
|
||||
This sum is then L1-normalized along the sequence dimension to produce a target distribution p t , : ∈ ℝ t p_{t,:}\in\mathbb{R}^{t} .
|
||||
Based on p t , : p_{t,:} , we set a KL-divergence loss as the training objective of the indexer:
|
||||
ℒ I = ∑ t 𝔻 KL ( p t , : ∥ Softmax ( I t , : ) ) . \mathcal{L}^{I}=\sum_{t}\mathbb{D}_{\mathrm{KL}}\left(p_{t,:}\,\middle\|\,\text{Softmax}\quantity({I}_{t,:})\right).
|
||||
(3)
|
||||
For warm-up, we use a learning rate of 10 − 3 10^{-3} .
|
||||
We train the indexer for only 1000 steps, with each step consisting of 16 sequences of 128K tokens, resulting in a total of 2.1B tokens.
|
||||
## Sparse Training Stage.
|
||||
Following indexer warm-up, we introduce the fine-grained token selection mechanism and optimize all model parameters to adapt the model to the sparse pattern of DSA.
|
||||
In this stage, we also keep aligning the indexer outputs to the main attention distribution, but considering only the selected token set 𝒮 t = { s | I t , s ∈ Top-k ( I t , : ) } \mathcal{S}_{t}=\quantity{s \, \middle| \, I_{t,s} \in\text{Top-k} \qty(I_{t,:})} :
|
||||
ℒ I = ∑ t 𝔻 KL ( p t , 𝒮 t ∥ Softmax ( I t , 𝒮 t ) ) . \mathcal{L}^{I}=\sum_{t}\mathbb{D}_{\mathrm{KL}}\left(p_{t,\mathcal{S}_{t}}\,\middle\|\,\text{Softmax}\quantity(I_{t,\mathcal{S}_t})\right).
|
||||
(4)
|
||||
It is worth noting that we detach the indexer input from the computational graph for separate optimization.
|
||||
The training signal of the indexer is from only ℒ I \mathcal{L}^{I} , while the optimization of the main model is according to only the language modeling loss.
|
||||
In this sparse training stage, we use a learning rate of 7.3 × 10 − 6 7.3\times 10^{-6} , and select 2048 key-value tokens for each query token.
|
||||
We train both the main model and the indexer for 15000 15000 steps, with each step consisting of 480 sequences of 128K tokens, resulting in a total of 943.7B tokens.
|
||||
## 2.2 Parity Evaluation
|
||||
## Standard Benchmark
|
||||
In September 2025, we evaluate DeepSeek-V3.2-Exp on a suite of benchmarks, which focus on diverse capabilities, and compare it with DeepSeek-V3.1-Terminus showing similar performance.
|
||||
While DeepSeek V3.2 Exp significantly improves computational efficiency on long sequences, we do not observe substantial performance degradation compared with DeepSeek-V3.1-Terminus, on both short- and long-context tasks.
|
||||
## Human Preference
|
||||
Given that direct human preference assessments are inherently susceptible to bias, we employ ChatbotArena as an indirect evaluation framework to approximate user preferences for the newly developed base models. Both DeepSeek - V3.1 - Terminus and DeepSeek - V3.2 - Exp share an identical post - training strategy, and their Elo scores, obtained from evaluations conducted on 10 November 2025, are closely matched. These results suggest that the new base model achieves performance on par with the previous iteration, despite incorporating a sparse attention mechanism.
|
||||
## Long Context Eval
|
||||
Following the release of DeepSeek - V3.2 - Exp, several independent long - context evaluations were conducted using previously unseen test sets. A representative benchmark is AA - LCR 3 3
|
||||
3
|
||||
https://artificialanalysis.ai/evaluations/artificial-analysis-long-context-reasoning , in which DeepSeek - V3.2 - Exp scores four points higher than DeepSeek-V3.1-Terminus in reasoning mode. In the Fiction.liveBench evaluation 4 4
|
||||
4
|
||||
https://fiction.live/stories/Fiction-liveBench-April-6-2025/oQdzQvKHw8JyXbN87 , DeepSeek - V3.2 - Exp consistently outperforms DeepSeek-V3.1-Terminus across multiple metrics. This evidence indicates the base checkpoint of DeepSeek - V3.2 - Exp does not regress on long context tasks.
|
||||
## 2.3 Inference Costs
|
||||
DSA reduces the core attention complexity of the main model from 𝒪 ( L 2 ) \order{L^2} to 𝒪 ( L k ) \order{L k} , where k k ( ≪ L \ll L ) is the number of selected tokens.
|
||||
Although the lightning indexer still has a complexity of 𝒪 ( L 2 ) \order{L^2} , it requires much less computation compared with MLA in DeepSeek-V3.1-Terminus.
|
||||
Combined with our optimized implementation, DSA achieves a significant end-to-end speedup in long-context scenarios.
|
||||
Figure 3 presents how token costs of DeepSeek-V3.1-Terminus and DeepSeek-V3.2 vary with the token position in the sequence.
|
||||
These costs are estimated from benchmarking the actual service deployed on H800 GPUs, at a rental price of 2 USD per GPU hour.
|
||||
Note that for short-sequence prefilling, we specially implement a masked MHA mode to simulate DSA, which can achieve higher efficiency under short-context conditions.
|
||||
(a) Prefilling
|
||||
(b) Decoding
|
||||
Figure 3:
|
||||
Inference costs of DeepSeek-V3.1-Terminus and DeepSeek-V3.2 on H800 clusters.
|
||||
## 3 Post-Training
|
||||
After continued pre-training, we perform post-training to create the final DeepSeek-V3.2.
|
||||
The post-training of DeepSeek-V3.2 also employs sparse attention in the same way as the sparse continued pre-training stage.
|
||||
For DeepSeek-V3.2, we maintain the same post-training pipeline as in DeepSeek-V3.2-Exp, which includes specialist distillation and mixed RL training.
|
||||
## Specialist Distillation
|
||||
For each task, we initially develop a specialized model dedicated exclusively to that particular domain, with all specialist models being fine-tuned from the same pre-trained DeepSeek-V3.2 base checkpoint.
|
||||
In addition to writing tasks and general question-answering, our framework encompasses six specialized domains: mathematics, programming, general logical reasoning, general agentic tasks, agentic coding, and agentic search, with all the domains supporting both thinking and non-thinking modes.
|
||||
Each specialist is trained with large-scale Reinforcement Learning (RL) computing.
|
||||
Furthermore, we employ different models to generate training data for long chain-of-thought reasoning (thinking mode) and direct response generation (non-thinking mode).
|
||||
Once the specialist models are prepared, they are used to produce the domain-specific data for the final checkpoint.
|
||||
Experimental results demonstrate that models trained on the distilled data achieve performance levels only marginally below those of domain-specific specialists, with the performance gap being effectively eliminated through subsequent RL training.
|
||||
## Mixed RL Training
|
||||
For DeepSeek-V3.2, we still adopt Group Relative Policy Optimization (GRPO) ( Shao et al., 2024 ; DeepSeek-AI, 2025 ) as the RL training algorithm.
|
||||
As DeepSeek-V3.2-Exp, we merge reasoning, agent, and human alignment training into one RL stage.
|
||||
This approach effectively balances performance across diverse domains while circumventing the catastrophic forgetting issues commonly associated with multi-stage training paradigms.
|
||||
For reasoning and agent tasks, we employ rule-based outcome reward, length penalty, and language consistency reward.
|
||||
For general tasks, we employ a generative reward model where each prompt has its own rubrics for evaluation.
|
||||
## DeepSeek-V3.2 and DeepSeek-V3.2-Speciale
|
||||
DeepSeek-V3.2 integrates reasoning, agent, and human alignment data distilled from specialists, undergoing thousands of steps of continued RL training to reach the final checkpoints.
|
||||
To investigate the potential of extended thinking, we also developed an experimental variant, DeepSeek-V3.2-Speciale. This model was trained exclusively on reasoning data with a reduced length penalty during RL. Additionally, we incorporated the dataset and reward method from DeepSeekMath-V2 ( Shao et al., 2025 ) to enhance capabilities in mathematical proofs.
|
||||
We would like to highlight our efforts in how to create a stable recipe to scale up RL compute in Section 3.1 , and how to integrate thinking into agentic tasks in Section 3.2
|
||||
## 3.1 Scaling GRPO
|
||||
We first review the objective of GRPO. GRPO optimizes the policy model π θ \pi_{\theta} by maximizing the following objective on a group of responses { o 1 , ⋯ , o G } \{o_{1},\cdots,o_{G}\} sampled from the old policy π old \pi_{\mathrm{old}} given each question q q :
|
||||
𝒥 GRPO ( θ ) = \displaystyle\mathcal{J}_{\mathrm{GRPO}}(\theta)=\kern 5.0pt
|
||||
𝔼 q ∼ P ( Q ) , { o i } i = 1 G ∼ π old ( ⋅ | q ) [ 1 G ∑ i = 1 G 1 | o i | ∑ t = 1 | o i | \displaystyle\mathbb{E}_{q\sim P(Q),\{o_{i}\}_{i=1}^{G}\sim\pi_{\mathrm{old}}(\cdot|q)}\Bigg[\frac{1}{G}\sum_{i=1}^{G}\frac{1}{|o_{i}|}\sum_{t=1}^{|o_{i}|}
|
||||
min ( r i , t ( θ ) A ^ i , t , clip ( r i , t ( θ ) , 1 − ε , 1 + ε ) A ^ i , t ) − β 𝔻 KL ( π θ ( o i , t ) ∥ π ref ( o i , t ) ) ] , \displaystyle\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\text{clip}\left(r_{i,t}(\theta),1-\varepsilon,1+\varepsilon\right)\hat{A}_{i,t}\right)-\beta\mathbb{D}_{\mathrm{KL}}\left(\pi_{\theta}(o_{i,t})\,\middle\|\,\pi_{\mathrm{ref}}(o_{i,t})\right)\Bigg],
|
||||
(5)
|
||||
where
|
||||
r i , t ( θ ) = π θ ( o i , t | q , o i , < t ) π old ( o i , t | q , o i , < t ) r_{i,t}(\theta)=\frac{\pi_{\theta}(o_{i,t}|q,o_{i,<t})}{\pi_{\mathrm{old}}(o_{i,t}|q,o_{i,<t})}
|
||||
(6)
|
||||
is the importance sampling ratio between the current and old policy. ε \varepsilon and β \beta are hyper-parameters controlling the clipping range and KL penalty strength, respectively. A ^ i , t \hat{A}_{i,t} is the advantage of o i , t o_{i,t} which is estimated by normalizing the outcome reward within each group. Specifically, a set of reward models are used to score an outcome reward R i R_{i} for each output o i o_{i} in the group, yielding G G rewards 𝑹 = { R 1 , ⋯ , R G } \boldsymbol{R}=\{R_{1},\cdots,R_{G}\} respectively. The advantage of o i , t o_{i,t} is calculated by subtracting the average reward of the group from the reward of output o i o_{i} , i.e., A ^ i , t = R i − mean ( 𝑹 ) \hat{A}_{i,t}=R_{i}-\text{mean}(\boldsymbol{R}) .
|
||||
In the following, we outline additional strategies that stabilize RL scaling, directly building on the GRPO algorithm.
|
||||
## Unbiased KL Estimate
|
||||
Given o i , t o_{i,t} is sampled from the old policy π old ( ⋅ | q , o i , < t ) \pi_{\mathrm{old}}(\cdot|q,o_{i,<t}) , we correct the K3 estimator ( Schulman, 2020 ) to obtain an unbiased KL estimate
|
||||
using the importance-sampling ratio between the current policy π θ \pi_{\theta} and the old policy π old \pi_{\mathrm{old}} .
|
||||
𝔻 KL ( π θ ( o i , t ) ∥ π ref ( o i , t ) ) = π θ ( o i , t | q , o i , < t ) π old ( o i , t | q , o i , < t ) ( π ref ( o i , t | q , o i , < t ) π θ ( o i , t | q , o i , < t ) − log π ref ( o i , t | q , o i , < t ) π θ ( o i , t | q , o i , < t ) − 1 ) . \mathbb{D}_{\mathrm{KL}}\left(\pi_{\theta}(o_{i,t})\,\middle\|\,\pi_{\mathrm{ref}}(o_{i,t})\right)=\frac{\pi_{\theta}(o_{i,t}|q,o_{i,<t})}{\pi_{\mathrm{old}}(o_{i,t}|q,o_{i,<t})}\left(\frac{\pi_{\mathrm{ref}}(o_{i,t}|q,o_{i,<t})}{\pi_{\theta}(o_{i,t}|q,o_{i,<t})}-\log\frac{\pi_{\mathrm{ref}}(o_{i,t}|q,o_{i,<t})}{\pi_{\theta}(o_{i,t}|q,o_{i,<t})}-1\right).
|
||||
(7)
|
||||
As a direct result of this adjustment, the gradient of this KL estimator becomes unbiased, which eliminates systematic estimation errors, thereby facilitating stable convergence. This contrasts sharply with the original K3 estimator, particularly when the sampled tokens have substantially lower probabilities under the current policy than the reference policy, i.e., π θ ≪ π ref \pi_{\theta}\ll\pi_{\mathrm{ref}} . In such cases, the gradient of the K3 estimator assigns disproportionately large, unbounded weights to maximize the likelihood of these tokens, resulting in noisy gradient updates that accumulate to degrade sample quality in subsequent iterations and lead to unstable training dynamics. In practice, we find that different domains benefit from varying strengths of KL regularization. For certain domains, such as mathematics, applying a relatively weak KL penalty or even omitting it entirely can yield improved performance.
|
||||
## Off-Policy Sequence Masking
|
||||
To improve the efficiency of RL systems, we typically generate a large batch of rollout data, which is subsequently split into multiple mini-batches for several gradient update steps. This practice inherently introduces off-policy behavior. Additionally, inference frameworks used for efficient data generation are often highly optimized, which may differ in implementation details from training frameworks. Such training-inference inconsistency further exacerbates the degree of off-policyness. To stabilize training and improve tolerance for off-policy updates, we mask negative sequences that introduce significant policy divergence, as measured by the KL divergence between the data-sampling policy π old \pi_{\mathrm{old}} and the current policy π θ \pi_{\theta} . More specifically, we introduce a binary mask M M into the GRPO loss:
|
||||
𝒥 GRPO ( θ ) = \displaystyle\mathcal{J}_{\mathrm{GRPO}}(\theta)=\kern 5.0pt
|
||||
𝔼 q ∼ P ( Q ) , { o i } i = 1 G ∼ π old ( ⋅ | q ) [ 1 G ∑ i = 1 G 1 | o i | ∑ t = 1 | o i | \displaystyle\mathbb{E}_{q\sim P(Q),\{o_{i}\}_{i=1}^{G}\sim\pi_{\mathrm{old}}(\cdot|q)}\Bigg[\frac{1}{G}\sum_{i=1}^{G}\frac{1}{|o_{i}|}\sum_{t=1}^{|o_{i}|}
|
||||
min ( r i , t ( θ ) A ^ i , t , clip ( r i , t ( θ ) , 1 − ε , 1 + ε ) A ^ i , t ) M i , t − β 𝔻 KL ( π θ ( o i , t ) ∥ π ref ( o i , t ) ) ] , \displaystyle\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\text{clip}\left(r_{i,t}(\theta),1-\varepsilon,1+\varepsilon\right)\hat{A}_{i,t}\right)M_{i,t}-\beta\mathbb{D}_{\mathrm{KL}}\left(\pi_{\theta}(o_{i,t})\,\middle\|\,\pi_{\mathrm{ref}}(o_{i,t})\right)\Bigg],
|
||||
(8)
|
||||
where
|
||||
M i , t = { 0 A ^ i , t < 0 , 1 | o i | ∑ t = 1 | o i | log π old ( o i , t | q , o i , < t ) π θ ( o i , t | q , o i , < t ) > δ 1 otherwise , M_{i,t}=\begin{cases}0&{\hat{A}_{i,t}<0,\frac{1}{|o_{i}|}\sum_{t=1}^{|o_{i}|}\log\frac{\pi_{\mathrm{old}}(o_{i,t}|q,o_{i,<t})}{\pi_{\theta}(o_{i,t}|q,o_{i,<t})}>\delta}\\[4.30554pt]
|
||||
1&{\text{otherwise},}\end{cases}
|
||||
(9)
|
||||
and δ \delta is a hyper-parameter that controls the threshold of policy divergence. Note that π old \pi_{\mathrm{old}} here denotes the sampling probability directly returned by the inference framework, thus the KL divergence between the old and current policy accounts for both sources of off-policyness mentioned above. It is also worth noting that we only mask sequences with negative advantages.
|
||||
Intuitively, the model benefits the most by learning from its own mistakes, whereas highly off-policy negative samples can be detrimental, potentially misleading or destabilizing the optimization process. We empirically observe that this Off-Policy Sequence Masking operation improves stability in certain training scenarios that would otherwise exhibit instability.
|
||||
## Keep Routing
|
||||
Mixture-of-Experts (MoE) models improve computational efficiency by activating only a subset of expert modules during inference. However, discrepancies between inference and training frameworks, compounded by policy updates, can result in inconsistent expert routing during inference and training even for identical inputs. Such inconsistency induces abrupt shifts in the active parameter subspace, which destabilizes optimization and exacerbates off-policy issues. To mitigate this, we preserve the expert routing paths used during sampling in the inference framework and enforce the same routing paths during training, ensuring that identical expert parameters are optimized. This Keep Routing operation was found crucial for RL training stability of MoE models, and has been adopted in our RL training pipeline since DeepSeek-V3-0324.
|
||||
## Keep Sampling Mask
|
||||
Top-p and top-k sampling are widely used sampling strategies to enhance the quality of responses generated by LLMs. Employing these strategies in RL training is also advantageous, as it avoids sampling extremely low-probability tokens that would be used as optimization targets. While such truncation preserves sample quality, it introduces a mismatch between the action spaces of π old \pi_{\mathrm{old}} and π θ \pi_{\theta} , which violates the principles of importance sampling and destabilizes training. To address this, we preserve the truncation masks during sampling from π old \pi_{\mathrm{old}} and apply them to π θ \pi_{\theta} during training, ensuring both policies share identical action subspaces. Empirically, we find that combining top-p sampling with the Keep Sampling Mask strategy effectively preserves language consistency during RL training.
|
||||
## 3.2 Thinking in Tool-Use
|
||||
## 3.2.1 Thinking Context Management
|
||||
DeepSeek-R1 has demonstrated that incorporating a thinking process can significantly enhance a model’s ability to solve complex problems. Building on this insight, we aim to integrate thinking capabilities into tool-calling scenarios.
|
||||
We observed that replicating DeepSeek-R1’s strategy—discarding reasoning content upon the arrival of the second round of messages—results in significant token inefficiency. This approach forces the model to redundantly re-reason through the entire problem for each subsequent tool call. To mitigate this, we developed a context management strictly tailored for tool-calling scenarios as shown in Fig 4 :
|
||||
•
|
||||
Historical reasoning content is discarded only when a new user message is introduced to the conversation. If only tool-related messages (e.g., tool outputs) are appended, the reasoning content is retained throughout the interaction.
|
||||
•
|
||||
When reasoning traces are removed, the history of tool calls and their results remains preserved in the context.
|
||||
Notably, certain agent frameworks, such as Roo Code or Terminus, simulate tool interactions via user messages. These frameworks may not fully benefit from our enhanced reasoning persistence due to the context management rules outlined above. Therefore, we recommend utilizing non-thinking models for optimal performance with such architectures.
|
||||
Figure 4:
|
||||
Thinking retention mechanism in tool-calling scenarios.
|
||||
## 3.2.2 Cold-Start
|
||||
Given the availability of reasoning data (non-agentic) and non-reasoning agentic data, a straightforward strategy for integrating these two capabilities is through carefully designed prompting. We posit that the model possesses sufficient ability to accurately follow explicit instructions, thereby enabling the seamless incorporation of tool execution within the reasoning process.
|
||||
To demonstrate the operation of the cold-start mechanism, we selectively sample the training data as shown in Appendix Tables 8 – 8 . It is important to note that distinct task prompts are associated with different system prompts. Tables 8 – 8 present an illustrative example corresponding to a competitive programming prompt.
|
||||
Table 8 presents an example of our reasoning data, which uses a system prompt to explicitly asks the model to do reasoning before the final answer and uses a special tag <think></think> to label the reasoning path. Table 8 shows the prompt of non-reasoning agentic data, where the system prompt contains the guidance of toolcall. Table 8 presents the system prompt we designed to instruct the model to incorporate multiple tool calls within its reasoning process.
|
||||
In this manner, although the reasoning in tool - use patterns may lack robustness, the model is occasionally able to generate the desired trajectories, thereby providing a basis for subsequent reinforcement learning stages.
|
||||
## 3.2.3 Large-Scale Agentic Tasks
|
||||
A diverse set of RL tasks is crucial for enhancing model robustness. For tasks such as search, code engineering, and code interpretation, we employ real-world tools, including actual web search APIs, coding tools, and Jupyter Notebooks. While these RL environments are real, the prompts employed are either extracted from Internet sources or synthetically generated, rather than obtained from actual user interactions.
|
||||
For other tasks, the environment and prompts are both synthetically constructed. The agent tasks we used are described in Table 1 .
|
||||
Table 1: The description of different agent tasks, including the number of tasks, environment type (real or synthesized), and prompt source (extracted or synthesized).
|
||||
number of tasks
|
||||
environment
|
||||
prompt
|
||||
code agent
|
||||
24667
|
||||
real
|
||||
extracted
|
||||
search agent
|
||||
50275
|
||||
real
|
||||
synthesized
|
||||
general agent
|
||||
4417
|
||||
synthesized
|
||||
synthesized
|
||||
code interpreter
|
||||
5908
|
||||
real
|
||||
extracted
|
||||
## Search Agent
|
||||
We employ a multi-agent pipeline based on DeepSeek-V3.2 to generate diverse, high-quality training data.
|
||||
We first sample informative long-tail entities across diverse domains from large-scale web corpora.
|
||||
A question-construction agent then explores each entity using search tools with configurable depth and breadth parameters, consolidating the discovered information into question-answer pairs.
|
||||
Multiple answer-generation agents with heterogeneous configurations (different checkpoints, system prompts, etc.) produce diverse candidate responses for each proposed QA pair.
|
||||
A verification agent with search capabilities validates all answers through multiple passes, retaining only samples where the ground-truth is correct and all candidates are verifiably incorrect.
|
||||
These data spans multiple languages, domains, and difficulty levels.
|
||||
To complement these verifiable samples and better reflect real-world usage, we also augment the dataset with filtered instances from our existing helpful RL datasets, for which the search tool provides measurable benefits.
|
||||
We then develop detailed evaluation rubrics across multiple quality dimensions and employ a generative reward model to score responses based on these rubrics.
|
||||
This hybrid approach enables optimization for both factual reliability and practical helpfulness.
|
||||
## Code Agent
|
||||
We constructed large-scale, executable environments for software issue resolution by mining millions of issue-Pull Request (PR) pairs from GitHub. This dataset was rigorously filtered using heuristic rules and LLM-based judgments to ensure high quality, requiring that each entry contain a reasonable issue description, a correlated gold patch, and a test patch for validation. An automated environment-setup agent, powered by DeepSeek-V3.2, was employed to build executable environments for these pairs. This agent handles package installation, dependency resolution, and test execution. Test results are output in the standard JUnit format, ensuring consistent parsing across programming languages and test frameworks. An environment is deemed successfully built only when applying the gold patch results in a non-zero count of false-to-positive (F2P) test cases (indicating the issue is fixed) and a zero count of pass-to-fail (P2F) test cases (indicating no regressions). Using this pipeline, we successfully built tens of thousands of reproducible issue resolution environments spanning multiple programming languages, including Python, Java, JavaScript, TypeScript, C, C++, Go, and PHP.
|
||||
## Code Interpreter Agent
|
||||
We utilize Jupyter Notebook as a code interpreter to address complex reasoning tasks. To facilitate this, we curate a diverse set of problems spanning mathematics, logic, and data science, each requiring the model to leverage code execution capabilities to arrive at a solution.
|
||||
## General Agent
|
||||
To scale up agent environments and tasks in RL, we employ an automatic environment-synthesis agent that synthesizes 1,827 task-oriented environments. These tasks are hard to solve but easy to verify. The synthesis workflow primarily consists of environment and toolset construction, task synthesis, and solution generation. Specifically, the workflow proceeds as follows.
|
||||
1.
|
||||
Given a task category (e.g., planning a travel itinerary) and a sandbox equipped with a bash and a search tool, the agent first uses these tools to generate or retrieve relevant data from the Internet and store them in the sandbox database.
|
||||
2.
|
||||
The agent then synthesizes a set of task-specific tools, each implemented as a function.
|
||||
3.
|
||||
To create tasks that are both challenging and automatically verifiable, the agent initially proposes a simple task based on the current database, along with its solution and verification functions implemented in Python. The solution function is restricted to invoking tool functions or performing logical computations, and cannot call other functions or directly access the database, ensuring the task can only be solved through the tool interface. Additionally, the results produced by the solution function must be validated by the verification function. If the solution is not validated, the agent will modify the solution or verification functions until the solution’s output passes the verification. The agent then iteratively increases the difficulty of the task and updates the corresponding solution and verification functions. During this iterative process, if the current toolset is not sufficient to solve the task, the agent will augment the toolset.
|
||||
Following this workflow, we obtain thousands of ⟨ environment , tools , task , verifier ⟩ \textlangle\text{environment},\text{tools},\text{task},\text{verifier}\textrangle tuples. We then perform RL on this dataset using DeepSeek-V3.2 and retain only instances with non-zero pass@100, resulting in 1,827 environments and their corresponding tasks (4,417 in total). A synthetic trip-planning example is illustrated below. This example highlights that, while searching the large combinatorial space for a trip plan that satisfies all constraints is challenging, checking whether a given candidate solution satisfies these constraints is relatively straightforward.
|
||||
An Example of Synthesized Task: Trip Planning
|
||||
I’m planning a three-day trip starting from Hangzhou, and I need help creating an itinerary from October 1st to October 3rd, 2025.
|
||||
A few important requirements: I don’t want to repeat any cities, hotels, attractions, or restaurants during the entire trip. Also, please make sure that every hotel, restaurant, and attraction you recommend is actually located in the city where I’ll be staying that day.
|
||||
One more thing about the second day - I’m trying to be smart about my budget. If I end up booking a luxury hotel that costs 800 CNY or more per night, then I need to be more careful with other expenses: my total spending on both restaurants (lunch and dinner) should stay under 350 CNY, both restaurants should be rated at least 4.0 stars, and the afternoon attraction ticket needs to be less than 120 CNY.
|
||||
If the hotel on day 2 is in the mid-to-high range (500-800 CNY), then I have a bit more flexibility - I just need to make sure at least one of my restaurant choices is rated 4.0 or higher, and the attraction ticket should be below 180 CNY.
|
||||
For more affordable hotels (200-500 CNY range), I only need to ensure that at least one restaurant has a rating of 3.2 or above.
|
||||
Can you help me put together this itinerary?
|
||||
Submit Result Format
|
||||
[
|
||||
{
|
||||
"time": "2025-10-01",
|
||||
"city": "cite_name",
|
||||
"hotel": "hotel_name",
|
||||
"afternoon_restaurant": "restaurant_name",
|
||||
"afternoon_attraction": "attraction_name",
|
||||
"evening_restaurant": "restaurant_name"
|
||||
},
|
||||
{
|
||||
"time": "2025-10-02",
|
||||
"city": "cite_name",
|
||||
"hotel": "hotel_name",
|
||||
"afternoon_restaurant": "restaurant_name",
|
||||
"afternoon_attraction": "attraction_name",
|
||||
"evening_restaurant": "restaurant_name"
|
||||
},
|
||||
{
|
||||
"time": "2025-10-03",
|
||||
"city": "cite_name",
|
||||
"hotel": "hotel_name",
|
||||
"afternoon_restaurant": "restaurant_name",
|
||||
"afternoon_attraction": "attraction_name",
|
||||
"evening_restaurant": "restaurant_name"
|
||||
}
|
||||
]
|
||||
Tool Set for Trip Planning
|
||||
Function Name
|
||||
Description
|
||||
get_all_attractions_by_city(city)
|
||||
Get all attractions for given city.
|
||||
get_all_cities()
|
||||
Get all cities from the database.
|
||||
get_all_hotels_by_city(city)
|
||||
Get all hotels for given city.
|
||||
get_all_restaurants_by_city(city)
|
||||
Get all restaurants for given city.
|
||||
get_city_by_attraction(attraction)
|
||||
Get city for given attraction name.
|
||||
get_city_by_hotel(hotel)
|
||||
Get city for given hotel name.
|
||||
get_city_by_restaurant(restaurant)
|
||||
Get city for given restaurant name.
|
||||
get_city_transport(city)
|
||||
Get all intra-city transport options for given city.
|
||||
get_infos_by_attraction(info_keywords, attraction)
|
||||
Get specified infos for given attraction.
|
||||
get_infos_by_city(info_keywords, city)
|
||||
Get specified infos for given city.
|
||||
get_infos_by_hotel(info_keywords, hotel)
|
||||
Get specified infos for given hotel.
|
||||
get_infos_by_restaurant(info_keywords, restaurant)
|
||||
Get specified infos for given restaurant.
|
||||
get_inter_city_transport(from_city, to_city)
|
||||
Get all transports between given city pair.
|
||||
get_weather_by_city_date(city, date)
|
||||
Get weather for given city-date pair.
|
||||
submit_result(answer_text)
|
||||
Submit the final answer content.
|
||||
## 4 Evaluation
|
||||
## 4.1 Main Results
|
||||
We evaluate models on MMLU-Pro ( Wang et al., 2024 ) , GPQA Diamond ( Rein et al., 2023 ) , Human Last Exam (HLE) Text-only ( Phan et al., 2025 ) , LiveCodeBench (2024.08-2025.04), Codeforces, Aider-Polyglot, AIME 2025, HMMT Feb 2025, HMMT Nov 2025 ( Balunović et al., 2025 ) , IMOAnswerBench ( Luong et al., 2025 ) , Terminal Bench 2.0, SWE-Verified ( OpenAI, 2024a ) , SWE Multilingual ( Yang et al., 2025 ) , BrowseComp ( Wei et al., 2025 ) , BrowseCompZh ( Zhou et al., 2025 ) , τ 2 \tau^{2} -bench ( Barres et al., 2025 ) , MCP-Universe ( Luo et al., 2025 ) , MCP-Mark ( EvalSys, 2025 ) , and Tool-Decathlon ( Li et al., 2025 ) . Tool-use benchmarks are evaluated using the standard function call format, wherein models are configured to thinking mode.
|
||||
For MCP-Universe ( Luo et al., 2025 ) and MCP-Mark ( EvalSys, 2025 ) , we evaluate all models with our internal environment, because the search and playwright environment might be slightly different from the official setting.
|
||||
We set the temperature to 1.0, and the context window to 128K tokens.
|
||||
For math-related tasks such as AIME, HMMT, IMOAnswerBench, and HLE, we eval with the following template: "{question}\nPlease reason step by step, and put your final answer within \boxed{}."
|
||||
In the case of HLE, we additionally assessed DeepSeek-V3.2-Thinking using the official template, resulting in a score of 23.9 23.9 .
|
||||
Table 2: Comparison between DeepSeek-V3.2 and closed/open models. For open models, we just compare with models supports thinking in tooluse.
|
||||
Numbers in bold represent the best scores within each model class (open-source and closed-source). The τ 2 \tau^{2} -Bench result is computed by the average of each category. Regarding BrowseComp, the performance with the context management technique is noted with *.
|
||||
Benchmark (Metric)
|
||||
Claude-4.5-
|
||||
GPT-5
|
||||
Gemini-3.0
|
||||
Kimi-K2
|
||||
MiniMax
|
||||
DeepSeek-V3.2
|
||||
Sonnet
|
||||
High
|
||||
Pro
|
||||
Thinking
|
||||
M2
|
||||
Thinking
|
||||
English
|
||||
MMLU-Pro (EM)
|
||||
88.2
|
||||
87.5
|
||||
90.1
|
||||
84.6
|
||||
82.0
|
||||
85.0
|
||||
GPQA Diamond (Pass@1)
|
||||
83.4
|
||||
85.7
|
||||
91.9
|
||||
84.5
|
||||
77.7
|
||||
82.4
|
||||
HLE (Pass@1)
|
||||
13.7
|
||||
26.3
|
||||
37.7
|
||||
23.9
|
||||
12.5
|
||||
25.1
|
||||
Code
|
||||
LiveCodeBench (Pass@1-COT)
|
||||
64.0
|
||||
84.5
|
||||
90.7
|
||||
82.6
|
||||
83.0
|
||||
83.3
|
||||
Codeforces (Rating)
|
||||
1480
|
||||
2537
|
||||
2708
|
||||
-
|
||||
-
|
||||
2386
|
||||
Math
|
||||
AIME 2025 (Pass@1)
|
||||
87.0
|
||||
94.6
|
||||
95.0
|
||||
94.5
|
||||
78.3
|
||||
93.1
|
||||
HMMT Feb 2025 (Pass@1)
|
||||
79.2
|
||||
88.3
|
||||
97.5
|
||||
89.4
|
||||
-
|
||||
92.5
|
||||
HMMT Nov 2025 (Pass@1)
|
||||
81.7
|
||||
89.2
|
||||
93.3
|
||||
89.2
|
||||
-
|
||||
90.2
|
||||
IMOAnswerBench (Pass@1)
|
||||
-
|
||||
76.0
|
||||
83.3
|
||||
78.6
|
||||
-
|
||||
78.3
|
||||
Code Agent
|
||||
Terminal Bench 2.0 (Acc)
|
||||
42.8
|
||||
35.2
|
||||
54.2
|
||||
35.7
|
||||
30.0
|
||||
46.4
|
||||
SWE Verified (Resolved)
|
||||
77.2
|
||||
74.9
|
||||
76.2
|
||||
71.3
|
||||
69.4
|
||||
73.1
|
||||
SWE Multilingual (Resolved)
|
||||
68.0
|
||||
55.3
|
||||
-
|
||||
61.1
|
||||
56.5
|
||||
70.2
|
||||
Search Agent
|
||||
BrowseComp (Pass@1)
|
||||
24.1
|
||||
54.9
|
||||
-
|
||||
-/60.2*
|
||||
44.0
|
||||
51.4/67.6 *
|
||||
BrowseCompZh (Pass@1)
|
||||
42.4
|
||||
63.0
|
||||
-
|
||||
62.3
|
||||
48.5
|
||||
65.0
|
||||
HLE (Pass@1)
|
||||
32.0
|
||||
35.2
|
||||
45.8
|
||||
44.9
|
||||
31.8
|
||||
40.8
|
||||
ToolUse
|
||||
τ 2 \tau^{2} -Bench (Pass@1)
|
||||
84.7
|
||||
80.2
|
||||
85.4
|
||||
74.3
|
||||
76.9
|
||||
80.3
|
||||
MCP-Universe (Success Rate)
|
||||
46.5
|
||||
47.9
|
||||
50.7
|
||||
35.6
|
||||
29.4
|
||||
45.9
|
||||
MCP-Mark (Pass@1)
|
||||
33.3
|
||||
50.9
|
||||
43.1
|
||||
20.4
|
||||
24.4
|
||||
38.0
|
||||
Tool-Decathlon (Pass@1)
|
||||
38.6
|
||||
29.0
|
||||
36.4
|
||||
17.6
|
||||
16.0
|
||||
35.2
|
||||
DeepSeek-V3.2 achieves similar performance with GPT-5-high on reasoning tasks, but is slightly worse than Gemini-3.0-Pro. Compared to K2-Thinking, DeepSeek-V3.2 achieves comparable scores with substantially fewer output tokens, as shown in Table 3 . These performance gains can be attributed to the increased computational resources allocated to RL training. Over recent months, we have observed consistent performance improvements correlating with extended RL training budget, which already exceeds 10% of the pre-training cost. We hypothesize that reasoning capabilities could be further enhanced with additional computational budget allocation. Notably, the performance of DeepSeek-V3.2 presented herein is constrained by a length constraint reward model; upon removal of the restriction, we observe further improvement in model performance, as detailed in Section 4.2 .
|
||||
In code agent evaluations, DeepSeek-V3.2 significantly outperforms open-source LLMs on both SWE-bench Verified and Terminal Bench 2.0, demonstrating its potential within real-world coding workflows. Regarding Terminal Bench 2.0, as previously noted, our context management strategy for the ’thinking mode’ is currently incompatible with Terminus; consequently, the reported score of 46.4 was achieved using the Claude Code framework. We also evaluated DeepSeek-V3.2 with Terminus in non-thinking mode, yielding a score of 39.3. For SWE-bench Verified, the primary score was obtained using our internal framework. Robustness tests across other settings—including the Claude Code and RooCode frameworks, as well as non-thinking mode—produced consistent results, ranging from 72 to 74.
|
||||
For the search agent evaluation, we assess our models using a standard commercial search API. Since DeepSeek-V3.2 supports a maximum context length of only 128K, approximately 20%+ of the test cases exceed this limit. To address this, we employ a context management method to derive the final score. For reference, the score is 51.4 without context management. Further details are provided in Section 4.4 .
|
||||
On tool-use benchmarks, DeepSeek-V3.2 substantially narrows the performance gap between open-source and closed-source LLMs, though it remains below frontier models. For τ 2 \tau^{2} -bench, we employ the model itself as the user agent, achieving final category scores of 63.8 (Airline), 81.1 (Retail), and 96.2 (Telecom). For the MCP benchmarks, we employ the function calling format and place tool outputs within messages designated with the ’tool’ role, rather than the ’user’ role.
|
||||
During our testing, we observed that DeepSeek-V3.2 frequently engages in redundant self-verification, generating excessively long trajectories. This tendency often causes the context length to exceed the 128K limit, particularly in tasks such as MCP-Mark GitHub and Playwright evaluation. Consequently, this phenomenon hinders the final performance of DeepSeek-V3.2. However, integrating context management strategies can further enhance performance. We identify this as a direction for future work and a practical consideration for users. Even if DeepSeek-V3.2 suffers from the issue, it still significantly outperforms existing open models.
|
||||
Notably, since the environments and toolsets employed in these benchmarks were not encountered during RL training, the observed improvements demonstrate DeepSeek-V3.2’s capacity to generalize its reasoning strategies to out-of-domain agentic scenarios. The evaluation of non-thinking model in the agent scenario is shown in Appendix Table 9 .
|
||||
## 4.2 Results of DeepSeek-V3.2-Speciale
|
||||
Table 3:
|
||||
Benchmark performance and efficiency of reasoning models. For each benchmark, cells show accuracy and output token count (in thousands). The highest accuracy per benchmark is in bold; the second-highest is underlined.
|
||||
Benchmark
|
||||
GPT-5
|
||||
Gemini-3.0
|
||||
Kimi-K2
|
||||
DeepSeek-V3.2
|
||||
DeepSeek-V3.2
|
||||
High
|
||||
Pro
|
||||
Thinking
|
||||
Thinking
|
||||
Speciale
|
||||
AIME 2025 (Pass@1)
|
||||
94.6 (13k)
|
||||
95.0 (15k)
|
||||
94.5 (24k)
|
||||
93.1 (16k)
|
||||
96.0 (23k)
|
||||
HMMT Feb 2025 (Pass@1)
|
||||
88.3 (16k)
|
||||
97.5 (16k)
|
||||
89.4 (31k)
|
||||
92.5 (19k)
|
||||
99.2 (27k)
|
||||
HMMT Nov 2025 (Pass@1)
|
||||
89.2 (20k)
|
||||
93.3 (15k)
|
||||
89.2 (29k)
|
||||
90.2 (18k)
|
||||
94.4 (25k)
|
||||
IMOAnswerBench (Pass@1)
|
||||
76.0 (31k)
|
||||
83.3 (18k)
|
||||
78.6 (37k)
|
||||
78.3 (27k)
|
||||
84.5 (45k)
|
||||
LiveCodeBench (Pass@1-COT)
|
||||
84.5 (13k)
|
||||
90.7 (13k)
|
||||
82.6 (29k)
|
||||
83.3 (16k)
|
||||
88.7 (27k)
|
||||
CodeForces (Rating)
|
||||
2537 (29k)
|
||||
2708 (22k)
|
||||
-
|
||||
2386 (42k)
|
||||
2701 (77k)
|
||||
GPQA Diamond (Pass@1)
|
||||
85.7 (8k)
|
||||
91.9 (8k)
|
||||
84.5 (12k)
|
||||
82.4 (7k)
|
||||
85.7 (16k)
|
||||
HLE (Pass@1)
|
||||
26.3 (15k)
|
||||
37.7 (15k)
|
||||
23.9 (24k)
|
||||
25.1 (21k)
|
||||
30.6 (35k)
|
||||
Table 3 demonstrates that DeepSeek-V3.2-Speciale achieves superior performance by leveraging increased reasoning tokens, surpassing the state-of-the-art Gemini-3.0-Pro across multiple benchmarks.
|
||||
Remarkably, as shown in Table 4 , this general-purpose model attains gold-medal level performance in the 2025 International Olympiad in Informatics (IOI) and the ICPC World Finals (ICPC WF) without targeted training.
|
||||
Furthermore, by incorporating techniques from Shao et al. (2025) , the model excels in complex proof tasks, reaching gold-medal thresholds in the 2025 International Mathematical Olympiad (IMO) and China Mathematical Olympiad (CMO) 5 5
|
||||
5
|
||||
We evaluated the English version of CMO 2025. The IMO 2025 and CMO 2025 problems, together with the inference code, can be found at: https://github.com/deepseek-ai/DeepSeek-Math-V2 . . Detailed evaluation protocols are provided in Appendix D .
|
||||
However, the token efficiency of DeepSeek-V3.2-Speciale remains significantly inferior to that of Gemini-3.0-Pro.
|
||||
To mitigate deployment costs and latency, we imposed stricter token constraints during the training of the official DeepSeek-V3.2, aiming to optimize the trade-off between performance and cost.
|
||||
We believe that token efficiency remains a critical area for future investigation.
|
||||
Table 4: Performance of DeepSeek-V3.2-Speciale in top-tier mathematics and coding competitions.
|
||||
For ICPC WF 2025, we report the number of submissions for each successfully solved problem. DeepSeek-V3.2-Speciale ranked 2nd in ICPC WF 2025 and 10th in IOI 2025.
|
||||
Competition
|
||||
P1
|
||||
P2
|
||||
P3
|
||||
P4
|
||||
P5
|
||||
P6
|
||||
Overall
|
||||
Medal
|
||||
IMO 2025
|
||||
7
|
||||
7
|
||||
7
|
||||
7
|
||||
7
|
||||
0
|
||||
35/42
|
||||
Gold
|
||||
CMO 2025
|
||||
18
|
||||
18
|
||||
9
|
||||
21
|
||||
18
|
||||
18
|
||||
102/126
|
||||
Gold
|
||||
IOI 2025
|
||||
100
|
||||
82
|
||||
72
|
||||
100
|
||||
55
|
||||
83
|
||||
492/600
|
||||
Gold
|
||||
Competition
|
||||
A
|
||||
B
|
||||
C
|
||||
D
|
||||
E
|
||||
F
|
||||
G
|
||||
H
|
||||
I
|
||||
J
|
||||
K
|
||||
L
|
||||
Overall
|
||||
Medal
|
||||
ICPC WF 2025
|
||||
3
|
||||
-
|
||||
1
|
||||
1
|
||||
2
|
||||
2
|
||||
-
|
||||
1
|
||||
1
|
||||
1
|
||||
1
|
||||
1
|
||||
10/12
|
||||
Gold
|
||||
## 4.3 Synthesis Agentic Tasks
|
||||
In this section, we perform ablation experiments to study the effect of synthetic agentic tasks. We focus on two questions. First, are synthetic tasks sufficiently challenging for reinforcement learning? Second, how well do these synthetic tasks generalize, i.e., can they transfer to different downstream tasks or real-world environments?
|
||||
To address the first question, we randomly sample 50 instances from the general synthesized agentic tasks and evaluate both the model used for synthesis and frontier closed-source LLMs. As shown in Table 5 , DeepSeek-V3.2-Exp attains an accuracy of only 12%, while frontier closed-source models achieve at most 62%. These results indicate that the synthetic data include agentic tasks that are challenging for both DeepSeek-V3.2-Exp and frontier closed-source models.
|
||||
Table 5: Accuracy of general synthesized tasks on different models.
|
||||
Pass@K
|
||||
DeepSeek-v3.2-Exp
|
||||
Sonnet-4.5
|
||||
Gemini-3.0 Pro
|
||||
GPT-5-Thinking
|
||||
1
|
||||
12%
|
||||
34%
|
||||
51%
|
||||
62%
|
||||
2
|
||||
18%
|
||||
47%
|
||||
65%
|
||||
75%
|
||||
4
|
||||
26%
|
||||
62%
|
||||
74%
|
||||
82%
|
||||
To investigate whether RL on synthetic data can generalize to different tasks or real-world environments, we apply RL to the SFT checkpoint of DeepSeek-V3.2 (denoted DeepSeek-V3.2-SFT). To exclude the effects of long CoT and other RL data, we conduct RL only on synthetic agentic tasks in non-thinking mode. We then compare the model with DeepSeek-V3.2-SFT and DeepSeek-V3.2-Exp, where DeepSeek-V3.2-Exp is trained with RL only in search and code environments. As shown in Figure 5 , large-scale RL on synthetic data yields substantial improvements over DeepSeek-V3.2-SFT on Tau2Bench, MCP-Mark, and MCP-Universe benchmarks. In contrast, restricting RL to code and search scenarios does not improve performance on these benchmarks, further highlighting the potential of synthetic data.
|
||||
Figure 5:
|
||||
RL training of DeepSeek-V3.2-SFT using exclusively synthetic general agent data.
|
||||
## 4.4 Context Management of Search Agent
|
||||
Figure 6:
|
||||
Accuracy of Browsecomp with different test-time compute expansion strategies.
|
||||
Even with extended context windows such as 128k, agentic workflows, particularly in search-based scenarios, frequently encounter maximum length limitations that prematurely truncate the reasoning process.
|
||||
This bottleneck inhibits the full realization of test-time compute potential.
|
||||
To address this, we introduce context management employing simple strategies to extend token budgets at test time， when the token usage exceeds 80% of the context window length.
|
||||
These strategies include
|
||||
(1) Summary , which summarizes the overflowed trajectory and re-initiates the rollout;
|
||||
(2) Discard-75% , which discards the first 75% tool call history in the trajectory to free up spaces;
|
||||
(3) Discard-all , which resets the context by discarding all previous tool call history (similar to the new context tool ( Anthropic, 2025b ) ).
|
||||
For comparison, we also implement a parallel scaling baseline, Parallel-fewest-step , which samples N independent trajectories and selects the trajectory with the fewest steps.
|
||||
We evaluate these strategies on the BrowseComp benchmark ( Wei et al., 2025 ) .
|
||||
As illustrated in Figure 6 , under varying compute budgets, context management leads to significant performance gains by allowing the model to scale up test-time compute, providing more space to perform additional execution steps.
|
||||
For example, Summary extends the average steps to 364, achieving a performance improvement of up to 60.2. However, its overall efficiency is relatively low.
|
||||
Despite its simplicity, Discard-all performs well in both efficiency and scalability, achieving a score of 67.6, comparable to parallel scaling while using significantly fewer steps.
|
||||
In summary, test-time compute can be scaled either serially through context management or in parallel, both effectively extending the model’s problem-solving capacity.
|
||||
However, different strategies exhibit varying efficiency and scalability.
|
||||
Thus, it is crucial to account for actual compute costs when benchmarking model performance.
|
||||
Meanwhile, finding the optimal combination of serial and parallel scaling to maximize both efficiency and scalability remains a crucial direction for future work.
|
||||
## 5 Conclusion, Limitation, and Future Work
|
||||
In this work, we introduced DeepSeek-V3.2, a framework that effectively bridges the gap between computational efficiency and advanced reasoning capabilities. Using DSA, we addressed critical computation complexity without sacrificing long-context performance. By increasing computational budget, DeepSeek-V3.2 achieves comparable performance with GPT-5 on reasoning benchmarks. Finally, the integration of our large-scale agentic task synthesis pipeline significantly enhances tool-use proficiency, unlocking new possibilities for robust and generalizable AI agents with open LLM. Furthermore, our high-compute variant, DeepSeek-V3.2-Speciale, validated by gold-medal achievements in the IMO and IOI, sets a milestone for open LLMs.
|
||||
Despite these achievements, we acknowledge certain limitations when compared to frontier closed-source models such as Gemini-3.0-Pro. First, due to fewer total training FLOPs, the breadth of world knowledge in DeepSeek-V3.2 still lags behind that of leading proprietary models. We plan to address this knowledge gap in future iterations by scaling up the pre-training compute. Second, token efficiency remains a challenge; DeepSeek-V3.2 typically requires longer generation trajectories (i.e., more tokens) to match the output quality of models like Gemini-3.0-Pro. Future work will focus on optimizing the intelligence density of the model’s reasoning chains to improve efficiency. Third, solving complex tasks is still inferior to frontier models, motivating us to further refine our foundation model and post-training recipe.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,462 @@
|
||||
TITLE: mHC: Manifold-Constrained Hyper-Connections
|
||||
ARXIV: 2512.24880
|
||||
AUTHORS: Zhenda Xie , Yixuan Wei , Huanqi Cao , Chenggang Zhao , Chengqi Deng , Jiashi Li , Damai Dai , Huazuo Gao , Jiang Chang , Kuai Yu , Liang Zhao , Shangyan Zhou , Zhean Xu , Zhengyan Zhang , Wangding Zeng , Shengding Hu , Yuqing Wang , Jingyang Yuan , Lean Wang , Wenfeng Liang
|
||||
DATE: [Submitted on 31 Dec 2025 ( v1 ), last revised 5 Jan 2026 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Related Works
|
||||
## 2.1 Micro Design
|
||||
## 2.2 Macro Design
|
||||
## 3 Preliminary
|
||||
## 3.1 Numerical Instability
|
||||
## 3.2 System Overhead
|
||||
## 4 Method
|
||||
## 4.1 Manifold-Constrained Hyper-Connections
|
||||
## 4.2 Parameterization and Manifold Projection
|
||||
## 4.3 Efficient Infrastructure Design
|
||||
## 4.3.1 Kernel Fusion
|
||||
## 4.3.2 Recomputing
|
||||
## 4.3.3 Overlapping Communication in DualPipe
|
||||
## 5 Experiments
|
||||
## 5.1 Experimental Setup
|
||||
## 5.2 Main Results
|
||||
## 5.3 Scaling Experiments
|
||||
## 5.4 Stability Analysis
|
||||
## 6 Conclusion and Outlook
|
||||
|
||||
=== BODY (head) ===
|
||||
mHC: Manifold-Constrained Hyper-Connections
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Micro Design
|
||||
2.2 Macro Design
|
||||
3 Preliminary
|
||||
3.1 Numerical Instability
|
||||
3.2 System Overhead
|
||||
4 Method
|
||||
4.1 Manifold-Constrained Hyper-Connections
|
||||
4.2 Parameterization and Manifold Projection
|
||||
4.3 Efficient Infrastructure Design
|
||||
4.3.1 Kernel Fusion
|
||||
4.3.2 Recomputing
|
||||
4.3.3 Overlapping Communication in DualPipe
|
||||
5 Experiments
|
||||
5.1 Experimental Setup
|
||||
5.2 Main Results
|
||||
5.3 Scaling Experiments
|
||||
5.4 Stability Analysis
|
||||
6 Conclusion and Outlook
|
||||
References
|
||||
A Appendix
|
||||
A.1 Detailed Model Specifications and Hyper-parameters.
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2512.24880v2 [cs.CL] 05 Jan 2026
|
||||
\reportnumber
|
||||
001
|
||||
##
|
||||
m HC: Manifold-Constrained Hyper-Connections
|
||||
Zhenda Xie* † , Yixuan Wei*, Huanqi Cao*,
|
||||
Chenggang Zhao, Chengqi Deng, Jiashi Li, Damai Dai, Huazuo Gao, Jiang Chang,
|
||||
Kuai Yu, Liang Zhao, Shangyan Zhou, Zhean Xu, Zhengyan Zhang, Wangding Zeng,
|
||||
Shengding Hu, Yuqing Wang, Jingyang Yuan, Lean Wang, Wenfeng Liang
|
||||
DeepSeek-AI
|
||||
Corresponding author: Core contributors. † Corresponding author: xie.zhenda@deepseek.com
|
||||
## Abstract
|
||||
Recently, studies exemplified by Hyper-Connections (HC) have extended the ubiquitous residual connection paradigm established over the past decade by expanding the residual stream width and diversifying connectivity patterns.
|
||||
While yielding substantial performance gains, this diversification fundamentally compromises the identity mapping property intrinsic to the residual connection, which causes severe training instability and restricted scalability, and additionally incurs notable memory access overhead.
|
||||
To address these challenges, we propose Manifold-Constrained Hyper-Connections ( m HC ), a general framework that projects the residual connection space of HC onto a specific manifold to restore the identity mapping property, while incorporating rigorous infrastructure optimization to ensure efficiency.
|
||||
Empirical experiments demonstrate that m HC is effective for training at scale, offering tangible performance improvements and superior scalability.
|
||||
We anticipate that m HC, as a flexible and practical extension of HC, will contribute to a deeper understanding of topological architecture design and suggest promising directions for the evolution of foundational models.
|
||||
Figure 1 :
|
||||
Illustrations of Residual Connection Paradigms. This figure compares the structural design of (a) standard Residual Connection, (b) Hyper-Connections (HC), and (c) our proposed Manifold-Constrained Hyper-Connections ( m HC ). Unlike the unconstrained HC, m HC focuses on optimizing the residual connection space by projecting the matrices onto a constrained manifold to ensure stability.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Micro Design
|
||||
2.2 Macro Design
|
||||
3 Preliminary
|
||||
3.1 Numerical Instability
|
||||
3.2 System Overhead
|
||||
4 Method
|
||||
4.1 Manifold-Constrained Hyper-Connections
|
||||
4.2 Parameterization and Manifold Projection
|
||||
4.3 Efficient Infrastructure Design
|
||||
4.3.1 Kernel Fusion
|
||||
4.3.2 Recomputing
|
||||
4.3.3 Overlapping Communication in DualPipe
|
||||
5 Experiments
|
||||
5.1 Experimental Setup
|
||||
5.2 Main Results
|
||||
5.3 Scaling Experiments
|
||||
5.4 Stability Analysis
|
||||
6 Conclusion and Outlook
|
||||
References
|
||||
A Appendix
|
||||
A.1 Detailed Model Specifications and Hyper-parameters.
|
||||
## 1 Introduction
|
||||
Deep neural network architectures have undergone rapid evolution since the introduction of ResNets ( He et al., 2016a ) .
|
||||
As illustrated in Fig. 1 (a), the structure of a single-layer can be formulated as follows:
|
||||
𝐱 l + 1 = 𝐱 l + ℱ ( 𝐱 l , 𝒲 l ) , \mathbf{x}_{l+1}=\mathbf{x}_{l}+\mathcal{F}(\mathbf{x}_{l},\mathcal{W}_{l}),
|
||||
(1)
|
||||
where 𝐱 l \mathbf{x}_{l} and 𝐱 l + 1 \mathbf{x}_{l+1} denote the C C -dimensional input and output of the l l -th layer, respectively, and ℱ \mathcal{F} represents the residual function.
|
||||
Although the residual function ℱ \mathcal{F} has evolved over the past decade to include various operations such as convolution, attention mechanisms, and feed forward networks, the paradigm of the residual connection has maintained its original form. Accompanying the progression of Transformer ( Vaswani et al., 2017 ) architecture, this paradigm has currently established itself as a fundamental design element in large language models (LLMs) ( Brown et al., 2020 ; Touvron et al., 2023 ; Liu et al., 2024b ) .
|
||||
This success is primarily attributed to the concise form of the residual connection. More importantly, early research ( He et al., 2016b ) revealed that the identity mapping property of the residual connection maintains stability and efficiency during large-scale training. By recursively extending the residual connection across multiple layers, Eq. ( 1 ) yields:
|
||||
𝐱 L = 𝐱 l + ∑ i = l L − 1 ℱ ( 𝐱 i , 𝒲 i ) , \mathbf{x}_{L}=\mathbf{x}_{l}+\sum_{i=l}^{L-1}\mathcal{F}(\mathbf{x}_{i},\mathcal{W}_{i}),
|
||||
(2)
|
||||
where L L and l l correspond to deeper and shallower layers, respectively.
|
||||
The term identity mapping refers to the component 𝐱 l \mathbf{x}_{l} itself, which emphasizes the property that the signal from the shallower layer maps directly to the deeper layer without any modification.
|
||||
Recently, studies exemplified by Hyper-Connections (HC) ( Zhu et al., 2024 ) have introduced a new dimension to the residual connection and empirically demonstrated its performance potential. The single-layer architecture of HC is illustrated in Fig. 1 (b).
|
||||
By expanding the width of the residual stream and enhancing connection complexity, HC significantly increases topological complexity without altering the computational overhead of individual units regarding FLOPs.
|
||||
Formally, single-layer propagation in HC is defined as:
|
||||
𝐱 l + 1 = ℋ l res 𝐱 l + ℋ l post ⊤ ℱ ( ℋ l pre 𝐱 l , 𝒲 l ) , \mathbf{x}_{l+1}=\mathcal{H}_{l}^{\mathrm{res}}\mathbf{x}_{l}+\mathcal{H}_{l}^{\mathrm{post}\,\top}\mathcal{F}(\mathcal{H}_{l}^{\mathrm{pre}}\mathbf{x}_{l},\mathcal{W}_{l}),
|
||||
(3)
|
||||
where 𝐱 l \mathbf{x}_{l} and 𝐱 l + 1 \mathbf{x}_{l+1} denote the input and output of the l l -th layer, respectively.
|
||||
Unlike the formulation in Eq. ( 1 ), the feature dimension of 𝐱 l \mathbf{x}_{l} and 𝐱 l + 1 \mathbf{x}_{l+1} is expanded from C C to n × C n\times C , where n n is the expansion rate.
|
||||
The term ℋ l res ∈ ℝ n × n \mathcal{H}_{l}^{\mathrm{res}}\in\mathbb{R}^{n\times n} represents a learnable mapping that mixes features within the residual stream.
|
||||
Also as a learnable mapping, ℋ l pre ∈ ℝ 1 × n \mathcal{H}_{l}^{\mathrm{pre}}\in\mathbb{R}^{1\times n} aggregates features from the n C nC -dim stream into a C C -dim layer input, and conversely, ℋ l post ∈ ℝ 1 × n \mathcal{H}_{l}^{\mathrm{post}}\in\mathbb{R}^{1\times n} maps the layer output back onto the stream.
|
||||
However, as the training scale increases, HC introduces potential risks of instability.
|
||||
The primary concern is that the unconstrained nature of HC compromises the identity mapping property when the architecture extends across multiple layers.
|
||||
In architectures comprising multiple parallel streams, an ideal identity mapping serves as a conservation mechanism. It ensures that the average signal intensity across streams remains invariant during both forward and backward propagation.
|
||||
Recursively extending HC to multiple layers via Eq. ( 3 ) yields:
|
||||
𝐱 L = ( ∏ i = 1 L − l ℋ L − i res ) 𝐱 l + ∑ i = l L − 1 ( ∏ j = 1 L − 1 − i ℋ L − j res ) ℋ i post ⊤ ℱ ( ℋ i pre 𝐱 i , 𝒲 i ) , \mathbf{x}_{L}=\left(\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}\right)\mathbf{x}_{l}+\sum_{i=l}^{L-1}\left(\prod_{j=1}^{L-1-i}\mathcal{H}_{L-j}^{\mathrm{res}}\right)\mathcal{H}_{i}^{\mathrm{post}\,\top}\mathcal{F}(\mathcal{H}_{i}^{\mathrm{pre}}\mathbf{x}_{i},\mathcal{W}_{i}),
|
||||
(4)
|
||||
where L L and l l represent a deeper layer and a shallower layer, respectively.
|
||||
In contrast to Eq. ( 2 ), the composite mapping ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} in HC fails to preserve the global mean of the features. This discrepancy leads to unbounded signal amplification or attenuation, resulting in instability during large-scale training.
|
||||
A further consideration is that, while HC preserves computational efficiency in terms of FLOPs, the hardware efficiency concerning memory access costs for the widened residual stream remains unaddressed in the original design. These factors collectively restrict the practical scalability of HC and hinder its application in large-scale training.
|
||||
To address these challenges, we propose Manifold-Constrained Hyper-Connections ( m HC ), as shown in Fig. 1 (c), a general framework that projects the residual connection space of HC onto a specific manifold to restore the identity mapping property, while incorporating rigorous infrastructure optimization to ensure efficiency.
|
||||
Specifically, m HC utilizes the Sinkhorn-Knopp algorithm ( Sinkhorn and Knopp, 1967 ) to entropically project ℋ l res \mathcal{H}_{l}^{\mathrm{res}} onto the Birkhoff polytope. This operation effectively constrains the residual connection matrices within the manifold that is constituted by doubly stochastic matrices.
|
||||
Since the row and column sums of these matrices equal to 1 1 , the operation ℋ l res 𝐱 l \mathcal{H}_{l}^{\mathrm{res}}\mathbf{x}_{l} functions as a convex combination of the input features. This characteristic facilitates a well-conditioned signal propagation where the feature mean is conserved, and the signal norm is strictly regularized, effectively mitigating the risk of vanishing or exploding signals. Furthermore, due to the closure of matrix multiplication for doubly stochastic matrices, the composite mapping ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} retains this conservation property. Consequently, m HC effectively maintains the stability of identity mappings between arbitrary depths. To ensure efficiency, we employ kernel fusion and develop mixed precision kernels utilizing TileLang ( Wang et al., 2025 ) . Furthermore, we mitigate the memory footprint through selective recomputing and carefully overlap communication within the DualPipe schedule ( Liu et al., 2024b ) .
|
||||
Extensive experiments on language model pretraining demonstrate that m HC exhibits exceptional stability and scalability while maintaining the performance advantages of HC. In-house large-scale training indicates that m HC supports training at scale and introduces only a 6.7% additional time overhead when expansion rate n = 4 n=4 .
|
||||
## 2 Related Works
|
||||
Architectural advancements in deep learning can be primarily classified into micro-design and macro-design .
|
||||
Micro-design concerns the internal architecture of computational blocks, specifying how features are processed across spatial, temporal, and channel dimensions.
|
||||
In contrast, macro-design establishes the inter-block topological structure, thereby dictating how feature representations are propagated, routed, and merged across distinct layers.
|
||||
## 2.1 Micro Design
|
||||
Driven by parameter sharing and translation invariance, convolution initially dominated the processing of structured signals. While subsequent variations such as depthwise separable ( Chollet, 2017 ) and grouped convolutions ( Xie et al., 2017 ) optimized efficiency, the advent of Transformers ( Vaswani et al., 2017 ) established Attention and Feed-Forward Networks (FFNs) as the fundamental building blocks of modern architecture. Attention mechanisms facilitate global information propagation, while FFNs enhance the representational capacity of individual features. To balance performance with the computational demands of LLMs, attention mechanisms have evolved towards efficient variants such as Multi-Query Attention (MQA) ( Shazeer, 2019 ) , Grouped-Query Attention (GQA) ( Ainslie et al., 2023 ) , and Multi-Head Latent Attention (MLA) ( Liu et al., 2024a ) . Simultaneously, FFNs have been generalized into sparse computing paradigms via Mixture-of-Experts (MoE) ( Shazeer et al., 2017 ; Lepikhin et al., 2020 ; Fedus et al., 2022 ) , allowing for massive parameter scaling without proportional computational costs.
|
||||
## 2.2 Macro Design
|
||||
Macro-design governs the global topology of the network ( Srivastava et al., 2015 ) . Following ResNet ( He et al., 2016a ) , architectures such as DenseNet ( Huang et al., 2017 ) and FractalNet ( Larsson et al., 2016 ) aimed to enhance performance by increasing topological complexity through dense connectivity and multi-path structures, respectively. Deep Layer Aggregation (DLA) ( Yu et al., 2018 ) further extended this paradigm by recursively aggregating features across various depths and resolutions.
|
||||
More recently, the focus of macro-design has shifted toward expanding the width of the residual stream ( Chai et al., 2020 ; Fang et al., 2023 ; Xie et al., 2023 ; Pagliardini et al., 2024 ; Menghani et al., 2025 ; Heddes et al., 2025 ; Zhu et al., 2024 ; Mak and Flanigan, 2025 ; Xiao et al., 2025 ) . Hyper-Connections (HC) ( Zhu et al., 2024 ) introduced learnable matrices to modulate connection strengths among features at varying depths, while the Residual Matrix Transformer (RMT) ( Mak and Flanigan, 2025 ) replaced the standard residual stream with an outer-product memory matrix to facilitate feature storage. Similarly, MUDDFormer ( Xiao et al., 2025 ) employs multiway dynamic dense connections to optimize cross-layer information flow.
|
||||
Despite their potential, these approaches compromise the inherent identity mapping property of the residual connection, thereby introducing instability and hindering scalability. Furthermore, they incur significant memory access overhead due to expanded feature widths.
|
||||
Building upon HC, the proposed m HC restricts the residual connection space onto a specific manifold to restore the identity mapping property, while also incorporating rigorous infrastructure optimizations to ensure efficiency. This approach enhances stability and scalability while maintaining the topological benefits of expanded connections.
|
||||
## 3 Preliminary
|
||||
We first establish the notation used in this work. In the HC formulation, the input to the l l -th layer, x l ∈ ℝ 1 × C \textbf{x}_{l}\in\mathbb{R}^{1\times C} , is expanded by a factor of n n to construct a hidden matrix x l = ( x l , 0 ⊤ , … , x l , n − 1 ⊤ ) ⊤ ∈ ℝ n × C \textbf{x}_{l}=(\textbf{x}^{\top}_{l,0},\ldots,\textbf{x}^{\top}_{l,n-1})^{\top}\in\mathbb{R}^{n\times C} which can be viewed as n n -stream residual. This operation effectively broadens the width of the residual stream. To govern the read-out, write-in, and updating processes of this stream, HC introduces three learnable linear mappings— ℋ l pre , ℋ l post ∈ ℝ 1 × n \mathcal{H}^{\mathrm{pre}}_{l},\mathcal{H}^{\mathrm{post}}_{l}\in\mathbb{R}^{1\times n} , and ℋ l res ∈ ℝ n × n \mathcal{H}^{\mathrm{res}}_{l}\in\mathbb{R}^{n\times n} . These mappings modify the standard residual connection shown in Eq. ( 1 ), resulting in the formulation given in Eq. ( 3 ).
|
||||
In the HC formulation, learnable mappings are composed of two parts of coefficients: the input-dependent one and the global one, referred to as dynamic mappings and static mappings, respectively. Formally, HC computes the coefficients as follows:
|
||||
{ 𝐱 ~ l = RMSNorm ( 𝐱 l ) ℋ l pre = α l pre ⋅ tanh ( θ l pre 𝐱 ~ l ⊤ ) + 𝐛 l pre ℋ l post = α l post ⋅ tanh ( θ l post 𝐱 ~ l ⊤ ) + 𝐛 l post ℋ l res = α l res ⋅ tanh ( θ l res 𝐱 ~ l ⊤ ) + 𝐛 l res , \begin{cases}\tilde{\mathbf{x}}_{l}=\text{RMSNorm}(\mathbf{x}_{l})\\
|
||||
\mathcal{H}^{\mathrm{pre}}_{l}=\alpha_{l}^{\mathrm{pre}}\cdot\tanh(\theta^{\mathrm{pre}}_{l}\tilde{\mathbf{x}}^{\top}_{l})+\mathbf{b}_{l}^{\mathrm{pre}}\\
|
||||
\mathcal{H}^{\mathrm{post}}_{l}=\alpha_{l}^{\mathrm{post}}\cdot\tanh(\theta^{\mathrm{post}}_{l}\tilde{\mathbf{x}}^{\top}_{l})+\mathbf{b}_{l}^{\mathrm{post}}\\
|
||||
\mathcal{H}^{\mathrm{res}}_{l}=\alpha_{l}^{\mathrm{res}}\cdot\tanh(\theta^{\mathrm{res}}_{l}\tilde{\mathbf{x}}^{\top}_{l})+\mathbf{b}_{l}^{\mathrm{res}},\\
|
||||
\end{cases}
|
||||
(5)
|
||||
where RMSNorm ( ⋅ ) \text{RMSNorm}(\cdot) ( Zhang and Sennrich, 2019 ) is applied to the last dimension, and the scalars α l pre , α l post \alpha_{l}^{\mathrm{pre}},\alpha_{l}^{\mathrm{post}} and α l res ∈ ℝ \alpha_{l}^{\mathrm{res}}\in\mathbb{R} are learnable gating factors initialized to small values. The dynamic mappings are derived via linear projections parameterized by θ l pre , θ l post ∈ ℝ 1 × C \theta^{\mathrm{pre}}_{l},\theta^{\mathrm{post}}_{l}\in\mathbb{R}^{1\times C} and θ l res ∈ ℝ n × C \theta^{\mathrm{res}}_{l}\in\mathbb{R}^{n\times C} , while the static mappings are represented by learnable biases 𝐛 l pre , 𝐛 l post ∈ ℝ 1 × n \mathbf{b}_{l}^{\mathrm{pre}},\mathbf{b}_{l}^{\mathrm{post}}\in\mathbb{R}^{1\times n} and 𝐛 l res ∈ ℝ n × n \mathbf{b}_{l}^{\mathrm{res}}\in\mathbb{R}^{n\times n} .
|
||||
It is worth noting that the introduction of these mappings— ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , and ℋ l res \mathcal{H}^{\mathrm{res}}_{l} —incurs negligible computational overhead, as the typical expansion rate n n , e.g. 4, is much smaller than the input dimension C C . With this design, HC effectively decouples the information capacity of the residual stream from the layer’s input dimension, which is strongly correlated with the model’s computational complexity (FLOPs). Consequently, HC offers a new avenue for scaling by adjusting the residual stream width, complementing the traditional scaling dimensions of model FLOPs and training data size discussed in pre-training scaling laws ( Hoffmann et al., 2022 ) .
|
||||
Although HC necessitates three mappings to manage the dimensional mismatch between the residual stream and the layer input, preliminary experiments presented in Tab. 1 indicate that the residual mapping ℋ l res \mathcal{H}^{\mathrm{res}}_{l} yields the most significant performance gain. This finding underscores the critical importance of effective information exchange within the residual stream.
|
||||
Table 1 : Ablation Study of HC Components. When a specific mapping ( ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , or ℋ l res \mathcal{H}^{\mathrm{res}}_{l} ) is disabled, we employ a fixed mapping to maintain dimensional consistency: uniform weights of 1 / n 1/n for ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , uniform weights of ones for ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , and the identity matrix for ℋ l res \mathcal{H}^{\mathrm{res}}_{l} .
|
||||
ℋ l res \mathcal{H}^{\mathrm{res}}_{l}
|
||||
ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l}
|
||||
ℋ l post \mathcal{H}^{\mathrm{post}}_{l}
|
||||
Absolute Loss Gap
|
||||
0.0 \phantom{-\ }0.0\phantom{00}
|
||||
✓
|
||||
− 0.022 -\ 0.022
|
||||
✓
|
||||
✓
|
||||
− 0.025 -\ 0.025
|
||||
✓
|
||||
✓
|
||||
✓
|
||||
− 0.027 -\ 0.027
|
||||
## 3.1 Numerical Instability
|
||||
While the residual mapping ℋ l res \mathcal{H}^{\mathrm{res}}_{l} is instrumental for performance, its sequential application poses a significant risk to numerical stability. As detailed in Eq. ( 4 ), when HC is extended across multiple layers, the effective signal propagation from layer l l to L L is governed by the composite mapping ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} .
|
||||
Since the learnable mapping ℋ l res \mathcal{H}^{\mathrm{res}}_{l} is unconstrained, this composite mapping inevitably deviates from the identity mapping. Consequently, the signal magnitude is prone to explosion or vanishing during both the forward pass and backpropagation. This phenomenon undermines the fundamental premise of residual learning, which relies on unimpeded signal flow, thereby destabilizing the training process in deeper or larger-scale models.
|
||||
Empirical evidence supports this analysis. We observe unstable loss behavior in large-scale experiments, as illustrated in Fig. 2 . Taking m HC as the baseline, HC exhibits an unexpected loss surge around the 12k step, which is highly correlated with the instability in the gradient norm. Furthermore, the analysis on ℋ l res \mathcal{H}^{\mathrm{res}}_{l} validates the mechanism of this instability.
|
||||
To quantify how the composite mapping ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} amplifies signals along the residual stream, we utilize two metrics. The first, based on the maximum absolute value of the row sums of the composite mapping, captures the worst-case expansion in the forward pass. The second, based on the maximum absolute column sum, corresponds to the backward pass. We refer to these metrics as the Amax Gain Magnitude of the composite mapping. As shown in Fig. 3 (b), the Amax Gain Magnitude yields extreme values with peaks of 3000, a stark divergence from 1 that confirms the presence of exploding residual streams.
|
||||
Figure 2 : Training Instability of Hyper-Connections (HC). This figure illustrates (a) the absolute loss gap of HC relative to m HC, and (b) the comparisons of gradient norms. All results are based on 27B models.
|
||||
Figure 3 : Propagation Instability of Hyper-Connections (HC). This figure illustrates the propagation dynamics of (a) the single-layer mapping ℋ l res \mathcal{H}^{\mathrm{res}}_{l} and (b) the composite mapping ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} within the 27B model. The layer index l l (x-axis) unrolls each standard Transformer block into two independent layers (Attention and FFN). The Amax Gain Magnitude (y-axis) is calculated as the maximum absolute row sum (for the forward signal) and column sum (for the backward gradient), averaged over all tokens in a selected sequence.
|
||||
## 3.2 System Overhead
|
||||
While the computational complexity of HC remains manageable due to the linearity of the additional mappings, the system-level overhead prevents a non-negligible challenge. Specifically, memory access (I/O) costs often constitute one of the primary bottlenecks in modern model architectures, which is widely referred to as the “memory wall” ( Dao et al., 2022 ) . This bottleneck is frequently overlooked in architectural design, yet it decisively impacts runtime efficiency.
|
||||
Table 2 : Comparison of Memory Access Costs Per Token. This analysis accounts for the overhead introduced by the residual stream maintenance in the forward pass, excluding the internal I/O of the layer function ℱ \mathcal{F} .
|
||||
Method
|
||||
Operation
|
||||
Read (Elements)
|
||||
Write (Elements)
|
||||
Residual
|
||||
Connection
|
||||
Residual Merge
|
||||
2 C 2C
|
||||
C C
|
||||
Total I/O
|
||||
𝟐 𝐂 \mathbf{2C}
|
||||
𝐂 \mathbf{C}
|
||||
Hyper-
|
||||
Connections
|
||||
Calculate ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , ℋ l res \mathcal{H}^{\mathrm{res}}_{l}
|
||||
n C nC
|
||||
n 2 + 2 n n^{2}+2n
|
||||
ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l}
|
||||
n C + n nC+n
|
||||
C C
|
||||
ℋ l post \mathcal{H}^{\mathrm{post}}_{l}
|
||||
C + n C+n
|
||||
n C nC
|
||||
ℋ l res \mathcal{H}^{\mathrm{res}}_{l}
|
||||
n C + n 2 nC+n^{2}
|
||||
n C nC
|
||||
Residual Merge
|
||||
2 n C 2nC
|
||||
n C nC
|
||||
Total I/O
|
||||
( 𝟓 𝐧 + 𝟏 ) 𝐂 + 𝐧 𝟐 + 𝟐 𝐧 \mathbf{(5n+1)C+n^{2}+2n}
|
||||
( 𝟑 𝐧 + 𝟏 ) 𝐂 + 𝐧 𝟐 + 𝟐 𝐧 \mathbf{(3n+1)C+n^{2}+2n}
|
||||
Focusing on the widely adopted pre-norm Transformer ( Vaswani et al., 2017 ) architecture, we analyze the I/O patterns inherent to HC. Tab. 2 summarizes the per token memory access overhead in a single residual layer introduced by the n n -stream residual design. The analysis reveals that HC increases the memory access cost by a factor approximately proportional to n n . This excessive I/O demand significantly degrades training throughput without the mitigation of fused kernels. Besides, since ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , and ℋ l res \mathcal{H}^{\mathrm{res}}_{l} involve learnable parameters, their intermediate activations are required for backpropagation. This results in a substantial increase in the GPU memory footprint, often necessitating gradient checkpointing to maintain feasible memory usage. Furthermore, HC requires n n -fold more communication cost in pipeline parallelism ( Qi et al., 2024 ) , leading to larger bubbles and decreasing the training throughput.
|
||||
## 4 Method
|
||||
## 4.1 Manifold-Constrained Hyper-Connections
|
||||
Drawing inspiration from the identity mapping principle ( He et al., 2016b ) , the core premise of m HC is to constrain the residual mapping ℋ l res \mathcal{H}^{\mathrm{res}}_{l} onto a specific manifold.
|
||||
While the original identity mapping ensures stability by enforcing ℋ l res = 𝐈 \mathcal{H}^{\mathrm{res}}_{l}=\mathbf{I} , it fundamentally precludes information exchange within the residual stream, which is critical for maximizing the potential of multi-stream architectures. Therefore, we propose projecting the residual mapping onto a manifold that simultaneously maintains the stability of signal propagation across layers and facilitates mutual interaction among residual streams to preserve the model’s expressivity.
|
||||
To this end, we restrict ℋ l res \mathcal{H}^{\mathrm{res}}_{l} to be a doubly stochastic matrix, which has non-negative entries where both the rows and columns sum to 1. Formally, let ℳ res \mathcal{M}^{\mathrm{res}} denote the manifold of doubly stochastic matrices (also known as the Birkhoff polytope).
|
||||
We constrain ℋ l res \mathcal{H}^{\mathrm{res}}_{l} to 𝒫 ℳ res ( ℋ l res ) \mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}^{\mathrm{res}}_{l}) , defined as:
|
||||
𝒫 ℳ res ( ℋ l res ) :- { ℋ l res ∈ ℝ n × n ∣ ℋ l res 𝟏 n = 𝟏 n , 1 n ⊤ ℋ l res = 𝟏 n ⊤ , ℋ l res ⩾ 0 } , \mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}^{\mathrm{res}}_{l})\coloneq\left\{\mathcal{H}^{\mathrm{res}}_{l}\in\mathbb{R}^{n\times n}\mid\mathcal{H}^{\mathrm{res}}_{l}\mathbf{1}_{n}=\mathbf{1}_{n},\ \mathbf{1}^{\top}_{n}\mathcal{H}^{\mathrm{res}}_{l}=\mathbf{1}^{\top}_{n},\ \mathcal{H}^{\mathrm{res}}_{l}\geqslant 0\right\},
|
||||
(6)
|
||||
where 𝟏 n \mathbf{1}_{n} represents the n n -dimensional vector of all ones.
|
||||
It is worth noting that when n = 1 n=1 , the doubly stochastic condition degenerates to the scalar 1 1 , thereby recovering the original identity mapping. The choice of double stochasticity confers several rigorous theoretical properties beneficial for large-scale model training:
|
||||
1.
|
||||
Norm Preservation: The spectral norm of a doubly stochastic matrix is bounded by 1 (i.e., ‖ ℋ l res ‖ 2 ≤ 1 \|\mathcal{H}^{\mathrm{res}}_{l}\|_{2}\leq 1 ). This implies that the learnable mapping is non-expansive, effectively mitigating the gradient explosion problem.
|
||||
2.
|
||||
Compositional Closure: The set of doubly stochastic matrices is closed under matrix multiplication. This ensures that the composite residual mapping across multiple layers, ∏ i = 1 L − l ℋ L − i res \prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}} , remains doubly stochastic, thereby preserving stability throughout the entire depth of the model.
|
||||
3.
|
||||
Geometric Interpretation via the Birkhoff Polytope: The set ℳ res \mathcal{M}^{\mathrm{res}} forms the Birkhoff polytope, which is the convex hull of the set of permutation matrices. This provides a clear geometric interpretation: the residual mapping acts as a convex combination of permutations. Mathematically, the repeated application of such matrices tends to increase the mixing of information across streams monotonically, effectively functioning as a robust feature fusion mechanism.
|
||||
Additionally, we impose non-negativity constraints on the input mappings ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} and output mappings ℋ l post \mathcal{H}^{\mathrm{post}}_{l} . This constrain prevents signal cancellation arising from the composition of positive and negative coefficients, which can also be considered as a special manifold projection.
|
||||
## 4.2 Parameterization and Manifold Projection
|
||||
In this section, we detail the calculation process of ℋ l pre , ℋ l post , and ℋ l res \mathcal{H}^{\mathrm{pre}}_{l},\mathcal{H}^{\mathrm{post}}_{l},\text{and }\mathcal{H}^{\mathrm{res}}_{l} in m HC. Given the input hidden matrix 𝐱 l ∈ ℝ n × C \mathbf{x}_{l}\in\mathbb{R}^{n\times C} at the l l -th layer, we first flatten it into a vector 𝐱 → l = vec ( 𝐱 l ) ∈ ℝ 1 × n C \vec{\mathbf{x}}_{l}=\text{vec}(\mathbf{x}_{l})\in\mathbb{R}^{1\times nC} to preserve full context information. Then, we follow the original HC formulation to get the dynamic mappings and the static mappings as follows:
|
||||
{ 𝐱 → l ′ = RMSNorm ( 𝐱 → l ) ℋ ~ l pre = α l pre ⋅ ( 𝐱 → l ′ φ l pre ) + 𝐛 l pre ℋ ~ l post = α l post ⋅ ( 𝐱 → l ′ φ l post ) + 𝐛 l post ℋ ~ l res = α l res ⋅ mat ( 𝐱 → l ′ φ l res ) + 𝐛 l res , \begin{cases}\vec{\mathbf{x}}^{\prime}_{l}=\text{RMSNorm}(\vec{\mathbf{x}}_{l})\\
|
||||
\tilde{\mathcal{H}}^{\mathrm{pre}}_{l}=\alpha_{l}^{\mathrm{pre}}\cdot(\vec{\mathbf{x}}^{\prime}_{l}\varphi^{\mathrm{pre}}_{l})+\mathbf{b}_{l}^{\mathrm{pre}}\\
|
||||
\tilde{\mathcal{H}}^{\mathrm{post}}_{l}=\alpha_{l}^{\mathrm{post}}\cdot(\vec{\mathbf{x}}^{\prime}_{l}\varphi^{\mathrm{post}}_{l})+\mathbf{b}_{l}^{\mathrm{post}}\\
|
||||
\tilde{\mathcal{H}}^{\mathrm{res}}_{l}=\alpha_{l}^{\mathrm{res}}\cdot\text{mat}(\vec{\mathbf{x}}^{\prime}_{l}\varphi^{\mathrm{res}}_{l})+\mathbf{b}_{l}^{\mathrm{res}},\\
|
||||
\end{cases}
|
||||
(7)
|
||||
where φ l pre , φ l post ∈ ℝ n C × n \varphi^{\mathrm{pre}}_{l},\varphi^{\mathrm{post}}_{l}\in\mathbb{R}^{nC\times n} and φ l res ∈ ℝ n C × n 2 \varphi^{\mathrm{res}}_{l}\in\mathbb{R}^{nC\times n^{2}} are linear projections for dynamic mappings and mat ( ⋅ ) \text{mat}(\cdot) is a reshape function from ℝ 1 × n 2 \mathbb{R}^{1\times n^{2}} to ℝ n × n \mathbb{R}^{n\times n} .
|
||||
Then, the final constrained mappings are obtained via:
|
||||
{ ℋ l pre = σ ( ℋ ~ l pre ) ℋ l post = 2 σ ( ℋ ~ l post ) ℋ l res = Sinkhorn-Knopp ( ℋ ~ l res ) , \begin{cases}\mathcal{H}^{\mathrm{pre}}_{l}=\sigma(\tilde{\mathcal{H}}^{\mathrm{pre}}_{l})\\
|
||||
\mathcal{H}^{\mathrm{post}}_{l}=2\sigma(\tilde{\mathcal{H}}^{\mathrm{post}}_{l})\\
|
||||
\mathcal{H}^{\mathrm{res}}_{l}=\text{Sinkhorn-Knopp}(\tilde{\mathcal{H}}^{\mathrm{res}}_{l}),\end{cases}
|
||||
(8)
|
||||
where σ ( ⋅ ) \sigma(\cdot) denotes the Sigmoid function. The Sinkhorn-Knopp ( ⋅ ) \text{Sinkhorn-Knopp}(\cdot) operator firstly makes all elements to be positive via an exponent operator and then conducts iterative normalization process that alternately rescales rows and columns to sum to 1. Specifically, given a positive matrix 𝐌 ( 0 ) = exp ( ℋ ~ l res ) \mathbf{M}^{(0)}=\exp(\tilde{\mathcal{H}}^{\mathrm{res}}_{l}) as the start point, the normalization iteration proceeds as:
|
||||
𝐌 ( t ) = 𝒯 r ( 𝒯 c ( 𝐌 ( t − 1 ) ) ) , \mathbf{M}^{(t)}=\mathcal{T}_{r}\left(\mathcal{T}_{c}(\mathbf{M}^{(t-1)})\right),
|
||||
(9)
|
||||
where 𝒯 r \mathcal{T}_{r} and 𝒯 c \mathcal{T}_{c} denote row and column normalization, respectively. This process converges to a doubly stochastic matrix ℋ l res = 𝐌 ( t max ) \mathcal{H}^{\mathrm{res}}_{l}=\mathbf{M}^{(t_{\text{max}})} as t max → ∞ t_{\text{max}}\to\infty . We choose t max = 20 t_{\text{max}}=20 as a practical value in our experiments.
|
||||
## 4.3 Efficient Infrastructure Design
|
||||
In this section, we detail the infrastructure design tailored for m HC.
|
||||
Through rigorous optimization, we implement m HC (with n = 4 n=4 ) in large-scale models with a marginal training overhead of only 6.7%.
|
||||
## 4.3.1 Kernel Fusion
|
||||
Observing that RMSNorm in m HC imposes significant latency when operating on the high-dimensional hidden state 𝐱 → l ∈ ℝ 1 × n C \vec{\mathbf{x}}_{l}\in\mathbb{R}^{1\times nC} , we reorder the dividing-by-norm operation to follow the matrix multiplication. This optimization maintains mathematical equivalence while improving efficiency.
|
||||
Furthermore, we employ mixed-precision strategies to maximize numerical accuracy without compromising speed, and fuse multiple operations with shared memory access into unified compute kernels to reduce memory bandwidth bottlenecks.
|
||||
Based on the inputs and parameters detailed in Eq. ( 10 ) to ( 13 ), we implement three specialized m HC kernels to compute ℋ l pre \mathcal{H}^{\mathrm{pre}}_{l} , ℋ l post \mathcal{H}^{\mathrm{post}}_{l} , and ℋ l res \mathcal{H}^{\mathrm{res}}_{l} . In these kernels, the biases and linear projections are consolidated into 𝐛 l \mathbf{b}_{l} and φ l \varphi_{l} , and the RMSNorm weight is also absorbed in φ l \varphi_{l} .
|
||||
•
|
||||
Eq. ( 14 ) to ( 15 ): We develop a unified kernel that fuses two scans on 𝐱 → l \vec{\mathbf{x}}_{l} , leveraging matrix multiplication units to maximize memory bandwidth utilization. The backward pass—comprising two matrix multiplications—is similarly consolidated into a single kernel, eliminating redundant reloading of 𝐱 → l \vec{\mathbf{x}}_{l} . Both kernels feature a finely tuned pipeline (load, cast, compute, store) to efficiently handle mixed-precision processing.
|
||||
•
|
||||
Eq. ( 16 ) to ( 18 ): These lightweight operations on small coefficients are opportunistically fused into a single kernel, significantly reducing kernel launch overhead.
|
||||
•
|
||||
Eq. ( 19 ): We implement the Sinkhorn-Knopp iteration within a single kernel. For the backward pass, we derive a custom backward kernel that recomputes the intermediate results on-chip and traverses the entire iteration.
|
||||
φ l \displaystyle\varphi_{l}
|
||||
: tfloat32 \displaystyle:\text{tfloat32}
|
||||
[ n C , n 2 + 2 n ] \displaystyle[nC,n^{2}+2n]
|
||||
(10)
|
||||
𝐱 → l \displaystyle\vec{\mathbf{x}}_{l}
|
||||
: bfloat16 \displaystyle:\text{bfloat16}
|
||||
[ 1 , n C ] \displaystyle[1,nC]
|
||||
(11)
|
||||
α l pre , α l post , α l res \displaystyle\alpha_{l}^{\mathrm{pre}},\alpha_{l}^{\mathrm{post}},\alpha_{l}^{\mathrm{res}}
|
||||
: float32 \displaystyle:\text{float32}
|
||||
Scalars
|
||||
(12)
|
||||
𝐛 l \displaystyle\mathbf{b}_{l}
|
||||
: float32 \displaystyle:\text{float32}
|
||||
[ 1 , n 2 + 2 n ] \displaystyle[1,n^{2}+2n]
|
||||
(13)
|
||||
[ ℋ ~ ~ l pre , ℋ ~ ~ l post , ℋ ~ ~ l res ] \displaystyle\left[{\tilde{\tilde{\mathcal{H}}}^{\mathrm{pre}}_{l}},{\tilde{\tilde{\mathcal{H}}}^{\mathrm{post}}_{l}},{\tilde{\tilde{\mathcal{H}}}^{\mathrm{res}}_{l}}\right]
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= 𝐱 → l φ l \displaystyle=\vec{\mathbf{x}}_{l}\varphi_{l}
|
||||
(14)
|
||||
r \displaystyle r
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= ‖ 𝐱 → l ‖ 2 / n C \displaystyle=\left\|\vec{\mathbf{x}}_{l}\right\|_{2}/\sqrt{nC}
|
||||
(15)
|
||||
[ ℋ ~ l pre , ℋ ~ l post , ℋ ~ l res ] \displaystyle\left[\tilde{\mathcal{H}}^{\mathrm{pre}}_{l},\tilde{\mathcal{H}}^{\mathrm{post}}_{l},\tilde{\mathcal{H}}^{\mathrm{res}}_{l}\right]
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= 1 / r [ α l pre ℋ ~ ~ l pre , α l post ℋ ~ ~ l post , α l res ℋ ~ ~ l res ] + 𝐛 l \displaystyle=1/r\left[\alpha_{l}^{\mathrm{pre}}{\tilde{\tilde{\mathcal{H}}}^{\mathrm{pre}}_{l}},\alpha_{l}^{\mathrm{post}}{\tilde{\tilde{\mathcal{H}}}^{\mathrm{post}}_{l}},\alpha_{l}^{\mathrm{res}}{\tilde{\tilde{\mathcal{H}}}^{\mathrm{res}}_{l}}\right]+\mathbf{b}_{l}
|
||||
(16)
|
||||
ℋ l pre \displaystyle\mathcal{H}^{\mathrm{pre}}_{l}
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= σ ( ℋ ~ l pre ) \displaystyle=\sigma\left(\tilde{\mathcal{H}}^{\mathrm{pre}}_{l}\right)
|
||||
(17)
|
||||
ℋ l post \displaystyle\mathcal{H}^{\mathrm{post}}_{l}
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= 2 σ ( ℋ ~ l post ) \displaystyle=2\sigma\left(\tilde{\mathcal{H}}^{\mathrm{post}}_{l}\right)
|
||||
(18)
|
||||
ℋ l res \displaystyle\mathcal{H}^{\mathrm{res}}_{l}
|
||||
: float32 \displaystyle:\text{float32}
|
||||
= Sinkhorn-Knopp ( ℋ ~ l res ) \displaystyle=\text{Sinkhorn-Knopp}\left(\tilde{\mathcal{H}}^{\mathrm{res}}_{l}\right)
|
||||
(19)
|
||||
Using the coefficients derived from the aforementioned kernels, we introduce two additional kernels to apply these mappings: one for ℱ pre :- ℋ l pre 𝐱 l \mathcal{F}_{\mathrm{pre}}\coloneq\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_{l} and another for ℱ post , res :- ℋ l res 𝐱 l + ℋ l post ⊤ ℱ ( ⋅ , ⋅ ) \mathcal{F}_{\mathrm{post,res}}\coloneq\mathcal{H}^{\mathrm{res}}_{l}\mathbf{x}_{l}+\mathcal{H}_{l}^{\mathrm{post}\,\top}\mathcal{F}(\cdot,\cdot) . Through fusing the application of ℋ l post \mathcal{H}^{\mathrm{post}}_{l} and ℋ l res \mathcal{H}^{\mathrm{res}}_{l} with residual merging, we reduce the number of elements read from ( 3 n + 1 ) C (3n+1)C to ( n + 1 ) C (n+1)C and the number of elements written from 3 n C 3nC to n C nC for this kernel.
|
||||
We efficiently implement the majority of kernels (excluding Eq. ( 14 ) to ( 15 )) using TileLang ( Wang et al., 2025 ) . This framework streamlines the implementation of kernels with complex calculation process and allows us to fully utilize the memory bandwidth with minimal engineering effort.
|
||||
## 4.3.2 Recomputing
|
||||
The n n -stream residual design introduces substantial memory overhead during training.
|
||||
To mitigate this, we discard the intermediate activations of the m HC kernels after the forward pass and recompute them on-the-fly in the backward pass, through re-executing the m HC kernels without the heavy layer function ℱ \mathcal{F} .
|
||||
Consequently, for a block of L r L_{r} consecutive layers, we need only store the input 𝐱 l 0 \mathbf{x}_{l_{0}} to the first layer.
|
||||
Excluding lightweight coefficients while accounting for the pre-norm with in ℱ \mathcal{F} , Tab. 3 summarizes the intermediate activations preserved for the backward pass.
|
||||
Table 3 : Stored and Recomputed Intermediate Activations We list per token activation preserved for the backward pass and the transient activation recomputed in L r L_{r} consecutive layers. Layer l 0 l_{0} represents the first layer in L r L_{r} layers and layer l l is in [ l 0 , l 0 + L r − 1 ] [l_{0},l_{0}+L_{r}-1] .
|
||||
Activations
|
||||
𝐱 l 0 \mathbf{x}_{l_{0}}
|
||||
ℱ ( ℋ l pre 𝐱 l , 𝒲 l ) \mathcal{F}(\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_{l},\mathcal{W}_{l})
|
||||
𝐱 l \mathbf{x}_{l}
|
||||
ℋ l pre 𝐱 l \mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_{l}
|
||||
RMSNorm ( ℋ l pre 𝐱 l ) \text{RMSNorm}(\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_{l})
|
||||
Size (Elements)
|
||||
n C nC
|
||||
C C
|
||||
n C nC
|
||||
C C
|
||||
C C
|
||||
Stored Method
|
||||
Every L r L_{r} layers
|
||||
Every layer
|
||||
Transient inside L r L_{r} layers
|
||||
Since m HC kernels recomputation is performed for blocks of L r L_{r} consecutive layers, given a total of L L layers, we must persistently store the first layer input 𝐱 l 0 \mathbf{x}_{l_{0}} for all ⌈ L L r ⌉ \lceil\tfrac{L}{L_{r}}\rceil blocks for the backward pass.
|
||||
In addition to this resident memory, the recomputation process introduces a transient memory overhead of ( n + 2 ) C × L r (n+2)C\times L_{r} elements for the active block, which determines the peak memory usage during backpropagation. Consequently, we determine the optimal block size L r ∗ L_{r}^{*} by minimizing the total memory footprint corresponded to L r L_{r} :
|
||||
L r ∗ = arg min L r [ n C × ⌈ L L r ⌉ + ( n + 2 ) C × L r ] ≈ n L n + 2 . L_{r}^{*}=\arg\min_{L_{r}}\left[nC\times\left\lceil\frac{L}{L_{r}}\right\rceil+(n+2)C\times L_{r}\right]\approx\sqrt{\frac{nL}{n+2}}.
|
||||
(20)
|
||||
Furthermore, pipeline parallelism in large-scale training imposes a constraint: recomputation blocks must not cross pipeline stage boundaries. Observing that the theoretical optimum L r ∗ L_{r}^{*} typically aligns with the number of layers per pipeline stage, we choose to synchronize the recomputation boundaries with the pipeline stages.
|
||||
## 4.3.3 Overlapping Communication in DualPipe
|
||||
In large-scale training, pipeline parallelism is the standard practice for mitigating parameter and gradient memory footprints.
|
||||
Specifically, we adopt the DualPipe schedule ( Liu et al., 2024b ) , which effectively overlaps scale-out interconnected communication traffic, such as those in expert and pipeline parallelism.
|
||||
However, compared to the single-stream design, the proposed n n -stream residual in m HC incurs substantial communication latency across pipeline stages.
|
||||
Furthermore, at stage boundaries, the recomputation of m HC kernels for all L r L_{r} layers introduces non-negligible computational overhead.
|
||||
To address these bottlenecks, we extend the DualPipe schedule (see Fig. 4 ) to facilitate improved overlapping of communication and computation at pipeline stage boundaries.
|
||||
Figure 4 : Communication-Computation Overlapping for m HC.
|
||||
We extend the DualPipe schedule to handle the overhead introduced by m HC. Lengths of each block are illustrative only and do not represent actual duration. (F), (B), (W) refers to forward pass, backward pass, weight gradient computation, respectively. ℱ A and ℱ M \mathcal{F}^{\mathrm{A}}\ \text{and}\ \mathcal{F}^{\mathrm{M}} represents kernels corresponded to Attention and MLP, respectively.
|
||||
Notably, to prevent blocking the communication stream, we execute the ℱ post , res \mathcal{F}_{\mathrm{post,res}} kernels of MLP (i.e. FFN) layers on a dedicated high-priority compute stream. We further refrain from employing persistent kernels for long-running operations in attention layers, thereby preventing extended stalls. This design enables the preemption of overlapped attention computations, allowing for flexible scheduling while maintaining high utilization of the compute device’s processing units.
|
||||
Furthermore, the recomputation process is decoupled from pipeline communication dependencies, as the initial activation of each stage 𝐱 l 0 \mathbf{x}_{l_{0}} is already cached locally.
|
||||
## 5 Experiments
|
||||
## 5.1 Experimental Setup
|
||||
We validate the proposed method via language model pre-training, conducting a comparative analysis between the baseline, HC, and our proposed m HC.
|
||||
Utilizing MoE architectures inspired by DeepSeek-V3 ( Liu et al., 2024b ) , we train four distinct model variants to cover different evaluation regimes.
|
||||
Specifically, the expansion rate n n for both HC and m HC is set to 4.
|
||||
Our primary focus is a 27B model trained with a dataset size proportional to its parameters, which serves as the subject for our system-level main results.
|
||||
Expanding on this, we analyze the compute scaling behavior by incorporating smaller 3B and 9B models trained with proportional data, which allows us to observe performance trends across varying compute.
|
||||
Additionally, to specifically investigate the token scaling behavior, we train a separate 3B model on a fixed corpus of 1 trillion tokens.
|
||||
Detailed model configurations and training hyper-parameters are provided in Appendix A.1 .
|
||||
## 5.2 Main Results
|
||||
Figure 5 :
|
||||
Training Stability of Manifold-Constrained Hyper-Connections ( m HC). This figure illustrates (a) the absolute training loss gap of m HC and HC relative to the baseline, and (b) the gradient norm of the three methods. All experiments utilize the 27B model. The results demonstrate that m HC exhibits improved stability in terms of both loss and gradient norm.
|
||||
Table 4 :
|
||||
System-level Benchmark Results for 27B Models. This table compares the zero-shot and few-shot performance of the Baseline, HC, and m HC across 8 diverse downstream benchmarks. m HC consistently outperforms the Baseline and surpasses HC on the majority of benchmarks, demonstrating its effectiveness in large-scale pre-training.
|
||||
Benchmark
|
||||
BBH
|
||||
DROP
|
||||
GSM8K
|
||||
HellaSwag
|
||||
MATH
|
||||
MMLU
|
||||
PIQA
|
||||
TriviaQA
|
||||
(Metric)
|
||||
(EM)
|
||||
(F1)
|
||||
(EM)
|
||||
(Acc.)
|
||||
(EM)
|
||||
(Acc.)
|
||||
(Acc.)
|
||||
(EM)
|
||||
# Shots
|
||||
3-shot
|
||||
3-shot
|
||||
8-shot
|
||||
10-shot
|
||||
4-shot
|
||||
5-shot
|
||||
0-shot
|
||||
5-shot
|
||||
27B Baseline
|
||||
43.8
|
||||
47.0
|
||||
46.7
|
||||
73.7
|
||||
22.0
|
||||
59.0
|
||||
78.5
|
||||
54.3
|
||||
27B w/ HC
|
||||
48.9
|
||||
51.6
|
||||
53.2
|
||||
74.3
|
||||
26.4
|
||||
63.0
|
||||
79.9
|
||||
56.3
|
||||
27B w/ m HC
|
||||
51.0
|
||||
53.9
|
||||
53.8
|
||||
74.7
|
||||
26.0
|
||||
63.4
|
||||
80.5
|
||||
57.6
|
||||
We begin by examining the training stability and convergence of the 27B models. As illustrated in Fig. 5 (a), m HC effectively mitigates the training instability observed in HC, achieving a final loss reduction of 0.021 compared to the baseline. This improved stability is further corroborated by the gradient norm analysis in Fig. 5 (b), where m HC exhibits significantly better behavior than HC, maintaining a stable profile comparable to the baseline.
|
||||
Tab. 4 presents the downstream performance across a diverse set of benchmarks ( Hendrycks et al., 2020 ; Cobbe et al., 2021 ; Zellers et al., 2019 ; Hendrycks et al., 2021 ; Bisk et al., 2020 ; Joshi et al., 2017 ) . m HC yields comprehensive improvements, consistently outperforming the baseline and surpassing HC on the majority of tasks. Notably, compared to HC, m HC further enhances the model’s reasoning capabilities, delivering performance gains of 2.1% on BBH ( Suzgun et al., 2022 ) and 2.3% on DROP ( Dua et al., 2019 ) .
|
||||
## 5.3 Scaling Experiments
|
||||
Figure 6 :
|
||||
Scaling properties of m HC compared to the Baseline.
|
||||
(a) Compute Scaling Curve.
|
||||
Solid lines depict the performance gap across different compute budgets. Each point represents a specific compute-optimal configuration of model size and dataset size, scaling from 3B and 9B to 27B parameters.
|
||||
(b) Token Scaling Curve.
|
||||
Trajectory of the 3B model during training. Each point represents the model’s performance at different training tokens.
|
||||
Detailed architectures and training configurations are provided in Appendix A.1 .
|
||||
To assess the scalability of our approach, we report the relative loss improvement of m HC against the baseline across different scales.
|
||||
In Fig. 6 (a), we plot the compute scaling curve spanning 3B, 9B, and 27B parameters.
|
||||
The trajectory indicates that the performance advantage is robustly maintained even at higher computational budgets, showing only marginal attenuation.
|
||||
Furthermore, we examine the within-run dynamics in Fig. 6 (b), which presents the token scaling curve for the 3B model.
|
||||
Collectively, these findings validate the effectiveness of m HC in large-scale scenarios. This conclusion is further corroborated by our in-house large-scale training experiments.
|
||||
## 5.4 Stability Analysis
|
||||
Figure 7 :
|
||||
Propagation Stability of Manifold-Constrained Hyper-Connections ( m HC). This figure illustrates the propagation dynamics of (a) the single-layer mapping 𝒫 ℳ res ( ℋ l res ) \mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}^{\mathrm{res}}_{l}) and (b) the composite mapping ∏ i = 1 L − l 𝒫 ℳ res ( ℋ L − i res ) \prod_{i=1}^{L-l}\mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}_{L-i}^{\mathrm{res}}) within the 27B model. The results demonstrate that m HC significantly enhances propagation stability compared to HC.
|
||||
Figure 8 :
|
||||
Visualizations of Learnable Mappings. This figure displays representative single-layer and composite mappings for HC (first row) and m HC (second row). Each matrix is computed by averaging over all tokens within a selected sequence. The labels annotated along the y-axis and x-axis indicate the forward signal gain (row sum) and the backward gradient gain (column sum), respectively.
|
||||
Similar to Fig. 3 , Fig. 7 illustrates the propagation stability of m HC. Ideally, the single-layer mapping satisfies the doubly stochastic constraint, implying that both the forward signal gain and the backward gradient gain should equal to 1. However, practice implementations utilizing the Sinkhorn-Knopp algorithm must limit the number of iterations to achieve computational efficiency. In our settings, we use 20 iterations to obtain an approximate solution. Consequently, as shown in Fig. 7 (a), the backward gradient gain deviates slightly from 1. In the composite case shown in Fig. 7 (b), the deviation increases but remains bounded, reaching a maximum value of approximately 1.6.
|
||||
Notably, compared to the maximum gain magnitude of nearly 3000 in HC, m HC significantly reduces it by three orders of magnitude.
|
||||
These results demonstrate that m HC significantly enhances propagation stability compared to HC, ensuring stable forward signal and backward gradient flows.
|
||||
Additionally, Fig. 8 displays representative mappings. We observe that for HC, when the maximum gain is large, other values also tend to be significant, which indicates general instability across all propagation paths. In contrast, m HC consistently yields stable results.
|
||||
## 6 Conclusion and Outlook
|
||||
In this paper, we identify that while expanding the width of residual stream and diversifying connections yields performance gains as proposed in Hyper-Connections (HC), the unconstrained nature of these connections leads to signal divergence.
|
||||
This disruption compromises the conservation of signal energy across layers, inducing training instability and hindering the scalability of deep networks.
|
||||
To address these challenges, we introduce Manifold-Constrained Hyper-Connections ( m HC ), a generalized framework that projects the residual connection space onto a specific manifold.
|
||||
By employing the Sinkhorn-Knopp algorithm to enforce a doubly stochastic constraint on residual mappings, m HC transforms signal propagation into a convex combination of features.
|
||||
Empirical results confirm that m HC effectively restores the identity mapping property, enabling stable large-scale training with superior scalability compared to conventional HC.
|
||||
Crucially, through efficient infrastructure-level optimizations, m HC delivers these improvements with negligible computational overhead.
|
||||
As a generalized extension of the HC paradigm, m HC opens several promising avenues for future research.
|
||||
Although this work utilizes doubly stochastic matrices to ensure stability, the framework accommodates the exploration of diverse manifold constraints tailored to specific learning objectives.
|
||||
We anticipate that further investigation into distinct geometric constraints could yield novel methods that better optimize the trade-off between plasticity and stability.
|
||||
Furthermore, we hope m HC rejuvenates community interest in macro-architecture design. By deepening the understanding of how topological structures influence optimization and representation learning, m HC will help address current limitations and potentially illuminate new pathways for the evolution of next-generation foundational architectures.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,802 @@
|
||||
TITLE: Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
|
||||
ARXIV: 2601.07372
|
||||
AUTHORS: Xin Cheng , Rui Tian , Wangding Zeng , Damai Dai , Qinyu Chen , Bingxuan Wang , Zhenda Xie , Kezhao Huang , Xingkai Yu , Chengqi Deng , Shangyan Zhou , Chenggang Zhao , Zhewen Hao , Yukun Li , Han Zhang , Zhengyan Zhang , Yixu Wei , M.Y Xu , Huishuai Zhang , Dongyan Zhao , Wenfeng Liang
|
||||
DATE: [Submitted on 12 Jan 2026 ( v1 ), last revised 12 Jul 2026 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Abstract
|
||||
## 1 Introduction
|
||||
## 2 Architecture
|
||||
## 2.1 Overview
|
||||
## 2.2 Sparse Retrieval via Hashed N N -grams
|
||||
## Tokenizer Compression
|
||||
## Multi-Head Hashing.
|
||||
## 2.3 Context-aware Gating
|
||||
## 2.4 Integration with Multi-branch Architecture
|
||||
## 2.5 System Efficiency: Decoupling Compute and Memory
|
||||
## 3 Scaling Laws and Sparsity Allocation
|
||||
## 3.1 Optimal Allocation Ratio Between MoE and Engram
|
||||
## Compute-matched formulation.
|
||||
## Allocation ratio.
|
||||
## Experimental protocol.
|
||||
## Results and Analysis.
|
||||
## 3.2 Engram under Infinite Memory Regime
|
||||
## Experimental protocol.
|
||||
## Results.
|
||||
## 4 Large Scale Pre-training
|
||||
## 4.1 Experimental Setup
|
||||
## Training Data and Model Configurations
|
||||
## Evaluation Protocol
|
||||
## 4.2 Experimental Results
|
||||
## 5 Long Context Training
|
||||
## 5.1 Experimental Setup
|
||||
## Training Details.
|
||||
## Model Configurations.
|
||||
## Evaluation Benchmarks.
|
||||
## 5.2 Experimental Results
|
||||
## 6 Analysis
|
||||
## 6.1 Is Engram functionally equivalent to increasing the model’s depth?
|
||||
## 6.1.1 Accelerated Prediction Convergence
|
||||
## 6.1.2 Representational Alignment and Effective Depth
|
||||
## 6.2 Structural Ablation and Layer Sensitivity
|
||||
## Reference configuration.
|
||||
## Where should memory be injected?
|
||||
## Which components matter?
|
||||
## 6.3 Sensitivity Analysis
|
||||
## 6.4 System Efficiency
|
||||
## Experimental Setup.
|
||||
## Results.
|
||||
## 6.5 Case Study: Gating Visualization
|
||||
## 7 Related Work
|
||||
## N N -gram Modeling and Embedding Scaling.
|
||||
## High-Cardinality Categorical Embeddings.
|
||||
## Mixture-of-Experts.
|
||||
## Memory Network.
|
||||
## Mechanisms of Knowledge Storage.
|
||||
## 8 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Overview
|
||||
2.2 Sparse Retrieval via Hashed N N -grams
|
||||
Tokenizer Compression
|
||||
Multi-Head Hashing.
|
||||
2.3 Context-aware Gating
|
||||
2.4 Integration with Multi-branch Architecture
|
||||
2.5 System Efficiency: Decoupling Compute and Memory
|
||||
3 Scaling Laws and Sparsity Allocation
|
||||
3.1 Optimal Allocation Ratio Between MoE and Engram
|
||||
Compute-matched formulation.
|
||||
Allocation ratio.
|
||||
Experimental protocol.
|
||||
Results and Analysis.
|
||||
3.2 Engram under Infinite Memory Regime
|
||||
Experimental protocol.
|
||||
Results.
|
||||
4 Large Scale Pre-training
|
||||
4.1 Experimental Setup
|
||||
Training Data and Model Configurations
|
||||
Evaluation Protocol
|
||||
4.2 Experimental Results
|
||||
5 Long Context Training
|
||||
5.1 Experimental Setup
|
||||
Training Details.
|
||||
Model Configurations.
|
||||
Evaluation Benchmarks.
|
||||
5.2 Experimental Results
|
||||
6 Analysis
|
||||
6.1 Is Engram functionally equivalent to increasing the model’s depth?
|
||||
6.1.1 Accelerated Prediction Convergence
|
||||
6.1.2 Representational Alignment and Effective Depth
|
||||
6.2 Structural Ablation and Layer Sensitivity
|
||||
Reference configuration.
|
||||
Where should memory be injected?
|
||||
Which components matter?
|
||||
6.3 Sensitivity Analysis
|
||||
6.4 System Efficiency
|
||||
Experimental Setup.
|
||||
Results.
|
||||
6.5 Case Study: Gating Visualization
|
||||
7 Related Work
|
||||
N N -gram Modeling and Embedding Scaling.
|
||||
High-Cardinality Categorical Embeddings.
|
||||
Mixture-of-Experts.
|
||||
Memory Network.
|
||||
Mechanisms of Knowledge Storage.
|
||||
8 Conclusion
|
||||
References
|
||||
A Detailed Model Architecture and Hyper Parameters
|
||||
B Full Benchmark Curves
|
||||
C Case Study of Tokenizer Compression
|
||||
License: CC BY 4.0
|
||||
arXiv:2601.07372v2 [cs.CL] 12 Jul 2026
|
||||
\reportnumber
|
||||
001
|
||||
##
|
||||
Conditional Memory via Scalable Lookup:
|
||||
A New Axis of Sparsity for Large Language Models
|
||||
Xin Cheng 1,2∗ , Rui Tian 2∗ , Wangding Zeng 2 , Damai Dai 2 , Qinyu Chen 2 ,
|
||||
Bingxuan Wang 2 , Zhenda Xie 2 , Kezhao Huang 2 , Xingkai Yu 2
|
||||
Chengqi Deng 2 , Shangyan Zhou 2 , Chenggang Zhao 2 , Zhewen Hao 2
|
||||
Yukun Li 2 , Han Zhang 2 , Zhengyan Zhang 2 , Yixuan Wei 2 , M.Y Xu 2
|
||||
Huishuai Zhang 1 , Dongyan Zhao 1 , Wenfeng Liang 2
|
||||
1 Peking University 2 DeepSeek-AI
|
||||
{zhanghuishuai, zhaody}@pku.edu.cn
|
||||
{chengxin, tianr22, zengwangding, damai.dai}@deepseek.com
|
||||
## Abstract
|
||||
While Mixture-of-Experts (MoE) scales capacity via conditional computation, Transformers lack a native primitive for knowledge lookup, forcing them to inefficiently simulate retrieval through computation. To address this, we introduce conditional memory as a complementary sparsity axis, instantiated via Engram , a module that modernizes classic N N -gram embedding for 𝒪 ( 1 ) \mathcal{O}(1) lookup.
|
||||
By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law that optimizes the trade-off between neural computation (MoE) and static memory (Engram).
|
||||
Guided by this law, we scale Engram to 27B parameters, achieving superior performance over a strictly iso-parameter and iso-FLOPs MoE baseline.
|
||||
Most notably, while the memory module is expected to aid knowledge retrieval (e.g., MMLU + 3.4 +3.4 ; CMMLU + 4.0 +4.0 ), we observe even larger gains in general reasoning (e.g., BBH + 5.0 +5.0 ; ARC-Challenge + 3.7 +3.7 ) and code/math domains (HumanEval + 3.0 +3.0 ; MATH + 2.4 +2.4 ).
|
||||
Mechanistic analyses reveal that Engram relieves the backbone’s early layers from static reconstruction, effectively deepening the network for complex reasoning. Furthermore, by delegating local dependencies to lookups, it frees up attention capacity for global context, substantially boosting long-context retrieval (e.g., Multi-Query NIAH: 84.2 → 97.0 84.2\to 97.0 ).
|
||||
Finally, Engram establishes infrastructure-aware efficiency: its deterministic addressing enables runtime prefetching from host memory, incurring negligible overhead.
|
||||
We envision conditional memory as an indispensable modeling primitive for next-generation sparse models. Code available at: https://github.com/deepseek-ai/Engram .
|
||||
* * footnotetext: Equal contribution.
|
||||
## 1 Introduction
|
||||
Sparsity is a recurring design principle for intelligent systems, spanning from biological neural circuits ( Olshausen and Field, 1997 ; Lennie, 2003 ) to modern Large Language Models (LLMs). Currently, this principle is primarily realized through Mixture-of-Experts (MoE) ( Shazeer et al., 2017 ; Dai et al., 2024 ) , which scales capacity via conditional computation.
|
||||
Owing to its ability to drastically increase model size without proportional increases in compute, MoE has become the de facto standard for frontier models ( Guo et al., 2025 ; Comanici et al., 2025 ; Kimi, 2025 ) .
|
||||
Despite the success of this conditional computation paradigm, the intrinsic heterogeneity of linguistic signals suggests significant room for structural optimization . Specifically, language modeling entails two qualitatively different sub-tasks: compositional reasoning and knowledge retrieval. While the former demands deep, dynamic computation, a substantial portion of text—such as named entities and formulaic patterns—is local, static, and highly stereotyped ( Erman, 2000 ; Constant et al., 2017 ) . The effectiveness of classical N N -gram models ( Liu et al., 2024b ; Nguyen, 2024 ; Brants et al., 2007 ) in capturing such local dependencies implies that these regularities are naturally represented as computationally inexpensive lookups. Since standard Transformers ( Vaswani et al., 2017 ) lack a native knowledge lookup primitive, current LLMs are forced to simulate retrieval through computation . For instance, resolving a common multi-token entity requires consuming multiple early layers of attention and feed-forward networks ( Ghandeharioun et al., 2024 ; Jin et al., 2025 ) (see Table 3 ). This process essentially amounts to an expensive runtime reconstruction of a static lookup table, wasting valuable sequential depth on trivial operations that could otherwise be allocated to higher-level reasoning.
|
||||
To align model architecture with this linguistic duality, we advocate for a complementary axis of sparsity: conditional memory. Whereas conditional computation sparsely activates parameters to process dynamic logic ( Shazeer et al., 2017 ; Bengio et al., 2013 ) , conditional memory relies on sparse lookup operations to retrieve static embeddings for fixed knowledge. As a preliminary exploration of this paradigm, we revisit N N -gram embeddings ( Bojanowski et al., 2017 ) as a canonical instantiation: local context serves as a key to index a massive embedding table via constant-time 𝒪 ( 1 ) \mathcal{O}(1) lookups ( Tito Svenstrup et al., 2017 ; Huang et al., 2025a ; Pagnoni et al., 2025 ; Yu et al., 2025 ) . Our investigation reveals that, perhaps surprisingly, this static retrieval mechanism can serve as an ideal complement to modern MoE architecture—but only if it is properly designed. In this paper, we propose Engram , a conditional memory module grounded in the classic N N -gram structure but equipped with modern adaptations such as tokenizer compression, multi-head hashing, contextualized gating, and multi-branch integration (detailed in Section 2 ).
|
||||
To quantify the synergy between these two primitives, we formulate the Sparsity Allocation problem: given a fixed total parameter budget, how should capacity be distributed between MoE experts and Engram memory? Our experiments uncover a distinct U-shaped scaling law, revealing that even simple lookup mechanisms, when treated as a first-class modeling primitive, act as essential complements to neural computation.
|
||||
Guided by this allocation law, we scale Engram to a 27B-parameter model. Compared to a strictly iso-parameter and iso-FLOPs MoE-27B baseline, Engram-27B achieves superior efficiency across diverse domains. Crucially, the gains are not limited to knowledge-intensive tasks (e.g., MMLU: + 3.4 +3.4 ; CMMLU: + 4.0 +4.0 ; MMLU-Pro: + 1.8 +1.8 ), where memory capacity is intuitively beneficial; we observe even more significant improvements in general reasoning (e.g., BBH: + 5.0 +5.0 ; ARC-Challenge: + 3.7 +3.7 ; DROP: + 3.3 +3.3 ) and code/math domains (e.g., HumanEval: + 3.0 +3.0 ; MATH: + 2.4 +2.4 ; GSM8K: + 2.2 +2.2 ).
|
||||
Mechanistic analysis via LogitLens ( nostalgebraist, 2020 ) and CKA ( Hendrycks et al., 2021a ) reveals the source of these gains: Engram relieves the backbone from reconstructing static knowledge in early layers, thereby increasing effective depth available for complex reasoning.
|
||||
Furthermore, by delegating local dependencies to lookups, Engram frees up attention capacity to focus on global context, enabling exceptional performance in long-context scenarios—substantially outperforming baselines on LongPPL ( Fang et al., ) and RULER ( Hsieh et al., ) (e.g., Multi-Query NIAH: 97.0 97.0 vs. 84.2 84.2 ; Variable Tracking: 89.0 89.0 vs. 77.0 77.0 ).
|
||||
Finally, we establish infrastructure-aware efficiency as a first-class principle. Unlike MoE’s dynamic routing, Engram employs deterministic IDs to enable runtime prefetching, overlapping communication with computation. Empirical results show that offloading a 100B-parameter table to host memory incurs negligible overhead ( < 3 % <3\% ). This demonstrates that Engram effectively bypasses GPU memory constraints, facilitating aggressive parameter expansion.
|
||||
Figure 1 :
|
||||
The Engram Architecture. The module augments the backbone by retrieving static N N -gram memory and fusing it with dynamic hidden states via context-aware gating. This module is applied only to specific layers to decouple memory from compute, leaving the standard input embedding and un-embedding module intact.
|
||||
## 2 Architecture
|
||||
## 2.1 Overview
|
||||
As shown in Figure 1 , Engram is a conditional memory module designed to augment the Transformer backbone by structurally separating static pattern storage from dynamic computation. Formally, given an input sequence X = ( x 1 , … , x T ) X=(x_{1},\dots,x_{T}) and hidden states 𝐇 ( ℓ ) ∈ ℝ T × d \mathbf{H}^{(\ell)}\in\mathbb{R}^{T\times d} at layer ℓ \ell , the module processes each position t t in two functional phases: retrieval and fusion .
|
||||
First, as detailed in Section 2.2 , we extract and compress suffix N N -grams to deterministically retrieve static embedding vectors via hashing.
|
||||
Subsequently, in Section 2.3 , these retrieved embeddings are dynamically modulated by the current hidden state and refined via a lightweight convolution.
|
||||
Finally, we discuss the integration with multi-branch architectures in Section 2.4 and the system-level design in Section 2.5 .
|
||||
## 2.2 Sparse Retrieval via Hashed N N -grams
|
||||
The first phase maps local contexts to static memory entries, involving tokenizer compression and retrieving embeddings via deterministic hashing.
|
||||
## Tokenizer Compression
|
||||
While N N -gram models typically operate directly on tokenizer outputs, standard subword tokenizers prioritize lossless reconstruction , often assigning disjoint IDs to semantically equivalent terms (e.g., Apple vs. ␣apple ) ( Kudo and Richardson, 2018 ; Li et al., 2023b ) . To maximize semantic density, we implement a vocabulary projection layer. Specifically, we pre-compute a surjective function 𝒫 : V → V ′ \mathcal{P}:V\to V^{\prime} that collapses raw token IDs into canonical identifiers based on normalized textual equivalence (using NFKC ( Whistler, 2025 ) , lowercasing, etc.). In practice, this process achieves a 23% reduction in the effective vocabulary size for a 128k tokenizer (see Appendix C ).
|
||||
Formally, for a token at position t t , we map its raw ID x t x_{t} to a canonical ID x t ′ = 𝒫 ( x t ) x^{\prime}_{t}=\mathcal{P}(x_{t}) to form the suffix N N -gram g t , n = ( x t − n + 1 ′ , … , x t ′ ) g_{t,n}=(x^{\prime}_{t-n+1},\dots,x^{\prime}_{t}) .
|
||||
## Multi-Head Hashing.
|
||||
Directly parameterizing the combinatorial space of all possible N N -grams is intractable. Following Tito Svenstrup et al. (2017) , we adopt a hashing-based approach. To mitigate collisions, we employ K K distinct hash heads for each N N -gram order n n . Each head k k maps the compressed context to an index within an embedding table 𝐄 n , k \mathbf{E}_{n,k} (of prime size M n , k M_{n,k} ) via a deterministic function φ n , k \varphi_{n,k} :
|
||||
z t , n , k ≜ φ n , k ( g t , n ) , 𝐞 t , n , k = 𝐄 n , k [ z t , n , k ] . z_{t,n,k}\triangleq\varphi_{n,k}(g_{t,n}),\quad\mathbf{e}_{t,n,k}=\mathbf{E}_{n,k}[z_{t,n,k}].
|
||||
(1)
|
||||
In practice, φ n , k \varphi_{n,k} is implemented as a lightweight multiplicative-XOR hash. We construct the final memory vector 𝐞 t ∈ ℝ d mem \mathbf{e}_{t}\in\mathbb{R}^{d_{\text{mem}}} by concatenating all retrieved embeddings:
|
||||
𝐞 t ≜ ∥ n = 2 N ∥ k = 1 K 𝐞 t , n , k . \mathbf{e}_{t}\triangleq\mathop{\|}_{n=2}^{N}\mathop{\|}_{k=1}^{K}\mathbf{e}_{t,n,k}.
|
||||
(2)
|
||||
## 2.3 Context-aware Gating
|
||||
The retrieved embeddings 𝐞 t \mathbf{e}_{t} serve as context-independent priors. Being static, however, they inherently lack contextual adaptability and may suffer from noise due to hash collisions or polysemy ( Haber and Poesio, 2024 ) .
|
||||
To enhance expressivity and resolve this ambiguity, we employ a context-aware gating mechanism inspired by Attention ( Vaswani et al., 2017 ; Bahdanau et al., 2015 ) . Specifically, we utilize the current hidden state 𝐡 t \mathbf{h}_{t} —which has aggregated global context via preceding attention layers—as a dynamic Query, while the retrieved memory 𝐞 t \mathbf{e}_{t} serves as the source for both Key and Value projections:
|
||||
𝐤 t = 𝐖 K 𝐞 t , 𝐯 t = 𝐖 V 𝐞 t \mathbf{k}_{t}=\mathbf{W}_{K}\mathbf{e}_{t},\quad\mathbf{v}_{t}=\mathbf{W}_{V}\mathbf{e}_{t}
|
||||
(3)
|
||||
where 𝐖 K , 𝐖 V \mathbf{W}_{K},\mathbf{W}_{V} are learnable projection matrices. To ensure gradient stability ( Dehghani et al., 2023 ) , we apply RMSNorm ( Zhang and Sennrich, 2019 ) to the Query and Key before computing the scalar gate α t ∈ ( 0 , 1 ) \alpha_{t}\in(0,1) :
|
||||
α t = σ ( RMSNorm ( 𝐡 t ) ⊤ RMSNorm ( 𝐤 t ) d ) . \alpha_{t}=\sigma\left(\frac{\text{RMSNorm}(\mathbf{h}_{t})^{\top}\text{RMSNorm}(\mathbf{k}_{t})}{\sqrt{d}}\right).
|
||||
(4)
|
||||
The gated output is defined as 𝐯 ~ t = α t ⋅ 𝐯 t \tilde{\mathbf{v}}_{t}=\alpha_{t}\cdot\mathbf{v}_{t} . This design enforces semantic alignment: if the retrieved memory 𝐞 t \mathbf{e}_{t} contradicts the current context 𝐡 t \mathbf{h}_{t} , the gate α t \alpha_{t} tends toward zero, effectively suppressing the noise.
|
||||
Finally, to expand the receptive field and enhance the model’s non-linearity, we introduce a short, depthwise causal convolution ( Gu et al., 2022 ; Peng et al., 2023 ) . Let 𝐕 ~ ∈ ℝ T × d \tilde{\mathbf{V}}\in\mathbb{R}^{T\times d} denote the sequence of gated values. Using a kernel size w w (set to 4), dilation δ \delta (set to the max N N -gram order) and SiLU activation ( Elfwing et al., 2018 ) , the final output 𝐘 \mathbf{Y} is computed as:
|
||||
𝐘 = SiLU ( Conv1D ( RMSNorm ( 𝐕 ~ ) ) ) + 𝐕 ~ , \mathbf{Y}=\text{SiLU}\left(\text{Conv1D}(\text{RMSNorm}(\tilde{\mathbf{V}}))\right)+\tilde{\mathbf{V}},
|
||||
(5)
|
||||
The Engram module is integrated into the backbone via a residual connection: 𝐇 ( ℓ ) ← 𝐇 ( ℓ ) + 𝐘 \mathbf{H}^{(\ell)}\leftarrow\mathbf{H}^{(\ell)}+\mathbf{Y} , followed by the standard Attention and MoE.
|
||||
Crucially, Engram is not applied to every layer; its specific placement is optimized to balance modeling effectiveness against the system-level latency constraints detailed in Section 2.5 .
|
||||
Figure 2 :
|
||||
System implementation of Engram.
|
||||
(a) Training Phase: The massive embedding tables are sharded across available GPUs. An All-to-All communication primitive is employed to retrieve active embedding rows across devices.
|
||||
(b) Inference Phase: Engram tables are offloaded to host memory. By exploiting the deterministic retrieval logic, the host asynchronously prefetches and transfers embeddings, overlapping communication with the on-device computation of preceding Transformer blocks.
|
||||
## 2.4 Integration with Multi-branch Architecture
|
||||
In this work, rather than standard single-stream connections ( He et al., 2016 ) , we adopt the advanced multi-branch architecture as our default backbone, chosen for its superior modeling capabilities ( Larsson et al., 2017 ; Zhu et al., 2025 ; Xie et al., 2025 ; Szegedy et al., 2015 ) . A defining characteristic of this architecture is the expansion of the residual stream into M M parallel branches, where information flow is modulated by learnable connection weights.
|
||||
Although the Engram module is inherently topology-agnostic, adapting it to this multi-branch framework necessitates structural optimization to balance efficiency and expressivity.
|
||||
Specifically, we implement a parameter-sharing strategy: a single sparse embedding table and a Value projection matrix 𝐖 V \mathbf{W}_{V} are shared across all M M branches, whereas M M distinct Key projection matrices { 𝐖 K ( m ) } m = 1 M \{\mathbf{W}_{K}^{(m)}\}_{m=1}^{M} are employed to enable branch-specific gating behaviors.
|
||||
For the m m -th branch with hidden state 𝐡 t ( m ) \mathbf{h}_{t}^{(m)} , the branch-specific gating signal is computed as:
|
||||
α t ( m ) = σ ( RMSNorm ( 𝐡 t ( m ) ) ⊤ RMSNorm ( 𝐖 K ( m ) 𝐞 t ) d ) . \alpha_{t}^{(m)}=\sigma\left(\frac{\text{RMSNorm}(\mathbf{h}_{t}^{(m)})^{\top}\text{RMSNorm}(\mathbf{W}_{K}^{(m)}\mathbf{e}_{t})}{\sqrt{d}}\right).
|
||||
(6)
|
||||
The retrieved memory is then modulated by these independent gates applied to the shared value vector: 𝐮 t ( m ) = α t ( m ) ⋅ ( 𝐖 V 𝐞 t ) \mathbf{u}_{t}^{(m)}=\alpha_{t}^{(m)}\cdot(\mathbf{W}_{V}\mathbf{e}_{t}) . This design allows the linear projections (one 𝐖 V \mathbf{W}_{V} and M M distinct 𝐖 K ( m ) \mathbf{W}_{K}^{(m)} ) to be fused into a single dense FP8 matrix multiplication, maximizing the compute utilization of modern GPUs. Unless otherwise stated, all experiments utilize this integration with Manifold-Constrained Hyper-Connections ( M = 4 M=4 ) ( Xie et al., 2025 ) .
|
||||
## 2.5 System Efficiency: Decoupling Compute and Memory
|
||||
Scaling model parameters is often constrained by the limited capacity of GPU high-bandwidth memory (HBM). However, unlike MoE which relies on runtime hidden states for dynamic routing, Engram employs a deterministic retrieval mechanism based solely on input token IDs, naturally decoupling parameter storage from computation. This predictability facilitates specialized optimization strategies for both training and inference, as illustrated in Figure 2 .
|
||||
During training, to accommodate large-scale embedding tables, we employ standard model parallelism by sharding the tables across available GPUs. An All-to-All communication primitive is used to gather active rows in the forward pass and dispatch gradients in the backward pass, enabling the total memory capacity to scale linearly with the number of accelerators.
|
||||
During inference, this deterministic nature enables a prefetch-and-overlap strategy. Since memory indices are known prior to the forward pass, the system can asynchronously retrieve embeddings from abundant host memory via PCIe. To effectively mask communication latency, the Engram module is placed at specific layers within the backbone, leveraging the computation of preceding layers as a buffer to prevent GPU stalls.
|
||||
This necessitates a hardware-algorithm co-design strategy: while placing Engram deeper extends the compute window available for hiding latency, our ablation in Section 6.2 shows that modeling performance favors early intervention to offload local pattern reconstruction. Therefore, the optimal placement must simultaneously satisfy both modeling and system latency constraints.
|
||||
Furthermore, natural language N N -grams inherently follow a Zipfian distribution ( Piantadosi, 2014 ; Chao and Zipf, 1950 ) , where a small fraction of patterns accounts for the vast majority of memory accesses. This statistical property motivates a Multi-Level Cache Hierarchy: frequently accessed embeddings can be cached in faster storage tiers (e.g., GPU HBM or Host DRAM), while the long tail of rare patterns resides in slower, high-capacity media (e.g., NVMe SSD). This stratification allows Engram to scale to massive memory capacities with minimal impact on effective latency.
|
||||
Figure 3 : Sparsity allocation and Engram scaling.
|
||||
Left: Validation loss across allocation ratios ρ \rho . Two compute budgets are shown ( 2 e 20 2\text{e}20 and 6 e 20 6\text{e}20 FLOPs). Both regimes exhibit a U-shape, with hybrid allocation surpassing Pure MoE.
|
||||
Right: Scaling behavior in the infinite-memory regime. Validation loss exhibits a log-linear trend with respect to the number of embeddings.
|
||||
## 3 Scaling Laws and Sparsity Allocation
|
||||
Engram, as an instantiation of conditional memory, is structurally complementary to the conditional computation provided by MoE experts.
|
||||
This section investigates the scaling properties of this duality and how to optimally allocate sparse capacity. Specifically, two key questions drive our research:
|
||||
1.
|
||||
Allocation under Finite Constraints. When total parameters and training compute are fixed (Iso-parameter and Iso-FLOPs), how should we split the sparse capacity between MoE experts and Engram embeddings?
|
||||
2.
|
||||
Infinite Memory Regime. Considering the non-scaling 𝒪 ( 1 ) \mathcal{O}(1) overhead of Engram, if the memory budget is relaxed or scaled aggressively, what scaling behavior does Engram exhibit by itself?
|
||||
## 3.1 Optimal Allocation Ratio Between MoE and Engram
|
||||
## Compute-matched formulation.
|
||||
We analyze the trade-off using three parameter metrics:
|
||||
•
|
||||
P tot P_{\mathrm{tot}} : total trainable parameters, excluding vocabulary embedding and LM head.
|
||||
•
|
||||
P act P_{\mathrm{act}} : activated parameters per token. This quantity determines the training cost (FLOPs).
|
||||
•
|
||||
P sparse ≜ P tot − P act P_{\mathrm{sparse}}\triangleq P_{\mathrm{tot}}-P_{\mathrm{act}} : the inactive parameters, which represents the “free” parameter budget available for scaling model size without incurring computational cost (e.g., unselected experts or unretrieved embeddings).
|
||||
We keep P tot P_{\mathrm{tot}} and P act P_{\mathrm{act}} fixed within each FLOPs budget, so that models have the same number of parameters and the same per-token FLOPs. For MoE, P act P_{\mathrm{act}} is determined by the top- k k selected experts, while the parameters of non-selected experts contribute to P sparse P_{\mathrm{sparse}} . For Engram, only a constant number of slots are retrieved per token, so scaling the number of embedding slots increases P tot P_{\mathrm{tot}} without increasing per-token FLOPs.
|
||||
## Allocation ratio.
|
||||
We define the allocation ratio ρ ∈ [ 0 , 1 ] \rho\in[0,1] as the fraction of the inactive-parameter budget assigned to MoE expert capacity:
|
||||
P MoE ( sparse ) = ρ P sparse , P Engram = ( 1 − ρ ) P sparse . P_{\mathrm{MoE}}^{(\mathrm{sparse})}=\rho\,P_{\mathrm{sparse}},\qquad P_{\mathrm{Engram}}=(1-\rho)\,P_{\mathrm{sparse}}.
|
||||
(7)
|
||||
Intuitively:
|
||||
•
|
||||
ρ = 1 \rho=1 corresponds to a pure MoE model (all inactive parameters are routed experts).
|
||||
•
|
||||
ρ < 1 \rho<1 reduces the number of routed experts and reallocates the freed parameters to Engram embedding slots.
|
||||
## Experimental protocol.
|
||||
We evaluate this trade-off at two compute regimes and maintain a constant sparsity ratio P tot / P act ≈ 10 P_{\mathrm{tot}}/P_{\mathrm{act}}\approx 10 across both settings:
|
||||
•
|
||||
C = 2 × 10 20 C=2\times 10^{20} FLOPs: P tot ≈ 5.7 P_{\mathrm{tot}}\approx 5.7 B and P act = 568 P_{\mathrm{act}}=568 M. The baseline ( ρ = 1 \rho=1 ) has a total of 106 experts.
|
||||
•
|
||||
C = 6 × 10 20 C=6\times 10^{20} FLOPs: P tot ≈ 9.9 P_{\mathrm{tot}}\approx 9.9 B and P act = 993 P_{\mathrm{act}}=993 M. The baseline ( ρ = 1 \rho=1 ) has a total of 99 experts.
|
||||
For different ρ \rho , we construct the corresponding model only by adjusting the number of routed experts and the number of Engram embedding slots.
|
||||
All runs use the identical training pipeline and optimization hyperparameters.
|
||||
## Results and Analysis.
|
||||
Figure 3 (left) reveals a consistent U-shaped relationship between validation loss and the allocation ratio ρ \rho .
|
||||
Remarkably, the Engram model achieves comparable performance to the pure MoE baseline ( ρ = 100 % \rho=100\% ) even when the MoE allocation is reduced to just ρ ≈ 40 % \rho\approx 40\% (i.e., a total of 46 experts for the 5.7B model and 43 experts for the 9.9B model).
|
||||
Furthermore, the pure MoE baseline proves suboptimal: reallocating roughly 20 % – 25 % 20\%\text{--}25\% of the sparse parameter budget to Engram yields the best performance.
|
||||
Quantitatively, in the 10B regime ( C = 6 × 10 20 C=6\times 10^{20} ), validation loss improves from 1.7248 1.7248 (at ρ = 100 % \rho=100\% ) to 1.7109 1.7109 near the optimum of ρ ≈ 80 % \rho\approx 80\% ( Δ = 0.0139 \Delta=0.0139 ).
|
||||
Crucially, the location of this optimum is stable across regimes ( ρ ≈ 75 % – 80 % \rho\approx 75\%\text{--}80\% ), suggesting a robust allocation preference across the examined scales (under fixed sparsity). This observed U-shape confirms the structural complementarity between the two modules:
|
||||
•
|
||||
MoE-dominated ( ρ → 100 % \rho\to 100\% ): The model lacks dedicated memory for static patterns, forcing it to inefficiently reconstruct them through depth and computation.
|
||||
•
|
||||
Engram-dominated ( ρ → 0 % \rho\to 0\% ): The model loses conditional computation capacity, hurting tasks that require dynamic, context-dependent reasoning; memory cannot replace computation in this regime.
|
||||
## 3.2 Engram under Infinite Memory Regime
|
||||
In Section 3.1 , we optimized the allocation under a fixed parameter budget. We now explore the complementary setting: aggressive memory scaling.
|
||||
This investigation is motivated by Engram’s unique ability to decouple storage from compute detailed in Section 2.5 .
|
||||
## Experimental protocol.
|
||||
We utilize a fixed MoE backbone with P tot ≈ 3 P_{\mathrm{tot}}\approx 3 B and P act = 568 P_{\mathrm{act}}=568 M, trained for 100B tokens to ensure convergence.
|
||||
On top of this backbone, we attach an Engram table and sweep the number of slots M M from 2.58 × 10 5 2.58\times 10^{5} to 1.0 × 10 7 1.0\times 10^{7} (adding up to ≈ 13 \approx 13 billion parameters).
|
||||
For baselines, we compare against OverEncoding ( Huang et al., 2025a ) , which integrates N N -gram embeddings via averaging with the vocabulary embedding.
|
||||
We note that while other work such as SCONE ( Yu et al., 2025 ) also investigates large-scale embeddings, it is primarily inference-focused and includes extra module ( f-gram model ) and additional training FLOPs, rendering it incompatible with the strict iso-compute constraints of this study.
|
||||
## Results.
|
||||
Figure 3 (right) demonstrates that scaling the number of memory slots yields a clear and consistent improvement in validation loss.
|
||||
Across the explored range, the curve follows a strict power law (linear in log-space), indicating that Engram provides a predictable scaling knob: larger memory continues to pay off without requiring additional computation.
|
||||
Crucially, regarding scaling efficiency: while the direct averaging approach of OverEncoding benefits from larger memory tables, Engram unlocks much larger scaling potential from the same memory budget. Together with the allocation law in Section 3.1 , these results validate that conditional memory serves as a distinct, scalable axis of sparse capacity that complements the conditional computation of MoE.
|
||||
## 4 Large Scale Pre-training
|
||||
Table 1 : Pre-training performance comparison between dense, MoE, and Engram models . All models are trained for 262B tokens and are matched in activated parameters (3.8B). Engram-27B is iso-parameters with MoE-27B by reallocating parameters from routed experts (72 → \to 55) to a 5.7B-parameter Engram memory. Engram-40B further increases Engram memory (18.5B parameters) while keeping the activated-parameter budget fixed. Full training-time benchmark trajectories are reported in Appendix B .
|
||||
Benchmark (Metric)
|
||||
# Shots
|
||||
Dense-4B
|
||||
MoE-27B
|
||||
Engram-27B
|
||||
Engram-40B
|
||||
# Total Params
|
||||
4.1B
|
||||
26.7B
|
||||
26.7B
|
||||
39.5B
|
||||
# Activated (w/o token embed)
|
||||
3.8B
|
||||
3.8B
|
||||
3.8B
|
||||
3.8B
|
||||
# Trained Tokens
|
||||
262B
|
||||
262B
|
||||
262B
|
||||
262B
|
||||
# Experts (shared + routed, top- k k )
|
||||
-
|
||||
2 + 72 2+72 (top-6)
|
||||
2 + 55 2+55 (top-6)
|
||||
2 + 55 2+55 (top-6)
|
||||
# Engram Params
|
||||
-
|
||||
-
|
||||
5.7B
|
||||
18.5B
|
||||
Language
|
||||
Modeling
|
||||
Pile (loss)
|
||||
-
|
||||
2.091
|
||||
1.960
|
||||
1.950
|
||||
1.942
|
||||
Validation Set (loss)
|
||||
-
|
||||
1.768
|
||||
1.634
|
||||
1.622
|
||||
1.610
|
||||
Knowledge
|
||||
&
|
||||
Reasoning
|
||||
MMLU (Acc.)
|
||||
5-shot
|
||||
48.6
|
||||
57.4
|
||||
60.4
|
||||
60.6
|
||||
MMLU-Redux (Acc.)
|
||||
5-shot
|
||||
50.7
|
||||
60.6
|
||||
64.0
|
||||
64.5
|
||||
MMLU-Pro (Acc.)
|
||||
5-shot
|
||||
21.1
|
||||
28.3
|
||||
30.1
|
||||
31.3
|
||||
CMMLU (Acc.)
|
||||
5-shot
|
||||
47.9
|
||||
57.9
|
||||
61.9
|
||||
63.4
|
||||
C-Eval (Acc.)
|
||||
5-shot
|
||||
46.9
|
||||
58.0
|
||||
62.7
|
||||
63.3
|
||||
AGIEval (Acc.)
|
||||
0-shot
|
||||
29.1
|
||||
38.6
|
||||
41.8
|
||||
45.9
|
||||
ARC-Easy (Acc.)
|
||||
25-shot
|
||||
76.8
|
||||
86.5
|
||||
89.0
|
||||
90.1
|
||||
ARC-Challenge (Acc.)
|
||||
25-shot
|
||||
59.3
|
||||
70.1
|
||||
73.8
|
||||
76.4
|
||||
TriviaQA (EM)
|
||||
5-shot
|
||||
33.0
|
||||
48.8
|
||||
50.7
|
||||
51.8
|
||||
TriviaQA-ZH (EM)
|
||||
5-shot
|
||||
62.8
|
||||
74.8
|
||||
76.3
|
||||
77.9
|
||||
PopQA (EM)
|
||||
15-shot
|
||||
15.1
|
||||
19.2
|
||||
19.4
|
||||
21.2
|
||||
CCPM (Acc.)
|
||||
0-shot
|
||||
72.2
|
||||
79.6
|
||||
87.1
|
||||
87.7
|
||||
BBH (EM)
|
||||
3-shot
|
||||
42.8
|
||||
50.9
|
||||
55.9
|
||||
57.5
|
||||
HellaSwag (Acc.)
|
||||
0-shot
|
||||
64.3
|
||||
71.8
|
||||
72.7
|
||||
73.1
|
||||
PIQA (Acc.)
|
||||
0-shot
|
||||
63.8
|
||||
71.9
|
||||
73.5
|
||||
76.5
|
||||
WinoGrande (Acc.)
|
||||
5-shot
|
||||
64.0
|
||||
67.6
|
||||
67.8
|
||||
68.1
|
||||
Reading
|
||||
Comprehension
|
||||
DROP (F1)
|
||||
1-shot
|
||||
41.6
|
||||
55.7
|
||||
59.0
|
||||
60.7
|
||||
RACE-Middle (Acc.)
|
||||
5-shot
|
||||
72.4
|
||||
80.9
|
||||
82.8
|
||||
83.3
|
||||
RACE-High (Acc.)
|
||||
5-shot
|
||||
66.0
|
||||
75.4
|
||||
78.2
|
||||
79.2
|
||||
C3 (Acc.)
|
||||
0-shot
|
||||
57.7
|
||||
60.1
|
||||
63.6
|
||||
61.8
|
||||
Code & Math
|
||||
HumanEval (Pass@1)
|
||||
0-shot
|
||||
26.8
|
||||
37.8
|
||||
40.8
|
||||
38.4
|
||||
MBPP (Pass@1)
|
||||
3-shot
|
||||
35.4
|
||||
46.6
|
||||
48.2
|
||||
46.2
|
||||
CruxEval-i (EM)
|
||||
0-shot
|
||||
27.6
|
||||
30.7
|
||||
32.2
|
||||
36.2
|
||||
CruxEval-o (EM)
|
||||
0-shot
|
||||
28.7
|
||||
34.1
|
||||
35.0
|
||||
35.3
|
||||
GSM8K (EM)
|
||||
8-shot
|
||||
35.5
|
||||
58.4
|
||||
60.6
|
||||
62.6
|
||||
MGSM (EM)
|
||||
8-shot
|
||||
27.0
|
||||
46.8
|
||||
49.4
|
||||
52.4
|
||||
MATH (EM)
|
||||
4-shot
|
||||
15.2
|
||||
28.3
|
||||
30.7
|
||||
30.6
|
||||
With the proposed Engram architecture and the empirically derived allocation law, we scale Engram to the multi-billion parameter to validate its efficacy in real-world language model pre-training. Specifically, we train four models:
|
||||
(1) Dense-4B (4.1B total parameters),
|
||||
(2) MoE-27B (26.7B total parameters),
|
||||
(3) Engram-27B (26.7B total parameters), and
|
||||
(4) Engram-40B (39.5B total parameters).
|
||||
All models are trained using an identical data curriculum (same token budget and order) and are strictly matched in the number of activated parameters.
|
||||
## 4.1 Experimental Setup
|
||||
## Training Data and Model Configurations
|
||||
All models are pre-trained on a corpus of 262 billion tokens and we utilize the tokenizer from DeepSeek-v3 ( Liu et al., 2024a ) with a vocabulary size of 128k. For modeling, to ensure a controlled comparison, we adhere to a consistent default setting across all models unless explicitly stated otherwise. We utilize a 30-block Transformer with a hidden size of 2560. Each block integrates a Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024 ) with 32 heads, connected to FFNs via mHC ( Xie et al., 2025 ) with an expansion rate of 4. All models are optimized using Muon ( Jordan et al., 2024 ; Liu et al., 2025b ) ; detailed hyperparameters are listed in the Appendix A . We instantiate four distinct models:
|
||||
•
|
||||
Dense-4B serves as the baseline model. It utilizes the backbone architecture described above, incorporating a standard dense FFN into every block.
|
||||
•
|
||||
MoE-27B replaces the standard dense FFN with a DeepSeekMoE module ( Dai et al., 2024 ) . Configured with 72 routed experts and 2 shared experts (activating the top- k = 6 k=6 routed experts per token), this model scales to 26.7B total parameters while maintaining the same activated parameters as Dense-4B.
|
||||
•
|
||||
Engram-27B is strictly derived from the MoE-27B architecture to ensure fair comparison. We reduce the number of routed experts from 72 to 55 and reallocate the freed parameters to a 5.7B-parameter embedding module ( ρ = 74.3 % \rho=74.3\% ), keeping the total model size constant at 26.7B. Regarding the Engram configuration, we instantiate the module at layers 2 and 15 and set the maximum N N -gram size to 3, the number of heads to 8, and the dimension to 1280. For optimization, the embedding parameters are updated using Adam ( Kingma, 2014 ) with a learning rate scaled by 5 × 5\times and no weight decay, while the convolution parameters are initialized to zero to strictly preserve the identity mapping at the start of training.
|
||||
•
|
||||
Engram-40B retains the same backbone and computation budget as Engram-27B but scales the sparse embedding module to 18.5B parameters (totaling 39.5B parameters). This model is designed to investigate the scaling properties of Engram.
|
||||
## Evaluation Protocol
|
||||
We evaluate models on a diverse suite of benchmarks spanning language modeling, knowledge, reasoning, reading comprehension, and code/math.
|
||||
For each benchmark, we follow standard prompting protocols and evaluation metrics.
|
||||
•
|
||||
Language Modeling:
|
||||
We report loss on the test set of The Pile ( Gao et al., 2020 ) and an validation set drawn from the same distribution as the training data.
|
||||
•
|
||||
Knowledge & Reasoning:
|
||||
MMLU ( Hendrycks et al., 2021a ) , MMLU-Redux ( Gema et al., 2025 ) , MMLU-Pro ( Wang et al., 2024b ) , CMMLU ( Li et al., 2024 ) , C-Eval ( Huang et al., 2023 ) , AGIEval ( Zhong et al., 2024 ) , ARC-Easy/Challenge ( Clark et al., 2018 ) , TriviaQA ( Joshi et al., 2017 ) , TriviaQA-ZH (internal), PopQA ( Mallen et al., 2023 ) , CCPM ( Li et al., 2021 ) , BBH ( Suzgun et al., 2023 ) , HellaSwag ( Zellers et al., 2019 ) , PIQA ( Bisk et al., 2020 ) , and WinoGrande ( Sakaguchi et al., 2021 ) .
|
||||
•
|
||||
Reading Comprehension:
|
||||
DROP ( Dua et al., 2019 ) , RACE (Middle/High) ( Lai et al., 2017 ) , and C3 ( Sun et al., 2020 ) .
|
||||
•
|
||||
Code & Math:
|
||||
HumanEval ( Chen et al., 2021 ) , MBPP ( Austin et al., 2021 ) , CruxEval ( Gu et al., 2024 ) , GSM8K ( Cobbe et al., 2021 ) , MGSM ( Shi et al., 2023 ) , and MATH ( Hendrycks et al., 2021b ) .
|
||||
## 4.2 Experimental Results
|
||||
Table 1 summarizes the main results.
|
||||
First, consistent with prior literature ( Shazeer et al., 2017 ; He, 2024 ; Borgeaud et al., 2022 ) , sparse architectures demonstrate superior scaling laws compared to dense models. Under the same training compute budget, all three sparse variants (MoE-27B, Engram-27B/40B) significantly outperform the iso-FLOPs Dense-4B baseline across all benchmarks.
|
||||
More importantly, Engram-27B consistently improves over the iso-parameter and iso-FLOPs MoE-27B baseline. Interestingly, these gains are not limited to knowledge-intensive tasks (e.g., MMLU: +3.0, MMLU-Pro: +1.8, CMMLU: +4.0), where memory capacity is intuitively beneficial. We observe even more significant improvements in general-reasoning domains (e.g., BBH: +5.0, ARC-Challenge: +3.7, DROP: +3.3), as well as code and mathematical reasoning (e.g., HumanEval: +3.0, MBPP: +1.6, GSM8K: +2.2, MATH: +2.4). To reduce the impact of benchmark noise and to visualize training dynamics, we provide full benchmark trajectories during pre-training in Appendix B .
|
||||
These results support our hypothesis that introducing a dedicated knowledge lookup primitive improves representation efficiency beyond what can be achieved by allocating the entire sparse budget to conditional computation.
|
||||
Finally, scaling to Engram-40B further reduces pre-training loss and improves performance across most benchmarks. Although it does not yet strictly dominate Engram-27B on every task, this is likely an artifact of under-training. We observe that the training loss gap between Engram-40B and the baselines continues to widen towards the end of training, suggesting that the expanded memory capacity has not yet fully saturated within the current token budget.
|
||||
Table 2 :
|
||||
Long-context performance comparison. Parenthetical values (e.g. (50k, 1.62) ) denote the pre-training steps and the corresponding loss prior to the long-context extension. Two key findings: (1) With only 82% of the pre-training FLOPs (41k vs. 50k), Engram-27B matches the baseline’s LongPPL ( Fang et al., ) performance while achieving significantly higher accuracy on RULER ( Hsieh et al., ) ; (2) Under both iso-pretraining-loss (46k) and iso-pretraining-FLOPs (50k) settings, Engram-27B substantially outperforms the baseline across all metrics.
|
||||
Bold indicates the best and underline the second.
|
||||
Model
|
||||
LongPPL (32k)
|
||||
RULER (32k)
|
||||
Perplexity ( ↓ \downarrow )
|
||||
NIAH Accuracy ( ↑ \uparrow )
|
||||
Other Tasks ( ↑ \uparrow )
|
||||
Book
|
||||
Paper
|
||||
Code
|
||||
L-CoT
|
||||
S
|
||||
MK
|
||||
MV
|
||||
MQ
|
||||
VT
|
||||
CWE
|
||||
FWE
|
||||
QA
|
||||
MoE-27B (50k, 1.63)
|
||||
4.38
|
||||
2.91
|
||||
2.49
|
||||
14.16
|
||||
100.0
|
||||
88.0
|
||||
92.7
|
||||
84.2
|
||||
77.0
|
||||
4.5
|
||||
73.0
|
||||
34.5
|
||||
Engram-27B (41k, 1.66)
|
||||
4.37
|
||||
2.92
|
||||
2.50
|
||||
14.26
|
||||
99.6
|
||||
88.3
|
||||
93.0
|
||||
89.5
|
||||
83.2
|
||||
3.8
|
||||
99.6
|
||||
44.0
|
||||
Engram-27B (46k, 1.63)
|
||||
4.19
|
||||
2.84
|
||||
2.45
|
||||
13.59
|
||||
97.6
|
||||
89.0
|
||||
95.5
|
||||
97.0
|
||||
87.2
|
||||
4.3
|
||||
98.6
|
||||
37.5
|
||||
Engram-27B (50k, 1.62)
|
||||
4.14
|
||||
2.82
|
||||
2.44
|
||||
13.41
|
||||
99.3
|
||||
89.3
|
||||
96.5
|
||||
97.0
|
||||
89.0
|
||||
5.9
|
||||
99.3
|
||||
40.5
|
||||
## 5 Long Context Training
|
||||
By offloading local dependency modeling to static lookups, the Engram architecture preserves valuable attention capacity for managing global context. In this section, we empirically verify this structural advantage by conducting long-context extension training ( Gao et al., 2025 ; Peng et al., 2024 ) . Through a rigorous evaluation protocol that isolates architectural contributions from base model capabilities, we demonstrate that Engram yields significant gains in long-range retrieval and reasoning tasks.
|
||||
## 5.1 Experimental Setup
|
||||
## Training Details.
|
||||
To enable long-context capabilities, we adopt the context expansion strategy introduced in DeepSeek-V3 ( Liu et al., 2024a ) . Following the pre-training stage, we apply YaRN ( Peng et al., 2024 ) for context window extension in a 32768-token context training stage for 5,000 steps (30B tokens of high-quality, long-context data). The hyper-parameters are scale s = 10 , α = 1 , β = 32 s=10,\alpha=1,\beta=32 and the scaling factor f = 0.707 f=0.707 .
|
||||
## Model Configurations.
|
||||
We compare context extensions across four distinct model configurations. We utilize the final pre-training checkpoints (50k steps) for both MoE-27B and Engram-27B. Additionally, to rigorously benchmark architectural efficiency, we select two intermediate checkpoints for Engram-27B at 41k and 46k steps.
|
||||
Despite differing initialization stages, all variants undergo the exact same context extension training protocol.
|
||||
Crucially, Engram-27B (46k) is selected because it exhibits the same pre-training loss as the fully trained MoE-27B (50k). This creates a controlled "Iso-Loss" setting, ensuring that any performance divergence during context extension is attributable to the architecture rather than the starting quality of the model.
|
||||
## Evaluation Benchmarks.
|
||||
We assess long-context performance using LongPPL ( Fang et al., ) and RULER ( Hsieh et al., ) . For LongPPL, we construct evaluation sets spanning four categories: long books, research papers, code repositories, and long chain-of-thought (CoT) trajectories. For RULER, we evaluate on 14 subsets aggregated into 8 categories: Single (S), Multi-keys (MK), Multi-values (MV) and Multi-queries (MQ) Needle-in-a-Haystack; Multi-hop Variable Tracking (VT), Common Words Extraction (CWE), Frequent Words Extraction (FWE), and Question Answering (QA).
|
||||
## 5.2 Experimental Results
|
||||
The evaluation results are summarized in Table 2 . To accurately assess the contribution of the Engram architecture, our analysis proceeds in two steps: first, decoupling the impact of base model capability from architectural design, and second, conducting a controlled analysis.
|
||||
1. Long-Context Capability Beyond Attention Mechanics.
|
||||
While attention mechanisms and positional encoding provide the structural basis for context processing ( Su et al., 2024 ; Press et al., 2022 ; Yang et al., 2025 ; Xiao et al., 2024 ) , our results indicate that long-context performance is not solely determined by architectural priors. Observing the trajectory of Engram (41k → \rightarrow 50k), we find that long-context performance improves monotonically with pre-training progression, even when controlling for identical model architecture and a fixed computational budget during the context extension stage.
|
||||
This suggests that long-context performance is intrinsically coupled with the general modeling ability of the base model. Consequently, a rigorous architectural comparison must control for this confounding variable by aligning base model loss, rather than merely aligning training steps.
|
||||
2. Architectural Superiority under Controlled Settings.
|
||||
Guided by the principle above, we benchmark Engram against the MoE baseline. When controlling for base capability, the efficiency gains of the Engram module become evident:
|
||||
•
|
||||
Iso-Loss Setting (46k vs. Baseline): This setting strictly isolates architectural efficiency. When comparing Engram-27B (46k) against the fully trained MoE-27B (50k)—models aligned on pre-training loss—Engram demonstrates significant gains. Specifically, it outperforms the baseline on complex retrieval tasks (e.g., Multi-Query NIAH: 97.0 97.0 vs. 84.2 84.2 ; VT: 87.2 87.2 vs. 77.0 77.0 ).
|
||||
•
|
||||
Iso-FLOPs Setting (50k vs. Baseline): Under the standard iso-compute budget, Engram-27B (50k) further widens this gap, establishing the highest performance across the board.
|
||||
•
|
||||
Extreme Setting ( ≈ 82 % \approx 82\% Compute): Even the early-stopped Engram-27B (41k) remains highly competitive against the fully trained MoE-27B (50k). It matches the baseline on LongPPL and surpasses it on RULER, underscoring the intrinsic superiority of the Engram architecture.
|
||||
Figure 4 : Analysis of representational alignment and convergence speed. (a) Layer-wise KL Divergence via LogitLens ( nostalgebraist, 2020 ) . The consistently lower divergence in early layers indicates that Engram accelerates prediction convergence. (b-c) Similarity heatmap computed by CKA ( Kornblith et al., 2019 ) . The distinct upward shift of the high-similarity diagonal demonstrates that =Engram’s shallow layers are functionally equivalent to deeper layers of the MoE model, effectively increasing the model’s depth.
|
||||
Table 3:
|
||||
Entity resolution example reproduced from Ghandeharioun et al. (2024) .
|
||||
This table illustrates how LLMs gradually integrate context tokens through layers of attention and FFNs to construct the internal representation of the entity: ‘‘Diana, Princess of Wales’’ . The “Latent State Translation” column displays the automatically generated text for the last token: ‘‘Wales’’ by PatchScope ( Ghandeharioun et al., 2024 ) , while the “Explanation” column presents the manual interpretation provided by the original authors.
|
||||
Layer
|
||||
Latent State Translation
|
||||
Explanation
|
||||
1-2
|
||||
: Country in the United Kingdom
|
||||
Wales
|
||||
3
|
||||
: Country in Europe
|
||||
Wales
|
||||
4
|
||||
: Title held by female sovereigns in their own right or by queens consort
|
||||
Princess of Wales
|
||||
(unspecific)
|
||||
5
|
||||
: Title given to the wife of the Prince of Wales (and later King)
|
||||
Princess of Wales
|
||||
(unspecific)
|
||||
6
|
||||
: Diana, Princess of Wales (1961-1997), the first wife of Prince Charles, Prince of Wales, who was famous for her beauty and humanitarian work
|
||||
Diana,
|
||||
Princess of Wales
|
||||
## 6 Analysis
|
||||
In this section, we investigate the internal mechanisms of Engram, including its effective depth ( Section 6.1 ), core module design ( Section 6.2 ), and parametric sensitivity ( Section 6.3 ). Additionally, we evaluate the inference throughput with offloading ( Section 6.4 ) and conclude with a case study ( Section 6.5 ).
|
||||
## 6.1 Is Engram functionally equivalent to increasing the model’s depth?
|
||||
Current LLMs lack a dedicated knowledge lookup primitive and they rely on computation to simulate memory recall. As shown in Table 3 , to recognize the entity "Diana, Princess of Wales" , an LLM must consume multiple layers of Attention and FFNs to progressively compose features ( Li and Subramani, 2025 ; Ghandeharioun et al., 2024 ; Jin et al., 2025 ) , a process that could theoretically be identified via a knowledge lookup operation.
|
||||
Given this, we posit that by equipping the model with an explicit knowledge lookup capability, Engram effectively mimics an increase in model depth by relieving the model of the early stages of feature composition. To validate this hypothesis, we employ two mechanistic interpretability tools: LogitLens ( nostalgebraist, 2020 ; Belrose et al., 2023 ) and Centered Kernel Alignment analysis (CKA) ( Kornblith et al., 2019 ; Davari et al., 2023 ) .
|
||||
## 6.1.1 Accelerated Prediction Convergence
|
||||
We first analyze the evolution of predictions across layers using LogitLens ( nostalgebraist, 2020 ) . By projecting each intermediate layer’s hidden state with the final LM Head, we compute the Kullback–Leibler divergence ( Kullback and Leibler, 1951 ) between the intermediate output distribution and the model’s final output distribution. This metric quantifies how close a latent representation is to being “prediction-ready” ( Csordás et al., 2025 ; Belrose et al., 2023 ) .
|
||||
Figure 4 (a) reports the layer-wise KL divergence. Compared to the MoE baseline, both Engram variants exhibit systematically smaller KL divergence, with the most pronounced gap appearing in the early blocks. The steeper descent in the Engram curves indicates that the model finishes feature composition much faster. This observation aligns with our hypothesis: by accessing external knowledge explicitly, Engram reduces the computational steps required, thereby reaching high-confidence, valid predictions earlier in the network hierarchy.
|
||||
## 6.1.2 Representational Alignment and Effective Depth
|
||||
To further investigate whether Engram layers semantically correspond to deeper layers of the baseline, we employ Centered Kernel Alignment (CKA), a widely established metric for comparing representational structures ( Kornblith et al., 2019 ; Kriegeskorte et al., 2008 ) . Given two sets of representations X X and Y Y (e.g., activations from different models or layers), CKA is defined as:
|
||||
CKA ( K , L ) = HSIC ( K , L ) HSIC ( K , K ) HSIC ( L , L ) \text{CKA}(K,L)=\frac{\text{HSIC}(K,L)}{\sqrt{\text{HSIC}(K,K)\text{HSIC}(L,L)}}
|
||||
(8)
|
||||
where K = X X ⊤ K=XX^{\top} and L = Y Y ⊤ L=YY^{\top} denote the Gram matrices (using a linear kernel) and HSIC is Hilbert-Schmidt Independence Criterion ( Gretton et al., 2005 ) . We employ a minibatch implementation with an unbiased estimator of HSIC ( Davari et al., 2023 ) and evaluate on the Few-NERD dataset ( Ding et al., 2021 ) , extracting hidden states corresponding to the final token of named entities.
|
||||
To rigorously quantify the layer-wise correspondence,
|
||||
we first compute the pairwise CKA similarity matrix S ∈ [ 0 , 1 ] L × L S\in[0,1]^{L\times L} , where L L is the number of layers. We then introduce a soft alignment index a j a_{j} , defined as the weighted centroid of the top- k k most similar MoE layers for each Engram layer j j :
|
||||
a j = ∑ i ∈ ℐ j S i , j ⋅ i ∑ i ∈ ℐ j S i , j , where ℐ j = argtop k i ( S i , j ) . a_{j}=\frac{\sum_{i\in\mathcal{I}_{j}}S_{i,j}\cdot i}{\sum_{i\in\mathcal{I}_{j}}S_{i,j}},\quad\text{where }\mathcal{I}_{j}=\mathop{\text{argtop}k}_{i}(S_{i,j}).
|
||||
(9)
|
||||
Here, S i , j S_{i,j} denotes the similarity score between MoE layer i i and Engram layer j j . The index a j a_{j} serves as a robust proxy for the “effective MoE depth” corresponding to Engram layer j j , utilizing top- k k filtering (with k = 5 k=5 ) to mitigate low-similarity noise.
|
||||
Figure 4 (b)–(c) visualize the similarity heatmaps overlayed with the soft alignment curve (dashed white line). We observe a distinct upward shift from the diagonal, meaning that a j > j a_{j}>j for a wide range of layers. For instance, the representations formed at layer 5 of Engram-27B align most closely with those at approximately layer 12 of the MoE baseline.
|
||||
The consistent off-diagonal shift, which aligns with the LogitLens results ( Section 6.1.1 ), confirms that Engram achieves deeper representations at earlier layers. This validates our central hypothesis: by bypassing early-stage feature composition via explicit lookups, Engram is functionally equivalent to increasing the model’s effective depth.
|
||||
Figure 5 :
|
||||
Architecture ablation results.
|
||||
We compare the 3B MoE baseline against Engram variations in two settings:
|
||||
(1) Layer Sensitivity (dark blue curve): Sweeping the insertion depth of a single Engram module confirms that early injection (Layer 2) is optimal, whereas efficacy degrades in deeper layers.
|
||||
(2) Component Ablation (Right Markers): Removing sub-modules from the reference configuration demonstrates the importance of multi-branch integration, tokenizer compression, and context-aware gating.
|
||||
## 6.2 Structural Ablation and Layer Sensitivity
|
||||
In this section, we ablate Engram under a controlled setting to investigate the effectiveness of each key module design.
|
||||
Unless otherwise specified, the backbone is a 12-layer 3B MoE model (0.56B activated parameters) trained for 100B tokens.
|
||||
Figure 5 reports validation loss.
|
||||
The dashed orange line denotes the 3B MoE baseline (Val Loss = 1.808 =1.808 ).
|
||||
## Reference configuration.
|
||||
We augment the backbone with a fixed 1.6B-parameter Engram memory.
|
||||
Our reference model uses { 2 , 3 } \{2,3\} -grams and inserts Engram at Layers 2 and 6, achieving Val Loss = 1.768 =1.768 ,
|
||||
a substantial improvement over the MoE baseline ( Δ = 0.04 \Delta=0.04 ).
|
||||
All structural ablations below are defined relative to this reference.
|
||||
## Where should memory be injected?
|
||||
To study depth sensitivity, we keep the Engram budget fixed (1.6B) but consolidate it into a single Engram module,
|
||||
and sweep its insertion layer from 1 to 12 (dark blue “Layer Sweep” curve in Figure 5 ).
|
||||
This experiment exposes an inherent trade-off in Engram placement.
|
||||
A placement trade-off.
|
||||
Injecting Engram early allows it to offload local pattern reconstruction before the backbone expends computational depth, aligning with the backbone’s natural hierarchical processing ( Tenney et al., 2019 ; Ghandeharioun et al., 2024 ; Li and Subramani, 2025 ; Jin et al., 2025 ) .
|
||||
However, this incurs a cost in gating precision: early hidden states have not yet aggregated sufficient global context via attention, and the parallel branches lack the representational divergence required for fine-grained modulation ( Xie et al., 2025 ; Zhu et al., 2025 ) .
|
||||
Consequently, optimal placement requires balancing (i) offloading static local patterns early and (ii) utilizing stronger contextual queries for gating later.
|
||||
The sweep shows that Layer 2 achieves the best single-layer performance (Val Loss = 1.770 =1.770 ), outperforming Layer 1 and
|
||||
degrading as the insertion point moves deeper.
|
||||
This indicates that one round of attention is already sufficient to provide a meaningfully contextualized 𝐡 t \mathbf{h}_{t} for gating,
|
||||
while still being early enough to replace the backbone’s bottom-layer local aggregation.
|
||||
While Layer 2 is optimal under a single injection constraint, we find that dividing the same 1.6B memory into two smaller
|
||||
modules (achieved by reducing the embedding dimension d mem d_{\text{mem}} ) and placing them at Layers 2 and 6 performs even better (Val Loss = 1.768 =1.768 ).
|
||||
This layered design reconciles the trade-off by combining early intervention with rich, late-stage contextual gating.
|
||||
More importantly, layered insertion also provides a practical system advantage, enabling better utilization of the memory hierarchy as discussed in Section 2.5 .
|
||||
## Which components matter?
|
||||
Starting from the reference configuration, we ablate individual design choices while
|
||||
keeping the Engram parameter budget fixed. Results are denoted by markers in Figure 5 .
|
||||
We find that three components yield the most significant gains:
|
||||
(i) branch-specific fusion within the multi-branch backbone,
|
||||
(ii) context-aware gating, and
|
||||
(iii) tokenizer compression.
|
||||
Removing any of these causes the largest regressions in validation loss.
|
||||
Specifically, for the “w/o multi branch” ablation, we retain the mHC backbone structure but replace the branch-specific gating with a single Engram fusion applied to the hidden states after the pre-mapping ℋ p r e \mathcal{H}^{pre} ( Xie et al., 2025 ) .
|
||||
Other changes have smaller effects: removing the lightweight depthwise convolution only marginally degrades performance.
|
||||
Allocating capacity to 4-grams is slightly suboptimal under a fixed 1.6B budget—likely because it dilutes capacity
|
||||
from the more frequent 2/3-gram patterns—though we do not rule out that higher-order N N -grams become beneficial at larger memory scales.
|
||||
## 6.3 Sensitivity Analysis
|
||||
Figure 6 : Retained performance under Engram ablation. Factual knowledge relies heavily on the Engram module, whereas reading comprehension is largely preserved by the backbone.
|
||||
To characterize the functional contribution of the Engram module, we evaluate the model by completely suppressing the sparse embedding output during inference while keeping the backbone unchanged. Crucially, this post-hoc ablation induces a training–inference inconsistency , potentially introducing noise in complex, mixed-capability tasks. Consequently, we prioritize the analysis of Factual Knowledge and Reading Comprehension —the two extremes of the sensitivity spectrum—which exhibit the highest signal-to-noise ratio under this stress test.
|
||||
As shown in Figure 6 , the results reveal a sharp functional dichotomy. Factual knowledge benchmarks suffer a catastrophic collapse, retaining only 29–44% of the original performance (e.g., TriviaQA at 29%), confirming that the Engram module acts as the primary repository for parametric knowledge. Conversely, reading comprehension tasks are remarkably resilient, retaining 81–93% (e.g., C3 at 93%), suggesting that context-grounded tasks rely primarily on the backbone’s attention mechanism rather than Engram.
|
||||
## 6.4 System Efficiency
|
||||
A pivotal system advantage of Engram over routing-based MoE is that its sparse activations are addressed by explicit, static hash IDs. This yields a strictly deterministic memory access pattern: indices for the next Engram lookup are fixed once the token sequence is known and can be computed before the corresponding layer executes.
|
||||
## Experimental Setup.
|
||||
We implemented an inference harness based on nano-vLLM 1 1
|
||||
1
|
||||
https://github.com/GeeeekExplorer/nano-vllm —a streamlined prototype of the industry-standard vLLM engine ( Kwon et al., 2023 ) . To obtain a clean latency baseline without the confounding communication patterns of Expert Parallel in MoE, we benchmark on two dense backbones (Dense-4B and Dense-8B).
|
||||
We insert a massive 100B-parameter Engram layer into the second Transformer block, with the entire embedding table resident in host DRAM.
|
||||
During inference, the system prefetches embeddings for the Engram layer asynchronously, overlapping the PCIe transfer with the computation of the first block.
|
||||
Table 4 : End-to-end Inference Throughput . We measure infernece throughput with a 100B-parameter Engram layer entirely offloaded to host memory.
|
||||
Experimental Setup
|
||||
Hardware
|
||||
NVIDIA H800
|
||||
Workload
|
||||
512 Sequences
|
||||
Sequence Length
|
||||
Uniform ( 100 , 1024 ) (100,1024)
|
||||
Throughput Results
|
||||
Base Model
|
||||
Configuration
|
||||
Throughput (tok/s)
|
||||
4B-Dense
|
||||
Baseline
|
||||
9,031.62
|
||||
+ 100B Engram (CPU Offload)
|
||||
8,858.28
|
||||
8B-Dense
|
||||
Baseline
|
||||
6,315.52
|
||||
+ 100B Engram (CPU Offload)
|
||||
6,140.02
|
||||
## Results.
|
||||
As detailed in Table 4 , offloading a 100B-parameter embedding table incurs a negligible throughput penalty, peaking at only 2.8 % 2.8\% on the 8B backbone.
|
||||
This confirms that the compute intensity of early dense blocks provides a sufficient temporal window to mask the retrieval latency. Crucially, the effective communication volume per step scales with the number of activated slots rather than the total embedding table size.
|
||||
Crucially, this experiment serves as a conservative baseline. While the hierarchical design in Section 2.5 exploits Zipfian locality to cache frequent items in HBM, our experimental setup forces all retrievals to traverse the PCIe bus from host memory. The fact that this baseline retrieval strategy yields minimal overhead strongly suggests that a fully optimized, locality-aware implementation would incur negligible throughput penalty.
|
||||
## 6.5 Case Study: Gating Visualization
|
||||
In Section 2.3 , we introduced the context-aware gating mechanism, designed to dynamically modulate the integration of retrieved static memory into the backbone. To empirically validate whether Engram behaves as intended, we visualize the gating scalar α t \alpha_{t} of Engram-27B 2 2
|
||||
2
|
||||
As detailed in our architecture setup, this model utilizes a mHC ( M = 4 M=4 ) with Engram modules inserted at layers 2 and 15. Consequently, for any given token, the model computes a total of 8 distinct gating scalars. We observe that not every branch encodes interpretable activation patterns. For the clarity of this visualization, we select and display the gating values most strongly correlated with semantic pattern matching. across various samples in Figure 7 .
|
||||
The results demonstrate a distinct pattern of selectivity. The gating mechanism consistently activates (shown in red) upon completing local, static patterns. In English, we observe strong activations on multi-token named entities (e.g., “Alexander the Great”, “the Milky Way”) and formulaic phrases (e.g., “By the way”, “Princess of Wales”). This behavior generalizes effectively across languages. In the Chinese examples, Engram identifies and retrieves distinct idiomatic expressions and historical entities, such as “Four Great Inventions” ( 四大发明 ) and “Zhang Zhongjing” ( 张仲景 ). These qualitative results confirm that Engram successfully identifies and handles stereotyped linguistic dependencies, effectively relieving the Transformer backbone from memorizing these static associations.
|
||||
Figure 7 :
|
||||
Visualization of the gating mechanism of Engram. The heatmap intensity corresponds to the magnitude of the gating scalar α t ∈ [ 0 , 1 ] \alpha_{t}\in[0,1] , where darker red indicates stronger activation. Because Engram operates on suffix N N -grams (here N = 3 N=3 ), a high activation on a specific token x t x_{t} implies that the preceding tokens culminating in that token (e.g., the phrase ending at t t ) are recognized as a static pattern effectively retrieved from memory.
|
||||
## 7 Related Work
|
||||
## N N -gram Modeling and Embedding Scaling.
|
||||
Originating from Shannon’s framework ( Shannon, 1948 ) , N N -gram models rely on local history to predict tokens, traditionally employing smoothing techniques ( Kneser and Ney, 1995 ; Katz, 1987 ) to mitigate data sparsity. Despite the paradigm shift toward neural architectures ( Bengio et al., 2003 ) for capturing long-range dependencies, the computational efficiency of N N -gram lookups has been preserved in modern representation learning, as exemplified by seminal works like FastText ( Bojanowski et al., 2017 ) .
|
||||
Recently, this paradigm has resurged as embedding scaling . While architectures such as Per-Layer Embeddings ( Team, 2025 ) , STEM ( Sadhukhan et al., 2026 ) , L3 ( Tseng and Sa, 2026 ) and DeepEmbed ( RWKV Team, 2025 ) expand capacity via massive tables, a distinct line of pioneering research—most relevant to our approach—integrates compositional N N -gram structures directly into the representation space.
|
||||
N-Grammer ( Roy et al., 2022 ) directly augments the transformer architecture with N N -gram that are constructed from a discrete latent text representation, while Feng et al. (2023) demonstrates the effectiveness of N N -gram for speech recognition.
|
||||
SuperBPE ( Liu et al., 2025a ) and SCONE ( Yu et al., 2025 ) explicitly target high-frequency patterns: the former by merging multi-word expressions into “superword” tokens, and the latter via an auxiliary encoding model.
|
||||
In parallel, OverEncoding ( Huang et al., 2025a ) and Byte Latent Transformer (BLT) ( Pagnoni et al., 2025 ) adopt hash N N -gram embeddings to capture local dependencies at the token and byte levels, respectively. These studies collectively demonstrate the efficacy of scaling parameters through N N -gram representations with minimal computational overhead. While these approaches offer significant gains in their respective settings, our work diverges fundamentally in two key dimensions.
|
||||
•
|
||||
First, regarding modeling and evaluation protocols. Prior approaches often treat N N -gram embeddings as external augmentations without validating their efficiency under strictly fair comparison protocols. For instance, SCONE ( Yu et al., 2025 ) is inference-focused and relies on auxiliary modules that incur addition
|
||||
|
||||
=== BODY (tail) ===
|
||||
al training FLOPs. Similarly, OverEncoding ( Huang et al., 2025a ) fails to yield meaningful improvements on sparse MoE backbones even under a non-isoparametric setting. In contrast, we treat conditional memory as a first-class modeling primitive instantiated via the carefully designed Engram module. By rigorously evaluating this design within our Sparsity Allocation framework, we demonstrate its clear advantage over strictly iso-parameter and iso-FLOPs MoE baselines.
|
||||
•
|
||||
Second, from a system perspective, we advocate for algorithm-system co-design. Existing approaches place embeddings strictly at the input layer (Layer 0), which inherently serializes memory access and computation ( Yu et al., 2025 ; Huang et al., 2025a ) . Engram, conversely, strategically injects memory into deeper layers to enable communication-computation overlap. Furthermore, by exploiting the inherent Zipfian distribution of N N -grams, we could maximize the utility of the hardware memory hierarchy. This holistic design allows Engram to scale to massive parameters with negligible inference overhead.
|
||||
## High-Cardinality Categorical Embeddings.
|
||||
A closely related representation problem arises in large-scale recommender systems, where models learn embeddings for hundreds of categorical features whose vocabularies may contain millions to billions of IDs ( Coleman et al., 2023 ) .
|
||||
Prior work improves the parameter–accuracy and memory-access trade-offs through multi-hash and compositional representations ( Tito Svenstrup et al., 2017 ; Shi et al., 2020 ) ,
|
||||
frequency-aware hashing and collision management
|
||||
( Zhang et al., 2020 ; Tsang and Ahle, 2022 ; Zhang et al., 2024 ) ,
|
||||
frequency- or importance-adaptive capacity allocation
|
||||
( Ginart et al., 2021 ; Joglekar et al., 2020 ; Liu et al., 2021 ) ,
|
||||
and structured compression based on quantization, tensor decomposition,
|
||||
or shared memory layouts
|
||||
( Kang et al., 2020 ; Yin et al., 2021 ; Desai et al., 2022 ) .
|
||||
Other approaches share representations across multiple categorical fields
|
||||
( Coleman et al., 2023 ) , generate embeddings without explicit tables
|
||||
( Kang et al., 2021 ) , or replace arbitrary item IDs with learned
|
||||
semantic identifiers
|
||||
( Rajput et al., 2023 ; Singh et al., 2024 ) .
|
||||
Engram shares with this literature the challenge of representing an
|
||||
extremely large discrete key space under a highly skewed access
|
||||
distribution, but differs in constructing keys from ordered textual
|
||||
N N -grams and injecting the retrieved representations into intermediate
|
||||
Transformer layers through context-aware gating.
|
||||
## Mixture-of-Experts.
|
||||
MoE architectures decouple model capacity from computational cost by conditionally activating a sparse subset of experts per token, a paradigm introduced by Shazeer et al. (2017) . Subsequent innovations such as GShard ( Lepikhin et al., 2020 ) , BASE ( Lewis et al., 2021 ) , Switch Transformer ( Fedus et al., 2022 ) and GLaM ( Du et al., 2022 ) enabled super-linear parameter scaling while maintaining constant inference costs. More recently, DeepSeek-MoE ( Dai et al., 2024 ) demonstrated superior efficiency, significantly outperforming dense models with equivalent active parameters via fine-grained expert segmentation and shared expert isolation. Adopting this architecture, state-of-the-art models such as DeepSeek-V3 ( Liu et al., 2024a ) and Kimi-k2 ( Kimi, 2025 ) have further pushed total parameters to hundreds of billions scale.
|
||||
## Memory Network.
|
||||
Research on memory-augmented networks aims to expand model capacity without a proportional increase in computational cost, broadly categorized into parametric and non-parametric approaches. Parametric memory methods, such as PKM ( Lample et al., 2019 ) , PEER ( He, 2024 ) , Selfmem ( Cheng et al., 2023b ) , Memory+ ( Berges et al., 2025 ) and UltraMem ( Huang et al., 2025c ; Huang et al., 2025b ) , integrate large-scale, sparse key-value stores directly into the model layers, thereby significantly increasing capacity with negligible impact on FLOPs.
|
||||
Conversely, non-parametric memory approaches like REALM ( Guu et al., 2020 ) , RETRO ( Borgeaud et al., 2022 ; Wang et al., 2023 ) , CoG ( Lan et al., 2023 ; Cao et al., ) and PlugLM ( Cheng et al., 2023a ) decouple knowledge storage from model processing, treating the external memory as an editable and scalable key-value store that allows the model to adapt to evolving information without expensive retraining.
|
||||
## Mechanisms of Knowledge Storage.
|
||||
Parallel to capacity scaling, substantial research has scrutinized the internal mechanisms governing how Transformers encode and retrieve factual knowledge. The Feed-Forward Networks (FFNs) are widely hypothesized to function as Key-Value memories ( Geva et al., 2021 ) . Under this framework, the first layer acts as a pattern detector ("keys") while the second layer projects specific information into the residual stream ("values"). This modularity is evidenced by the identification of specific “knowledge neurons” responsible for storing distinct facts ( Dai et al., 2022 ) . Further validation is provided by causal tracing methodologies, which map the information flow of factual recall to specific FFN layers ( Meng et al., 2022 ) . These insights have enabled precise model editing algorithms such as ROME ( Meng et al., 2022 ) and MEMIT ( Meng et al., 2023 ) , which allow for the direct update of factual associations without retraining. Moreover, investigations into internal representations, such as those in Othello-GPT ( Li et al., 2023a ) , suggest that these storage mechanisms may facilitate the emergence of structured “world models” rather than mere statistical memorization.
|
||||
## 8 Conclusion
|
||||
In this work, we introduce conditional memory as a complementary sparsity axis to the prevailing conditional computation paradigm (MoE), aiming to resolve the inefficiency of simulating knowledge retrieval through dynamic computation. We instantiate this concept via Engram, a module that modernizes classic N N -gram embeddings to enable scalable, constant-time O ( 1 ) O(1) lookups for static patterns.
|
||||
By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law, demonstrating that a hybrid allocation of sparse capacity between MoE experts and Engram memory strictly outperforms pure MoE baselines. Guided by this law, we scale Engram to 27B parameters, achieving superior performance across diverse domains. Notably, while the memory module intuitively aids knowledge retrieval, we observe even larger gains in general reasoning, code, and mathematics.
|
||||
Our mechanistic analysis reveals that Engram effectively “deepen” the network by relieving early layers from static reconstruction tasks, thereby freeing up attention capacity to focus on global context and complex reasoning. This architectural shift translates into substantial improvements in long-context capabilities, as evidenced by performance gains in LongPPL and RULER. Finally, Engram advocates for infrastructure-aware efficiency as a first-class design principle. Its deterministic addressing allows for the decoupling of storage and compute, enabling the offloading of massive parameter tables to host memory with negligible inference overhead. We envision conditional memory functions as an indispensable modeling primitive for next-generation sparse models.
|
||||
@@ -0,0 +1,526 @@
|
||||
TITLE: DeepSeek-OCR 2: Visual Causal Flow
|
||||
ARXIV: 2601.20552
|
||||
AUTHORS: Haoran Wei , Yaofeng Sun , Yukun Li
|
||||
DATE: [Submitted on 28 Jan 2026]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-OCR 2: Visual Causal Flow
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Related Works
|
||||
## 2.1 Parallelized Queries in Decoder
|
||||
## 2.2 Parallelized Queries in Projector
|
||||
## 2.3 LLM-based Multimodal Initialization
|
||||
## 3 Methodology
|
||||
## 3.1 Architecture
|
||||
## 3.2 DeepEncoder V2
|
||||
## 3.2.1 Vision tokenizer
|
||||
## 3.2.2 Language model as vision encoder
|
||||
## 3.2.3 Causal flow query
|
||||
## 3.2.4 Attention mask
|
||||
## 3.3 DeepSeek-MoE Decoder
|
||||
## 4 Experimental Settings
|
||||
## 4.1 Data Engine
|
||||
## 4.2 Training Pipelines
|
||||
## 4.2.1 Training DeepEncoder V2
|
||||
## 4.2.2 Query enhancement
|
||||
## 4.2.3 Continue-training LLM
|
||||
## 5 Evaluation
|
||||
## 5.1 Main Results
|
||||
## 5.2 Improvement Headroom
|
||||
## 5.3 Practical Readiness
|
||||
## 6 Discussion and Future Works
|
||||
## 6.1 Towards Genuine 2D Reasoning
|
||||
## 6.2 Towards Native Multimodality
|
||||
## 7 Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-OCR 2: Visual Causal Flow
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Parallelized Queries in Decoder
|
||||
2.2 Parallelized Queries in Projector
|
||||
2.3 LLM-based Multimodal Initialization
|
||||
3 Methodology
|
||||
3.1 Architecture
|
||||
3.2 DeepEncoder V2
|
||||
3.2.1 Vision tokenizer
|
||||
3.2.2 Language model as vision encoder
|
||||
3.2.3 Causal flow query
|
||||
3.2.4 Attention mask
|
||||
3.3 DeepSeek-MoE Decoder
|
||||
4 Experimental Settings
|
||||
4.1 Data Engine
|
||||
4.2 Training Pipelines
|
||||
4.2.1 Training DeepEncoder V2
|
||||
4.2.2 Query enhancement
|
||||
4.2.3 Continue-training LLM
|
||||
5 Evaluation
|
||||
5.1 Main Results
|
||||
5.2 Improvement Headroom
|
||||
5.3 Practical Readiness
|
||||
6 Discussion and Future Works
|
||||
6.1 Towards Genuine 2D Reasoning
|
||||
6.2 Towards Native Multimodality
|
||||
7 Conclusion
|
||||
References
|
||||
License: CC BY 4.0
|
||||
arXiv:2601.20552v1 [cs.CV] 28 Jan 2026
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-OCR 2: Visual Causal Flow
|
||||
Haoran Wei
|
||||
Yaofeng Sun
|
||||
Yukun Li
|
||||
Affiliation: DeepSeek-AI
|
||||
## Abstract
|
||||
We present DeepSeek-OCR 2 to investigate the feasibility of a novel encoder—DeepEncoder V2—capable of dynamically reordering visual tokens upon image semantics. Conventional vision-language models (VLMs) invariably process visual tokens in a rigid raster-scan order (top-left to bottom-right) with fixed positional encoding when fed into LLMs. However, this contradicts human visual perception, which follows flexible yet semantically coherent scanning patterns driven by inherent logical structures. Particularly for images with complex layouts, human vision exhibits causally-informed sequential processing. Inspired by this cognitive mechanism, DeepEncoder V2 is designed to endow the encoder with causal reasoning capabilities, enabling it to intelligently reorder visual tokens prior to LLM-based content interpretation. This work explores a novel paradigm: whether 2D image understanding can be effectively achieved through two-cascaded 1D causal reasoning structures, thereby offering a new architectural approach with the potential to achieve genuine 2D reasoning. Codes and model weights are publicly accessible at http://github.com/deepseek-ai/DeepSeek-OCR-2 .
|
||||
Figure 1 : We substitute the CLIP component in DeepEncoder with an LLM-style architecture. By customizing the attention mask, visual tokens utilize bidirectional attention while learnable queries adopt causal attention. Each query token can thus attend to all visual tokens and preceding queries, allowing progressive causal reordering over visual information.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Related Works
|
||||
2.1 Parallelized Queries in Decoder
|
||||
2.2 Parallelized Queries in Projector
|
||||
2.3 LLM-based Multimodal Initialization
|
||||
3 Methodology
|
||||
3.1 Architecture
|
||||
3.2 DeepEncoder V2
|
||||
3.2.1 Vision tokenizer
|
||||
3.2.2 Language model as vision encoder
|
||||
3.2.3 Causal flow query
|
||||
3.2.4 Attention mask
|
||||
3.3 DeepSeek-MoE Decoder
|
||||
4 Experimental Settings
|
||||
4.1 Data Engine
|
||||
4.2 Training Pipelines
|
||||
4.2.1 Training DeepEncoder V2
|
||||
4.2.2 Query enhancement
|
||||
4.2.3 Continue-training LLM
|
||||
5 Evaluation
|
||||
5.1 Main Results
|
||||
5.2 Improvement Headroom
|
||||
5.3 Practical Readiness
|
||||
6 Discussion and Future Works
|
||||
6.1 Towards Genuine 2D Reasoning
|
||||
6.2 Towards Native Multimodality
|
||||
7 Conclusion
|
||||
References
|
||||
## 1 Introduction
|
||||
The human visual system closely mirrors transformer-based vision encoders [ 12 , 10 ] : foveal fixations function as visual tokens, locally sharp yet globally aware. However, unlike existing encoders that rigidly scan tokens from top-left to bottom-right, human vision follows a causally-driven flow guided by semantic understanding. Consider tracing a spiral—our eye movements follow inherent logic where each subsequent fixation causally depends on previous ones. By analogy, visual tokens in models should be selectively processed with ordering highly contingent on visual semantics rather than spatial coordinates.
|
||||
This insight motivates us to fundamentally reconsider the architectural design of vision-language models (VLMs), particularly the encoder component. LLMs are inherently trained on 1D sequential data, while images are 2D structures. Directly flattening image patches in a predefined raster-scan order introduces unwarranted inductive bias that ignores semantic relationships. To address this, we present DeepSeek-OCR 2 with a novel encoder design—DeepEncoder V2—to advance toward more human-like visual encoding. Following DeepSeek-OCR [ 53 ] , we adopt document reading as our primary experimental testbed. Documents present rich challenges including complex layout orders, intricate formulas, and tables. These structured elements inherently carry causal visual logic, demanding sophisticated reasoning capabilities that make document OCR an ideal platform for validating our approach.
|
||||
Our main contributions are threefold:
|
||||
First, we present DeepEncoder V2, featuring several key innovations: (1) we replace the CLIP [ 37 ] component in DeepEncoder [ 53 ] with a compact LLM [ 48 ] architecture, as illustrated in Figure 1 , to achieve visual causal flow; (2) to enable parallelized processing, we introduce learnable queries [ 6 ] , termed causal flow tokens, with visual tokens prepended as a prefix—through a customized attention mask, visual tokens maintain global receptive fields, while causal flow tokens can obtain visual token reordering ability; (3) we maintain equal cardinality between causal and visual tokens (with redundancy such as padding and borders) to provide sufficient capacity for re-fixation; (4) only the causal flow tokens—the latter half of the encoder outputs—are fed to the LLM [ 20 ] decoder, enabling cascade causal-aware visual understanding.
|
||||
Second, leveraging DeepEncoder V2, we present DeepSeek-OCR 2, which preserves the image compression ratio and decoding efficiency of DeepSeek-OCR while achieving substantial performance improvements. We constrain visual tokens fed to the LLM between 256 and 1120. The lower bound (256) corresponds to DeepSeek-OCR’s tokenization of 1024×1024 images, while the upper bound (1120) matches Gemini-3 pro’s [ 44 ] maximum visual token budget. This design positions DeepSeek-OCR 2 as both a novel VLM architecture for research exploration and a practical tool for generating high-quality training data for LLM pretraining.
|
||||
Finally, we provide preliminary validation for employing language model architectures as VLM encoders—a promising pathway toward unified omni-modal encoding. This framework enables feature extraction and token compression across diverse modalities (images, audio, text [ 24 ] ) by simply configuring modality-specific learnable queries. Crucially, it naturally succeeds to advanced infrastructure optimizations from the LLM community, including Mixture-of-Experts (MoE) architectures, efficient attention mechanisms [ 22 ] , and so on.
|
||||
In summary, we propose DeepEncoder V2 for DeepSeek-OCR 2, employing specialized attention mechanisms to effectively model the causal visual flow of document reading. Compared to the DeepSeek-OCR baseline, DeepSeek-OCR 2 achieves 3.73% performance gains on OmniDocBench v1.5 [ 34 ] and yields considerable advances in visual reading logic.
|
||||
Figure 2 : This figure shows two computer vision models with parallelized queries: DETR’s decoder [ 6 ] for object detection and BLIP2’s Q-former [ 18 ] for visual token compression. Both employ bidirectional self-attention among queries.
|
||||
## 2 Related Works
|
||||
## 2.1 Parallelized Queries in Decoder
|
||||
DETR [ 6 ] pioneered the integration of transformer architecture into object detection, fundamentally breaking away from traditional detection paradigms [ 40 , 38 ] . To overcome the efficiency limitations of serial decoding in transformer blocks, DETR introduced preset parallelized learnable queries—a set of 100 object queries that encode object priors such as shape and position through training. These queries interact with feature maps [ 14 ]
|
||||
via cross-attention mechanisms, while simultaneously engaging in bidirectional information exchange among themselves through self-attention.
|
||||
DETR established a foundational paradigm that enables transformers to handle parallelized tokens.
|
||||
The object query design has since become the de facto standard architectural component in subsequent transformer-based detection methods [ 56 , 25 ] .
|
||||
## 2.2 Parallelized Queries in Projector
|
||||
In recent years, vision-language models [ 18 , 3 , 5 , 50 ] have developed rapidly, with architectures converging toward the encoder-projector-LLM paradigm. The projector aligns visual tokens with the LLM’s embedding space, serving as a critical bridge that enables LLMs to understand visual content. Q-former, introduced in BLIP-2 [ 18 ] , exemplifies an effective projector design that employs learnable queries for visual token compression. Adopting a BERT-like [ 11 ] architecture and drawing inspiration from DETR’s object queries [ 6 ] , Q-former utilizes 32 learnable queries that interact with hundreds of CLIP [ 37 ] visual tokens through cross-attention. These compressed query representations are subsequently fed into the LLM, achieving effective mapping from visual to language space. The success of Q-former demonstrates that parallelized learnable queries are effective not only for feature decoding in detection tasks but also for token compression in multimodal alignment.
|
||||
## 2.3 LLM-based Multimodal Initialization
|
||||
LLMs trained on large-scale internet data have proven effective as initialization for multimodal models. Pang et al. [ 35 ] demonstrated that frozen LLM transformer layers enhance visual discriminative tasks. Moreover, encoder-free or lightweight-encoder models such as Fuyu [ 1 ] and Chameleon [ 43 ] in vision, as well as VALL-E [ 47 ] in speech, further validate the potential of LLM pretrained weights for
|
||||
multimodal initialization.
|
||||
Figure 3 : DeepSeek-OCR 2 adopts the visual token compression mechanism from DeepEncoder, employing an 80M-parameter image compressor that reduces visual tokens by a factor of 16. DeepEncoder V2 differs by replacing DeepEncoder’s CLIP module with a compact language model architecture. Through customized attention masks, this LM-style vision encoder acquires CLIP’s knowledge compression capabilities while initiating causal modeling of visual sequences.
|
||||
## 3 Methodology
|
||||
## 3.1 Architecture
|
||||
As shown in Figure 3 , DeepSeek-OCR 2 inherits the overall architecture of DeepSeek-OCR, which consists of an encoder and a decoder. The encoder discretizes images into visual tokens, while the decoder generates outputs conditioned on these visual tokens and text prompts. The key distinction lies in the encoder: we upgrade DeepEncoder to DeepEncoder V2, which retains all capabilities of its predecessor while introducing causal reasoning through a novel architectural design. We elaborate on the details of DeepSeek-OCR 2 in the following sections.
|
||||
## 3.2 DeepEncoder V2
|
||||
The vanilla encoder serves as an important component that extracts and compresses image features through attention mechanisms, where each token attends to all others, achieving full-image receptive fields analogous to human foveal and peripheral vision. However, flattening 2D image patches into a 1D sequence imposes a rigid ordering bias through text-oriented positional encodings (e.g., RoPE [ 42 ] ). This contradicts natural visual reading patterns, especially non-linear layouts in optical texts, forms and tables.
|
||||
Figure 4 : Token count calculation in DeepEncoder V2. DeepEncoder V2 outputs 256 − - 1120 tokens per image using a multi-crop strategy with 0 − - 6 local views. With 0 local views, only the global view produces 256 tokens; with 6 local views, the count reaches 1120 (6 × \times 144+256).
|
||||
## 3.2.1 Vision tokenizer
|
||||
The first component of DeepEncoder V2 is a vision tokenizer. Following DeepEncoder, we employ an architecture combining an 80M-parameter SAM-base [ 17 ] along with two convolutional layers [ 50 ] . The output dimension of the final convolutional layer is reduced from 1024 in DeepEncoder to 896 to align with the subsequent pipeline. Note that this compression-based tokenizer is not mandatory and can be replaced with simple patch embedding. We retain it because it achieves 16 × \times token compression [ 52 , 45 , 51 , 15 ] through window attention with minimal parameters, significantly reducing both computational cost and activation memory for the subsequent global attention module. Moreover, its parameter count ( 80M) remains comparable to the typical 100M parameters used for text input embeddings in LLMs.
|
||||
## 3.2.2 Language model as vision encoder
|
||||
In DeepEncoder, a CLIP ViT follows the vision tokenizer to compress visual knowledge. DeepEncoder V2 redesigns this component into an LLM-style architecture with a dual-stream attention mechanism. Visual tokens utilize bidirectional attention to preserve CLIP’s global modeling capability, while newly introduced causal flow queries employ causal attention. These learnable queries are appended after visual tokens as a suffix, where each query attends to all visual tokens and preceding queries. By maintaining equal cardinality between queries and visual tokens, this design imposes semantic ordering and distilling on visual features without altering token count. Finally, only the causal query outputs are fed to the LLM decoder.
|
||||
We instantiate this architecture using Qwen2-0.5B [ 48 ] , whose 500M parameters are comparable to CLIP ViT (300M) without introducing excessive computational overhead. The decoder-only architecture with prefix-concatenation of visual tokens proves crucial: extra experiments with cross-attention in an mBART-style [ 26 ] encoder-decoder structure fail to converge. We hypothesize this failure stems from insufficient visual token interaction when isolated in a separate encoder. In contrast, the prefix design keeps visual tokens active throughout all layers, fostering effective visual information exchange with causal queries.
|
||||
This architecture actually establishes two-stage cascade causal reasoning: the encoder semantically reorders visual tokens through learnable queries, while the LLM decoder performs autoregressive reasoning over the ordered sequence. Unlike vanilla encoders that impose rigid spatial ordering through positional encodings, our causally-ordered queries adapt to smooth visual semantics while naturally aligning with the LLM’s unidirectional attention pattern. This design may bridge the gap between 2D spatial structure and 1D causal language modeling.
|
||||
Figure 5 : Attention mask architecture of DeepEncoder V2. Concatenation of bidirectional mask (vision tokens, ViT-like) and causal triangular mask (flow tokens, LLM decoder-style).
|
||||
## 3.2.3 Causal flow query
|
||||
As aforementioned, the number of causal query tokens equals the number of visual tokens, computed as W × H 16 2 × 16 \frac{W\times H}{16^{2}\times 16} , where W W and H H denote the width and height of the image input to the encoder. To avoid maintaining multiple query sets for different resolutions, we adopt a multi-crop strategy with fixed query configurations at predefined resolutions.
|
||||
Specifically, the global view uses a resolution of 1024 × 1024 1024\times 1024 , corresponding to 256 query embeddings denoted as query global \text{query}_{\text{global}} . Local crops adopt a resolution of 768 × 768 768\times 768 , with the number of crops k k ranging from 0 to 6 (no cropping is applied when both image dimensions are smaller than 768). All local views share a unified set of 144 query embeddings, denoted as query local \text{query}_{\text{local}} . Therefore, the total number of reordered visual tokens fed to the LLM is k × 144 + 256 k\times 144+256 , ranging from [ 256 , 1120 ] [256,1120] . This maximum token count (1120) is lower than DeepSeek-OCR’s 1156 (Gundam mode) and matches Gemini-3-Pro’s maximum visual token budget.
|
||||
## 3.2.4 Attention mask
|
||||
To better illustrate the attention mechanism of DeepEncoder V2, we visualize the attention mask in Figure 5 . The attention mask is composed of two distinct regions. The left region applies bidirectional attention (similar to ViT) to original visual tokens, allowing full token-to-token visibility. The right region employs causal attention (triangular mask, identical to decoder-only LLMs) for causal flow tokens, where each token attends only to previous tokens. These two components are concatenated along the sequence dimension to construct DeepEncoder V2’s attention mask (M), as follows:
|
||||
M = [ 𝟏 m × m 𝟎 m × n 𝟏 n × m LowerTri ( n ) ] , where n = m M=\begin{bmatrix}\mathbf{1}_{m\times m}&\mathbf{0}_{m\times n}\\
|
||||
\mathbf{1}_{n\times m}&\text{LowerTri}(n)\end{bmatrix},\quad\text{where }n=m
|
||||
(1)
|
||||
where n n is the number of causal query tokens, m m represents vanilla visual tokens number, and LowerTri denotes a lower triangular matrix (with ones on and below the diagonal, zeros above).
|
||||
## 3.3 DeepSeek-MoE Decoder
|
||||
Since DeepSeek-OCR 2 primarily focuses on encoder improvements, we do not upgrade the decoder component. Following this design principle, we retain DeepSeek-OCR’s decoder − - a 3B-parameter MoE structure with about 500M active parameters. The core forward pass of DeepSeek-OCR 2 can be formulated as:
|
||||
𝐎 = 𝒟 ( π Q ( 𝒯 L ( ℰ ( 𝐈 ) ⊕ 𝐐 0 , 𝐌 ) ) ) \mathbf{O}=\mathcal{D}\left(\pi_{Q}\left(\mathcal{T}^{L}\left(\mathcal{E}(\mathbf{I})\oplus\mathbf{Q}_{0};\mathbf{M}\right)\right)\right)
|
||||
(2)
|
||||
where 𝐈 ∈ ℝ H × W × 3 \mathbf{I}\in\mathbb{R}^{H\times W\times 3} is the input image, ℰ \mathcal{E} is the vision tokenizer mapping images to m m visual tokens 𝐕 ∈ ℝ m × d \mathbf{V}\in\mathbb{R}^{m\times d} , 𝐐 0 ∈ ℝ n × d \mathbf{Q}_{0}\in\mathbb{R}^{n\times d} are learnable causal query embeddings, ⊕ \oplus denotes sequence concatenation, 𝒯 L \mathcal{T}^{L} represents an L L -layer Transformer with masked attention, 𝐌 ∈ { 0 , 1 } 2 n × 2 n \mathbf{M}\in\{0,1\}^{2n\times 2n} is the block causal attention mask defined in Equation 1 , π Q \pi_{Q} is the projection operator that extracts the last n n tokens (i.e., 𝐙 = 𝐗 m + 1 : m + n \mathbf{Z}=\mathbf{X}_{m+1:m+n} ), 𝒟 \mathcal{D} is the language decoder, and 𝐎 ∈ ℝ n × | 𝒱 | \mathbf{O}\in\mathbb{R}^{n\times|\mathcal{V}|} is the output logits over LLM vocabulary.
|
||||
## 4 Experimental Settings
|
||||
## 4.1 Data Engine
|
||||
DeepSeek-OCR 2 employs the same data sources as DeepSeek-OCR, comprising OCR 1.0, OCR 2.0 [ 7 , 54 , 23 ] , and general vision data [ 53 ] , with OCR data constituting 80% of the training mixture. We also introduce two modifications: (1) a more balanced sampling strategy for OCR 1.0 data, partitioning pages by content type (text, formulas, tables) with a 3:1:1 ratio, and (2) label refinement for layout detection by merging semantically similar categories (e.g., unifying "figure caption" and "figure title"). Given these minimal differences, we consider DeepSeek-OCR a valid baseline for comparison.
|
||||
## 4.2 Training Pipelines
|
||||
We train DeepSeek-OCR 2 in three stages: (1) encoder pretraining, (2) query enhancement, and (3) decoder specialization. The stage-1 enables the vision tokenizer and LLM-style encoder to acquire fundamental capabilities in feature extraction, token compression, and token reordering capabilities. The stage-2 further strengthens the token reordering capability of the encoder while enhancing visual knowledge compression. The stage-3 freezes the encoder parameters and optimizes only the decoder, enabling higher data throughput under the same FLOPs.
|
||||
## 4.2.1 Training DeepEncoder V2
|
||||
Following DeepSeek-OCR and Vary [ 50 ] , we train DeepEncoder V2 using a language modeling objective, coupling the encoder with a lightweight decoder [ 16 ] for joint optimization via next token prediction. We employ two dataloaders at resolutions of 768 × \times 768 and 1024 × \times 1024. The vision tokenizer is initialized from DeepEncoder, and the LLM-like encoder from Qwen2-0.5B-base [ 48 ] . After pretraining, only the encoder parameters are retained for subsequent stages. We use the AdamW [ 28 ] optimizer with cosine learning rate decay from 1e-4 to 1e-6, training on 160 A100 GPUs (20 nodes × \times 8 GPUs) with batch size 640 for 40k iterations (with sequence packing at 8K length, about 100M image-text pair samples).
|
||||
## 4.2.2 Query enhancement
|
||||
After DeepEncoder V2 pretraining, we integrate it with DeepSeek-3B-A500M [ 20 , 21 ] as our final pipeline. We freeze the visual tokenizer (SAM-conv structure) while jointly optimizing the LLM encoder and LLM decoder to enhance query representations. At this stage, we unify the two resolutions into a single dataloader via multi-crop strategy. We adopt 4-stage pipeline parallelism: vision tokenizer (PP0), LLM-style encoder (PP1), and DeepSeek-LLM layers (6 layers per stage on PP2-3). With 160 GPUs (40GB/per-GPU), we configure 40 data parallel replicas (4 GPUs per replica) and train with global batch size 1280 using the same optimizer and learning rate decay from 5e-5 to 1e-6 over 15k iterations.
|
||||
## 4.2.3 Continue-training LLM
|
||||
To rapidly consume training data, we freeze all DeepEncoder V2 parameters in this stage and only update the DeepSeek-LLM parameters. This stage accelerates training (more than doubles the training speed under the same global batch size) while helping the LLM better understand DeepEncoder V2’s reordered visual tokens. Continuing from stage-2, we perform another learning rate decay from 1e-6 to 5e-8 training for 20k iterations in this stage.
|
||||
Table 1 : Comprehensive evaluation of document reading on OmniDocBench v1.5. V-token max represents the maximum number of visual tokens used per page in this benchmark. R-order denotes reading order. Except for DeepSeek OCR and DeepSeek OCR 2, all other model results in this table are sourced from the OmniDocBench repository.
|
||||
Model
|
||||
V-token max ↓ \downarrow
|
||||
Overall ↑ \uparrow
|
||||
Text Edit ↓ \downarrow
|
||||
Formula CDM ↑ \uparrow
|
||||
Table TEDs ↑ \uparrow
|
||||
Table T E D S s {}^{TEDS_{s}} ↑ \uparrow
|
||||
R-order Edit ↑ \uparrow
|
||||
Pipline
|
||||
Marker-1.8.2 [ 29 ]
|
||||
-
|
||||
71.30
|
||||
0.206
|
||||
76.66
|
||||
57.88
|
||||
71.17
|
||||
0.250
|
||||
MinerU2-pp [ 46 ]
|
||||
-
|
||||
71.51
|
||||
0.209
|
||||
76.55
|
||||
70.90
|
||||
79.11
|
||||
0.225
|
||||
Dolphin [ 13 ]
|
||||
-
|
||||
74.67
|
||||
0.125
|
||||
67.85
|
||||
68.70
|
||||
77.77
|
||||
0.124
|
||||
Dolphin-1.5 [ 13 ]
|
||||
-
|
||||
83.21
|
||||
0.092
|
||||
80.78
|
||||
78.06
|
||||
84.10
|
||||
0.080
|
||||
PP-StructureV3 [ 9 ]
|
||||
-
|
||||
86.73
|
||||
0.073
|
||||
85.79
|
||||
81.68
|
||||
89.48
|
||||
0.073
|
||||
MonkeyOCR-pro-1.2B [ 19 ]
|
||||
-
|
||||
86.96
|
||||
0.084
|
||||
85.02
|
||||
84.24
|
||||
89.02
|
||||
0.130
|
||||
MonkeyOCR-3B [ 19 ]
|
||||
-
|
||||
87.13
|
||||
0.075
|
||||
87.45
|
||||
81.39
|
||||
85.92
|
||||
0.129
|
||||
MonkeyOCR-pro-3B [ 19 ]
|
||||
-
|
||||
88.85
|
||||
0.075
|
||||
87.25
|
||||
86.78
|
||||
90.63
|
||||
0.128
|
||||
MinerU2.5 [ 46 ]
|
||||
-
|
||||
90.67
|
||||
0.047
|
||||
88.46
|
||||
88.22
|
||||
92.38
|
||||
0.044
|
||||
PaddleOCR-VL [ 8 ]
|
||||
-
|
||||
92.86
|
||||
0.035
|
||||
91.22
|
||||
90.89
|
||||
94.76
|
||||
0.043
|
||||
End-to-end Model
|
||||
OCRFlux [ 31 ]
|
||||
>6000
|
||||
74.82
|
||||
0.193
|
||||
68.03
|
||||
75.75
|
||||
80.23
|
||||
0.202
|
||||
GPT-4o [ 33 ]
|
||||
-
|
||||
75.02
|
||||
0.217
|
||||
79.70
|
||||
67.07
|
||||
76.09
|
||||
0.148
|
||||
InternVL3 [ 55 ]
|
||||
>7000
|
||||
80.33
|
||||
0.131
|
||||
83.42
|
||||
70.64
|
||||
77.74
|
||||
0.113
|
||||
POINTS-Reader [ 27 ]
|
||||
>6000
|
||||
80.98
|
||||
0.134
|
||||
79.20
|
||||
77.13
|
||||
81.66
|
||||
0.145
|
||||
olmOCR [ 36 ]
|
||||
>6000
|
||||
81.79
|
||||
0.096
|
||||
86.04
|
||||
68.92
|
||||
74.77
|
||||
0.121
|
||||
InternVL3.5-241B [ 49 ]
|
||||
>7000
|
||||
82.67
|
||||
0.142
|
||||
87.23
|
||||
75.00
|
||||
81.28
|
||||
0.125
|
||||
MinerU2-VLM [ 46 ]
|
||||
>7000
|
||||
85.56
|
||||
0.078
|
||||
80.95
|
||||
83.54
|
||||
87.66
|
||||
0.086
|
||||
Nanonets-OCR-s [ 30 ]
|
||||
>7000
|
||||
85.59
|
||||
0.093
|
||||
85.90
|
||||
80.14
|
||||
85.57
|
||||
0.108
|
||||
Qwen2.5-VL-72B [ 5 ]
|
||||
>6000
|
||||
87.02
|
||||
0.094
|
||||
88.27
|
||||
82.15
|
||||
86.22
|
||||
0.102
|
||||
Gemini-2.5 Pro [ 2 ]
|
||||
-
|
||||
88.03
|
||||
0.075
|
||||
85.82
|
||||
85.71
|
||||
90.29
|
||||
0.097
|
||||
dots.ocr [ 39 ]
|
||||
>6000
|
||||
88.41
|
||||
0.048
|
||||
83.22
|
||||
86.78
|
||||
90.62
|
||||
0.053
|
||||
OCRVerse [ 32 ]
|
||||
>6000
|
||||
88.56
|
||||
0.058
|
||||
86.91
|
||||
84.55
|
||||
88.45
|
||||
0.071
|
||||
Qwen3-VL-235B [ 4 ]
|
||||
>6000
|
||||
89.15
|
||||
0.069
|
||||
88.14
|
||||
86.21
|
||||
90.55
|
||||
0.068
|
||||
DeepSeek-OCR (9-crops)
|
||||
1156
|
||||
87.36
|
||||
0.073
|
||||
84.14
|
||||
85.25
|
||||
89.01
|
||||
0.085
|
||||
DeepSeek-OCR 2
|
||||
1120
|
||||
91.09
|
||||
0.048
|
||||
90.31
|
||||
87.75
|
||||
92.06
|
||||
0.057
|
||||
↓ 36 \downarrow 36
|
||||
↑ 3.73 \uparrow 3.73
|
||||
↓ 0.025 \downarrow 0.025
|
||||
↑ 6.17 \uparrow 6.17
|
||||
↑ 2.5 \uparrow 2.5
|
||||
↑ 3.05 \uparrow 3.05
|
||||
↓ 0.028 \downarrow 0.028
|
||||
## 5 Evaluation
|
||||
We select OmniDocBench v1.5 [ 34 ] as our primary benchmark for evaluation. This benchmark comprises 1,355 document pages spanning 9 major categories (including magazines, academic papers, research reports, and so on) in both Chinese and English. With its diverse test samples and robust evaluation criteria, OmniDocBench provides an effective framework for validating the performance of DeepSeek-OCR 2, particularly the effectiveness of DeepEncoder V2.
|
||||
## 5.1 Main Results
|
||||
As shown in Table 1 , DeepSeek-OCR 2 achieves advanced performance of 91.09% while using the smallest upper limit of visual tokens (V-token max ). Compared to the DeepSeek-OCR baseline, it demonstrates a 3.73% improvement under similar train data sources, validating the effectiveness of our newly designed architecture. Beyond the overall improvement, the Edit Distance (ED) for reading order (R-order) has also significantly decreased (from 0.085 to 0.057), indicating that the new DeepEncoder V2 can effectively select and arrange initial visual tokens based on image information. As illustrated in Table 2 , DeepSeek-OCR 2 (0.100) achieves lower ED in document parsing compared to Gemini-3 Pro (0.115) under a similar visual token budget (1120), further demonstrating that our new model maintains high compression rates of visual tokens while ensuring superior performance, with exceptionally high potential.
|
||||
Table 2: Edit Distances for different categories of document-elements in OmniDocBench v1.5. V-token max denotes the lowest maximum number of visual tokens.
|
||||
Model
|
||||
V-token max ↓ \downarrow
|
||||
Text Edit ↓ \downarrow
|
||||
Formula Edit ↓ \downarrow
|
||||
Table Edit ↓ \downarrow
|
||||
R-order Edit ↓ \downarrow
|
||||
Overall Edit ↓ \downarrow
|
||||
Gemini-3 pro [ 44 ]
|
||||
1120
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.115
|
||||
Seed-1.8 [ 41 ]
|
||||
5120
|
||||
-
|
||||
-
|
||||
-
|
||||
-
|
||||
0.106
|
||||
DeepSeek-OCR
|
||||
1156
|
||||
0.073
|
||||
0.236
|
||||
0.123
|
||||
0.085
|
||||
0.129
|
||||
DeepSeek-OCR 2
|
||||
1120
|
||||
0.048
|
||||
0.198
|
||||
0.096
|
||||
0.057
|
||||
0.100
|
||||
## 5.2 Improvement Headroom
|
||||
We conduct a detailed performance comparison between DeepSeek-OCR and DeepSeek-OCR 2 across 9 document types and found that DeepSeek-OCR 2 still has considerable room for improvement, as shown in Table 3 . For text recognition Edit Distance (ED), DeepSeek-OCR 2 outperforms DeepSeek-OCR in most cases, but there are also notable weaknesses, such as newspapers, where it performs > 0.13 >0.13 ED. We believe there are two main reasons: (1) the lower upper limit of visual tokens may affect the recognition of text-super-rich newspapers, which can be simply addressed in the future by increasing the number of local crops; (2) insufficient newspaper data − - our training data contains only 250k relevant samples, which is inadequate for training DeepEncoder V2 for this class. Of course, for the reading order (R-order) metric, DeepSeek-OCR 2 consistently outperforms DeepSeek-OCR across the board, further validating the effectiveness of our visual causal flow encoder design.
|
||||
Table 3: Detailed comparison between DeepSeek-OCR 2 and DeepSeek-OCR across 9 document types. R-order denotes reading order. All metrics are Edit Distances, where lower is better.
|
||||
Model
|
||||
Edit ↓ \downarrow
|
||||
PPT
|
||||
Academic
|
||||
Paper
|
||||
Book
|
||||
Colorful
|
||||
Textbook
|
||||
Exam
|
||||
Paper
|
||||
Magazine
|
||||
Newspaper
|
||||
Note
|
||||
Research
|
||||
Report
|
||||
DS-OCR
|
||||
Text
|
||||
0.052
|
||||
0.028
|
||||
0.022
|
||||
0.130
|
||||
0.074
|
||||
0.049
|
||||
0.131
|
||||
0.145
|
||||
0.015
|
||||
R-order
|
||||
0.052
|
||||
0.021
|
||||
0.040
|
||||
0.125
|
||||
0.083
|
||||
0.101
|
||||
0.217
|
||||
0.089
|
||||
0.016
|
||||
DS-OCR 2
|
||||
Text
|
||||
0.031
|
||||
0.013
|
||||
0.033
|
||||
0.053
|
||||
0.047
|
||||
0.026
|
||||
0.139
|
||||
0.068
|
||||
0.008
|
||||
R-order
|
||||
0.025
|
||||
0.013
|
||||
0.027
|
||||
0.066
|
||||
0.048
|
||||
0.100
|
||||
0.176
|
||||
0.035
|
||||
0.011
|
||||
Table 4: Production performance comparison between DeepSeek-OCR and DeepSeek-OCR 2. For OCR models serving LLM pipelines, ground truth is not accessible in production environments. Therefore, Repetition rate constitutes the primary observable quality metric.
|
||||
Model
|
||||
Metric
|
||||
online-user-logs (image)
|
||||
pretrain-data (PDF)
|
||||
DeepSeek-OCR
|
||||
Repeat ↓ \downarrow
|
||||
6.25%
|
||||
3.69%
|
||||
DeepSeek-OCR 2
|
||||
4.17% ↓ 2.08 % \downarrow 2.08\%
|
||||
2.88% ↓ 0.81 % \downarrow 0.81\%
|
||||
## 5.3 Practical Readiness
|
||||
DeepSeek-OCR serves two primary production use cases: an online OCR service that reads image/documents for DeepSeek-LLMs, and a pretraining data pipeline that performs batch PDF processing. We compare the production performance between DeepSeek-OCR 2 and DeepSeek-OCR. Since ground truth is unavailable in production environments, we focus primarily on repetition rate as our key metric. As shown in Table 4 , DeepSeek-OCR 2 demonstrates markedly improved practical readiness compared to its predecessor (DeepSeek-OCR), reducing the repetition rate from 6.25% to 4.17% for online user-log images, and from 3.69% to 2.88% for PDF data production. These results further validate the effectiveness of the DeepSeek-OCR 2 architecture, particularly its logical visual comprehension capabilities.
|
||||
## 6 Discussion and Future Works
|
||||
## 6.1 Towards Genuine 2D Reasoning
|
||||
DeepSeek-OCR 2 presents a novel architectural paradigm with an LLM-style encoder cascaded with an LLM decoder. This cascade of two 1D causal reasoners holds promise for genuine 2D reasoning: the encoder performs reading logic reasoning (causally reordering visual information through query tokens), while the decoder executes visual task reasoning over these causally-ordered representations. Decomposing 2D understanding into two complementary/orthogonal 1D causal reasoning subtasks may represent a breakthrough toward genuine 2D reasoning. Of course, achieving this goal remains a long journey. For instance, to enable multiple re-examinations and multi-hop reordering of visual content, we may need substantially longer causal flow tokens than the original visual token sequence. We will continue to refine this architecture and explore its effectiveness on general visual reasoning tasks in future work.
|
||||
## 6.2 Towards Native Multimodality
|
||||
DeepEncoder V2 provides initial validation of the LLM-style encoder’s viability for visual tasks. More importantly, this architecture enjoys the potential to evolve into a unified omni-modal encoder: a single encoder with shared W k , W v Wk,Wv projections, attention mechanisms, and FFNs can process multiple modalities through modality-specific learnable query embeddings. Such an encoder could compress text, extract speech features, and reorganize visual content within the same parameter space, differing only in the learned weights of their query embeddings. DeepSeek-OCR’s optical compression represents an initial exploration toward native multimodality, while we believe DeepSeek-OCR 2’s LLM-style encoder architecture marks our further step in this direction. We will also continue exploring the integration of additional modalities through this shared encoder framework in the future.
|
||||
## 7 Conclusion
|
||||
In this technical report, we present DeepSeek-OCR 2, a significant upgrade to DeepSeek-OCR, that maintains high visual token compression while achieving meaningfully performance improvements. This advancement is powered by the newly proposed DeepEncoder V2, which implicitly distills causal understanding of the visual world through the integration of both bidirectional and causal attention mechanisms, leading to causal reasoning capabilities in the vision encoder and, consequently, marked lifts in visual reading logic.
|
||||
While optical text reading, particularly document parsing, represents one of the most practical vision tasks in the LLM era, it constitutes only a small part of the broader visual understanding landscape. Looking ahead, we will refine and adapt this architecture to more diverse scenarios, seeking deeper toward a more comprehensive vision of multimodal intelligence.
|
||||
|
||||
=== BODY (tail) ===
|
||||
@@ -0,0 +1,556 @@
|
||||
TITLE: DualPath: Breaking the Storage Bandwidth Bottleneck in Agentic LLM Inference
|
||||
ARXIV: 2602.21548
|
||||
AUTHORS: Yongtong Wu , Shaoyuan Chen , Yinmin Zhong , Rilin Huang , Yixuan Tan , Wentao Zhang , Liyue Zhang , Shangyan Zhou , Yuxuan Liu , Shunfeng Zhou , Mingxing Zhang , Xin Jin , Panpan Huang
|
||||
DATE: [Submitted on 25 Feb 2026 ( v1 ), last revised 26 Feb 2026 (this version, v2)]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DualPath: Breaking the Storage Bandwidth
|
||||
## Abstract.
|
||||
## 1. Introduction
|
||||
## 2. Background
|
||||
## 2.1. LLM Inference Preliminary
|
||||
## 2.2. Agentic Use of LLMs
|
||||
## 2.3. Modern AI Data Center Architecture
|
||||
## 3. Bottleneck & Motivation
|
||||
## 4. DualPath System Overview
|
||||
## 4.1. Dual-Path Loading
|
||||
## 4.2. Bottleneck-Free Analysis
|
||||
## 4.3. Practical Challenges
|
||||
## 5. CNIC-Centric Traffic Manager
|
||||
## 5.1. Traffic Isolation
|
||||
## 5.2. CNIC-Assisted KV-Cache Copy
|
||||
## 6. Adaptive Request Scheduler
|
||||
## 6.1. Inter-Engine Scheduling
|
||||
## 6.2. Intra-Engine Scheduling
|
||||
## 7. Evaluation
|
||||
## 7.1. Implementation
|
||||
## 7.2. Experimental Setup
|
||||
## 7.3. Offline Batch Inference
|
||||
## 7.4. Online Serving
|
||||
## 7.5. Ablation Study
|
||||
## 7.6. Large-Scale Scalability
|
||||
## 8. Discussion
|
||||
## 8.1. Potential Future Work
|
||||
## 8.2. Working Set Analysis
|
||||
## 9. Related Work
|
||||
## 10. Conclusion
|
||||
|
||||
=== BODY (head) ===
|
||||
DualPath: Breaking the Storage Bandwidth Bottleneck in Agentic LLM Inference
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract.
|
||||
1 Introduction
|
||||
2 Background
|
||||
2.1 LLM Inference Preliminary
|
||||
2.2 Agentic Use of LLMs
|
||||
2.3 Modern AI Data Center Architecture
|
||||
3 Bottleneck & Motivation
|
||||
4 DualPath System Overview
|
||||
4.1 Dual-Path Loading
|
||||
4.2 Bottleneck-Free Analysis
|
||||
4.3 Practical Challenges
|
||||
5 CNIC-Centric Traffic Manager
|
||||
5.1 Traffic Isolation
|
||||
5.2 CNIC-Assisted KV-Cache Copy
|
||||
6 Adaptive Request Scheduler
|
||||
6.1 Inter-Engine Scheduling
|
||||
6.2 Intra-Engine Scheduling
|
||||
7 Evaluation
|
||||
7.1 Implementation
|
||||
7.2 Experimental Setup
|
||||
7.3 Offline Batch Inference
|
||||
7.4 Online Serving
|
||||
7.5 Ablation Study
|
||||
7.6 Large-Scale Scalability
|
||||
8 Discussion
|
||||
8.1 Potential Future Work
|
||||
8.2 Working Set Analysis
|
||||
9 Related Work
|
||||
10 Conclusion
|
||||
References
|
||||
A Appendix
|
||||
A.1 Traffic Isolation Configuration Details
|
||||
A.2 27B Model Specifications
|
||||
A.3 Agent Task Structure
|
||||
A.4 Experimental Configurations
|
||||
A.5 KV-Cache Block Layout
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2602.21548v2 [cs.DC] 26 Feb 2026
|
||||
## DualPath: Breaking the Storage Bandwidth
|
||||
Bottleneck in Agentic LLM Inference
|
||||
Yongtong Wu 1,3 Shaoyuan Chen 2,3 Yinmin Zhong 1,3 Rilin Huang 1
|
||||
Yixuan Tan 3 Wentao Zhang 3 Liyue Zhang 3 Shangyan Zhou 3 Yuxuan Liu 3 Shunfeng Zhou 3 Mingxing Zhang 2 Xin Jin 1 Panpan Huang 3
|
||||
Affiliation:
|
||||
1 School of Computer Science, Peking University 2 Tsinghua University 3 DeepSeek-AI
|
||||
## Abstract.
|
||||
The performance of multi-turn, agentic LLM inference is increasingly dominated by KV-Cache storage I/O rather than computation. In prevalent disaggregated architectures, loading the massive KV-Cache from external storage creates a fundamental imbalance: storage NICs on prefill engines become bandwidth-saturated, while those on decoding engines remain idle. This asymmetry severely constrains overall system throughput.
|
||||
We present DualPath, an inference system that breaks this bottleneck by introducing dual-path KV-Cache loading. Beyond the traditional storage-to-prefill path, DualPath enables a novel storage-to-decode path, in which the KV-Cache is loaded into decoding engines and then efficiently transferred to prefill engines via RDMA over the compute network. DualPath combines this optimized data path — which inherently avoids network congestion and avoids interference with latency-critical model execution communications — with a global scheduler that dynamically balances load across prefill and decode engines.
|
||||
Our evaluation on three models with realistic agentic workloads demonstrates that DualPath improves offline inference throughput by up to 1.87 × \times on our in-house inference system. It can also improve online serving throughput by an average factor of 1.96 × \times without violating SLO.
|
||||
## 1. Introduction
|
||||
Large Language Models (LLMs) are rapidly evolving from single-turn chatbots ( OpenAI, 2025b ; DeepSeek-AI, 2025d ) and standalone reasoners ( OpenAI, 2025b ) into agentic systems that can autonomously plan, invoke tools, and solve real-world tasks through multi-turn interactions ( Chowa et al., 2026 ; Wang et al., 2024 ; Xi et al., 2025 ; Jiang et al., 2024 ; Mohammadi et al., 2025 ) . In such settings, an LLM no longer serves isolated prompts; instead, it participates in long-running sessions where context accumulates over time ( Lin et al., 2025 ) . As agentic applications become increasingly prevalent, multi-turn LLM inference has emerged as a critical workload in production systems, ranging from coding assistants ( Yang et al., 2024 ; Wu et al., 2023 ) to autonomous task agents ( Zhou et al., 2023 ; Li et al., 2024 ) .
|
||||
This paradigm shift in applications has driven a significant transformation in LLM inference workloads: from traditional human-LLM interaction to human-LLM-environment interaction, called the agentic paradigm . The typical pattern of human-model interaction involves users providing input, engaging in a few rounds of interaction with the LLM, and consuming the results generated by the LLM. By contrast, an agentic LLM may interact with an external environment, through tools such as a web browser and Python interpreter, over dozens or even hundreds of turns. Although each individual tool call or feedback is short (often hundreds of tokens), the context accumulates across turns and can grow to extreme lengths. As a result, agentic workloads become highly I/O-bound: the multi-turn, short-append pattern leads to very high KV-Cache hit rates — typically ≥ 95 % \geq 95\% ( Chen et al., 2026 ) — making the efficiency of KV-Cache loading, rather than pure computation, the dominant performance factor.
|
||||
Figure 1 . Existing bottleneck (left) and DualPath (right).
|
||||
To improve throughput under agentic workloads, existing LLM inference systems have converged on a common set of architectural patterns: layer-wise prefill ( Xiong et al., 2024 ; Du et al., 2025 ) , prefill–decode (PD) disaggregation ( Zhong et al., 2024 ; Patel et al., 2025 ; Zhao et al., 2025a ) , and external KV-Cache storage ( Gao et al., 2024 ; Liu et al., 2025 ; Qin et al., 2025 ) .
|
||||
In these systems, prefill engines load the KV-Cache in a layer-wise manner to accommodate as many requests as possible within a single batch. When prefill completes, decoding engines typically receive KV-Cache from prefill engines via a high-performance RDMA network. The decoding engines then generate tokens and store their KV-Cache in distributed storage to enable reuse across turns.
|
||||
However, this architecture also introduces a critical limitation. As shown in Figure 1 , prefill engines must load large volumes of KV-Cache from remote storage.
|
||||
As a result, prefill-side storage network bandwidth becomes the throughput bottleneck of the entire system, even though decoding engines often have substantial unused storage network bandwidth.
|
||||
This imbalance reveals a fundamental inefficiency in existing designs: storage network bandwidth is unevenly utilized across engines. The bandwidth of prefill engines are persistently saturated, while decoding engines remain underutilized.
|
||||
Simply provisioning more bandwidth to prefill engines is costly and often impractical in general-purpose clusters. Therefore, it is promising to exploit and combine the available I/O bandwidth of all engines, rather than overloading prefill engines alone, to accelerate KV-Cache loading for agentic LLM workloads.
|
||||
Prior studies have attempted to alleviate the KV-Cache loading bottleneck.
|
||||
Mooncake ( Qin et al., 2025 ) caches KV-Cache in a distributed DRAM pool and employs an affinity-aware scheduler to maximize the DRAM KV-Cache hit rate.
|
||||
However, it cannot be used in memory-constrained scenarios, such as the rollout phase in RL, where DRAM is occupied to hold large training state that is offloaded from HBM.
|
||||
It is also not cost-effective in scenarios with enormous working sets (e.g., online serving), considering the cost comparison between DRAM and SSD.
|
||||
Other attempts reduce the amount of KV-Cache data to retrieve ( Gao et al., 2025 ) and reduce the retrieval overhead ( Hu et al., 2025 ; Yan et al., 2025 ) .
|
||||
However, they do not solve the inherent inefficiency caused by storage I/O imbalance between different engines.
|
||||
In this paper, we present DualPath, a new LLM inference system that rethinks KV-Cache loading in modern inference architectures for agentic workloads. The key insight behind DualPath is that KV-Cache loading does not have to be prefill-centric. While existing systems always load KV-Cache directly from storage into prefill engines, they cannot utilize the remote storage bandwidth of decoding engines. DualPath leverages this observation by enabling dual-path KV-Cache loading : in addition to the conventional storage-to-prefill path, KV-Cache can be loaded into decoding engines and then transferred to prefill engines via high-performance RDMA. By dynamically selecting between these paths, DualPath redistributes network load and alleviates prefill-side bandwidth pressure.
|
||||
Realizing this design raises two challenges.
|
||||
First, introducing an extra loading path introduces complex traffic patterns and potential interference with collective primitives in model execution, which can degrade overall performance if unmanaged.
|
||||
Second, the system must decide online which loading path to use under dynamic and heterogeneous workloads, and ensure load balance across both GPUs and NICs simultaneously.
|
||||
To address these challenges, DualPath adopts (1) an optimized dual-path loading data path design, which introduces no inherent congestion under common P/D ratios, (2) a NIC-centric traffic management approach to isolate KV-Cache traffic from latency-sensitive model inference communications, and (3) a dynamic scheduling policy that jointly balances computation and network utilization across prefill and decoding engines.
|
||||
We implement DualPath on top of a modern inference stack and evaluate it using representative agentic workloads with long contexts and high cache reuse.
|
||||
Experiments show that DualPath significantly improves system throughput and the first token latency, while maintaining the latency between tokens. In agentic inference scenarios, DualPath increases end-to-end throughput by up to 1.87 × \times for offline inference, and improves the online serving throughput by 1.96 × \times on average.
|
||||
In summary, this paper makes three contributions:
|
||||
•
|
||||
We identify the I/O-bound nature of multi-turn, agentic LLM workloads and show that KV-Cache loading dominates system performance under modern LLM inference architectures.
|
||||
•
|
||||
We present DualPath, an inference system that introduces dual-path KV-Cache loading and leverages decoding-engine bandwidth to resolve prefill-side bottlenecks.
|
||||
•
|
||||
We design and evaluate a workload-aware scheduling algorithm that dynamically balances computation and network resources, significantly improving balance on realistic workloads.
|
||||
## 2. Background
|
||||
## 2.1. LLM Inference Preliminary
|
||||
LLM inference is becoming one of the most important system workloads recently.
|
||||
Popular LLMs utilize decoder-only transformer architecture, comprising stacked blocks with attention layers and feed-forward networks (FFNs).
|
||||
Attention layers enable token interactions within requests, while FFNs process tokens independently.
|
||||
The model predicts subsequent tokens based on preceding ones, storing attention keys and values as KV-Cache in HBM to avoid recomputations.
|
||||
PD-disaggregated Inference. Prefill–decode (PD) disaggregation ( Zhong et al., 2024 ; Patel et al., 2025 ) separates the prefill phase from the decode phase, assigning them to dedicated prefill engines (PEs) and decode engines (DEs), respectively. The two phases exhibit distinct compute and memory patterns: prefill is compute-intensive and batched, while decode is memory-bound and latency-sensitive. With PD disaggregation, PEs load the hit KV-Cache and perform prefilling; then, they transfer the KV-Cache to DEs, which perform autoregressive decoding. This design reduces interference between phases, enables stage-specific optimizations, and improves scalability, making it the de facto architecture for modern LLM serving. To support multi-turn conversations, the KV-Cache is often stored in distributed storage for reuse across turns.
|
||||
Layerwise Prefill. Long-context prefill is bottlenecked by HBM capacity, as both activations and the KV-Cache for the entire batch must reside within it, forcing limited batch sizes and leading to poor GPU utilization. LayerKV ( Xiong et al., 2024 ) and PrefillOnly ( Du et al., 2025 ) address this problem by exploiting the strong locality in prefill computations: each layer requires only its own layer-specific KV-Cache. Consequently, the KV-Cache can be allocated and freed per layer, and the GPU holds only one layer’s KV-Cache for the forward batch. This increases the effective batch size (in tokens) by approximately a factor equal to the number of layers, boosting prefill throughput.
|
||||
## 2.2. Agentic Use of LLMs
|
||||
Figure 2 . Agent trajectory example.
|
||||
LLMs increasingly power agentic applications that perform multi-turn reasoning and interact with an environment (via e.g., terminal commands, code execution, or asking for human feedback) over long sessions.
|
||||
As shown in Figure 2 , in a typical turn , the model receives a prompt composed by the previous context plus some newly appended tokens (often tool output or user input) and generates the next action or response. A single agent run is a trajectory of dozens or even hundreds of turns: the context grows turn-by-turn and can reach up to one million tokens ( Anthropic, 2026 ; DeepMind, 2026 ) .
|
||||
Because most of the context, typically >95% tokens in our traces, is reused across rounds, the vast majority of tokens in each round can hit the KV-Cache; only the newly appended context needs prefill computation. Due to the extreme length of agent trajectories, DRAM and HBM-based KV-Cache storage like Mooncake ( Qin et al., 2025 ) can only store a small proportion of KV-Caches, necessitating the use of larger yet cheaper external SSD-based KV-Cache storage ( DeepSeek-AI, 2025a ) .
|
||||
The agentic LLM inference workload is also prevalent in agent LLM training, which often adopts reinforcement learning (RL) approaches. In a typical RL training loop, the agent LLM first undergoes a rollout phase, where it is prompted to generate a large number of multi-step agent trajectories. These trajectories are then scored by a separate reward model. Finally, the LLM parameters are updated to increase the likelihood of high-scoring outputs and reduce the likelihood of low-scoring ones. During the rollout phase, substantial data (like reward model and optimizer states) is offloaded to host DRAM, further constraining the available DRAM for KV-Cache. This reinforces the need for external, high-capacity KV-Cache storage that can accommodate long agentic rollout contexts efficiently.
|
||||
## 2.3. Modern AI Data Center Architecture
|
||||
Modern AI data centers are purpose-built logical supercomputers engineered to handle large-scale generative AI training and inference workloads. For example, in a standard NVIDIA DGX SuperPOD ( NVIDIA, 2023 ) , each node is equipped with 8 Hopper GPUs interconnected via high-speed NVLink. Each GPU is paired with a dedicated 400 Gbps compute NIC ( CNIC , also known as east-west NIC), which maximizes inter-node communication bandwidth. Independent of the compute fabric, each node also features a storage NIC ( SNIC , also known as south-north NIC) up to 400 Gbps, providing fast access to datasets, model checkpoints, and on-disk KV cache.
|
||||
A fundamental principle of this architecture is that the compute network and the storage network are isolated from each other ( Zhao et al., 2025a ) . This separation is essential to maximize both storage and application performance. By isolating high-intensity east-west compute traffic between GPUs from storage traffic, the architecture prevents interference between them, and drastically reduces compute communication latency. This design also ensures that the inter-GPU communication remains highly reliable and predictable even when performing data-intensive tasks such as reading large datasets or writing multi-terabyte model checkpoints.
|
||||
## 3. Bottleneck & Motivation
|
||||
We observe severe GPU underutilization during agentic inference tasks.
|
||||
Our investigation reveals that KV-Cache loading speed is the bottleneck due to the limited bandwidth of the single storage NIC on each node.
|
||||
Analysis demonstrates that three decisive factors jointly cause this bottleneck, as discussed below.
|
||||
First, agentic workloads exhibit high KV-Cache hit rates, which require more I/O and less computation , thus creating a severe I/O bottleneck.
|
||||
Agentic workloads are naturally long-context, short-append, and multi-turn.
|
||||
On each turn, the GPU needs to read the KV-Cache of the entire context from persistent storage and perform prefill computation for appended tokens.
|
||||
Our trace collected from representative coding tasks shows the mean number of rounds is 157, demonstrating the tendency of LLMs to engage in multi-turn interactions. The average context length is 32.7k, while the append length mean is only 429, which means a KV-Cache hit rate of 98.7%.
|
||||
In such a scenario, the cache-compute ratio, defined as the ratio of KV-Cache to load and the computation needed, is approximately 22 GB/PFLOP for DeepSeek-V3.2 ( DeepSeek-AI, 2025d ) , posing a significant bottleneck on storage bandwidth.
|
||||
Note that the KV-Cache size of DeepSeek MLA model is already highly optimized; for models with larger KV-Cache sizes (see Table 1 ), the situation is even worse.
|
||||
The ratio of DeepSeek-V3.2 is higher than DeepSeek-V3 ( DeepSeek-AI, 2025c ) , benefiting from its sparse attention design, lowering computation demands.
|
||||
Table 1 . Cache-compute ratio with append length 429, across context lengths (16k–64k). KV-Cache data type defaults to FP8 unless specified.
|
||||
Model
|
||||
GB/PFLOP (16K–64K)
|
||||
Qwen2.5-32B ( Team, 2025a ) (FP16)
|
||||
117-267
|
||||
GPT-OSS-120B ( OpenAI, 2025a )
|
||||
47–95
|
||||
Qwen3-235B-A22B ( Team, 2025b )
|
||||
39–60
|
||||
DeepSeek-V3.2 660B ( DeepSeek-AI, 2025d )
|
||||
13–36
|
||||
DeepSeek-V3 660B ( DeepSeek-AI, 2025c )
|
||||
4.8–5.8
|
||||
Second, the hardware evolution trend is not well suited for agentic inference workloads.
|
||||
In recent years, network bandwidth and HBM capacity have lagged behind the growth of GPU FLOPS, which drives us to run into memory and communication walls under agentic workloads.
|
||||
As shown in Figure 3 , from NVIDIA Ampere to Blackwell, the I/O-compute ratio decreases by 14.4 × \times .
|
||||
Low NIC bandwidth limits KV-Cache loading speed, making GPUs idle.
|
||||
In addition, small HBM capacity limits the token batch size for GPU kernels ( Dao, 2024 ; Ye et al., 2025 ; DeepSeek-AI, 2025b ; Li and Liu, 2025 ) to compute at the same time, hindering full utilization of compute units such as Tensor Core ( Du et al., 2025 ) .
|
||||
Third, existing LLM inference systems exhibit severe storage network utilization imbalance across different engine types. In prevalent PD-disaggregated systems, the KV-Cache for hit tokens is loaded exclusively by prefill engines directly from remote storage. This design centralizes all storage I/O pressure on the prefill-side SNICs, while the SNICs on decode engines remain largely idle. Consequently, the aggregate storage network bandwidth cannot be fully harnessed.
|
||||
Figure 3 . Left: Hardware trends of NVIDIA GPUs. Right: Relative token throughput with varying request batch size (each request has 30K context with 300 tokens appended).
|
||||
The above analysis demonstrates that the fundamental performance issue for agentic inference on PD-disaggregated architecture is the high I/O demand for KV-Cache retrieval and unbalanced storage network bandwidth utilization across inference engines. Meanwhile, we observe that the network traffic of the compute network, which has much larger aggregate bandwidth than the storage network, exhibits an intermittent pattern: collective operations used in model inference burst in sub-millisecond intervals. Therefore, an opportunity naturally emerges: we can utilize the SNIC bandwidth of decode nodes to load KV-Cache from storage, and transfer it back to the prefill nodes, utilizing the spare bandwidth of the faster compute network.
|
||||
## 4. DualPath System Overview
|
||||
To break the prefill-side storage I/O bottleneck, we propose a dual-path loading architecture that fundamentally rethinks how KV-Cache is retrieved in PD-disaggregated inference. Based on this architecture, we design and implement DualPath.
|
||||
DualPath adopts two widely-adopted techniques demonstrated in § 2 :
|
||||
(1) PD Disaggregation ( Zhong et al., 2024 ; Patel et al., 2025 ) , which separates prompt and decode processing for better efficiency.
|
||||
(2) Layerwise prefill , which avoids HBM bottlenecks recognized by LayerKV ( Xiong et al., 2024 ) and PrefillOnly ( Du et al., 2025 ) on prefill engines and improves GPU utilization.
|
||||
Our system consists of the following components:
|
||||
•
|
||||
Inference Engines. Each engine manages one GPU.
|
||||
Engines are categorized into prefill engines (PEs) for prefill and decoding engines (DEs) for decode.
|
||||
•
|
||||
Traffic Manager ( § 5 ). Each engine contains a traffic manager to conduct (1) Host-Device memory copies (H2D & D2H), (2) KV-Cache transfers between PEs and DEs, and (3) KV-Cache reads/writes from/to storage via the storage NIC. We adopt a CNIC-centric traffic management approach, detailed in § 5 , to prevent KV-Cache traffic from affecting communications in model inference.
|
||||
•
|
||||
Request Scheduler ( § 6 ). A central scheduler that receives client requests and distributes them across engines. It is also responsible for dynamically distributing data traffic between two paths ( Figure 4 ).
|
||||
(a) PE Read Path
|
||||
(b) DE Read Path
|
||||
Figure 4 . Dual-path loading illustration. The scheduler dynamically distributes data traffic between the two paths.
|
||||
## 4.1. Dual-Path Loading
|
||||
In addition to the conventional storage-to-prefill path, DualPath introduces a novel storage-to-decode path, allowing KV-Cache to be loaded first into a decode engine and then transferred to the prefill engine via high-bandwidth RDMA over the compute network.
|
||||
By dynamically distributing load across both paths, the system aggregates the storage NIC bandwidth of all engines — including otherwise-idle decode-side NICs — and eliminates the asymmetric bandwidth saturation that limits existing systems. This approach transforms the storage I/O from a single-bottleneck resource into a globally pooled and schedulable capacity. The exact data flows of the dual-path are described below.
|
||||
To implement dual-path loading, DualPath allocates a small amount of DRAM as buffers on each PE and DE, called PE buffer and DE Buffer .
|
||||
Prefill PE read path.
|
||||
First, the KV-Caches of hit tokens are read from persistent storage into the PE buffer (as Label 1 and 2 shown in 4(a) ).
|
||||
Before the computation of an attention layer, those KV-Caches of that layer are transferred to PE HBM (3 and 4) to compute the KV-Cache of cache-miss prompt tokens.
|
||||
Then, all KV-Caches of both hit and miss tokens are transferred to the DE buffer to form the complete prompt KV-Cache (5-7). This process (3-7) repeats n l a y e r n_{layer} times.
|
||||
During the prefill forward pass, transfers overlap with computation.
|
||||
Prefill DE read path. The KV-Caches of hit tokens are first read into DE buffer (as Label 1 and 2 shown in 4(b) ).
|
||||
During PE prefill, KV-Cache for the corresponding layer is read from the DE buffer, also overlapping with computation (3-5). This process repeats n l a y e r n_{layer} times.
|
||||
After a layer’s computation completes, only the KV-Caches of miss tokens are transferred to DE buffer and merged with the existing hit token KV-Cache.
|
||||
Decode Phase. After receiving the complete prompt KV-Cache in DE buffer (including loaded KV-Cache via PE read path and the KV-Cache of newly appended tokens), the decode phase begins.
|
||||
The DE first allocates HBM and performs host-to-device (H2D) transfers (Label 8 and 9 in 4(a) ; Label 6 and 7 in 4(b) ), then releases CPU memory before starting decode.
|
||||
The design of DE buffer imposes bandwidth pressure on DRAM and CNIC (an extra H2D), which could be avoided by directly bypassing it via GPU Direct RDMA.
|
||||
However, since the generation length is typically short in this scenario, time-to-first-token (TTFT) accounts for a non-negligible portion of the total end-to-end request time.
|
||||
Introducing DE buffer helps reduce GPU memory usage.
|
||||
During decode, whenever a full block of tokens (e.g., 64 tokens) is accumulated, it is immediately persisted to disk.
|
||||
Different Block Layouts.
|
||||
We adopt two different block layouts: Full Block and Layer Block , which contain all layers and a single layer, respectively.
|
||||
Detailed layout can be found in § A.5 .
|
||||
For all interactions with storage, we adopt Full Blocks.
|
||||
In the PE read case, KV-Cache loading to PE HBM and transfer to DE Buffer occur in a layerwise streaming fashion, both using Layer Blocks.
|
||||
Similarly, for the DE read path, transfers from the DE Buffer to the PE HBM use Layer Blocks.
|
||||
## 4.2. Bottleneck-Free Analysis
|
||||
We demonstrate that the system can fully saturate all storage NICs without introducing compute-NIC or DRAM bottlenecks, under most reasonable P/D ratios.
|
||||
We assume a well-configured PCIe topology (each pair of GPU–NIC is under the same PCIe switch), load-balanced task scheduling, no congestion on the computation network, and that storage read bandwidth is fully utilized.
|
||||
Notation. Let P P and D D denote the number of prefill and decode nodes, respectively. Each node has g g GPUs, each with one compute NIC of bandwidth B B . The storage bandwidth per machine is s × B s\times B (shared by all engines on that machine); M M is the memory bandwidth per machine.
|
||||
Traffic per PE-DE pair. We assume that the storage read bandwidth is fully utilized and that task scheduling is load-balanced. Under load balancing, storage NIC bandwidth is evenly shared. The traffic per pair for the PE read path (all steps in 4(a) ) is T p = B s / ( D g 2 ) T_{p}=Bs/(Dg^{2}) ; for the DE read path ( 4(b) ) it is T c = B s / ( P g 2 ) T_{c}=Bs/(Pg^{2}) . Link traffic is the sum over all pairs using that link.
|
||||
PE CNIC Bandwidth Analysis.
|
||||
For PE CNIC, loopback traffic (i.e., H2D and D2H that does not traverse switches) exists, so the total traffic on the PCIe side is always greater than or equal to the switch-direction traffic, regardless of read or write operations.
|
||||
Therefore, we only need to compute the pressure on the PCIe side.
|
||||
Read operations include PE paths (3) and (5), with total traffic over all pairs:
|
||||
(1)
|
||||
2 × T p × D g = 2 B s / g ≤ B \displaystyle 2\times T_{p}\times Dg=2Bs/g\leq B
|
||||
Since s ≤ g s\leq g always holds in practice, the read direction is always bottleneck-free.
|
||||
Write operations include PE path (4) and DE path (5), with total traffic:
|
||||
(2)
|
||||
( T p + T c ) × D g = B s / g × ( 1 + D / P ) ≤ B \displaystyle(T_{p}+T_{c})\times Dg=Bs/g\times(1+D/P)\leq B
|
||||
Then, we obtain:
|
||||
(3)
|
||||
P / D ≥ s g − s \displaystyle P/D\geq\frac{s}{g-s}
|
||||
DE CNIC Bandwidth Analysis.
|
||||
For DE CNIC, read operations include PE path 8 and DE paths 3/6, with traffic:
|
||||
(4)
|
||||
( T p + T c × 2 ) × P g = s / g × ( P / D + 2 ) × B ≤ B \displaystyle(T_{p}+T_{c}\times 2)\times Pg=s/g\times(P/D+2)\times B\leq B
|
||||
Then, we obtain:
|
||||
(5)
|
||||
P / D ≤ g − 2 s s \displaystyle P/D\leq\frac{g-2s}{s}
|
||||
Write operations include PE paths 7/9 and DE path 7, with traffic:
|
||||
(6)
|
||||
( 2 T p + T c ) × P g ≤ B \displaystyle(2T_{p}+T_{c})\times Pg\leq B
|
||||
This implies:
|
||||
(7)
|
||||
P / D ≤ g − s 2 s \displaystyle P/D\leq\frac{g-s}{2s}
|
||||
DRAM Pressure Analysis.
|
||||
DRAM is half-duplex, so we sum the read and write pressures.
|
||||
For PE MEM, the pressure is 2 s B 2sB , which generally does not exceed memory bandwidth.
|
||||
For DE MEM, following the similar analysis above, we can get the pressure is ( 3 + 2 P / D ) B s (3+2P/D)Bs .
|
||||
Requiring the DE MEM pressure to be less than or equal to M M , we obtain:
|
||||
(8)
|
||||
P / D ≤ M / B s − 3 2 \displaystyle P/D\leq\frac{M/Bs-3}{2}
|
||||
Summary. Combining all the above analyses, we have:
|
||||
(9)
|
||||
s g − s ≤ P / D ≤ min { g − 2 s s , g − s 2 s , M / B s − 3 2 } . \displaystyle\frac{s}{g-s}\leq P/D\leq\min\left\{\frac{g-2s}{s},\frac{g-s}{2s},\frac{M/Bs-3}{2}\right\}.
|
||||
For ( g = 8 , s = 1 ) (g=8,s=1) with M ≈ 500 M\approx 500 GB/s and B s ≈ 50 Bs\approx 50 GB/s, the bottleneck-free range is 1 7 ≤ \frac{1}{7}\leq P/D ≤ 7 2 \leq\frac{7}{2} , which covers most practical configurations.
|
||||
## 4.3. Practical Challenges
|
||||
The dual-path architecture fundamentally reorients data movement: KV-Cache can be loaded either directly from storage into prefill engines or indirectly via decode engines, thereby aggregating storage bandwidth across all engines and breaking the prefill-side I/O bottleneck. However, realizing this high-level design in a practical system introduces three interrelated challenges. We briefly outline these challenges below and refer the reader to the corresponding sections for details.
|
||||
Fine-grained data transfer ( § 5 ).
|
||||
The layer-wise execution paradigm, while essential for overcoming HBM capacity limits, fragments the KV-Cache into numerous fine-grained blocks ( Patel et al., 2025 ) . Transferring this multitude of fine-grained data chunks between storage, host DRAM, and GPU HBM must incur minimal overhead and seamlessly overlap with computation to realize throughput gains.
|
||||
Traffic isolation ( § 5 ).
|
||||
The complex data path in DualPath introduces additional KV-Cache transfer traffic on both the compute network and PCIe links. A primary concern is that this traffic may interfere with existing latency-sensitive collective communication operations essential for model execution — such as AllToAll in expert parallel ( Zhao et al., 2025b ) and ReduceScatter/AllGather in tensor/context parallel. Since these collective communications are critical to end-to-end inference latency, a key challenge lies in exploiting spare I/O bandwidth without degrading model inference performance.
|
||||
Dynamic load balancing ( § 6 ).
|
||||
As we are adopting two different paths for KV-cache loading, the system must promptly decide which path to use for each request. A naive policy could overload one path, recreating the original bottleneck. The traffic scheduler must balance multiple factors in real-time: storage NIC queue lengths, computational load on GPUs, and request workload characteristics.
|
||||
## 5. CNIC-Centric Traffic Manager
|
||||
Modern LLM inference systems employ a range of advanced data transfer technologies — such as on-chip CUDA copy engine and GPUDirect Storage ( NVIDIA, 2026b ) — to move data efficiently between storage, host memory, and GPU HBM. However, all these mechanisms can interfere with latency-sensitive collective communications (e.g., EP AllToAll) during model execution. This arises for two primary reasons: (1) such transfer technologies often operate over separate paths that do not share the same QoS controls as the compute network, and (2) existing GPUs do not support PCIe QoS ( Richter et al., 2016 ) , making it difficult to shield model inference communication from other traffic contending for PCIe bandwidth. Additionally, because the collective communications occur in rapid, sub-millisecond-level bursts, it is impractical to rely on a software-based traffic shaper to interleave lower-priority I/O operations between these high-priority traffic windows.
|
||||
To address this, we propose a CNIC–centric data transfer approach which is widely adopted in our production deployment: all data traffic in or out of a GPU, including local H2D/D2H copy, must go through the GPU’s paired CNIC with a GPUDirect RDMA ( NVIDIA, 2026a ) data path. By consolidating all traffic onto the compute network, we can leverage the native QoS capabilities of compute network to enforce strict traffic differentiation.
|
||||
## 5.1. Traffic Isolation
|
||||
For the InfiniBand-based network, we leverage virtual lanes (VLs) ( Association, 2007 ) to enforce isolation between different traffic classes.
|
||||
All model inference communication traffic is assigned to a dedicated high-priority VL, while all other traffic, including KV-Cache transfer, is mapped to a separate low-priority VL.
|
||||
We configure the VL arbiters of all network switches and NICs with a weighted round-robin policy that reserves approximately 99% of total bandwidth to high-priority VL.
|
||||
The remaining bandwidth is allocated to the low-priority VL to prevent starvation.
|
||||
This configuration ensures that model execution traffic is virtually unaffected by KV-cache transfers, while still allowing KV-cache traffic to opportunistically utilize otherwise idle bandwidth in the compute network.
|
||||
Detailed configurations are described in § A.1 .
|
||||
Although our experiments are conducted on an InfiniBand-based network, the same design principles naturally extend to other interconnect technologies.
|
||||
DualPath can be conducted on RDMA over Converged Ethernet (RoCE) by leveraging Traffic Class (TC) and Differentiated Services Code Point (DSCP) markings ( Guo et al., 2016 ; Carpenter and Nichols, 2002 ) in conjunction with hardware packet queues.
|
||||
Emerging technologies such as UnifiedBus ( , 2026 ) and Ultra Ethernet ( Consortium, 2026 ) are likewise converging on QoS mechanisms for heterogeneous traffic, which can directly support the requirements of DualPath.
|
||||
## 5.2. CNIC-Assisted KV-Cache Copy
|
||||
Existing GPU data transfer technologies include GPUDirect Storage ( NVIDIA, 2026b ) , which loads KV-Cache from storage backend to GPU HBM, and CUDA copy engine, which directly copies host DRAM to GPU via PCIe. However, these methods fail to isolate KV-Cache traffic from high-priority latency-sensitive collective communications in model execution, severely degrading inference performance.
|
||||
To solve the limitations of existing approaches, we adopt a CNIC-assisted H2D/D2H data path. For KV-Cache loading, we first read the KV-Cache into host DRAM from the storage backend. Then, we submit an RDMA Write request to the GPU’s paired CNIC to perform local H2D copy. Storing newly generated KV-Cache follows a symmetric process: it is first transferred to host DRAM via CNIC, then persisted to the storage backend over the storage network. This design establishes the CNIC as the central QoS scheduler for all GPU PCIe traffic, allowing its VL arbiter to prioritize the inference communication traffic and perform KV-Cache transfer using spare PCIe bandwidth.
|
||||
Although this approach may appear to take a detour compared to GPUDirect Storage (which directly reads KV-Cache to GPU HBM) and CUDA copy engine (which directly copies host memory to GPU HBM), to our best knowledge, this is currently the only practical method to ensure that KV-Cache load/store does not degrade the performance of critical model–execution communication.
|
||||
We also observe that CNIC-assisted H2D and D2H outperform the CUDA copy engine when handling a large number of small data chunks. Our measurements show that submitting a single copy operation via cudaMemcpyAsync incurs a latency overhead of approximately 5-7 μ s \mu s . We failed to further break down this overhead due to the closed-source nature of CUDA driver. In contrast, submitting one RDMA Write work request involves only a few mmio writes to NIC registers in user space and takes only around 1 μ s \mu s . Furthermore, the RDMA work submission overhead can be significantly amortized by leveraging doorbell batching ( Kalia et al., 2016 ) .
|
||||
## 6. Adaptive Request Scheduler
|
||||
Although our theoretical analysis shows promising results, imbalanced load reduces hardware utilization, and in this scenario, we need to consider two dimensions of balance simultaneously: (1) NIC traffic, and (2) the utilization balance of GPUs.
|
||||
We divide scheduling into two levels: inter-engine scheduling, which assigns requests to a (PE, DE) pair and selects the read path (PE or DE) for each request; and intra-engine scheduling, which determines which requests are included in each forward batch for computation.
|
||||
## 6.1. Inter-Engine Scheduling
|
||||
We organize engines into groups to reduce the scheduler pressure.
|
||||
Only the engine rank 0 0 , called Leader Engine , interacts with the scheduler.
|
||||
All engines of a group are all PEs or all DEs.
|
||||
All engines on one node are guaranteed to be in the same group.
|
||||
All engines in one group proactively fetch tasks together regularly.
|
||||
When fetching new requests, each engine e e reports (1) s e q e seq_{e} , the number of requests assigned to it that have not yet completed; (2) the total token count t o k e tok_{e} over those s e q e seq_{e} requests; and (3) the disk reading queue length r e a d _ q n ( e ) read\_q_{n(e)} of the node n ( e ) n(e) that engine e e belongs to.
|
||||
GPU load, disk read load, and network load are all strongly correlated with token count. We therefore use token count as a proxy and aim to balance it across engines.
|
||||
Figure 5 . An illustration of Inter-Engine PE Scheduling. All eight GPUs are in the same PE engine group and the scheduler will choose the best.
|
||||
PE Scheduling. All requests arriving at the scheduler enter a waiting queue and are scheduled in a FIFO order. The scheduling algorithm is invoked when a PE group initiates a fetch request.
|
||||
An illustration of the inter-engine scheduling process is shown in Figure 5 .
|
||||
We define two constants, short reading queue threshold α \alpha , and unfinished token upper limit β \beta , measured in tokens.
|
||||
All engines are split into three categories:
|
||||
(1) overloaded engines where t o k e > β tok_{e}>\beta ;
|
||||
(2) engines on nodes with short disk reading queues where r e a d _ q n ( e ) ≤ α read\_q_{n(e)}\leq\alpha and t o k e ≤ β tok_{e}\leq\beta ; and
|
||||
(3) engines on nodes with longer disk reading queues where r e a d _ q n ( e ) > α read\_q_{n(e)}>\alpha and t o k e ≤ β tok_{e}\leq\beta .
|
||||
We do not assign new requests to overloaded engines.
|
||||
Second-category engines are prioritized over third-category engines because they reside on nodes with shorter disk reading queues, and lack of subsequent requests would easily lead to storage NIC underutilization.
|
||||
We assign the current request to the PE with minimum t o k e tok_{e} in the second category if non-empty, otherwise in the third category if non-empty.
|
||||
After assignment, we update the selected PE’s t o k e tok_{e} , then proceed to the next request in the waiting queue.
|
||||
If both categories are empty, we terminate this fetch request and return the already-assigned requests to the Leader Engine.
|
||||
Algorithm 1 Inter-PE Scheduling Algorithm
|
||||
Data: Waiting queue Q Q , PE group G P E G_{PE} , load metrics ( r e a d _ q n ( e ) , t o k e ) (read\_q_{n(e)},tok_{e}) reported by each engine e e , where n ( e ) n(e) denotes the node that engine e e belongs to, constants α \alpha and β \beta
|
||||
Result: Assigned requests to PEs
|
||||
Each engine e e reports ( r e a d _ q n ( e ) , t o k e ) (read\_q_{n(e)},tok_{e}) ;
|
||||
Classify all PEs into three categories:;
|
||||
C 1 ← { e ∈ G P E : t o k e > β } C_{1}\leftarrow\{e\in G_{PE}:tok_{e}>\beta\} ;
|
||||
C 2 ← { e ∈ G P E : r e a d _ q n ( e ) ≤ α ∧ t o k e ≤ β } C_{2}\leftarrow\{e\in G_{PE}:read\_q_{n(e)}\leq\alpha\wedge tok_{e}\leq\beta\} ;
|
||||
C 3 ← { e ∈ G P E : r e a d _ q n ( e ) > α ∧ t o k e ≤ β } C_{3}\leftarrow\{e\in G_{PE}:read\_q_{n(e)}>\alpha\wedge tok_{e}\leq\beta\} ;
|
||||
while Q Q is not empty do
|
||||
r ← r\leftarrow head of Q Q ;
|
||||
if C 2 ≠ ∅ C_{2}\neq\emptyset then
|
||||
p e ∗ ← arg min e ∈ C 2 t o k e pe^{*}\leftarrow\arg\min_{e\in C_{2}}tok_{e} ;
|
||||
else
|
||||
if C 3 ≠ ∅ C_{3}\neq\emptyset then
|
||||
p e ∗ ← arg min e ∈ C 3 t o k e pe^{*}\leftarrow\arg\min_{e\in C_{3}}tok_{e} ;
|
||||
else
|
||||
Terminate this fetch request;
|
||||
Return assigned requests to Leader Engine;
|
||||
break ;
|
||||
end if
|
||||
end if
|
||||
Assign request r r to PE p e ∗ pe^{*} ;
|
||||
Update t o k p e ∗ ← t o k p e ∗ + tokens ( r ) tok_{pe^{*}}\leftarrow tok_{pe^{*}}+\text{tokens}(r) ;
|
||||
Remove r r from Q Q ;
|
||||
end while
|
||||
DE Scheduling Phase 1: across groups.
|
||||
DE scheduling is two-level and does not preserve global FIFO. There is a global waiting queue and a private queue per DE engine group. Incoming requests first enter the global queue. When a DE group fetches, group-level scheduling drains the global queue and assigns each request to the group whose total t o k e tok_{e} (sum over its engines) is minimum; this balances token count across groups and thus NIC and GPU load.
|
||||
DE Scheduling Phase 2: within a group.
|
||||
Then we calculate the sum of remaining HBM for all DEs in the group and traverse from the head of the private queue to calculate how many requests can be scheduled assuming no HBM fragmentation. These requests form the set R R .
|
||||
It is an upper bound that can be scheduled.
|
||||
Then, we calculate a high token threshold Z = 1.05 × ( ∑ r ∈ R l e n r + ∑ e ∈ E t o k e ) / | E | Z=1.05\times(\sum_{r\in R}{len_{r}}+\sum_{e\in E}{tok_{e}})/|E| .
|
||||
Next, we try to pop the head of private queue and schedule it to a DE.
|
||||
Among DEs with sufficient remaining HBM for the request, we partition into (1) high-token DEs, where t o k e + len ( r ) > Z tok_{e}+\text{len}(r)>Z , and (2) the rest. We prefer category (2) to keep token count balanced; category (1) DEs already have higher GPU and NIC pressure. Within category (2) we choose the DE with minimum s e q e seq_{e} to balance request count; if category (2) is empty we choose the DE with minimum t o k e tok_{e} in category (1) to reduce HBM exhaustion and preemption risk. If no DE has sufficient HBM, the fetch ends and already-assigned requests are returned.
|
||||
KV-Cache Read Task Scheduling.
|
||||
After selecting PE and DE for a request, we choose to read on the side with the shorter reading queue. It is probably better to split the request into two parts and read them from both sides, and we leave it as future work.
|
||||
## 6.2. Intra-Engine Scheduling
|
||||
Only PEs require intra-engine scheduling, as DEs always place all requests into the forward batch.
|
||||
An illustration of the intra-engine scheduling process is shown in Figure 6 .
|
||||
Data parallelism is widely adopted for attention layers, especially for MLA models.
|
||||
Under such a parallel configuration, each GPU serves a different set of requests.
|
||||
It may lead to workload imbalance among all GPUs that must synchronize after the attention stage and enter the FFN stage together, causing GPU bubbles waiting for other peers.
|
||||
Therefore, we need to make sure they have similar attention layer execution times to minimize the waiting bubbles.
|
||||
Figure 6 . Intra-Engine Schedule. Left: compute-quota-based batch selection. Right: GPU timeline before and after applying compute quota.
|
||||
Layer Time Estimation. We use FIFO packing to decide how many requests to include in a forward batch.
|
||||
Each request in a forward batch is described by a pair ( c a c h e d , b s z ) (cached,bsz) , where c a c h e d cached is the number of tokens with KV-Cache already available (from storage hits or previous forward passes), and b s z bsz is the number of tokens requiring KV-Cache computation in this forward batch.
|
||||
From these pairs, we compute the total theoretical computation for the attention layer and estimate its execution time.
|
||||
The relationship between theoretical computation and wall-clock time depends on hardware and parallel configuration, and can be fitted in advance through profiling as previous works ( Du et al., 2025 ) and ( Agrawal et al., 2024 ) .
|
||||
Algorithm. We keep adding requests in FIFO order as long as the predicted attention layer execution time does not exceed a predefined upper bound, called the compute quota .
|
||||
If adding a request would exceed this bound, we perform binary search on b s z bsz to find a smaller b s z ′ bsz^{\prime} to fit in the remaining compute quota and perform chunked prefill for that request.
|
||||
## 7. Evaluation
|
||||
Figure 7 . Offline inference performance under varying numbers of agents and maximum agent context lengths. Top: DS 27B. Middle: DS 660B. Bottom: Qwen 32B. N/A for running into an error before finishing.
|
||||
## 7.1. Implementation
|
||||
We implement DualPath based on our in-house inference framework.
|
||||
For CUDA kernels, our in-house framework adopt the combination of FlashMLA ( Li and Liu, 2025 ) , DeepGEMM ( DeepSeek-AI, 2025b ) , and DeepEP ( Zhao et al., 2025b ) , which aligns with the current mainstream open-source framework ( Zheng et al., 2024 ; Kwon et al., 2023 ) .
|
||||
The DualPath implementation involves approximately 5K lines of modifications on top of it.
|
||||
We adopt 3FS ( DeepSeek-AI, 2025a ) as distributed storage and use an io_uring -like interface for kernel bypass.
|
||||
## 7.2. Experimental Setup
|
||||
Testbed.
|
||||
We conduct our experiments on a cluster of GPU servers with InfiniBand interconnection. Each server has 8 NVIDIA Hopper GPUs and dual processors. Additionally, each node is provisioned with eight 400Gbps RDMA NICs connected to InfiniBand network and one additional storage NIC connected to 3FS. The computation and storage networks are physically isolated. Our cluster-wide 3FS has no internal DRAM cache and can saturate the 400Gbps bandwidth of the storage NIC.
|
||||
Models.
|
||||
We evaluate on three models: (1) DeepSeek V3.2 ( DeepSeek-AI, 2025d ) 660B, an MoE model with DeepSeek Sparse Attention, denoted as DS 660B , (2) a 27B downscaled version of DS 660B, denoted as DS 27B , and (3) Qwen2.5-32B ( Team, 2025a ) , a dense model with GQA, denoted as Qwen 32B .
|
||||
DS 660B and Qwen 32B correspond to the publicly released checkpoint on HuggingFace.
|
||||
DS 27B is our internal experimental model with a similar architecture to DS 660B.
|
||||
Detailed specifications are provided in § A.2 .
|
||||
Datasets.
|
||||
We collected three agent trace datasets from our production agentic RL training workloads with varying maximum context lengths (MaxLen). Each dataset contains 500 trajectories. The average interaction turns (Turns), average appended and generated tokens per turn (Append and Gen), average number of total tokens (Total), and average number of context tokens (Context) are summarized in Table 2 .
|
||||
Table 2 . Statistics of agent trace datasets.
|
||||
MaxLen
|
||||
Turns
|
||||
Append
|
||||
Gen
|
||||
Total
|
||||
Context
|
||||
32K
|
||||
60
|
||||
608
|
||||
148
|
||||
28639
|
||||
17183
|
||||
48K
|
||||
106
|
||||
474
|
||||
172
|
||||
42607
|
||||
25120
|
||||
64K
|
||||
157
|
||||
429
|
||||
176
|
||||
55958
|
||||
32721
|
||||
Baselines.
|
||||
We compare DualPath, denoted as Ours, against the following baselines:
|
||||
•
|
||||
SGL(MC) : SGLang ( Zheng et al., 2024 ) (commit 19089aa) with HiCache ( SGLang, 2026 ) , Mooncake ( Qin et al., 2025 ) Store enabled and 3FS as the storage backend, and Mooncake Transfer Engine for prefill-decode disaggregation.
|
||||
We did not run SGL(MC) for DS 27B because SGLang lacks support for this downscaled version.
|
||||
•
|
||||
Basic : Our unmodified internal inference framework (detailed in § 7.1 ).
|
||||
Comparing DualPath and SGL(MC) is unfair due to implementation differences.
|
||||
Therefore, we only report performance improvements from Basic to Ours.
|
||||
•
|
||||
Oracle : Based on DualPath, we bypass all disk reads, D2H & H2D transfers, and inter-PD KV-Cache transfers.
|
||||
This configuration represents the theoretical performance upper bound assuming zero I/O overhead.
|
||||
P/D Ratio and Parallelism.
|
||||
We default to 2P4D for DS 660B, 1P2D for Qwen 32B, and 1P1D for DS 27B (where 1P1D means one node for each side).
|
||||
For DS models, we use EP and DP.
|
||||
For Qwen 32B, we use DP only in DualPath, while SGL(MC) uses TP=8 since DP attention is not supported for this model in SGLang.
|
||||
Detailed configuration is provided in § A.4 .
|
||||
Metrics. For offline inference scenarios, we measure job completion time (JCT) for the entire task.
|
||||
For online serving scenarios, we measure TTFT, TTST (Time to the second token), and TPOT.
|
||||
## 7.3. Offline Batch Inference
|
||||
This section evaluates throughput performance in offline batch inference, which is the case of the rollout phase in RL training.
|
||||
In this scenario, n n agents start to rollout simultaneously, and we measure the JCT when all requests have finished.
|
||||
Varying Agents Batch Size & Max Agent Length (MAL).
|
||||
DualPath benefits more from larger batch sizes and longer MALs.
|
||||
Figure 7 reports JCT under different batch sizes and MALs.
|
||||
SGL(MC) encountered errors in our setup and failed to complete some large configurations (marked as N/A).
|
||||
On DS 660B, DualPath achieves up to 1.87 × 1.87\times over Basic, and demonstrates performance with Oracle, indicating that KV-cache I/O is largely eliminated.
|
||||
On DS 27B, DualPath improves over Basic by up to 1.78 × 1.78\times but remains 1.09 1.09 – 1.85 × 1.85\times slower than Oracle due to limited storage bandwidth in 1P1D ( Figure 8 ).
|
||||
For Qwen 32B, it shows similar trends as DS 27B.
|
||||
Figure 8 . Impact of prefill-decode ratio on offline inference performance (DS 27B).
|
||||
Figure 9 . Left: varying append lengths (DS 660B, 64K context, 1024 agents). Right: varying generation lengths (DS 660B, 64K, 1024 agents)
|
||||
Figure 10 . TTFT, TTST, and TPOT as functions of agent arrival rate (APS). Shadow means the fluctuation in the last 150s before experiments finish. Top: DS 27B, Bottom: DS 660B.
|
||||
Varying Append Length & Generation Length.
|
||||
DualPath has more advantages when append and generation tokens are short.
|
||||
Longer append lengths imply greater GPU compute pressure, and longer generation lengths lower KV-Cache loading pressure due to larger prefill gap time.
|
||||
To investigate the impact of this factor, we scale each round’s append length by a constant factor, and then truncate the whole trajectory at given MAL.
|
||||
The same holds for generation length.
|
||||
As shown in Figure 9 , with append length increases, Basic performance gradually approaches DualPath and Oracle, while DualPath and Oracle performance changes only slightly, indicating that the bottleneck consistently lies in GPU compute pressure.
|
||||
Compared to Basic, DualPath achieves 1.82 − 1.99 × 1.82-1.99\times speedup at different append scales.
|
||||
The trend for generation length scaling is similar.
|
||||
Varying Prefill-Decode Ratio.
|
||||
Across all ratios, DualPath demonstrates substantial performance gains compared to Basic.
|
||||
We conduct rollout experiments on DS 27B with 1P1D, 2P1D, and 1P2D prefill-decode ratios to characterize the impact of resource allocation between prefill and decode stages on overall system performance.
|
||||
As shown in Figure 8 , DualPath achieves an average speedup of 1.64 × 1.64\times across all configurations (up to 2.46 × 2.46\times ).
|
||||
Basic 1P1D and Basic 1P2D perform comparably;
|
||||
so do DualPath 1P1D and Basic 2P1D,
|
||||
as well as DualPath 2P1D and DualPath 1P2D.
|
||||
This occurs because each pair of systems has equivalent available storage bandwidth (Basic can only use prefill node storage bandwidth, while DualPath can utilize all nodes), which confirms that storage bandwidth is the dominant bottleneck in agentic scenarios.
|
||||
## 7.4. Online Serving
|
||||
Methodology. We evaluate system latency characteristics under varying agent arrival rates per second (APS).
|
||||
Agents arrive according to a Poisson process at a specified rate, with each agent commencing replay from round zero to its last round upon arrival.
|
||||
For our experiments, the SLO is set as TTFT ≤ \leq 4 seconds and TPOT ≤ \leq 50ms.
|
||||
In the TPOT and TTST figures, data points exceeding the SLO threshold are omitted.
|
||||
Experiment termination is triggered when either: (1) TTFT exceeds 4 seconds, or (2) the system reaches steady state, defined as TTFT variation within a 150-second sliding window remaining below 5% compared to that 30 minutes prior.
|
||||
As shown in Figure 10 , DualPath achieves higher APS capacity than Basic (1.67 × \times for DS 27B, 2.25 × \times for DS 660B).
|
||||
DualPath’s TTST is comparable to Basic, while TPOT shows that DualPath does not introduce additional decoding overhead compared to Basic.
|
||||
SGL(MC) exhibits anomalously low TTST, likely due to implementation issues where the first two tokens arrive at the
|
||||
client almost simultaneously.
|
||||
For DS 27B, all metrics exhibit trends similar to DS 660B.
|
||||
However, both Basic and DualPath show significantly higher TPOT than Oracle, suggesting the overhead of basic P-D transferring is considerable in small model cases. We leave it as future work.
|
||||
Figure 11 . Average completion time of all trajectories versus arrival rate for online serving.
|
||||
Average JCT for both models are presented in Figure 11 .
|
||||
A detailed analysis of working set implications is discussed in § 8 .
|
||||
As shown in Figure 12 (left), DualPath maintains stable TTFT components across different APS, while Basic’s queuing time grows dramatically due to insufficient storage bandwidth.
|
||||
## 7.5. Ablation Study
|
||||
Figure 12 . Left ( § 7.4 ): TTFT breakdown for online serving (DS 660B) across APSs, Sch. for scheduling, A. for allocating, R. for reading KV-cache, PF. for prefill. In each pair of pillars, the first is for DualPath and the second is for Basic. Right ( § 7.5 ): Offline inference ablation results (DS 660B, 64K context length). Layer, DPL, Sched stands for Layerwise prefill, Dual-Path Loading, and scheduling, respectively.
|
||||
We conduct an ablation study to quantify the contribution of each technical component in DualPath.
|
||||
Experiments are performed under the offline inference setting with 64K MAL and agent batch size 1024 and 2048.
|
||||
The differences between Basic and Ours are grouped into three techniques: layerwise prefill, dual-path loading, and scheduling algorithm.
|
||||
We add the techniques gradually to demonstrate individual contribution.
|
||||
As shown in Figure 12 , compared to Basic, adding layerwise prefill reduces JCT by 17.21 % 17.21\% on average, alleviating PE HBM bottlenecks and hiding transfer overhead.
|
||||
Adding Dual-path loading on top of layerwise prefill delivers the primary performance gains, reducing JCT by 38.19 % 38.19\% on average compared to Basic, as it enables requests to read KV-Cache from either PE or DE, fully utilizing distributed storage bandwidth.
|
||||
Finally, employing our scheduling algorithm on top of dual-path loading to decide KV-Cache loading paths achieves the best performance, reducing JCT by 45.62 % 45.62\% compared to Basic, demonstrating the effectiveness of load-balanced scheduling across storage NICs.
|
||||
Figure 13 . Load balance of storage NICs traffic
|
||||
Figure 14 . Load balance of attention execution time.
|
||||
Load Balance.
|
||||
DualPath’s scheduling algorithm improves load balance for both storage NICs and attention layer execution times.
|
||||
For storage NICs, our scheduling algorithm improves load balance from 1.53 to 1.18 compared to round robin scheduling ( Figure 13 ).
|
||||
For attention layers, DualPath maintains the Max/Avg ratio as low as 1.06 during the first 5% of the task, reducing GPU idle bubbles ( Figure 14 ).
|
||||
The storage NIC metric is the ratio of maximum to average traffic across all storage NICs on three machines within a small time window, where 1.0 represents perfect balance.
|
||||
The attention layer metric is calculated among all GPUs in an expert parallel group for each forward.
|
||||
As the task progresses, both ratios become meaningless due to underloaded system.
|
||||
Therefore, we do not show the tail phase of the workload.
|
||||
Table 3 . Large-scale experiment results.
|
||||
Setting
|
||||
JCT
|
||||
TTFT
|
||||
TTST
|
||||
TPOT
|
||||
Offline
|
||||
2P4D, 2K agents
|
||||
3,167s
|
||||
–
|
||||
–
|
||||
–
|
||||
48P96D, 48K agents
|
||||
3,201s
|
||||
–
|
||||
–
|
||||
–
|
||||
Online
|
||||
2P4D, 0.4 APS
|
||||
–
|
||||
1.739s
|
||||
0.228s
|
||||
0.039s
|
||||
44P88D, 8.8 APS
|
||||
–
|
||||
1.847s
|
||||
0.194s
|
||||
0.036s
|
||||
Figure 15 . 48P96D offline inference metrics. 1e7 is the scaling factor of Prompt TPS.
|
||||
## 7.6. Large-Scale Scalability
|
||||
We conduct both offline and online experiments using up to 1,152 GPUs to demonstrate large-scale scalability ( Table 3 ).
|
||||
For offline inference, scaling from 2P4D (2K agents) to 48P96D (48K agents) achieves near-linear speedup with comparable JCT (3,167s vs. 3,201s).
|
||||
For online serving, the 44P88D configuration achieves 22 × \times throughput (8.8 vs. 0.4 APS) while maintaining similar latency.
|
||||
Across all experiments, scheduler CPU usage remains below 10 cores, confirming it is not a bottleneck. Some detailed metrics over offline inference process are shown in Figure 15 .
|
||||
Due to the lack of fine-tuned parallelism settings and P/D ratios (which requires a substantial experimentation budget), the large-scale experiments do not demonstrate additional JCT or serving capacity gains compared to multiple small-scale units with equivalent cost.
|
||||
However, large-scale setting remains important for the following reasons. First, it reduces fragmentation and provides greater flexibility for fine-tuning parallelism and P/D ratios. Second, large-scale deployment offers more scheduling opportunities to mitigate queuing latency under unpredictable bursty online requests. These observations suggest several directions for future work ( § 8.1 ).
|
||||
## 8. Discussion
|
||||
## 8.1. Potential Future Work
|
||||
The workload of offline inference is highly dynamic. For example, in our agentic RL tasks, the workload depends heavily on the researchers’ algorithm design, and the prefill stage typically experiences significantly higher pressure in the first half of execution than in the second half. Meanwhile, profiling these tentative experiments is costly, as some experiments are only run a limited number of times. Therefore, more adaptive and flexible approaches for parallelism and P/D ratio configuration are needed, such as simulators or online adjustment mechanisms. Second, the scheduling algorithm still has room for improvement, as we expect to achieve lower TTFT percentiles under large-scale deployment.
|
||||
## 8.2. Working Set Analysis
|
||||
As shown in Figure 11 , given an arrival rate λ \lambda (i.e., new trajectories per second) and mean JCT T ¯ \bar{T} , the working set of the KV-Cache can thus be approximated as λ T ¯ × t o t a l _ l e n a v g / 2 \lambda\bar{T}\times total\_len_{avg}/2 .
|
||||
In our setting of DS 660B serving, this value of DualPath ranges from 69 GB at APS 0.1 to 681 GB at APS 0.45.
|
||||
In realistic setting, the working set would be larger since our evaluation assumes zero inter-arrival time and zero tool call latency.
|
||||
If JCT increases by r r times due to these gaps, the system’s APS capacity increases by r r times (gaps do not stress LLM inference), causing the working set to expand by r 2 r^{2} times.
|
||||
This would exceed available memory and reduce the distributed memory pool’s hit rate.
|
||||
Such experiments require r r times more machine hours and r 2 r^{2} times more storage (cost scaling as r 3 r^{3} ), which we cannot afford given limited resources.
|
||||
## 9. Related Work
|
||||
Distributed Memory Cache Pools.
|
||||
Mooncake ( Qin et al., 2025 ) builds a distributed DRAM pool for KV-Cache.
|
||||
TokenLake ( Wu et al., 2025 ) introduces a unified segment-level prefix cache pool.
|
||||
Compared to them, DualPath targets storage backend directly, balancing the traffic among all SNICs, and reduces DRAM usage greatly without harming performance.
|
||||
DualPath can also be combined with a middle DRAM cache, but the performance gain is marginal.
|
||||
KV-Cache I/O Optimization.
|
||||
Efficiently loading the massive KV-Cache from other caching tiers is a fundamental bottleneck in disaggregated LLM serving architectures.
|
||||
Prior work has approached this problem primarily from the perspective of a single data path.
|
||||
Strata ( Xie et al., 2025 ) tackles I/O bottlenecks in hierarchical storage by co-designing GPU-assisted I/O with cache-aware scheduling.
|
||||
Other work, like KVPR ( Jiang et al., 2025 ) and TailorKV ( Yao et al., 2025 ) , mitigates bandwidth constraints (e.g., PCIe) on this path via recomputation overlapping and layer-granular hybrid quantization.
|
||||
LLM Inference System.
|
||||
Recent years have seen many inference acceleration techniques, such as paged attention ( Kwon et al., 2023 ) , chunked prefill, and hybrid batching ( Agrawal et al., 2024 ; Holmes et al., 2024 ) . Prefill-decode disaggregated inference ( Patel et al., 2025 ; Zhong et al., 2024 ) separates the prefill and decode stages onto different GPUs, reducing performance interference between them and allowing each stage to adopt distinct parallel strategies and hardware configurations, which unlocks substantial optimization opportunities. It has largely become the de facto standard for inference serving.
|
||||
Attention Mechanisms.
|
||||
Attention mechanism allows tokens to interact with previous tokens in the sequence.
|
||||
There are many variants such as Multi-Head Attention (MHA) ( Vaswani et al., 2017 ) , Multi-Query Attention (MQA) ( Shazeer, 2019 ) and Grouped-Query Attention (GQA) ( Ainslie et al., 2023 ) , Multi-head Latent Attention (MLA) ( DeepSeek-AI, 2024 ) .
|
||||
For those attention mechanisms (denoted as dense attention ), the ratio of computation and KV-Cache size for one token is a constant since both scale linearly with sequence length.
|
||||
## 10. Conclusion
|
||||
This paper presents DualPath, an agentic LLM inference framework that addresses the imbalance of KV-Cache reading under PD-disaggregated architecture through dual-path KV-Cache loading. By redistributing storage network load with workload-aware scheduling, DualPath achieves up to 1.87 × \times throu
|
||||
|
||||
=== BODY (tail) ===
|
||||
ghput improvement for offline inference. It also achieves 1.96 × \times higher agent runs per second on average in online serving.
|
||||
@@ -0,0 +1,876 @@
|
||||
TITLE: DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence
|
||||
ARXIV: 2606.19348
|
||||
AUTHORS: DeepSeek-AI , Anyi Xu , Bangcai Lin , Bing Xue , Bingxuan Wang , Bingzheng Xu , Bochao Wu , Bowei Zhang , Chaofan Lin , Chen Dong , Chenchen Ling , Chengda Lu , Chenggang Zhao , Chengqi Deng , Chengyu Hou , Chenhao Xu , Chenze Shao , Chong Ruan , Conner Sun , Damai Dai , Daya Guo , Dejian Yang , Deli Chen , Donghao Li , Dongjie Ji , Erhang Li , Fang Wei , Fangyun Lin , Fangzhou Yuan , Feiyu Xia , Fucong Dai , Guangbo Hao , Guanting Chen , Guoai Cao , Guolai Meng , Guowei Li , Han Yu , Han Zhang , Hanwei Xu , Hao Li , Haofen Liang , Haoling Zhang , Haoming Luo , Haoran Wei , Haotian Yuan , Haowei Zhang , Haowen Luo , Haoyu Chen , Haozhe Ji , Hengqing Zhang , Honghui Ding , Hongxuan Tang , Huanqi Cao , Huazuo Gao , Hui Qu , Hui Zeng , J Yang , JQ Zhu , Jia Luo , Jia Song , Jia Yu , Jialiang Huang , Jialu Cai , Jian Liang , Jiangting Zhou , Jiasheng Ye , Jiashi Li , Jiaxin Xu , Jiewen Hu , Jieyu Yang , Jin Chen , Jin Yan , Jingchang Chen , Jingli Zhou , Jingting Xiang , Jingyang Yuan , Jingyuan Cheng , Jingzi Zhou , Jinhua Zhu , Jiping Yu , Joseph Sun , Jun Ran , Junguang Jiang , Junjie Qiu , Junlong Li , Junmin Zheng , Junxiao Song , Kai Dong , Kaige Gao , Kang Guan , Kexing Zhou , Kezhao Huang , Kuai Yu , Lean Wang , Lecong Zhang , Lei Wang , Leyi Xia , Li Zhang , Liang Zhao , Lihua Guo , Lingxiao Luo , Linwang Ma , Linyan Zhu , Litong Wang , Liyu Cai , Liyue Zhang , Longhao Chen , MS Di , MY Xu , Max Mei , Miaojun Wang , Mingchuan Zhang , Minghua Zhang , Minghui Tang , Mingmi
|
||||
DATE: [Submitted on 26 Apr 2026]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## DeepSeek-V4:
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Architecture
|
||||
## 2.1 Designs Inherited from DeepSeek-V3
|
||||
## Mixture-of-Experts.
|
||||
## Multi-Token Prediction.
|
||||
## 2.2 Manifold-Constrained Hyper-Connections
|
||||
## Standard Hyper-Connections.
|
||||
## Manifold-Constrained Residual Mapping.
|
||||
## Dynamic Parameterization.
|
||||
## Applying Parameter Constraints.
|
||||
## 2.3 Hybrid Attention with CSA and HCA
|
||||
## 2.3.1 Compressed Sparse Attention
|
||||
## Compressed Key-Value Entries.
|
||||
## Lightning Indexer for Sparse Selection.
|
||||
## Shared Key-Value MQA.
|
||||
## Grouped Output Projection.
|
||||
## 2.3.2 Heavily Compressed Attention
|
||||
## Compressed Key-Value Entries.
|
||||
## Shared Key-Value MQA and Grouped Output Projection.
|
||||
## 2.3.3 Other Details
|
||||
## Query and Key-Value Entry Normalization.
|
||||
## Partial Rotary Positional Embedding.
|
||||
## Additional Branch of Sliding Window Attention.
|
||||
## Attention Sink.
|
||||
## 2.3.4 Efficiency Discussion
|
||||
## 2.4 Muon Optimizer
|
||||
## Basic Configurations.
|
||||
## Hybrid Newton-Schulz Iterations.
|
||||
## Avoiding Exploding Attention Logits.
|
||||
## 3 General Infrastructures
|
||||
## 3.1 Fine-Grained Communication-Computation Overlap in Expert Parallelism
|
||||
## Communication Latency Can Be Hidden.
|
||||
## Fine-Grained EP Scheme.
|
||||
## Performance and Open-Sourced Mega-Kernel.
|
||||
## Observations and Proposals.
|
||||
## 3.2 Flexible and Efficient Kernel Development with TileLang
|
||||
## Reducing Invocation Overhead with Host Codegen.
|
||||
## SMT-Solver-Assisted Formal Integer Analysis.
|
||||
## Numerical Precision and Bitwise Reproducibility.
|
||||
## 3.3 High-Performance Batch-Invariant and Deterministic Kernel Libraries
|
||||
## Batch Invariance.
|
||||
## Determinism.
|
||||
## 3.4 Training Framework
|
||||
## 3.4.1 Efficient Implementation of Muon
|
||||
## 3.4.2 Cost-Effective and Memory-Efficient Implementation of m HC
|
||||
## 3.4.3 Contextual Parallelism for Long-Context Attention
|
||||
## 3.4.4 Extended Automatic Differentiation for Flexible Activation Checkpointing
|
||||
## 3.5 Inference Framework
|
||||
## 3.5.1 KV Cache Structure and Management
|
||||
## Heterogeneous KV Entries in DeepSeek-V4.
|
||||
## Challenges in Managing Hybrid Attention KV Cache.
|
||||
## State Cache for SWA and Uncompressed Tail Tokens.
|
||||
## Sparse Attention Kernel Co-Design.
|
||||
## 3.5.2 On-Disk KV Cache Storage
|
||||
## 4 Pre-Training
|
||||
## 4.1 Data Construction
|
||||
## 4.2 Pre-Training Setups
|
||||
## 4.2.1 Model Setups
|
||||
## DeepSeek-V4-Flash.
|
||||
## DeepSeek-V4-Pro.
|
||||
## 4.2.2 Training Setups
|
||||
## DeepSeek-V4-Flash.
|
||||
## DeepSeek-V4-Pro.
|
||||
## 4.2.3 Mitigating Training Instability
|
||||
## Anticipatory Routing.
|
||||
## SwiGLU Clamping.
|
||||
## 4.3 Evaluations
|
||||
## 4.3.1 Evaluation Benchmarks
|
||||
## 4.3.2 Evaluation Results
|
||||
## 5 Post-Training
|
||||
## 5.1 Post-Training Pipeline
|
||||
## 5.1.1 Specialist Training
|
||||
## Reasoning Efforts.
|
||||
## Generative Reward Model.
|
||||
## Tools
|
||||
## Tool-Call Schema and Special Token.
|
||||
## Interleaved Thinking.
|
||||
## Quick Instruction.
|
||||
## 5.1.2 On-Policy Distillation
|
||||
## 5.2 Post-Training Infrastructures
|
||||
## 5.2.1 FP4 Quantization-Aware Training
|
||||
## 5.2.2 Efficient Teacher Scheduling for Full-Vocabulary OPD
|
||||
## 5.2.3 Preemptible and Fault-Tolerant Rollout Service
|
||||
## 5.2.4 Scaling RL Framework for Million-Token Context
|
||||
## 5.2.5 Sandbox Infrastructure for Agentic AI
|
||||
## Four Execution Substrates Behind One Unified Interface.
|
||||
## Fast Image Loading via Layered Storage.
|
||||
## Density Optimizations Under Massive Concurrency.
|
||||
## Trajectory Logging and Preemption-Safe Resumption.
|
||||
## 5.3 Standard Benchmark Evaluation
|
||||
## 5.3.1 Evaluation Setup
|
||||
## Knowledge and Reasoning.
|
||||
## 1M-Token Context.
|
||||
## Agent.
|
||||
## 5.3.2 Evaluation Results
|
||||
## Knowledge.
|
||||
## Reasoning.
|
||||
## Agent.
|
||||
## 1M-Token Context.
|
||||
## Reasoning Effort.
|
||||
## 5.4 Performance on Real-World Tasks
|
||||
## 5.4.1 Chinese Writing
|
||||
## 5.4.2 Search
|
||||
## Retrieval Augmented Search.
|
||||
## Agentic Search.
|
||||
## 5.4.3 White-Collar Task
|
||||
## 5.4.4 Code Agent
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Designs Inherited from DeepSeek-V3
|
||||
Mixture-of-Experts.
|
||||
Multi-Token Prediction.
|
||||
2.2 Manifold-Constrained Hyper-Connections
|
||||
Standard Hyper-Connections.
|
||||
Manifold-Constrained Residual Mapping.
|
||||
Dynamic Parameterization.
|
||||
Applying Parameter Constraints.
|
||||
2.3 Hybrid Attention with CSA and HCA
|
||||
2.3.1 Compressed Sparse Attention
|
||||
Compressed Key-Value Entries.
|
||||
Lightning Indexer for Sparse Selection.
|
||||
Shared Key-Value MQA.
|
||||
Grouped Output Projection.
|
||||
2.3.2 Heavily Compressed Attention
|
||||
Compressed Key-Value Entries.
|
||||
Shared Key-Value MQA and Grouped Output Projection.
|
||||
2.3.3 Other Details
|
||||
Query and Key-Value Entry Normalization.
|
||||
Partial Rotary Positional Embedding.
|
||||
Additional Branch of Sliding Window Attention.
|
||||
Attention Sink.
|
||||
2.3.4 Efficiency Discussion
|
||||
2.4 Muon Optimizer
|
||||
Basic Configurations.
|
||||
Hybrid Newton-Schulz Iterations.
|
||||
Avoiding Exploding Attention Logits.
|
||||
3 General Infrastructures
|
||||
3.1 Fine-Grained Communication-Computation Overlap in Expert Parallelism
|
||||
Communication Latency Can Be Hidden.
|
||||
Fine-Grained EP Scheme.
|
||||
Performance and Open-Sourced Mega-Kernel.
|
||||
Observations and Proposals.
|
||||
3.2 Flexible and Efficient Kernel Development with TileLang
|
||||
Reducing Invocation Overhead with Host Codegen.
|
||||
SMT-Solver-Assisted Formal Integer Analysis.
|
||||
Numerical Precision and Bitwise Reproducibility.
|
||||
3.3 High-Performance Batch-Invariant and Deterministic Kernel Libraries
|
||||
Batch Invariance.
|
||||
Determinism.
|
||||
3.4 Training Framework
|
||||
3.4.1 Efficient Implementation of Muon
|
||||
3.4.2 Cost-Effective and Memory-Efficient Implementation of m HC
|
||||
3.4.3 Contextual Parallelism for Long-Context Attention
|
||||
3.4.4 Extended Automatic Differentiation for Flexible Activation Checkpointing
|
||||
3.5 Inference Framework
|
||||
3.5.1 KV Cache Structure and Management
|
||||
Heterogeneous KV Entries in DeepSeek-V4.
|
||||
Challenges in Managing Hybrid Attention KV Cache.
|
||||
State Cache for SWA and Uncompressed Tail Tokens.
|
||||
Sparse Attention Kernel Co-Design.
|
||||
3.5.2 On-Disk KV Cache Storage
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
4.2 Pre-Training Setups
|
||||
4.2.1 Model Setups
|
||||
DeepSeek-V4-Flash.
|
||||
DeepSeek-V4-Pro.
|
||||
4.2.2 Training Setups
|
||||
DeepSeek-V4-Flash.
|
||||
DeepSeek-V4-Pro.
|
||||
4.2.3 Mitigating Training Instability
|
||||
Anticipatory Routing.
|
||||
SwiGLU Clamping.
|
||||
4.3 Evaluations
|
||||
4.3.1 Evaluation Benchmarks
|
||||
4.3.2 Evaluation Results
|
||||
5 Post-Training
|
||||
5.1 Post-Training Pipeline
|
||||
5.1.1 Specialist Training
|
||||
Reasoning Efforts.
|
||||
Generative Reward Model.
|
||||
Tool-Call Schema and Special Token.
|
||||
Interleaved Thinking.
|
||||
Quick Instruction.
|
||||
5.1.2 On-Policy Distillation
|
||||
5.2 Post-Training Infrastructures
|
||||
5.2.1 FP4 Quantization-Aware Training
|
||||
5.2.2 Efficient Teacher Scheduling for Full-Vocabulary OPD
|
||||
5.2.3 Preemptible and Fault-Tolerant Rollout Service
|
||||
5.2.4 Scaling RL Framework for Million-Token Context
|
||||
5.2.5 Sandbox Infrastructure for Agentic AI
|
||||
Four Execution Substrates Behind One Unified Interface.
|
||||
Fast Image Loading via Layered Storage.
|
||||
Density Optimizations Under Massive Concurrency.
|
||||
Trajectory Logging and Preemption-Safe Resumption.
|
||||
5.3 Standard Benchmark Evaluation
|
||||
5.3.1 Evaluation Setup
|
||||
Knowledge and Reasoning.
|
||||
1M-Token Context.
|
||||
Agent.
|
||||
5.3.2 Evaluation Results
|
||||
Knowledge.
|
||||
Reasoning.
|
||||
Agent.
|
||||
1M-Token Context.
|
||||
Reasoning Effort.
|
||||
5.4 Performance on Real-World Tasks
|
||||
5.4.1 Chinese Writing
|
||||
5.4.2 Search
|
||||
Retrieval Augmented Search.
|
||||
Agentic Search.
|
||||
5.4.3 White-Collar Task
|
||||
5.4.4 Code Agent
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Author List and Acknowledgment
|
||||
A.1 Author List
|
||||
A.2 Acknowledgment
|
||||
B Evaluation Details
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2606.19348v1 [cs.CL] 26 Apr 2026
|
||||
\reportnumber
|
||||
001
|
||||
## DeepSeek-V4:
|
||||
Towards Highly Efficient Million-Token Context Intelligence
|
||||
DeepSeek-AI
|
||||
research@deepseek.com
|
||||
## Abstract
|
||||
We present a preview version of DeepSeek-V4 series, including two strong Mixture-of-Experts (MoE) language models — DeepSeek-V4-Pro with 1.6T parameters (49B activated) and DeepSeek-V4-Flash with 284B parameters (13B activated) — both supporting a context length of one million tokens.
|
||||
DeepSeek-V4 series incorporate several key upgrades in architecture and optimization:
|
||||
(1) a hybrid attention architecture that combines Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA) to improve long-context efficiency;
|
||||
(2) Manifold-Constrained Hyper-Connections ( m HC) that enhance conventional residual connections;
|
||||
(3) and the Muon optimizer for faster convergence and greater training stability.
|
||||
We pre-train both models on more than 32T diverse and high-quality tokens, followed by a comprehensive post-training pipeline that unlocks and further enhances their capabilities.
|
||||
DeepSeek-V4-Pro-Max, the maximum reasoning effort mode of DeepSeek-V4-Pro, redefines the state-of-the-art for open models, outperforming its predecessors in core tasks.
|
||||
Meanwhile, DeepSeek-V4 series are highly efficient in long-context scenarios.
|
||||
In the one-million-token context setting, DeepSeek-V4-Pro requires only 27% of single-token inference FLOPs and 10% of KV cache compared with DeepSeek-V3.2.
|
||||
This enables us to routinely support one-million-token contexts, thereby making long-horizon tasks and further test-time scaling more feasible.
|
||||
The model checkpoints are available at https://huggingface.co/collections/deepseek-ai/deepseek-v4 .
|
||||
Figure 1:
|
||||
Left : benchmark performance of DeepSeek-V4-Pro-Max and its counterparts.
|
||||
Right : inference FLOPs and KV cache size of DeepSeek-V4 series and DeepSeek-V3.2.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Designs Inherited from DeepSeek-V3
|
||||
2.2 Manifold-Constrained Hyper-Connections
|
||||
2.3 Hybrid Attention with CSA and HCA
|
||||
2.3.1 Compressed Sparse Attention
|
||||
2.3.2 Heavily Compressed Attention
|
||||
2.3.3 Other Details
|
||||
2.3.4 Efficiency Discussion
|
||||
2.4 Muon Optimizer
|
||||
3 General Infrastructures
|
||||
3.1 Fine-Grained Communication-Computation Overlap in Expert Parallelism
|
||||
3.2 Flexible and Efficient Kernel Development with TileLang
|
||||
3.3 High-Performance Batch-Invariant and Deterministic Kernel Libraries
|
||||
3.4 Training Framework
|
||||
3.4.1 Efficient Implementation of Muon
|
||||
3.4.2 Cost-Effective and Memory-Efficient Implementation of m HC
|
||||
3.4.3 Contextual Parallelism for Long-Context Attention
|
||||
3.4.4 Extended Automatic Differentiation for Flexible Activation Checkpointing
|
||||
3.5 Inference Framework
|
||||
3.5.1 KV Cache Structure and Management
|
||||
3.5.2 On-Disk KV Cache Storage
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
4.2 Pre-Training Setups
|
||||
4.2.1 Model Setups
|
||||
4.2.2 Training Setups
|
||||
4.2.3 Mitigating Training Instability
|
||||
4.3 Evaluations
|
||||
4.3.1 Evaluation Benchmarks
|
||||
4.3.2 Evaluation Results
|
||||
5 Post-Training
|
||||
5.1 Post-Training Pipeline
|
||||
5.1.1 Specialist Training
|
||||
5.1.2 On-Policy Distillation
|
||||
5.2 Post-Training Infrastructures
|
||||
5.2.1 FP4 Quantization-Aware Training
|
||||
5.2.2 Efficient Teacher Scheduling for Full-Vocabulary OPD
|
||||
5.2.3 Preemptible and Fault-Tolerant Rollout Service
|
||||
5.2.4 Scaling RL Framework for Million-Token Context
|
||||
5.2.5 Sandbox Infrastructure for Agentic AI
|
||||
5.3 Standard Benchmark Evaluation
|
||||
5.3.1 Evaluation Setup
|
||||
5.3.2 Evaluation Results
|
||||
5.4 Performance on Real-World Tasks
|
||||
5.4.1 Chinese Writing
|
||||
5.4.2 Search
|
||||
5.4.3 White-Collar Task
|
||||
5.4.4 Code Agent
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Author List and Acknowledgment
|
||||
A.1 Author List
|
||||
A.2 Acknowledgment
|
||||
B Evaluation Details
|
||||
## 1 Introduction
|
||||
The emergence of reasoning models ( OpenAI, 2024c ; DeepSeek-AI, 2025a ) has established a new paradigm of test-time scaling, driving substantial performance gains for Large Language Models (LLMs).
|
||||
However, this scaling paradigm is fundamentally constrained by the quadratic computational complexity of the vanilla attention mechanism ( Vaswani et al., 2017 ) , which creates a prohibitive bottleneck for ultra-long contexts and reasoning processes.
|
||||
Concurrently, the emergence of long-horizon scenarios and tasks — from complex agentic workflows to massive cross-document analysis — has also made efficient support for ultra-long contexts critical for future progress.
|
||||
While recent open-source efforts ( DeepSeek-AI, 2024c ; Qwen, 2025 ; MiniMax, 2025 ; Bai et al., 2025a ) have advanced general capabilities, this core architectural inefficiency in handling ultra-long sequences remains a key impediment, limiting further gains from test-time scaling and hindering further exploration into long-horizon scenarios and tasks.
|
||||
In order to break the efficiency barrier in ultra-long contexts, we develop the DeepSeek-V4 series, including the preview versions of DeepSeek-V4-Pro with 1.6T parameters (49B activated) and DeepSeek-V4-Flash with 284B parameters (13B activated).
|
||||
Through architectural innovations, DeepSeek-V4 series achieve a dramatic leap in computational efficiency for processing ultra-long sequences.
|
||||
This breakthrough enables efficient support for a context length of one million tokens, ushering in a new era of million-length contexts for next-generation LLMs.
|
||||
We believe our capability to efficiently handle ultra-long sequences unlocks the next frontier of test-time scaling, paves the way for deeper research into long-horizon tasks, and establishes a necessary foundation for exploring future paradigms like online learning.
|
||||
Compared with the DeepSeek-V3 architecture ( DeepSeek-AI, 2024c ) , DeepSeek-V4 series retain the DeepSeekMoE framework ( Dai et al., 2024 ) and Multi-Token Prediction (MTP) strategy, while introducing several key innovations in architecture and optimization.
|
||||
To enhance long-context efficiency, we design a hybrid attention mechanism combining Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA).
|
||||
CSA compresses the KV caches along the sequence dimension and then performs DeepSeek Sparse Attention (DSA) ( DeepSeek-AI, 2025b ) , whereas HCA applies more aggressive compression to the KV caches but keeps dense attention.
|
||||
To strengthen modeling capability, we incorporate Manifold-Constrained Hyper-Connections ( m HC) ( Xie et al., 2026 ) that upgrade conventional residual connections.
|
||||
Additionally, we introduce the Muon ( Jordan et al., 2024 ; Liu et al., 2025 ) optimizer to the training of DeepSeek-V4 series, leading to faster convergence and improved training stability.
|
||||
To enable efficient training and inference for DeepSeek-V4 series as well as productive development, we introduce several infrastructure optimizations.
|
||||
First, we design and implement a single fused kernel for MoE modules that fully overlaps computation, communication, and memory access.
|
||||
Second, we employ TileLang ( Wang et al., 2026 ) , a Domain-Specific Language (DSL) to balance development productivity and runtime efficiency.
|
||||
Third, we provide efficient batch-invariant and deterministic kernel libraries to ensure bitwise reproducibility across training and inference.
|
||||
Fourth, for the training framework, we extend the autograd framework with tensor-level checkpointing for fine-grained recomputation control; and we enhance training efficiency with a hybrid ZeRO strategy for the Muon optimizer, cost-effective m HC implementations via recomputation and fused kernels, and two-stage contextual parallelism to manage compressed attention.
|
||||
Fifth, for the inference framework, we design a heterogeneous KV cache structure with on-disk storage strategies to enable efficient shared-prefix reuse.
|
||||
In addition, during the post-training stage, we incorporate FP4 quantization-aware training for MoE expert weights and the indexer QK path to reduce memory and computation.
|
||||
By employing hybrid CSA and HCA, along with precision optimizations on computation and storage, DeepSeek-V4 series achieve significantly lower inference FLOPs and a substantially reduced KV cache size compared with DeepSeek-V3.2, especially in long-context settings.
|
||||
The right part of Figure 1 demonstrates the estimated single-token inference FLOPs and accumulated KV cache size of DeepSeek-V3.2 and DeepSeek-V4 series.
|
||||
In the scenario of 1M-token context, even DeepSeek-V4-Pro, which has a larger number of activated parameters, attains only 27% of the single-token FLOPs (measured in equivalent FP8 FLOPs) and 10% of the KV cache size relative to DeepSeek-V3.2.
|
||||
Furthermore, DeepSeek-V4-Flash, with its smaller number of activated parameters, pushes efficiency even further: in the 1M-token context setting, it achieves only 10% of the single-token FLOPs and 7% of the KV cache size compared with DeepSeek-V3.2.
|
||||
Additionally, for DeepSeek-V4 series, the routed expert parameters utilize FP4 precision.
|
||||
While the peak FLOPs for FP4 × \times FP8 operations are currently the same as FP8 × \times FP8 on existing hardware, they can theoretically be implemented to be 1 / 3 1/3 more efficient on future hardware, which will further enhance the efficiency of DeepSeek-V4 series.
|
||||
During pre-training, we train DeepSeek-V4-Flash on 32T tokens and DeepSeek-V4-Pro on 33T tokens, respectively.
|
||||
After pre-training, these two models can natively and efficiently support 1M-length contexts.
|
||||
In our internal evaluations, DeepSeek-V4-Flash-Base already surpasses DeepSeek-V3.2-Base across a majority of benchmarks with its more parameter-efficient design.
|
||||
DeepSeek-V4-Pro-Base further extends this advantage to set a new performance standard among DeepSeek foundation models, achieving comprehensive superiority across reasoning, coding, long-context, and world knowledge tasks.
|
||||
The post-training pipeline of DeepSeek-V4 series features a two-stage paradigm: the independent cultivation of domain-specific experts, followed by unified model consolidation via on-policy distillation ( Lu and Lab, 2025 ; Gu et al., 2024 ) .
|
||||
Initially, for each target domain — such as mathematics, coding, agent, and instruction following — a separate expert model is trained independently.
|
||||
The base model first undergoes Supervised Fine-Tuning (SFT) on high-quality, domain-specific data to establish foundational capabilities.
|
||||
Subsequently, Reinforcement Learning (RL) is applied using Group Relative Policy Optimization (GRPO) ( DeepSeek-AI, 2025a ) , which further optimizes the model for domain-aligned behaviors guided by reward models tailored to specific success criteria.
|
||||
This phase yields a diverse set of specialized experts, each excelling in its respective field.
|
||||
Finally, to integrate these distinct proficiencies, a single unified model is trained through on-policy distillation, wherein the unified model acts as the student learning to optimize the reverse KL loss with teacher models.
|
||||
Summary of Core Evaluation Results
|
||||
•
|
||||
Knowledge :
|
||||
In assessments of broad world knowledge, DeepSeek-V4-Pro-Max, the maximum reasoning effort mode of DeepSeek-V4-Pro, significantly outperforms leading open-source models on the SimpleQA ( OpenAI, 2024a ) and Chinese-SimpleQA ( He et al., 2024 ) benchmarks.
|
||||
Regarding educational knowledge — evaluated via MMLU-Pro ( Wang et al., 2024b ) , HLE ( Phan et al., 2025 ) , and GPQA ( Rein et al., 2023 ) — DeepSeek-V4-Pro-Max shows a marginal lead over its open-source counterparts.
|
||||
DeepSeek-V4-Pro-Max has significantly closed the gap with the leading proprietary model, Gemini-3.1-Pro, despite still trailing it in these knowledge-based evaluations.
|
||||
•
|
||||
Reasoning :
|
||||
Through the expansion of reasoning tokens, DeepSeek-V4-Pro-Max demonstrates superior performance relative to GPT-5.2 and Gemini-3.0-Pro on standard reasoning benchmarks.
|
||||
Nevertheless, its performance falls marginally short of GPT-5.4 and Gemini-3.1-Pro, suggesting a developmental trajectory that trails state-of-the-art frontier models by approximately 3 to 6 months.
|
||||
Furthermore, DeepSeek-V4-Flash-Max achieves comparable performance to GPT-5.2 and Gemini-3.0-Pro, establishing itself as a highly cost-effective architecture for complex reasoning tasks.
|
||||
•
|
||||
Agent :
|
||||
On public benchmarks, DeepSeek-V4-Pro-Max is on par with leading open-source models, such as Kimi-K2.6 and GLM-5.1, but slightly worse than frontier closed models.
|
||||
In our internal evaluation, DeepSeek-V4-Pro-Max outperforms Claude Sonnet 4.5 and approaches the level of Opus 4.5.
|
||||
•
|
||||
Long-Context : DeepSeek-V4-Pro-Max delivers strong results on synthetic and real use cases with a 1-million-token context window, surpassing even Gemini-3.1-Pro on academic benchmarks.
|
||||
•
|
||||
DeepSeek-V4-Pro v.s. DeepSeek-V4-Flash :
|
||||
DeepSeek-V4-Flash-Max exhibits lower performance in knowledge evaluations due to its smaller parameter scale.
|
||||
However, it achieves comparable results on reasoning tasks when allocated a larger thinking budget.
|
||||
In agent evaluations, while DeepSeek-V4-Flash-Max matches the performance of DeepSeek-V4-Pro-Max on several benchmarks, it still trails its larger counterpart on more complex, high-difficulty tasks.
|
||||
Figure 2:
|
||||
Overall architecture of DeepSeek-V4 series.
|
||||
We use hybrid CSA (Compressed Sparse Attention) and HCA (Heavily Compressed Attention) for attention layers, DeepSeekMoE for feed-forward layers, and strengthen conventional residual connections with m HC.
|
||||
## 2 Architecture
|
||||
Overall, DeepSeek-V4 series retain the Transformer ( Vaswani et al., 2017 ) architecture and Multi-Token Prediction (MTP) modules ( Gloeckle et al., 2024 ; DeepSeek-AI, 2024c ) , while introducing several key upgrades over DeepSeek-V3:
|
||||
(1) firstly, we introduce the Manifold-Constrained Hyper-Connections ( m HC) ( Xie et al., 2026 ) to strengthen conventional residual connections;
|
||||
(2) secondly, we design a hybrid attention architecture, which greatly improves long-context efficiency through Compressed Sparse Attention and Heavily Compressed Attention.
|
||||
(3) thirdly, we employ Muon ( Jordan et al., 2024 ; Liu et al., 2025 ) as the optimizer.
|
||||
For the Mixture-of-Experts (MoE) components, we still adopt the DeepSeekMoE ( Dai et al., 2024 ) architecture, with only minor adjustments from DeepSeek-V3.
|
||||
The Multi-Token Prediction (MTP) ( Qi et al., 2020 ; Gloeckle et al., 2024 ; Li et al., 2024 ; DeepSeek-AI, 2024c ) configuration remains identical to that of DeepSeek-V3.
|
||||
All other unspecified details follow the settings established in DeepSeek-V3 ( DeepSeek-AI, 2024c ) .
|
||||
Figure 2 illustrates the overall architecture of DeepSeek-V4, and the details are described below.
|
||||
## 2.1 Designs Inherited from DeepSeek-V3
|
||||
## Mixture-of-Experts.
|
||||
As previous DeepSeek-series models ( DeepSeek-AI, 2024b ; DeepSeek-AI, 2024c ) , DeepSeek-V4 series also adopt the DeepSeekMoE paradigm ( Dai et al., 2024 ) for Feed-Forward Networks (FFNs), which sets fine-grained routed experts and shared experts.
|
||||
Different from DeepSeek-V3, we change the activation function that computes the affinity scores from Sigmoid ( ⋅ ) \operatorname{Sigmoid}(\cdot) into Sqrt ( Softplus ( ⋅ ) ) \operatorname{Sqrt}(\operatorname{Softplus}(\cdot)) .
|
||||
For load balancing, we also employ the auxiliary-loss-free strategy ( Wang et al., 2024a ; DeepSeek-AI, 2024c ) , augmented by a slight sequence-wise balance loss that prevents extreme imbalance within individual sequences.
|
||||
For DeepSeek-V4, we remove the constraint on the number of routing target nodes, and carefully redesign the parallelism strategy to maintain training efficiency.
|
||||
Furthermore, compared with DeepSeek-V3, we replace the dense FFN layers in the initial several Transformer blocks with MoE layers that employ Hash routing ( Roller et al., 2021 ) .
|
||||
The Hash routing strategy determines the target experts of each token according to a predefined hash function with regard to the input token ID.
|
||||
## Multi-Token Prediction.
|
||||
As DeepSeek-V3, DeepSeek-V4 series also set MTP modules and objectives.
|
||||
Given that the MTP strategy has been validated in DeepSeek-V3, we adopt the same strategy for DeepSeek-V4 series without modification.
|
||||
## 2.2 Manifold-Constrained Hyper-Connections
|
||||
As shown in Figure 2 , DeepSeek-V4 series incorporate Manifold-Constrained Hyper-Connections ( m HC) ( Xie et al., 2026 ) to strengthen the conventional residual connections between adjacent Transformer blocks.
|
||||
Compared with naive Hyper-Connections (HC) ( Zhu et al., 2025 ) , the core idea of m HC is to constrain the residual mapping onto a specific manifold, and thus enhance the stability of signal propagation across layers while preserving model expressivity.
|
||||
This subsection briefly introduces the standard HC and describes how we design m HC for stable training.
|
||||
## Standard Hyper-Connections.
|
||||
The standard HC expands the width of the residual stream by a factor of n hc n_{\text{hc}} .
|
||||
Specifically, the shape of the residual stream is expanded from ℝ d \mathbb{R}^{d} to ℝ n hc × d \mathbb{R}^{n_{\text{hc}}\times d} , where d d is the hidden size of the actual layer input.
|
||||
Let X l = [ 𝐱 l , 1 ; … ; 𝐱 l , n hc ] T ∈ ℝ n hc × d X_{l}=[\mathbf{x}_{l,1};\ldots;\mathbf{x}_{l,n_{\text{hc}}}]^{T}\in\mathbb{R}^{n_{\text{hc}}\times d} be the residual state before the l l -th layer.
|
||||
HC introduces three linear mappings: an input mapping A l ∈ ℝ 1 × n hc A_{l}\in\mathbb{R}^{1\times n_{\text{hc}}} , a residual transformation B l ∈ ℝ n hc × n hc B_{l}\in\mathbb{R}^{n_{\text{hc}}\times n_{\text{hc}}} , and an output mapping C l ∈ ℝ n hc × 1 C_{l}\in\mathbb{R}^{n_{\text{hc}}\times 1} .
|
||||
The update of the residual state is then formulated as:
|
||||
X l + 1 = B l X l + C l ℱ l ( A l X l ) , X_{l+1}=B_{l}X_{l}+C_{l}\mathcal{F}_{l}(A_{l}X_{l}),
|
||||
(1)
|
||||
where ℱ l \mathcal{F}_{l} denotes the l l -th layer (e.g., an MoE layer), whose input and output shapes are both ℝ d \mathbb{R}^{d} .
|
||||
Note that the actual layer input A l X l ∈ ℝ d A_{l}X_{l}\in\mathbb{R}^{d} is also d d -dimensional, so the expanded residual width does not influence the design of the inner layers.
|
||||
HC decouples the residual width from the actual hidden size, offering a complementary scaling axis with minimal computational overhead, as n hc n_{\text{hc}} is typically much smaller than the hidden size d d .
|
||||
However, even though HC has demonstrated potential in improving model performance, we find that the training will frequently exhibit numerical instability when stacking multiple layers, which hinders the scaling of HC.
|
||||
## Manifold-Constrained Residual Mapping.
|
||||
The core innovation of m HC is to constrain the residual mapping matrix B l B_{l} to the manifold of doubly stochastic matrices (the Birkhoff polytope) ℳ \mathcal{M} , and thus enhance the stability of signal propagation across layers:
|
||||
B l ∈ ℳ :- { M ∈ ℝ n × n ∣ M 𝟏 n = 𝟏 n , 1 n T M = 𝟏 n T , M ⩾ 0 } . B_{l}\in\mathcal{M}\coloneq\{M\in\mathbb{R}^{n\times n}\mid M\mathbf{1}_{n}=\mathbf{1}_{n},\;\mathbf{1}_{n}^{T}M=\mathbf{1}_{n}^{T},\;M\geqslant 0\}.
|
||||
(2)
|
||||
This constraint ensures that the spectral norm of the mapping matrix ‖ B l ‖ 2 \|B_{l}\|_{2} is bounded by 1, so the residual transformation is non-expansive, which increases the numerical stability during both the forward pass and backpropagation.
|
||||
Besides, the set ℳ \mathcal{M} is closed under multiplication, which guarantees stability in the scenarios of deep stacks of m HC.
|
||||
In addition, the input transformation A l A_{l} and output transformation C l C_{l} are also constrained to be non-negative and bounded via a Sigmoid function to avoid the risk of signal cancellation.
|
||||
## Dynamic Parameterization.
|
||||
The parameters of three linear mappings are dynamically generated, which are decomposed into a dynamic (input-dependent) component and a static (input-independent) component.
|
||||
Given the input X l ∈ ℝ n hc × d X_{l}\in\mathbb{R}^{n_{\text{hc}}\times d} , it is first flattened and normalized: X ^ l = RMSNorm ( vec ( X l ) ) ∈ ℝ 1 × n hc d \hat{X}_{l}=\operatorname{RMSNorm}(\operatorname{vec}(X_{l}))\in\mathbb{R}^{1\times n_{\text{hc}}d} .
|
||||
Then, we follow the conventional HC to generate the unconstrained raw parameters A ~ l ∈ ℝ 1 × n hc \tilde{A}_{l}\in\mathbb{R}^{1\times n_{\text{hc}}} , B ~ l ∈ ℝ n hc × n hc \tilde{B}_{l}\in\mathbb{R}^{n_{\text{hc}}\times n_{\text{hc}}} , and C ~ l ∈ ℝ n hc × 1 \tilde{C}_{l}\in\mathbb{R}^{n_{\text{hc}}\times 1} :
|
||||
A ~ l \displaystyle\tilde{A}_{l}
|
||||
= α l pre ⋅ ( X ^ l W l pre ) + S l pre , \displaystyle=\alpha_{l}^{\mathrm{pre}}\cdot(\hat{X}_{l}W^{\mathrm{pre}}_{l})+S_{l}^{\mathrm{pre}},
|
||||
(3)
|
||||
B ~ l \displaystyle\tilde{B}_{l}
|
||||
= α l res ⋅ Mat ( X ^ l W l res ) + S l res , \displaystyle=\alpha_{l}^{\mathrm{res}}\cdot\operatorname{Mat}(\hat{X}_{l}W^{\mathrm{res}}_{l})+S_{l}^{\mathrm{res}},
|
||||
(4)
|
||||
C ~ l \displaystyle\tilde{C}_{l}
|
||||
= α l post ⋅ ( X ^ l W l post ) T + S l post , \displaystyle=\alpha_{l}^{\mathrm{post}}\cdot(\hat{X}_{l}W^{\mathrm{post}}_{l})^{T}+S_{l}^{\mathrm{post}},
|
||||
(5)
|
||||
where W l pre , W l post ∈ ℝ n hc d × n hc W^{\mathrm{pre}}_{l},W^{\mathrm{post}}_{l}\in\mathbb{R}^{n_{\text{hc}}d\times n_{\text{hc}}} and W l res ∈ ℝ n hc d × n hc 2 W^{\mathrm{res}}_{l}\in\mathbb{R}^{n_{\text{hc}}d\times n_{\text{hc}}^{2}} are learnable parameters for generating the dynamic components;
|
||||
Mat ( ⋅ ) \operatorname{Mat}(\cdot) reshapes a vector of size 1 × n hc 2 1\times n_{\text{hc}}^{2} into a matrix of size n hc × n hc n_{\text{hc}}\times n_{\text{hc}} ;
|
||||
S l pre ∈ ℝ 1 × n hc S_{l}^{\mathrm{pre}}\in\mathbb{R}^{1\times n_{\text{hc}}} , S l post ∈ ℝ n hc × 1 S_{l}^{\mathrm{post}}\in\mathbb{R}^{n_{\text{hc}}\times 1} , and S l res ∈ ℝ n hc × n hc S_{l}^{\mathrm{res}}\in\mathbb{R}^{n_{\text{hc}}\times n_{\text{hc}}} are learnable static biases;
|
||||
and α l pre \alpha_{l}^{\mathrm{pre}} , α l res \alpha_{l}^{\mathrm{res}} , α l post ∈ ℝ \alpha_{l}^{\mathrm{post}}\in\mathbb{R} are learnable gating factors initialized to small values.
|
||||
## Applying Parameter Constraints.
|
||||
After obtaining the unconstrained raw parameters A ~ l , B ~ l , C ~ l \tilde{A}_{l},\tilde{B}_{l},\tilde{C}_{l} , we then apply constraints described earlier to them to enhance the numerical stability.
|
||||
To be specific, for the input and output mappings, we employ a Sigmoid function σ ( ⋅ ) \sigma(\cdot) to ensure their non-negativity and boundedness:
|
||||
A l \displaystyle A_{l}
|
||||
= σ ( A ~ l ) , \displaystyle=\sigma(\tilde{A}_{l}),
|
||||
(6)
|
||||
C l \displaystyle C_{l}
|
||||
= 2 σ ( C ~ l ) . \displaystyle=2\sigma(\tilde{C}_{l}).
|
||||
(7)
|
||||
As for the residual mapping B ~ l \tilde{B}_{l} , we project it onto the manifold of doubly stochastic matrices ℳ \mathcal{M} .
|
||||
This is achieved by the Sinkhorn-Knopp algorithm, which first applies an exponential function to B ~ l \tilde{B}_{l} to ensure positivity, getting M ( 0 ) = exp ( B ~ l ) M^{(0)}=\exp(\tilde{B}_{l}) , and then iteratively performs column and row normalization:
|
||||
M ( t ) = 𝒯 r ( 𝒯 c ( M ( t − 1 ) ) ) , M^{(t)}=\mathcal{T}_{r}(\mathcal{T}_{c}(M^{(t-1)})),
|
||||
(8)
|
||||
where 𝒯 r \mathcal{T}_{r} and 𝒯 c \mathcal{T}_{c} denote row and column normalization, respectively.
|
||||
This iteration converges to a constrained doubly stochastic matrix B l = M ( t max ) B_{l}=M^{(t_{\text{max}})} .
|
||||
We choose t max = 20 t_{\text{max}}=20 as a practical value.
|
||||
## 2.3 Hybrid Attention with CSA and HCA
|
||||
As the context length reaches extreme scales, the attention mechanism emerges as the dominant computational bottleneck in a model.
|
||||
For DeepSeek-V4, we design two efficient attention architectures — Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA) — and employ their interleaved hybrid configuration, which substantially reduces the computational cost of attention in long-text scenarios.
|
||||
CSA integrates both compression and sparse attention strategies: it first compresses the Key-Value (KV) cache of every m m tokens into one entry, and then applies DeepSeek Sparse Attention (DSA) ( DeepSeek-AI, 2025b ) where each query token attends to only k k compressed KV entries.
|
||||
HCA aims for extreme compression by consolidating the KV cache of every m ′ m^{\prime} ( ≫ m \gg m ) tokens into a single entry.
|
||||
The hybrid architecture of CSA and HCA remarkably improves the long-context efficiency of DeepSeek-V4 series, making one-million-token context feasible in practice.
|
||||
This subsection describes the core techniques of our hybrid attention architecture, and we also provide an open-source implementation 1 1
|
||||
1
|
||||
https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/tree/main/inference to specify more details unambiguously.
|
||||
Figure 3:
|
||||
Core architectures of CSA.
|
||||
It compresses the number of KV entries to 1 m \frac{1}{m} times, and then applies DeepSeek Sparse Attention for further acceleration.
|
||||
Additionally, a small set of sliding window KV entries is combined with the selected compressed KV entries to enhance local fine-grained dependencies.
|
||||
## 2.3.1 Compressed Sparse Attention
|
||||
The core architecture of CSA is illustrated in Figure 3 , which first compresses the KV cache of each m m tokens into one entry, and then applies DeepSeek Sparse Attention for further acceleration.
|
||||
## Compressed Key-Value Entries.
|
||||
Let H ∈ ℝ n × d H\in\mathbb{R}^{n\times d} be a sequence of input hidden states, where n n is the sequence length and d d is the hidden size.
|
||||
CSA first computes two series of KV entries C a , C b ∈ ℝ n × c C^{a},C^{b}\in\mathbb{R}^{n\times c} and their corresponding compression weights Z a , Z b ∈ ℝ n × c Z^{a},Z^{b}\in\mathbb{R}^{n\times c} , where c c is the head dimension:
|
||||
C a \displaystyle C^{a}
|
||||
= H ⋅ W a K V , C b = H ⋅ W b K V , \displaystyle=H\cdot W^{aKV},\quad C^{b}=H\cdot W^{bKV},
|
||||
(9)
|
||||
Z a \displaystyle Z^{a}
|
||||
= H ⋅ W a Z , Z b = H ⋅ W b Z , \displaystyle=H\cdot W^{aZ},\quad\penalty\ \penalty\ Z^{b}=H\cdot W^{bZ},
|
||||
(10)
|
||||
where W a K V , W b K V , W a Z , W b Z ∈ ℝ d × c W^{aKV},W^{bKV},W^{aZ},W^{bZ}\in\mathbb{R}^{d\times c} are trainable parameters.
|
||||
Next, each m m KV entries in C a C^{a} and C b C^{b} will be compressed into one entry according to their compression weights and learnable positional biases B a , B b ∈ ℝ m × c B^{a},B^{b}\in\mathbb{R}^{m\times c} , producing C Comp ∈ ℝ n m × c C^{\text{Comp}}\in\mathbb{R}^{\frac{n}{m}\times c} .
|
||||
Each compressed entry C i Comp ∈ ℝ c C^{\text{Comp}}_{i}\in\mathbb{R}^{c} is computed by
|
||||
[ S m i : m ( i + 1 ) − 1 a ; S m ( i − 1 ) : m i − 1 b ] \displaystyle[S^{a}_{mi:m(i+1)-1};S^{b}_{m(i-1):mi-1}]
|
||||
= Softmax row ( [ Z m i : m ( i + 1 ) − 1 a + B a ; Z m ( i − 1 ) : m i − 1 b + B b ] ) , \displaystyle=\operatorname{Softmax}_{\text{row}}([Z^{a}_{mi:m(i+1)-1}+B^{a};Z^{b}_{m(i-1):mi-1}+B^{b}]),
|
||||
(11)
|
||||
C i Comp \displaystyle C^{\text{Comp}}_{i}
|
||||
= ∑ j = m i m ( i + 1 ) − 1 S j a ⊙ C j a + ∑ j = m ( i − 1 ) m i − 1 S j b ⊙ C j b , \displaystyle=\sum_{j=mi}^{m(i+1)-1}S^{a}_{j}\odot C^{a}_{j}+\sum_{j=m(i-1)}^{mi-1}S^{b}_{j}\odot C^{b}_{j},
|
||||
(12)
|
||||
where ⊙ \odot denotes the Hadamard product;
|
||||
Softmax row ( ⋅ ) \operatorname{Softmax}_{\text{row}}(\cdot) denotes the softmax operation along the row dimension, which performs normalization across the total of 2 m 2m elements from both Z a Z^{a} and Z b Z^{b} .
|
||||
When i = 0 i=0 , Z b m ( i − 1 ) : m i − 1 Z^{b}_{m(i-1):mi-1} is padded with negative infinity and C b m ( i − 1 ) : m i − 1 C^{b}_{m(i-1):mi-1} is padded with zeros.
|
||||
Note that each C i Comp C^{\text{Comp}}_{i} is derived from 2 m 2m KV entries, but the indexes of C b C^{b} used for C i Comp C^{\text{Comp}}_{i} and the indexes of C a C^{a} used for C i − 1 Comp C^{\text{Comp}}_{i-1} are overlapped.
|
||||
Therefore, CSA in fact compresses the sequence length to 1 m \frac{1}{m} times.
|
||||
## Lightning Indexer for Sparse Selection.
|
||||
After obtaining the compressed KV entries C Comp C^{\text{Comp}} , CSA applies the DSA strategy to select top-k compressed KV entries for core attention.
|
||||
First, CSA performs the same compression operation used for C Comp C^{\text{Comp}} to get compressed indexer keys K IComp ∈ ℝ n m × c I K^{\text{IComp}}\in\mathbb{R}^{\frac{n}{m}\times c^{I}} , where c I c^{I} is the indexer head dimension.
|
||||
Then, for a query token t t , we produce the indexer queries { 𝐪 t , 1 I ; 𝐪 t , 2 I ; … ; 𝐪 t , n h I I } \{\mathbf{q}_{t,1}^{I};\mathbf{q}_{t,2}^{I};...;\mathbf{q}_{t,n_{h}^{I}}^{I}\} in a low-rank manner:
|
||||
𝐜 t Q \displaystyle\mathbf{c}_{t}^{Q}
|
||||
= 𝐡 t ⋅ W D Q , \displaystyle=\mathbf{h}_{t}\cdot W^{DQ},
|
||||
(13)
|
||||
[ 𝐪 t , 1 I ; 𝐪 t , 2 I ; … ; 𝐪 t , n h I I ] = 𝐪 t I \displaystyle[\mathbf{q}_{t,1}^{I};\mathbf{q}_{t,2}^{I};...;\mathbf{q}_{t,n_{h}^{I}}^{I}]=\mathbf{q}_{t}^{I}
|
||||
= 𝐜 t Q ⋅ W I U Q , \displaystyle=\mathbf{c}_{t}^{Q}\cdot W^{IUQ},
|
||||
(14)
|
||||
where 𝐡 t ∈ ℝ d \mathbf{h}_{t}\in\mathbb{R}^{d} is the input hidden state of the query token t t ; 𝐜 t Q ∈ ℝ d c \mathbf{c}_{t}^{Q}\in\mathbb{R}^{d_{c}} is the compressed latent vector for queries; d c d_{c} denotes the query compression dimension; n h I n_{h}^{I} denotes the number of indexer query heads; W D Q ∈ ℝ d × d c W^{DQ}\in\mathbb{R}^{d\times d_{c}} and W I U Q ∈ ℝ d c × c I n h I W^{IUQ}\in\mathbb{R}^{d_{c}\times c^{I}n_{h}^{I}} are the down-projection and up-projection matrices for indexer queries, respectively.
|
||||
Next, the index score I t , s ∈ ℝ I_{t,s}\in\mathbb{R} between the query token t t and a preceding compressed block s s ( s s < Floor ( t m ) \operatorname{Floor}(\frac{t}{m}) ) is computed by
|
||||
[ w t , 1 I ; w t , 2 I ; … ; w t , n h I I ] = 𝐰 t I \displaystyle[w_{t,1}^{I};w_{t,2}^{I};...;w_{t,n_{h}^{I}}^{I}]=\mathbf{w}_{t}^{I}
|
||||
= 𝐡 t ⋅ W w , \displaystyle=\mathbf{h}_{t}\cdot W^{w},
|
||||
(15)
|
||||
I t , s \displaystyle I_{t,s}
|
||||
= ∑ h = 1 n h I w t , h I ⋅ ReLU ( 𝐪 t , h I ⋅ K s IComp ) , \displaystyle=\sum_{h=1}^{n_{h}^{I}}w_{t,h}^{I}\cdot\text{ReLU}\left(\mathbf{q}^{I}_{t,h}\cdot K^{\text{IComp}}_{s}\right),
|
||||
(16)
|
||||
where W w ∈ ℝ d × n h I W^{w}\in\mathbb{R}^{d\times n_{h}^{I}} is a learnable matrix; w t , h I ∈ ℝ w_{t,h}^{I}\in\mathbb{R} is the weight of the h h -th indexer head.
|
||||
For a query token t t , given its index scores I t , : I_{t,:} , we employ a top-k selector to selectively retain a subset of compressed KV entries 𝒞 t SprsComp \mathcal{C}^{\text{SprsComp}}_{t} for subsequent core attention:
|
||||
𝒞 t SprsComp = { C s Comp | I t , s ∈ Top − k ( I t , : ) } . \mathcal{C}^{\text{SprsComp}}_{t}=\left\{C^{\text{Comp}}_{s}\penalty\ \Big|\penalty\ I_{t,s}\in\operatorname{Top-k}(I_{t,:})\right\}.
|
||||
(17)
|
||||
## Shared Key-Value MQA.
|
||||
After selecting the sparse KV entries, CSA then performs core attention in a Multi-Query Attention (MQA) ( Shazeer, 2019 ) manner, where each compressed KV entry in 𝒞 t SprsComp \mathcal{C}^{\text{SprsComp}}_{t} serves as both attention key and value.
|
||||
To be specific, for a query token t t , we first produce attention queries { 𝐪 t , 1 ; 𝐪 t , 2 ; … ; 𝐪 t , n h } \{\mathbf{q}_{t,1};\mathbf{q}_{t,2};...;\mathbf{q}_{t,n_{h}}\} from the compressed latent vector 𝐜 t Q \mathbf{c}_{t}^{Q} :
|
||||
[ 𝐪 t , 1 ; 𝐪 t , 2 ; … ; 𝐪 t , n h ] = 𝐪 t = 𝐜 t Q ⋅ W U Q , [\mathbf{q}_{t,1};\mathbf{q}_{t,2};...;\mathbf{q}_{t,n_{h}}]=\mathbf{q}_{t}=\mathbf{c}_{t}^{Q}\cdot W^{UQ},
|
||||
(18)
|
||||
where n h n_{h} denotes the number of query heads; W U Q ∈ ℝ d c × c n h W^{UQ}\in\mathbb{R}^{d_{c}\times cn_{h}} is the up-projection matrices for queries.
|
||||
Note that the latent query vector 𝐜 t Q \mathbf{c}_{t}^{Q} is shared with that used for the indexer queries.
|
||||
Next, we perform MQA on { 𝐪 t , i } \{\mathbf{q}_{t,i}\} and 𝒞 t SprsComp \mathcal{C}^{\text{SprsComp}}_{t} :
|
||||
𝐨 t , i = CoreAttn ( query= 𝐪 t , i , key= 𝒞 t SprsComp , value= 𝒞 t SprsComp ) , \mathbf{o}_{t,i}=\operatorname{CoreAttn}\left(\texttt{query=}\mathbf{q}_{t,i},\texttt{key=}\mathcal{C}^{\text{SprsComp}}_{t},\texttt{value=}\mathcal{C}^{\text{SprsComp}}_{t}\right),
|
||||
(19)
|
||||
where 𝐨 t , i ∈ ℝ c \mathbf{o}_{t,i}\in\mathbb{R}^{c} is the core attention output of the i i -th head at the t t -th token; CoreAttn ( ⋅ ) \operatorname{CoreAttn}(\cdot) denotes the core attention operation.
|
||||
## Grouped Output Projection.
|
||||
In the configuration of DeepSeek-V4, c n h cn_{h} is quite large.
|
||||
Therefore, directly projecting the outputs of the core attention operation [ 𝐨 t , 1 ; 𝐨 t , 2 ; … ; 𝐨 t , n h ] = 𝐨 t ∈ ℝ c n h [\mathbf{o}_{t,1};\mathbf{o}_{t,2};...;\mathbf{o}_{t,n_{h}}]=\mathbf{o}_{t}\in\mathbb{R}^{cn_{h}} to a d d -dimensional hidden state will impose a substantial computational burden.
|
||||
To mitigate this cost, we design a grouped output projection strategy.
|
||||
To be specific, we first split n h n_{h} outputs into g g groups, and then for each group of output 𝐨 t , i G ∈ ℝ c n h g \mathbf{o}^{G}_{t,i}\in\mathbb{R}^{c\frac{n_{h}}{g}} , we project it to a d g d_{g} -dimensional intermediate output 𝐨 t , i G ′ ∈ ℝ d g \mathbf{o}^{G^{\prime}}_{t,i}\in\mathbb{R}^{d_{g}} , where d g < c n h g d_{g}<c\frac{n_{h}}{g} .
|
||||
Finally, we project the intermediate output [ 𝐨 t , 1 G ′ ; 𝐨 t , 2 G ′ ; … ; 𝐨 t , g G ′ ] ∈ ℝ d g g [\mathbf{o}^{G^{\prime}}_{t,1};\mathbf{o}^{G^{\prime}}_{t,2};...;\mathbf{o}^{G^{\prime}}_{t,g}]\in\mathbb{R}^{d_{g}g} to the final attention output 𝐨 ^ t ∈ ℝ d \mathbf{\hat{o}}_{t}\in\mathbb{R}^{d} .
|
||||
Figure 4:
|
||||
Core architectures of HCA.
|
||||
It performs heavier compression, where the KV entries of m ′ m^{\prime} ( ≫ m \gg m ) tokens will be consolidated into one.
|
||||
Also, we additionally introduce a small set of sliding window KV entries to enhance local fine-grained dependencies.
|
||||
## 2.3.2 Heavily Compressed Attention
|
||||
The core architecture of HCA is illustrated in Figure 4 , which compresses the KV cache in a heavier manner, but does not employ sparse attention.
|
||||
## Compressed Key-Value Entries.
|
||||
By and large, the compression strategy of HCA is similar to that of CSA, but employs a larger compression rate m ′ m^{\prime} ( ≫ m \gg m ) and does not perform overlapped compression.
|
||||
Let H ∈ ℝ n × d H\in\mathbb{R}^{n\times d} be a sequence of input hidden states, HCA first computes the original KV entries C ∈ ℝ n × c C\in\mathbb{R}^{n\times c} and their corresponding compression weights Z ∈ ℝ n × c Z\in\mathbb{R}^{n\times c} :
|
||||
C \displaystyle C
|
||||
= H ⋅ W K V , \displaystyle=H\cdot W^{KV},
|
||||
(20)
|
||||
Z \displaystyle Z
|
||||
= H ⋅ W Z , \displaystyle=H\cdot W^{Z},
|
||||
(21)
|
||||
where W K V , W Z ∈ ℝ d × c W^{KV},W^{Z}\in\mathbb{R}^{d\times c} are trainable parameters.
|
||||
Next, each m ′ m^{\prime} KV entries in C C will be compressed into one according to the compression weights and learnable positional biases B ∈ ℝ m ′ × c B\in\mathbb{R}^{m^{\prime}\times c} , producing C Comp ∈ ℝ n m ′ × c C^{\text{Comp}}\in\mathbb{R}^{\frac{n}{m^{\prime}}\times c} .
|
||||
Each compressed entry C i Comp ∈ ℝ c C^{\text{Comp}}_{i}\in\mathbb{R}^{c} is computed by
|
||||
S m ′ i : m ′ ( i + 1 ) − 1 \displaystyle S_{m^{\prime}i:m^{\prime}(i+1)-1}
|
||||
= Softmax row ( Z m ′ i : m ′ ( i + 1 ) − 1 + B ) , \displaystyle=\operatorname{Softmax}_{\text{row}}(Z_{m^{\prime}i:m^{\prime}(i+1)-1}+B),
|
||||
(22)
|
||||
C i Comp \displaystyle C^{\text{Comp}}_{i}
|
||||
= ∑ j = m ′ i m ′ ( i + 1 ) − 1 S j ⊙ C j . \displaystyle=\sum_{j=m^{\prime}i}^{m^{\prime}(i+1)-1}S_{j}\odot C_{j}.
|
||||
(23)
|
||||
Through this compression operation, HCA compresses the sequence length to 1 m ′ \frac{1}{m^{\prime}} times.
|
||||
## Shared Key-Value MQA and Grouped Output Projection.
|
||||
HCA also employs the shared KV MQA and grouped output projection strategies as CSA does.
|
||||
After the KV compression, for a query token t t , HCA first produces attention queries { 𝐪 t , 1 ; 𝐪 t , 2 ; … ; 𝐪 t , n h } \{\mathbf{q}_{t,1};\mathbf{q}_{t,2};...;\mathbf{q}_{t,n_{h}}\} in a low-rank manner:
|
||||
𝐜 t Q \displaystyle\mathbf{c}_{t}^{Q}
|
||||
= 𝐡 t ⋅ W D Q , \displaystyle=\mathbf{h}_{t}\cdot W^{DQ},
|
||||
(24)
|
||||
[ 𝐪 t , 1 ; 𝐪 t , 2 ; … ; 𝐪 t , n h ] = 𝐪 t \displaystyle[\mathbf{q}_{t,1};\mathbf{q}_{t,2};...;\mathbf{q}_{t,n_{h}}]=\mathbf{q}_{t}
|
||||
= 𝐜 t Q ⋅ W U Q , \displaystyle=\mathbf{c}_{t}^{Q}\cdot W^{UQ},
|
||||
(25)
|
||||
where 𝐡 t ∈ ℝ d \mathbf{h}_{t}\in\mathbb{R}^{d} is the input hidden state of the query token t t ; n h n_{h} denotes the number of query heads; W D Q ∈ ℝ d × d c W^{DQ}\in\mathbb{R}^{d\times d_{c}} and W U Q ∈ ℝ d c × c n h W^{UQ}\in\mathbb{R}^{d_{c}\times cn_{h}} are the down-projection and up-projection matrices for queries, respectively.
|
||||
Next, we perform MQA on { 𝐪 t , i } \{\mathbf{q}_{t,i}\} and C Comp C^{\text{Comp}} :
|
||||
𝐨 t , i = CoreAttn ( query= 𝐪 t , i , key= C Comp , value= C Comp ) , \mathbf{o}_{t,i}=\operatorname{CoreAttn}\left(\texttt{query=}\mathbf{q}_{t,i},\texttt{key=}C^{\text{Comp}},\texttt{value=}C^{\text{Comp}}\right),
|
||||
(26)
|
||||
where 𝐨 t , i ∈ ℝ c \mathbf{o}_{t,i}\in\mathbb{R}^{c} is the core attention output of the i i -th head at the t t -th token.
|
||||
Next, as CSA does, HCA splits n h n_{h} outputs into g g groups, and for each group of output 𝐨 t , i G ∈ ℝ c n h g \mathbf{o}^{G}_{t,i}\in\mathbb{R}^{c\frac{n_{h}}{g}} , HCA projects it to a d g d_{g} -dimensional intermediate output 𝐨 t , i G ′ ∈ ℝ d g \mathbf{o}^{G^{\prime}}_{t,i}\in\mathbb{R}^{d_{g}} , where d g < c n h g d_{g}<c\frac{n_{h}}{g} .
|
||||
Finally, HCA projects the intermediate output [ 𝐨 t , 1 G ′ ; 𝐨 t , 2 G ′ ; … ; 𝐨 t , g G ′ ] ∈ ℝ d g g [\mathbf{o}^{G^{\prime}}_{t,1};\mathbf{o}^{G^{\prime}}_{t,2};...;\mathbf{o}^{G^{\prime}}_{t,g}]\in\mathbb{R}^{d_{g}g} to the final attention output 𝐨 ^ t ∈ ℝ d \mathbf{\hat{o}}_{t}\in\mathbb{R}^{d} .
|
||||
## 2.3.3 Other Details
|
||||
In addition to the core architectures of CSA and HCA described above, our hybrid attention incorporates several other techniques.
|
||||
For writing clarity, we omit these additional techniques from the above introduction and will briefly describe them in this subsection.
|
||||
Also, this subsection focuses only on the core ideas of them and may omit some tiny details for simplicity.
|
||||
We encourage the readers to refer to our open-source implementation for unambiguous details.
|
||||
## Query and Key-Value Entry Normalization.
|
||||
For both CSA and HCA, we perform an additional RMSNorm operation on each head of the queries and the only head of the compressed KV entries, just before the core attention operation.
|
||||
This normalization avoids exploding attention logits and may improve training stability.
|
||||
## Partial Rotary Positional Embedding.
|
||||
For both CSA and HCA, we partially employ the Rotary Positional Embedding (RoPE) ( Su et al., 2024 ) to the attention queries, KV entries, and the core attention outputs.
|
||||
To be specific, for each query vector and KV entry vector used in CSA and HCA, we apply RoPE to its last 64 dimensions.
|
||||
Since the KV entries serve as both attention keys and values, the naive core attention outputs { 𝐨 t , i } \{\mathbf{o}_{t,i}\} will carry absolute position embeddings, derived from the weighted sum of KV entries.
|
||||
As a countermeasure, we also apply RoPE with position − i -i on the last 64 dimensions of each 𝐨 t , i \mathbf{o}_{t,i} .
|
||||
In this way, the output of the core attention will also carry relative position embeddings — the contribution of each KV entry to the core attention outputs will also be related to the distance between the query and the KV entry.
|
||||
## Additional Branch of Sliding Window Attention.
|
||||
In order to strictly preserve causality in CSA and HCA, each query attends to only preceding compressed KV blocks.
|
||||
Consequently, a query cannot access information from other tokens within its own compressed block.
|
||||
Meanwhile, recent tokens usually possess greater relevance to the query token in language modeling.
|
||||
For these reasons, we introduce a supplementary attention branch to both CSA and HCA in a sliding window manner, for better modeling of local dependencies.
|
||||
To be specific, for each query token, we additionally produce n win n_{\text{win}} uncompressed KV entries corresponding to the recent n win n_{\text{win}} tokens.
|
||||
In the core attention of CSA and HCA, these KV entries in the sliding window will be used along with the compressed KV entries.
|
||||
## Attention Sink.
|
||||
In the core attention of CSA and HCA, we employ the trick of attention sink ( Xiao et al., 2024 ; OpenAI, 2025 ) .
|
||||
To be specific, we set a series of learnable sink logits { z 1 ′ , z 2 ′ , … , z n h ′ } \{z^{\prime}_{1},z^{\prime}_{2},...,z^{\prime}_{n_{h}}\} .
|
||||
For the h h -th attention head, Exp ( z h ′ ) \operatorname{Exp}(z^{\prime}_{h}) will be added to the denominator of the attention score:
|
||||
s h , i , j = Exp ( z h , i , j ) ∑ k Exp ( z h , i , k ) + Exp ( z h ′ ) , s_{h,i,j}=\frac{\operatorname{Exp}(z_{h,i,j})}{\sum_{k}\operatorname{Exp}(z_{h,i,k})+\operatorname{Exp}(z^{\prime}_{h})},
|
||||
(27)
|
||||
where s h , i , j , z h , i , j ∈ ℝ s_{h,i,j},z_{h,i,j}\in\mathbb{R} denote the attention score and attention logit of the h h -th attention head between the i i -th query token and the j j -th preceding token or compressed block.
|
||||
This technique allows each query head to adjust its total attention scores to be not equal to 1, and even to be near 0.
|
||||
## 2.3.4 Efficiency Discussion
|
||||
Due to the employment of hybrid CSA and HCA, together with low-precision computation and storage, the attention module of DeepSeek-V4 series achieves remarkable efficiency in both attention FLOPs and KV cache size, especially in long-context scenarios.
|
||||
First, we adopt a mixed storage format for KV entries: BF16 precision is used for the rotary positional embedding (RoPE) dimensions, while FP8 precision is applied to the remaining dimensions.
|
||||
This hybrid representation reduces the KV cache size by nearly half compared with pure BF16 storage.
|
||||
Second, attention computation within the lightning indexer is performed in FP4 precision, which accelerates the attention operation under extremely long contexts.
|
||||
Third, relative to DeepSeek-V3.2, a smaller attention top-k is chosen in DeepSeek-V4 series, thereby improving model efficiency on short- and medium-length texts.
|
||||
Finally, and most importantly, compressed attention and hybrid attention techniques substantially reduce both the KV cache size and the computational FLOPs.
|
||||
Taking BF16 GQA8 ( Ainslie et al., 2023 ) with a head dimension of 128 as the baseline — one of the common configurations of LLM attention — the KV cache size of DeepSeek-V4 series can be dramatically reduced to approximately 2 % 2\% times of that baseline in the 1M-context setting.
|
||||
Moreover, even when compared with DeepSeek-V3.2 ( DeepSeek-AI, 2025b ) — already an efficient baseline — DeepSeek-V4 series still exhibits substantial advantages in efficiency.
|
||||
A comparison of their inference FLOPs and KV cache size is provided in the right part of Figure 1 .
|
||||
Algorithm 1 Muon Optimizer for DeepSeek-V4
|
||||
1:
|
||||
Learning rate η \eta , momentum μ \mu , weight decay λ \lambda , update rescaling factor γ \gamma
|
||||
2:
|
||||
for each training step t t do
|
||||
3:
|
||||
for each logically independent weight W ∈ ℝ n × m W\in\mathbb{R}^{n\times m} do
|
||||
4:
|
||||
G t = ∇ W ℒ t ( W t − 1 ) G_{t}=\nabla_{W}\mathcal{L}_{t}(W_{t-1})
|
||||
⊳ \triangleright Compute gradients
|
||||
5:
|
||||
M t = μ M t − 1 + G t M_{t}=\mu M_{t-1}+G_{t}
|
||||
⊳ \triangleright Accumulate momentum buffer
|
||||
6:
|
||||
O t ′ = HybridNewtonSchulz ( μ M t + G t ) O^{\prime}_{t}=\operatorname{HybridNewtonSchulz}(\mu M_{t}+G_{t})
|
||||
⊳ \triangleright Nesterov trick and hybrid Newton-Schulz
|
||||
7:
|
||||
O t = O t ′ ⋅ max ( n , m ) ⋅ γ O_{t}=O^{\prime}_{t}\cdot\sqrt{\max(n,m)}\cdot\gamma
|
||||
⊳ \triangleright Rescale the update RMS
|
||||
8:
|
||||
W t = W t − 1 ⋅ ( 1 − η λ ) − η O t W_{t}=W_{t-1}\cdot\left(1-\eta\lambda\right)-\eta O_{t}
|
||||
⊳ \triangleright Perform weight decay and update
|
||||
9:
|
||||
end for
|
||||
10:
|
||||
end for
|
||||
## 2.4 Muon Optimizer
|
||||
We employ the Muon ( Jordan et al., 2024 ; Liu et al., 2025 ) optimizer for the majority of modules in DeepSeek-V4 series due to its faster convergence and improved training stability.
|
||||
The full algorithm of our Muon optimization is summarized in Algorithm 1 .
|
||||
## Basic Configurations.
|
||||
We maintain the AdamW ( Loshchilov and Hutter, 2017 ) optimizer for the embedding module, the prediction head module, the static biases and gating factors of m HC modules, and the weights of all RMSNorm modules.
|
||||
All other modules are updated with Muon.
|
||||
Following Liu et al. (2025) , we also apply weight decay to Muon parameters, use the Nesterov ( Nesterov, 1983 ; Jordan et al., 2024 ) trick, and rescale the Root Mean Square (RMS) of the update matrix for reutilization of our AdamW hyper-parameters.
|
||||
Different from them, we use hybrid Newton-Schulz iterations for orthogonalization.
|
||||
## Hybrid Newton-Schulz Iterations.
|
||||
For a given matrix M M , let its Singular Value Decomposition (SVD) be M = U Σ V T M=U\Sigma V^{T} .
|
||||
The Newton-Schulz iterations aim to approximately orthogonalize M M to be U V T UV^{T} .
|
||||
Usually, M M will be first normalized as M 0 = M / ‖ M ‖ F M_{0}=M/||M||_{F} to ensure its maximum singular value does not exceed 1.
|
||||
Then, each Newton-Schulz iteration performs the following operation:
|
||||
M k = a M k − 1 + b ( M k − 1 M k − 1 T ) M k − 1 + c ( M k − 1 M k − 1 T ) 2 M k − 1 . M_{k}=aM_{k-1}+b(M_{k-1}M_{k-1}^{T})M_{k-1}+c(M_{k-1}M_{k-1}^{T})^{2}M_{k-1}.
|
||||
(28)
|
||||
Our hybrid Newton-Schulz performs 10 iterations over two distinct stages.
|
||||
During the first 8 steps, we use coefficients ( a , b , c ) = ( 3.4445 , − 4.7750 , 2.0315 ) (a,b,c)=(3.4445,-4.7750,2.0315) to drive rapid convergence, bringing the singular values close to 1.
|
||||
In the final 2 steps, we switch to coefficients ( a , b , c ) = ( 2 , − 1.5 , 0.5 ) (a,b,c)=(2,-1.5,0.5) , which stabilize the singular values precisely at 1.
|
||||
## Avoiding Exploding Attention Logits.
|
||||
The attention architecture of DeepSeek-V4 series allows us to directly apply RMSNorm on the attention queries and KV entries, which effectively prevents attention logits from exploding.
|
||||
Consequently, we do not employ the QK-Clip technique ( Liu et al., 2025 ) in our Muon optimizer.
|
||||
## 3 General Infrastructures
|
||||
## 3.1 Fine-Grained Communication-Computation Overlap in Expert Parallelism
|
||||
Mixture-of-Experts (MoE) can be accelerated via Expert Parallelism (EP).
|
||||
However, EP requires complex inter-node communication and imposes substantial demands on interconnect bandwidth and latency.
|
||||
To alleviate the communication bottleneck in EP and achieve higher end-to-end performance under lower interconnection bandwidth requirements, we propose a fine-grained EP scheme that fuses communication and computation into a single pipelined kernel for communication-computation overlapping.
|
||||
## Communication Latency Can Be Hidden.
|
||||
The key insight of our EP scheme is that the communication latency can be effectively hidden beneath computation in MoE layers.
|
||||
As shown in Figure 5 , in DeepSeek-V4 series, each MoE layer can be decomposed mainly into four stages: two communication-bound stages, Dispatch and Combine , and two computation-bound stages, Linear-1 and Linear-2 .
|
||||
Our profiling reveals that within a single MoE layer, the total time of communication is less than that of the computation.
|
||||
Therefore, after fusing communication and computation into a unified pipeline, computation remains the dominant bottleneck, implying that the system can tolerate lower interconnect bandwidth without degrading end-to-end performance.
|
||||
Figure 5:
|
||||
Illustration of our EP scheme with related works.
|
||||
Comet ( Zhang et al., 2025b ) overlaps Dispatch with Linear-1, and Linear-2 with Combine, separately.
|
||||
Our EP scheme achieves a finer-grained overlapping by splitting and scheduling experts into waves.
|
||||
The theoretical speedup is evaluated in the configuration of the DeepSeek-V4-Flash architecture.
|
||||
## Fine-Grained EP Scheme.
|
||||
To further lower the interconnect bandwidth requirement and amplify the benefits of overlapping, we introduce a finer-grained expert partitioning scheme.
|
||||
Inspired by many related works ( Aimuyo et al., 2025 ; Zhang et al., 2025b ) , we split and schedule the experts into waves .
|
||||
Each wave consists of a small portion of experts.
|
||||
As soon as all experts within the wave have completed their communication, computation can commence immediately without waiting for other experts.
|
||||
In steady state, computation of current wave, token transfer for the next wave, and result sending of completed experts all proceed concurrently, as demonstrated in Figure 5 .
|
||||
This forms a fine-grained pipeline among experts, keeping both computation and communication continuous throughout the wave.
|
||||
The wave-based scheduling speeds up the performance on extreme cases such as Reinforcement Learning (RL) rollout, which usually encounters long-tail small batches.
|
||||
## Performance and Open-Sourced Mega-Kernel.
|
||||
We validated the fine-grained EP scheme on both NVIDIA GPUs and HUAWEI Ascend NPUs platforms.
|
||||
Compared against strong non-fused baselines, it achieves 1.50 ∼ 1.73 × 1.50\sim 1.73\times speedup for general inference workloads, and up to 1.96 × 1.96\times for latency-sensitive scenarios such as RL rollouts and high-speed agent serving.
|
||||
We have open-sourced the CUDA-based mega-kernel implementation named MegaMoE 2 2
|
||||
2
|
||||
https://github.com/deepseek-ai/DeepGEMM/pull/304 as a component of DeepGEMM.
|
||||
## Observations and Proposals.
|
||||
We share observations and lessons from kernel development and offer some proposals to hardware vendors, in the hope of aiding efficient hardware design and achieving better software-hardware co-design:
|
||||
•
|
||||
Computation-Communication Ratio.
|
||||
Full communication-computation overlap hinges on the computation-communication ratio, rather than the bandwidth solely.
|
||||
Denoting peak compute throughput as C C and interconnect bandwidth as B B , communication can be fully hidden when C / B ⩽ V comp / V comm C/B\leqslant V_{\text{comp}}/V_{\text{comm}} , where V comp V_{\text{comp}} denotes the computation volume and V comm V_{\text{comm}} refers to the communication volume.
|
||||
For DeepSeek-V4-Pro, where each token-expert pair requires 6 h d 6hd FLOPs (SwiGLU gate, up, and down projections) but only 3 h 3h bytes of communication (FP8 Dispatch + BF16 Combine), this simplifies to:
|
||||
C B ⩽ 2 d = 6144 FLOPs / Byte . \frac{C}{B}\leqslant 2d=6144\;\mathrm{FLOPs/Byte}.
|
||||
That is, each GBps of interconnect bandwidth suffices to hide the communication for 6.1 TFLOP/s of compute.
|
||||
Once bandwidth meets this threshold, it ceases to be the bottleneck, and devoting additional silicon area to further bandwidth brings diminishing returns.
|
||||
We encourage future hardware designs to target such balance points rather than scale bandwidth unconditionally.
|
||||
•
|
||||
Power Budget.
|
||||
Extreme kernel fusion drives compute, memory, and network to high load simultaneously, making power throttling a key performance limiter.
|
||||
We suggest that future hardware designs provide sufficient power headroom for such fully concurrent workloads.
|
||||
•
|
||||
Communication Primitives.
|
||||
In the dispatch stage, we adopt a pull-based approach where each GPU actively reads activations from remote GPUs, avoiding the high notification latency that fine-grained push entails.
|
||||
Future hardware with lower-latency cross-GPU signaling would make push viable and enable more natural communication patterns.
|
||||
•
|
||||
Activation Function.
|
||||
We propose replacing SwiGLU with a low-cost element-wise activation that involves no exponential or division operations.
|
||||
This directly reduces the overhead of post-GEMM processing, preventing the GEMM pipeline from being stalled by activation function computation, thereby enhancing overall computational throughput and resource utilization.
|
||||
## 3.2 Flexible and Efficient Kernel Development with TileLang
|
||||
In practice, our elaborate model architecture would have resulted in hundreds of fine-grained Torch ATen operators.
|
||||
We adopt TileLang ( Wang et al., 2026 ) to develop a set of fused kernels to replace the vast majority of them, delivering optimal performance with minimal effort.
|
||||
It also allows us to quickly prototype operators like attention variants during validation.
|
||||
These kernels play critical roles in model architecture development, large - scale training, and ultimately production deployment of inference services.
|
||||
As a Domain-Specific Language (DSL), TileLang balances development productivity with runtime efficiency, enabling rapid development while supporting deep, iterative optimizations within the same codebase.
|
||||
Additionally, we collaborate closely with the TileLang community to foster a more agile, efficient, and stable kernel development workflow.
|
||||
## Reducing Invocation Overhead with Host Codegen.
|
||||
As accelerators continue to grow in performance, CPU-side orchestration overhead becomes increasingly prominent. For small, highly optimized kernels, such fixed host overhead can easily cap utilization and throughput. A common source of this overhead is that host-side logic, such as runtime contract checks, is typically written in Python for flexibility and thus incurs a fixed per-invocation cost.
|
||||
We mitigate this overhead with Host Codegen , which moves most host-side logic into generated host code. Specifically, we first co-generate the device kernel and a lightweight host launcher at the IR (Intermediate Representation) level, embedding the necessary metadata—such as data types, rank/shape constraints, and stride/layout assumptions—parsed from the language frontend. The launcher is then lowered to the host source code built on top of the TVM-FFI ( Chen et al., 2018 ) framework, whose compact calling convention and zero-copy tensor interop together minimize host-side overhead. At runtime, this generated host code performs validation and argument marshaling, shifting all per - invocation checks out of the Python execution path. Our measurements show that CPU-side validation overhead drops from tens or hundreds of microseconds to less than one microsecond per invocation.
|
||||
## SMT-Solver-Assisted Formal Integer Analysis.
|
||||
TileLang kernels involve complex tensor index arithmetic that requires strong formal integer analysis.
|
||||
During compilation passes such as layout inference, memory hazard detection, and bound analysis, the compiler must verify whether integer expressions satisfy specific properties to enable the corresponding optimizations.
|
||||
Therefore, stronger formal analysis capabilities can unlock more advanced and complex optimization opportunities.
|
||||
To this end, we integrate the Z3 SMT solver ( De Moura and Bjørner, 2008 ) into TileLang’s algebraic system, providing formal analysis capability for most integer expressions in tensor programs.
|
||||
We strike a balance between computational overhead and formal expressiveness by translating TileLang’s integer expressions into Z3’s quantifier-free non-linear integer arithmetic (QF_NIA).
|
||||
Based on Integer Linear Programming (ILP) solvers, QF_NIA seamlessly resolves standard linear integer expressions common in kernels.
|
||||
Furthermore, its inherent non-linear reasoning capacity effectively addresses advanced challenges like vectorization over variable tensor shapes.
|
||||
Under reasonable resource limits, Z3 elevates overall optimization performance while restricting compilation time overhead to just a few seconds.
|
||||
The impact is substantial across multiple passes, including vectorization, barrier insertion, and code simplification.
|
||||
## Numerical Precision and Bitwise Reproducibility.
|
||||
In production settings, numerical correctness and reproducibility are as critical as raw throughput.
|
||||
We therefore prioritize accuracy by default: fast - math optimizations are disabled at the compiler level, and precision - affecting approximations are provided only as explicit, opt-in frontend operators (e.g., T.__exp , T.__log , and T.__sin ).
|
||||
Conversely, when strict IEEE-754 semantics are required, TileLang provides IEEE-compliant intrinsics with explicit rounding modes (e.g., T.ieee_fsqrt , T.ieee_fdiv , and T.ieee_add ), enabling developers to precisely specify numerical behavior.
|
||||
We also target bitwise reproducibility for validating kernels against hand - written CUDA baselines.
|
||||
We align TileLang’s algebraic simplification and lowering rules with mainstream CUDA toolchains (e.g., NVCC) to avoid transformations that introduce unintended bit - level differences.
|
||||
Layout annotations (e.g., T.annotate_layout ) further allow users to pin down layout-dependent lowering decisions, keeping evaluation and accumulation order consistent with the reference CUDA implementation and thus enabling bit-identical outputs when desired.
|
||||
Our evaluation shows that these accuracy- and reproducibility-oriented design choices do not sacrifice performance: under conservative defaults, TileLang kernels remain competitive, while exposing knobs to selectively relax numerical constraints for higher speed.
|
||||
## 3.3 High-Performance Batch-Invariant and Deterministic Kernel Libraries
|
||||
To enable efficient training an
|
||||
|
||||
=== BODY (tail) ===
|
||||
, DeepSeek-V4-Pro and DeepSeek-V4-Flash excel in coding competitions. According to our evaluation, their performance is comparable to GPT-5.4, making this the first time an open model has matched a closed model on this task. On the Codeforces leaderboard, DeepSeek-V4-Pro-Max currently ranks 23rd among human candidates.
|
||||
DeepSeek-V4 also demonstrates strong performance on formal mathematical task under both agentic and compute-intensive settings. Under an agentic setup, it achieves state-of-the-art results, shown in Figure 8 , outperforming prior models such as Seed Prover ( Chen et al., 2025 ) . With a more compute-intensive pipeline, performance further improves, surpassing systems including Aristotle ( Achim et al., 2025 ) and matching the best known results under this setting.
|
||||
## Agent.
|
||||
The DeepSeek-V4 series demonstrates strong agent performance in evaluations.
|
||||
For code agent tasks, DeepSeek-V4-Pro achieves results comparable to K2.6 and GLM-5.1, though all these open models still lag behind their closed-source counterparts.
|
||||
DeepSeek-V4-Flash underperforms DeepSeek-V4-Pro on coding tasks, particularly on Terminal Bench 2.0.
|
||||
A similar trend is observed across other agent evaluations. It is worth noting that DeepSeek-V4-Pro performs well on MCPAtlas and Toolathlon—two evaluation test sets that include a wide range of tools and MCP services—indicating that our model has excellent generalization capability and does not perform well only on internal frameworks.
|
||||
Figure 9:
|
||||
DeepSeek-V4 series performance on the MRCR task.
|
||||
## 1M-Token Context.
|
||||
DeepSeek-V4-Pro outperforms Gemini-3.1-Pro on the MRCR task, which measures in-context retrieval, but remains behind Claude Opus 4.6.
|
||||
As illustrated in Figure 9 , retrieval performance remains highly stable within a 128K context window.
|
||||
While a performance degradation becomes visible beyond the 128K mark, the model’s retrieval capabilities at 1M tokens remain remarkably strong compared to both proprietary and open-source counterparts.
|
||||
Unlike MRCR, CorpusQA is similar to real scenarios.
|
||||
The evaluation results also indicate that DeepSeek-V4-Pro is better than Gemini-3.1-Pro.
|
||||
Figure 10:
|
||||
HLE and Terminal Bench 2.0 performance by reasoning effort. “None” indicates Non-think mode, and “Speciale” indicates DeepSeek-V3.2-Speciale model.
|
||||
## Reasoning Effort.
|
||||
As shown in Table 7 , the Max mode, which employs longer contexts and reduced length penalties in RL, outperforms the High mode on the most challenging tasks.
|
||||
Figure 10 presents a comparison of performance and cost among DeepSeek-V4-Pro, DeepSeek-V4-Flash, and DeepSeek-V3.2 on representative reasoning and agentic tasks.
|
||||
By scaling test-time compute, DeepSeek-V4 series achieve substantial improvements over the predecessor.
|
||||
Furthermore, on reasoning tasks like HLE, DeepSeek-V4-Pro demonstrates higher token efficiency than DeepSeek-V3.2.
|
||||
## 5.4 Performance on Real-World Tasks
|
||||
Standardized benchmarks often struggle to capture the complexities of diverse, real-world tasks, creating a gap between test results and actual user experience.
|
||||
To bridge this, we have developed proprietary internal metrics that prioritize real-world usage patterns over traditional benchmarks.
|
||||
This approach ensures that our optimizations translate into tangible benefits.
|
||||
Our evaluation framework specifically targets the primary use cases of the DeepSeek API and Chatbot, aligning model performance with practical demands.
|
||||
## 5.4.1 Chinese Writing
|
||||
One of the primary use cases for DeepSeek is Chinese writing.
|
||||
We conducted a rigorous evaluation on functional writing and creative writing.
|
||||
Table 12 presents a pairwise comparison between DeepSeek-V4-Pro and Gemini-3.1-Pro on functional writing tasks.
|
||||
These tasks consist of common daily writing queries, where prompts are typically concise and straightforward.
|
||||
Gemini-3.1-Pro was selected as the baseline, as it stands as the top-performing external model for Chinese writing in our evaluations.
|
||||
The results indicate that DeepSeek-V4-Pro outperforms the baseline with an overall win rate of 62.7% versus 34.1%; this is primarily because Gemini occasionally allows its inherent stylistic preferences to override the user’s explicit requirements in Chinese writing scenarios.
|
||||
Table 13 presents the creative writing comparison, which is evaluated along two axes: instruction following and writing quality.
|
||||
Compared with Gemini-3.1-Pro, DeepSeek-V4-Pro achieves a 60.0% win rate in instruction following and 77.5% in writing quality, demonstrating a marginal improvement in instruction following and a substantial gain in writing quality.
|
||||
Although DeepSeek-V4-Pro yields superior results in aggregate user case analysis, an evaluation restricted to the most challenging prompts — specifically those involving high-complexity constraints or multi-turn scenarios — reveals that Claude Opus 4.5 retains a performance advantage over DeepSeek-V4-Pro.
|
||||
As shown in Table 14 , Claude Opus 4.5 achieves a 52.0% win rate versus 45.9%.
|
||||
## 5.4.2 Search
|
||||
Search-augmented question answering is a core capability of the DeepSeek chatbot.
|
||||
On the DeepSeek web and app, the "non-think" mode employs Retrieval-Augmented Search (RAG), whereas the "thinking" mode utilizes agentic search.
|
||||
## Retrieval Augmented Search.
|
||||
We conducted a pairwise evaluation comparing DeepSeek-V4-Pro and DeepSeek-V3.2 across both objective and subjective Q&A categories.
|
||||
As presented in Table 11 , DeepSeek-V4-Pro outperforms DeepSeek-V3.2 by a substantial margin, demonstrating a consistent advantage across both categories.
|
||||
The most pronounced gains are observed in single-value search and planning & strategy tasks, suggesting that DeepSeek-V4-Pro excels at locating precise factual answers and synthesizing structured plans from retrieved context.
|
||||
However, DeepSeek-V3.2 remains relatively competitive on comparison and recommendation tasks, indicating potential room for improvement for DeepSeek-V4-Pro in scenarios requiring balanced, multi-perspective reasoning over search results.
|
||||
## Agentic Search.
|
||||
Unlike standard RAG, agentic search empowers the model to iteratively invoke search and fetch tools per query, significantly enhancing overall search performance.
|
||||
For the thinking mode in DeepSeek-Chat, we optimized the agentic search function to maximize response accuracy within a predefined "thinking budget".
|
||||
As shown in Table 9 , agentic search consistently outperforms RAG, particularly on complex tasks.
|
||||
Furthermore, its cost remains highly efficient, with agentic search being only marginally more expensive than standard RAG (see Table 10 ).
|
||||
## 5.4.3 White-Collar Task
|
||||
To rigorously evaluate the model’s utility in sophisticated enterprise productivity scenarios, we constructed a comprehensive suite of 30 advanced Chinese professional tasks.
|
||||
These workflows deliberately encompass high-level cognitive demands, including in-depth information analysis, comprehensive document generation, and nuanced document editing, spanning a diverse spectrum of 13 critical industries (e.g., finance, education, law, and technology).
|
||||
The evaluation was conducted within an in-house agent harness equipped with basic tools, including Bash and web search.
|
||||
Given the open-ended nature of these tasks, automated metrics usually fall short in capturing the nuances of a high-quality response.
|
||||
Therefore, we conducted human evaluations to compare the performance of DeepSeek-V4-Pro-Max against Opus-4.6-Max.
|
||||
Annotators blindly assessed the model outputs across four dimensions:
|
||||
•
|
||||
Task Completion: Whether the core problem was successfully resolved.
|
||||
•
|
||||
Instruction Following: Adherence to specific constraints and directives.
|
||||
•
|
||||
Content Quality: Factual accuracy, logical coherence, and professional tone.
|
||||
•
|
||||
Formatting Aesthetics: Layout readability and visual presentation.
|
||||
As illustrated in Figure 12 , DeepSeek-V4-Pro-Max outperforms Opus-4.6-Max on diverse Chinese white-collar tasks, achieving an impressive non-loss rate of 63%, and demonstrating consistent advantages across analysis, generation, and editing tasks.
|
||||
The detailed dimension scores shown in Figure 12 highlight the model’s primary strengths in Task Completion and Content Quality.
|
||||
Specifically, DeepSeek-V4-Pro-Max proactively anticipates implicit user intents by frequently providing supplementary insights and self-verification steps.
|
||||
It also excels in long-form generation, delivering in-depth, coherent narratives rather than relying on the overly simplistic bullet points frequently produced by Opus-4.6-Max.
|
||||
Additionally, the model strictly conforms to formal professional conventions, such as standardized Chinese hierarchical numbering.
|
||||
However, in terms of Instruction Following, it occasionally overlooks specific formatting constraints and slightly trails Opus.
|
||||
Furthermore, the model is less proficient at condensing extensive text inputs into succinct summaries.
|
||||
Finally, its Formatting Aesthetics still have substantial room for improvement regarding the overall visual design of presentation slides.
|
||||
Figure 13 , 14 , and 15 present several test cases; due to the extensive length of certain outputs, only partial pages are displayed.
|
||||
Figure 11: Win-rate comparison across analysis, generation, editing tasks, and the overall performance.
|
||||
Figure 12: Detailed dimension scores including Task Completion, Content Quality, Formatting Aesthetics, and Instruction Following.
|
||||
Figure 13: Example output of a task which requires drafting a joint marketing proposal for a popular bubble tea brand and the Beijing Subway.
|
||||
## 5.4.4 Code Agent
|
||||
To benchmark our coding agent capability, we curate tasks from real internal R&D workloads
|
||||
We collect ∼ {\sim} 200 challenging tasks from 50+ internal engineers,
|
||||
spanning feature development, bug fixing, refactoring, and diagnostics
|
||||
across diverse technology stacks including PyTorch, CUDA, Rust, and C++.
|
||||
Each task is accompanied by its original repository, the corresponding execution environment, and human-annotated scoring rubrics; after rigorous quality filtering, 30 tasks are retained as the evaluation set.
|
||||
As shown in Table 8 , DeepSeek-V4-Pro significantly outperforms Claude Sonnet 4.5 and approaches the level of Claude Opus 4.5.
|
||||
Table 8: Comparison on R&D Coding Benchmark (external models included strictly for evaluation purposes).
|
||||
Model
|
||||
Haiku 4.5
|
||||
Sonnet 4.5
|
||||
DeepSeek-V4-Pro-Max
|
||||
Opus 4.5
|
||||
Opus 4.5
|
||||
Thinking
|
||||
Opus 4.6
|
||||
Thinking
|
||||
Pass Rate (%)
|
||||
13
|
||||
47
|
||||
67
|
||||
70
|
||||
73
|
||||
80
|
||||
In a survey asking DeepSeek developers and researchers ( N = 85 N=85 ) — all with experience of using DeepSeek-V4-Pro for agentic coding in their daily work — whether DeepSeek-V4-Pro is ready to serve as their default and primary coding model compared to other frontier models, 52% said yes, 39% leaned toward yes, and fewer than 9% said no.
|
||||
Respondents find DeepSeek-V4-Pro to deliver satisfactory results across most tasks, but note trivial mistakes, misinterpretation of vague prompts, and occasional over-thinking.
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
In this work, we present a preview version of DeepSeek-V4 series, aiming at next-generation large language models that break the efficiency barrier of ultra-long-context processing.
|
||||
By combining a hybrid attention architecture that integrates CSA and HCA, DeepSeek-V4 series achieve a dramatic leap in long-sequence efficiency.
|
||||
The architectural innovations, together with extensive infrastructure optimization, enable efficient native support for million-token contexts and establish a necessary foundation for future test-time scaling, long-horizon tasks, and emerging paradigms such as online learning.
|
||||
Evaluation results demonstrate that DeepSeek-V4-Pro-Max, the maximum reasoning effort mode of DeepSeek-V4-Pro, redefines the state-of-the-art for open models.
|
||||
It substantially outperforms prior open-source models on knowledge benchmarks, achieves superior reasoning performance close to the frontier proprietary models, and delivers competitive agent capabilities.
|
||||
Meanwhile, DeepSeek-V4-Flash-Max attains comparable reasoning performance to leading closed models while maintaining a highly cost-efficient architecture.
|
||||
We believe DeepSeek-V4 series usher in a new era of million-length contexts for open models and pave the way toward better efficiency, scale, and intelligence.
|
||||
In pursuit of extreme long-context efficiency, DeepSeek-V4 series adopted a bold architectural design.
|
||||
To minimize risk, we retained many preliminarily validated components and tricks, which, while effective, made the architecture relatively complex.
|
||||
In future iterations, we will carry out more comprehensive and principled investigations to distill the architecture down to its most essential designs, making it more elegant without sacrificing performance.
|
||||
Meanwhile, although Anticipatory Routing and SwiGLU Clamping have been proven effective in mitigating training instabilities, their underlying principles remain insufficiently understood.
|
||||
We will actively study foundational problems on training stability and strengthen internal metric monitoring, aiming for a more principled and predictive approach to stable large-scale training.
|
||||
In addition, beyond the MoE and sparse attention architecture, we will also proactively explore model sparsity along new dimensions — such as more sparse embedding modules ( Cheng et al., 2026 ) — to further improve computational and memory efficiency without compromising capability.
|
||||
We will also continuously investigate low-latency architectures and system techniques to make long-context deployment and interaction more responsive.
|
||||
Furthermore, we recognize the importance and practical value of long-horizon, multi-round agentic tasks, and will continue to iterate and explore in this direction.
|
||||
We are also working on incorporating multimodal capabilities to our models.
|
||||
Finally, we are committed to developing better data curation and synthesis strategies to consistently enhance model intelligence, robustness, and practical usability across an increasingly broad range of scenarios and tasks.
|
||||
@@ -0,0 +1,533 @@
|
||||
TITLE: DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression
|
||||
ARXIV: 2609.19969
|
||||
AUTHORS: DeepSeek-AI : Anyi Xu , B. Li , Bangcai Lin , Bing Xue , BingCheng Xian , Bingzheng Xu , Bochao Wu , Bowei Zhang , Boyi Deng , C.C. Yu , Chao Jin , Chaofan Lin , Chen Dong , Chenbing Wang , Chenfan Feng , Chengda Lu , Chenggang Zhao , Chengqi Deng , Chengyuan Zhang , Chenhao Xu , Chenqi Zhao , Chenze Shao , Chuhao Wang , Chuqi Zhang , Damai Dai , Dejian Yang , Deli Chen , Di Huang , Di Wu , Donghao Li , Erhang Li , Eric Fu , F. Zhou , Fangwei Zhou , Fangyun Lin , Fangzhou Yuan , Feiyu Xia , Fucong Dai , Guangbo Hao , Guanglin Li , Guanting Chen , Guoai Cao , Guofan Fan , Guolai Meng , Guowei Li , Haichuan Zhang , Haiyang Ma , Haiyang Shen , Han Li , Han Yu , Han Zhang , Hangyuan Deng , Hanwei Xu , Hanxiang Xu , Hanxun Zhong , Hao Guo , Hao Jiang , Hao Li , Hao Qin , Haodong Wen , Haofen Liang , Haofeng Huang , Haohua Liu , Haoling Zhang , Haoming Luo , Haoran Yang , Haotian Xu , Haotian Yuan , Haoting Huang , Haowen Luo , Haoyang Cai , Haoyu Chen , Haozhe Ji , Hengran Zhang , Hengrui Wang , Hengxu Wu , Honghui Ding , Hongxuan Tang , Huadong Wang , Huanqi Cao , Huazuo Gao , Hui Qu , Hui Zeng , J. Yang , J.H. Jin , J.H. Zhang , J.X. Zou , Jia Yu , Jiahui Zhou , Jiajun Chen , Jialiang Huang , Jialin Zhao , Jiamin Tang , Jian Zhou , Jianan Tong , Jianwen Li , Jiaqi Zhu , Jiarui Wang , Jiasheng Ye , Jiashi Li , Jiaxin Xu , Jiaying Ding , Jibai Lu , Jiewen Hu , Jin Yan , Jincheng Zhai , Jingchang Chen , Jingcheng Hu , Jingli Zhou , Jingsheng Xu , Jingting Xiang , Jingyan Yun , Jing
|
||||
DATE: [Submitted on 17 Sep 2026]
|
||||
|
||||
=== SECTION HEADINGS ===
|
||||
## Report GitHub Issue
|
||||
## Abstract
|
||||
## Contents
|
||||
## 1 Introduction
|
||||
## 2 Architecture
|
||||
## 2.1 Overview
|
||||
## 2.1.1 Multimodal Architecture
|
||||
## DeepSeek-ViT
|
||||
## Multimodal Auxiliary-loss-free Load Balancing for MoEs
|
||||
## 2.2 Causal Encoder-Decoder (CED)
|
||||
## 2.3 Compressed Sparse Attention 2 (CSA2)
|
||||
## 2.3.1 Cross-Layer KV and Index Reuse
|
||||
## 2.3.2 Hierarchical Sparse Indexer
|
||||
## 2.4 Efficient Architectural Extensions
|
||||
## 2.4.1 Single-Pass m HC
|
||||
## 2.4.2 Engram
|
||||
## 2.4.3 DSpark
|
||||
## 2.4.4 FP4 Main KV Cache
|
||||
## 2.5 Optimization
|
||||
## Basic Configurations.
|
||||
## Sinkhorn-Balanced Updates for Engram / Embedding / Prediction Head.
|
||||
## 3 General Infrastructures
|
||||
## 3.1 Training Infrastructure
|
||||
## 3.1.1 Multimodal Training Infrastructure
|
||||
## Communication-Computation Overlap in Contrastive Learning.
|
||||
## End-to-End Parallelism.
|
||||
## Long-Sequence Multimodal Training Optimization.
|
||||
## 3.1.2 Attention Sharing Training for CSA2
|
||||
## 3.1.3 Engram
|
||||
## 3.2 Inference System
|
||||
## 3.2.1 Persistent KV Cache Management
|
||||
## 3.2.2 SWA Bounded Replay
|
||||
## Encoder SWA Bounded Replay.
|
||||
## Decoder SWA Bounded Replay.
|
||||
## 4 Pre-Training
|
||||
## 4.1 Data Construction
|
||||
## Text Data Curation
|
||||
## Multimodal Data Curation
|
||||
## Data Integration and Deduplication
|
||||
## 4.2 Pre-Training Setups
|
||||
## 4.2.1 Model Setups
|
||||
## 4.2.2 Training Setups
|
||||
## Vision Encoder Training.
|
||||
## 4.3 Evaluations
|
||||
## 4.3.1 Evaluation Benchmarks
|
||||
## 4.3.2 Evaluation Results
|
||||
## 5 Post-Training
|
||||
## 5.1 Post-Training Pipeline
|
||||
## 5.1.1 Large-Scale Agent Task Synthesis
|
||||
## 5.1.2 RL in Synthesized Tasks
|
||||
## 5.1.3 Running Agents at Massive Scale: DSec
|
||||
## 5.1.4 Controllable Reasoning Effort in RL
|
||||
## 5.2 Asynchronous Post-training Infrastructure
|
||||
## 5.2.1 Overall Workflow
|
||||
## 5.2.2 Mitigating Length Bias and Off-Policy Effects
|
||||
## 5.2.3 Performance Optimization
|
||||
## 5.2.4 Large-Scale On-Policy Distillation
|
||||
## 5.3 Evaluation
|
||||
## 5.3.1 Evaluation Setup
|
||||
## 5.3.2 Evaluation Results
|
||||
## 5.3.3 Performance across reasoning efforts
|
||||
## 5.3.4 Performance across agent scaffolds
|
||||
## 5.3.5 Multi-Agent
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
|
||||
=== BODY (head) ===
|
||||
DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression
|
||||
## Report GitHub Issue
|
||||
×
|
||||
Title:
|
||||
Content selection saved. Describe the issue below:
|
||||
Description:
|
||||
Submit without GitHub
|
||||
Submit in GitHub
|
||||
arXiv is now an independent nonprofit!
|
||||
Learn more
|
||||
×
|
||||
Back to arXiv
|
||||
Why HTML?
|
||||
Report Issue
|
||||
Back to Abstract
|
||||
Download PDF
|
||||
Abstract
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Overview
|
||||
2.1.1 Multimodal Architecture
|
||||
DeepSeek-ViT
|
||||
Multimodal Auxiliary-loss-free Load Balancing for MoEs
|
||||
2.2 Causal Encoder-Decoder (CED)
|
||||
2.3 Compressed Sparse Attention 2 (CSA2)
|
||||
2.3.1 Cross-Layer KV and Index Reuse
|
||||
2.3.2 Hierarchical Sparse Indexer
|
||||
2.4 Efficient Architectural Extensions
|
||||
2.4.1 Single-Pass m HC
|
||||
2.4.2 Engram
|
||||
2.4.3 DSpark
|
||||
2.4.4 FP4 Main KV Cache
|
||||
2.5 Optimization
|
||||
Basic Configurations.
|
||||
Sinkhorn-Balanced Updates for Engram / Embedding / Prediction Head.
|
||||
3 General Infrastructures
|
||||
3.1 Training Infrastructure
|
||||
3.1.1 Multimodal Training Infrastructure
|
||||
Communication-Computation Overlap in Contrastive Learning.
|
||||
End-to-End Parallelism.
|
||||
Long-Sequence Multimodal Training Optimization.
|
||||
3.1.2 Attention Sharing Training for CSA2
|
||||
3.1.3 Engram
|
||||
3.2 Inference System
|
||||
3.2.1 Persistent KV Cache Management
|
||||
3.2.2 SWA Bounded Replay
|
||||
Encoder SWA Bounded Replay.
|
||||
Decoder SWA Bounded Replay.
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
Text Data Curation
|
||||
Multimodal Data Curation
|
||||
Data Integration and Deduplication
|
||||
4.2 Pre-Training Setups
|
||||
4.2.1 Model Setups
|
||||
4.2.2 Training Setups
|
||||
Vision Encoder Training.
|
||||
4.3 Evaluations
|
||||
4.3.1 Evaluation Benchmarks
|
||||
4.3.2 Evaluation Results
|
||||
5 Post-Training
|
||||
5.1 Post-Training Pipeline
|
||||
5.1.1 Large-Scale Agent Task Synthesis
|
||||
5.1.2 RL in Synthesized Tasks
|
||||
5.1.3 Running Agents at Massive Scale: DSec
|
||||
5.1.4 Controllable Reasoning Effort in RL
|
||||
5.2 Asynchronous Post-training Infrastructure
|
||||
5.2.1 Overall Workflow
|
||||
5.2.2 Mitigating Length Bias and Off-Policy Effects
|
||||
5.2.3 Performance Optimization
|
||||
5.2.4 Large-Scale On-Policy Distillation
|
||||
5.3 Evaluation
|
||||
5.3.1 Evaluation Setup
|
||||
5.3.2 Evaluation Results
|
||||
5.3.3 Performance across reasoning efforts
|
||||
5.3.4 Performance across agent scaffolds
|
||||
5.3.5 Multi-Agent
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Author List
|
||||
B Evaluation Details
|
||||
B.1 Scaffold Configurations
|
||||
B.2 Reasoning Efforts across Scaffolds
|
||||
B.3 Detailed Results of Reasoning Benchmarks across Reasoning Efforts
|
||||
C Exponential Token Penalty in Reasoning Effort Control
|
||||
License: arXiv.org perpetual non-exclusive license
|
||||
arXiv:2609.19969v1 [cs.CL] 17 Sep 2026
|
||||
\reportnumber
|
||||
001
|
||||
##
|
||||
DeepSeek-V4.1-Flash:
|
||||
Pushing the Limits of KV Cache Compression
|
||||
DeepSeek-AI
|
||||
research@deepseek.com
|
||||
## Abstract
|
||||
The widespread adoption of long-horizon agents has made model workloads increasingly input-heavy. Although prior work has substantially reduced the cost of long-context computation, prefill remains computationally expensive, and large KV caches continue to strain HBM and SSD capacity and data-transfer bandwidth. Together, these compute, storage, and bandwidth demands constitute the primary bottleneck to further lowering deployment costs. To address this challenge, we introduce DeepSeek-V4.1-Flash, a multimodal Mixture-of-Experts (MoE) model with 552B backbone parameters and support for contexts of up to one million tokens. With its Causal Encoder-Decoder (CED) architecture, the model activates 16B parameters per token during decode but only 8B parameters during prefill, substantially improving cost efficiency for agentic workloads.
|
||||
To push the limits of KV cache compression, DeepSeek-V4.1-Flash combines cross-layer KV cache reuse in Compressed Sparse Attention 2 (CSA2) with FP4 KV caching.
|
||||
These designs reduce its global KV cache footprint (always in HBM) to 890 bytes per token, roughly 1/4 of the corresponding footprint of DeepSeek-V4-Flash.
|
||||
Further, through a dedicated deployment optimization known as SWA Bounded Replay, DeepSeek-V4.1-Flash reduces its persistent KV cache footprint (always on SSD or in host memory) to roughly 1/8 of that of DeepSeek-V4-Flash.
|
||||
Despite its much smaller KV cache footprint, the model delivers substantially better performance than the baseline.
|
||||
In addition, we streamline the DeepSeek-V4 architecture and introduce several efficient architectural extensions.
|
||||
We pretrain DeepSeek-V4.1-Flash on a multimodal corpus comprising 45T tokens and conduct comprehensive post-training, yielding strong performance across diverse text-based and multimodal agentic scenarios. Model checkpoints are available at https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash .
|
||||
(a) (b)
|
||||
Figure 1:
|
||||
(a) Performance of DeepSeek-V4.1-Flash and its counterparts on agentic benchmarks.
|
||||
(b) Global KV cache size per token (in bytes) across generations of DeepSeek models, highlighting DeepSeek’s sustained efforts to reduce context memory requirements. DeepSeek-V4.1-Flash achieves approximately 4-fold and 437-fold reductions in per-token global KV cache size relative to DeepSeek-V4-Flash and DeepSeek-V1, respectively.
|
||||
## Contents
|
||||
1 Introduction
|
||||
2 Architecture
|
||||
2.1 Overview
|
||||
2.1.1 Multimodal Architecture
|
||||
2.2 Causal Encoder-Decoder (CED)
|
||||
2.3 Compressed Sparse Attention 2 (CSA2)
|
||||
2.3.1 Cross-Layer KV and Index Reuse
|
||||
2.3.2 Hierarchical Sparse Indexer
|
||||
2.4 Efficient Architectural Extensions
|
||||
2.4.1 Single-Pass m HC
|
||||
2.4.2 Engram
|
||||
2.4.3 DSpark
|
||||
2.4.4 FP4 Main KV Cache
|
||||
2.5 Optimization
|
||||
3 General Infrastructures
|
||||
3.1 Training Infrastructure
|
||||
3.1.1 Multimodal Training Infrastructure
|
||||
3.1.2 Attention Sharing Training for CSA2
|
||||
3.1.3 Engram
|
||||
3.2 Inference System
|
||||
3.2.1 Persistent KV Cache Management
|
||||
3.2.2 SWA Bounded Replay
|
||||
4 Pre-Training
|
||||
4.1 Data Construction
|
||||
4.2 Pre-Training Setups
|
||||
4.2.1 Model Setups
|
||||
4.2.2 Training Setups
|
||||
4.3 Evaluations
|
||||
4.3.1 Evaluation Benchmarks
|
||||
4.3.2 Evaluation Results
|
||||
5 Post-Training
|
||||
5.1 Post-Training Pipeline
|
||||
5.1.1 Large-Scale Agent Task Synthesis
|
||||
5.1.2 RL in Synthesized Tasks
|
||||
5.1.3 Running Agents at Massive Scale: DSec
|
||||
5.1.4 Controllable Reasoning Effort in RL
|
||||
5.2 Asynchronous Post-training Infrastructure
|
||||
5.2.1 Overall Workflow
|
||||
5.2.2 Mitigating Length Bias and Off-Policy Effects
|
||||
5.2.3 Performance Optimization
|
||||
5.2.4 Large-Scale On-Policy Distillation
|
||||
5.3 Evaluation
|
||||
5.3.1 Evaluation Setup
|
||||
5.3.2 Evaluation Results
|
||||
5.3.3 Performance across reasoning efforts
|
||||
5.3.4 Performance across agent scaffolds
|
||||
5.3.5 Multi-Agent
|
||||
6 Conclusion, Limitations, and Future Directions
|
||||
References
|
||||
A Author List
|
||||
B Evaluation Details
|
||||
B.1 Scaffold Configurations
|
||||
B.2 Reasoning Efforts across Scaffolds
|
||||
B.3 Detailed Results of Reasoning Benchmarks across Reasoning Efforts
|
||||
C Exponential Token Penalty in Reasoning Effort Control
|
||||
## 1 Introduction
|
||||
Applications of long-horizon agents have expanded rapidly in recent years, making ultra-long-context processing an increasingly important model workload. Supporting such workloads requires not only efficient long-sequence processing, but also the persistent storage, reuse, and transfer of large KV caches. KV cache management has therefore become a foundational capability for model deployment, while introducing substantial challenges across computation, storage, and communication. Prior advances in sparse attention ( DeepSeek-AI, 2025 ; DeepSeek-AI, 2026b ) have significantly reduced the computational cost of long-sequence processing, making persistent storage and data movement increasingly prominent bottlenecks.
|
||||
Specifically, DeepSeek-V4 ( DeepSeek-AI, 2026b ) combines a global attention branch spanning the full context with local Sliding-Window Attention (SWA). The global branch maintains global KV, comprising main KV and indexer K, while SWA maintains local KV states. For a fixed window size, SWA KV storage is bounded independently of sequence length. For sufficiently long sequences, global KV therefore dominates the runtime KV footprint, which is constrained by HBM capacity. In addition, certain KV are persisted for prefix reuse, referred to as persistent KV caches, which are constrained by SSD and host memory capacity. I/O and interconnect bandwidth also limit cache migration and loading. Together, these constraints limit serving throughput, increase deployment costs, and ultimately hinder the deployment and adoption of agents over longer task horizons and across broader application scenarios.
|
||||
Further reducing the KV cache footprint is therefore critical to alleviating storage and communication bottlenecks and lowering the cost of long-context serving. To this end, we develop DeepSeek-V4.1-Flash, a multimodal Mixture-of-Experts (MoE) model designed for more aggressive KV cache compression. DeepSeek-V4.1-Flash has 552B backbone parameters, natively supports multimodal inputs, and accommodates contexts of up to one million tokens. We adopt a Causal Encoder-Decoder (CED) architecture, in which decoder global KV is projected from the final encoder hidden states. This design enables the model to activate 8B parameters per token during prefill and 16B during decode, which is particularly cost-effective for input-heavy agentic scenarios. Despite being considerably larger than DeepSeek-V4-Flash, DeepSeek-V4.1-Flash requires only approximately 1/4 as much runtime KV cache storage and 1/8 as much persistent KV cache storage at the same sequence length. Moreover, DeepSeek-V4.1-Flash delivers better overall performance than DeepSeek-V4-Flash.
|
||||
This level of KV cache compression is achieved through joint optimizations in model architecture, cache precision, and deployment strategy. Conceptually, DeepSeek-V4 can be viewed as an SWA-based local-processing backbone augmented with compressed global context. This perspective motivates us to focus on simplifying the global branch while largely preserving the local attention design. At the architectural level, we design Compressed Sparse Attention 2 (CSA2), which applies cross-layer reuse to global KV (including main KV and indexer K) and Top-K indices to substantially reduce KV cache storage.
|
||||
CSA2 has three statically assigned modes: Full, Reindex, and Reuse. Full Mode generates global KV and performs indexing. Reindex Mode reuses the global KV from a preceding layer, and uses its own indexer Q to rescore the shared indexer K and select fresh Top-K indices. Reuse Mode reuses both global KV and the Top-K indices in a preceding layer, and directly performs sparse attention. In all three modes, each layer retains its own global Q and SWA KV. Sharing global KV and indexer K reduces duplicated cache storage. In addition, different from DeepSeek-V4 that employs the Compressed Sparse Attention (CSA)–Heavily Compressed Attention (HCA) hybrid architecture, DeepSeek-V4.1-Flash uses pure CSA2.
|
||||
At the cache-precision level, we use FP4 global KV caches during training with only marginal performance degradation. Together, CSA2 and FP4 KV caching reduce global KV cache storage to approximately 1/4 of that of DeepSeek-V4-Flash, as shown in Figure 1 (b). At the deployment level, DeepSeek-V4.1-Flash, like DeepSeek-V4, uses Sliding-Window Attention (SWA) in every layer. In DeepSeek-V4, we use a hybrid strategy to balance the storage cost of persisting SWA KV caches against the computation required for exact reconstruction. Exact reconstruction requires replaying the most recent L × n win L\times n_{\mathrm{win}} tokens, where L L is the number of layers and n win n_{\mathrm{win}} is the SWA window size. In DeepSeek-V4.1-Flash, we introduce SWA Bounded Replay, which approximately reconstructs the required SWA KV states by replaying only the most recent n win n_{\mathrm{win}} tokens. Our experiments show that this incurs only negligible performance degradation. This finding establishes a new storage–computation trade-off, allowing us to avoid persisting SWA KV cache to SSD while incurring a small amount of prefill recomputation. With SWA Bounded Replay, the persistent KV cache footprint is further reduced to approximately 1/8 of that of DeepSeek-V4-Flash. Together, these optimizations greatly ease pressure on HBM and SSD capacity, reduce deployment costs, and pave the way for deployment at a larger scale.
|
||||
Figure 2: Single-token Decode FLOPs versus context length across generations of DeepSeek models. We account for compute precision by weighting BF16, FP8, and FP4 operations by 1 1 , 0.5 0.5 , and 0.25 0.25 , respectively. DeepSeek-V4.1-Flash maintains nearly constant Decode FLOPs as context length increases, substantially reducing the computational cost of long-context scenarios.
|
||||
Complementing CED and CSA2, we further streamline the original DeepSeek-V4 architecture. Additionally, we upgrade the original m HC ( Xie et al., 2026 ) design to Single-Pass m HC, with an accompanying Mega- m HC deployment kernel that halves activation memory traffic relative to the original four-kernel implementation. Furthermore, we integrate the Engram ( Cheng et al., 2026b ) conditional memory module to strengthen model capabilities. We also introduce the DSpark ( Cheng et al., 2026a ) speculative decoding architecture to improve decoding efficiency through semi-autoregressive draft generation and confidence-scheduled verification.
|
||||
With all the architectural improvements combined, the single-token Decode FLOPs of DeepSeek-V4.1-Flash remain nearly constant across context lengths. Figure 2 shows that extending the context length 256-fold, from 4K to 1M, increases its Decode FLOPs by only 1/4, significantly less than the growth observed for DeepSeek-V4-Flash.
|
||||
To fully realize the KV cache compression benefits of these architectural designs and further improve training and inference efficiency, we systematically co-optimize the training infrastructure and inference system for DeepSeek-V4.1-Flash, ensuring efficient and scalable large-scale multimodal training and long-context deployment. Training infrastructure supports disaggregated vision-encoder execution, balanced image sharding for long sequences, and cross-stage shared-state management for attention reuse. The inference system implements Encoder and Decoder SWA Bounded Replay paths.
|
||||
Further optimizations include communication–computation overlap, sharded Engram embedding tables, and inference kernel fusion. In particular, each CSA2 Reuse Mode layer executes with only 15 kernels during prefill and 11 during decode. We also separate long-lived global KV storage from short-lived encoder SWA KV in host memory, using bounded replay to approximately reconstruct missing encoder SWA states.
|
||||
During pre-training, we train DeepSeek-V4.1-Flash on a large-scale multimodal corpus comprising 45T tokens. Sparse attention is trained from scratch at a sequence length of 64K, without any dense attention warmup stages. After pre-training, the model possesses native multimodal capabilities and supports contexts of up to one million tokens.
|
||||
In our evaluations, DeepSeek-V4.1-Flash-Base achieves world knowledge, reasoning and coding abilities comparable to DeepSeek-V4-Pro-Base, and delivers 5%–10% improvements on held-out evaluations, using only 1/3 total parameters and 1/4 activated parameters.
|
||||
Together, these results highlight its strong parameter efficiency and reflect improvements in training data quality for real-world deployment.
|
||||
Building on this base model, we conduct post-training to elicit its reasoning and agentic capabilities. In contrast to the architectural innovations described above, our post-training introduces no algorithmic innovation: the recipe follows the standard paradigm of supervised fine-tuning (SFT) followed by reinforcement learning (RL) and on-policy distillation (OPD), without any modification beyond well-established practice used in DeepSeek-V4 development ( DeepSeek-AI, 2026b ) . All substantive changes lie instead in the data pipeline. We develop large-scale automated pipelines for data synthesis and environment construction, and progressively scale the data, tasks, and rollouts employed during RL, thereby extending the model’s capabilities across textual, multimodal, and agentic domains.
|
||||
Figure 1 (a) summarizes DeepSeek-V4.1-Flash’s performance on core agentic benchmarks. Our evaluation shows that, despite its compact size, DeepSeek-V4.1-Flash exhibits a distinctive capability profile:
|
||||
•
|
||||
Reasoning. The model delivers strong reasoning ability, sustaining high accuracy on reasoning-intensive benchmarks such as mathematics and competitive programming, showing comparable performance with top open-source models, such as Kimi-K3 ( Team et al., 2026a ) and DeepSeek-V4-Pro.
|
||||
•
|
||||
Agent.
|
||||
DeepSeek-V4.1-Flash achieves performance on par with closed-source frontier models across standard agentic benchmarks like Terminal-Bench 2.1 ( Merrill et al., 2026 ) , DeepSWE v1.1 ( DataCurve, 2026 ) , and AutomationBench ( Shepard and Salimans, 2026 ) . It has proven fully capable of handling everyday coding tasks and white-collar workflows.
|
||||
However, a gap with giant models remains on science-oriented agentic tasks, such as Terminal-Bench 4.0 ( Marten et al., 2026a ) , that require expert-level domain knowledge.
|
||||
•
|
||||
Multimodal. Within the multimodal domain, the model surpasses top-tier open-source competitors like Kimi-K3 specifically on benchmarks evaluating visual reasoning and the interpretation of professional charts. Beyond formal metrics, it also exhibits practical utility in real-world visual agentic workflows, such as frontend development and office automation, where it can utilize rendered screen captures for visual inspection and self-correction. Nevertheless, we acknowledge that a distinct overall performance gap remains when compared to giant closed-source systems.
|
||||
These results indicate that DeepSeek-V4.1-Flash can already match closed-source frontier models on the vast majority of benchmarks, and is capable of completing over 95% of real-world tasks. Meanwhile, its small activation footprint yields low inference latency and serving cost. We therefore believe that DeepSeek-V4.1-Flash offers a favorable trade-off between capability and efficiency, and can serve as a fast, affordable assistant supporting the daily work of a broad population of users. In summary, DeepSeek-V4.1-Flash simultaneously improves model intelligence and inference efficiency while reducing deployment costs. It substantially lowers the cost barrier to deploying long-horizon agents at scale and creates new opportunities for their adoption across a broader range of scenarios. DeepSeek-V4.1-Flash also serves as a new starting point for our continued scaling efforts. Building on this foundation, we will pursue the joint scaling of model architecture, pre-training, and post-training to further explore the frontier of model intelligence.
|
||||
Figure 3:
|
||||
Overall architecture of DeepSeek-V4.1-Flash.
|
||||
The 40-layer network is divided into a causal encoder and a decoder, each with 20 layers.
|
||||
All feed-forward layers use standard DeepSeekMoE.
|
||||
The first two encoder layers use sliding window attention (SWA); the rest use Compressed Sparse Attention 2 (CSA2), with CSA2(ratio, mode) specifying the compression ratio and mode.
|
||||
The model also uses Single-Pass m HC, Engram, DSpark, and a Hierarchical Sparse Indexer.
|
||||
## 2 Architecture
|
||||
## 2.1 Overview
|
||||
DeepSeek-V4.1-Flash is a multimodal mixture-of-experts (MoE) Transformer that takes images and text as input and generates text autoregressively. Its language backbone comprises 40 causal Transformer layers, organized into a 20-layer causal encoder followed by a 20-layer decoder. Each layer incorporates both global attention and sliding window attention (SWA), except for the first two layers, which use SWA only. A vision encoder and an MLP projector convert images into visual embeddings that are processed jointly with text embeddings, with multimodal data incorporated from the start of language-model pre-training. Overall, DeepSeek-V4.1-Flash has 552B backbone parameters and 196B Engram parameters, activating 8B parameters per token during prefill and 16B during decode. Figure 3 illustrates the overall architecture of DeepSeek-V4.1-Flash.
|
||||
The Causal Encoder–Decoder (CED) architecture and Compressed Sparse Attention 2 (CSA2) address complementary costs of long-context inference. CED constructs the decoder’s global key-value (KV) cache from encoder outputs, allowing most prompt tokens to bypass full decoder computation while retaining layer-local sliding-window attention. This nearly halves prefill computation, lowering the cost of processing new or uncached inputs in agentic workloads with growing contexts. CSA2 shares global KV across layers to reduce cache storage and reuses sparse selections to reduce indexing work. In the decoder, a Hierarchical Sparse Indexer restricts later indexers to a candidate pool selected by an earlier indexer, further reducing the number of entries scored per query.
|
||||
We retain the shared and fine-grained routed experts of DeepSeekMoE ( Dai et al., 2024 ) , and introduce modality-specific load balancing ( Wang et al., 2024a ) for image and text tokens. Single-Pass m HC ( Xie et al., 2026 ) revises residual-stream mixing to enable more efficient kernel fusion, and Engram ( Cheng et al., 2026b ) adds sparsely accessed conditional memory. We omit the MTP module during backbone pre-training and use DSpark ( Cheng et al., 2026a ) for speculative decoding. We train DSpark separately after the backbone pre-training stage. Additionally, we compress the main KV cache to FP4 to further reduce storage overhead. The following sections describe these components and the corresponding optimization changes.
|
||||
## 2.1.1 Multimodal Architecture
|
||||
The multimodal input pathway comprises a vision encoder and an MLP projector. For each input image, the vision encoder produces a spatial grid of visual features. A 3 × 3 3\times 3 pixel-unshuffle operation then rearranges each local neighborhood along the channel dimension, reducing the spatial resolution before the MLP projector maps the features to the hidden dimension of the language backbone. Finally, the resulting visual embeddings are inserted at the corresponding image-token positions in the input embedding sequence and processed jointly with text embeddings by the language backbone.
|
||||
## DeepSeek-ViT
|
||||
We train a vision encoder named DeepSeek-ViT from scratch to natively process images at varying resolutions. We build DeepSeek-ViT on the Vision Transformer ( Dosovitskiy et al., 2021 ) architecture with several modifications. To accommodate inputs of arbitrary resolutions, we replace standard absolute positional embeddings with 2D-RoPE. To align the ViT more closely with LLM design principles, we replace the patch embedding layer’s convolution with a linear projection to ensure compatibility with the Muon optimizer. We also adopt RMSNorm ( Zhang and Sennrich, 2019 ) for normalization and SwiGLU ( Shazeer, 2020 ) as the activation function. Before feeding visual features into the LLM, we apply a pixel-unshuffle operation with 3 × 3 3\times 3 downsampling to reduce the visual token count by a factor of nine, effectively supporting input resolutions up to approximately 1344 × 1344 1344\times 1344 pixels.
|
||||
## Multimodal Auxiliary-loss-free Load Balancing for MoEs
|
||||
Image and text tokens exhibit distinct representation distributions and may induce different expert-routing preferences in MoEs. Balancing their aggregate load may therefore obscure modality-specific imbalance. To address this issue, we extend auxiliary-loss-free load balancing ( Wang et al., 2024a ) by maintaining separate expert-wise correction biases for text and image tokens. During routing, each token uses the correction biases associated with its modality for expert selection, while the original routing scores are retained for weighting the selected expert outputs. After each training step, the two sets of biases are updated independently according to their respective expert loads. This design balances expert utilization within each modality and contributes to stable and efficient multimodal training.
|
||||
## 2.2 Causal Encoder-Decoder (CED)
|
||||
In agentic workflows, frequent tool calls generate extensive prefill requests, imposing severe computational overhead when KV caches miss. To alleviate this prefill bottleneck, we propose the Causal Encoder-Decoder (CED) architecture, inspired by YOCO ( Sun et al., 2024 ) . YOCO reduces prefill computation by allowing the upper half of the layers to directly share the KV cache generated by the lower half. Building upon this concept, CED introduces a series of structural improvements to enhance both the overall KV cache capacity and the computational depth of KV generation. Consequently, CED successfully reduces nearly half of the prefill computation while maintaining performance comparable to the baseline.
|
||||
For global attention, CED treats the bottom L / 2 L/2 layers of the Transformer as the causal encoder. For the upper half layers (i.e., the decoder, l > L / 2 l>L/2 ), the KV entries are not derived from their respective hidden states H l H_{l} . Instead, they are projected directly from the hidden state of the ( L / 2 ) (L/2) -th layer, H L / 2 H_{L/2} , using layer-dependent projection weights ( W l K V W_{l}^{KV} and W l Z W_{l}^{Z} ):
|
||||
C l = H L / 2 W l K V , Z l = H L / 2 W l Z , l > L 2 , C_{l}=H_{L/2}W_{l}^{KV},\quad Z_{l}=H_{L/2}W_{l}^{Z},\quad l>\frac{L}{2},
|
||||
(1)
|
||||
where C C and Z Z represent the KV entries and their corresponding compression weights, respectively. This design allows CED to compute only the first half of the layers during the prefill phase, acquiring the upper-layer global KV cache with minimal computational cost.
|
||||
For sliding window attention (SWA), CED maintains the conventional layer-wise computation across all layers. Specifically, for any layer l l , the local keys and values are derived directly from the current layer’s hidden state H l H_{l} .
|
||||
This design effectively increases the computational depth of local KV generation. However, maintaining this layer-wise computation necessitates an SWA replay process. During the prefill phase, computing the SWA KV cache for the decoder requires processing an additional n win × L / 2 n_{\mathrm{win}}\times L/2 tokens (where n win n_{\mathrm{win}} denotes the window size). For multi-turn interactions with short prompts per turn, this computational overhead in the decoder becomes non-negligible. Fortunately, prior work ( Chen et al., 2025 ) has shown that the actual effective receptive field of SWA is much smaller than the theoretical n win × L / 2 n_{\mathrm{win}}\times L/2 . Motivated by this observation, we introduce Decoder SWA Bounded Replay, which only prefills the last n win n_{\mathrm{win}} tokens of the prompt for the SWA computation, thereby significantly reducing the computational cost. Further details are provided in Section 3.2.2 .
|
||||
Overall, for a sequence length N ≫ n win N\gg n_{\mathrm{win}} , CED reduces the prefill complexity from 𝒪 ( N L ) \mathcal{O}(NL) to 𝒪 ( N L / 2 + n win × L / 2 ) ≈ 𝒪 ( N L / 2 ) \mathcal{O}(NL/2+n_{\mathrm{win}}\times L/2)\approx\mathcal{O}(NL/2) , effectively halving the overall computation.
|
||||
Figure 4:
|
||||
Three operating modes of CSA2.
|
||||
The modes differ in how they obtain main KV, indexer K, and Top-K indices. Green blocks indicate quantities computed in the current layer; yellow blocks indicate main KV and indexer K reused from the most recent Full Mode layer; while red blocks indicate Top-K indices reused from the most recent index-producing (Full or Reindex Mode) layer. All three modes compute main Q and SWA KV in the current layer.
|
||||
## 2.3 Compressed Sparse Attention 2 (CSA2)
|
||||
Serving long contexts requires controlling both KV cache storage and attention computation. These costs can be reduced along three multiplicative dimensions: the entry size, where GQA ( Ainslie et al., 2023 ) reduces the number of KV heads and MLA ( DeepSeek-AI, 2024a ) shares a small latent across heads; the sequence dimension, where every m m tokens are compressed into one entry, like CSA and HCA in DeepSeek-V4 ( DeepSeek-AI, 2026b ) ; and the layer dimension, where some layers reuse the caches ( Brandon et al., 2024 ) and selections of other layers instead of keeping their own, or are replaced altogether by more efficient layers. Prior work has shown that compression along the layer dimension is effective: IndexCache ( Bai et al., 2026 ) reuses Top-K indices across layers to cut indexer computation; YOIO ( Sun et al., 2026b ) computes the sparse routing once and shares it across all layers; and HySparse ( Gao et al., 2026 ) lets sparse layers reuse the KV cache of dense layers. However, index reuse alone saves no main KV storage, network-wide routing sharing limits performance, and hybrid designs still retain full attention layers; more importantly, none of these methods covers all three multiplicative dimensions.
|
||||
CSA2 exploits the three dimensions jointly: it shares main KV and indexer K across layers and allows layers to reuse Top-K indices, with cache sharing and index reuse decoupled. It combines these reuse strategies with a simplified compressor and a Hierarchical Sparse Indexer that narrows the search domain of subsequent indexing layers in the Decoder.
|
||||
Similar to CSA, CSA2 includes a lightweight indexer that scores the main KV entries using indexer Q and indexer K and selects the Top-K entries for each query. Each Q attends to the selected entries together with the layer-local sliding-window KV (SWA KV). CSA2 also includes the uncompressed main KV setting as a special case with a compression ratio of 1. Meanwhile, CSA2 simplifies both the compressor and the indexer. In CSA, a compression ratio of m m produces each main KV entry from 2 m 2m original KV cache entries, with overlapping source entries for adjacent compressed entries. It also includes absolute positional embedding to encode the positions of these 2 m 2m entries during compression. CSA2 removes this overlap and absolute positional embedding. In addition, CSA2 obtains indexer K by projecting main KV entries, replacing CSA’s separate compression path from hidden states. Both designs simplify the implementation and increase the training efficiency.
|
||||
Sections 2.3.1 and 2.3.2 describe the cross-layer reuse strategies and the Hierarchical Sparse Indexer, respectively.
|
||||
## 2.3.1 Cross-Layer KV and Index Reuse
|
||||
Each CSA2 layer is statically assigned one of three modes: Full, Reindex, or Reuse. In all three modes, the layer computes its own query and SWA KV and uses them together with the selected main KV entries to produce a new attention output. The modes differ in how they obtain main KV, indexer K, and Top-K indices. Figure 4 illustrates the three modes.
|
||||
Full Mode. The layer computes its own main KV and indexer Q, projects indexer K from that main KV, and runs the indexer to produce fresh Top-K indices. It therefore executes the complete CSA2 computation path and has the same component responsibilities as a complete CSA layer in DeepSeek-V4.
|
||||
Reindex Mode. The layer reuses the most recent available main KV from a preceding layer together with its corresponding indexer K . The indexer computes its own query, rescores the reused keys, and produces fresh Top-K indices. This allows the sparse selection to change across layers while main KV and indexer K remain shared.
|
||||
Reuse Mode. The layer reuses the most recent available main KV and the latest Top-K indices computed against that main KV by a preceding layer in Full or Reindex Mode. It performs attention using this selection without computing indexer Q or evaluating index scores.
|
||||
Sharing main KV and indexer K reduces cache storage, while reusing Top-K indices avoids additional indexer computation. Reindex Mode preserves cache sharing while allowing the selected entries to change across layers.
|
||||
When CSA2 is combined with CED, the decoder layer assigned to Full Mode computes its own global KV from the hidden state of the ( L / 2 ) (L/2) -th layer, i.e. the last layer of the causal encoder. The Reindex and Reuse Modes are unchanged.
|
||||
Figure 5:
|
||||
Hierarchical Sparse Indexer.
|
||||
Each square represents a position; green squares mark selected indices, and blue rectangles mark blocks selected based on their maximum indexer scores. The decoder’s first CSA2 layer in Full mode selects its own Top-512 indices and builds a shared candidate pool from the selected blocks for subsequent layers. CSA2 layers in Reindex mode then select their Top-512 indices from this pool.
|
||||
## 2.3.2 Hierarchical Sparse Indexer
|
||||
Cross-layer index reuse reduces the number of indexer evaluations, but the remaining indexers still score the full causally visible context. For extremely long contexts, this cost remains a major computational bottleneck. Prior work introduced indexer sparsity by scoring and pruning pooled block representations before token-level indexing ( Xu et al., 2026b ) . We find that in the decoder, information from shallower indexers can naturally be used to restrict the candidates considered by deeper indexers without adding any extra state. We therefore introduce the Hierarchical Sparse Indexer, which is used only in the decoder of CED to reduce this repeated scoring during decode. For each query, the first layer assigned to Full Mode constructs a candidate pool that later re-indexing layers use as their search domain. For a fixed candidate-pool size, this changes the per-query cost of deeper indexers from linear in context length to constant. The mechanism is training-aware and introduced in post-training: the candidate restriction is applied identically during training and inference, so deeper indexers are optimized under the same search domain they use at inference. Figure 5 illustrates this process.
|
||||
This first Full Mode layer scores all causally visible main KV positions and produces the Top-K indices for its own attention. It also performs blockwise candidate selection: each block is assigned the maximum index score among its positions, and the blocks with the highest scores are selected. It then collects the positions covered by the selected blocks into a candidate pool larger than the final Top-K set. For example, selecting 2,048 blocks with 8 positions each yields 16,384 candidate positions. This pool defines where later indexers search; the final Top-K selection determines which main KV entries each layer reads.
|
||||
Subsequent layers in Reindex Mode score only the candidate positions for the corresponding query and select their own Top-K entries within that pool. Layers in Reuse Mode perform no new indexing and use the latest Top-K indices computed against the main KV they reuse. Thus, the candidate pool is shared across indexing layers, while their final selections can differ.
|
||||
For a fixed candidate-pool size, the number of positions scored per query by each subsequent indexer is bounded independently of context length. The first Full Mode layer still scans the entire causally visible range. Hierarchical indexing therefore reduces the cost of later indexer evaluations while retaining the initial full-range pass.
|
||||
## 2.4 Efficient Architectural Extensions
|
||||
## 2.4.1 Single-Pass m HC
|
||||
In DeepSeek-V4, we introduced m HC ( Xie et al., 2026 ) , which maintains n n residual streams between adjacent Transformer blocks. For each token, we denote these streams by X l ∈ ℝ n × d X_{l}\in\mathbb{R}^{n\times d} , where l l is the block index and d d is the hidden dimension. The streams are updated as follows:
|
||||
X l + 1 = B l X l + C l ℱ l ( A l X l ) , ( A l , B l , C l ) = ℋ ( X l ) , X_{l+1}=B_{l}X_{l}+C_{l}\,\mathcal{F}_{l}(A_{l}X_{l}),\qquad(A_{l},B_{l},C_{l})=\mathcal{H}(X_{l}),
|
||||
(2)
|
||||
where A l ∈ ℝ 1 × n A_{l}\in\mathbb{R}^{1\times n} , C l ∈ ℝ n × 1 C_{l}\in\mathbb{R}^{n\times 1} and B l ∈ ℝ n × n B_{l}\in\mathbb{R}^{n\times n} are token-wise coefficients predicted from X l X_{l} . The coefficient predictor ℋ \mathcal{H} includes normalization and projection.
|
||||
Ideally, the residual transformation between two blocks is a single map from ( X l − 1 , Y l − 1 ) (X_{l-1},Y_{l-1}) to ( X l , X ^ l ) (X_{l},\hat{X}_{l}) , where X ^ l = A l X l \hat{X}_{l}=A_{l}X_{l} is the current block input and Y l − 1 = ℱ l − 1 ( X ^ l − 1 ) Y_{l-1}=\mathcal{F}_{l-1}(\hat{X}_{l-1}) is the previous block output. Such a map requires ( n + 1 ) d (n+1)d reads and ( n + 1 ) d (n+1)d writes, giving a lower bound of ( 2 n + 2 ) d (2n+2)d on activation memory traffic.
|
||||
In practice, DeepSeek-V4 uses a multi-pass implementation of ( 2 ), with three kernels that execute sequentially due to data dependencies:
|
||||
X l \displaystyle X_{l}
|
||||
= B l − 1 X l − 1 + C l − 1 Y l − 1 \displaystyle=B_{l-1}X_{l-1}+C_{l-1}Y_{l-1}
|
||||
Residual update, contraction over n n
|
||||
(3)
|
||||
( A l , B l , C l ) \displaystyle(A_{l},B_{l},C_{l})
|
||||
= ℋ ( X l ) \displaystyle=\mathcal{H}(X_{l})
|
||||
Coefficients, contraction over n d nd
|
||||
(4)
|
||||
X ^ l \displaystyle\hat{X}_{l}
|
||||
= A l X l \displaystyle=A_{l}X_{l}
|
||||
Input mixing, contraction over n n
|
||||
(5)
|
||||
The three kernels read ( n + 1 ) d (n+1)d , n d nd and n d nd values respectively and write ( n + 1 ) d (n+1)d in total. Including the pre-norm in ℱ l \mathcal{F}_{l} , the total activation memory traffic is ( 4 n + 4 ) d (4n+4)d , twice the lower bound.
|
||||
In ℋ \mathcal{H} , the normalization weights are folded into the projection weights offline, and the RMS division is applied after the projection. Two of the three stages can therefore share one traversal of the residual: the residual update does not require a reduction across the hidden dimension, so each tile of X l X_{l} can be computed and immediately used to accumulate the projection outputs and the sum of squares needed to compute the RMS.
|
||||
Input mixing cannot be fused into this pass because A l A_{l} is not available until the reduction over all hidden tiles is complete. It therefore requires a second read of X l X_{l} . This second pass can also incorporate input pre-norm. This two-pass implementation would require ( 3 n + 2 ) d (3n+2)d activation reads and writes in total, one additional read of X l X_{l} compared with the lower bound.
|
||||
We therefore introduce Single-Pass m HC , which shifts the input-mixing coefficients by one block, i.e. every block consumes the mixing coefficients produced by the previous one, so that the above dependency disappears:
|
||||
X l + 1 = B l X l + C l ℱ l ( A l − 1 X l ) , ( A l , B l , C l ) = ℋ ( X l ) . X_{l+1}=B_{l}X_{l}+C_{l}\,\mathcal{F}_{l}(A_{l-1}X_{l}),\qquad(A_{l},B_{l},C_{l})=\mathcal{H}(X_{l}).
|
||||
(6)
|
||||
Input mixing now uses A l − 1 A_{l-1} instead of A l A_{l} , so it no longer depends on the coefficients computed from X l X_{l} . Each tile of X l X_{l} can therefore be used immediately for both input mixing and coefficient prediction, without waiting for the full reduction. Empirically, this shift incurs negligible performance degradation.
|
||||
For pre-training, we keep the existing multi-kernel implementation, since the shift only changes which mixing coefficients each block applies.
|
||||
For deployment, we fuse residual update, input mixing, and coefficient prediction into a single kernel, Mega- m HC . The kernel implements m HC with ( 3 n + 2 ) d (3n+2)d activation reads and writes and Single-Pass m HC with ( 2 n + 2 ) d (2n+2)d activation reads and writes.
|
||||
Mega- m HC processes X l X_{l} in tiles along the hidden dimension. Each tile is used to compute the mixed input and to accumulate the quantities needed to predict ( A l , B l , C l ) (A_{l},B_{l},C_{l}) for the next block. The kernel also incorporates input pre-norm and FP8 conversion.
|
||||
The residual is thereby read once and written once, attaining the ( n + 1 ) d (n+1)d reads and ( n + 1 ) d (n+1)d writes of the ideal map and halving the activation memory traffic of our original implementation.
|
||||
## 2.4.2 Engram
|
||||
We augment DeepSeek-V4.1-Flash with Engram ( Cheng et al., 2026c ) , the conditional memory module introduced in our previous work to decouple memorization from computation. We follow the original Engram design—tokenizer compression, multi-head hashing, context-aware gating, and multi-branch integration—with two modifications. First, we omit the short causal convolution because its performance gains do not justify the added complexity in our inference stack. Second, we optimize the Engram embedding with momentum-based update followed by Sinkhorn balancing, as detailed in Section 2.5 .
|
||||
We allocate 196B Engram parameters evenly across two modules. Each module uses N N -gram orders { 2 , 3 , 4 } \{2,3,4\} , with 8 hash heads and a total embedding dimension of 2048 per order. Each head indexes a table of approximately 16M entries, with table sizes chosen to be distinct primes. Both the embedding tables and the key/value projections use FP8 precision. The modules are placed at layers 1 and 14 (zero-indexed) to balance memory usage across training pipeline stages. During inference, deterministic addressing enables embeddings to be prefetched from host memory via background RDMA transfers, with prefetching for the first module overlapping computation in the first Transformer block. Further implementation details for Engram training and inference are discussed in Section 3.1.3 .
|
||||
## 2.4.3 DSpark
|
||||
We equip DeepSeek-V4.1-Flash with DSpark ( Cheng et al., 2026a ) , a speculative decoding module that combines semi-autoregressive drafting with confidence-scheduled verification. The drafter comprises three Transformer blocks with a sliding attention window of 128 tokens. A single forward pass through these blocks computes base logits for five draft positions in parallel, while a lightweight Markov head models dependencies among the draft tokens. A confidence head predicts per-position conditional acceptance probabilities, which are used to estimate prefix survival probabilities. The scheduler combines these estimates with profiled engine throughput curves to dynamically select the verification length for each request, aiming to maximize expected system-wide token throughput under the current system load.
|
||||
Unlike the MTP module in DeepSeek-V3 ( DeepSeek-AI, 2024b ) , which is trained jointly with the backbone throughout pre-training, DSpark is introduced in a dedicated stage after pre-training. In this stage, we train only DSpark while keeping the backbone frozen. During post-training, we continue to train DSpark alongside the backbone, without propagating gradients from the DSpark objective into the backbone. This keeps DSpark aligned with the evolving policy, enabling it to accelerate both online serving and rollout generation for RL and OPD.
|
||||
## 2.4.4 FP4 Main KV Cache
|
||||
Long-context agent workloads require large per-request KV caches, increasing serving costs. DeepSeek-V4 already uses quantization-aware training (QAT) ( Jacob et al., 2018 ) for FP4 indexer queries and keys, accelerating index computation and reducing the indexer cache size. We adopt the OCP-standard MXFP4 format ( Rouhani et al., 2023 ) to support as many hardware platforms as possible, despite the higher accuracy of alternative formats in our experiments.
|
||||
We now extend QAT to the main KV cache, where FP4 reduces storage rather than accelerates matrix multiplication. Dequantizing cached values before attention allows us to use a more accurate format without requiring native matrix-multiplication support for that format, preserving compatibility across hardware platforms.
|
||||
Among the approximately four-bit formats evaluated, we select E2M1 with one E4M3 scale per 16 channels, following NVFP4 ( Alvarez et al., 2025 ) but omitting its second-level global scale to balance accuracy and simplicity.
|
||||
Omitting this scale leaves ample dynamic range for the main KV cache: the format supports magnitudes up to 448 × 6 = 2688 448\times 6=2688 , far above the cache’s magnitude bound. In DeepSeek-V4.1-Flash, the largest trained RMSNorm weight magnitude is approximately 1 1 . After RMS normalization, the L2 norm of the 512 512 -channel KV latent is at most approximately 512 \sqrt{512} . RoPE preserves this norm, so the maximum absolute value across channels after rotation is also bounded by approximately 512 ≈ 22.6 \sqrt{512}\approx 22.6 . Besides, the maximum magnitude observed during training is around 10 10 . Therefore, omitting the global scale causes no measurable decrease in accuracy and simplifies the cache layout.
|
||||
To enable FP4 main KV cache storage in DeepSeek-V4.1-Flash, we introduce QAT during post-training. The non-RoPE and RoPE components use the same quantization format. We quantize the cache after RoPE: quantizing before RoPE yields only a marginal accuracy improvement in our experiments and would introduce additional overhead during decoding.
|
||||
We retain FP8 for the SWA KV cache due to its sensitivity to quantization.
|
||||
Compared with the FP8 main KV cache in DeepSeek-V4, this format nearly halves the storage footprint, both in HBM and when offloaded to SSD.
|
||||
## 2.5 Optimization
|
||||
Building upon the optimization configuration used in DeepSeek-V4, we make some new modifications to better align with the architectural design.
|
||||
First, we use head-wise Muon, where Query weights are split by head before applying the Muon update. Here, we briefly discuss the motivation for such a design.
|
||||
By viewing Muon as a preconditioned gradient descent, vanilla Muon uses one preconditioner for all heads, whereas head-wise Muon provides different preconditioners for different heads. This design can better handle the heterogeneity across attention heads Zhang et al., 2024 ; Zhang, 2026 , Section 3 .
|
||||
As a result, we observe that head-wise Muon outperforms vanilla Muon. The empirical advantage of head-wise Muon is also validated in GLM 5 ( Zeng et al., 2026 ) and Kimi-K3 ( Team et al., 2026a ) .
|
||||
Second, applying Adam to the newly introduced Engram parameters substantially increases the optimizer-state memory footprint.
|
||||
To reduce memory usage during training, we instead optimize the Engram embedding tables, token embedding, and prediction head using a momentum-based update followed by Sinkhorn balancing.
|
||||
Sinkhorn balancing has previously been applied to linear-layer weight matrices in SinkGD ( Scetbon et al., 2025 ) ; here, we extend it to these large parameter matrices.
|
||||
Like Muon, this approach requires only a momentum buffer while empirically outperforming Adam.
|
||||
## Basic Configurations.
|
||||
We retain AdamW ( Loshchilov and Hutter, 2019 ) for normalization-layer weights and other non-matrix parameters, including biases and scaling factors.
|
||||
We use Muon ( Jordan et al., 2024 ) for the weight matrices of linear transformations in the language-model backbone, the Engram projection layers, and the vision-language projector. We use head-wise Muon for Query and Key weights.
|
||||
We apply decoupled weight decay and Nesterov momentum to Muon ( Nesterov, 1983 ; Liu et al., 2025 ) ; normalization-layer weights are also subject to weight decay, whereas biases and scaling factors are not.
|
||||
The Sinkhorn-balanced update also uses Nesterov momentum but does not apply weight decay.
|
||||
During pre-training, we keep the vision encoder frozen until the learning-rate decay stage, while its final normalization layer and the vision–language projector remain trainable. At the onset of learning-rate decay, we unfreeze the vision encoder and optimize it jointly with the LLM with a smaller learning rate.
|
||||
## Sinkhorn-Balanced Updates for Engram / Embedding / Prediction Head.
|
||||
The complete procedure is summarized in Algorithm 1 .
|
||||
At a high level, it follows the same workflow as Muon, with Sinkhorn balancing taking the place of Newton–Schulz orthogonalization.
|
||||
We denote the larger matrix dimension by m m , which corresponds to the vocabulary size for embedding tables and prediction heads, and denote the hidden dimension by n n .
|
||||
Algorithm 1 Momentum update with Sinkhorn balancing
|
||||
1:
|
||||
Weight W t ∈ ℝ m × n W_{t}\in\mathbb{R}^{m\times n} , gradient G t G_{t} , momentum M t − 1 M_{t-1} , momentum coefficient β \beta , base learning rate η t \eta_{t} , learning-rate correction γ \gamma , numerical constants ε \varepsilon and τ \tau , and an odd number of normalization steps K K
|
||||
2:
|
||||
M t ← β M t − 1 + ( 1 − β ) G t M_{t}\leftarrow\beta M_{t-1}+(1-\beta)G_{t}
|
||||
3:
|
||||
G ^ t ← β M t + ( 1 − β ) G t \widehat{G}_{t}\leftarrow\beta M_{t}+(1-\beta)G_{t} ⊳ \triangleright Nesterov momentum
|
||||
4:
|
||||
ρ i ← ∥ G ^ t , i , : ∥ 2 \rho_{i}\leftarrow\lVert\widehat{G}_{t,i,:}\rVert_{2} and ρ ¯ ← 1 m ∑ i = 1 m ρ i \bar{\rho}\leftarrow\frac{1}{m}\sum_{i=1}^{m}\rho_{i}
|
||||
5:
|
||||
U ( 0 ) ← G ^ t U^{(0)}\leftarrow\widehat{G}_{t}
|
||||
6:
|
||||
Set U ( 0 ) i , : ← 0 U^{(0)}_{i,:}\leftarrow 0 if ρ i ⩽ τ ρ ¯ \rho_{i}\leqslant\tau\bar{\rho} ⊳ \triangleright Mask near-zero rows
|
||||
7:
|
||||
for k = 1 , … , K k=1,\ldots,K do
|
||||
8:
|
||||
if k k is odd then
|
||||
9:
|
||||
for all i = 1 , … , m i=1,\ldots,m do
|
||||
10:
|
||||
U i , : ( k ) ← U i , : ( k − 1 ) / ( ∥ U i , : ( k − 1 ) ∥ 2 + ε ) U^{(k)}_{i,:}\leftarrow U^{(k-1)}_{i,:}/(\lVert U^{(k-1)}_{i,:}\rVert_{2}+\varepsilon)
|
||||
11:
|
||||
end for
|
||||
12:
|
||||
else
|
||||
13:
|
||||
for all j = 1 , … , n j=1,\ldots,n do
|
||||
14:
|
||||
U : , j ( k ) ← U : , j ( k − 1 ) / ( ∥ U : , j ( k − 1 ) ∥ 2 + ε ) U^{(k)}_{:,j}\leftarrow U^{(k-1)}_{:,j}/(\lVert U^{(k-1)}_{:,j}\rVert_{2}+\varepsilon)
|
||||
15:
|
||||
end for
|
||||
16:
|
||||
end if
|
||||
17:
|
||||
end for
|
||||
18:
|
||||
Δ t ← n U ( K ) \Delta_{t}\leftarrow\sqrt{n}\,U^{(K)} ⊳ \triangleright Convert unit row ℓ 2 \ell_{2} norm to unit row RMS
|
||||
19:
|
||||
η ~ t ← γ η t \widetilde{\eta}_{t}\leftarrow\gamma\eta_{t} ⊳ \triangleright Match the update magnitude of Adam
|
||||
20:
|
||||
W t + 1 ← W t − η ~ t Δ t W_{t+1}\leftarrow W_{t}-\widetilde{\eta}_{t}\Delta_{t}
|
||||
Given the Nesterov momentum update G ^ t \widehat{G}_{t} , Sinkhorn balancing finds diagonal scaling matrices D r D_{r} and D c D_{c} such that
|
||||
Δ t = n U ( K ) = n D r G ^ t D c , 1 n ∑ j = 1 n ( Δ t ) i j 2 ≈ 1 , 1 m ∑ i = 1 m ( Δ t ) i j 2 ≈ 1 , \Delta_{t}=\sqrt{n}\,U^{(K)}=\sqrt{n}\,D_{r}\widehat{G}_{t}D_{c},\qquad\frac{1}{n}\sum_{j=1}^{n}(\Delta_{t})_{ij}^{2}\approx 1,\qquad\frac{1}{m}\sum_{i=1}^{m}(\Delta_{t})_{ij}^{2}\approx 1,
|
||||
(7)
|
||||
Thus, the procedure approximately equalizes the row-wise and column-wise RMS of the update matrix. Here, one row corresponds to one token index or n-gram identity; and one column encodes one hidden feature. Sinkhorn balancing exploits this token–feature structure by normalizing along both rows and columns.
|
||||
For numerical stability, rows satisfying ρ i ⩽ τ ρ ¯ \rho_{i}\leqslant\tau\bar{\rho} are masked. The factor n \sqrt{n} converts unit row ℓ 2 \ell_{2} norm into unit row-wise RMS. Separately, we adjust the effective learning rate as η ~ t = γ η t \widetilde{\eta}_{t}=\gamma\eta_{t} to match the update magnitude of Adam. We set γ = 0.18 \gamma=0.18 , which is close to the factor 0.2 0.2 used in Moonlight ( Liu et al., 2025 ) .
|
||||
More broadly, Sinkhorn balancing is closely related to optimizers that exploit matrix or tensor axis structure ( Shazeer and Stern, 2018 ; Zhang et al., 2025a ; Wen et al., 2025 ; Glentis et al., 2025 ; Deng et al., 2026 ; Yuan et al., 2026 ; Xu et al., 2026a ) . For example, Adafactor ( Shazeer and Stern, 2018 ) conducts row- and column-wise normalization in a different manner, and Adam-mini ( Zhang et al., 2025a ) uses an alternative row-wise normalization for embedding tables and prediction head. These normalization strategies may differ in optimization performance and communication overhead. We leave more detailed investigation as a future direction.
|
||||
## 3 General Infrastructures
|
||||
## 3.1 Training Infrastructure
|
||||
## 3.1.1 Multimodal Training Infrastructure
|
||||
## Communication-Computation Overlap in Contrastive Learning.
|
||||
The vision encoder is first optimized with a contrastive objective before being fine-tuned with a generative next-token prediction loss. In the contrastive phase, the loss is computed over a full batch of text and vision pairs, so the features of both modalities must be all-gathered across data-parallel ranks, incurring substantial communication. Because the gradient of the text features depends only on the gathered visual features—and, symmetrically, the gradient of the visual features depends only on the gathered text features—each all-gather can be overlapped with the forward or backward pass instead of stalling the pipeline:
|
||||
Forward ( V ) → ( Forward ( T ) ∥ AllGather ( V ) ) → ∇ Text → ( Backward ( T ) ∥ AllGather ( T ) ) → ∇ Vision → Backward ( V ) , \text{Forward}(V)\;\rightarrow\;\bigl(\text{Forward}(T)\;\|\;\text{AllGather}(V)\bigr)\;\rightarrow\;\nabla_{\text{Text}}\\
|
||||
\;\rightarrow\;\bigl(\text{Backward}(T)\;\|\;\text{AllGather}(T)\bigr)\;\rightarrow\;\nabla_{\text{Vision}}\;\rightarrow\;\text{Backward}(V),
|
||||
where V V and T T denote the visual and text features, ( A ∥ C ) (A\,\|\,C) denotes the overlap of computation A A with communication C C , and ∇ \nabla denotes the gradient computation. In this schedule, the visual features are gathered during the text forward pass and the text features during the text backward pass, so that both all-gathers are hidden entirely behind useful computation.
|
||||
## End-to-End Parallelism.
|
||||
To handle the model and data heterogeneity between the vision encoder and the LLM ( Zhang et al., 2025b ) , we adopt the disaggregated encoder design used in recent training systems ( Team et al., 2025 ; Team et al., 2026b ) . The vision encoder is replicated outside the LLM parameter tree, and each training step is divided into three phases: vision encoder forward, LLM forward/backward, and vision encoder backward. This separation prevents interference between vision encoder and LLM computation. Load-balanced vision processing is confined to the first and last phases, while the LLM phase remains free of vision computation and preserves the parallel strategy of text-only training.
|
||||
## Long-Sequence Multimodal Training Optimization.
|
||||
DeepSeek-V4.1-Flash is trained on sequences of up to one million tokens, with a substantial share of training occurring at ultra-long sequence lengths in both the pre-training and post-training stages. At these sequence lengths, multimodal samples create heavy I/O, CPU, and memory bottlenecks.
|
||||
•
|
||||
Balanced image sharding. During pre-training, a single ultra-long, image-dense sequence can exhaust one host’s I/O, CPU, and memory during loading, so the images of each sequence are sharded across the CP ranks with load balancing, and each image is loaded exactly once. With images read once, loading stays hidden behind compute whenever
|
||||
N × ρ B IO < N × C B GPU ⟺ ρ < B IO B GPU C , \frac{N\times\rho}{B_{\mathrm{IO}}}<\frac{N\times C}{B_{\mathrm{GPU}}}\quad\Longleftrightarrow\quad\rho<\frac{B_{\mathrm{IO}}}{B_{\mathrm{GPU}}}\,C,
|
||||
where N N is the token count, ρ \rho the raw bytes per token, C C the per-token compute, and B IO B_{\mathrm{IO}} , B GPU B_{\mathrm{GPU}} the file-system and GPU bandwidths. Since N N cancels, the criterion involves only per-token quantities ( ρ \rho and C C ), independent of sequence length and cluster size; ρ \rho is set by the vision-module configuration (e.g., the resolution cap or the spatial downsample). Storage throughput therefore becomes a bottleneck only for small models with low per-token compute, as in ablations, while production-scale models remain compute-bound.
|
||||
•
|
||||
Incremental image transfer. Besides the balanced sharding above, the reinforcement-learning rollout transfers images to the inference engine only incrementally, and caches the engine’s CPU-side decoding and preprocessing outputs on a distributed file system for reuse across rollouts and subsequent training.
|
||||
## 3.1.2 Attention Sharing Training for CSA2
|
||||
In Section 2.3 , we introduce CSA2, an attention sharing method with three modes, some of which involve sharing one or more of the main KV, the indexer K, and the Top-K indices across multiple layers. Supporting CSA2 in large-scale distributed training requires additional coordination beyond the attention computation itself. In particular, layers that share attention components may be placed on different pipeline stages, making direct module reuse incompatible with conventional stage-local execution. Therefore, we adopt several designs to support CSA2 training.
|
||||
Shadow indexers address this issue by placing a lightweight executable replica on each participating stage while retaining a single logical owner for the shared parameters. The owner remains responsible for optimization and checkpointing, whereas parameter synchronization and gradient aggregation keep the shadow replicas consistent throughout training. This design preserves the original model semantics without requiring the pipeline scheduler to treat shared layers as a special execution unit.
|
||||
Pipeline payload extensions provide the intermediate representations and sparse routing information required by downstream consumers when the source and consumer layers cross a pipeline boundary. These states are incorporated into the existing point-to-point communication path and partitioned consistently with context parallelism, avoiding unnecessary replication while maintaining the corresponding gradient flow.
|
||||
Micro-batch-level shared-state management tracks the states associated with concurrently active pipeline micro-batches and coordinates their lifetimes across forward execution, activation recomputation, and backward propagation. Shared states are retained until their final consumer has completed and are then released promptly to limit additional memory overhead. The same runtime abstraction also handles stage placement and source–consumer relationships, allowing the attention implementation to access shared states without depending on the physical pipeline layout.
|
||||
Together with lightweight adaptations to the optimizer, checkpointing, warm-up, and computation-graph tracing workflows, these mechanisms enable CSA2 to operate transparently under the existing distributed training interface and pipeline schedules.
|
||||
## 3.1.3 Engram
|
||||
Engram embedding tables are partitioned by row across dedicated process groups of engram parallel size.
|
||||
The group size controls the trade-off between per-device memory usage and the communication scope of embedding lookups.
|
||||
Optimizer states are further sharded across replicas of each table partition.
|
||||
Engram lookup indices depend solely on the input token sequence.
|
||||
Embedding prefetch is therefore initiated for the entire local batch before each pipeline stage begins processing microbatches for the current training step, minimizing interference with pipeline execution.
|
||||
Embedding gradients are buffered during backward and returned to their owning ranks after the backbone backward pass.
|
||||
For efficient integration with multimodal training, embedding prefetch and gradient transfers are scheduled to overlap with the vision encoder’s forward and backward computation.
|
||||
Embeddings are stored and fetched in FP8, with the retrieved values and scaling factors passed directly to the following GEMM. For Engram table updates, Sinkhorn normalization maintains row and column scaling vectors across iterations to avoid repeated writes of the full normalized matrix. Row normalization and the accumulation of partial column statistics are fused into a single kernel to further reduce memory traffic. During RL rollouts, Engram embedding tables remain resident in GPU memory.
|
||||
This placement reduces host memory pressure and helps avoid out-of-memory failures caused by host memory fragmentation.
|
||||
## 3.2 Inference System
|
||||
DeepSeek-V4.1-Flash is designed with inference efficiency as a first-class concern. Although its architecture is conceptually complex, the resulting inference kernel flow is remarkably concise. Through reasonable kernel fusion, we encapsulate the intricate operations and keep hardware resources fully pipelined inside a small number of fused kernels—including the fused-RoPE-attention-RoPE-cast kernel in FlashMLA ( Li and Liu, 2025 ) , the Mega-Gate, Mega- m HC, and Mega-MoE kernels in DeepGEMM ( Zhao et al., 2025 ) , the kernels in TileKernels ( Wang et al., 2026a ) , and the TopK kernel in DeepSelect ( Qian et al., 2026 ) . As a result, the vast majority of Transformer layers—those whose CSA2 operates in Reuse Mode—execute with only 15 kernels during prefill and 11 during decode, thereby achieving both high-throughput and low-l
|
||||
|
||||
=== BODY (tail) ===
|
||||
ities, substantially closing the gap with frontier closed-source models while establishing clear advantages among open-source alternatives.
|
||||
On cyber-security tasks, DeepSeek-V4.1-Flash establishes a new state of the art among open-source models. Given the dual-use nature of these capabilities, we encourage the community to apply them responsibly, such as for defensive security research and vulnerability remediation. In the domain of visual agent tasks, DeepSeek-V4.1-Flash demonstrates robust capabilities, particularly in scenarios requiring visual reasoning and the analysis of complex professional charts. While it outperforms the leading open-source model Kimi-K3, we acknowledge that a measurable gap still remains when benchmarked against the leading closed-source alternatives.
|
||||
Beyond peak performance, DeepSeek-V4.1-Flash exposes a reasoning-effort setting that allows
|
||||
users to trade inference cost for accuracy in a controllable manner. As illustrated in
|
||||
Figure 9 , both accuracy and output length increase steadily with
|
||||
the effort level across reasoning and agentic benchmarks alike. Raising the effort from 25 to
|
||||
100 improves the average Pass@1 on eight reasoning-intensive benchmarks from 67.1% to 76.3%,
|
||||
on DeepSWE v1.1 from 66.0% to 74.2%, and on Terminal-Bench 2.1 from 82.4% to 90.6%, at
|
||||
the cost of roughly 2.5 × 2.5\times more output tokens. Notably, the effort control learned on
|
||||
single-response reasoning transfers faithfully to long-horizon agentic trajectories, where it
|
||||
governs the total amount of exploration and verification across turns. The gains are
|
||||
front-loaded: the 60–80 range already recovers most of the accuracy of the maximum setting
|
||||
at less than half of its token budget, whereas the final step to effort 100 lengthens agent
|
||||
trajectories by 1.6 1.6 – 1.8 × 1.8\times for only marginal improvements. The maximum tier is thus
|
||||
best reserved for the most challenging tasks, while moderate effort levels offer a favorable
|
||||
cost–performance balance for everyday agentic use.
|
||||
## 5.3.3 Performance across reasoning efforts
|
||||
Figure 9 reports performance and average response length across a range of budget values. As the reasoning effort increases, the model produces progressively longer reasoning traces and accuracy improves monotonically on reasoning-intensive benchmarks and software engineering tasks, with the largest gains concentrated in the low-to-mid budget range and diminishing returns beyond. Although our RL involves only a limited number of effort levels, the use of scalar efforts achieves flexible, interpolated control of response length within a specific range. This allows practitioners to trade off quality against latency and token cost along a smooth continuum: latency-sensitive applications can operate at low effort with modest accuracy degradation, while difficult tasks can invoke high effort to recover the model’s full reasoning capability. In our public API service, we expose three preset effort levels that map onto this scale: low , high , and max correspond to effort values of 50, 75, and 100, respectively.
|
||||
Figure 9: Performance and output length as a function of reasoning effort. Each panel plots Pass@1 (solid, left axis) and mean output
|
||||
tokens per response (dashed, right axis) as the reasoning-effort value is varied from 25 to 100; The results of reasoning-intensive benchmarks are averaged over eight benchmarks (AIME 2026, Apex 2025 Shortlist, GPQA Diamond, HLE, IMO-AnswerBench, LiveCodeBench, MathArena-Apex, SimpleQA-Verified). DeepSWE v1.1 is evaluated based on mini-SWE and Terminal-Bench v2.1 is evaluated based on DeepSeek Harness (Minimal).
|
||||
## 5.3.4 Performance across agent scaffolds
|
||||
In practice, a model is rarely deployed within a single fixed agent framework; different scaffolds vary in their system prompts, tool definitions, context management strategies, and interaction protocols, and a model that overfits to one particular harness may degrade substantially when placed in another. To assess the robustness of our model to such variation, our comparison covers eight configurations from six scaffold families: Claude Code ( Anthropic, 2026 ) , Codex ( OpenAI, 2026 ) , OpenCode ( Anomaly, 2026 ) , Pi ( Zechner, 2026 ) , mini-SWE ( Yang et al., 2024 ) , and DeepSeek Harness (DSH) ( DeepSeek-AI, 2026a ) in Minimal, Standard, and PTC modes. For each scaffold, we keep the model checkpoint, decoding configuration, and task set identical, and only the surrounding harness, including its native system prompt, tool schema, and turn-taking logic, is changed. Table 4 reports performance at Max reasoning effort (100) on DeepSWE v1.1 and Terminal-Bench v2.1.
|
||||
Table 4: Performance across agent scaffolds at Max reasoning effort.
|
||||
Benchmark (Metric)
|
||||
Claude Code
|
||||
Codex
|
||||
OpenCode
|
||||
Pi
|
||||
mini-SWE
|
||||
DeepSeek Harness
|
||||
Minimal
|
||||
Standard
|
||||
PTC
|
||||
DeepSWE v1.1 (Resolved)
|
||||
69.8
|
||||
65.6
|
||||
65.5
|
||||
66.2
|
||||
74.2
|
||||
72.6
|
||||
70.5
|
||||
67.6
|
||||
Terminal-Bench v2.1 (Pass@1)
|
||||
88.0
|
||||
84.1
|
||||
85.0
|
||||
86.1
|
||||
90.3
|
||||
90.6
|
||||
85.8
|
||||
85.8
|
||||
Note. All scaffolds use N = 8 N=8 samples per task on DeepSWE v1.1 and N = 3 N=3 on Terminal-Bench v2.1.
|
||||
Runs use Linux containers, temperature 1.0, top-p 0.95, a 1M-token context window, and max_steps=500 model-generation rounds per agent on both benchmarks.
|
||||
Terminal-Bench v2.1 is evaluated without network access.
|
||||
We evaluate four Claude Code versions, with per-version results and their average reported in Appendix Table 5 ; this table reports v2.1.251.
|
||||
More scaffold-specific configurations are detailed in Appendix B.1 .
|
||||
The model’s agentic capabilities transfer well across scaffold families with different prompts and tool interfaces, rather than depending on conventions specific to a particular harness. Its performance remains robust as the surrounding interaction protocol and tool abstractions change, indicating that its agentic behavior is not tightly coupled to a single scaffold design. This robustness is consistent with the diversity of environments, tool schemas, and interaction formats in our synthesized training data (Section 5 ), which is designed to encourage generalization across agent scaffolds.
|
||||
## 5.3.5 Multi-Agent
|
||||
Figure 10: Test-time compute scaling for single-agent and multi-agent configurations on ProgramBench (Almost@1) and FrontierSWE v2 (Mean@5) as functions of the per-rollout wall-clock deadline.
|
||||
To explore multi-agent collaboration on complex tasks, we conduct preliminary experiments with DeepSeek-V4.1-Flash using DeepSeek Harness’s Agent Team mode.
|
||||
Multi-Agent Harnesses.
|
||||
We use DeepSeek Harness in Agent Team mode, where a lead agent can asynchronously create named, persistent teammates by default through spawn_teammate . Each receives a delegated task and starts in either fresh mode without lead history or fork mode with a one-time snapshot of the lead’s completed turns. All agents share one repository checkout, making edits immediately visible to one another.
|
||||
Agents communicate through a durable peer mailbox. A message sent through send_message reaches a running teammate at its next step boundary, starts a new turn for an idle teammate, or resumes an inactive teammate.
|
||||
Across each teammate’s task lifecycle, the lead monitors runtime status with list_agents and waits for status, mailbox, or shared-task changes with wait_agent . Task ownership, dependencies, and advisory write scopes are maintained on a shared task board (using the four team_task_* tools with revision checks on updates). When intervention is needed, only the lead can interrupt a teammate’s current turn through interrupt_agent .
|
||||
Once the required work is complete, the lead reviews and tests the combined changes and produces the final response.
|
||||
Training. We train Agent Team mode with an RL reward combining task performance, a collaboration bonus that encourages delegation and inter-agent communication, and a derived-latency penalty that promotes efficient coordination. Derived latency is computed by representing execution events and their collaboration dependencies as a directed acyclic graph (DAG), assigning costs from token counts at fixed prefill/decode rates plus measured tool-execution time, and taking the length of the critical path. This encourages useful parallelism while penalizing unnecessary sequential work and synchronization, with reduced sensitivity to serving-side batching and queuing delays.
|
||||
Performance.
|
||||
We construct a high-confidence subset of ProgramBench ( Yang et al., 2026 ) by retaining only tasks for which the reference solution achieves a pass rate of at least 95% on the hidden test suite. This filtering procedure leaves 172 "golden" tasks. We also evaluate on FrontierSWE v2 ( Kondra et al., 2026 ) , a larger and more challenging successor to FrontierSWE that uses a substantially improved methodology. We construct a no-GPU subset from the currently public tasks by excluding tasks that require GPU access. We evaluate single-agent and multi-agent configurations on both benchmarks under explicit per-rollout wall-clock deadlines. The results reported here are preliminary: we compare the strongest observed multi-agent configurations with the strongest available single-agent baselines.
|
||||
On ProgramBench, we run up to three rollouts per task, corresponding to 516 planned rollouts for each configuration. We report Almost@1, which measures the fraction of individual rollouts achieving a score of at least 0.95. On FrontierSWE v2, we report Mean@5. ProgramBench deadlines range from 1 to 12 hours, while FrontierSWE v2 is evaluated at deadlines ranging from 1 to 20 hours. At each deadline, metrics are computed from the outputs available when the deadline is reached.
|
||||
As shown in Figure 10 , multi-agent configurations outperform their single-agent counterparts at every deadline on both benchmarks. On ProgramBench, Almost@1 increases from 13.59% at 1 hour to a peak of 30.04% at 8 hours for the multi-agent configuration, compared with 12.79% and 20.39% for the single-agent configuration. On FrontierSWE v2, Mean@5 increases from 13.50% at 1 hour to 32.90% at 20 hours for the multi-agent configuration, compared with an increase from 10.50% to 28.20% for the single-agent configuration.
|
||||
## 6 Conclusion, Limitations, and Future Directions
|
||||
In this work, we introduce DeepSeek-V4.1-Flash, a multimodal Mixture-of-Experts (MoE) model with support for contexts of up to one million tokens. Through joint optimization of model architecture, cache precision, and deployment strategy, DeepSeek-V4.1-Flash pushes the limits of KV cache compression. Its Causal Encoder-Decoder (CED) architecture enables the model to activate only 8B parameters per token during prefill, compared with 16B during decode, improving cost efficiency for input-heavy agentic workloads. At equal sequence lengths, cross-layer KV cache reuse in Compressed Sparse Attention 2 (CSA2) and FP4 KV caching reduce its global KV cache footprint (always in HBM) to 890 bytes per token, roughly 1/4 of the corresponding footprint of DeepSeek-V4-Flash. SWA Bounded Replay further reduces its persistent KV cache footprint (always on SSD or in host memory) to roughly 1/8 of that of DeepSeek-V4-Flash. These reductions alleviate HBM and SSD capacity pressure while the model delivers substantially better overall performance than DeepSeek-V4-Flash. Despite possessing a significantly smaller parameter footprint than contemporary open-source models such as GLM-5.3 and Kimi-K3, DeepSeek-V4.1 achieves comparable—and in several tasks, superior—performance across key benchmarks.
|
||||
Although DeepSeek-V4.1-Flash substantially simplifies several architectural components relative to DeepSeek-V4-Flash, the newly introduced architectural changes also create robustness boundaries that have yet to be fully characterized. Our internal evaluations cover a diverse range of test cases and boundary conditions, and we have not observed any systematic degradation in model capabilities in the evaluated settings. Nevertheless, no finite test suite can cover every extreme input and deployment condition. Potential selection errors in CSA2 and approximate state reconstruction in SWA Bounded Replay may still cause capability degradation in untested boundary cases. Going forward, we will continue to expand our stress-testing and evaluation stack, with particular attention to sparse retrieval over long contexts and SWA state reconstruction at cache-resumption boundaries. We will also monitor real-world workloads, systematically characterize potential failure modes and robustness boundaries, and further improve model robustness under extreme conditions.
|
||||
As AI models achieve remarkable performance capabilities, standard evaluation benchmarks have increasingly reached saturation. While DeepSeek-V4.1-Flash demonstrates performance that closely approaches top-tier models like Fable-5 and GPT-6 Astra—offering a highly comparable user experience in daily applications—a performance gap remains on the most challenging tasks. Although benchmark scores show a narrow margin, this parity does not imply that the model matches the frontier capabilities of leading closed-source systems on complex, high-difficulty reasoning and edge cases.
|
||||
Consequently, we will continuously update our evaluation protocols to ensure rigorous assessment of state-of-the-art reasoning boundaries. Alongside continued efforts to reduce model costs, we believe that further advances in model intelligence will depend on the coordinated scaling of data, model capacity, and RL. With DeepSeek-V4.1-Flash as a new starting point, we will continue to explore the limits of model capabilities and systematically address key challenges in large-scale data synthesis and RL scaling. We will also actively integrate model–harness co-design, enabling the joint system to evolve and be optimized together. By advancing cost reduction and capability scaling in tandem, we hope to make highly capable agents more accessible and easier to deploy, further lowering the barriers to adopting AI technologies across a broader range of industries and scenarios.
|
||||
Reference in New Issue
Block a user