xT: Nested Tokenization for Larger Context in Large Images

TL;DR

The paper tackles the challenge of modeling extremely large images without down-sampling or cropping, which degrade high-frequency details and global context. They propose xT, a streaming, two-stage framework that uses nested tokenization to extract region-level features with a lightweight context encoder (Transformer-XL or Mamba variants) to aggregate global context across many regions. xT achieves near-linear memory growth and can handle images up to $29{,}000\times 29{,}000$ pixels, delivering up to $8.6\%$ accuracy gains in classification and $11.6$ in $F_1$ for context-dependent segmentation. The paper also analyzes effective receptive fields and throughput, showing improved context integration with minimal memory overhead.

Abstract

Modern computer vision pipelines handle large images in one of two sub-optimal ways: down-sampling or cropping. These two methods incur significant losses in the amount of information and context present in an image. There are many downstream applications in which global context matters as much as high frequency details, such as in real-world satellite imagery; in such cases researchers have to make the uncomfortable choice of which information to discard. We introduce xT, a simple framework for vision transformers which effectively aggregates global context with local details and can model large images end-to-end on contemporary GPUs. We select a set of benchmark datasets across classic vision tasks which accurately reflect a vision model's ability to understand truly large images and incorporate fine details over large scales and assess our method's improvement on them. xT is a streaming, two-stage architecture that adapts existing vision backbones and long sequence language models to effectively model large images without quadratic memory growth. We are able to increase accuracy by up to 8.6% on challenging classification tasks and $F_1$ score by 11.6 on context-dependent segmentation on images as large as 29,000 x 29,000 pixels.

Figures (6)

• Figure 1: xT allows large images to be modeled end-to-end on contemporary GPUs without compromising on high frequency features or global context.
• Figure 2: Overview of our methodology.xT provides a way for existing vision backbones trained on small images to work effectively with large images. The key is our nested tokenization of large images at multiple levels: at the region level as input $R^0, \dots, R^8$ ($R$, …, $R+8$ for readability) for the region encoders, and then at the patch level $R^i_0,\dots, R^i_{P-1}$ inside the encoders to understand local details. The image regions then undergo independent, hierarchical encoding, by passing through a vision backbone that serves as a region encoder. Hierarchical region encoders result in down-sampled features which, when combined with context encoders, allows us to process more regions at once than typically possible. One such context encoder, Transformer-XL, is illustrated in Stage 2. It recurrently processes previous prior sequence tokens using cross attention, extending its context range significantly with depth. The resulting sequence has assimilated both local and global context and is finally fed to a task-specific decoder.
• Figure 3: xT improves upon the accuracy-parameter frontier of existing methods for high-resolution classification on iNaturalist using nested tokenization on top of the origin architecture for larger context. See Table \ref{['tab:results_inat']} for detailed comparisons.
• Figure 4: Effective receptive fields of Swin-B and Swin-B 〈xT〉 Hyper. The center feature from the last layer of region encoder is used to assess sensitivity to areas across the image. Darker green signifies greater sensitivity.
• Figure 5: xT offers greatly increased accuracy per throughput. On iNaturalist classification, we find that our models only slightly diminished throughput (with the exception of Swin-T 〈xT〉 XL/Mamba) but achieved greater accuracies at each throughput threshold.