MassAlloc Attention: Let Attention Allocate Its Own Compute

MassAlloc Attention. Let Attention Allocate Its Own Compute. By Jingze Shi and eight others, from the Hong Kong University of Science and Technology in Guangzhou, the Beijing Academy of Artificial Intelligence, and one other institution. Published on the twenty-sixth of September. Its code repository is worked out: about two hundred source files, with a readme and a licence.

In full attention, FullAttn in the abstract's shorthand, much of the causal score space often ends up with negligible normalized mass. Most of the probability sits on a small part of it. Dense kernels ignore that. Once they have formed each tile of query-key scores, they execute the complete post-score path regardless, spending the same work on interactions that contribute almost nothing as on the ones that carry the output.

MALA is a fused attention primitive that keeps score access to every legal causal interaction, so nothing is ruled out in advance, and then uses each interaction's normalized contribution to decide how much post-score computation to allocate. The scores are still all computed. What changes is where the work after them goes.

The forward pass relies on its evolving online-softmax normalizer. The backward pass reuses the finalized normalizer to derive a nested retained support, and it does so using only the standard state that attention already keeps. A single common tolerance governs both training and inference, which lets the amount of retained work adapt. The result is that low-contribution post-score computation is reduced.

A matched-work study at eight thousand tokens isolates the value of allocating according to the distribution. Under exactly matched total post-score work, MALA approaches a per-instance oracle that knows the reference mass: its mean omitted mass is zero point zero one eight eight percent, against zero point zero one eight two percent for the oracle. Across context lengths from one thousand to thirty-two thousand tokens, the same tolerance keeps output and gradient errors low relative to the reference. In a broader controlled associative-recall comparison, MALA closely tracks FullAttn as context grows, reaching eighty-nine point six seven percent accuracy at eight thousand tokens compared with eighty-nine point nine seven percent.

On an attention-operator benchmark at one hundred and twenty-eight thousand tokens with tensor parallelism, MALA reduces training forward latency by two point two times and backward latency by three times. Decoding latency at inference drops by one point six times relative to FullAttn. Scaling-law training runs from zero point six billion up to fourteen billion parameters, and MALA closely tracks FullAttn in perplexity while reducing total training FLOPs. The fourteen billion parameter models that result, and thirty-two billion parameter models from separate continued training, achieve knowledge, reasoning, and long-context retrieval scores comparable to FullAttn.

Where a fixed sparse pattern would decide in advance which tokens to skip, MALA decides after the scores exist, by how much each interaction actually contributes. Nothing is hidden from the scoring step. Only the follow-on work is rationed, under one tolerance shared by training and inference. The authors read their results as showing that this kind of allocation can retain the evaluated capabilities of FullAttn while doing less attention computation.