PaperlayerView exact revision on arXiv

Source-linked AI summary

ProxyFormer: A Dual-Stream Proxy Architecture for Ultra-Long Context and High-Resolution Generation

Zhongpan Tang

arXiv:2608.23463v1cs.LG

TL;DR

Ultra-long contexts make attention computation and KV caching prohibitively costly. ProxyFormer uses persistent local and compressed proxy streams for global interaction, retaining 92%-95% retrieval accuracy at 1,048,576 tokens after 64K-window training.

  • Problem

    Attention computation and KV-cache size grow quadratically and linearly with sequence length, limiting training length and inference throughput for ultra-long contexts.

  • Method

    ProxyFormer maintains local and proxy streams, compressing local features for global proxy interactions before decompressing contextualized proxies back into the local stream.

  • Results

    92%-95% retrieval accuracy is maintained at 1,048,576 tokens after 64K-window training, while 8K-window training exceeds 94% at 256K tokens.

  • Takeaways & Limitations

    ProxyFormer offers a unified proxy-based architecture for reducing global-interaction computation and caching costs across sequences, images, point clouds, videos, and higher-dimensional tensors.

  • Takeaways & Limitations

    Experiments use validation-scale models with dmodel = 512 and 10 layers, leaving scaling laws and larger-model training stability unverified.

Abstract

from arXiv · show

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.

1 Introduction

ProxyFormer addresses the cost of long-context attention by separating fine-grained local feature preservation from low-cost global interaction in a dual-stream proxy architecture. Its layered compression–interaction–decompression loop preserves information across layers and supports efficient proxy-only inference across sequence and generative-model settings.

  • Motivation: Standard self-attention costs O(L^2d), while autoregressive decoding requires a single-layer KV cache of O(Ld).These costs become central bottlenecks as context lengths reach hundreds of thousands or millions of tokens.
  • Architecture: ProxyFormer uses parallel fine/local and proxy streams, preserving one-to-one local details while restricting global interactions to compressed proxy vectors.The architecture is designed to decouple information fidelity from global interaction cost.
  • Architecture: Each layer performs bottom-up compression, proxy-space interaction, and top-down decompression and injection, while residual local features remain available for later refinement.Global proxy information is injected back into the local stream through decompression and mechanisms such as gating and modulation.
  • Scope: ProxyFormer supports 1D sequences, 2D images, 3D point clouds, voxels, videos, and higher-dimensional tensors across language modeling and generative tasks.The listed applications include autoregressive language modeling, flow-matching image generation, and diffusion conditional generation.
  • Efficiency: The design adds bidirectional stream communication, factorized multi-level compression/decompression, dynamic layer-wise ratios, and proxy-level KV caching.During autoregressive inference, historical proxy states have length L/P, while original historical local features can be released or offloaded.

2 Related Work

ProxyFormer relates to efficient attention, KV-cache compression, and high-resolution generation methods while distinguishing itself through proxy-space interactions that preserve the global operator’s form. It also extends prior fixed-length or fixed-shape compression approaches with dynamic, symmetric, dual-stream processing.

  • Efficient attention: Linear, low-rank, and sparse-attention methods reduce long-sequence cost through feature mappings, projections, or restricted attention patterns.Linear attention and kernel methods target O(L) complexity; Linformer uses learnable projections, while Longformer and BigBird combine local and global patterns.
  • Efficient attention: ProxyFormer performs global interactions in extremely short proxy space without changing the global interaction operator’s mathematical form.The operator can be attention, SSM, RNN, or identity mapping.
  • KV-cache compression: KV-cache methods reduce autoregressive inference memory through attention sinks, token eviction, reusable states, or historical-block summary tokens.StreamingLLM retains sinks and recent tokens; H2O evicts low-scoring tokens; Infini-attention and Landmark Attention compress longer histories.
  • High-resolution generation: For high-resolution image generation, ProxyFormer offers a unified patch-to-proxy pathway for pixel-space JIT and latent-space JLT feature maps.It extracts proxy vectors from 2D patches alongside established latent-space and local-window approaches.
  • Compression architectures: Unlike fixed-length or fixed-shape compression studies, ProxyFormer supports input- and layer-dependent compressed lengths with symmetric decompression and top-down injection in a trainable dual-stream architecture.The architecture systematizes near-lossless fixed-shape compression into a stackable design applicable to autoregressive decoding.

