Autoregressive Sequence Modeling for 3D Medical Image Representation

TL;DR

This paper introduces an autoregressive sequence modeling framework for 3D medical images, enabling better contextual understanding and representation learning across diverse medical imaging tasks.

Contribution

It proposes a novel autoregressive pre-training method that sequences 3D medical images as interconnected tokens, capturing complex relationships for improved downstream task performance.

Findings

Outperforms existing methods on nine public datasets

Effectively models complex relationships among local regions

Enhances robustness with a random startup strategy

Abstract

Three-dimensional (3D) medical images, such as Computed Tomography (CT) and Magnetic Resonance Imaging (MRI), are essential for clinical applications. However, the need for diverse and comprehensive representations is particularly pronounced when considering the variability across different organs, diagnostic tasks, and imaging modalities. How to effectively interpret the intricate contextual information and extract meaningful insights from these images remains an open challenge to the community. While current self-supervised learning methods have shown potential, they often consider an image as a whole thereby overlooking the extensive, complex relationships among local regions from one or multiple images. In this work, we introduce a pioneering method for learning 3D medical image representations through an autoregressive pre-training framework. Our approach sequences various 3D medical images based on spatial, contrast, and semantic correlations, treating them as interconnected visual tokens within a token sequence. By employing an autoregressive sequence modeling task, we predict the next visual token in the sequence, which allows our model to deeply understand and integrate the contextual information inherent in 3D medical images. Additionally, we implement a random startup strategy to avoid overestimating token relationships and to enhance the robustness of learning. The effectiveness of our approach is demonstrated by the superior performance over others on nine downstream tasks in public datasets.