Visual Self-Refinement for Autoregressive Models

TL;DR

Autoregressive vision-language models generate a token sequence for images, but spatially structured visuals cause suboptimal outputs due to reliance on previous tokens and error accumulation. The paper introduces a plug-and-play self-refinement module that, after pretraining, refines the entire generated token sequence by learning a lightweight Transformer $g_\phi(\cdot)$ to produce refined embeddings $\mathbf{e'}_{\text{seq}}$ aligned to ground-truth embeddings via cosine-distance supervision, followed by nearest-neighbor decoding to discrete tokens and VQGAN-based image reconstruction while keeping the backbone frozen. The approach delivers three main contributions: (1) a post-pretraining finetuning operation preserving generalization, (2) a compact, plug-and-play refinement module with low data and compute cost, and (3) empirical gains across colorization, inpainting, and edge-detection tasks and across different pretrained AR backbones. The results show reduced error accumulation and improved visual fidelity, demonstrating a practical path to leverage global context in autoregressive visual generation with high efficiency and generalization potential.

Abstract

Autoregressive models excel in sequential modeling and have proven to be effective for vision-language data. However, the spatial nature of visual signals conflicts with the sequential dependencies of next-token prediction, leading to suboptimal results. This work proposes a plug-and-play refinement module to enhance the complex spatial correspondence modeling within the generated visual sequence. This module operates as a post-pretraining step to jointly refine all generated tokens of autoregressive model, enhancing vision-language modeling under a shared sequential prediction framework. By leveraging global context and relationship across the tokens, our method mitigates the error accumulation issue within the sequential generation. Experiments demonstrate that the proposed method improves the generation quality, enhancing the model's ability to produce semantically consistent results.

Figures (4)

• Figure 1: Training and generation pipeline of the proposed method. Input images ($K$ demonstrations pairs and a query image) are encoded and fed into $f_{AR}(\cdot)$ for generating tokens iteratively. A refinement network $g_\phi(\cdot)$ is introduced to improve the generated tokens $y_q$ in the embedding space, yielding $\mathbf{e}'_{seq}$. During training, we use ground truth embedding $\mathbf{e}^*_{seq}$ for supervision and minimize the cosine distance. During generation, the refined embedding $\mathbf{e}'_{seq}$ is decoded to the discrete tokens via nearest-neighbor search. Finally, the refined tokens $y'_q$ are passed to VQGAN decoder to generate the image. We take edge detection as an example.
• Figure 2: We compare the LVM with LoRA and the LVM with self-refinement in terms of the cosine distance among generated embeddings and the ground truth embeddings for $256$ tokens. Self-refinement reduces the error accumulation especially for latter tokens ($>160$). The accumulated errors (i.e., by AUC) are provided.
• Figure 3: Effect of demonstration image length for the proposed self-refinement (e.g., upon colorization). The input image sequence is composed of one or more image pairs plus one query RGB image. The proposed method enables notable performance boost under different context lengths, demonstrating the robustness to the context length.
• Figure : Training with Refinement Network