3 ProxyFormer Methodology · 3.1 Problem Definition and Overall Architecture · 3.2 Bottom-Up Compression

ProxyFormer represents arbitrary structured inputs as flattened feature sequences and replaces full-sequence Transformer processing with a dual-stream architecture that compresses local blocks into proxy states for global interaction. Its bottom-up compression supports modality-specific partitioning and projections, residual proxy updates, and asymmetric embeddings for autoregressive models.

  • 3.1 Problem Definition and Overall Architecture: ProxyFormer flattens general D-dimensional inputs into a sequence of length L, covering text, images, 3D data, video, and higher-dimensional tensors.The flattened length is L = ∏D_j=1 N_j, with D = 1 for text, D = 2 for images, and D = 3 for 3D or spatiotemporal data.
  • 3.1 Problem Definition and Overall Architecture: Standard Transformer attention has O(L2) computation and activation-cache costs, motivating ProxyFormer’s proxy-space interaction pathway.The local stream persists while global interaction is performed on the compressed proxy stream.
  • 3.1 Problem Definition and Overall Architecture: Each ProxyFormer layer partitions local features into blocks, compresses them into proxy vectors, updates proxies globally, and injects contextualized proxy information back into local features.The layer uses compression C, decompression D, proxy-state update U, proxy interaction F, and local injection and fusion Φ.
  • 3.2 Bottom-Up Compression: For 1D sequences, ProxyFormer uses contiguous non-overlapping blocks, while images, voxels, point clouds, videos, and higher-dimensional tensors receive modality-specific partitioning schemes.Supported alternatives include sliding windows, spatial patches, 3D grouping or clustering, and recursive flattening.
  • 3.2 Bottom-Up Compression: The compression module maps each local block Xi to a low-dimensional proxy ci, with experiments using reshape–linear projection.Convolutional or unfolding-based projections can be used for 2D, 3D, and video blocks without custom sparse indexers.
  • 3.2 Bottom-Up Compression: Proxy states use residual updates, with gated updates, concatenated projection, and direct substitution offered as alternatives; absent cross-layer input, the update becomes direct compression.The residual formulation allows proxy information to persist across layers.
  • 3.2 Bottom-Up Compression: In autoregressive models, asymmetric dual embeddings assign a low-dimensional historical embedding and a full-dimensional generation embedding with independent parameters.The experiments use dh = 64 and dmodel = 512, compressing historical channel memory by a factor of 8.

3.3 Global Interaction in Proxy Space · 3.4 Top-Down Decompression and Injection · 3.5 Multi-Level Factorized Compression and Decompression

ProxyFormer performs global interaction in a compressed proxy sequence, then decompresses and injects the result while preserving the local stream for iterative refinement. Factorized multi-level compression and decompression reduces the burden of large-ratio projections and extends the design to visual generation.

  • 3.3 Global Interaction in Proxy Space: Proxy interaction is a pluggable operator that can use pass-through, per-proxy MLPs, causal or bidirectional attention, recurrent operators, or cross-modal attention.Causal attention uses a causal mask for autoregressive tasks, while bidirectional attention uses a zero matrix.
  • 3.3 Global Interaction in Proxy Space: P^2 times lower attention cost follows because the proxy sequence has length m ≈ L/P, without relying on low-rank or kernel approximations.The reduction applies to proxy-space attention relative to full attention.
  • 3.4 Top-Down Decompression and Injection: Interacted proxy vectors are decompressed back into local space, where a selectable local fusion mechanism combines them with the local stream.The architecture explicitly separates proxy interaction from top-down local fusion.
  • 3.4 Top-Down Decompression and Injection: Gated cross-attention injects proxy history into the generation region in the language-model implementation.The local stream can subsequently continue refinement with a local Transformer or convolution.
  • 3.4 Top-Down Decompression and Injection: The local stream remains available across layers, allowing missed details to be recompressed later and correcting compression bias through cross-layer iteration.This differs from methods that discard original tokens after a single compression.
  • 3.5 Multi-Level Factorized Compression and Decompression: Conv/ConvTranspose downsampling and upsampling are 2D special cases of the factorized compression and decompression equations, with 3D convolutions extending the view to 3D data.Intermediate features exist only on shorter sequences with narrower channels.
  • 3.5 Multi-Level Factorized Compression and Decompression: Large compression ratios make single linear projections O(Pd_finedproxy) in parameters and computation, motivating factorization of P into M integer factors.The motivation is that such projections can overfit or become untrainable.
  • 3.5 Multi-Level Factorized Compression and Decompression: P = 64, M = 3, and k_j = 4 yields a reshape-projection pathway that significantly reduces single-layer matrix dimensionality.Cascaded compression groups elements, projects through narrower intermediate representations, and reverses the process during decompression.

