新增: DeepSeek 全部 33 篇官方论文(原文 PDF 走 Git LFS + 中文精读笔记与索引)
This commit is contained in:
ylhan
@@ -0,0 +1,525 @@
|
||||
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.
|
||||
## References
|
||||
Bautista et al. (2022)
|
||||
Miguel Angel Bautista, Pengsheng Guo, Samira Abnar, Walter Talbott, Alexander Toshev, Zhuoyuan Chen, Laurent Dinh, Shuangfei Zhai, Hanlin Goh, Daniel Ulbricht, et al.
|
||||
Gaudi: A neural architect for immersive 3d scene generation.
|
||||
Advances in Neural Information Processing Systems , 35:25102–25116, 2022.
|
||||
Chan et al. (2021)
|
||||
Eric R Chan, Marco Monteiro, Petr Kellnhofer, Jiajun Wu, and Gordon Wetzstein.
|
||||
pi-gan: Periodic implicit generative adversarial networks for 3d-aware image synthesis.
|
||||
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pp. 5799–5809, 2021.
|
||||
Chan et al. (2022)
|
||||
Eric R Chan, Connor Z Lin, Matthew A Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio Gallo, Leonidas J Guibas, Jonathan Tremblay, Sameh Khamis, et al.
|
||||
Efficient geometry-aware 3d generative adversarial networks.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 16123–16133, 2022.
|
||||
Chang et al. (2015)
|
||||
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al.
|
||||
Shapenet: An information-rich 3d model repository.
|
||||
arXiv preprint arXiv:1512.03012 , 2015.
|
||||
Chen et al. (2023)
|
||||
Rui Chen, Yongwei Chen, Ningxin Jiao, and Kui Jia.
|
||||
Fantasia3d: Disentangling geometry and appearance for high-quality text-to-3d content creation.
|
||||
arXiv preprint arXiv:2303.13873 , 2023.
|
||||
Cheng et al. (2023)
|
||||
Yen-Chi Cheng, Hsin-Ying Lee, Sergey Tulyakov, Alexander G Schwing, and Liang-Yan Gui.
|
||||
Sdfusion: Multimodal 3d shape completion, reconstruction, and generation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 4456–4465, 2023.
|
||||
Deitke et al. (2023)
|
||||
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar Michel, Eli VanderBilt, Ludwig Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi.
|
||||
Objaverse: A universe of annotated 3d objects.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 13142–13153, 2023.
|
||||
Deng et al. (2023)
|
||||
Congyue Deng, Chiyu Jiang, Charles R Qi, Xinchen Yan, Yin Zhou, Leonidas Guibas, Dragomir Anguelov, et al.
|
||||
Nerdi: Single-view nerf synthesis with language-guided diffusion as general image priors.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 20637–20647, 2023.
|
||||
Eftekhar et al. (2021)
|
||||
Ainaz Eftekhar, Alexander Sax, Jitendra Malik, and Amir Zamir.
|
||||
Omnidata: A scalable pipeline for making multi-task mid-level vision datasets from 3d scans.
|
||||
In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 10786–10796, 2021.
|
||||
Gao et al. (2022)
|
||||
Jun Gao, Tianchang Shen, Zian Wang, Wenzheng Chen, Kangxue Yin, Daiqing Li, Or Litany, Zan Gojcic, and Sanja Fidler.
|
||||
Get3d: A generative model of high quality 3d textured shapes learned from images.
|
||||
Advances In Neural Information Processing Systems , 35:31841–31854, 2022.
|
||||
Gu et al. (2022)
|
||||
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo.
|
||||
Vector quantized diffusion model for text-to-image synthesis.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 10696–10706, 2022.
|
||||
Guo et al. (2023)
|
||||
Yuan-Chen Guo, Ying-Tian Liu, Ruizhi Shao, Christian Laforte, Vikram Voleti, Guan Luo, Chia-Hao Chen, Zi-Xin Zou, Chen Wang, Yan-Pei Cao, and Song-Hai Zhang.
|
||||
threestudio: A unified framework for 3d content generation.
|
||||
https://github.com/threestudio-project/threestudio , 2023.
|
||||
Gupta et al. (2023)
|
||||
Anchit Gupta, Wenhan Xiong, Yixin Nie, Ian Jones, and Barlas Oğuz.
|
||||
3dgen: Triplane latent diffusion for textured mesh generation.
|
||||
arXiv preprint arXiv:2303.05371 , 2023.
|
||||
Hong et al. (2022)
|
||||
Fangzhou Hong, Mingyuan Zhang, Liang Pan, Zhongang Cai, Lei Yang, and Ziwei Liu.
|
||||
Avatarclip: Zero-shot text-driven generation and animation of 3d avatars.
|
||||
ACM Transactions on Graphics (TOG) , 41(4):1–19, 2022.
|
||||
Hu et al. (2021)
|
||||
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.
|
||||
Lora: Low-rank adaptation of large language models.
|
||||
arXiv preprint arXiv:2106.09685 , 2021.
|
||||
Huang et al. (2023)
|
||||
Yukun Huang, Jianan Wang, Yukai Shi, Xianbiao Qi, Zheng-Jun Zha, and Lei Zhang.
|
||||
Dreamtime: An improved optimization strategy for text-to-3d content creation.
|
||||
arXiv preprint arXiv:2306.12422 , 2023.
|
||||
Jain et al. (2022)
|
||||
Ajay Jain, Ben Mildenhall, Jonathan T. Barron, Pieter Abbeel, and Ben Poole.
|
||||
Zero-shot text-guided object generation with dream fields.
|
||||
2022.
|
||||
Jun & Nichol (2023)
|
||||
Heewoo Jun and Alex Nichol.
|
||||
Shap-e: Generating conditional 3d implicit functions.
|
||||
arXiv preprint arXiv:2305.02463 , 2023.
|
||||
Lee & Chang (2022)
|
||||
Han-Hung Lee and Angel X Chang.
|
||||
Understanding pure clip guidance for voxel grid nerf models.
|
||||
arXiv preprint arXiv:2209.15172 , 2022.
|
||||
Li et al. (2023)
|
||||
Muheng Li, Yueqi Duan, Jie Zhou, and Jiwen Lu.
|
||||
Diffusion-sdf: Text-to-shape via voxelized diffusion.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 12642–12651, 2023.
|
||||
Lin et al. (2023)
|
||||
Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming-Yu Liu, and Tsung-Yi Lin.
|
||||
Magic3d: High-resolution text-to-3d content creation.
|
||||
In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2023.
|
||||
Liu et al. (2023a)
|
||||
Minghua Liu, Chao Xu, Haian Jin, Linghao Chen, Zexiang Xu, Hao Su, et al.
|
||||
One-2-3-45: Any single image to 3d mesh in 45 seconds without per-shape optimization.
|
||||
arXiv preprint arXiv:2306.16928 , 2023a.
|
||||
Liu et al. (2023b)
|
||||
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick.
|
||||
Zero-1-to-3: Zero-shot one image to 3d object.
|
||||
arXiv preprint arXiv:2303.11328 , 2023b.
|
||||
Mechrez et al. (2018)
|
||||
Roey Mechrez, Itamar Talmi, and Lihi Zelnik-Manor.
|
||||
The contextual loss for image transformation with non-aligned data.
|
||||
In Proceedings of the European conference on computer vision (ECCV) , pp. 768–783, 2018.
|
||||
Metzer et al. (2023)
|
||||
Gal Metzer, Elad Richardson, Or Patashnik, Raja Giryes, and Daniel Cohen-Or.
|
||||
Latent-nerf for shape-guided generation of 3d shapes and textures.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 12663–12673, 2023.
|
||||
Mildenhall et al. (2021)
|
||||
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng.
|
||||
Nerf: Representing scenes as neural radiance fields for view synthesis.
|
||||
Communications of the ACM , 65(1):99–106, 2021.
|
||||
Mittal et al. (2022)
|
||||
Paritosh Mittal, Yen-Chi Cheng, Maneesh Singh, and Shubham Tulsiani.
|
||||
Autosdf: Shape priors for 3d completion, reconstruction and generation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 306–315, 2022.
|
||||
Müller et al. (2022)
|
||||
Thomas Müller, Alex Evans, Christoph Schied, and Alexander Keller.
|
||||
Instant neural graphics primitives with a multiresolution hash encoding.
|
||||
ACM Trans. Graph. , 41(4):102:1–102:15, July 2022.
|
||||
doi: 10.1145/3528223.3530127 .
|
||||
URL https://doi.org/10.1145/3528223.3530127 .
|
||||
Nam et al. (2022)
|
||||
Gimin Nam, Mariem Khlifi, Andrew Rodriguez, Alberto Tono, Linqi Zhou, and Paul Guerrero.
|
||||
3d-ldm: Neural implicit 3d shape generation with latent diffusion models.
|
||||
arXiv preprint arXiv:2212.00842 , 2022.
|
||||
Nichol et al. (2022)
|
||||
Alex Nichol, Heewoo Jun, Prafulla Dhariwal, Pamela Mishkin, and Mark Chen.
|
||||
Point-e: A system for generating 3d point clouds from complex prompts.
|
||||
arXiv preprint arXiv:2212.08751 , 2022.
|
||||
Poole et al. (2022)
|
||||
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall.
|
||||
Dreamfusion: Text-to-3d using 2d diffusion.
|
||||
arXiv preprint arXiv:2209.14988 , 2022.
|
||||
Qian et al. (2023)
|
||||
Guocheng Qian, Jinjie Mai, Abdullah Hamdi, Jian Ren, Aliaksandr Siarohin, Bing Li, Hsin-Ying Lee, Ivan Skorokhodov, Peter Wonka, Sergey Tulyakov, et al.
|
||||
Magic123: One image to high-quality 3d object generation using both 2d and 3d diffusion priors.
|
||||
arXiv preprint arXiv:2306.17843 , 2023.
|
||||
Radford et al. (2021)
|
||||
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.
|
||||
Learning transferable visual models from natural language supervision.
|
||||
In International conference on machine learning , pp. 8748–8763. PMLR, 2021.
|
||||
Raj et al. (2023)
|
||||
Amit Raj, Srinivas Kaza, Ben Poole, Michael Niemeyer, Nataniel Ruiz, Ben Mildenhall, Shiran Zada, Kfir Aberman, Michael Rubinstein, Jonathan Barron, et al.
|
||||
Dreambooth3d: Subject-driven text-to-3d generation.
|
||||
arXiv preprint arXiv:2303.13508 , 2023.
|
||||
Ramesh et al. (2022)
|
||||
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.
|
||||
Hierarchical text-conditional image generation with clip latents.
|
||||
arXiv preprint arXiv:2204.06125 , 1(2):3, 2022.
|
||||
Rombach et al. (2021)
|
||||
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
|
||||
High-resolution image synthesis with latent diffusion models, 2021.
|
||||
Rombach et al. (2022)
|
||||
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
|
||||
High-resolution image synthesis with latent diffusion models.
|
||||
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pp. 10684–10695, 2022.
|
||||
Ruiz et al. (2023)
|
||||
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman.
|
||||
Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 22500–22510, 2023.
|
||||
Saharia et al. (2022)
|
||||
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al.
|
||||
Photorealistic text-to-image diffusion models with deep language understanding.
|
||||
Advances in Neural Information Processing Systems , 35:36479–36494, 2022.
|
||||
Sanghi et al. (2022)
|
||||
Aditya Sanghi, Hang Chu, Joseph G Lambourne, Ye Wang, Chin-Yi Cheng, Marco Fumero, and Kamal Rahimi Malekshan.
|
||||
Clip-forge: Towards zero-shot text-to-shape generation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 18603–18613, 2022.
|
||||
Sargent et al. (2023)
|
||||
Kyle Sargent, Jing Yu Koh, Han Zhang, Huiwen Chang, Charles Herrmann, Pratul Srinivasan, Jiajun Wu, and Deqing Sun.
|
||||
Vq3d: Learning a 3d-aware generative model on imagenet.
|
||||
arXiv preprint arXiv:2302.06833 , 2023.
|
||||
Shen et al. (2021)
|
||||
Tianchang Shen, Jun Gao, Kangxue Yin, Ming-Yu Liu, and Sanja Fidler.
|
||||
Deep marching tetrahedra: a hybrid representation for high-resolution 3d shape synthesis.
|
||||
Advances in Neural Information Processing Systems , 34:6087–6101, 2021.
|
||||
Shonenkov et al. (2023)
|
||||
Alex Shonenkov, Misha Konstantinov, Daria Bakshandaeva, Christoph Schuhmann, Ksenia Ivanova, and Nadiia Klokova.
|
||||
DeepFloyd IF: A modular cascaded diffusion model.
|
||||
https://github.com/deep-floyd/IF/tree/develop , 2023.
|
||||
Skorokhodov et al. (2023)
|
||||
Ivan Skorokhodov, Aliaksandr Siarohin, Yinghao Xu, Jian Ren, Hsin-Ying Lee, Peter Wonka, and Sergey Tulyakov.
|
||||
3d generation on imagenet.
|
||||
arXiv preprint arXiv:2303.01416 , 2023.
|
||||
Sun et al. (2022)
|
||||
Jingxiang Sun, Xuan Wang, Yichun Shi, Lizhen Wang, Jue Wang, and Yebin Liu.
|
||||
Ide-3d: Interactive disentangled editing for high-resolution 3d-aware portrait synthesis.
|
||||
ACM Trans. Graph. , 41(6), nov 2022.
|
||||
ISSN 0730-0301.
|
||||
doi: 10.1145/3550454.3555506 .
|
||||
URL https://doi.org/10.1145/3550454.3555506 .
|
||||
Sun et al. (2023)
|
||||
Jingxiang Sun, Xuan Wang, Lizhen Wang, Xiaoyu Li, Yong Zhang, Hongwen Zhang, and Yebin Liu.
|
||||
Next3d: Generative neural texture rasterization for 3d-aware head avatars.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 20991–21002, June 2023.
|
||||
Tang et al. (2022)
|
||||
Junshu Tang, Bo Zhang, Binxin Yang, Ting Zhang, Dong Chen, Lizhuang Ma, and Fang Wen.
|
||||
Explicitly controllable 3d-aware portrait generation.
|
||||
arXiv preprint arXiv:2209.05434 , 2022.
|
||||
Tang et al. (2023)
|
||||
Junshu Tang, Tengfei Wang, Bo Zhang, Ting Zhang, Ran Yi, Lizhuang Ma, and Dong Chen.
|
||||
Make-it-3d: High-fidelity 3d creation from a single image with diffusion prior.
|
||||
arXiv preprint arXiv:2303.14184 , 2023.
|
||||
Wang et al. (2022)
|
||||
Haochen Wang, Xiaodan Du, Jiahao Li, Raymond A. Yeh, and Greg Shakhnarovich.
|
||||
Score jacobian chaining: Lifting pretrained 2d diffusion models for 3d generation.
|
||||
arXiv preprint arXiv:2212.00774 , 2022.
|
||||
Wang et al. (2021)
|
||||
Peng Wang, Lingjie Liu, Yuan Liu, Christian Theobalt, Taku Komura, and Wenping Wang.
|
||||
Neus: Learning neural implicit surfaces by volume rendering for multi-view reconstruction.
|
||||
arXiv preprint arXiv:2106.10689 , 2021.
|
||||
Wang et al. (2023a)
|
||||
Tengfei Wang, Bo Zhang, Ting Zhang, Shuyang Gu, Jianmin Bao, Tadas Baltrusaitis, Jingjing Shen, Dong Chen, Fang Wen, Qifeng Chen, et al.
|
||||
Rodin: A generative model for sculpting 3d digital avatars using diffusion.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 4563–4573, 2023a.
|
||||
Wang et al. (2023b)
|
||||
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu.
|
||||
Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation.
|
||||
arXiv preprint arXiv:2305.16213 , 2023b.
|
||||
Watson et al. (2022)
|
||||
Daniel Watson, William Chan, Ricardo Martin-Brualla, Jonathan Ho, Andrea Tagliasacchi, and Mohammad Norouzi.
|
||||
Novel view synthesis with diffusion models.
|
||||
arXiv preprint arXiv:2210.04628 , 2022.
|
||||
Wu et al. (2023)
|
||||
Tong Wu, Jiarui Zhang, Xiao Fu, Yuxin Wang, Jiawei Ren, Liang Pan, Wayne Wu, Lei Yang, Jiaqi Wang, Chen Qian, et al.
|
||||
Omniobject3d: Large-vocabulary 3d object dataset for realistic perception, reconstruction and generation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 803–814, 2023.
|
||||
Xiang et al. (2023)
|
||||
Jianfeng Xiang, Jiaolong Yang, Binbin Huang, and Xin Tong.
|
||||
3d-aware image generation using 2d diffusion models.
|
||||
arXiv preprint arXiv:2303.17905 , 2023.
|
||||
Xie et al. (2021)
|
||||
Chulin Xie, Chuxin Wang, Bo Zhang, Hao Yang, Dong Chen, and Fang Wen.
|
||||
Style-based point generator with adversarial rendering for point cloud completion.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 4619–4628, 2021.
|
||||
Yan et al. (2022)
|
||||
Xingguang Yan, Liqiang Lin, Niloy J Mitra, Dani Lischinski, Daniel Cohen-Or, and Hui Huang.
|
||||
Shapeformer: Transformer-based shape completion via sparse representation.
|
||||
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 6239–6249, 2022.
|
||||
Yu et al. (2023)
|
||||
Wang Yu, Xuelin Qian, Jingyang Huo, Tiejun Huang, Bo Zhao, and Yanwei Fu.
|
||||
Pushing the limits of 3d shape generation at scale.
|
||||
arXiv preprint arXiv:2306.11510 , 2023.
|
||||
Zeng et al. (2022)
|
||||
Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojcic, Or Litany, Sanja Fidler, and Karsten Kreis.
|
||||
Lion: Latent point diffusion models for 3d shape generation.
|
||||
arXiv preprint arXiv:2210.06978 , 2022.
|
||||
Zhang et al. (2022)
|
||||
Biao Zhang, Matthias Nießner, and Peter Wonka.
|
||||
3dilg: Irregular latent grids for 3d generative modeling.
|
||||
In NeurIPS , 2022.
|
||||
Zhang et al. (2023a)
|
||||
Biao Zhang, Jiapeng Tang, Matthias Niessner, and Peter Wonka.
|
||||
3dshape2vecset: A 3d shape representation for neural fields and generative diffusion models.
|
||||
arXiv preprint arXiv:2301.11445 , 2023a.
|
||||
Zhang et al. (2023b)
|
||||
Huichao Zhang, Bowen Chen, Hao Yang, Liao Qu, Xu Wang, Li Chen, Chao Long, Feida Zhu, Kang Du, and Min Zheng.
|
||||
Avatarverse: High-quality & stable 3d avatar creation from text and pose.
|
||||
arXiv preprint arXiv:2308.03610 , 2023b.
|
||||
Zhang et al. (2018)
|
||||
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.
|
||||
The unreasonable effectiveness of deep features as a perceptual metric.
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition , pp. 586–595, 2018.
|
||||
## Appendix A Appendix
|
||||
## A.1 Implementation details
|
||||
Algorithms for Bootstrapped Score Distillation. We provide a summarized algorithm of bootstrapped score distillation in Algorithm 1 . The “Bootstrapped Score Distillation” algorithm starts by initializing n n ( n = 1 n=1 in our case) meshes and a pretrained text-to-image diffusion model ϵ DreamBooth \epsilon_{\text{DreamBooth}} parameterized by ϕ \phi . The algorithm then enters an iterative loop: in each iteration, it renders the mesh to obtain multi-view images x x , augments these images with Gaussian noises to form x t ′ = α t ′ x + σ t ′ ϵ x_{t^{\prime}}=\alpha_{t^{\prime}}x+\sigma_{t^{\prime}}\epsilon , and fine-tunes the pretrained diffusion model ϵ DreamBooth \epsilon_{\text{DreamBooth}} based on these augmented image renderings. Within each iteration, there’s an inner loop running for n n steps, where a random mesh and camera pose are sampled, and a 2D image is rendered from the chosen pose. Then, updates are performed on θ \theta and ϕ \phi using gradients calculated from the difference between the pretrained score function and the predicted score function, and from the ℒ 2 \mathcal{L}_{2} norm between the predicted score and real noise, respectively. These iterations continue until convergence, and the final refined mesh(s) are returned.
|
||||
Algorithm 1 Bootstrapped Score Distillation
|
||||
Input: Number of particles n n ( ≥ 1 ) (\geq 1) . Pretrained text-to-image diffusion model
|
||||
ϵ DreamBooth \mathbf{\epsilon}_{\text{DreamBooth}} . Learning rate η 1 \eta_{1} and η 2 \eta_{2} for 3D structures and diffusion model parameters, respectively. A prompt y y . Number of images m m and Camera poses { c r ( i ) } i = 1 m \{c_{r}^{(i)}\}_{i=1}^{m} for the multi-view datasets.
|
||||
1:
|
||||
initialize n n meshes { θ ( i ) } i = 1 n \{\theta^{(i)}\}_{i=1}^{n} , a noise prediction model ϵ ϕ \mathbf{\epsilon}_{\phi} parameterized by ϕ \phi .
|
||||
2:
|
||||
while not converged do
|
||||
3:
|
||||
Render the mesh to get multi-view images 𝒙 = g ( θ , c r ) \bm{x}=g(\theta,c_{r}) .
|
||||
4:
|
||||
Augment image renderings with Gaussian noises: 𝒙 t ′ = α t ′ 𝒙 + σ t ′ ϵ \bm{x}_{t^{\prime}}=\alpha_{t^{\prime}}\bm{x}+\sigma_{t^{\prime}}\bm{\epsilon} .
|
||||
5:
|
||||
Finetune ϵ DreamBooth \epsilon_{\textnormal{DreamBooth}} on augmented image renderings 𝒙 𝒓 = r t ′ ( 𝒙 ) \bm{x_{r}}=r_{t^{\prime}}(\bm{x}) .
|
||||
6:
|
||||
for i in T T steps do
|
||||
7:
|
||||
Randomly sample θ ∼ { θ ( i ) } i = 1 n \theta\sim\{\theta^{(i)}\}_{i=1}^{n} and a camera pose c c .
|
||||
8:
|
||||
Render the 3D structure θ \theta at pose c c to get a 2D image 𝒙 0 = 𝒈 ( θ , c ) {\bm{x}}_{0}={\bm{g}}(\theta,c) .
|
||||
9:
|
||||
θ ← θ − η 1 𝔼 t , ϵ , c [ ω ( t ) ( ϵ DreamBooth ( 𝐱 t , t , y ) − ϵ ϕ ( 𝐱 t , t , c , y ) ) ∂ g ( θ , c ) ∂ θ ] \theta\leftarrow\theta-\eta_{1}\mathbb{E}_{t,\bm{\epsilon},c}\left[\omega(t)\left(\epsilon_{\text{DreamBooth}}(\mathbf{x}_{t},t,y)-\epsilon_{\phi}(\mathbf{x}_{t},t,c,y)\right)\frac{\partial g(\theta,c)}{\partial\theta}\right]
|
||||
10:
|
||||
ϕ ← ϕ − η 2 ∇ ϕ 𝔼 t , ϵ ‖ ϵ ϕ ( 𝐱 t , t , c , y ) − ϵ ‖ 2 2 . \phi\leftarrow\phi-\eta_{2}\nabla_{\phi}\mathbb{E}_{t,\mathbf{\epsilon}}||\bm{\epsilon}_{\phi}(\mathbf{x}_{t},t,c,y)-\epsilon||_{2}^{2}\mbox{.}
|
||||
11:
|
||||
end for
|
||||
12:
|
||||
end while
|
||||
13:
|
||||
return
|
||||
Structure-aware latent regularization. To maintain the high-quality output produced by BSD while reducing noise and inconsistencies, we further incorporate a control net-guided inpainting diffusion model that regularizes the generated textures. Specifically, for a rendered image x from an arbitrary viewpoint, the visible section under the reference view is initially computed. This invariant portion during the generation process allows our inpainting model to fill in the remaining segments. As these remaining parts adhere to geometric constraints, we integrate geometric normal information through a control net. Ultimately, this method permits us to enforce view consistency and generate realistic results using a control-net guided inpainting diffusion model. To preserve the high-quality generation output, we avoid utilizing this image directly as a loss against the rendered image. Instead, we subtly introduce it by constraining the norm of the latent variables:
|
||||
ℒ reg ( ϕ , g ( θ ) ) = Σ ( ‖ E ( x ) ‖ 2 − ‖ E ( x reg ) ‖ 2 ) 2 . \mathcal{L}_{\textnormal{reg}}(\phi,g(\theta))=\Sigma(\|E(x)\|_{2}-\|E(x_{\textnormal{reg}})\|_{2})^{2}.
|
||||
(8)
|
||||
Architectural details. In the Neus approach ( Wang et al., 2021 ) , we employ a single-layer Multi-Layer Perceptron (MLP) with 32 hidden units to simultaneously predict RGB color, volume density, and normal. The inputs to this MLP are the concatenated feature vectors derived from multi-resolution hash encoding sampled with trilinear interpolation. To sparsify the Instant NGP representation, we implement density-based pruning every 10 iterations within an octree structure, as suggested by Magic3D ( Lin et al., 2023 ) . In our experiments, we use a bounding sphere with a radius of 2. For the density prediction, we utilise the softplus activation function and, following the approach of Poole et al. 2022 , include an initial spatial density bias in order to encourage optimization in favor of the object-centric neural field.
|
||||
Camera and light augmentations. We follow Magic3D to add random augmentations to the camera and light sampling for rendering the shaded images. Differently, we sample the point light location such that the angular distance from the random camera center location (w.r.t. the origin) is sampled from ϕ cam ∼ 𝒰 ( 0 , π / 3 ) \phi_{\textnormal{cam}}\sim\mathcal{U}(0,\pi/3) with a random point light distance r cam ∼ ( 7.5 , 10 ) r_{\textnormal{cam}}\sim(7.5,10) , and (b) we freeze the material augmentation unlike Dreamfusion and Magic3D, as we found it is bad for training convergence (c) In the coarse neus stage, we propose a fixed-random mixed camera pose strategy. Specifically, following the common practice of the current text-to-3D methods, random camera view sampling benefits scene optimization. However, Zero1-to-3 needs fixed camera intrinsic parameters. Therefore, we let half of the GPUs sample the camera distance from 𝒰 ( 3.2 , 3.5 ) \mathcal{U}(3.2,3.5) , and the Field-of-View from 𝒰 ( 10 , 20 ) \mathcal{U}(10,20) , while the left GPUs are fixed to the default camera intrinsic.
|
||||
Time annealing. At the beginning of the geometry sculpting stage, We utilize a simple two-stage annealing of time step t in the score distillation objective. For the first several iterations we sample time steps t ∼ 𝒰 ( 0.7 , 0.85 ) t\sim\mathcal{U}(0.7,0.85) and then anneal into t ∼ 𝒰 ( 0.2 , 0.50 ) t\sim\mathcal{U}(0.2,0.50) . We refer the readers to ProlificDreamer ( Wang et al., 2023b ) . For the left iterations, we fix time steps to 𝒰 ( 0.2 , 0.50 ) \mathcal{U}(0.2,0.50) . We also utilize a simple two-stage time annealing for the multi-view dataset generation, that is, for the first updating step, we select a time step t = 0.5 t=0.5 for all rendered images and then anneal it into t = 0.1 t=0.1 along the later updating steps.
|
||||
## A.2 Additional experiments
|
||||
Figure 8: DreamCraft3D skillfully generates an assortment of visually compelling 3D models when provided with a textual description.
|
||||
Figure 9: Failure case. Our method learns incorrect geometry for elephant nose.
|
||||
Diversity. Prior studies frequently yield models of limited diversity with disproportionately smooth textures. Our approach to superior text-to-3D generation initially translates the text prompt into a reference image via 2D diffusion before implementing our proprietary image-based 3D creation methodology. Figure 8 demonstrates the proficiency of our method in generating an array of diverse models from a single text prompt, all characterized by their remarkable quality.
|
||||
## A.3 Limitations
|
||||
Our approach occasionally incorporates frontal-view geometric details into texture, as depicted in Figure 9 , due to depth ambiguity and inaccuracies in the depth prior. Furthermore, we do not expressly segregate material and lighting from the 2D reference image, an aspect deferred for future exploration.
|
||||
## A.4 Additional qualitative results
|
||||
Figure 10 − - Figure 13 provides more results produced by DreamCraft3D. Our method is able to produce photo-realistic 3D assets with compelling textural details. Moreover, our method shows significantly improved 3D consistency. Please find video results in the supplementary video.
|
||||
Figure 10: Additional results of DreamCraft3D.
|
||||
Figure 11: Additional results of DreamCraft3D.
|
||||
Figure 12: Additional results of DreamCraft3D.
|
||||
Figure 13: Additional results of DreamCraft3D.
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
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
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,750 @@
|
||||
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.
|
||||
## References
|
||||
Achiam et al. [2023]
|
||||
J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al.
|
||||
Gpt-4 technical report.
|
||||
arXiv preprint arXiv:2303.08774 , 2023.
|
||||
Avigad [2023]
|
||||
J. Avigad.
|
||||
Mathematics and the formal turn, 2023.
|
||||
Azerbayev et al. [2023]
|
||||
Z. Azerbayev, H. Schoelkopf, K. Paster, M. D. Santos, S. McAleer, A. Q. Jiang, J. Deng, S. Biderman, and S. Welleck.
|
||||
Llemma: An open language model for mathematics.
|
||||
arXiv preprint arXiv:2310.10631 , 2023.
|
||||
Bansal et al. [2019]
|
||||
K. Bansal, S. Loos, M. Rabe, C. Szegedy, and S. Wilcox.
|
||||
Holist: An environment for machine learning of higher order logic theorem proving.
|
||||
In International Conference on Machine Learning , pages 454–463. PMLR, 2019.
|
||||
Bibel [2013]
|
||||
W. Bibel.
|
||||
Automated theorem proving .
|
||||
Springer Science & Business Media, 2013.
|
||||
Crouse et al. [2021]
|
||||
M. Crouse, I. Abdelaziz, B. Makni, S. Whitehead, C. Cornelio, P. Kapanipathi, K. Srinivas, V. Thost, M. Witbrock, and A. Fokoue.
|
||||
A deep reinforcement learning approach to first-order logic theorem proving.
|
||||
In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 6279–6287, 2021.
|
||||
De Moura et al. [2015]
|
||||
L. De Moura, S. Kong, J. Avigad, F. Van Doorn, and J. von Raumer.
|
||||
The lean theorem prover (system description).
|
||||
In Automated Deduction-CADE-25: 25th International Conference on Automated Deduction, Berlin, Germany, August 1-7, 2015, Proceedings 25 , pages 378–388. Springer, 2015.
|
||||
First et al. [2023]
|
||||
E. First, M. N. Rabe, T. Ringer, and Y. Brun.
|
||||
Baldur: Whole-proof generation and repair with large language models, 2023.
|
||||
Han et al. [2021]
|
||||
J. M. Han, J. Rute, Y. Wu, E. W. Ayers, and S. Polu.
|
||||
Proof artifact co-training for theorem proving with language models.
|
||||
arXiv preprint arXiv:2102.06203 , 2021.
|
||||
Huang et al. [2024]
|
||||
Y. Huang, X. Lin, Z. Liu, Q. Cao, H. Xin, H. Wang, Z. Li, L. Song, and X. Liang.
|
||||
Mustard: Mastering uniform synthesis of theorem and proof data.
|
||||
arXiv preprint arXiv:2402.08957 , 2024.
|
||||
Jiang et al. [2021]
|
||||
A. Q. Jiang, W. Li, J. M. Han, and Y. Wu.
|
||||
Lisa: Language models of isabelle proofs.
|
||||
In 6th Conference on Artificial Intelligence and Theorem Proving , pages 378–392, 2021.
|
||||
Jiang et al. [2022a]
|
||||
A. Q. Jiang, W. Li, S. Tworkowski, K. Czechowski, T. Odrzygóźdź, P. Miłoś, Y. Wu, and M. Jamnik.
|
||||
Thor: Wielding hammers to integrate language models and automated theorem provers.
|
||||
Advances in Neural Information Processing Systems , 35:8360–8373, 2022a.
|
||||
Jiang et al. [2022b]
|
||||
A. Q. Jiang, S. Welleck, J. P. Zhou, W. Li, J. Liu, M. Jamnik, T. Lacroix, Y. Wu, and G. Lample.
|
||||
Draft, sketch, and prove: Guiding formal theorem provers with informal proofs.
|
||||
arXiv preprint arXiv:2210.12283 , 2022b.
|
||||
Jiang et al. [2023]
|
||||
A. Q. Jiang, W. Li, and M. Jamnik.
|
||||
Multilingual mathematical autoformalization.
|
||||
arXiv preprint arXiv:2311.03755 , 2023.
|
||||
Kaliszyk et al. [2018]
|
||||
C. Kaliszyk, J. Urban, H. Michalewski, and M. Olšák.
|
||||
Reinforcement learning of theorem proving.
|
||||
Advances in Neural Information Processing Systems , 31, 2018.
|
||||
Kovács and Voronkov [2013]
|
||||
L. Kovács and A. Voronkov.
|
||||
First-order theorem proving and vampire.
|
||||
In International Conference on Computer Aided Verification , pages 1–35. Springer, 2013.
|
||||
Lample et al. [2022]
|
||||
G. Lample, T. Lacroix, M.-A. Lachaux, A. Rodriguez, A. Hayat, T. Lavril, G. Ebner, and X. Martinet.
|
||||
Hypertree proof search for neural theorem proving.
|
||||
Advances in neural information processing systems , 35:26337–26349, 2022.
|
||||
Liu et al. [2023]
|
||||
C. Liu, J. Shen, H. Xin, Z. Liu, Y. Yuan, H. Wang, W. Ju, C. Zheng, Y. Yin, L. Li, et al.
|
||||
Fimo: A challenge formal dataset for automated theorem proving.
|
||||
arXiv preprint arXiv:2309.04295 , 2023.
|
||||
Loos et al. [2017]
|
||||
S. Loos, G. Irving, C. Szegedy, and C. Kaliszyk.
|
||||
Deep network guided proof search.
|
||||
arXiv preprint arXiv:1701.06972 , 2017.
|
||||
Moura and Ullrich [2021]
|
||||
L. d. Moura and S. Ullrich.
|
||||
The lean 4 theorem prover and programming language.
|
||||
In Automated Deduction–CADE 28: 28th International Conference on Automated Deduction, Virtual Event, July 12–15, 2021, Proceedings 28 , pages 625–635. Springer, 2021.
|
||||
Paulson [1994]
|
||||
L. C. Paulson.
|
||||
Isabelle a Generic Theorem Prover .
|
||||
Springer Verlag, 1994.
|
||||
Polu and Sutskever [2020]
|
||||
S. Polu and I. Sutskever.
|
||||
Generative language modeling for automated theorem proving.
|
||||
arXiv preprint arXiv:2009.03393 , 2020.
|
||||
Polu et al. [2022]
|
||||
S. Polu, J. M. Han, K. Zheng, M. Baksys, I. Babuschkin, and I. Sutskever.
|
||||
Formal mathematics statement curriculum learning.
|
||||
arXiv preprint arXiv:2202.01344 , 2022.
|
||||
Radford et al. [2019]
|
||||
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al.
|
||||
Language models are unsupervised multitask learners.
|
||||
OpenAI blog , 1(8):9, 2019.
|
||||
Schulz [2002]
|
||||
S. Schulz.
|
||||
E–a brainiac theorem prover.
|
||||
Ai Communications , 15(2-3):111–126, 2002.
|
||||
Shao et al. [2024]
|
||||
Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, M. Zhang, Y. Li, Y. Wu, and D. Guo.
|
||||
Deepseekmath: Pushing the limits of mathematical reasoning in open language models.
|
||||
arXiv preprint arXiv:2402.03300 , 2024.
|
||||
Shulman [2024]
|
||||
M. Shulman.
|
||||
Strange new universes: Proof assistants and synthetic foundations, 2024.
|
||||
Thakur et al. [2023]
|
||||
A. Thakur, Y. Wen, and S. Chaudhuri.
|
||||
A language-agent approach to formal theorem-proving.
|
||||
arXiv preprint arXiv:2310.04353 , 2023.
|
||||
[29]
|
||||
The Coq Development Team.
|
||||
Coq.
|
||||
URL https://coq.inria.fr .
|
||||
Vaswani et al. [2017]
|
||||
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.
|
||||
Attention is all you need.
|
||||
Advances in neural information processing systems , 30, 2017.
|
||||
Wang and Deng [2020]
|
||||
M. Wang and J. Deng.
|
||||
Learning to prove theorems by learning to generate theorems.
|
||||
Advances in Neural Information Processing Systems , 33:18146–18157, 2020.
|
||||
Wu et al. [2021]
|
||||
M. Wu, M. Norrish, C. Walder, and A. Dezfouli.
|
||||
Tacticzero: Learning to prove theorems from scratch with deep reinforcement learning.
|
||||
Advances in Neural Information Processing Systems , 34:9330–9342, 2021.
|
||||
Wu et al. [2020]
|
||||
Y. Wu, A. Q. Jiang, J. Ba, and R. Grosse.
|
||||
Int: An inequality benchmark for evaluating generalization in theorem proving.
|
||||
arXiv preprint arXiv:2007.02924 , 2020.
|
||||
Wu et al. [2022]
|
||||
Y. Wu, A. Q. Jiang, W. Li, M. Rabe, C. Staats, M. Jamnik, and C. Szegedy.
|
||||
Autoformalization with large language models.
|
||||
Advances in Neural Information Processing Systems , 35:32353–32368, 2022.
|
||||
Xin et al. [2023]
|
||||
H. Xin, H. Wang, C. Zheng, L. Li, Z. Liu, Q. Cao, Y. Huang, J. Xiong, H. Shi, E. Xie, et al.
|
||||
Lego-prover: Neural theorem proving with growing libraries.
|
||||
arXiv preprint arXiv:2310.00656 , 2023.
|
||||
Xiong et al. [2023]
|
||||
J. Xiong, J. Shen, Y. Yuan, H. Wang, Y. Yin, Z. Liu, L. Li, Z. Guo, Q. Cao, Y. Huang, et al.
|
||||
Trigo: Benchmarking formal mathematical proof reduction for generative language models.
|
||||
arXiv preprint arXiv:2310.10180 , 2023.
|
||||
Yang et al. [2024]
|
||||
K. Yang, A. Swope, A. Gu, R. Chalamala, P. Song, S. Yu, S. Godil, R. J. Prenger, and A. Anandkumar.
|
||||
Leandojo: Theorem proving with retrieval-augmented language models.
|
||||
Advances in Neural Information Processing Systems , 36, 2024.
|
||||
Zhao et al. [2023]
|
||||
X. Zhao, W. Li, and L. Kong.
|
||||
Decomposing the enigma: Subgoal-based demonstration learning for formal theorem proving.
|
||||
arXiv preprint arXiv:2305.16366 , 2023.
|
||||
Zheng et al. [2021]
|
||||
K. Zheng, J. M. Han, and S. Polu.
|
||||
Minif2f: a cross-system benchmark for formal olympiad-level mathematics.
|
||||
arXiv preprint arXiv:2109.00110 , 2021.
|
||||
## Appendix A Appendix / supplemental material
|
||||
## A.1 Prompts
|
||||
Specifically, we use the following format for scoring for the quality of the formalized statements:
|
||||
Prompt :
|
||||
To evaluate whether a formal Lean4 statement will be of interest to the community, consider the following criteria:
|
||||
1. Relevance to Current Research: Does the statement address a problem or concept that is actively being researched in mathematics or related fields? Higher relevance scores indicate greater potential interest.
|
||||
2. Complexity and Depth: Is the statement complex enough to challenge existing theories and methodologies, yet deep enough to provide significant insights or advancements? Complexity and depth showcase Lean4’s capabilities and attract interest.
|
||||
3. Interdisciplinary Potential: Does the statement offer opportunities for interdisciplinary research, connecting mathematics with other fields such as computer science, physics, or biology? Interdisciplinary projects often garner wide interest.
|
||||
4. Community Needs and Gaps: Does the statement fill an identified need or gap within the Lean4 community or the broader mathematical community? Addressing these needs directly correlates with interest.
|
||||
5. Innovativeness: How innovative is the statement? Does it propose new methods, concepts, or applications? Innovation drives interest and engagement.
|
||||
Customize your evaluation for each problem accordingly, assessing it as ‘excellent‘, ‘good‘, ‘above average‘, ‘fair‘ or ‘poor‘.
|
||||
You should respond in the following format for each statement:
|
||||
‘‘‘
|
||||
Translate the code to natural language: (Detailed explanation of the informal statement, including any relevant background information, assumptions, and definitions.)
|
||||
Analysis: (Provide a brief justification for each score, highlighting why the statement scored as it did across the criteria.)
|
||||
Assessment: (Based on the criteria, rate the statement as ‘excellent’, ‘good‘, ‘above average‘, ‘fair‘ or ‘poor‘.)
|
||||
‘‘‘
|
||||
## A.2 Case Studies of Autoformalization
|
||||
Example a.
|
||||
Problem in Natural Language : For a real number a a and a function f f defined on real numbers, where f ( x ) = x 3 − a x − 1 f(x)=x^{3}-ax-1 , if for all x x , f ( x ) ≤ 0 f(x)\leq 0 implies x x is either less than − 1 -1 or greater than 1 1 , then a a must equal 3 3 .
|
||||
Autoformalized Theorems with Complete Proofs :
|
||||
⬇
|
||||
example ( a : ℝ) ( f : ℝ → ℝ) ( h ₀ : ∀ x , f x = x ^ 3 - a * x - 1) :
|
||||
(∀ x , f x ≤ 0 → x ∈ Set . Iio (-1) ∪ Set . Ioi 1) → a = 3 := by
|
||||
intro h ₁
|
||||
have h ₂ := h ₁ 0
|
||||
simp [ h ₀] at h ₂
|
||||
have h ₃ := h ₁ 1
|
||||
simp [ h ₀] at h ₃
|
||||
have h ₄ := h ₁ (-1)
|
||||
simp [ h ₀] at h ₄
|
||||
linarith
|
||||
Analysis: This is a simple example to illustrate common cases in the autoformalization procedure. The model is required to recognize natural language descriptions in the problem statement and formalize them into proper mathematical operators. In this case, the statement " either less than − 1 -1 or greater than 1 1 " describes a union of interval ( − ∞ , − 1 ) ∪ ( 1 , + ∞ ) (-\infty,-1)\cup(1,+\infty) , which has been accurately formalized to Set.Iio (-1) ∪ \cup Set.Ioi 1 , where Set.Iio (-1) denotes an infinite-open interval ends at − 1 -1 , and Set.Ioi 1 denotes an open-infinite interval begins at 1 1 .
|
||||
Example b.
|
||||
Problem in Natural Language: Define two functions, F F and f f , where F ( k ) F(k) is the binomial coefficient of ( k + 2 ) (k+2) choose 2, and f ( k ) f(k) is the binomial coefficient of ( k + 1 ) (k+1) choose 1. For any natural number n n greater than or equal to 2, there exists a natural number k k such that f ( k ) f(k) equals ( n − 1 ) (n-1) times n n divided by 2. Furthermore, there exists a natural number m m such that F ( m ) F(m) equals ( n − 1 ) (n-1) times n n divided by 2, implying m m equals k k .
|
||||
Autoformalized Theorems with Complete Proofs:
|
||||
⬇
|
||||
example :
|
||||
let F := fun k => Nat . choose ( k + 2) 2;
|
||||
let f := fun k => Nat . choose ( k + 1) 1;
|
||||
∀ n : ℕ, 2 ≤ n → ∃ k : ℕ, f k = ( n - 1) * n / 2 → (∃ m : ℕ, F m = ( n - 1) * n / 2 → m = k ) := by
|
||||
simp [ Nat . choose , Nat . mul_sub_left_distrib , Nat . sub_sub , Nat . mul_one ]
|
||||
aesop
|
||||
Analysis: The formalization for this combinatorial problem involving binomial coefficients is semantically faithful and effective. The use of 𝙽𝚊𝚝 . 𝚌𝚑𝚘𝚘𝚜𝚎 \mathtt{Nat.choose} accurately represents the binomial coefficients, and the logical structure with existential quantifiers ( ∃ \exists ) and implications ( → \rightarrow ) aligns well with the problem’s requirements.
|
||||
Example c.
|
||||
Problem in Natural Language: Given three 3-dimensional vectors a a , b b , and c c , where the dot product of a a and b b equals − 3 -3 , the dot product of a a and c c equals 4 4 , and the dot product of b b and c c equals 6 6 . Prove that the dot product of b b with 7 c − 2 a 7c-2a equals 48 48 .
|
||||
Autoformalized Theorems with Complete Proofs:
|
||||
⬇
|
||||
example ( a b c : ℝ × ℝ × ℝ)
|
||||
( h ₀ : a .1 * b .1 + a .2.1 * b .2.1 + a .2.2 * b .2.2 = -3)
|
||||
( h ₁ : a .1 * c .1 + a .2.1 * c .2.1 + a .2.2 * c .2.2 = 4)
|
||||
( h ₂ : b .1 * c .1 + b .2.1 * c .2.1 + b .2.2 * c .2.2 = 6)
|
||||
: b .1 * (7 * c .1 - 2 * a .1) + b .2.1 * (7 * c .2.1 - 2 * a .2.1) + b .2.2 * (7 * c .2.2 - 2 * a .2.2) = 48 := by
|
||||
linarith [ h ₀, h ₁, h ₂]
|
||||
Analysis: This formalization result uses Cartesian product ℝ × ℝ × ℝ \mathbb{R}\times\mathbb{R}\times\mathbb{R} to represent a 3-dimensional vector, where the components of a 3-dimensional Cartesian product a are denoted by a.1 , a.2.1 , and a.2.2 in Lean. The notation of dot products is expanded to the sum of element-wise products by definition.
|
||||
## A.3 Example Lean Proofs Generated by Our Model
|
||||
## A.3.1 Results on MiniF2F-Test Dataset
|
||||
Example a.
|
||||
Problem in Natural Language:
|
||||
Show that for any real number x x and any natural number n n , if x > − 1 x>-1 , then ( 1 + n x ) ≤ ( 1 + x ) n (1+nx)\leq(1+x)^{n} .
|
||||
Formal Proof:
|
||||
⬇
|
||||
theorem induction_1pxpownlt1pnx ( x : ℝ) ( n : ℕ) ( h ₀ : -1 < x ) ( h ₁ : 0 < n )
|
||||
: 1 + ↑ n * x ≤ (1 + x ) ^ ( n : ℕ) := by
|
||||
induction ’ h ₁ with k h ₁ <;> simp_all [ Nat . succ_eq_add_one , pow_add , mul_add , add_assoc , add_comm , add_left_comm ]
|
||||
nlinarith [ mul_self_nonneg (1 + x - 1)]
|
||||
Example b.
|
||||
Problem in Natural Language:
|
||||
Ms. Blackwell gives an exam to two classes. The mean of the scores of the students in the morning class is 84 84 , and the afternoon class’s mean score is 70 70 . The ratio of the number of students in the morning class to the number of students in the afternoon class is 3 4 \frac{3}{4} . What is the mean of the scores of all the students?
|
||||
Show that it is 76.
|
||||
Formal Proof:
|
||||
⬇
|
||||
theorem amc12b_2021_p4 ( m a : ℕ) ( h ₀ : 0 < m ∧ 0 < a )
|
||||
( h ₁ : ↑ m / ↑ a = (3 : ℝ) / 4)
|
||||
: (84 * ↑ m + 70 * ↑ a ) / (↑ m + ↑ a ) = (76 : ℝ) := by
|
||||
have h ₂ := h ₀.1. ne ’
|
||||
have h ₃ := h ₀.2. ne ’
|
||||
field_simp at h ₂ h ₃ ⊢
|
||||
ring_nf
|
||||
norm_num
|
||||
rw [ div_eq_inv_mul ] at h ₁
|
||||
field_simp at h ₁
|
||||
linarith
|
||||
Example c.
|
||||
Problem in Natural Language:
|
||||
For how many positive integers m m does there exist at least one positive integer n such that m ⋅ n ≤ m + n m\cdot n\leq m+n ? Show that it is infinitely many.
|
||||
Formal Proof:
|
||||
⬇
|
||||
theorem amc12a_2002_p6 ( n : ℕ) ( h ₀ : 0 < n )
|
||||
: ∃ m , m > n ∧ ∃ p , m * p ≤ m + p := by
|
||||
simp_all only [ ge_iff_le , gt_iff_lt , mul_one , mul_add , mul_comm , mul_assoc , mul_left_comm ]
|
||||
use n + 1
|
||||
constructor
|
||||
exact Nat . lt_succ_self n
|
||||
use 1
|
||||
ring_nf
|
||||
nlinarith
|
||||
## A.3.2 Results on FIMO Dataset
|
||||
Example a.
|
||||
Problem in Natural Language:
|
||||
Determine all functions f f from the set of positive integers into the set of positive integers such that for all x x and y y there exists a non degenerated triangle with sides of lengths x x , f ( y ) f(y) and f ( y + f ( x ) − 1 ) f(y+f(x)-1) .
|
||||
The final answer is the identity function f ( x ) = x f(x)=x .
|
||||
Formal Proof:
|
||||
⬇
|
||||
theorem fimo_2009_algebra_p3
|
||||
( f : ℕ → ℕ)
|
||||
( h ₀ : ∀ x y , ∃ ( a b c : ℕ),
|
||||
a = x ∧
|
||||
b = f y ∧
|
||||
c = f ( y + f x - 1) ∧
|
||||
a + b > c ∧
|
||||
a + c > b ∧
|
||||
b + c > a ) :
|
||||
∀ x , f x = x := by
|
||||
have h ₁ : f 0 = 0 := by
|
||||
specialize h ₀ 0 0
|
||||
obtain ⟨ _ , _ , _ , h ₁, h ₂, h ₃⟩ := h ₀
|
||||
simp at *
|
||||
linarith
|
||||
intro x
|
||||
induction x with
|
||||
| zero => exact h ₁
|
||||
| succ x ih =>
|
||||
specialize h ₀ ( x + 1) 0
|
||||
obtain ⟨ _ , _ , _ , h ₁, h ₂, h ₃⟩ := h ₀
|
||||
simp at *
|
||||
linarith
|
||||
Example b.
|
||||
Problem in Natural Language:
|
||||
Prove that for every positive integer n n , there exists a fraction a b \frac{a}{b} where a a and b b are integers satisfying 0 < b ⩽ n + 1 0<b\leqslant\sqrt{n}+1 and n ⩽ a b ⩽ n + 1 \sqrt{n}\leqslant\frac{a}{b}\leqslant\sqrt{n+1} .
|
||||
Formal Proof:
|
||||
⬇
|
||||
theorem fimo_2016_algebra_p5_1
|
||||
( n : ℕ)
|
||||
( h ₀ : 0 < n ) :
|
||||
∃ a b : ℕ, 0 < b ∧ b ≤ n . sqrt + 1 ∧
|
||||
n . sqrt ≤ a / b ∧ a / b ≤ ( n + 1). sqrt := by
|
||||
use n . sqrt , 1
|
||||
constructor <;>
|
||||
simp [ h ₀, Nat . sqrt_le_sqrt ] <;>
|
||||
norm_num <;>
|
||||
linarith [ Nat . sqrt_le_sqrt n ]
|
||||
## A.4 Formal Verification Details
|
||||
We verify the generated Lean 4 code with the following code as the prefix:
|
||||
⬇
|
||||
import Mathlib . Algebra . Algebra . Basic
|
||||
import Mathlib . Algebra . Order . Floor
|
||||
import Mathlib . Algebra . Associated
|
||||
import Mathlib . Algebra . BigOperators . Basic
|
||||
import Mathlib . Algebra . BigOperators . Order
|
||||
import Mathlib . Algebra . BigOperators . Pi
|
||||
import Mathlib . Algebra . GeomSum
|
||||
import Mathlib . Algebra . Group . Pi . Basic
|
||||
import Mathlib . Algebra . Group . Commute . Basic
|
||||
import Mathlib . Algebra . GroupPower . Basic
|
||||
import Mathlib . Algebra . GroupPower . Identities
|
||||
import Mathlib . Algebra . Order . Floor
|
||||
import Mathlib . Algebra . QuadraticDiscriminant
|
||||
import Mathlib . Algebra . Ring . Basic
|
||||
import Mathlib . Analysis . Asymptotics . AsymptoticEquivalent
|
||||
import Mathlib . Analysis . NormedSpace . Basic
|
||||
import Mathlib . Analysis . SpecialFunctions . Log . Basic
|
||||
import Mathlib . Analysis . SpecialFunctions . Log . Base
|
||||
import Mathlib . Combinatorics . SimpleGraph . Basic
|
||||
import Mathlib . Data . Complex . Basic
|
||||
import Mathlib . Data . Complex . Exponential
|
||||
import Mathlib . Data . Finset . Basic
|
||||
import Mathlib . Data . Fintype . Card
|
||||
import Mathlib . Data . Int . Basic
|
||||
import Mathlib . Data . Int . GCD
|
||||
import Mathlib . Data . Int . ModEq
|
||||
import Mathlib . Data . Int . Parity
|
||||
import Mathlib . Data . List . Intervals
|
||||
import Mathlib . Data . List . Palindrome
|
||||
import Mathlib . Data . Multiset . Basic
|
||||
import Mathlib . Data . Nat . Basic
|
||||
import Mathlib . Data . Nat . Choose . Basic
|
||||
import Mathlib . Data . Nat . Digits
|
||||
import Mathlib . Data . Nat . Factorial . Basic
|
||||
import Mathlib . Data . Nat . ModEq
|
||||
import Mathlib . Data . Nat . Multiplicity
|
||||
import Mathlib . Data . Nat . Parity
|
||||
import Mathlib . Data . Nat . Prime
|
||||
import Mathlib . Data . PNat . Basic
|
||||
import Mathlib . Data . PNat . Prime
|
||||
import Mathlib . Data . Polynomial . Basic
|
||||
import Mathlib . Data . Polynomial . Eval
|
||||
import Mathlib . Data . Real . Basic
|
||||
import Mathlib . Data . Real . Irrational
|
||||
import Mathlib . Data . Real . NNReal
|
||||
import Mathlib . Data . Real . Sqrt
|
||||
import Mathlib . Data . Set . Finite
|
||||
import Mathlib . Data . Sym . Sym2
|
||||
import Mathlib . Data . ZMod . Basic
|
||||
import Mathlib . Dynamics . FixedPoints . Basic
|
||||
import Mathlib . LinearAlgebra . AffineSpace . AffineMap
|
||||
import Mathlib . LinearAlgebra . AffineSpace . Independent
|
||||
import Mathlib . LinearAlgebra . AffineSpace . Ordered
|
||||
import Mathlib . LinearAlgebra . FiniteDimensional
|
||||
import Mathlib . Logic . Equiv . Basic
|
||||
import Mathlib . Order . Filter . Basic
|
||||
import Mathlib . Order . LocallyFinite
|
||||
import Mathlib . Order . WellFounded
|
||||
import Mathlib . Topology . Basic
|
||||
import Mathlib . Topology . Instances . NNReal
|
||||
import Aesop
|
||||
set_option maxHeartbeats 0
|
||||
set_option trace . aesop true
|
||||
set_option trace . aesop . proof true
|
||||
open Nat Real Rat BigOperators
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
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,482 @@
|
||||
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.
|
||||
## References
|
||||
Dai et al. (2024)
|
||||
Damai Dai, Chengqi Deng, Chenggang Zhao, Runxin Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang.
|
||||
Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.
|
||||
ArXiv , abs/2401.06066, 2024.
|
||||
URL https://api.semanticscholar.org/CorpusID:266933338 .
|
||||
DeepSeek-AI et al. (2024)
|
||||
DeepSeek-AI, Qihao Zhu, Daya Guo, Zhihong Shao, Dejian Yang, Peiyi Wang, Runxin Xu, Y. Wu, Yukun Li, Huazuo Gao, Shirong Ma, Wangding Zeng, Xiao Bi, Zihui Gu, Hanwei Xu, Damai Dai, Kai Dong, Liyue Zhang, Yishi Piao, Zhibin Gou, Zhenda Xie, Zhewen Hao, Bing-Li Wang, Jun-Mei Song, Deli Chen, Xin Xie, Kang Guan, Yu mei You, Aixin Liu, Qiushi Du, Wenjun Gao, Xuan Lu, Qinyu Chen, Yaohui Wang, Chengqi Deng, Jiashi Li, Chenggang Zhao, Chong Ruan, Fuli Luo, and Wenfeng Liang.
|
||||
Deepseek-coder-v2: Breaking the barrier of closed-source models in code intelligence.
|
||||
ArXiv , abs/2406.11931, 2024.
|
||||
URL https://api.semanticscholar.org/CorpusID:270562723 .
|
||||
Fedus et al. (2021)
|
||||
William Fedus, Barret Zoph, and Noam M. Shazeer.
|
||||
Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.
|
||||
J. Mach. Learn. Res. , 23:120:1–120:39, 2021.
|
||||
URL https://api.semanticscholar.org/CorpusID:231573431 .
|
||||
Lepikhin et al. (2020)
|
||||
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam M. Shazeer, and Z. Chen.
|
||||
Gshard: Scaling giant models with conditional computation and automatic sharding.
|
||||
ArXiv , abs/2006.16668, 2020.
|
||||
URL https://api.semanticscholar.org/CorpusID:220265858 .
|
||||
Loshchilov & Hutter (2016)
|
||||
Ilya Loshchilov and Frank Hutter.
|
||||
Sgdr: Stochastic gradient descent with warm restarts.
|
||||
arXiv: Learning , 2016.
|
||||
URL https://api.semanticscholar.org/CorpusID:14337532 .
|
||||
Sennrich et al. (2015)
|
||||
Rico Sennrich, Barry Haddow, and Alexandra Birch.
|
||||
Neural machine translation of rare words with subword units.
|
||||
ArXiv , abs/1508.07909, 2015.
|
||||
URL https://api.semanticscholar.org/CorpusID:1114678 .
|
||||
Shao et al. (2024)
|
||||
Zhihong Shao, Damai Dai, Daya Guo, Bo Liu, and Zihan Wang.
|
||||
Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.
|
||||
ArXiv , abs/2405.04434, 2024.
|
||||
URL https://api.semanticscholar.org/CorpusID:269613809 .
|
||||
Shazeer et al. (2017)
|
||||
Noam M. Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean.
|
||||
Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.
|
||||
ArXiv , abs/1701.06538, 2017.
|
||||
URL https://api.semanticscholar.org/CorpusID:12462234 .
|
||||
Vaswani et al. (2017)
|
||||
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.
|
||||
Attention is all you need.
|
||||
In Neural Information Processing Systems , 2017.
|
||||
URL https://api.semanticscholar.org/CorpusID:13756489 .
|
||||
Zhou et al. (2022)
|
||||
Yan-Quan Zhou, Tao Lei, Han-Chu Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V. Le, and James Laudon.
|
||||
Mixture-of-experts with expert choice routing.
|
||||
ArXiv , abs/2202.09368, 2022.
|
||||
URL https://api.semanticscholar.org/CorpusID:247011948 .
|
||||
## Appendix A Model Architecture
|
||||
We employ the DeepSeekMoE ( Dai et al., 2024 ) architecture as the backbone, which introduces shared experts to mitigate knowledge redundancy among routed experts:
|
||||
𝐡 t = 𝐮 t + ∑ i = 1 N s FFN i ( s ) ( 𝐮 t ) + ∑ i = 1 N r g i , t FFN i ( r ) ( 𝐮 t ) , \displaystyle\mathbf{h}_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N_{s}}\operatorname{FFN}_{i}^{(s)}\left(\mathbf{u}_{t}\right)+\sum_{i=1}^{N_{r}}g_{i,t}\operatorname{FFN}_{i}^{(r)}\left(\mathbf{u}_{t}\right),
|
||||
(6)
|
||||
where r r denotes the routed experts, while s s the shared experts. DeepSeekMoE replaces all FFN layers with MoE layers, except the dense FFN layer just after the input embedding layer.
|
||||
The detailed architecture hyper-parameters are listed in Table 5 .
|
||||
Table 5: Model architecture.
|
||||
hyper-parameters
|
||||
1B
|
||||
3B
|
||||
Vocab size
|
||||
32064
|
||||
32064
|
||||
Hidden size
|
||||
1024
|
||||
1280
|
||||
Attention heads
|
||||
8
|
||||
10
|
||||
MoE layers
|
||||
9
|
||||
11
|
||||
Granularity ( d ff d expert \frac{d_{\text{ff}}}{d_{\text{expert}}} )
|
||||
16 3 \frac{\text{16}}{\text{3}}
|
||||
4
|
||||
Shared experts
|
||||
2
|
||||
2
|
||||
Routed experts
|
||||
64
|
||||
64
|
||||
Activated routed experts
|
||||
6
|
||||
6
|
||||
## Appendix B Training Settings
|
||||
Following the work of Dai et al. (2024) , we initialize all learnable parameters with a standard deviation of 0.006, and set the maximum training sequence length to 2048.
|
||||
For the 1B model, we employ a cosine learning rate scheduler with warmup, setting the learning rate to 1e-3, the minimum learning rate to 1e-4, and the warmup steps to 1000. The training batch size for the 1B model is set to 1152, resulting in a total of 40000 training steps (100B tokens).
|
||||
For the 3B model, we use a multistep learning rate scheduler with stage steps = [45211, 50862, 56514] and corresponding stage learning rates of [7.8e-4, 2.47e-4, 7.8e-5]. The warmup steps for the 3B model are set to 2000. We use a training batch size of 1728 for the 3B model, resulting in a total of 56514 training steps (200B tokens).
|
||||
For validation, we leave around 70M tokens from the training corpus as the validation set (30 * 1B_batch_size * max_seq_len = 20 * 3B_batch_size * max_seq_len = 71M tokens).
|
||||
## Appendix C Experiments with Softmax Gate
|
||||
## C.1 Comparison of Sigmoid Gate Baseline and Softmax Gate Baseline
|
||||
We compare the sigmoid gate baseline and the softmax gate baseline with varying auxiliary loss coefficients α \alpha on a 1B-sized model. As shown in Figure 7 , the softmax gate exhibits higher perplexity under similar load balance conditions, and its performance is more sensitive to load imbalance compared to the sigmoid gate.
|
||||
Figure 7: Comparison of the sigmoid gate baseline and the softmax gate baseline. The softmax gate exhibits higher perplexity under similar load balance conditions and is more sensitive to load imbalance compared to the sigmoid gate.
|
||||
## C.2 Loss-Free Load Balancing with Softmax Gate
|
||||
Adjusting the per-expert bias for the softmax gate is more challenging due to the normalization property of softmax, which makes the score gap between two experts sensitive to the scores of other experts. In such a situation, we choose the 𝐛 i = 𝐛 i + u ∗ e i \mathbf{b}_{i}=\mathbf{b}_{i}+u*e_{i} variant to maintain load balance, where u u is set to 1e-3. For the baseline, we choose α \alpha = 0.0003, which yields the lowest perplexity for the softmax gate. The results are presented in Table 6 , showing that Loss-Free Balancing achieves a slightly lower perplexity while maintaining significantly better load balance compared to the auxiliary-loss training method. Figure 8 confirms that Loss-Free Balancing maintains a superior load balance throughout most of the training process.
|
||||
Table 6: For softmax gate, Loss-Free Balancing achieves a slightly lower perplexity while reaching a significantly better load balance compared to the auxiliary-loss training method.
|
||||
Load Balancing
|
||||
Perplexity
|
||||
MaxVio global \textbf{MaxVio}_{\text{global}}
|
||||
Loss-Controlled
|
||||
9.604
|
||||
0.937
|
||||
Loss-Free
|
||||
9.599
|
||||
0.027
|
||||
Figure 8: For softmax gate, Loss-Free Balancing maintains a superior load balance throughout most of the training process.
|
||||
## Appendix D Future Token Leakage in Expert Choice
|
||||
## D.1 Proof for Theoretical Leakage Amount
|
||||
Let R = K N R=\frac{K}{N} denote the MoE sparsity. Here K K denotes the average number of experts activated per token, and N N is the total number of experts.
|
||||
For an MoE layer in Expert Choice, the maximum information leakage I I (in bits per token), i.e., the information that the combinations of routing allocation can carry is:
|
||||
I = \displaystyle I=
|
||||
log 2 ( K T N T ) N / T \displaystyle\left.\log_{2}\binom{\frac{KT}{N}}{T}^{N}\right/T
|
||||
(7)
|
||||
> \displaystyle>
|
||||
N log 2 ( ( 1 − K N ) T ) K N T ( K N T ) K N T / T \displaystyle N\left.\log_{2}\frac{((1-\frac{K}{N})T)^{\frac{K}{N}T}}{(\frac{K}{N}T)^{\frac{K}{N}T}}\right/T
|
||||
= \displaystyle=
|
||||
K log 2 1 − R R . \displaystyle K\log_{2}\frac{1-R}{R}.
|
||||
For a model with a sparse ratio R = 2 16 = 0.125 R=\frac{2}{16}=0.125 and 9 MoE layers, the total leakage information is more than 50 bits per token.
|
||||
## D.2 Experimental Evidence
|
||||
We investigate the potential future token leakage of the Expert Choice by varying the chunk size used for experts’ top- k k selection, ranging from 512 tokens to 8192 tokens. 3 3
|
||||
3
|
||||
A chunk size of 2048 tokens means performing top- k k selection inside a sentence, while 512 tokens correspond to a quarter of a sentence and 8192 tokens to four sentences. We train a 2B MoE model on 100B tokens. The results, shown in Table 9 , reveal two key findings:
|
||||
1.
|
||||
Using a small chunk size of 512 leads to an abnormal loss drop, which can be attributed to significant future token leakage. A smaller chunk size allows the model to more easily exploit information from future tokens within the chunk during training.
|
||||
2.
|
||||
Shuffling tokens within a batch before chunking and selecting mitigates the observed loss drop. Such shuffling makes it more challenging for the model to utilize information leakage, as the future tokens are no longer in their original context. This finding supports the hypothesis that the loss drop originates from the model’s accessing and exploiting future token information.
|
||||
Figure 9: Comparison of Expert Choice with different chunk sizes and shuffling. Expert Choice with a chunk size of 512 exhibits a significant loss drop compared to chunk sizes of 8192 or 2048. Shuffling tokens eliminates this loss drop, indicating the presence of future token leakage.
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
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
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,824 @@
|
||||
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.
|
||||
## References
|
||||
Ainslie et al. (2023)
|
||||
J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai.
|
||||
Gqa: Training generalized multi-query transformer models from multi-head checkpoints.
|
||||
arXiv preprint arXiv:2305.13245 , 2023.
|
||||
Austin et al. (2021)
|
||||
J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, et al.
|
||||
Program synthesis with large language models.
|
||||
arXiv preprint arXiv:2108.07732 , 2021.
|
||||
Bai et al. (2023)
|
||||
Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, et al.
|
||||
Longbench: A bilingual, multitask benchmark for long context understanding.
|
||||
arXiv preprint arXiv:2308.14508 , 2023.
|
||||
Beltagy et al. (2020)
|
||||
I. Beltagy, M. E. Peters, and A. Cohan.
|
||||
Longformer: The long-document transformer.
|
||||
arXiv preprint arXiv:2004.05150 , 2020.
|
||||
Chen et al. (2024a)
|
||||
G. Chen, H. Shi, J. Li, Y. Gao, X. Ren, Y. Chen, X. Jiang, Z. Li, W. Liu, and C. Huang.
|
||||
Sepllm: Accelerate large language models by compressing one segment into one separator.
|
||||
arXiv preprint arXiv:2412.12094 , 2024a.
|
||||
Chen et al. (2021)
|
||||
M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. D. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al.
|
||||
Evaluating large language models trained on code.
|
||||
arXiv preprint arXiv:2107.03374 , 2021.
|
||||
Chen et al. (2024b)
|
||||
Z. Chen, R. Sadhukhan, Z. Ye, Y. Zhou, J. Zhang, N. Nolte, Y. Tian, M. Douze, L. Bottou, Z. Jia, et al.
|
||||
Magicpig: Lsh sampling for efficient llm generation.
|
||||
arXiv preprint arXiv:2410.16179 , 2024b.
|
||||
Cobbe et al. (2021)
|
||||
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al.
|
||||
Training verifiers to solve math word problems, 2021.
|
||||
URL https://arxiv. org/abs/2110.14168 , 2021.
|
||||
Dai et al. (2024)
|
||||
D. Dai, C. Deng, C. Zhao, R. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y. Wu, et al.
|
||||
Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.
|
||||
arXiv preprint arXiv:2401.06066 , 2024.
|
||||
DeepSeek-AI (2024)
|
||||
DeepSeek-AI.
|
||||
Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.
|
||||
2024.
|
||||
URL https://arxiv.org/abs/2405.04434 .
|
||||
DeepSeek-AI (2025)
|
||||
DeepSeek-AI.
|
||||
Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025.
|
||||
URL https://arxiv.org/abs/2501.12948 .
|
||||
Desai et al. (2024)
|
||||
A. Desai, S. Yang, A. Cuadron, A. Klimovic, M. Zaharia, J. E. Gonzalez, and I. Stoica.
|
||||
Hashattention: Semantic sparsity for faster inference.
|
||||
arXiv preprint arXiv:2412.14468 , 2024.
|
||||
Dua et al. (2019)
|
||||
D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner.
|
||||
Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs.
|
||||
arXiv preprint arXiv:1903.00161 , 2019.
|
||||
Fu et al. (2024a)
|
||||
T. Fu, H. Huang, X. Ning, G. Zhang, B. Chen, T. Wu, H. Wang, Z. Huang, S. Li, S. Yan, et al.
|
||||
Moa: Mixture of sparse attention for automatic large language model compression.
|
||||
arXiv preprint arXiv:2406.14909 , 2024a.
|
||||
Fu et al. (2024b)
|
||||
Y. Fu, Z. Cai, A. Asi, W. Xiong, Y. Dong, and W. Xiao.
|
||||
Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning.
|
||||
arXiv preprint arXiv:2410.19258 , 2024b.
|
||||
Gao et al. (2024)
|
||||
Y. Gao, Z. Zeng, D. Du, S. Cao, H. K.-H. So, T. Cao, F. Yang, and M. Yang.
|
||||
Seerattention: Learning intrinsic sparse attention in your llms.
|
||||
arXiv preprint arXiv:2410.13276 , 2024.
|
||||
Ge et al. (2023)
|
||||
S. Ge, Y. Zhang, L. Liu, M. Zhang, J. Han, and J. Gao.
|
||||
Model tells you what to discard: Adaptive kv cache compression for llms.
|
||||
arXiv preprint arXiv:2310.01801 , 2023.
|
||||
Google et al. (2024)
|
||||
G. T. Google, P. Georgiev, V. I. Lei, R. Burnell, L. Bai, A. Gulati, G. Tanzer, D. Vincent, Z. Pan, S. Wang, et al.
|
||||
Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.
|
||||
arXiv preprint arXiv:2403.05530 , 2024.
|
||||
Hendrycks et al. (2020)
|
||||
D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt.
|
||||
Measuring massive multitask language understanding.
|
||||
arXiv preprint arXiv:2009.03300 , 2020.
|
||||
Jiang et al. (2023)
|
||||
H. Jiang, Q. Wu, C.-Y. Lin, Y. Yang, and L. Qiu.
|
||||
Llmlingua: Compressing prompts for accelerated inference of large language models.
|
||||
arXiv preprint arXiv:2310.05736 , 2023.
|
||||
Jiang et al. (2024)
|
||||
H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C.-Y. Lin, et al.
|
||||
Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention.
|
||||
arXiv preprint arXiv:2407.02490 , 2024.
|
||||
Kamradt (2023)
|
||||
G. Kamradt.
|
||||
LLMTest NeedleInAHaystack.
|
||||
GitHub repository, 2023.
|
||||
URL https://github.com/gkamradt/LLMTest_NeedleInAHaystack .
|
||||
Accessed: [Insert Access Date Here].
|
||||
Li et al. (2023)
|
||||
H. Li, Y. Zhang, F. Koto, Y. Yang, H. Zhao, Y. Gong, N. Duan, and T. Baldwin.
|
||||
Cmmlu: Measuring massive multitask language understanding in chinese.
|
||||
arXiv preprint arXiv:2306.09212 , 2023.
|
||||
Li et al. (2024)
|
||||
Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen.
|
||||
Snapkv: Llm knows what you are looking for before generation.
|
||||
arXiv preprint arXiv:2404.14469 , 2024.
|
||||
Liu et al. (2024)
|
||||
G. Liu, C. Li, J. Zhao, C. Zhang, and M. Guo.
|
||||
Clusterkv: Manipulating llm kv cache in semantic space for recallable compression.
|
||||
arXiv preprint arXiv:2412.03213 , 2024.
|
||||
Park et al. (2023)
|
||||
J. S. Park, J. C. O’Brien, C. J. Cai, M. R. Morris, P. Liang, and M. S. Bernstein.
|
||||
Generative agents: Interactive simulacra of human behavior.
|
||||
In S. Follmer, J. Han, J. Steimle, and N. H. Riche, editors, Proceedings of the 36th Annual ACM Symposium on User Interface Software and Technology, UIST 2023, San Francisco, CA, USA, 29 October 2023– 1 November 2023 , pages 2:1–2:22. ACM, 2023.
|
||||
Peng et al. (2024)
|
||||
B. Peng, J. Quesnelle, H. Fan, and E. Shippole.
|
||||
Yarn: Efficient context window extension of large language models.
|
||||
In ICLR . OpenReview.net, 2024.
|
||||
Shazeer (2019)
|
||||
N. Shazeer.
|
||||
Fast transformer decoding: One write-head is all you need.
|
||||
CoRR , abs/1911.02150, 2019.
|
||||
Suzgun et al. (2022)
|
||||
M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. V. Le, E. H. Chi, D. Zhou, et al.
|
||||
Challenging big-bench tasks and whether chain-of-thought can solve them.
|
||||
arXiv preprint arXiv:2210.09261 , 2022.
|
||||
Tang et al. (2024)
|
||||
J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han.
|
||||
Quest: Query-aware sparsity for efficient long-context llm inference.
|
||||
arXiv preprint arXiv:2406.10774 , 2024.
|
||||
Tillet et al. (2019)
|
||||
P. Tillet, H.-T. Kung, and D. Cox.
|
||||
Triton: an intermediate language and compiler for tiled neural network computations.
|
||||
In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages , pages 10–19, 2019.
|
||||
Vaswani et al. (2017)
|
||||
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin.
|
||||
Attention is all you need.
|
||||
Advances in Neural Information Processing Systems , 2017.
|
||||
Wang et al. (2024)
|
||||
Y. Wang, X. Ma, G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj, X. He, Z. Jiang, et al.
|
||||
Mmlu-pro: A more robust and challenging multi-task language understanding benchmark.
|
||||
arXiv preprint arXiv:2406.01574 , 2024.
|
||||
Wei et al. (2022)
|
||||
J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al.
|
||||
Chain-of-thought prompting elicits reasoning in large language models.
|
||||
Advances in neural information processing systems , 35:24824–24837, 2022.
|
||||
Wu et al. (2024)
|
||||
W. Wu, Z. Pan, C. Wang, L. Chen, Y. Bai, K. Fu, Z. Wang, and H. Xiong.
|
||||
Tokenselect: Efficient long-context inference and length extrapolation for llms via dynamic token-level kv cache selection.
|
||||
arXiv preprint arXiv:2411.02886 , 2024.
|
||||
Xiao et al. (2024a)
|
||||
C. Xiao, P. Zhang, X. Han, G. Xiao, Y. Lin, Z. Zhang, Z. Liu, and M. Sun.
|
||||
Infllm: Training-free long-context extrapolation for llms with an efficient context memory.
|
||||
In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024a.
|
||||
Xiao et al. (2023)
|
||||
G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis.
|
||||
Efficient streaming language models with attention sinks.
|
||||
arXiv preprint arXiv:2309.17453 , 2023.
|
||||
Xiao et al. (2024b)
|
||||
G. Xiao, J. Tang, J. Zuo, J. Guo, S. Yang, H. Tang, Y. Fu, and S. Han.
|
||||
Duoattention: Efficient long-context llm inference with retrieval and streaming heads.
|
||||
arXiv preprint arXiv:2410.10819 , 2024b.
|
||||
Zaheer et al. (2020)
|
||||
M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al.
|
||||
Big bird: Transformers for longer sequences.
|
||||
Advances in neural information processing systems , 33:17283–17297, 2020.
|
||||
Zelikman et al. (2022)
|
||||
E. Zelikman, Y. Wu, J. Mu, and N. D. Goodman.
|
||||
Star: Bootstrapping reasoning with reasoning.
|
||||
In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 – December 9, 2022 , 2022.
|
||||
Zhang et al. (2023a)
|
||||
F. Zhang, B. Chen, Y. Zhang, J. Keung, J. Liu, D. Zan, Y. Mao, J. Lou, and W. Chen.
|
||||
Repocoder: Repository-level code completion through iterative retrieval and generation.
|
||||
In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6–10, 2023 , pages 2471–2484. Association for Computational Linguistics, 2023a.
|
||||
(42)
|
||||
K. Zhang, J. Li, G. Li, X. Shi, and Z. Jin.
|
||||
Codeagent: Enhancing code generation with tool-integrated agent systems for real-world repo-level coding challenges.
|
||||
In L. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11–16, 2024 , pages 13643–13658.
|
||||
Zhang et al. (2023b)
|
||||
Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, et al.
|
||||
H2o: Heavy-hitter oracle for efficient generative inference of large language models.
|
||||
Advances in Neural Information Processing Systems , 36:34661–34710, 2023b.
|
||||
Zhao et al. (2024)
|
||||
J. Zhao, Z. Fang, S. Li, S. Yang, and S. He.
|
||||
Buzz: Beehive-structured sparse kv cache with segmented heavy hitters for efficient llm inference.
|
||||
arXiv preprint arXiv:2410.23079 , 2024.
|
||||
Zhou et al. (2024)
|
||||
Z. Zhou, C. Li, X. Chen, S. Wang, Y. Chao, Z. Li, H. Wang, R. An, Q. Shi, Z. Tan, et al.
|
||||
Llm × \times mapreduce: Simplified long-sequence processing using large language models.
|
||||
arXiv preprint arXiv:2410.09342 , 2024.
|
||||
## Appendix A Example of AIME Results
|
||||
Prompt:
|
||||
“|User|Let x , y x,y and z z be positive real numbers that satisfy the following system of equations:
|
||||
log 2 ( x y z ) = 1 2 \log_{2}\left(\frac{x}{yz}\right)=\frac{1}{2}
|
||||
log 2 ( y x z ) = 1 3 \log_{2}\left(\frac{y}{xz}\right)=\frac{1}{3}
|
||||
log 2 ( z x y ) = 1 4 \log_{2}\left(\frac{z}{xy}\right)=\frac{1}{4}
|
||||
Then the value of | log 2 ( x 4 y 3 z 2 ) | \left|\log_{2}(x^{4}y^{3}z^{2})\right| is m n \frac{m}{n} where m m and n n are relatively prime positive integers. Find m + n m+n .
|
||||
Please reason step by step, and put your final answer within \boxed{} .|Assistant|”
|
||||
Our Result:
|
||||
⟨ \langle think ⟩ \rangle … ⟨ \langle /think ⟩ \rangle (2275 tokens)
|
||||
Given the system of equations:
|
||||
1. log 2 ( x y z ) = 1 2 \log_{2}\left(\frac{x}{yz}\right)=\frac{1}{2}
|
||||
2. log 2 ( y x z ) = 1 3 \log_{2}\left(\frac{y}{xz}\right)=\frac{1}{3}
|
||||
3. log 2 ( z x y ) = 1 4 \log_{2}\left(\frac{z}{xy}\right)=\frac{1}{4}
|
||||
We convert these logarithmic equations to their exponential forms:
|
||||
1. x y z = 2 \frac{x}{yz}=\sqrt{2}
|
||||
2. y x z = 2 1 / 3 \frac{y}{xz}=2^{1/3}
|
||||
3. z x y = 2 1 / 4 \frac{z}{xy}=2^{1/4}
|
||||
Expressing each variable in terms of x x , y y , and z z :
|
||||
1. x = 2 ⋅ y z x=\sqrt{2}\cdot yz
|
||||
2. y = 2 1 / 3 ⋅ x z y=2^{1/3}\cdot xz
|
||||
3. z = 2 1 / 4 ⋅ x y z=2^{1/4}\cdot xy
|
||||
Substituting these expressions into each other, we find:
|
||||
- From x = 2 ⋅ y z x=\sqrt{2}\cdot yz and y = 2 1 / 3 ⋅ x z y=2^{1/3}\cdot xz , we get x = 2 ⋅ ( 2 1 / 3 ⋅ x z ) ⋅ z x=\sqrt{2}\cdot(2^{1/3}\cdot xz)\cdot z , leading to z = 2 − 5 / 12 z=2^{-5/12} .
|
||||
- Substituting z z into the equation for y y , we get y = 2 − 3 / 8 y=2^{-3/8} .
|
||||
- Substituting z z and y y into the equation for x x , we get x = 2 − 7 / 24 x=2^{-7/24} .
|
||||
Next, we compute x 4 y 3 z 2 x^{4}y^{3}z^{2} :
|
||||
- x 4 = ( 2 − 7 / 24 ) 4 = 2 − 7 / 6 x^{4}=(2^{-7/24})^{4}=2^{-7/6}
|
||||
- y 3 = ( 2 − 3 / 8 ) 3 = 2 − 9 / 8 y^{3}=(2^{-3/8})^{3}=2^{-9/8}
|
||||
- z 2 = ( 2 − 5 / 12 ) 2 = 2 − 5 / 6 z^{2}=(2^{-5/12})^{2}=2^{-5/6}
|
||||
Multiplying these together:
|
||||
x 4 y 3 z 2 = 2 − 7 / 6 ⋅ 2 − 9 / 8 ⋅ 2 − 5 / 6 = 2 − 75 / 24 = 2 − 25 / 8 x^{4}y^{3}z^{2}=2^{-7/6}\cdot 2^{-9/8}\cdot 2^{-5/6}=2^{-75/24}=2^{-25/8}
|
||||
Taking the logarithm base 2:
|
||||
log 2 ( x 4 y 3 z 2 ) = − 25 / 8 \log_{2}(x^{4}y^{3}z^{2})=-25/8
|
||||
The absolute value is:
|
||||
| log 2 ( x 4 y 3 z 2 ) | = 25 / 8 |\log_{2}(x^{4}y^{3}z^{2})|=25/8
|
||||
Thus, m = 25 m=25 and n = 8 n=8 , giving m + n = 33 m+n=33 .
|
||||
33 \boxed{33}
|
||||
Baseline Result:
|
||||
⟨ \langle think ⟩ \rangle … ⟨ \langle /think ⟩ \rangle (9392 tokens)
|
||||
Given the system of logarithmic equations:
|
||||
log 2 ( x y z ) = 1 2 \log_{2}\left(\frac{x}{yz}\right)=\frac{1}{2}
|
||||
log 2 ( y x z ) = 1 3 \log_{2}\left(\frac{y}{xz}\right)=\frac{1}{3}
|
||||
log 2 ( z x y ) = 1 4 \log_{2}\left(\frac{z}{xy}\right)=\frac{1}{4}
|
||||
We convert these equations to exponential form:
|
||||
2 = x y z ⟹ x = 2 ⋅ y z \sqrt{2}=\frac{x}{yz}\implies x=\sqrt{2}\cdot yz
|
||||
2 1 3 = y x z ⟹ y = 2 1 3 ⋅ x z 2^{\frac{1}{3}}=\frac{y}{xz}\implies y=2^{\frac{1}{3}}\cdot xz
|
||||
2 1 4 = z x y ⟹ z = 2 1 4 ⋅ x y 2^{\frac{1}{4}}=\frac{z}{xy}\implies z=2^{\frac{1}{4}}\cdot xy
|
||||
Substituting these into each other, we find:
|
||||
1. x = 2 ⋅ y z x=\sqrt{2}\cdot yz
|
||||
2. y = 2 1 3 ⋅ x z y=2^{\frac{1}{3}}\cdot xz
|
||||
3. z = 2 1 4 ⋅ x y z=2^{\frac{1}{4}}\cdot xy
|
||||
By solving these equations step-by-step, we find:
|
||||
- z = 2 − 5 / 12 z=2^{-5/12}
|
||||
- x y = 2 − 2 / 3 xy=2^{-2/3}
|
||||
- y = 2 − 1 / 24 y=2^{-1/24}
|
||||
- x = 2 1 / 24 x=2^{1/24}
|
||||
However, upon verifying the third equation, we realize the system might not have a solution. Therefore, we use matrix inversion to solve the system:
|
||||
{ log 2 ( x ) − log 2 ( y ) − log 2 ( z ) = 1 2 log 2 ( y ) − log 2 ( x ) − log 2 ( z ) = 1 3 log 2 ( z ) − log 2 ( x ) − log 2 ( y ) = 1 4 \begin{cases}\log_{2}(x)-\log_{2}(y)-\log_{2}(z)=\frac{1}{2}\\
|
||||
\log_{2}(y)-\log_{2}(x)-\log_{2}(z)=\frac{1}{3}\\
|
||||
\log_{2}(z)-\log_{2}(x)-\log_{2}(y)=\frac{1}{4}\end{cases}
|
||||
Let A = log 2 ( x ) A=\log_{2}(x) , B = log 2 ( y ) B=\log_{2}(y) , C = log 2 ( z ) C=\log_{2}(z) . The system becomes:
|
||||
{ A − B − C = 1 2 B − A − C = 1 3 C − A − B = 1 4 \begin{cases}A-B-C=\frac{1}{2}\\
|
||||
B-A-C=\frac{1}{3}\\
|
||||
C-A-B=\frac{1}{4}\end{cases}
|
||||
We solve this system using Cramer’s rule and matrix inversion. The determinant of the matrix is -2, and the solutions are:
|
||||
A = 61 48 , B = 13 16 , C = − 1 6 A=\frac{61}{48},\quad B=\frac{13}{16},\quad C=-\frac{1}{6}
|
||||
Thus, we have:
|
||||
x = 2 61 / 48 , y = 2 13 / 16 , z = 2 − 1 / 6 x=2^{61/48},\quad y=2^{13/16},\quad z=2^{-1/6}
|
||||
We need to find | log 2 ( x 4 y 3 z 2 ) | \left|\log_{2}(x^{4}y^{3}z^{2})\right| :
|
||||
4 log 2 x + 3 log 2 y + 2 log 2 z = 4 ⋅ 61 48 + 3 ⋅ 13 16 + 2 ⋅ ( − 1 6 ) 4\log_{2}x+3\log_{2}y+2\log_{2}z=4\cdot\frac{61}{48}+3\cdot\frac{13}{16}+2\cdot\left(-\frac{1}{6}\right)
|
||||
Converting to a common denominator (48):
|
||||
4 ⋅ 61 48 = 244 48 , 3 ⋅ 13 16 = 117 48 , 2 ⋅ ( − 1 6 ) = − 16 48 4\cdot\frac{61}{48}=\frac{244}{48},\quad 3\cdot\frac{13}{16}=\frac{117}{48},\quad 2\cdot\left(-\frac{1}{6}\right)=-\frac{16}{48}
|
||||
Summing these:
|
||||
244 48 + 117 48 − 16 48 = 345 48 = 115 16 \frac{244}{48}+\frac{117}{48}-\frac{16}{48}=\frac{345}{48}=\frac{115}{16}
|
||||
The fraction 115 16 \frac{115}{16} is in simplest form, so m = 115 m=115 and n = 16 n=16 . Therefore, m + n = 115 + 16 = 131 m+n=115+16=131 .
|
||||
131 \boxed{131}
|
||||
Prompt:
|
||||
“|User |Real numbers x x and y y with x , y > 1 x,y>1 satisfy log x ( y x ) = log y ( x 4 y ) = 10 . \log_{x}(y^{x})=\log_{y}(x^{4y})=10.
|
||||
What is the value of x y xy ?
|
||||
Please reason step by step, and put your final answer within \boxed{} . |Assistant|”
|
||||
Our Result:
|
||||
⟨ \langle think ⟩ \rangle … ⟨ \langle /think ⟩ \rangle (15147 tokens)
|
||||
Given the equations:
|
||||
1. log x ( y x ) = 10 \log_{x}(y^{x})=10
|
||||
2. log y ( x 4 y ) = 10 \log_{y}(x^{4y})=10
|
||||
We start by applying the logarithm power rule to both equations.
|
||||
For the first equation:
|
||||
log x ( y x ) = x ⋅ log x ( y ) = 10 \log_{x}(y^{x})=x\cdot\log_{x}(y)=10
|
||||
Let z = log y ( x ) z=\log_{y}(x) . Then, log x ( y ) = 1 z \log_{x}(y)=\frac{1}{z} . Substituting this into the equation, we get:
|
||||
x ⋅ 1 z = 10 ⟹ x = 10 z x\cdot\frac{1}{z}=10\implies x=10z
|
||||
For the second equation:
|
||||
log y ( x 4 y ) = y ⋅ log y ( x 4 ) = 10 \log_{y}(x^{4y})=y\cdot\log_{y}(x^{4})=10
|
||||
Using the power rule, we get:
|
||||
y ⋅ 4 log y ( x ) = 10 ⟹ 4 y log y ( x ) = 10 y\cdot 4\log_{y}(x)=10\implies 4y\log_{y}(x)=10
|
||||
Since log y ( x ) = z \log_{y}(x)=z , this simplifies to:
|
||||
4 y z = 10 ⟹ y z = 5 2 4yz=10\implies yz=\frac{5}{2}
|
||||
We now have two expressions:
|
||||
1. x = 10 z x=10z
|
||||
2. y = 5 / 2 z y=\frac{5/2}{z}
|
||||
Substituting x = 10 z x=10z and y = 5 / 2 z y=\frac{5/2}{z} into the product x y xy , we get:
|
||||
x y = ( 10 z ) ⋅ ( 5 / 2 z ) = 10 z ⋅ 5 2 z = 10 ⋅ 5 2 = 25 xy=(10z)\cdot\left(\frac{5/2}{z}\right)=10z\cdot\frac{5}{2z}=10\cdot\frac{5}{2}=25
|
||||
Thus, the value of x y xy is 25 \boxed{25} .
|
||||
Baseline Result:
|
||||
⟨ \langle think ⟩ \rangle … ⟨ \langle /think ⟩ \rangle (16223 tokens)
|
||||
Given the problem where real numbers x x and y y with x , y > 1 x,y>1 satisfy the equations log x ( y x ) = 10 \log_{x}(y^{x})=10 and log y ( x 4 y ) = 10 \log_{y}(x^{4}y)=10 , we need to find the value of x y xy .
|
||||
First, we convert the logarithmic equations to their exponential forms using the change of base formula.
|
||||
1. log x ( y x ) = 10 \log_{x}(y^{x})=10
|
||||
x log x ( y ) = 10 x\log_{x}(y)=10
|
||||
Using the change of base formula, log x ( y ) = ln y ln x \log_{x}(y)=\frac{\ln y}{\ln x} . Since log y x = 9 4 \log_{y}x=\frac{9}{4} , we have log x y = 4 9 \log_{x}y=\frac{4}{9} . Thus,
|
||||
x ⋅ 4 9 = 10 ⟹ x = 10 ⋅ 9 4 = 90 4 = 45 2 x\cdot\frac{4}{9}=10\implies x=\frac{10\cdot 9}{4}=\frac{90}{4}=\frac{45}{2}
|
||||
2. log y ( x 4 y ) = 10 \log_{y}(x^{4}y)=10
|
||||
4 log y x + 1 = 10 ⟹ 4 log y x = 9 ⟹ log y x = 9 4 4\log_{y}x+1=10\implies 4\log_{y}x=9\implies\log_{y}x=\frac{9}{4}
|
||||
Again, using the change of base formula, log x y = 4 9 \log_{x}y=\frac{4}{9} .
|
||||
Next, we express x x and y y in terms of each other:
|
||||
x = 45 2 and y = ( 45 2 ) 4 / 9 x=\frac{45}{2}\quad\text{and}\quad y=\left(\frac{45}{2}\right)^{4/9}
|
||||
We then compute the product x y xy :
|
||||
x y = ( 45 2 ) ⋅ ( 45 2 ) 4 / 9 = ( 45 2 ) 1 + 4 / 9 = ( 45 2 ) 13 / 9 xy=\left(\frac{45}{2}\right)\cdot\left(\frac{45}{2}\right)^{4/9}=\left(\frac{45}{2}\right)^{1+4/9}=\left(\frac{45}{2}\right)^{13/9}
|
||||
However, upon re-evaluating and using another approach by setting x = 2 a x=2^{a} and y = 2 b y=2^{b} , we find:
|
||||
log x ( y x ) = 10 ⟹ x log x y log x y = 10 ⟹ x ⋅ 4 9 = 10 ⟹ x = 45 2 \log_{x}(y^{x})=10\implies\frac{x\log_{x}y}{\log_{x}y}=10\implies x\cdot\frac{4}{9}=10\implies x=\frac{45}{2}
|
||||
log y ( x 4 y ) = 10 ⟹ 4 log y x + 1 log y x = 10 ⟹ 4 ⋅ 9 4 + 1 = 10 ⟹ 9 + 1 = 10 consistent \log_{y}(x^{4}y)=10\implies\frac{4\log_{y}x+1}{\log_{y}x}=10\implies 4\cdot\frac{9}{4}+1=10\implies 9+1=10\quad\text{consistent}
|
||||
Finally, we find that the product x y xy simplifies to:
|
||||
x y = ( 45 2 ) xy=\left(\frac{45}{2}\right)
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
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,875 @@
|
||||
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.
|
||||
## References
|
||||
[1]
|
||||
Marker.
|
||||
URL https://github.com/datalab-to/marker .
|
||||
[2]
|
||||
Mathpix.
|
||||
URL https://mathpix.com/ .
|
||||
[3]
|
||||
Ocrflux, 2025.
|
||||
URL https://github.com/chatdoc-com/OCRFlux .
|
||||
[4]
|
||||
G. AI.
|
||||
Gemini 2.5-pro, 2025.
|
||||
URL https://gemini.google.com/ .
|
||||
[5]
|
||||
S. Bai, K. Chen, X. Liu, J. Wang, W. Ge, S. Song, K. Dang, P. Wang, S. Wang, J. Tang, H. Zhong, Y. Zhu, M. Yang, Z. Li, J. Wan, P. Wang, W. Ding, Z. Fu, Y. Xu, J. Ye, X. Zhang, T. Xie, Z. Cheng, H. Zhang, Z. Yang, H. Xu, and J. Lin.
|
||||
Qwen2.5-vl technical report.
|
||||
arXiv preprint arXiv:2502.13923 , 2025.
|
||||
[6]
|
||||
L. Blecher, G. Cucurull, T. Scialom, and R. Stojnic.
|
||||
Nougat: Neural optical understanding for academic documents.
|
||||
arXiv preprint arXiv:2308.13418 , 2023.
|
||||
[7]
|
||||
J. Chen, L. Kong, H. Wei, C. Liu, Z. Ge, L. Zhao, J. Sun, C. Han, and X. Zhang.
|
||||
Onechart: Purify the chart structural extraction via one auxiliary token.
|
||||
In Proceedings of the 32nd ACM International Conference on Multimedia , pages 147–155, 2024a.
|
||||
[8]
|
||||
Z. Chen, W. Wang, H. Tian, S. Ye, Z. Gao, E. Cui, W. Tong, K. Hu, J. Luo, Z. Ma, et al.
|
||||
How far are we to gpt-4v? closing the gap to commercial multimodal models with open-source suites.
|
||||
arXiv preprint arXiv:2404.16821 , 2024b.
|
||||
[9]
|
||||
C. Cui, T. Sun, M. Lin, T. Gao, Y. Zhang, J. Liu, X. Wang, Z. Zhang, C. Zhou, H. Liu, et al.
|
||||
Paddleocr 3.0 technical report.
|
||||
arXiv preprint arXiv:2507.05595 , 2025.
|
||||
[10]
|
||||
M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al.
|
||||
Patch n’ pack: Navit, a vision transformer for any aspect ratio and resolution.
|
||||
Advances in Neural Information Processing Systems , 36:3632–3656, 2023.
|
||||
[11]
|
||||
H. Feng, S. Wei, X. Fei, W. Shi, Y. Han, L. Liao, J. Lu, B. Wu, Q. Liu, C. Lin, et al.
|
||||
Dolphin: Document image parsing via heterogeneous anchor prompting.
|
||||
arXiv preprint arXiv:2505.14059 , 2025.
|
||||
[12]
|
||||
Y. Goyal, T. Khot, D. Summers-Stay, D. Batra, and D. Parikh.
|
||||
Making the v in vqa matter: Elevating the role of image understanding in visual question answering.
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 6904–6913, 2017.
|
||||
[13]
|
||||
J. Gu, X. Meng, G. Lu, L. Hou, N. Minzhe, X. Liang, L. Yao, R. Huang, W. Zhang, X. Jiang, et al.
|
||||
Wukong: A 100 million large-scale chinese cross-modal pre-training benchmark.
|
||||
Advances in Neural Information Processing Systems , 35:26418–26431, 2022.
|
||||
[14]
|
||||
High-flyer.
|
||||
HAI-LLM: Efficient and lightweight training tool for large models, 2023.
|
||||
URL https://www.high-flyer.cn/en/blog/hai-llm .
|
||||
[15]
|
||||
S. Iyer, X. V. Lin, R. Pasunuru, T. Mihaylov, D. Simig, P. Yu, K. Shuster, T. Wang, Q. Liu, P. S. Koura, et al.
|
||||
Opt-iml: Scaling language model instruction meta learning through the lens of generalization.
|
||||
arXiv preprint arXiv:2212.12017 , 2022.
|
||||
[16]
|
||||
S. Kazemzadeh, V. Ordonez, M. Matten, and T. Berg.
|
||||
Referitgame: Referring to objects in photographs of natural scenes.
|
||||
In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP) , pages 787–798, 2014.
|
||||
[17]
|
||||
A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, et al.
|
||||
Segment anything.
|
||||
arXiv preprint arXiv:2304.02643 , 2023.
|
||||
[18]
|
||||
Z. Li, Y. Liu, Q. Liu, Z. Ma, Z. Zhang, S. Zhang, Z. Guo, J. Zhang, X. Wang, and X. Bai.
|
||||
Monkeyocr: Document parsing with a structure-recognition-relation triplet paradigm.
|
||||
arXiv preprint arXiv:2506.05218 , 2025.
|
||||
[19]
|
||||
A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al.
|
||||
Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.
|
||||
arXiv preprint arXiv:2405.04434 , 2024a.
|
||||
[20]
|
||||
A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al.
|
||||
Deepseek-v3 technical report.
|
||||
arXiv preprint arXiv:2412.19437 , 2024b.
|
||||
[21]
|
||||
C. Liu, H. Wei, J. Chen, L. Kong, Z. Ge, Z. Zhu, L. Zhao, J. Sun, C. Han, and X. Zhang.
|
||||
Focus anywhere for fine-grained multi-page document understanding.
|
||||
arXiv preprint arXiv:2405.14295 , 2024c.
|
||||
[22]
|
||||
I. Loshchilov and F. Hutter.
|
||||
Sgdr: Stochastic gradient descent with warm restarts.
|
||||
arXiv preprint arXiv:1608.03983 , 2016.
|
||||
[23]
|
||||
I. Loshchilov and F. Hutter.
|
||||
Decoupled weight decay regularization.
|
||||
In ICLR , 2019.
|
||||
[24]
|
||||
A. Masry, D. X. Long, J. Q. Tan, S. Joty, and E. Hoque.
|
||||
Chartqa: A benchmark for question answering about charts with visual and logical reasoning.
|
||||
arXiv preprint arXiv:2203.10244 , 2022.
|
||||
[25]
|
||||
A. Nassar, A. Marafioti, M. Omenetti, M. Lysak, N. Livathinos, C. Auer, L. Morin, R. T. de Lima, Y. Kim, A. S. Gurbuz, et al.
|
||||
Smoldocling: An ultra-compact vision-language model for end-to-end multi-modal document conversion.
|
||||
arXiv preprint arXiv:2503.11576 , 2025.
|
||||
[26]
|
||||
OpenAI.
|
||||
Gpt-4 technical report, 2023.
|
||||
[27]
|
||||
L. Ouyang, Y. Qu, H. Zhou, J. Zhu, R. Zhang, Q. Lin, B. Wang, Z. Zhao, M. Jiang, X. Zhao, et al.
|
||||
Omnidocbench: Benchmarking diverse pdf document parsing with comprehensive annotations.
|
||||
In Proceedings of the Computer Vision and Pattern Recognition Conference , pages 24838–24848, 2025.
|
||||
[28]
|
||||
J. Poznanski, A. Rangapur, J. Borchardt, J. Dunkelberger, R. Huff, D. Lin, C. Wilhelm, K. Lo, and L. Soldaini.
|
||||
olmocr: Unlocking trillions of tokens in pdfs with vision language models.
|
||||
arXiv preprint arXiv:2502.18443 , 2025.
|
||||
[29]
|
||||
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al.
|
||||
Learning transferable visual models from natural language supervision.
|
||||
In International conference on machine learning , pages 8748–8763. PMLR, 2021.
|
||||
[30]
|
||||
Rednote.
|
||||
dots.ocr, 2025.
|
||||
URL https://github.com/rednote-hilab/dots.ocr .
|
||||
[31]
|
||||
C. Schuhmann, R. Vencu, R. Beaumont, R. Kaczmarczyk, C. Mullis, A. Katta, T. Coombes, J. Jitsev, and A. Komatsuzaki.
|
||||
Laion-400m: Open dataset of clip-filtered 400 million image-text pairs.
|
||||
arXiv preprint arXiv:2111.02114 , 2021.
|
||||
[32]
|
||||
A. Singh, V. Natarajan, M. Shah, Y. Jiang, X. Chen, D. Batra, D. Parikh, and M. Rohrbach.
|
||||
Towards vqa models that can read.
|
||||
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 8317–8326, 2019.
|
||||
[33]
|
||||
T. Sun, C. Cui, Y. Du, and Y. Liu.
|
||||
Pp-doclayout: A unified document layout detection model to accelerate large-scale data construction.
|
||||
arXiv preprint arXiv:2503.17213 , 2025.
|
||||
[34]
|
||||
B. Wang, C. Xu, X. Zhao, L. Ouyang, F. Wu, Z. Zhao, R. Xu, K. Liu, Y. Qu, F. Shang, et al.
|
||||
Mineru: An open-source solution for precise document content extraction.
|
||||
arXiv preprint arXiv:2409.18839 , 2024a.
|
||||
[35]
|
||||
P. Wang, S. Bai, S. Tan, S. Wang, Z. Fan, J. Bai, K. Chen, X. Liu, J. Wang, W. Ge, et al.
|
||||
Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution.
|
||||
arXiv preprint arXiv:2409.12191 , 2024b.
|
||||
[36]
|
||||
H. Wei, L. Kong, J. Chen, L. Zhao, Z. Ge, J. Yang, J. Sun, C. Han, and X. Zhang.
|
||||
Vary: Scaling up the vision vocabulary for large vision-language model.
|
||||
In European Conference on Computer Vision , pages 408–424. Springer, 2024a.
|
||||
[37]
|
||||
H. Wei, L. Kong, J. Chen, L. Zhao, Z. Ge, E. Yu, J. Sun, C. Han, and X. Zhang.
|
||||
Small language model meets with reinforced vision vocabulary.
|
||||
arXiv preprint arXiv:2401.12503 , 2024b.
|
||||
[38]
|
||||
H. Wei, C. Liu, J. Chen, J. Wang, L. Kong, Y. Xu, Z. Ge, L. Zhao, J. Sun, Y. Peng, et al.
|
||||
General ocr theory: Towards ocr-2.0 via a unified end-to-end model.
|
||||
arXiv preprint arXiv:2409.01704 , 2024c.
|
||||
[39]
|
||||
H. Wei, Y. Yin, Y. Li, J. Wang, L. Zhao, J. Sun, Z. Ge, X. Zhang, and D. Jiang.
|
||||
Slow perception: Let’s perceive geometric figures step-by-step.
|
||||
arXiv preprint arXiv:2412.20631 , 2024d.
|
||||
[40]
|
||||
Z. Wu, X. Chen, Z. Pan, X. Liu, W. Liu, D. Dai, H. Gao, Y. Ma, C. Wu, B. Wang, et al.
|
||||
Deepseek-vl2: Mixture-of-experts vision-language models for advanced multimodal understanding.
|
||||
arXiv preprint arXiv:2412.10302 , 2024.
|
||||
[41]
|
||||
W. Yu, Z. Yang, L. Li, J. Wang, K. Lin, Z. Liu, X. Wang, and L. Wang.
|
||||
Mm-vet: Evaluating large multimodal models for integrated capabilities.
|
||||
arXiv preprint arXiv:2308.02490 , 2023.
|
||||
[42]
|
||||
J. Zhu, W. Wang, Z. Chen, Z. Liu, S. Ye, L. Gu, H. Tian, Y. Duan, W. Su, J. Shao, et al.
|
||||
Internvl3: Exploring advanced training and test-time recipes for open-source multimodal models.
|
||||
arXiv preprint arXiv:2504.10479 , 2025.
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
@@ -0,0 +1,579 @@
|
||||
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.
|
||||
## References
|
||||
AlphaProof and teams (2024)
|
||||
AlphaProof and A. teams
|
||||
AI achieves silver-medal standard solving international mathematical olympiad problems .
|
||||
External Links: Link
|
||||
Cited by: §4 .
|
||||
Chen et al. (2025)
|
||||
L. Chen, J. Gu, L. Huang, W. Huang, Z. Jiang, A. Jie, X. Jin, X. Jin, C. Li, K. Ma, C. Ren, J. Shen, W. Shi, T. Sun, H. Sun, J. Wang, S. Wang, Z. Wang, C. Wei, S. Wei, Y. Wu, Y. Wu, Y. Xia, H. Xin, F. Yang, H. Ying, H. Yuan, Z. Yuan, T. Zhan, C. Zhang, Y. Zhang, G. Zhang, T. Zhao, J. Zhao, Y. Zhou, and T. H. Zhu
|
||||
Seed-prover: deep and broad reasoning for automated theorem proving .
|
||||
CoRR abs/2507.23726 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2507.23726
|
||||
Cited by: §4 .
|
||||
Chervonyi et al. (2025)
|
||||
Y. Chervonyi, T. H. Trinh, M. Olsák, X. Yang, H. Nguyen, M. Menegali, J. Jung, V. Verma, Q. V. Le, and T. Luong
|
||||
Gold-medalist performance in solving olympiad geometry with alphageometry2 .
|
||||
CoRR abs/2502.03544 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2502.03544
|
||||
Cited by: §4 .
|
||||
de Moura et al. (2015)
|
||||
L. M. de Moura, S. Kong, J. Avigad, F. van Doorn, and J. von Raumer
|
||||
The lean theorem prover (system description) .
|
||||
In Automated Deduction - CADE-25 - 25th International Conference on
|
||||
Automated Deduction, Berlin, Germany, August 1-7, 2015, Proceedings , A. P. Felty and A. Middeldorp (Eds.) ,
|
||||
Lecture Notes in Computer Science , Vol. 9195 , pp. 378–388 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §4 .
|
||||
DeepMind (2025)
|
||||
G. DeepMind
|
||||
Gemini 2.5 pro .
|
||||
External Links: Link
|
||||
Cited by: §3.3.1 .
|
||||
DeepSeek-AI (2025)
|
||||
DeepSeek-AI
|
||||
DeepSeek-v3.2-exp: boosting long-context efficiency with deepseek sparse attention .
|
||||
Cited by: §1 .
|
||||
Dekoninck et al. (2025)
|
||||
J. Dekoninck, I. Petrov, K. Minchev, M. Balunovic, M. T. Vechev, M. Marinov, M. Drencheva, L. Konova, M. Shumanov, K. Tsvetkov, N. Drenchev, L. Todorov, K. Nikolova, N. Georgiev, V. Kalinkova, and M. Ismoldayev
|
||||
The open proof corpus: A large-scale study of llm-generated mathematical proofs .
|
||||
CoRR abs/2506.21621 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2506.21621
|
||||
Cited by: §4 .
|
||||
Guo et al. (2025)
|
||||
D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, Q. Zhu, R. Xu, R. Zhang, S. Ma, X. Bi, X. Zhang, X. Yu, Y. Wu, Z. F. Wu, Z. Gou, Z. Shao, Z. Li, Z. Gao, A. Liu, B. Xue, B. Wang, B. Wu, B. Feng, C. Lu, C. Zhao, C. Deng, C. Ruan, D. Dai, D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang, H. Xu, H. Ding, H. Gao, H. Qu, H. Li, J. Guo, J. Li, J. Chen, J. Yuan, J. Tu, J. Qiu, J. Li, J. L. Cai, J. Ni, J. Liang, J. Chen, K. Dong, K. Hu, K. You, K. Gao, K. Guan, K. Huang, K. Yu, L. Wang, L. Zhang, L. Zhao, L. Wang, L. Zhang, L. Xu, L. Xia, M. Zhang, M. Zhang, M. Tang, M. Zhou, M. Li, M. Wang, M. Li, N. Tian, P. Huang, P. Zhang, Q. Wang, Q. Chen, Q. Du, R. Ge, R. Zhang, R. Pan, R. Wang, R. J. Chen, R. L. Jin, R. Chen, S. Lu, S. Zhou, S. Chen, S. Ye, S. Wang, S. Yu, S. Zhou, S. Pan, S. S. Li, S. Zhou, S. Wu, T. Yun, T. Pei, T. Sun, T. Wang, W. Zeng, W. Liu, W. Liang, W. Gao, W. Yu, W. Zhang, W. L. Xiao, W. An, X. Liu, X. Wang, X. Chen, X. Nie, X. Cheng, X. Liu, X. Xie, X. Liu, X. Yang, X. Li, X. Su, X. Lin, X. Q. Li, X. Jin, X. Shen, X. Chen, X. Sun, X. Wang, X. Song, X. Zhou, X. Wang, X. Shan, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. Zhang, Y. Xu, Y. Li, Y. Zhao, Y. Sun, Y. Wang, Y. Yu, Y. Zhang, Y. Shi, Y. Xiong, Y. He, Y. Piao, Y. Wang, Y. Tan, Y. Ma, Y. Liu, Y. Guo, Y. Ou, Y. Wang, Y. Gong, Y. Zou, Y. He, Y. Xiong, Y. Luo, Y. You, Y. Liu, Y. Zhou, Y. X. Zhu, Y. Huang, Y. Li, Y. Zheng, Y. Zhu, Y. Ma, Y. Tang, Y. Zha, Y. Yan, Z. Z. Ren, Z. Ren, Z. Sha, Z. Fu, Z. Xu, Z. Xie, Z. Zhang, Z. Hao, Z. Ma, Z. Yan, Z. Wu, Z. Gu, Z. Zhu, Z. Liu, Z. Li, Z. Xie, Z. Song, Z. Pan, Z. Huang, Z. Xu, Z. Zhang, and Z. Zhang
|
||||
DeepSeek-r1 incentivizes reasoning in llms through reinforcement learning .
|
||||
Nat. 645 ( 8081 ), pp. 633–638 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §1 ,
|
||||
§4 .
|
||||
Huang and Yang (2025)
|
||||
Y. Huang and L. F. Yang
|
||||
Gemini 2.5 pro capable of winning gold at IMO 2025 .
|
||||
CoRR abs/2507.15855 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2507.15855
|
||||
Cited by: §4 .
|
||||
Jiang et al. (2023)
|
||||
A. Q. Jiang, S. Welleck, J. P. Zhou, T. Lacroix, J. Liu, W. Li, M. Jamnik, G. Lample, and Y. Wu
|
||||
Draft, sketch, and prove: guiding formal theorem provers with informal proofs .
|
||||
In The Eleventh International Conference on Learning Representations,
|
||||
ICLR 2023, Kigali, Rwanda, May 1-5, 2023 ,
|
||||
External Links: Link
|
||||
Cited by: §4 .
|
||||
Luong et al. (2025)
|
||||
T. Luong, D. Hwang, H. H. Nguyen, G. Ghiasi, Y. Chervonyi, I. Seo, J. Kim, G. Bingham, J. Lee, S. Mishra, A. Zhai, C. H. Hu, H. Michalewski, J. Kim, J. Ahn, J. Bae, X. Song, T. H. Trinh, Q. V. Le, and J. Jung
|
||||
Towards robust mathematical reasoning .
|
||||
In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing ,
|
||||
External Links: Link
|
||||
Cited by: Figure 3 ,
|
||||
5th item ,
|
||||
§4 .
|
||||
Luong and Lockhart (2025)
|
||||
T. Luong and E. Lockhart
|
||||
Advanced version of gemini with deep think officially achieves gold-medal standard at the international mathematical olympiad .
|
||||
External Links: Link
|
||||
Cited by: 5th item ,
|
||||
§4 .
|
||||
OpenAI (2024)
|
||||
OpenAI
|
||||
Learning to reason with llms .
|
||||
External Links: Link
|
||||
Cited by: §4 .
|
||||
OpenAI (2025)
|
||||
OpenAI
|
||||
Introducing gpt-5 .
|
||||
External Links: Link
|
||||
Cited by: §3.3.1 .
|
||||
Paulson (1994)
|
||||
L. C. Paulson
|
||||
Isabelle - A generic theorem prover (with a contribution by t. nipkow) .
|
||||
Lecture Notes in Computer Science , Vol. 828 , Springer .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
ISBN 3-540-58244-4
|
||||
Cited by: §4 .
|
||||
Ren et al. (2025)
|
||||
Z. Z. Ren, Z. Shao, J. Song, H. Xin, H. Wang, W. Zhao, L. Zhang, Z. Fu, Q. Zhu, D. Yang, Z. F. Wu, Z. Gou, S. Ma, H. Tang, Y. Liu, W. Gao, D. Guo, and C. Ruan
|
||||
DeepSeek-prover-v2: advancing formal mathematical reasoning via reinforcement learning for subgoal decomposition .
|
||||
CoRR abs/2504.21801 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2504.21801
|
||||
Cited by: §4 .
|
||||
Shao et al. (2024)
|
||||
Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, M. Zhang, Y. K. Li, Y. Wu, and D. Guo
|
||||
DeepSeekMath: pushing the limits of mathematical reasoning in open language models .
|
||||
CoRR abs/2402.03300 .
|
||||
External Links: Link ,
|
||||
Document ,
|
||||
2402.03300
|
||||
Cited by: §3.1 .
|
||||
Trinh et al. (2024)
|
||||
T. H. Trinh, Y. Wu, Q. V. Le, H. He, and T. Luong
|
||||
Solving olympiad geometry without human demonstrations .
|
||||
Nat. 625 ( 7995 ), pp. 476–482 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §4 .
|
||||
## Appendix A Prompt Templates
|
||||
## A.1 Proof Generation Prompt
|
||||
⬇
|
||||
Your task is to solve a given problem . The problem may ask you to prove a statement , or ask for an answer . If finding an answer is required , you should come up with the answer , and your final solution should also be a rigorous proof of that answer being valid .
|
||||
Your final solution to the problem should be exceptionally comprehensive and easy - to - follow , which will be rated according to the following evaluation instruction :
|
||||
‘‘‘ txt
|
||||
Here is the instruction to evaluate the quality of a solution to a problem . The problem may ask for a proof of statement , or ask for an answer . If finding an answer is required , the solution should present the answer , and it should also be a rigorous proof of that answer being valid .
|
||||
Please evaluate the solution and score it according to the following criteria :
|
||||
- If the solution is completely correct , with all steps executed properly and clearly demonstrated , then the score is 1
|
||||
- If the solution is generally correct , but with some details omitted or minor errors , then the score is 0.5
|
||||
- If the solution does not actually address the required problem , contains fatal errors , or has severe omissions , then the score is 0
|
||||
Additionally , referencing anything from any paper does not save the need to prove the reference . It ’ s okay IF AND ONLY IF the solution also presents a valid proof of the reference argument ( s ); otherwise , if the solution omits the proof or if the proof provided is not completely correct , the solution should be scored according to the criteria above , and definitely not with a score of 1
|
||||
‘‘‘
|
||||
In fact , you already have the ability to rate your solution yourself , so you are expected to reason carefully about how to solve a given problem , evaluate your method according to the instruction , and refine your solution by fixing issues identified until you can make no further progress .
|
||||
In your final response , you should present a detailed solution to the problem followed by your evaluation of that solution .
|
||||
- To give a good final response , you should try your best to locate potential issues in your own ( partial ) solution according to the evaluation instruction above , and fix them as many as you can .
|
||||
- A good final response should just faithfully present your progress , including the best solution you can give , as well as a faithful evaluation of that solution .
|
||||
- Only when you fail to locate any issues in your solution should you score it with 1.
|
||||
- If you do notice some issues in your solution but fail to resolve them with your best efforts , it ’ s totally ok to faithfully present the issues in your final response .
|
||||
- The worst final response would provide a wrong solution but lie that it ’ s correct or claim that it ’ s correct without careful error checking . A better version should faithfully identify errors in the solution . Remember ! You CAN ’ T cheat ! If you cheat , we will know , and you will be penalized !
|
||||
Your final response should be in the following format :
|
||||
## Solution // Your final solution should start with this exact same markdown title
|
||||
... // Your final solution to the problem here . You should try your best to optimize the quality of your solution according to the evaluation instruction above before finalizing it here .
|
||||
## Self Evaluation // Your evaluation of your own solution above should start with this exact same markdown title
|
||||
Here is my evaluation of the solution : // Your analysis should start with this exact same phrase
|
||||
... // Your evaluation here . You are required to present in detail the key steps of the solution or the steps for which you had doubts regarding their correctness , and explicitly analyze whether each step is accurate : for correct steps , explain why you initially doubted their correctness and why they are indeed correct ; for erroneous steps , explain the reason for the error and the impact of that error on the solution . You should analyze your solution faithfully . E . g ., if there are issues in your final solution , you should point it out .
|
||||
Based on my evaluation , the final overall score should be :
|
||||
\\ boxed {{...}} // where ... should be the final overall score (0, 0.5, or 1, and nothing else ) based on the evaluation instruction above . You should reach this score ONLY AFTER careful RE - examination of your own solution above
|
||||
---
|
||||
Here is your task input :
|
||||
## Problem
|
||||
{ question }
|
||||
## A.2 Proof Verification Prompt
|
||||
⬇
|
||||
## Instruction
|
||||
Your task is to evaluate the quality of a solution to a problem . The problem may ask for a proof of statement , or ask for an answer . If finding an answer is required , the solution should present the answer , and it should also be a rigorous proof of that answer being valid .
|
||||
Please evaluate the solution and score it according to the following criteria :
|
||||
- If the solution is completely correct , with all steps executed properly and clearly demonstrated , then the score is 1
|
||||
- If the solution is generally correct , but with some details omitted or minor errors , then the score is 0.5
|
||||
- If the solution does not actually address the required problem , contains fatal errors , or has severe omissions , then the score is 0
|
||||
- Additionally , referencing anything from any paper does not save the need to prove the reference . It ’ s okay IF AND ONLY IF the solution also presents a valid proof of the reference argument ( s ); otherwise , if the solution omits the proof or if the proof provided is not completely correct , the solution should be scored according to the criteria above , and definitely not with a score of 1
|
||||
Please carefully reason out and analyze the quality of the solution below , and in your final response present a detailed evaluation of the solution ’ s quality followed by your score . Therefore , your response should be in the following format :
|
||||
Here is my evaluation of the solution :
|
||||
... // Your evaluation here . You are required to present in detail the key steps of the solution or the steps for which you had doubts regarding their correctness , and explicitly analyze whether each step is accurate : for correct steps , explain why you initially doubted their correctness and why they are indeed correct ; for erroneous steps , explain the reason for the error and the impact of that error on the solution .
|
||||
Based on my evaluation , the final overall score should be :
|
||||
\\ boxed {{...}} // where ... should be the final overall score (0, 0.5, or 1, and nothing else ) based on the above criteria
|
||||
---
|
||||
Here is your task input :
|
||||
## Problem
|
||||
{ question }
|
||||
## Solution
|
||||
{ proof }
|
||||
## A.3 Meta-Verification Prompt
|
||||
⬇
|
||||
You are given a " problem ", " solution ", and " solution evaluation ", and you need to assess the whether this " solution evaluation " is reasonable .
|
||||
First , " solution evaluation " is generated to evaluate the quality of the " solution ", by prompting a verifier with the rules below ( these are not your rules ):
|
||||
‘‘‘
|
||||
Please evaluate the solution and score it according to the following criteria :
|
||||
- If the solution is completely correct , with all steps executed properly and clearly demonstrated , then the score is 1
|
||||
- If the solution is generally correct , but with some details omitted or minor errors , then the score is 0.5
|
||||
- If the solution does not actually address the required problem , contains fatal errors , or has severe omissions , then the score is 0
|
||||
Additionally , referencing anything from any paper does not save the need to prove the reference . It ’ s okay IF AND ONLY IF the solution also presents a valid proof of the reference argument ( s ); otherwise , if the solution omits the proof or if the proof provided is not completely correct , the solution should be scored according to the criteria above , and definitely not with a score of 1
|
||||
‘‘‘
|
||||
Next , I will introduce the rules for you to analyze the quality of the " solution evaluation ":
|
||||
1. Your task is to analyze the " solution evaluation ". You do not need to solve the " problem ", nor do you need to strictly assess whether the " solution " is accurate . Your only task is to strictly follow the rules below to evaluate whether the " solution evaluation " is reasonable .
|
||||
2. You need to analyze the content of the " solution evaluation " from three aspects :
|
||||
Step Restatement : In the " solution evaluation ", certain behaviors of the " solution " may be restated . You need to return to the original text of the " solution " and check whether the " solution " actually has these behaviors mentioned in the " solution evaluation ".
|
||||
Defect Analysis : " solution evaluation " may point out errors or defects in the " solution ". You need to carefully analyze whether the mentioned errors and defects are indeed valid .
|
||||
Expression Analysis : Whether the " solution evaluation "’ s expressions are accurate .
|
||||
Score Analysis : Whether the final score given by the " solution evaluation " matches the defects it found . You need to analyze according to the scoring rules given above .
|
||||
3. The most important part is ** defect analysis **: In this part , your core task is to check whether the errors or defects of the " solution " pointed out in the " solution evaluation " are reasonable . In other words , any positive components about the " solution " in the " solution evaluation ", regardless of whether they are reasonable , are not within your evaluation scope .
|
||||
- For example : If the " solution evaluation " says that a certain conclusion in the " solution " is correct , but actually this conclusion is incorrect , then you do not need to care about this point . All parts that the " solution evaluation " considers correct do not belong to your evaluation scope .
|
||||
- Specifically : If the " solution evaluation " believes that the " solution " is completely accurate and has not found any errors or defects , then regardless of whether the " solution " itself is actually accurate , even if there are obvious errors , you should still consider its analysis of errors to be reasonable .
|
||||
** Importantly **, for defects found by the " solution evaluation ", you need to analyze two points simultaneously :
|
||||
- whether this defect actually exists
|
||||
- whether the " solution evaluation "’ s analysis of this defect is accurate
|
||||
These two aspects constitute the analysis of defects .
|
||||
4. About ** expression analysis **, if there are certain expression errors in the " solution evaluation ", even minor errors in details , you need to identify them . However , please note that identifying incorrect steps in the " solution " as correct steps does not constitute an ** expression error **.
|
||||
In practice , expression errors include but are not limited to :
|
||||
- If the " solution evaluation " identifies some reasoning step ( s ) in the " solution " as incorrect , then it cannot further indicate that subsequent conclusion ( s ) depending on those reasoning step ( s ) are wrong , but can only indicate that subsequent conclusion ( s ) are " not rigorously demonstrated ."
|
||||
- Typos and calculation errors made by " solution evaluation "
|
||||
- Inaccurate restatement of content from " solution "
|
||||
5. Finally , you need to present your analysis of the " solution evaluation " in your output and also rate its quality based on the rules below :
|
||||
First , if there is at least one unreasonable defect among the defects found by the " solution evaluation ", then you only need to do ** defect analysis **:
|
||||
- If all defects found by the " solution evaluation " are unreasonable , then you should rate it with \(0\)
|
||||
- If some defects found by the " solution evaluation " are reasonable and some are unreasonable , then your rating should be \(0.5\)
|
||||
Next , if the " solution evaluation " points out no errors or defects , or all defects found by the evaluation are reasonable , then you should do the following things :
|
||||
- Analyze whether " expression errors " exist in the " solution evaluation " (** expression analysis **) or whether " solution evaluation " gives a wrong score according to the rules for " solution evaluation " (** score analysis **). If yes , you should rate the " solution evaluation " with \(0.5\); if no , your rating should be \(1\)
|
||||
Your output should follow the format below :
|
||||
Here is my analysis of the " solution evaluation ":
|
||||
... // Your analysis here .
|
||||
Based on my analysis , I will rate the " solution evaluation " as :
|
||||
\\ boxed {{...}} // where ... should be a numerical rating of the " solution evaluation " (0, 0.5, or 1, and nothing else ) based on the criteria above .
|
||||
---
|
||||
Here is your task input :
|
||||
## Problem
|
||||
{ question }
|
||||
## Solution
|
||||
{ proof }
|
||||
## Solution Evaluation
|
||||
{ proof analysis }
|
||||
## A.4 Proof Refinement Prompt
|
||||
⬇
|
||||
{ proof_generation_prompt }
|
||||
## Candidate Solution ( s ) to Refine
|
||||
Here are some solution sample ( s ) along with their correctness evaluation ( s ). You should provide a better solution by solving issues mentioned in the evaluation ( s ), or by re - using promising ideas mentioned in the solution sample ( s ), or by doing both .
|
||||
{ proof }
|
||||
{ proof analyses }
|
||||
## Final Instruction
|
||||
Your final response should follow the format above , including a ‘## Solution ‘ section followed by a ‘## Self Evaluation ‘ section
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,889 @@
|
||||
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.
|
||||
## References
|
||||
Ainslie et al. (2023)
|
||||
J. Ainslie, J. Lee-Thorp, M. De Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai
|
||||
Gqa: training generalized multi-query transformer models from multi-head checkpoints .
|
||||
arXiv preprint arXiv:2305.13245 .
|
||||
Cited by: §2.1 .
|
||||
Bisk et al. (2020)
|
||||
Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi
|
||||
PIQA: reasoning about physical commonsense in natural language .
|
||||
In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI
|
||||
2020, The Thirty-Second Innovative Applications of Artificial Intelligence
|
||||
Conference, IAAI 2020, The Tenth AAAI Symposium on Educational
|
||||
Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA,
|
||||
February 7-12, 2020 ,
|
||||
pp. 7432–7439 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §5.2 .
|
||||
Brown et al. (2020)
|
||||
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al.
|
||||
Language models are few-shot learners .
|
||||
Advances in neural information processing systems 33 , pp. 1877–1901 .
|
||||
Cited by: §1 .
|
||||
Chai et al. (2020)
|
||||
Y. Chai, S. Jin, and X. Hou
|
||||
Highway transformer: self-gating enhanced self-attentive networks .
|
||||
In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics , D. Jurafsky, J. Chai, N. Schluter, and J. Tetreault (Eds.) ,
|
||||
Online , pp. 6887–6900 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §2.2 .
|
||||
Chollet (2017)
|
||||
F. Chollet
|
||||
Xception: deep learning with depthwise separable convolutions .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 1251–1258 .
|
||||
Cited by: §2.1 .
|
||||
Cobbe et al. (2021)
|
||||
K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al.
|
||||
Training verifiers to solve math word problems .
|
||||
arXiv preprint arXiv:2110.14168 .
|
||||
Cited by: §5.2 .
|
||||
Dao et al. (2022)
|
||||
T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré
|
||||
FlashAttention: fast and memory-efficient exact attention with IO-awareness .
|
||||
In Advances in Neural Information Processing Systems (NeurIPS) ,
|
||||
Cited by: §3.2 .
|
||||
Dua et al. (2019)
|
||||
D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner
|
||||
DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs .
|
||||
In Proceedings of the 2019 Conference of the North American Chapter of
|
||||
the Association for Computational Linguistics: Human Language Technologies,
|
||||
NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long
|
||||
and Short Papers) , J. Burstein, C. Doran, and T. Solorio (Eds.) ,
|
||||
pp. 2368–2378 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §5.2 .
|
||||
Fang et al. (2023)
|
||||
Y. Fang, Y. CAI, J. Chen, J. Zhao, G. Tian, and G. Li
|
||||
Cross-layer retrospective retrieving via layer attention .
|
||||
In The Eleventh International Conference on Learning Representations ,
|
||||
External Links: Link
|
||||
Cited by: §2.2 .
|
||||
Fedus et al. (2022)
|
||||
W. Fedus, B. Zoph, and N. Shazeer
|
||||
Switch transformers: scaling to trillion parameter models with simple and efficient sparsity .
|
||||
Journal of Machine Learning Research 23 ( 120 ), pp. 1–39 .
|
||||
Cited by: §2.1 .
|
||||
He et al. (2016a)
|
||||
K. He, X. Zhang, S. Ren, and J. Sun
|
||||
Deep residual learning for image recognition .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 770–778 .
|
||||
Cited by: §1 ,
|
||||
§2.2 .
|
||||
He et al. (2016b)
|
||||
K. He, X. Zhang, S. Ren, and J. Sun
|
||||
Identity mappings in deep residual networks .
|
||||
In European conference on computer vision ,
|
||||
pp. 630–645 .
|
||||
Cited by: §1 ,
|
||||
§4.1 .
|
||||
Heddes et al. (2025)
|
||||
M. Heddes, A. Javanmard, K. Axiotis, G. Fu, M. Bateni, and V. Mirrokni
|
||||
DeepCrossAttention: supercharging transformer residual connections .
|
||||
In Forty-second International Conference on Machine Learning ,
|
||||
External Links: Link
|
||||
Cited by: §2.2 .
|
||||
Hendrycks et al. (2020)
|
||||
D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt
|
||||
Measuring massive multitask language understanding .
|
||||
arXiv preprint arXiv:2009.03300 .
|
||||
Cited by: §5.2 .
|
||||
Hendrycks et al. (2021)
|
||||
D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt
|
||||
Measuring mathematical problem solving with the math dataset .
|
||||
arXiv preprint arXiv:2103.03874 .
|
||||
Cited by: §5.2 .
|
||||
Hoffmann et al. (2022)
|
||||
J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. Rae, and L. Sifre
|
||||
An empirical analysis of compute-optimal large language model training .
|
||||
In Advances in Neural Information Processing Systems , S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (Eds.) ,
|
||||
Vol. 35 , pp. 30016–30030 .
|
||||
External Links: Link
|
||||
Cited by: §3 .
|
||||
Huang et al. (2017)
|
||||
G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger
|
||||
Densely connected convolutional networks .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 4700–4708 .
|
||||
Cited by: §2.2 .
|
||||
Joshi et al. (2017)
|
||||
M. Joshi, E. Choi, D. Weld, and L. Zettlemoyer
|
||||
TriviaQA: a large scale distantly supervised challenge dataset for reading comprehension .
|
||||
In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , R. Barzilay and M. Kan (Eds.) ,
|
||||
Vancouver, Canada , pp. 1601–1611 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §5.2 .
|
||||
Larsson et al. (2016)
|
||||
G. Larsson, M. Maire, and G. Shakhnarovich
|
||||
Fractalnet: ultra-deep neural networks without residuals .
|
||||
arXiv preprint arXiv:1605.07648 .
|
||||
Cited by: §2.2 .
|
||||
Lepikhin et al. (2020)
|
||||
D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen
|
||||
Gshard: scaling giant models with conditional computation and automatic sharding .
|
||||
arXiv preprint arXiv:2006.16668 .
|
||||
Cited by: §2.1 .
|
||||
Liu et al. (2024a)
|
||||
A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al.
|
||||
Deepseek-v2: a strong, economical, and efficient mixture-of-experts language model .
|
||||
arXiv preprint arXiv:2405.04434 .
|
||||
Cited by: Table 5 ,
|
||||
§2.1 .
|
||||
Liu et al. (2024b)
|
||||
A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al.
|
||||
Deepseek-v3 technical report .
|
||||
arXiv preprint arXiv:2412.19437 .
|
||||
Cited by: Table 5 ,
|
||||
Table 5 ,
|
||||
§1 ,
|
||||
§1 ,
|
||||
§4.3.3 ,
|
||||
§5.1 .
|
||||
Loshchilov and Hutter (2017)
|
||||
I. Loshchilov and F. Hutter
|
||||
Decoupled weight decay regularization .
|
||||
arXiv preprint arXiv:1711.05101 .
|
||||
Cited by: Table 5 .
|
||||
Mak and Flanigan (2025)
|
||||
B. Mak and J. Flanigan
|
||||
Residual matrix transformers: scaling the size of the residual stream .
|
||||
arXiv preprint arXiv:2506.22696 .
|
||||
Cited by: §2.2 .
|
||||
Menghani et al. (2025)
|
||||
G. Menghani, R. Kumar, and S. Kumar
|
||||
LAurel: learned augmented residual layer .
|
||||
In Forty-second International Conference on Machine Learning ,
|
||||
External Links: Link
|
||||
Cited by: §2.2 .
|
||||
Pagliardini et al. (2024)
|
||||
M. Pagliardini, A. Mohtashami, F. Fleuret, and M. Jaggi
|
||||
DenseFormer: enhancing information flow in transformers via depth weighted averaging .
|
||||
In The Thirty-eighth Annual Conference on Neural Information Processing Systems ,
|
||||
External Links: Link
|
||||
Cited by: §2.2 .
|
||||
Qi et al. (2024)
|
||||
P. Qi, X. Wan, G. Huang, and M. Lin
|
||||
Zero bubble (almost) pipeline parallelism .
|
||||
In The Twelfth International Conference on Learning Representations ,
|
||||
External Links: Link
|
||||
Cited by: §3.2 .
|
||||
Shazeer et al. (2017)
|
||||
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean
|
||||
Outrageously large neural networks: the sparsely-gated mixture-of-experts layer .
|
||||
arXiv preprint arXiv:1701.06538 .
|
||||
Cited by: §2.1 .
|
||||
Shazeer (2019)
|
||||
N. Shazeer
|
||||
Fast transformer decoding: one write-head is all you need .
|
||||
arXiv preprint arXiv:1911.02150 .
|
||||
Cited by: §2.1 .
|
||||
Sinkhorn and Knopp (1967)
|
||||
R. Sinkhorn and P. Knopp
|
||||
Concerning nonnegative matrices and doubly stochastic matrices .
|
||||
Pacific Journal of Mathematics 21 ( 2 ), pp. 343–348 .
|
||||
Cited by: §1 .
|
||||
Srivastava et al. (2015)
|
||||
R. K. Srivastava, K. Greff, and J. Schmidhuber
|
||||
Training very deep networks .
|
||||
In Advances in Neural Information Processing Systems , C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (Eds.) ,
|
||||
Vol. 28 , pp. .
|
||||
External Links: Link
|
||||
Cited by: §2.2 .
|
||||
Su et al. (2024)
|
||||
J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu
|
||||
Roformer: enhanced transformer with rotary position embedding .
|
||||
Neurocomputing 568 , pp. 127063 .
|
||||
Cited by: Table 5 .
|
||||
Suzgun et al. (2022)
|
||||
M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. V. Le, E. H. Chi, D. Zhou, et al.
|
||||
Challenging big-bench tasks and whether chain-of-thought can solve them .
|
||||
arXiv preprint arXiv:2210.09261 .
|
||||
Cited by: §5.2 .
|
||||
Touvron et al. (2023)
|
||||
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al.
|
||||
Llama: open and efficient foundation language models .
|
||||
arXiv preprint arXiv:2302.13971 .
|
||||
Cited by: §1 .
|
||||
Vaswani et al. (2017)
|
||||
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin
|
||||
Attention is all you need .
|
||||
Advances in neural information processing systems 30 .
|
||||
Cited by: §1 ,
|
||||
§2.1 ,
|
||||
§3.2 .
|
||||
Wang et al. (2024)
|
||||
L. Wang, H. Gao, C. Zhao, X. Sun, and D. Dai
|
||||
Auxiliary-loss-free load balancing strategy for mixture-of-experts .
|
||||
arXiv preprint arXiv:2408.15664 .
|
||||
Cited by: Table 5 .
|
||||
Wang et al. (2025)
|
||||
L. Wang, Y. Cheng, Y. Shi, Z. Tang, Z. Mo, W. Xie, L. Ma, Y. Xia, J. Xue, F. Yang, et al.
|
||||
TileLang: a composable tiled programming model for ai systems .
|
||||
arXiv preprint arXiv:2504.17577 .
|
||||
Cited by: §1 ,
|
||||
§4.3.1 .
|
||||
Xiao et al. (2025)
|
||||
D. Xiao, Q. Meng, S. Li, and X. Yuan
|
||||
Muddformer: breaking residual bottlenecks in transformers via multiway dynamic dense connections .
|
||||
arXiv preprint arXiv:2502.12170 .
|
||||
Cited by: §2.2 .
|
||||
Xie et al. (2017)
|
||||
S. Xie, R. Girshick, P. Dollár, Z. Tu, and K. He
|
||||
Aggregated residual transformations for deep neural networks .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 1492–1500 .
|
||||
Cited by: §2.1 .
|
||||
Xie et al. (2023)
|
||||
S. Xie, H. Zhang, J. Guo, X. Tan, J. Bian, H. H. Awadalla, A. Menezes, T. Qin, and R. Yan
|
||||
ResiDual: transformer with dual residual connections .
|
||||
External Links: 2304.14802 ,
|
||||
Link
|
||||
Cited by: §2.2 .
|
||||
Yu et al. (2018)
|
||||
F. Yu, D. Wang, E. Shelhamer, and T. Darrell
|
||||
Deep layer aggregation .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 2403–2412 .
|
||||
Cited by: §2.2 .
|
||||
Zellers et al. (2019)
|
||||
R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi
|
||||
HellaSwag: can a machine really finish your sentence? .
|
||||
In Proceedings of the 57th Conference of the Association for Computational
|
||||
Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019,
|
||||
Volume 1: Long Papers , A. Korhonen, D. R. Traum, and L. Màrquez (Eds.) ,
|
||||
pp. 4791–4800 .
|
||||
External Links: Link ,
|
||||
Document
|
||||
Cited by: §5.2 .
|
||||
Zhang and Sennrich (2019)
|
||||
B. Zhang and R. Sennrich
|
||||
Root mean square layer normalization .
|
||||
Advances in neural information processing systems 32 .
|
||||
Cited by: Table 5 ,
|
||||
§3 .
|
||||
Zhu et al. (2024)
|
||||
D. Zhu, H. Huang, Z. Huang, Y. Zeng, Y. Mao, B. Wu, Q. Min, and X. Zhou
|
||||
Hyper-connections .
|
||||
arXiv preprint arXiv:2409.19606 .
|
||||
Cited by: §1 ,
|
||||
§2.2 .
|
||||
## Appendix A Appendix
|
||||
## A.1 Detailed Model Specifications and Hyper-parameters.
|
||||
Table 5 :
|
||||
Detailed Model Specifications and Hyper-parameters. This table presents the architectural configurations for the 3B, 9B, and 27B models based on the DeepSeek-V3 ( Liu et al., 2024b ) architecture. It outlines the specific hyper-parameters for m HC and HC, including the residual stream expansion and Sinkhorn-Knopp settings, alongside the optimization and training protocols used in the experiments.
|
||||
Attribute
|
||||
3B
|
||||
9B
|
||||
27B
|
||||
3B
|
||||
1T Tokens
|
||||
Vocab Params
|
||||
331M
|
||||
496M
|
||||
662M
|
||||
331M
|
||||
Active Params
|
||||
612M
|
||||
1.66B
|
||||
4.14B
|
||||
612M
|
||||
Total Params
|
||||
2.97B
|
||||
9.18B
|
||||
27.0B
|
||||
2.97B
|
||||
Layers
|
||||
12
|
||||
18
|
||||
30
|
||||
12
|
||||
Leading Dense Layers
|
||||
1
|
||||
1
|
||||
Routed Experts
|
||||
64
|
||||
64
|
||||
72
|
||||
64
|
||||
Active Experts
|
||||
6
|
||||
6
|
||||
Shared Experts
|
||||
2
|
||||
2
|
||||
Dimension
|
||||
1280
|
||||
1920
|
||||
2560
|
||||
1280
|
||||
FFN Dimension
|
||||
896
|
||||
1280
|
||||
1536
|
||||
896
|
||||
Load Balancing Method
|
||||
Loss-Free ( Wang et al., 2024 )
|
||||
Loss-Free
|
||||
Attention Heads
|
||||
16
|
||||
24
|
||||
32
|
||||
16
|
||||
Attention Dimension
|
||||
128
|
||||
128
|
||||
Attention Variant
|
||||
MLA ( Liu et al., 2024a )
|
||||
MLA
|
||||
KV Rank
|
||||
512
|
||||
512
|
||||
Position Embedding
|
||||
RoPE ( Su et al., 2024 )
|
||||
RoPE
|
||||
RoPE Dimension
|
||||
64
|
||||
64
|
||||
RoPE θ \theta
|
||||
10000
|
||||
10000
|
||||
Layer Norm Type
|
||||
RMSNorm ( Zhang and Sennrich, 2019 )
|
||||
RMSNorm
|
||||
Layer Norm ε \varepsilon
|
||||
1e-20
|
||||
1e-20
|
||||
m HC/HC Expansion Rate n n
|
||||
4
|
||||
4
|
||||
m HC/HC Gating Factor Init α \alpha
|
||||
0.01
|
||||
0.01
|
||||
m HC Sinkhorn-Knopp t max t_{\text{max}}
|
||||
20
|
||||
20
|
||||
Sequence Length
|
||||
4096
|
||||
4096
|
||||
Vocab Size
|
||||
129280
|
||||
129280
|
||||
Batch Size
|
||||
320
|
||||
512
|
||||
1280
|
||||
2560
|
||||
Training Steps
|
||||
30000
|
||||
50000
|
||||
50000
|
||||
100000
|
||||
Training Tokens
|
||||
39.3B
|
||||
105B
|
||||
262B
|
||||
1.05T
|
||||
Warmup Steps
|
||||
2000
|
||||
2000
|
||||
Optimizer
|
||||
AdamW ( Loshchilov and Hutter, 2017 )
|
||||
AdamW
|
||||
AdamW Betas
|
||||
(0.9, 0.95)
|
||||
(0.9, 0.95)
|
||||
AdamW ε \varepsilon
|
||||
1e-20
|
||||
1e-20
|
||||
Base Learning Rate
|
||||
8.6e-4
|
||||
5.9e-4
|
||||
4.0e-4
|
||||
9.0e-4
|
||||
Lr Scheduler
|
||||
Step
|
||||
Step
|
||||
Lr Decay Step Ratio
|
||||
[0.8 × \times , 0.9 × \times ]
|
||||
[0.8 × \times , 0.9 × \times ]
|
||||
Lr Decay Rate
|
||||
[0.316, 0.1]
|
||||
[0.316, 0.1]
|
||||
Weight Decay
|
||||
0.1
|
||||
0.1
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,842 @@
|
||||
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.
|
||||
## References
|
||||
[1]
|
||||
Adept (2023)
|
||||
Fuyu-8b .
|
||||
Hugging Face .
|
||||
Note: https://huggingface.co/adept/fuyu-8b
|
||||
Cited by: §2.3 .
|
||||
[2]
|
||||
G. AI (2025)
|
||||
Gemini 2.5-pro .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[3]
|
||||
J. Bai, S. Bai, S. Yang, S. Wang, S. Tan, P. Wang, J. Lin, C. Zhou, and J. Zhou (2023)
|
||||
Qwen-vl: a versatile vision-language model for understanding, localization, text reading, and beyond .
|
||||
arXiv preprint arXiv:2308.12966 .
|
||||
Cited by: §2.2 .
|
||||
[4]
|
||||
S. Bai, Y. Cai, R. Chen, et al. (2025)
|
||||
Qwen3-vl technical report .
|
||||
arXiv preprint arXiv:2511.21631 .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[5]
|
||||
S. Bai, K. Chen, X. Liu, J. Wang, W. Ge, S. Song, K. Dang, P. Wang, S. Wang, J. Tang, H. Zhong, Y. Zhu, M. Yang, Z. Li, J. Wan, P. Wang, W. Ding, Z. Fu, Y. Xu, J. Ye, X. Zhang, T. Xie, Z. Cheng, H. Zhang, Z. Yang, H. Xu, and J. Lin (2025)
|
||||
Qwen2.5-vl technical report .
|
||||
arXiv preprint arXiv:2502.13923 .
|
||||
Cited by: §2.2 ,
|
||||
Table 1 .
|
||||
[6]
|
||||
N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko (2020)
|
||||
End-to-end object detection with transformers .
|
||||
In European conference on computer vision ,
|
||||
pp. 213–229 .
|
||||
Cited by: Figure 2 ,
|
||||
Figure 2 ,
|
||||
§1 ,
|
||||
§2.1 ,
|
||||
§2.2 .
|
||||
[7]
|
||||
J. Chen, L. Kong, H. Wei, C. Liu, Z. Ge, L. Zhao, J. Sun, C. Han, and X. Zhang (2024)
|
||||
Onechart: purify the chart structural extraction via one auxiliary token .
|
||||
In Proceedings of the 32nd ACM International Conference on Multimedia ,
|
||||
pp. 147–155 .
|
||||
Cited by: §4.1 .
|
||||
[8]
|
||||
C. Cui, T. Sun, S. Liang, et al. (2025)
|
||||
Paddleocr-vl: boosting multilingual document parsing via a 0.9 b ultra-compact vision-language model .
|
||||
arXiv preprint arXiv:2510.14528 .
|
||||
Cited by: Table 1 .
|
||||
[9]
|
||||
C. Cui, T. Sun, M. Lin, T. Gao, Y. Zhang, J. Liu, X. Wang, Z. Zhang, C. Zhou, H. Liu, et al. (2025)
|
||||
Paddleocr 3.0 technical report .
|
||||
arXiv preprint arXiv:2507.05595 .
|
||||
Cited by: Table 1 .
|
||||
[10]
|
||||
M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al. (2023)
|
||||
Patch n’ pack: navit, a vision transformer for any aspect ratio and resolution .
|
||||
Advances in Neural Information Processing Systems 36 , pp. 3632–3656 .
|
||||
Cited by: §1 .
|
||||
[11]
|
||||
J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019)
|
||||
Bert: pre-training of deep bidirectional transformers for language understanding .
|
||||
In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers) ,
|
||||
pp. 4171–4186 .
|
||||
Cited by: §2.2 .
|
||||
[12]
|
||||
A. Dosovitskiy (2020)
|
||||
An image is worth 16x16 words: transformers for image recognition at scale .
|
||||
arXiv preprint arXiv:2010.11929 .
|
||||
Cited by: §1 .
|
||||
[13]
|
||||
H. Feng, S. Wei, X. Fei, W. Shi, Y. Han, L. Liao, J. Lu, B. Wu, Q. Liu, C. Lin, et al. (2025)
|
||||
Dolphin: document image parsing via heterogeneous anchor prompting .
|
||||
arXiv preprint arXiv:2505.14059 .
|
||||
Cited by: Table 1 ,
|
||||
Table 1 .
|
||||
[14]
|
||||
K. He, X. Zhang, S. Ren, and J. Sun (2016)
|
||||
Deep residual learning for image recognition .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 770–778 .
|
||||
Cited by: §2.1 .
|
||||
[15]
|
||||
A. Huang, C. Yao, C. Han, F. Wan, H. Guo, H. Lv, H. Zhou, J. Wang, J. Zhou, J. Sun, et al. (2026)
|
||||
STEP3-vl-10b technical report .
|
||||
arXiv preprint arXiv:2601.09668 .
|
||||
Cited by: §3.2.1 .
|
||||
[16]
|
||||
S. Iyer, X. V. Lin, R. Pasunuru, T. Mihaylov, D. Simig, P. Yu, K. Shuster, T. Wang, Q. Liu, P. S. Koura, et al. (2022)
|
||||
OPT-iml: scaling language model instruction meta learning through the lens of generalization .
|
||||
arXiv preprint arXiv:2212.12017 .
|
||||
Cited by: §4.2.1 .
|
||||
[17]
|
||||
A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W. Lo, et al. (2023)
|
||||
Segment anything .
|
||||
arXiv preprint arXiv:2304.02643 .
|
||||
Cited by: §3.2.1 .
|
||||
[18]
|
||||
J. Li, D. Li, S. Savarese, and S. Hoi (2023)
|
||||
Blip-2: bootstrapping language-image pre-training with frozen image encoders and large language models .
|
||||
In International conference on machine learning ,
|
||||
pp. 19730–19742 .
|
||||
Cited by: Figure 2 ,
|
||||
Figure 2 ,
|
||||
§2.2 .
|
||||
[19]
|
||||
Z. Li, Y. Liu, Q. Liu, Z. Ma, Z. Zhang, S. Zhang, Z. Guo, J. Zhang, X. Wang, and X. Bai (2025)
|
||||
MonkeyOCR: document parsing with a structure-recognition-relation triplet paradigm .
|
||||
arXiv preprint arXiv:2506.05218 .
|
||||
Cited by: Table 1 ,
|
||||
Table 1 ,
|
||||
Table 1 .
|
||||
[20]
|
||||
A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al. (2024)
|
||||
Deepseek-v2: a strong, economical, and efficient mixture-of-experts language model .
|
||||
arXiv preprint arXiv:2405.04434 .
|
||||
Cited by: §1 ,
|
||||
§4.2.2 .
|
||||
[21]
|
||||
A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al. (2024)
|
||||
Deepseek-v3 technical report .
|
||||
arXiv preprint arXiv:2412.19437 .
|
||||
Cited by: §4.2.2 .
|
||||
[22]
|
||||
A. Liu, A. Mei, B. Lin, B. Xue, B. Wang, B. Xu, B. Wu, B. Zhang, C. Lin, C. Dong, et al. (2025)
|
||||
Deepseek-v3. 2: pushing the frontier of open large language models .
|
||||
arXiv preprint arXiv:2512.02556 .
|
||||
Cited by: §1 .
|
||||
[23]
|
||||
C. Liu, H. Wei, J. Chen, L. Kong, Z. Ge, Z. Zhu, L. Zhao, J. Sun, C. Han, and X. Zhang (2024)
|
||||
Focus anywhere for fine-grained multi-page document understanding .
|
||||
arXiv preprint arXiv:2405.14295 .
|
||||
Cited by: §4.1 .
|
||||
[24]
|
||||
F. Liu and H. Qiu (2025)
|
||||
Context cascade compression: exploring the upper limits of text compression .
|
||||
arXiv preprint arXiv:2511.15244 .
|
||||
Cited by: §1 .
|
||||
[25]
|
||||
F. Liu, H. Wei, W. Zhao, G. Li, J. Peng, and Z. Li (2021)
|
||||
WB-detr: transformer-based detector without backbone .
|
||||
In Proceedings of the IEEE/CVF International Conference on Computer Vision ,
|
||||
pp. 2979–2987 .
|
||||
Cited by: §2.1 .
|
||||
[26]
|
||||
Y. Liu, J. Gu, N. Goyal, X. Li, S. Edunov, M. Ghazvininejad, M. Lewis, and L. Zettlemoyer (2020)
|
||||
Multilingual denoising pre-training for neural machine translation .
|
||||
Transactions of the Association for Computational Linguistics 8 , pp. 726–742 .
|
||||
Cited by: §3.2.2 .
|
||||
[27]
|
||||
Y. Liu, Z. Zhao, L. Tian, et al. (2025)
|
||||
POINTS-reader: distillation-free adaptation of vision-language models for document conversion .
|
||||
In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing ,
|
||||
pp. 1576–1601 .
|
||||
Cited by: Table 1 .
|
||||
[28]
|
||||
I. Loshchilov and F. Hutter (2019)
|
||||
Decoupled weight decay regularization .
|
||||
In ICLR ,
|
||||
Cited by: §4.2.1 .
|
||||
[29]
|
||||
()
|
||||
Marker .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[30]
|
||||
(2025)
|
||||
Nanonets-ocr-s .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[31]
|
||||
(2025)
|
||||
OCRFlux .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[32]
|
||||
(2025)
|
||||
OCRVerse .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[33]
|
||||
OpenAI (2023)
|
||||
GPT-4 technical report .
|
||||
External Links: arXiv preprint arXiv:2303.08774
|
||||
Cited by: Table 1 .
|
||||
[34]
|
||||
L. Ouyang, Y. Qu, H. Zhou, J. Zhu, R. Zhang, Q. Lin, B. Wang, Z. Zhao, M. Jiang, X. Zhao, et al. (2025)
|
||||
Omnidocbench: benchmarking diverse pdf document parsing with comprehensive annotations .
|
||||
In Proceedings of the Computer Vision and Pattern Recognition Conference ,
|
||||
pp. 24838–24848 .
|
||||
Cited by: §1 ,
|
||||
§5 .
|
||||
[35]
|
||||
Z. Pang, Z. Xie, Y. Man, and Y. Wang (2023)
|
||||
Frozen transformers in language models are effective visual encoder layers .
|
||||
arXiv preprint arXiv:2310.12973 .
|
||||
Cited by: §2.3 .
|
||||
[36]
|
||||
J. Poznanski, A. Rangapur, J. Borchardt, J. Dunkelberger, R. Huff, D. Lin, C. Wilhelm, K. Lo, and L. Soldaini (2025)
|
||||
Olmocr: unlocking trillions of tokens in pdfs with vision language models .
|
||||
arXiv preprint arXiv:2502.18443 .
|
||||
Cited by: Table 1 .
|
||||
[37]
|
||||
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. (2021)
|
||||
Learning transferable visual models from natural language supervision .
|
||||
In International conference on machine learning ,
|
||||
pp. 8748–8763 .
|
||||
Cited by: §1 ,
|
||||
§2.2 .
|
||||
[38]
|
||||
J. Redmon and A. Farhadi (2017)
|
||||
YOLO9000: better, faster, stronger .
|
||||
In Proceedings of the IEEE conference on computer vision and pattern recognition ,
|
||||
pp. 7263–7271 .
|
||||
Cited by: §2.1 .
|
||||
[39]
|
||||
Rednote (2025)
|
||||
Dots.ocr .
|
||||
External Links: Link
|
||||
Cited by: Table 1 .
|
||||
[40]
|
||||
S. Ren, K. He, R. Girshick, and J. Sun (2015)
|
||||
Faster r-cnn: towards real-time object detection with region proposal networks .
|
||||
In Advances in neural information processing systems ,
|
||||
pp. 91–99 .
|
||||
Cited by: §2.1 .
|
||||
[41]
|
||||
B. Seed (2025)
|
||||
Seed1. 8 model card: towards generalized real-world agency .
|
||||
Cited by: Table 2 .
|
||||
[42]
|
||||
J. Su, Y. Lu, S. Pan, B. Wen, and Y. Liu (2021)
|
||||
RoFormer: enhanced transformer with rotary position embedding .
|
||||
arXiv preprint arXiv:2104.09864 .
|
||||
Cited by: §3.2 .
|
||||
[43]
|
||||
C. Team (2024)
|
||||
Chameleon: mixed-modal early-fusion foundation models .
|
||||
arXiv preprint arXiv:2405.09818 .
|
||||
Cited by: §2.3 .
|
||||
[44]
|
||||
G. Team, R. Anil, S. Borgeaud, J. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, K. Millican, et al. (2023)
|
||||
Gemini: a family of highly capable multimodal models .
|
||||
arXiv preprint arXiv:2312.11805 .
|
||||
Cited by: §1 ,
|
||||
Table 2 .
|
||||
[45]
|
||||
B. Wang, B. Wang, C. Wan, G. Huang, H. Hu, H. Jia, H. Nie, M. Li, N. Chen, S. Chen, et al. (2025)
|
||||
Step-3 is large yet affordable: model-system co-design for cost-effective decoding .
|
||||
arXiv preprint arXiv:2507.19427 .
|
||||
Cited by: §3.2.1 .
|
||||
[46]
|
||||
B. Wang, C. Xu, X. Zhao, L. Ouyang, F. Wu, Z. Zhao, R. Xu, K. Liu, Y. Qu, F. Shang, et al. (2024)
|
||||
Mineru: an open-source solution for precise document content extraction .
|
||||
arXiv preprint arXiv:2409.18839 .
|
||||
Cited by: Table 1 ,
|
||||
Table 1 ,
|
||||
Table 1 .
|
||||
[47]
|
||||
C. Wang, S. Chen, Y. Wu, Z. Zhang, L. Zhou, S. Liu, Z. Chen, Y. Liu, H. Wang, J. Li, et al. (2023)
|
||||
Neural codec language models are zero-shot text to speech synthesizers .
|
||||
arXiv preprint arXiv:2301.02111 .
|
||||
Cited by: §2.3 .
|
||||
[48]
|
||||
P. Wang, S. Bai, S. Tan, S. Wang, Z. Fan, J. Bai, K. Chen, X. Liu, J. Wang, W. Ge, et al. (2024)
|
||||
Qwen2-vl: enhancing vision-language model’s perception of the world at any resolution .
|
||||
arXiv preprint arXiv:2409.12191 .
|
||||
Cited by: §1 ,
|
||||
§3.2.2 ,
|
||||
§4.2.1 .
|
||||
[49]
|
||||
W. Wang, Z. Gao, L. Gu, et al. (2025)
|
||||
InternVL3.5: advancing open-source multimodal models in versatility, reasoning, and efficiency .
|
||||
arXiv preprint arXiv:2508.18265 .
|
||||
Cited by: Table 1 .
|
||||
[50]
|
||||
H. Wei, L. Kong, J. Chen, L. Zhao, Z. Ge, J. Yang, J. Sun, C. Han, and X. Zhang (2024)
|
||||
Vary: scaling up the vision vocabulary for large vision-language model .
|
||||
In European Conference on Computer Vision ,
|
||||
pp. 408–424 .
|
||||
Cited by: §2.2 ,
|
||||
§3.2.1 ,
|
||||
§4.2.1 .
|
||||
[51]
|
||||
H. Wei, L. Kong, J. Chen, L. Zhao, Z. Ge, E. Yu, J. Sun, C. Han, and X. Zhang (2024)
|
||||
Small language model meets with reinforced vision vocabulary .
|
||||
arXiv preprint arXiv:2401.12503 .
|
||||
Cited by: §3.2.1 .
|
||||
[52]
|
||||
H. Wei, C. Liu, J. Chen, J. Wang, L. Kong, Y. Xu, Z. Ge, L. Zhao, J. Sun, Y. Peng, et al. (2024)
|
||||
General ocr theory: towards ocr-2.0 via a unified end-to-end model .
|
||||
arXiv preprint arXiv:2409.01704 .
|
||||
Cited by: §3.2.1 .
|
||||
[53]
|
||||
H. Wei, Y. Sun, and Y. Li (2025)
|
||||
Deepseek-ocr: contexts optical compression .
|
||||
arXiv preprint arXiv:2510.18234 .
|
||||
Cited by: §1 ,
|
||||
§1 ,
|
||||
§4.1 .
|
||||
[54]
|
||||
H. Wei, Y. Yin, Y. Li, J. Wang, L. Zhao, J. Sun, Z. Ge, X. Zhang, and D. Jiang (2024)
|
||||
Slow perception: let’s perceive geometric figures step-by-step .
|
||||
arXiv preprint arXiv:2412.20631 .
|
||||
Cited by: §4.1 .
|
||||
[55]
|
||||
J. Zhu, W. Wang, Z. Chen, Z. Liu, S. Ye, L. Gu, H. Tian, Y. Duan, W. Su, J. Shao, et al. (2025)
|
||||
Internvl3: exploring advanced training and test-time recipes for open-source multimodal models .
|
||||
arXiv preprint arXiv:2504.10479 .
|
||||
Cited by: Table 1 .
|
||||
[56]
|
||||
X. Zhu, W. Su, L. Lu, B. Li, X. Wang, and J. Dai (2020)
|
||||
Deformable detr: deformable transformers for end-to-end object detection .
|
||||
arXiv preprint arXiv:2010.04159 .
|
||||
Cited by: §2.1 .
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
@@ -0,0 +1,982 @@
|
||||
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 throughput improvement for offline inference. It also achieves 1.96 × \times higher agent runs per second on average in online serving.
|
||||
## References
|
||||
Agrawal et al. (2024)
|
||||
A. Agrawal, N. Kedia, A. Panwar, J. Mohan, N. Kwatra, B. Gulavani, A. Tumanov, and R. Ramjee
|
||||
Taming Throughput-Latency tradeoff in LLM inference with Sarathi-Serve .
|
||||
In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24) ,
|
||||
pp. 117–134 .
|
||||
Cited by: §6.2 ,
|
||||
§9 .
|
||||
Ainslie et al. (2023)
|
||||
J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebron, and S. Sanghai
|
||||
GQA: training generalized multi-query transformer models from multi-head checkpoints .
|
||||
In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing ,
|
||||
pp. 4895–4901 .
|
||||
Cited by: §9 .
|
||||
Anthropic (2026)
|
||||
Anthropic
|
||||
Introducing Claude Opus 4.6 .
|
||||
Note: https://www.anthropic.com/news/claude-opus-4-6
|
||||
Cited by: §2.2 .
|
||||
Association (2007)
|
||||
I. T. Association
|
||||
InfiniBand Architecture Specification Volume 1, Release 1.2.1 .
|
||||
Cited by: §A.1 ,
|
||||
§5.1 .
|
||||
Carpenter and Nichols (2002)
|
||||
B. E. Carpenter and K. M. Nichols
|
||||
Differentiated services in the internet .
|
||||
Proc. IEEE 90 , pp. 1479–1494 .
|
||||
External Links: Link
|
||||
Cited by: §5.1 .
|
||||
Chen et al. (2026)
|
||||
Q. Chen, Z. Ye, T. Tang, P. Sun, B. Tian, G. Wang, S. Li, Y. Wen, Z. Han, and T. Zhang
|
||||
CONCUR: high-throughput agentic batch inference of llm via congestion-based concurrency control .
|
||||
External Links: 2601.22705 ,
|
||||
Link
|
||||
Cited by: §1 .
|
||||
Chowa et al. (2026)
|
||||
S. S. Chowa, R. Alvi, S. S. Rahman, M. A. Rahman, M. A. K. Raiaan, M. R. Islam, M. Hussain, and S. Azam
|
||||
From language to action: a review of large language models as autonomous agents and tool users .
|
||||
Artificial Intelligence Review .
|
||||
Cited by: §1 .
|
||||
Consortium (2026)
|
||||
U. E. Consortium
|
||||
Ultra ethernet specification v1.0.2 .
|
||||
Cited by: §5.1 .
|
||||
Dao (2024)
|
||||
T. Dao
|
||||
FlashAttention-2: faster attention with better parallelism and work partitioning .
|
||||
In International Conference on Learning Representations ,
|
||||
pp. 35549–35562 .
|
||||
Cited by: §3 .
|
||||
DeepMind (2026)
|
||||
G. DeepMind
|
||||
Gemini 3 Pro .
|
||||
Note: https://deepmind.google/models/gemini/pro/
|
||||
Cited by: §2.2 .
|
||||
DeepSeek-AI (2024)
|
||||
DeepSeek-AI
|
||||
DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model .
|
||||
External Links: 2405.04434
|
||||
Cited by: §9 .
|
||||
DeepSeek-AI (2025a)
|
||||
DeepSeek-AI
|
||||
3FS .
|
||||
Note: https://github.com/deepseek-ai/3FS
|
||||
Cited by: §2.2 ,
|
||||
§7.1 .
|
||||
DeepSeek-AI (2025b)
|
||||
DeepSeek-AI
|
||||
DeepGEMM .
|
||||
Note: https://github.com/deepseek-ai/DeepGEMM
|
||||
Cited by: §3 ,
|
||||
§7.1 .
|
||||
DeepSeek-AI (2025c)
|
||||
DeepSeek-AI
|
||||
DeepSeek-v3 technical report .
|
||||
External Links: 2412.19437 ,
|
||||
Link
|
||||
Cited by: Table 1 ,
|
||||
§3 .
|
||||
DeepSeek-AI (2025d)
|
||||
DeepSeek-AI
|
||||
DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models .
|
||||
External Links: 2512.02556
|
||||
Cited by: §1 ,
|
||||
Table 1 ,
|
||||
§3 ,
|
||||
§7.2 .
|
||||
Du et al. (2025)
|
||||
K. Du, B. Wang, C. Zhang, Y. Cheng, Q. Lan, H. Sang, Y. Cheng, J. Yao, X. Liu, Y. Qiao, I. Stoica, and J. Jiang
|
||||
PrefillOnly: An Inference Engine for Prefill-only Workloads in Large Language Model Applications .
|
||||
In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles ,
|
||||
SOSP ’25 , pp. 399–414 .
|
||||
Cited by: §1 ,
|
||||
§2.1 ,
|
||||
§3 ,
|
||||
§4 ,
|
||||
§6.2 .
|
||||
Gao et al. (2024)
|
||||
B. Gao, Z. He, P. Sharma, Q. Kang, D. Jevdjic, J. Deng, X. Yang, Z. Yu, and P. Zuo
|
||||
Cost-Efficient large language model serving for multi-turn conversations with CachedAttention .
|
||||
In 2024 USENIX Annual Technical Conference (USENIX ATC 24) ,
|
||||
pp. 111–126 .
|
||||
Cited by: §1 .
|
||||
Gao et al. (2025)
|
||||
S. Gao, Y. Chen, and J. Shu
|
||||
Fast State Restoration in LLM Serving with HCache .
|
||||
In Proceedings of the 20th European Conference on Computer Systems ,
|
||||
pp. 128–143 .
|
||||
Cited by: §1 .
|
||||
Guo et al. (2016)
|
||||
C. Guo, H. Wu, Z. Deng, G. Soni, J. Ye, J. Padhye, and M. Lipshteyn
|
||||
RDMA over commodity ethernet at scale .
|
||||
In Proceedings of the 2016 ACM SIGCOMM Conference ,
|
||||
SIGCOMM ’16 , New York, NY, USA , pp. 202–215 .
|
||||
External Links: ISBN 9781450341936 ,
|
||||
Link ,
|
||||
Document
|
||||
Cited by: §5.1 .
|
||||
Holmes et al. (2024)
|
||||
C. Holmes, M. Tanaka, M. Wyatt, A. A. Awan, J. Rasley, S. Rajbhandari, R. Y. Aminabadi, H. Qin, A. Bakhtiari, L. Kurilenko, and Y. He
|
||||
DeepSpeed-FastGen: high-throughput text generation for llms via mii and deepspeed-inference .
|
||||
External Links: 2401.08671
|
||||
Cited by: §9 .
|
||||
Hu et al. (2025)
|
||||
Y. Hu, S. Qiu, J. Yan, H. Chen, X. Wang, T. Lu, G. Xue, and Y. Zhang
|
||||
TARDIS: a gpu-centric kv cache service for efficient llm inference .
|
||||
In Proceedings of the 16th ACM SIGOPS Asia-Pacific Workshop on Systems ,
|
||||
APSys ’25 , New York, NY, USA , pp. 46–53 .
|
||||
External Links: ISBN 9798400715723 ,
|
||||
Link ,
|
||||
Document
|
||||
Cited by: §1 .
|
||||
Jiang et al. (2025)
|
||||
C. Jiang, L. Gao, H. E. Zarch, and M. Annavaram
|
||||
KVPR: efficient LLM inference with I/O-aware KV cache partial recomputation .
|
||||
In Findings of the Association for Computational Linguistics: ACL 2025 ,
|
||||
pp. 19474–19488 .
|
||||
Cited by: §9 .
|
||||
Jiang et al. (2024)
|
||||
J. Jiang, F. Wang, J. Shen, S. Kim, and S. Kim
|
||||
A survey on large language models for code generation .
|
||||
ACM Transactions on Software Engineering and Methodology .
|
||||
Cited by: §1 .
|
||||
Kalia et al. (2016)
|
||||
A. Kalia, M. Kaminsky, and D. G. Andersen
|
||||
Design guidelines for high performance RDMA systems .
|
||||
In 2016 USENIX annual technical conference (USENIX ATC 16) ,
|
||||
pp. 437–450 .
|
||||
Cited by: §5.2 .
|
||||
Kwon et al. (2023)
|
||||
W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. Gonzalez, H. Zhang, and I. Stoica
|
||||
Efficient Memory Management for Large Language Model Serving with PagedAttention .
|
||||
In Proceedings of the 29th Symposium on Operating Systems Principles ,
|
||||
pp. 611–626 .
|
||||
Cited by: §7.1 ,
|
||||
§9 .
|
||||
Li and Liu (2025)
|
||||
J. Li and S. Liu
|
||||
FlashMLA: Efficient Multi-head Latent Attention Kernels .
|
||||
GitHub .
|
||||
Note: https://github.com/deepseek-ai/FlashMLA
|
||||
Cited by: §3 ,
|
||||
§7.1 .
|
||||
Li et al. (2024)
|
||||
Y. Li, H. Wen, W. Wang, X. Li, Y. Yuan, G. Liu, J. Liu, W. Xu, X. Wang, Y. Sun, R. Kong, Y. Wang, H. Geng, J. Luan, X. Jin, Z. Ye, G. Xiong, F. Zhang, X. Li, M. Xu, Z. Li, P. Li, Y. Liu, Y. Zhang, and Y. Liu
|
||||
Personal llm agents: insights and survey about the capability, efficiency and security .
|
||||
External Links: 2401.05459 ,
|
||||
Link
|
||||
Cited by: §1 .
|
||||
Lin et al. (2025)
|
||||
W. Lin, H. Zhen, S. Yang, X. Wang, R. Liu, H. Chen, W. Zhang, C. Zhou, Y. Li, C. Chen, X. Li, Z. Yang, X. Li, X. Yu, Z. Dong, M. Yuan, and Y. Wang
|
||||
Towards efficient agents: a co-design of inference architecture and system .
|
||||
External Links: 2512.18337 ,
|
||||
Link
|
||||
Cited by: §1 .
|
||||
Liu et al. (2025)
|
||||
Y. Liu, Y. Cheng, J. Yao, Y. An, X. Chen, S. Feng, Y. Huang, S. Shen, R. Zhang, K. Du, and J. Jiang
|
||||
LMCache: an efficient kv cache layer for enterprise-scale llm inference .
|
||||
External Links: 2510.09665
|
||||
Cited by: §1 .
|
||||
Mohammadi et al. (2025)
|
||||
M. Mohammadi, Y. Li, J. Lo, and W. Yip
|
||||
Evaluation and benchmarking of llm agents: a survey .
|
||||
In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 2 ,
|
||||
pp. 6129–6139 .
|
||||
Cited by: §1 .
|
||||
NVIDIA (2023)
|
||||
NVIDIA
|
||||
SuperPOD: next generation scalable infrastructure for ai leadership .
|
||||
Note: http://docs.nvidia.com/https:/docs.nvidia.com/dgx-superpod-reference-architecture-dgx-h100.pdf
|
||||
Cited by: §2.3 .
|
||||
NVIDIA (2026a)
|
||||
NVIDIA
|
||||
Developing a linux kernel module using gpudirect rdma .
|
||||
Note: https://docs.nvidia.com/cuda/gpudirect-rdma/index.html
|
||||
Cited by: §5 .
|
||||
NVIDIA (2026b)
|
||||
NVIDIA
|
||||
GPUDirect storage overview guide .
|
||||
Note: https://docs.nvidia.com/gpudirect-storage/overview-guide/index.html
|
||||
Cited by: §5.2 ,
|
||||
§5 .
|
||||
OpenAI (2025a)
|
||||
OpenAI
|
||||
gpt-oss-120b & \& gpt-oss-20b Model Card .
|
||||
External Links: 2508.10925
|
||||
Cited by: Table 1 .
|
||||
OpenAI (2025b)
|
||||
OpenAI
|
||||
Introducing GPT-5.2 .
|
||||
Note: https://openai.com/index/introducing-gpt-5-2/
|
||||
Cited by: §1 .
|
||||
Patel et al. (2025)
|
||||
P. Patel, E. Choukse, C. Zhang, A. Shah, Í. Goiri, S. Maleki, and R. Bianchini
|
||||
Splitwise: efficient generative llm inference using phase splitting .
|
||||
In Proceedings of the 51st Annual International Symposium on Computer Architecture ,
|
||||
pp. 118–132 .
|
||||
Cited by: §1 ,
|
||||
§2.1 ,
|
||||
§4.3 ,
|
||||
§4 ,
|
||||
§9 .
|
||||
Qin et al. (2025)
|
||||
R. Qin, Z. Li, W. He, J. Cui, F. Ren, M. Zhang, Y. Wu, W. Zheng, and X. Xu
|
||||
Mooncake: Trading More Storage for Less Computation — A KVCache-centric Architecture for Serving LLM Chatbot .
|
||||
In Proceedings of the 23rd USENIX Conference on File and Storage Technologies ,
|
||||
pp. 155–170 .
|
||||
Cited by: §1 ,
|
||||
§1 ,
|
||||
§2.2 ,
|
||||
1st item ,
|
||||
§9 .
|
||||
Richter et al. (2016)
|
||||
A. Richter, C. Herber, T. Wild, and A. Herkersdorf
|
||||
Resolving performance interference in sr-iov setups with pcie quality-of-service extensions .
|
||||
In 2016 Euromicro Conference on Digital System Design (DSD) ,
|
||||
Vol. , pp. 454–462 .
|
||||
External Links: Document
|
||||
Cited by: §5 .
|
||||
SGLang (2026)
|
||||
SGLang
|
||||
SGLang HiCache .
|
||||
Note: https://docs.sglang.io/advanced_features/hicache.html
|
||||
Cited by: 1st item .
|
||||
Shazeer (2019)
|
||||
N. Shazeer
|
||||
Fast transformer decoding: one write-head is all you need .
|
||||
External Links: 1911.02150
|
||||
Cited by: §9 .
|
||||
Team (2025a)
|
||||
Q. Team
|
||||
Qwen2.5 technical report .
|
||||
External Links: 2412.15115
|
||||
Cited by: Table 1 ,
|
||||
§7.2 .
|
||||
Team (2025b)
|
||||
Q. Team
|
||||
Qwen3 technical report .
|
||||
External Links: 2505.09388 ,
|
||||
Link
|
||||
Cited by: Table 1 .
|
||||
[43]
|
||||
(2026)
|
||||
UnifiedBus .
|
||||
Note: https://www.unifiedbus.com/en
|
||||
Cited by: §5.1 .
|
||||
Vaswani et al. (2017)
|
||||
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin
|
||||
Attention is All You Need .
|
||||
In Proceedings of the 31st International Conference on Neural Information Processing Systems ,
|
||||
pp. 6000–6010 .
|
||||
Cited by: §9 .
|
||||
Wang et al. (2024)
|
||||
L. Wang, C. Ma, X. Feng, Z. Zhang, H. Yang, J. Zhang, Z. Chen, J. Tang, X. Chen, Y. Lin, et al.
|
||||
A survey on large language model based autonomous agents .
|
||||
Frontiers of Computer Science 18 ( 6 ), pp. 186345 .
|
||||
Cited by: §1 .
|
||||
Wu et al. (2025)
|
||||
B. Wu, Z. Zhang, Y. Zhong, G. Huang, Y. Zhu, X. Liu, and X. Jin
|
||||
TokenLake: a unified segment-level prefix cache pool for fine-grained elastic long-context llm serving .
|
||||
External Links: 2508.17219 ,
|
||||
Link
|
||||
Cited by: §9 .
|
||||
Wu et al. (2023)
|
||||
Q. Wu, G. Bansal, J. Zhang, Y. Wu, B. Li, E. Zhu, L. Jiang, X. Zhang, S. Zhang, J. Liu, A. H. Awadallah, R. W. White, D. Burger, and C. Wang
|
||||
AutoGen: enabling next-gen llm applications via multi-agent conversation .
|
||||
External Links: 2308.08155 ,
|
||||
Link
|
||||
Cited by: §1 .
|
||||
Xi et al. (2025)
|
||||
Z. Xi, W. Chen, X. Guo, W. He, Y. Ding, B. Hong, M. Zhang, J. Wang, S. Jin, E. Zhou, et al.
|
||||
The rise and potential of large language model based agents: a survey .
|
||||
Science China Information Sciences 68 ( 2 ), pp. 121101 .
|
||||
Cited by: §1 .
|
||||
Xie et al. (2025)
|
||||
Z. Xie, Z. Xu, M. Zhao, Y. An, V. S. Mailthody, S. Mahlke, M. Garland, and C. Kozyrakis
|
||||
Strata: hierarchical context caching for long context language model serving .
|
||||
External Links: 2508.18572
|
||||
Cited by: §9 .
|
||||
Xiong et al. (2024)
|
||||
Y. Xiong, H. Wu, C. Shao, Z. Wang, R. Zhang, Y. Guo, J. Zhao, K. Zhang, and Z. Pan
|
||||
LayerKV: optimizing large language model serving with layer-wise kv cache management .
|
||||
External Links: 2410.00428
|
||||
Cited by: §1 ,
|
||||
§2.1 ,
|
||||
§4 .
|
||||
Yan et al. (2025)
|
||||
J. Yan, S. Qiu, Y. Lv, Y. Hu, H. Chen, Z. Shen, X. Yao, R. Chen, J. Shu, G. Zhang, and Y. Zhang
|
||||
Phoenix: a refactored i/o stack for gpu direct storage without phony buffers .
|
||||
In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis ,
|
||||
SC ’25 , pp. 1267–1283 .
|
||||
Cited by: §1 .
|
||||
Yang et al. (2024)
|
||||
J. Yang, C. E. Jimenez, A. Wettig, K. Lieret, S. Yao, K. Narasimhan, and O. Press
|
||||
Swe-agent: agent-computer interfaces enable automated software engineering .
|
||||
Advances in Neural Information Processing Systems 37 , pp. 50528–50652 .
|
||||
Cited by: §1 .
|
||||
Yao et al. (2025)
|
||||
D. Yao, B. Shen, Z. Lin, W. Liu, J. Luan, B. Wang, and W. Wang
|
||||
TailorKV: a hybrid framework for long-context inference via tailored KV cache optimization .
|
||||
In Findings of the Association for Computational Linguistics: ACL 2025 ,
|
||||
pp. 20340–20359 .
|
||||
Cited by: §9 .
|
||||
Ye et al. (2025)
|
||||
Z. Ye, L. Chen, R. Lai, W. Lin, Y. Zhang, S. Wang, T. Chen, B. Kasikci, V. Grover, A. Krishnamurthy, and L. Ceze
|
||||
FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving .
|
||||
External Links: 2501.01005
|
||||
Cited by: §3 .
|
||||
Zhao et al. (2025a)
|
||||
C. Zhao, C. Deng, C. Ruan, D. Dai, H. Gao, J. Li, L. Zhang, P. Huang, S. Zhou, S. Ma, et al.
|
||||
Insights into deepseek-v3: scaling challenges and reflections on hardware for ai architectures .
|
||||
In Proceedings of the 52nd Annual International Symposium on Computer Architecture ,
|
||||
pp. 1731–1745 .
|
||||
Cited by: §1 ,
|
||||
§2.3 .
|
||||
Zhao et al. (2025b)
|
||||
C. Zhao, S. Zhou, L. Zhang, C. Deng, Z. Xu, Y. Liu, K. Yu, J. Li, and L. Zhao
|
||||
DeepEP: an efficient expert-parallel communication library .
|
||||
GitHub .
|
||||
Note: https://github.com/deepseek-ai/DeepEP
|
||||
Cited by: §4.3 ,
|
||||
§7.1 .
|
||||
Zheng et al. (2024)
|
||||
L. Zheng, L. Yin, Z. Xie, C. Sun, J. Huang, C. H. Yu, S. Cao, C. Kozyrakis, I. Stoica, J. E. Gonzalez, C. Barrett, and Y. Sheng
|
||||
SGLang: Efficient Execution of Structured Language Model Programs .
|
||||
In Proceedings of the 38th International Conference on Neural Information Processing Systems ,
|
||||
Cited by: 1st item ,
|
||||
§7.1 .
|
||||
Zhong et al. (2024)
|
||||
Y. Zhong, S. Liu, J. Chen, J. Hu, Y. Zhu, X. Liu, X. Jin, and H. Zhang
|
||||
DistServe: disaggregating prefill and decoding for goodput-optimized large language model serving .
|
||||
In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24) ,
|
||||
pp. 193–210 .
|
||||
Cited by: §1 ,
|
||||
§2.1 ,
|
||||
§4 ,
|
||||
§9 .
|
||||
Zhou et al. (2023)
|
||||
S. Zhou, F. F. Xu, H. Zhu, X. Zhou, R. Lo, A. Sridhar, X. Cheng, T. Ou, Y. Bisk, D. Fried, et al.
|
||||
Webarena: a realistic web environment for building autonomous agents .
|
||||
arXiv preprint arXiv:2307.13854 .
|
||||
Cited by: §1 .
|
||||
## Appendix A Appendix
|
||||
## A.1. Traffic Isolation Configuration Details
|
||||
InfiniBand.
|
||||
The InfiniBand QoS mechanism employs two arbitrators: high-priority and low-priority.
|
||||
Traffic is scheduled using Weighted Round Robin (WRR) in the high-priority arbitrator, then steered to the low-priority arbitrator according to qos_high_limit ; setting it to 255 disables the low-priority arbitrator entirely.
|
||||
Detailed scheduling algorithm can be found in ( Association, 2007 ) .
|
||||
Our configuration:
|
||||
•
|
||||
qos_max_vls 4
|
||||
•
|
||||
qos_high_limit 240
|
||||
•
|
||||
qos_vlarb_high 0:192,1:192,2:0,3:192
|
||||
•
|
||||
qos_vlarb_low 0:192,1:192,2:64,3:192
|
||||
RoCE.
|
||||
RoCE enforces QoS through DSCP-based traffic classification and hardware traffic classes (TC). Packets are first mapped from DSCP values to TCs, each backed by a dedicated hardware queue on NICs and switches (typically up to eight). To match the four-VL configuration in InfiniBand, we configure four lossless RDMA TCs with Priority Flow Control (PFC) enabled. Bandwidth isolation is achieved by assigning proportional scheduling weights to these TCs on both NICs and switches, reserving the majority of bandwidth for model inference traffic while allocating a small fraction to KV-cache traffic to prevent starvation.
|
||||
## A.2. 27B Model Specifications
|
||||
In terms of overall model scale,
|
||||
the hidden dimension 2560,
|
||||
the intermediate size of dense layers is 12288,
|
||||
the number of hidden layers is 30,
|
||||
the number of attention heads is 32,
|
||||
the number of routed experts is 72,
|
||||
the MoE intermediate size is 1536,
|
||||
the number of activated experts per token is 6,
|
||||
the number of shared experts is 2,
|
||||
and the number of initial dense layer is 1.
|
||||
Regarding the index attention mechanism,
|
||||
the number of attention heads is 32,
|
||||
the head dimension is 64,
|
||||
the topk tokens for sparse attention is 1024.
|
||||
The LoRA compression for the Q matrix of both indexer and main attention is removed.
|
||||
## A.3. Agent Task Structure
|
||||
To provide context for the dataset characteristics, we briefly describe the agent task structure, though this background is orthogonal to our system design.
|
||||
Each agent operates within a sandbox environment containing a code repository with known bugs and associated error messages.
|
||||
The agent is instructed via prompt to diagnose and fix the bug.
|
||||
The model possesses tool-use capabilities, invoking bash commands in the sandbox by emitting structured outputs.
|
||||
The agent and environment engage in multi-turn interactions, where each turn consists of a prompt (previous context concatenated with new information, most of which is tool output) and the model generating a subsequent tool invocation by decoding.
|
||||
Each trajectory is a sequence of rounds; round i i consists of appended tokens A i A_{i} and the number of generated tokens g i g_{i} .
|
||||
We use G i G_{i} to indicate the tokens generated in round i i , which are not presented in our dataset.
|
||||
We define C o n t e x t i + 1 Context_{i+1} as the concatenated list of A 1 , G 1 , A 2 , G 2 , … , A i , G i A_{1},G_{1},A_{2},G_{2},...,A_{i},G_{i} .
|
||||
In round i + 1 i+1 of our replay, the agent concatenates the prompt as C o n t e x t i + 1 + A i + 1 Context_{i+1}+A_{i+1} , and then sets proper sampling parameters to ensure it generates exactly g i + 1 g_{i+1} tokens, i.e., G i + 1 G_{i+1} .
|
||||
To generate additional agent trajectories, we sample an existing trajectory and prepend a synthetic round with random tokens as A 1 A_{1} and g 1 = 1 g_{1}=1 .
|
||||
## A.4. Experimental Configurations
|
||||
Configuration Parameters.
|
||||
For DeepSeek models, DualPath allocates 80GB DRAM each node, and SGL(MC) uses totally 1.5TB DRAM on every node.
|
||||
For Qwen 32B, due to the larger KV-Cache, DualPath allocates 320GB.
|
||||
Speculative decoding is disabled for all settings.
|
||||
We use 3FS as the storage backend for all configurations.
|
||||
The short reading queue threshold α \alpha (described in § 6 ) is set to the number of tokens we can read during 3 3 seconds, and the unfinished token upper limit β \beta is set to the number of tokens one GPU can process for 5 5 seconds. Those values are profiled in advance.
|
||||
Compute Quota Threshold is set to 300ms for all DualPath and Oracle baselines.
|
||||
KV-Cache Hit Length Calculation.
|
||||
For all systems except SGL(MC), we limit the KV-Cache hits to only occur within a trajectory, and the hit length is calculated in the client because no eviction is needed.
|
||||
For SGL(MC), hit lengths are computed internally based on HiCache and Mooncake Store cache states.
|
||||
## A.5. KV-Cache Block Layout
|
||||
Layerwise prefill reduces KV-Cache block size to 1 / l a y e r 1/layer of the original, and makes the number of blocks larger to l a y e r × layer\times , posing challenges to transfer and storage performance.
|
||||
To overcome this, we design two distinct block types: Layer Block and Full Block .
|
||||
A Layer Block is a byte tensor with shape [ 1 , t o k e n s , b y t e s ] [1,tokens,bytes] and stores one-layer KV-Cache for some tokens.
|
||||
The number of tokens is called b l o c k _ s i z e block\_size .
|
||||
b y t e s bytes indicates the cache bytes needed per layer per token.
|
||||
Meanwhile, a Full Block has shape [ l a y e r , t o k e n s , b y t e s ] [layer,tokens,bytes] .
|
||||
This design enables us to avoid manual KV-Cache memory layout conversion throughout inference by simply concatenating n n Layer Blocks to yield a Full Block.
|
||||
KV-Cache is stored in distributed storage using a trie structure, where each tree node corresponds to a Full Block.
|
||||
Experimental support, please
|
||||
view the build logs
|
||||
for errors. Generated by
|
||||
L
|
||||
A
|
||||
T
|
||||
E
|
||||
xml
|
||||
.
|
||||
## Instructions for reporting errors
|
||||
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile
|
||||
support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the
|
||||
methods listed below:
|
||||
Click the "Report Issue" (
|
||||
) button, located in the page header.
|
||||
Tip: You can select the relevant text first, to include it in your report.
|
||||
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we
|
||||
may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability
|
||||
should not be a barrier to accessing research. Thank you for your continued support in championing open access for
|
||||
all.
|
||||
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .
|
||||
We gratefully acknowledge support from
|
||||
our major funders ,
|
||||
member institutions , ,
|
||||
and all contributors.
|
||||
About
|
||||
·
|
||||
Help
|
||||
·
|
||||
Contact
|
||||
·
|
||||
Subscribe
|
||||
·
|
||||
Copyright
|
||||
·
|
||||
Privacy
|
||||
·
|
||||
Accessibility
|
||||
·
|
||||
Operational Status (opens in new tab)
|
||||
Major funding support from
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user