HYDEN: Hyperbolic Density Representations for Medical Images and Reports

In our model, the features $[\hat{f^{v}},f^{t}]$ are derived from respective image and text encoders. For text data, we employ BioClinicalBERT (Alsentzer et al., 2019a), a model that has been pre-trained on the MIMIC III dataset (Shen et al., 2016), to generate token-level embeddings. Consistent with practices outlined in (Cheng et al., 2023b), the output of the $[CLS]$ token is used as the medical text feature $f^{t}$, encapsulating the overall semantic content of the input text.

For image encoding, we utilize the widely-used Vision Transformer (ViT) architecture (Mu et al., 2022). We assume that $\hat{f^{v}}=(\hat{f^{v}_{0}},\hat{f^{v}_{1}},\dots,\hat{f^{v}_{n}})$ captures the outputs from the image encoder. Recognizing that pathological symptoms often occupy only a portion of a medical image, relying solely on global representations may not adequately capture essential local semantic features. Thus, similar to approaches in (Huang et al., 2021; Cheng et al., 2023b; Müller et al., 2022b), we enhance the global features by integrating text-aware image local representations.

Specifically, we implement a Self-attention module (Vaswani et al., 2017), widely used in cross-modal feature extraction. In this setup, $\hat{f^{v}}$ acts as both the keys ($K$) and values ($V$), while the text embedding $f^{t}$ functions as the query ($Q$). This configuration allows us to derive a text-aware image local representation, denoted as $\bar{\hat{f^{v}}}$. By synthesizing both global and local representations, we achieve the final image feature $f^{v}=\hat{f^{v}}+\bar{\hat{f^{v}}}$, effectively capturing both the broad and nuanced semantic details present in medical images.

Our objective is to transform image-text features into density representations within hyperbolic space. Previous studies such as (Nagano et al., 2019) and (Mathieu et al., 2019) proposed methods like the pseudo-hyperbolic Gaussian distribution based on the Lorentz manifold. Due to the computational demands and numerical instabilities of the Poincaré-disk model, we opt for the more stable pseudo-hyperbolic Gaussian distribution for our hyperbolic density embedding. The tangent space $T_{\mu}\mathbb{H}^{n}$ of hyperbolic space $\mathbb{H}^{n}$ is a Euclidean space, and in $T_{\mu_{0}}\mathbb{H}^{n}$, vectors $\mathrm{v}$ satisfy $\mathrm{v}=\{v_{0},v_{1},...,v_{n}\}\in\mathbb{R}^{n+1}$ where $v_{0}=0$, aligning with the dimensional properties.

To begin, we introduce separate deep nonlinear network blocks, $B_{density}$, for processing image and text features independently. These blocks do not share parameters, ensuring distinct representations for each modality. As in Figure 2, for text features, $\hat{\mu^{t}}$ and $\acute{\beta^{t}}$ are the outputs of $B_{density}(f^{t})$.

Instead of generating covariance matrices directly, which can introduce numerical instability, we use matrices based on diagonal or spherical assumptions. These are known for their computational efficiency and effectiveness in embedding tasks, particularly in the context of word distribution embedding where spherical covariance matrices have been shown to better model distributional partial order relationships (Vilnis and McCallum, 2014). We thus employ a covariance matrix based on the spherical assumption: $\Sigma^{t}=\acute{\beta^{t}}\cdot\mathrm{I}\in\mathbb{R}^{(n+1)\times(n+1)}$.

To ensure that our covariance matrix is positively definite, necessary for the stability of the pseudo-hyperbolic Gaussian distribution, we modify $\beta^{t}$ using the expression $\beta^{t}=\text{exp}(\acute{\beta^{t}})$ referring to solution in VAE(Kingma and Welling, 2019). This adjustment is crucial for maintaining the mathematical integrity of our model when dealing with real-world data. For the embedding vector $\hat{\mu^{t}}\in\mathbb{R}^{n}$, our aim is to project this vector onto hyperboloid space, which is achieved by mapping it through the exponential function as detailed in Equation 2.

The vector $\mu^{t}_{tan}=[0,\hat{\mu^{t}}]$ resides in $\mathbb{R}^{n+1}$ and belongs to the tangent space $T_{\mu_{0}}\mathbb{H}^{n}$ at the origin of the hyperboloid, $\mathcal{O}$. The norm $||\mu^{t}_{tan}||_{\mathcal{L}}$, which equals $||\hat{\mu^{t}}||_{2}$, ensures that the mapping preserves the distances inherent to the model’s geometric structure. We apply the exponential map to $\mu^{t}_{tan}$, decomposing the transformation into two parts:

Upon applying the exponential map, we derive the expectation of the hyperbolic density representation:

This projection results in the hyperbolic density representation $\mathcal{G}^{t}(\mu^{t},\beta^{t}\cdot\mathrm{I})$. Following a similar procedure, we also derive $\mathcal{G}^{v}(\mu^{v},\beta^{v}\cdot\mathrm{I})$ for the image features, thereby ensuring a uniform approach to handling different modalities within our framework.