3.6 Layer-wise Dynamic Compression Ratios and Memory Hierarchy … 3.9 Visual Generation Instantiation

ProxyFormer extends its proxy architecture with layer- and history-dependent compression, proxy-based language-model decoding and caching, and visual diffusion/flow-matching instantiations. These designs preserve local features while restricting global conditioning and historical access to compressed proxy representations.

  • 3.6 Layer-wise Dynamic Compression Ratios and Memory Hierarchy: ProxyFormer supports layer-specific compression ratios and re-partitions proxy states when adjacent layers use different ratios.Different ratios can also be configured for historical blocks at different temporal distances.
  • 3.6 Layer-wise Dynamic Compression Ratios and Memory Hierarchy: Distant history uses coarser proxies while recent history retains finer tokens, controlled through partitioning and compression hyperparameters.Unlike standard Transformers, this configuration does not require changing attention masks or cache layouts.
  • 3.7 Language Model Instantiation: The language model compresses historical features into causal proxy states and combines generation-side self-attention with cross-attention to those proxies.Historical proxy vectors depend only on the current and past blocks, enabling safe caching.
  • 3.7 Language Model Instantiation: Per-head gating lets the language model adaptively determine how much information to read from historical proxies.The model uses a shared RMSNorm-plus-linear output head, standard next-token cross-entropy, and intermediate auxiliary-head losses.
  • 3.8 Proxy KV Cache and Streaming Decoding: During decoding, ProxyFormer caches only historical proxies, while current tokens use self-attention and cross-attention with the frozen proxy cache.When new generation tokens fill a block of size P, that block is compressed and appended to the proxy cache; original local features can be released, quantized, or offloaded.
  • 3.9 Visual Generation Instantiation: In text-to-image and diffusion conditional generation, global conditions apply cross-modal cross-attention only in the compressed image proxy space.The visual experiments use both JIT and JLT paradigms.
  • 3.9 Visual Generation Instantiation: ProxyFormer adapts to visual diffusion and flow matching by compressing image or latent patches into a 2D proxy stream.For patch size p×p, the spatial compression ratio is P = p^2, with local features produced by convolutional embedding.

4 Computational and Memory Complexity Analysis

ProxyFormer replaces full-sequence attention with compressed proxy attention plus linear block-wise compression and decompression, reducing dominant long-sequence costs while retaining a local stream. Its benefits are strongest for long sequences and high-resolution inputs, with further compression enabling channel-width reinvestment but short or wide-local-stream settings limiting gains.

  • Complexity formulation: ProxyFormer uses proxy length m = L/P, with block-wise linear compression and decompression alongside proxy attention.The analysis sets full attention to Cfull(L) = Θ(L2d) and the full KV cache to Mfull(L) = Θ(LlayerLd).
  • Asymptotic benefits: 4096× theoretical proxy attention speedup and 64× theoretical historical KV-cache compression follow from P = 64.These are theoretical values; actual gains depend on implementation, block size, and local stream maintenance cost.
  • Asymptotic benefits: 1/P 2 of full-attention computation is required when L ≫P dfine, while compression and decompression grow only linearly for fixed P.Under this condition, the proxy attention term dominates the computation.
  • Space-for-channel reinvestment: β/α2 scales proxy-attention cost when compression increases from P to αP and proxy channels increase from dproxy to βdproxy.Further compression does not increase the proxy-attention budget as long as β ≤α2.
  • Applicability boundaries: Short sequences or wide local streams where dfine ≈d can make the linear compression term exceed quadratic attention, favoring long sequences and high-resolution inputs.Asymmetric dual embeddings with dh = 64 ≪d = 512 further suppress the linear-term constant.

5 Experiments

