Jiangxia Cao

Author directory

2026

Long-context capability has become increasingly critical for next-generation Large Language Models (LLMs). However, the quadratic time complexity of standard softmax attention introduces substantial computation and memory overhead as sequence length grows, limiting the scalability of LLMs in long-context scenarios. Beyond existing approaches such as linear recurrent architectures and KV-cache compression methods, we propose Kwai Summary Attention (KSA), a novel attention mechanism that reduces sequence modeling cost by compressing historical contexts into learnable summary tokens. Specifically, summary tokens are periodically inserted into the input sequence and interleaved with text tokens, partitioning the sequence into distinct chunks while enabling text tokens and summary tokens to operate under different attention scopes. This design preserves dense local expressiveness through full attention among neighboring text tokens, while modeling long-range dependencies efficiently via summary tokens. Empirical results demonstrate that hybrid-KSA consistently outperforms existing hybrid attention variants in long-context modeling. Moreover, KSA is fully orthogonal to GQA and MLA during inference, and combining KSA with these methods achieves up to 8x KV-cache compression while preserving model performance.

2022

Chinese pre-trained language models usually exploit contextual character information to learn representations, while ignoring the linguistics knowledge, e.g., word and sentence information. Hence, we propose a task-free enhancement module termed as Heterogeneous Linguistics Graph (HLG) to enhance Chinese pre-trained language models by integrating linguistics knowledge. Specifically, we construct a hierarchical heterogeneous graph to model the characteristics linguistics structure of Chinese language, and conduct a graph-based method to summarize and concretize information on different granularities of Chinese linguistics hierarchies. Experimental results demonstrate our model has the ability to improve the performance of vanilla BERT, BERTwwm and ERNIE 1.0 on 6 natural language processing tasks with 10 benchmark datasets. Further, the detailed experimental analyses have proven that this kind of modelization achieves more improvements compared with previous strong baseline MWA. Meanwhile, our model introduces far fewer parameters (about half of MWA) and the training/inference speed is about 7x faster than MWA.