Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Zhang, Wenhu, Wu, Yiming, Wang, Huanyu, Liu, Yaoyang, Dou, Huanzhang, Yang, Senqiao, Wu, Sitong, Zhao, Hanbin, Jia, Jiaya
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866914580391788544
author Zhang, Wenhu
Wu, Yiming
Wang, Huanyu
Liu, Yaoyang
Dou, Huanzhang
Yang, Senqiao
Wu, Sitong
Zhao, Hanbin
Jia, Jiaya
author_facet Zhang, Wenhu
Wu, Yiming
Wang, Huanyu
Liu, Yaoyang
Dou, Huanzhang
Yang, Senqiao
Wu, Sitong
Zhao, Hanbin
Jia, Jiaya
contents Diffusion Language Models (DLMs) enable globally coherent, bidirectional, and controllable text generation, offering advantages over traditional autoregressive LLMs, while scaling to ultra-long sequences remains costly. Many existing block-sparse attention methods select blocks by fixed sampling patterns over the high-resolution attention space, such as tail regions or anti-diagonal stripes. Such prior-driven sampling can miss salient tokens and introduce instability under distribution shifts. In this paper, we propose the Block Approximate Sparse Attention framework (BA-Att) with block-wise pre-downsampled operation, which identifies informative regions within a compact downsampled space, avoiding reliance on brittle positional priors. To analyze its theoretical behavior, we define an oracle post-downsample attention map and formalize the approximation error between pre- and post-downsample schemes. Based on this insight, we introduce a lightweight norm-sorting module and a covariance-compensated correction that approximates full covariance using diagonal QK variances, reducing computational complexity. Extensive experiments show that our operator achieves up to 6.95x acceleration over FlashAttention in attention computation, and maintains near full-attention performance at 50% sparsity across language models, multimodal language models, and video generation models, demonstrating strong efficiency and generalization.
format Preprint
id arxiv_https___arxiv_org_abs_2605_19726
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention
Zhang, Wenhu
Wu, Yiming
Wang, Huanyu
Liu, Yaoyang
Dou, Huanzhang
Yang, Senqiao
Wu, Sitong
Zhao, Hanbin
Jia, Jiaya
Computer Vision and Pattern Recognition
Diffusion Language Models (DLMs) enable globally coherent, bidirectional, and controllable text generation, offering advantages over traditional autoregressive LLMs, while scaling to ultra-long sequences remains costly. Many existing block-sparse attention methods select blocks by fixed sampling patterns over the high-resolution attention space, such as tail regions or anti-diagonal stripes. Such prior-driven sampling can miss salient tokens and introduce instability under distribution shifts. In this paper, we propose the Block Approximate Sparse Attention framework (BA-Att) with block-wise pre-downsampled operation, which identifies informative regions within a compact downsampled space, avoiding reliance on brittle positional priors. To analyze its theoretical behavior, we define an oracle post-downsample attention map and formalize the approximation error between pre- and post-downsample schemes. Based on this insight, we introduce a lightweight norm-sorting module and a covariance-compensated correction that approximates full covariance using diagonal QK variances, reducing computational complexity. Extensive experiments show that our operator achieves up to 6.95x acceleration over FlashAttention in attention computation, and maintains near full-attention performance at 50% sparsity across language models, multimodal language models, and video generation models, demonstrating strong efficiency and generalization.
title Efficient Long-Context Modeling in Diffusion Language Models via Block Approximate Sparse Attention
topic Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2605.19726