Global Context Compression with Interleaved Vision-Text Transformation

TL;DR

This paper introduces VIST2, a Transformer model that interleaves visual and textual information for efficient global context compression, significantly reducing computational costs and memory usage in OCR tasks.

Contribution

VIST2 is a novel Transformer architecture that interleaves visual encodings with text chunks, enabling effective global context compression during both prefilling and inference stages.

Findings

Achieves 3x speedup in first-token generation

Reduces memory usage by 77%

Demonstrates superior performance on long writing tasks

Abstract

Recent achievements of vision-language models in end-to-end OCR point to a new avenue for low-loss compression of textual information. This motivates earlier works that render the Transformer's input into images for prefilling, which effectively reduces the number of tokens through visual encoding, thereby alleviating the quadratically increased Attention computations. However, this partial compression fails to save computational or memory costs at token-by-token inference. In this paper, we investigate global context compression, which saves tokens at both prefilling and inference stages. Consequently, we propose VIST2, a novel Transformer that interleaves input text chunks alongside their visual encoding, while depending exclusively on visual tokens in the pre-context to predict the next text token distribution. Around this idea, we render text chunks into sketch images and train VIST2 in multiple stages, starting from curriculum-scheduled pretraining for optical language modeling, followed by modal-interleaved instruction tuning. We conduct extensive experiments using VIST2 families scaled from 0.6B to 8B to explore the training recipe and hyperparameters. With a 4$\times$ compression ratio, the resulting models demonstrate significant superiority over baselines on long writing tasks, achieving, on average, a 3$\times$ speedup in first-token generation, 77% reduction in memory usage, and 74% reduction in FLOPS. Our codes and datasets will be public to support further studies.