Experiments show that ProxyFormer preserves or improves language-modeling quality while substantially reducing training memory and extending long-context retrieval. Image-generation results further demonstrate applicability to both pixel-space and latent-space flow matching.

  • Language Modeling: 21.01 PPL with compressed history improves over the history-less ProxyFormer result of 22.10 and the identically sized decoder-only baseline at 21.36.All models use no positional encoding on WikiText-103-v1.
  • Needle-in-a-Haystack Retrieval: 5,000 retrievals per length evaluate 50 randomly positioned passkeys across 100 independently sampled documents, using greedy decoding and exact 5-digit-password matching.The background text is sampled from WikiText-103 and divided into 50 segments before needle insertion.
  • Needle-in-a-Haystack Retrieval: Over 94% retrieval accuracy is maintained when the 8K-trained model extrapolates to 256K tokens, while the 64K-trained model achieves 92%–95% accuracy across depth intervals at 1,048,576 tokens.Neither model shows obvious lost-in-the-middle behavior within its effective extrapolation range.
  • Memory and Speed: 0.7M-token training remains feasible on a 16 GB GPU: ProxyFormer uses 15.1 GB at 716,800 tokens, while Base cannot run at that length.At 20,992 tokens, ProxyFormer uses 2.9 GB and reaches 16.10 it/s, compared with 15.6 GB and 1.3 it/s for the standard decoder-only model.
  • Image Generation: FID 16.21 and IS 8.42 on CIFAR-10, plus FID 2.06 and IS 1.95 on MNIST, are obtained by ProxyFormer-JIT, while JLT reaches MNIST FID 4.36 and IS 1.93.These preliminary results demonstrate proxy compression and top-down reconstruction in both pixel-space and latent-space flow matching, though they are not state of the art.

6 Discussion and Limitations

ProxyFormer’s preliminary mechanistic assessments attribute its gains to cheaper proxy-space communication, recoverability through the local residual stream, and proxy-only KV caching. The paper’s main limitations concern benchmark coverage, model and visual-task scale, compression choices, and incomplete theory.

  • Mechanistic assessments: Proxy-space interaction reduces long-distance communication cost from L2 to (L/P)2, driving VRAM and speed gains.The local residual stream preserves fine-grained information for later refinement after lossy compression, supporting stable accuracy across depth in NIAH.
  • Limitations: Current public evaluation covers WikiText PPL and multi-needle NIAH, but not large-scale LongBench, RULER, MMLU, or general dialogue.
  • Limitations: The validation models use dmodel = 512 and 10 layers, leaving scaling laws and larger-model training stability unverified.
  • Limitations: Main results use fixed partitioning with linear/reshape compression, while dynamic routing, overlapping patches, attention compression, and heterogeneous dimensions remain unabated.
  • Limitations: Image experiments are limited to MNIST and CIFAR-10 without systematic comparison against DiT/LDM-class high-resolution text-to-image models.
  • Limitations: Theoretical formalization remains needed for expressive capacity and compression-error propagation bounds as functions of depth and compression ratio.

7 Conclusion

ProxyFormer is a dual-stream architecture centered on proxy vectors, connecting local and proxy streams through bottom-up compression, global proxy-space interaction, and top-down decompression and injection. This closed loop reduces global-interaction computation and caching costs for ultra-long sequences while preserving the local residual stream.

  • 7 Conclusion: ProxyFormer establishes a closed loop between local and proxy streams using bottom-up compression, global interaction in proxy space, and top-down decompression and injection.The architecture is centered on proxy vectors.
  • 7 Conclusion: Approximately P 2 and P times reductions in computation and caching costs, respectively, are achieved for global interaction in ultra-long sequences.These reductions occur without discarding the local residual stream.
  • 7 Conclusion: The local residual stream remains available alongside compressed proxy representations.This preserves local information while global interactions occur in proxy space.

A Implementation and Reproduction Details · A.2 Intellectual Property and Licensing

ProxyFormer’s core architecture is patent pending, and public source code does not imply a commercial license. Academic and non-commercial research follows the repository license, while commercial use requires separate authorization.

  • A.2 Intellectual Property and Licensing: Patent protection has been submitted for the core ProxyFormer architecture, which is currently patent pending.
  • A.2 Intellectual Property and Licensing: Publicly available source code does not constitute an implied license for commercial implementation.
  • A.2 Intellectual Property and Licensing: Academic and non-commercial research use follows the license specified in the repository.
  • A.2 Intellectual Property and Licensing: Commercial deployment, for-profit corporate research and development, or product integration requires separate negotiation and authorization from the author.
Loading 2608.23463v1…