SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention

• URL: http://arxiv.org/abs/2604.13847v1
• Date: Wed, 15 Apr 2026 13:18:07 GMT
• Title: SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention
• Abstract summary: We propose SparseBalance, which exploits the sparsity and sequence heterogeneity to optimize model accuracy and system efficiency jointly.<n>First, we propose workload-aware dynamic sparsity tuning, which employs a bidirectional sparsity adjustment to eliminate stragglers inherent to free accuracy.<n>Second, we propose a sparityaware strategy to achieve coarse-grained sparsity tuning, which exploits sparse up to 1.33$times$ end-to-end speedup while still improving long-context capability by 0.46% on the LongBench benchmark.
• Score: 19.076502375242352
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: While sparse attention mitigates the computational bottleneck of long-context LLM training, its distributed training process exhibits extreme heterogeneity in both \textit{1)} sequence length and \textit{2)} sparsity sensitivity, leading to a severe imbalance problem and sub-optimal model accuracy. Existing algorithms and training frameworks typically focus on single issue, failing to systematically co-optimize these two problems. Therefore, we propose SparseBalance, a novel algorithm-system co-design framework, which exploits the sparsity and sequence heterogeneity to optimize model accuracy and system efficiency jointly. First, we propose workload-aware dynamic sparsity tuning, which employs a bidirectional sparsity adjustment to eliminate stragglers and exploit inherent bubbles for free accuracy. Second, we propose a sparsity-aware batching strategy to achieve coarse-grained balance, which complements dynamic sparsity tuning. Experimental results demonstrate that SparseBalance achieves up to a 1.33$\times$ end-to-end speedup while still improving the long-context capability by 0.46\% on the LongBench benchmark.

Related papers

• When Does Scale-Invariant Optimization Become Unstable? An Exact Schedule Law with Weight Decay [6.625270757109195]
  We show that a single quantity captures all schedule and decay forcing, while norm growth induces an opposing geometric self-quenching effect.<n>This yields a sharp boundary that cleanly separates contraction- and expansion-dominated effective learning rate regimes.
  arXiv  Detail & Related papers  (2026-09-08T17:46:17Z)
• Precise: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models [56.67321805551389]
  Reinforcement learning (RL) has become an effective way to improve prompt alignment and perceptual quality in diffusion and flow-matching generators.<n>The sampler, controlling the exploration behavior and denoising dynamics, is thus part of the policy.<n>We propose a new sampler that balances effective exploration with stability.
  arXiv  Detail & Related papers  (2026-05-22T11:37:22Z)
• $φ$-Balancing for Mixture-of-Experts Training [20.061328126956028]
  Mixture-of-Experts (MoE) models rely on balanced expert utilization to fully realize their scalability.<n>We propose $$-balancing, a principled framework that directly targets population-level expert balance by minimizing a strictly convex, symmetric, and differentiable potential of the expected routing distribution.
  arXiv  Detail & Related papers  (2026-05-14T20:39:28Z)
• Learning-Augmented Scalable Linear Assignment Problem Optimization via Neural Dual Warm-Starts [19.540758462427878]
  We propose a learning-augmented framework that accelerates exact assignment solvers while maintaining optimality and worst-case guarantees.<n>We introduce RowDualNet, a lightweight row-independent architecture that avoids the $mathcalO(N2)$ memory bottleneck of graph-based models.
  arXiv  Detail & Related papers  (2026-05-10T07:15:49Z)
• Nora: Normalized Orthogonal Row Alignment for Scalable Matrix Optimizer [52.9566744833465]
  Matrix-based approaches to training Large Language Models (LLMs) have demonstrated immense potential.<n>A superior must satisfy three core desiderata: efficiency, achieving Muon-like preconditioning to accelerate optimization; stability, strictly adhering to the inherent in neural networks; and speed, minimizing computational overhead.<n>We propose Nora, which rigorously satisfies all three requirements and is highly promising for large-scale training.
  arXiv  Detail & Related papers  (2026-05-05T14:00:27Z)