Traditional point vector embedding often utilizes entailment angle constraints to define relationships between entities (Desai et al., 2023). However, when dealing with probability densities, the notion of partial order can be more complexly captured through the concept of encapsulation. Specifically, a density $f$ is considered more specific than another density $g$ if $f$ is entirely encompassed by $g$, formally expressed as $f\preceq g\Leftrightarrow\{x:f(x)>\eta\}\subseteq\{x:g(x)>\eta\}$, for any $\eta\geq 0$, where $\eta$ indicates the degree of encapsulation necessary for one distribution to entail another.

Imposing such partial order constraints on distributions poses significant challenges. Drawing inspiration from Athiwaratkun and Wilson (2018), we employ asymmetric divergence measures between probability densities to address this. We introduce a simple penalty function, $d_{\gamma}(f,g)=\max(0,D(f\parallel g)-\gamma)$, which serves as a violation penalty rather than as a strict constraint of encapsulation. Here, $D(\parallel)$ represents the divergence measure used to quantify the extent of difference between distributions, and $\gamma$ is a threshold defining the acceptable range of difference.

Among the choices for divergence measures, $\alpha$-divergence provides a more flexible and generalized asymmetric measure (Renyi, 1961), allowing for adjustments in the zero-force penalty. This flexibility means that higher $\alpha$ values can enforce stricter encapsulation conditions $f\preceq g$. The general form of $\alpha$-divergence, for $\alpha\neq 0,1$, is given by:

This equation not only quantifies the differences between distributions but also facilitates a deeper understanding of the encapsulation relationships critical for effective density embedding.

We observe that as $\alpha$ approaches 0 or 1, it governs the degree of zero forcing, where minimizing $D_{\alpha}(f\parallel g)$ for high $\alpha$ values results in $f$ becoming more concentrated in regions of $g$ with high density. Conversely, for low $\alpha$ values, $f$ tends to be mass-covering, encompassing regions of $g$ even including those with low density. Notably, there exists a mathematical relationship between KL divergence and $\alpha$-divergence, as indicated by: $\lim_{\alpha\to 1}D_{\alpha}(f\parallel g)=D_{KL}(f\parallel g)$ and $\lim_{\alpha\to 0}D_{\alpha}(f\parallel g)=D_{KL}(g\parallel f)$ (Pardo, 2006). Therefore, in our model, we opt for the more flexible and robust $\alpha$-divergence as our metric.

For image-text embedded density $\mathcal{G}^{v}(\mu^{v},\beta^{v}\cdot\mathrm{I})$ and $\mathcal{G}^{t}(\mu^{t},\beta^{t}\cdot\mathrm{I})$, the encapsulation loss can be expressed as follows:

Let the batch sample $\mathbb{B}=\{\mathbb{B}^{P},\mathbb{B}^{N}\}$, where $\mathbb{B}^{P}$ denotes the positive image-text sample set, and $\mathbb{B}^{N}$ represents the negative set. We define the encapsulation loss function as follows:

For the positive samples, a definite partial order relationship exists, enabling the direct application of the density penalty $d_{\gamma}()$. For the negative samples, we enforce the penalty to exceed a margin $m$ due to the absence of an order relationship.

Our goal is to enhance the similarity of semantic distributions between image-text pairs. Therefore, we also employ the classic CLIP contrastive solution (Radford et al., 2021) to compute the geodesic distance between the expectation values of image and text in hyperbolic densities as defined in Equation 1, applying Softmax normalization. We define $\mathfrak{L}_{con}$ as the contrastive loss, which is computed as an average of the contrastive losses from both image and text perspectives.

Datasets: We train our alignment model using the MIMIC-CXR v2 dataset (Johnson et al., 2019), comprising over 227,000 studies of paired image-report data sourced from 65,379 patients undergoing various scans. Each study may contain one or two images, representing different scan views, resulting in a total of 377,110 images. During training, we perform random cropping, flipping, rotation, and other data augmentation techniques on the images, while also resizing them to a [224,224] dimension. Additionally, for the text data, we augment the reports by randomly adding medical entity prefixes to enhance semantic information, such as ’event_list: report’.

Settings: We employ ViT-B (Mu et al., 2022) with a patch size of 16 as the image encoder, as it has demonstrated competitive performance in hyperbolic space (Desai et al., 2023). Our initialization strategy for image/text encoders follows a similar style to MERU, with the exception of utilizing ClinicalBERT (Alsentzer et al., 2019b) as the pre-trained text encoder, which has been pre-trained on large-scale medical text data. For HYDEN, we initialize the learnable curvature parameter $c$ to 1.0 and clamp it within the range of [0.1, 10.0] to prevent training instability. All experiments were conducted using two NVIDIA A40 GPU and the PyTorch framework

Optimization: We adopt the AdamW optimizer with a weight decay of 0.2 and $(\beta_{1},\beta_{2})=(0.9,0.98)$. Weight decay is disabled for all gains, biases, and learnable scalars. Models are trained for 13,000 iterations with a batch size of 256. The maximum learning rate is set to $1\times 10^{-5}$, linearly increased for the first 500 iterations, followed by cosine decay to zero. We leverage mixed precision to expedite training, except when computing exponential maps and losses, where FP32 precision is used for numerical stability.

