Search papers, labs, and topics across Lattice.
This paper introduces ProxyFormer, a dual-stream architecture that addresses the computational bottleneck of attention mechanisms in ultra-long-context language models and high-resolution generative models by utilizing proxy tokens for efficient information processing. By compressing fine-grained local features into a small set of proxy states and performing global interactions in this compressed space, ProxyFormer significantly extends the trainable sequence length to approximately 0.7 million tokens while maintaining high retrieval accuracy. The architecture's innovative approach to multi-level compression and dynamic ratios allows for effective information retention and refinement across layers, showcasing its potential in both text and image generation tasks.
ProxyFormer achieves a staggering 0.7 million token context length while retaining up to 95% retrieval accuracy, revolutionizing the scalability of generative models.
The quadratic growth of attention computation and key-value (KV) cache with respect to sequence length is a central bottleneck for ultra-long-context language models and high-resolution generative models. We propose ProxyFormer, a general dual-stream architecture built upon proxy tokens. In each layer, fine-grained local features are compressed bottom-up into a small set of proxy states; expensive global interactions are performed only in the compressed proxy space; the globally contextualized proxies are then decompressed and injected top-down back into the local stream. Because the local stream persists across layers, fine-grained information that is not captured by one compression step remains accessible for later refinement, alleviating the irreversible information loss of conventional one-shot compression. We further introduce factorized multi-level compression/decompression, layer-wise dynamic compression ratios, asymmetric dual embeddings, and a proxy-only KV-cache inference scheme. On a 16GB GPU with batch size 1, a standard decoder-only model can train sequences of only about 20K tokens, whereas ProxyFormer with a compression ratio of 64 extends the trainable sequence length to about 0.7M. A model trained with a 64K window retains 92%-95% retrieval accuracy on a multi-needle retrieval task with 1,048,576 tokens, and a model trained with an 8K window exceeds 94% accuracy when extrapolated to 256K tokens. Preliminary image-generation experiments demonstrate the feasibility of ProxyFormer for both pixel-space and latent-space flow matching.
