ProxyFormer: A Dual-Stream Proxy Architecture for Ultra-Long Context and High-Resolution Generation
2026-08-24 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem that handling very long sequences in language and image models is slow and memory-heavy. They propose ProxyFormer, which compresses detailed information into smaller proxy tokens for global processing and then decompresses it back, keeping important details accessible across layers. This method allows much longer sequences to be trained and used than usual, with good accuracy in retrieval tasks. They also show it could work for image generation, suggesting broad usefulness.
Attention MechanismKey-Value CacheSequence LengthProxy TokensCompressionDecompressionDual-Stream ArchitectureRetrieval AccuracyTransformer ModelsFlow Matching
Authors
Zhongpan Tang
Abstract
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.