• Neural Network Pruning via QUBO Optimization [0.08796261172196741]
  We propose a unified Hybrid QUBO framework that bridges the importance of estimation with global optimization.<n>Our formulation integrates gradient-aware sensitivity metrics into the linear term, while utilizing data-driven similarity in the quadratic term.<n> Experiments on the SIDD image denoising dataset demonstrate that the proposed Hybrid QUBO significantly outperforms both greedy Taylor pruning and traditional L1-based QUBO.
  arXiv  Detail & Related papers  (2026-04-07T13:16:40Z)
• Taming the Instability: A Robust Second-Order Optimizer for Federated Learning over Non-IID Data [77.19117648492104]
  We present Federated Robust Curvature Optimization (FedRCO), a novel second-order optimization framework designed to improve convergence speed and reduce communication cost.<n>FedRCO addresses these challenges by integrating an efficient approximate curvature with a provable stability mechanism.<n>We show that FedRCO achieves superior robustness against diverse non-IID scenarios while achieving higher accuracy and faster convergence than both state-of-the-art first-order and second-order methods.
  arXiv  Detail & Related papers  (2026-03-30T11:37:46Z)
• Sparse Training of Neural Networks based on Multilevel Mirror Descent [0.688204255655161]
  We introduce a dynamic sparse training algorithm based on linearized Bregman iterations / mirror descent.<n>We empirically show that our algorithm can produce highly sparse and accurate models on standard benchmarks.
  arXiv  Detail & Related papers  (2026-02-03T13:51:45Z)
• Trajectory Consistency for One-Step Generation on Euler Mean Flows [24.038760671907024]
  We propose emphEuler Mean Flows (EMF), a flow-based generative framework for one-step and few-step generation.<n>EMF enforces long-range trajectory consistency with minimal sampling cost.
  arXiv  Detail & Related papers  (2026-01-31T04:32:32Z)
• ROOT: Robust Orthogonalized Optimizer for Neural Network Training [47.05662448082334]
  Large language models (LLMs) remain a critical challenge, particularly as model scaling exacerbates sensitivity to imprecision and training instability.<n>We develop a dimension-robustization scheme that enhances robustness through iterations tailored to specific matrix sizes.<n>Second, we introduce an optimization-robustization framework that suppresses outliers noise while preserving meaningful directions.
  arXiv  Detail & Related papers  (2025-11-25T18:48:05Z)
• Seer Self-Consistency: Advance Budget Estimation for Adaptive Test-Time Scaling [55.026048429595384]
  Test-time scaling improves the inference performance of Large Language Models (LLMs) but also incurs substantial computational costs.<n>We propose SeerSC, a dynamic self-consistency framework that simultaneously improves token efficiency and latency.
  arXiv  Detail & Related papers  (2025-11-12T13:57:43Z)
• PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
  Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
  arXiv  Detail & Related papers  (2025-09-27T03:01:48Z)
• Generalized Linear Bandits: Almost Optimal Regret with One-Pass Update [70.38810219913593]
  We study the generalized linear bandit (GLB) problem, a contextual multi-armed bandit framework that extends the classical linear model by incorporating a non-linear link function.<n>GLBs are widely applicable to real-world scenarios, but their non-linear nature introduces significant challenges in achieving both computational and statistical efficiency.<n>We propose a jointly efficient algorithm that attains a nearly optimal regret bound with $mathcalO(1)$ time and space complexities per round.
  arXiv  Detail & Related papers  (2025-07-16T02:24:21Z)
• Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization [11.223375172715722]
  We propose Mist, a memory, overlap, and imbalance-aware automatic distributed training system.<n>We show that Mist achieves an average of 1.28$times$ (up to 1.73$times$) speedup compared to state-of-the-art manual system Megatron-LM and state-of-the-art automatic system Aceso.
  arXiv  Detail & Related papers  (2025-03-24T18:21:08Z)
• Tree ensemble kernels for Bayesian optimization with known constraints over mixed-feature spaces [54.58348769621782]
  Tree ensembles can be well-suited for black-box optimization tasks such as algorithm tuning and neural architecture search. Two well-known challenges in using tree ensembles for black-box optimization are (i) effectively quantifying model uncertainty for exploration and (ii) optimizing over the piece-wise constant acquisition function. Our framework performs as well as state-of-the-art methods for unconstrained black-box optimization over continuous/discrete features and outperforms competing methods for problems combining mixed-variable feature spaces and known input constraints.
  arXiv  Detail & Related papers  (2022-07-02T16:59:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.