Query-Aware Cross-Segment KV Compression with Attention-Mass Calibration: An Operator-Level Study
DOI:
https://doi.org/10.54097/vzmvf788Keywords:
Masked diffusion language models, KV cache, Attention pooling, Adaptive allocation, Reproducibility, Long contextAbstract
Cross-segment KV compression reduces the history accessible during masked denoising, but the number of summaries retained from each completed segment can distort attention mass. We implement the learned-query CSKV pooling operation of SW-SpeedDLM using Q/K/V extracted from an actual four-head bidirectional masked encoder, then examine query-aware integer budget allocation and an analytic log(W/c) multiplicity correction. Evaluation uses 170 held-out 320-character spans from an Opticks transcription and a Wikipedia extract. With four 64-character segments, uniform pooling at the compressor’s training budget (B=16) is hard to beat with prefix-based allocation (NMSE 0.0054 versus 0.0050), whereas at untrained budgets such as B=8 adaptive allocation avoids the collapse of uniform pooling (0.0389 versus 1.8093). With query-ranked slots, calibrated adaptive routing at B=16 achieves 0.0044 versus 0.0051 without calibration; however, a paired interval for its difference against uniform pooling crosses zero. Further independently initialized and variable-budget-trained compressors show that allocation effects depend on learned bank quality. The block-constant attention identity is verified numerically and final-bank storage is exactly 93.75% smaller at B=16 than 256 uncompressed positions. The experiment measures cross-window attention fidelity, not end-to-end diffusion generation or GPU throughput.
Downloads
References
[1] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 30, 5998-6008. https://doi.org/10.48550/arXiv.1706.03762
[2] Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. (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 (NAACL-HLT) (pp. 4171-4186). Association for Computational Linguistics. https://doi.org/10.18653/v1/N19-1423
[3] Austin, J., Johnson, D. D., Ho, J., Tarlow, D., & van den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34, 17981-17993. https://doi.org/10.48550/arXiv.2107.03006
[4] Sahoo, S. S., Arora, S., Goyal, A., & Arora, S. (2024). Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37. https://doi.org/10.52202/079017-4135
[5] Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., & Li, C. (2025). Large language diffusion models. Advances in Neural Information Processing Systems, 38. https://doi.org/10.52202/085713-1689
[6] Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The long-document transformer. arXiv, arXiv:2004.05150. https://doi.org/10.48550/arXiv.2004.05150
[7] Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontañón, S., Pham, P., Ravula, A., Wang, Q., Yang, L., & Ahmed, A. (2020). Big Bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 17283-17297. https://doi.org/10.48550/arXiv.2007.14062
[8] Chen, J., Zhang, X., & Han, X. (2023). Explainable identification of tax-related financial reporting anomalies in U.S. public companies: Cross-fitted tax-consistency residuals and Owen-value Shapley propagation. World Journal of Economics and Business Research, 1(1), 33-45.
[9] Dao, T. (2024). FlashAttention-2: Faster attention with better parallelism and work partitioning. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2307.08691
[10] Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024). Efficient streaming language models with attention sinks. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2309.17453
[11] Yue, X., & Zhu, R. (2024). An explainable hybrid machine learning framework for cash-flow-conditioned multi-horizon liquidity risk early warning: An SME-oriented benchmark study. World Journal of Management Science, 2(1), 63-72.
[12] Li, Y., Huang, Y., Yang, Y., Zhang, J., Luo, M., Wang, H., & Zhang, Y. (2024). SnapKV: LLM knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37. https://doi.org/10.52202/079017-0722
[13] Cai, Z., Zhang, Y., Gao, L., & Liu, T. (2025). PyramidKV: Dynamic KV cache compression based on pyramidal information funneling. In Proceedings of the Conference on Language Modeling (COLM). https://doi.org/10.48550/arXiv.2406.02069
[14] Chevalier, A., Wettig, A., Ajith, A., & Chen, D. (2023). Adapting language models to compress contexts. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP) (pp. 3829-3846). Association for Computational Linguistics. https://doi.org/10.18653/v1/2023.emnlp-main.232
[15] Rae, J. W., Potapenko, A., Jayakumar, S. M., & Lillicrap, T. P. (2020). Compressive transformers for long-range sequence modelling. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.1911.05507
[16] Jaegle, A., Borgeaud, S., Alayrac, J.-B., Doersch, C., Ionescu, C., Ding, D., Koppula, S., Zoran, D., Brock, A., Shelhamer, E., Hénaff, O., Botvinick, M. M., Zisserman, A., Vinyals, O., & Carreira, J. (2022). Perceiver IO: A general architecture for structured inputs & outputs. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2107.14795
[17] Ghazvininejad, M., Levy, O., Liu, Y., & Zettlemoyer, L. (2019). Mask-Predict: Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) (pp. 6112-6121). Association for Computational Linguistics. https://doi.org/10.18653/v1/D19-1633
[18] Lou, A., Meng, C., & Ermon, S. (2024). Discrete diffusion modeling by estimating the ratios of the data distribution. In Proceedings of the 41st International Conference on Machine Learning (ICML), PMLR 235 (pp. 32819-32848). https://doi.org/10.48550/arXiv.2310.16834
[19] Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., & Li, J. (2024). LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (ACL) (pp. 3119-3137). Association for Computational Linguistics. https://doi.org/10.18653/v1/2024.acl-long.172
[20] Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning (ICML), PMLR 202 (pp. 19274-19286). https://doi.org/10.48550/arXiv.2211.17192
[21] Chen, J., Han, X., Zhang, X., & Yang, Y. (2023). Exact concept-level explanations for early detection of financial misstatements. Eurasia Journal of Science and Technology, 5(2), 34-48.
[22] Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L., & Weller, A. (2021). Rethinking attention with performers. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2009.14794
[23] Press, O., Smith, N. A., & Lewis, M. (2022). Train short, test long: Attention with linear biases enables input length extrapolation. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2108.12409
[24] Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., & Liu, Y. (2024). RoFormer: Enhanced transformer with rotary position embedding. Neurocomputing, 568, Article 127063. https://doi.org/10.1016/j.neucom.2023.127063
[25] Teng, D., Rhee, M., Qin, Y., Zi, B., & Liu, W. (2026). SW-SpeedDLM: Sliding window speculative decoding for diffusion language models under long context constraints. Mathematics, 14(12), Article 2137. https://doi.org/10.3390/math14122137
[26] Bolya, D., Fu, C.-Y., Dai, X., Zhang, P., Feichtenhofer, C., & Hoffman, J. (2023). Token merging: Your ViT but faster. In Proceedings of the International Conference on Learning Representations (ICLR). https://doi.org/10.48550/arXiv.2210.09461
[27] Shang, W., Teng, D., Chen, T., Zhang, F., & Ding, J. (2026). Neuro-Elastic: A unified framework for hardware-aware adaptive quantization and dynamic sparsity in real-time edge intent prediction. IEEE Access, 14, 1-16. https://doi.org/10.1109/ACCESS.2026.3715276
[28] Shen, Z., Zhao, W., Wang, B., Wang, Z., & Shang, W. (2026). CAGR: A cross-accelerator graph optimization framework for efficient recommender system inference. IEEE Access, 14, 1-15. https://doi.org/10.1109/ACCESS.2026.3671798
Downloads
Published
Issue
Section
License
Copyright (c) 2026 Computer Life
This work is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.
How to Cite
