OWT: A Foundational Organ-Wise Tokenization Framework for Medical Imaging

TL;DR

The paper tackles the interpretability and generalization limitations of holistic embeddings in medical imaging by introducing Organ-Wise Tokenization (OWT) with Token Group-based Reconstruction (TGR). OWT decomposes images into Semantically Disentangled Token Groups (SDTGs) via an Organ Collector and Adaptive Holistic Embedding Restorer (AHER), enabling organ-specific analysis, reconstruction, and downstream tasks from a single training run. Key contributions include a concrete tokenization framework, a semantically grounded reconstruction loss, and evidence of organ-specific tumor identification, organ-level retrieval, and semantic generation across CT, MRI, and even CelebAMaskHQ data, with favorable efficiency (lower GFLOPs) and data-efficiency (semi-supervised training). The approach offers a scalable, interpretable foundation for semantically disentangled representations in medical imaging and beyond, with practical implications for targeted clinical workflows and multimodal integration.

Abstract

Recent advances in representation learning often rely on holistic embeddings that entangle multiple semantic components, limiting interpretability and generalization. These issues are especially critical in medical imaging, where downstream tasks depend on anatomically interpretable features. To address these limitations, we propose an Organ-Wise Tokenization (OWT) framework with a Token Group-based Reconstruction (TGR) training paradigm. Unlike conventional approaches, OWT explicitly disentangles an image into separable token groups, each corresponding to a distinct organ or semantic entity. Our design ensures each token group encapsulates organ-specific information, boosting interpretability, generalization, and efficiency while enabling fine-grained control for targeted clinical applications. Experiments on CT and MRI datasets demonstrate OWT's power: it not only achieves strong performance on standard tasks like image reconstruction and segmentation, but also unlocks novel, high-impact clinical capabilities including organ-specific tumor identification, organ-level retrieval and semantic-level generation, without requiring any additional training. These findings underscore the potential of OWT as a foundational framework for semantically disentangled representation learning, offering broad scalability and a new perspective on how representations can be leveraged.

Figures (13)

• Figure 1: Objective comparison of (a) holistic embedding-based representation learning, (b) disentangled representation learning, and (c) our OWT framework. Unlike prior methods, our foundational tokenization framework enables the extraction of disentangled token groups containing semantically meaningful organ-specific representations. These token groups can be leveraged separately or in combination for downstream analyses, such as semantic-level reconstruction, segmentation, and retrieval.
• Figure 2: Overall architecture of Organ-Wise Tokenization (OWT). The input 2D or 3D images are first encoded into holistic embeddings ($X_H$). Then, the holistic tokens are disentangled and embedded in an organ-wise manner using the Organ Collector, forming Semantically Disentangled Token Groups (SDTGs, $X_G$). The SDTGs with solid lines represent randomly retained token groups ($\tilde{X}_G$). The Token Group Encoder further processes these retained groups, capturing both inter- and intra-organ relationships, and transforming them into $\tilde{X}'_G$. Finally, an Adaptive Holistic Embedding Restorer and a decoder sequentially restores the holistic representation ($X'_H$) and reconstructs the final output ($\hat{I}_T$), enabling semantic-level reconstruction.
• Figure 3: Comparison of the overall process of holistic embedding-based and token group-based reconstruction.
• Figure 4: Samples of holistic-based and semantic-based medical image reconstruction from an MRI dataset.
• Figure 5: Semantic-based organ-level retrieval of MRI images.