Representing Part-Whole Hierarchies in Foundation Models by Learning Localizability, Composability, and Decomposability from Anatomy via Self-Supervision

TL;DR

This work tackles the lack of explicit part–whole hierarchy modeling in foundation models for medical imaging. It proposes Adam--v2, a self supervised framework with three dedicated branches localizability composability and decomposability that learn a hierarchy of anatomical embeddings from unlabeled images in a coarse to fine fashion. Across zero shot few shot and full transfer settings on chest X-ray and fundus tasks, Adam--v2 achieves state-of-the-art or near state of the art results including a ChestX-ray14 mAUC of 83.4 and demonstrates emergent anatomy understanding and interpretable embedding properties. The results, ablations and qualitative analyses show that explicit anatomy aware representations improve generalization robustness and enable downstream utility with limited annotations, with code and pretrained models released as Eden.

Abstract

Humans effortlessly interpret images by parsing them into part-whole hierarchies; deep learning excels in learning multi-level feature spaces, but they often lack explicit coding of part-whole relations, a prominent property of medical imaging. To overcome this limitation, we introduce Adam-v2, a new self-supervised learning framework extending Adam [79] by explicitly incorporating part-whole hierarchies into its learning objectives through three key branches: (1) Localizability, acquiring discriminative representations to distinguish different anatomical patterns; (2) Composability, learning each anatomical structure in a parts-to-whole manner; and (3) Decomposability, comprehending each anatomical structure in a whole-to-parts manner. Experimental results across 10 tasks, compared to 11 baselines in zero-shot, few-shot transfer, and full fine-tuning settings, showcase Adam-v2's superior performance over large-scale medical models and existing SSL methods across diverse downstream tasks. The higher generality and robustness of Adam-v2's representations originate from its explicit construction of hierarchies for distinct anatomical structures from unlabeled medical images. Adam-v2 preserves a semantic balance of anatomical diversity and harmony in its embedding, yielding representations that are both generic and semantically meaningful, yet overlooked in existing SSL methods. All code and pretrained models are available at https://github.com/JLiangLab/Eden.

Figures (10)

• Figure 1: Human perception effortlessly organizes objects into hierarchies to understand their part-whole relationships in images. Taking lungs as an example in (a), even a non-radiologist can form a hierarchy of the right and left lungs, whereas a radiologist can further "see’’ the lobes in sub-hierarchies. To emulate this ability, we introduce a self-supervised learning framework that explicitly learns to encode inherent part-whole hierarchies within medical images into an embedding space, leading to the development of a powerful model (Adam--v2) that is foundational to medical imaging. Adam-v2 can transform each pixel in medical images (e.g., chest radiographs in (b)) into semantically meaningful embeddings (Eve--v2), forming multiple "echo chambers’’ (produced via co-segmentation amir2021deepzhou2023learning)---different anatomical structures are associated with distinct embeddings, and the same anatomical structures have (nearly) identical embeddings across patients.
• Figure 2: Adam--v2 learns hierarchical representations in a coarse-to-fine-manner via three branches: localizability, composability, and decomposability. Given an anchor whole $w$ randomly sampled from image $I$, the localizability branch augment and process $w$ and its multi-scale views, and enforce consistency between their embeddings, yielding distinct features for different anatomical structures. The composability branch decomposes $w$ into a set of parts, and enforces consistency between the embedding of $w$ and the aggregated embeddings of its parts, encoding part-whole relations. The decomposability branch decomposes the embedding of $w$ to acquire the embeddings of its constituent parts, and enforce consistency between the embeddings of parts and their decomposed counterparts, capturing whole-part relations.
• Figure 3: Adam--v2 learns localizability of anatomical structures, providing discriminative features for different landmarks. Same-colored points are instances of the same landmark across images.
• Figure 4: Adam--v2 balances diversity and harmony in embeddings of similar anatomical structures across patients and scales.
• Figure 5: Adam--v2's embeddings (Eve--v2) encode part-whole relations of anatomical structures.