We evaluate all baselines and HYDEN on three categories of zero-shot downstream tasks, classification, text-image retrieval and image-image retrieval. We use three public datasets for the evaluation, where both RSNA Pneumonia (Shih et al., 2019) and SIIM-ACR Pneumothorax (Kaggle, 2019) are used for binary classification, ChestXray14(Wang et al., 2017a) is used for multi-label classification, text-image retrieval and image-image retrieval. For the two binary classification tasks, we report the Area Under the Curve (AUC) and F1 score; for the multi-label task, we provide the Micro-AUC and Micro-F1. For the retrieval task, Top-k Precision (abbreviated as Prec@k) and Tok-k Normalized Discounted Cumulative Gain (abbreviated as NDCG@k) are used to evaluate the retrieval performance. Refer Appendix B for details about our evaluation tasks and datasets.

Zero-shot Image Classification Table 1 presents the performance of the baselines and HYDEN across three classification datasets. The results indicate that HYDEN consistently demonstrates robust transfer classification performance, both in binary classification tasks and multi-label classification task. Compared to CLIP, both MERU and HyperMed achieved improved accuracy. This suggests that using hyperbolic space for text-image representations, especially for medical data characterized by a visual semantic hierarchy, is more effective. Relative to MERU, HYDEN achieved the highest accuracy across almost all of metrics, highlighting the advantages of density embedding-based representation methods over point vector embedding, particularly in addressing the challenges of semantic uncertainty.

Zero-shot Retrieval Table 2 displays the performance of two baseline models and HYDEN in "image-to-text" and "image-to-image" retrieval tasks. The results demonstrate that representation learning in hyperbolic space mostly outperforms that in Euclidean space; among the methods, HYDEN exhibits the best retrieval performance. Furthermore, we observed a significant enhancement in the ranking quality of HYDEN’s retrieval results compared to the two baseline methods. We hypothesize that this improvement is linked to the method of density embedding. Similar to findings in the recommendation systems domain (Dos Santos et al., 2017), unlike point vector embeddings, density embeddings enable better handling of uncertainties, information sparsity, ambiguity, and even contradictions, which are common challenges in medical image-text data.

Ablation Studies In this section, we examine the impact of different design choices using HYDEN. Specifically, we trained three ablation models with default hyperparameters, and the results are presented in Table 3. From Table 3, we observe that: (1) Using $\alpha$-divergence in the loss function instead of KL divergence better aligns with the encapsulation’s partial order properties of text-image distribution embeddings. The experimental results also indicate that replacing $\alpha$-divergence with KL divergence leads to performance degradation across all tasks. (2) Omitting the encapsulation loss, i.e., not using $\mathfrak{L}_{order}$ as defined in Equation 8 and relying primarily on $\mathfrak{L}_{con}$, results in performance degradation across all tasks. This is because not using encapsulation loss implies that the prior partial order of text and image cannot be utilized in hyperbolic space, thus losing the benefits introduced by hyperbolic geometry. (3) The model experiences a performance drop across all tasks when not using text-aware image representation. This is primarily due to the nature of medical image-text features. As discussed in the Introduction, most regions in medical images may differ in texture and morphology but not in clinical significance, while actual pathological changes are localized. The results also show that enhancing text-aware local features is meaningful for medical image-text alignment.

Parameters Impacts In this section, we focus on adjusting the crucial predefined parameter $\alpha$ within the HyperMed model to observe its impact on the results. The experimental outcomes are displayed in Figure 4. The table reveals that different $\alpha$ values directly influence the evaluation outcomes. The performance does not continually increase with larger $\alpha$ values; instead, it exhibits a quasi-convex distribution, initially increasing and then decreasing, which may be related to the distribution of the training data. In the experiments mentioned above, we employ a hyperparameter setting of $\alpha=0.7$.

In this section, we explore the trained models to deduce the characteristics of the model in capturing the visual semantic hierarchy structure. The concept of ’Embedding distances from [ROOT]’ was introduced by Desai et al. (2023) to depict the generality differences between text and image embeddings in hyperbolic space. This concept highlights that in a representation space that effectively captures the visual semantic hierarchy, text embeddings are typically more general than image embeddings and, therefore, should be closer to the root node [ROOT].

Here, we visualize the differences in distance distributions between text and image embeddings. Given that our approach utilizes distribution embeddings, we specifically visualize the expectations of the distance distributions of text and image density embeddings. Figure 3 demonstrates that the distribution differences generated by our model lie between those produced by MERU and CLIP, with some overlapping distribution areas. This suggests that our model is capable of capturing the visual semantic hierarchy. Compared to the diversity of natural text, medical image-text data is relatively uniform, and the introduction of distributions diminishes the effect of prior entailments, explaining why our model achieves significantly higher accuracy and ranking quality in both text-image and image-image retrieval tasks.