UniMod: Enhancing Multi-Modal Medical Diagnosis through Cross-Modality and Within-Modality Alignment

UniMod: Enhancing Multi-Modal Medical Diagnosis through Cross-Modality and Within-Modality Alignment

Problem. Multi-modal learning combining medical images and clinical text is promising for disease diagnosis. However, standard multi-modal training leads to shortcut learning: models exploit the easier modality (e.g., diagnostic cues in text) while neglecting harder-to-learn features (e.g., subtle visual patterns). The paper states: "directly training multi-modal models often leads to shortcut learning [12, 16], which means the model exploits whichever modality provides the easiest path to correct predictions, while neglecting the other. This phenomenon is particularly prevalent in medical diagnosis, where clinical text frequently contains explicit diagnostic cues (e.g., 'suspected glaucoma,' 'abnormal finding'), whereas images require learning subtle visual patterns (e.g., optic disc cupping ratio, retinal nerve fiber layer defects)." This causes two consequences: the model does not fully exploit each modality's diagnostic potential, and performance degrades sharply when text is ambiguous, incomplete, or unavailable.

Proposed method. The authors propose UniMod, a framework that mitigates shortcut learning by requiring each modality to predict the diagnosis on its own. The key insight is: if a model can make accurate predictions using each modality independently, it must have learned genuine diagnostic features from that modality. UniMod consists of four complementary components:

1. Independent Modality Feature Extraction (Section 4.1): The framework supervises image-only, text-only, and multi-modal classification simultaneously. Classification losses are applied to all three prediction paths: L cls img = L CE(ŷ img, y), L cls txt = L CE(ŷ txt, y), and L cls mm = L CE(ŷ mm, y), combined as L cls = L cls img + L cls txt + L cls mm. By supervising each modality independently, we remove the lowest-loss route to a shortcut solution.

2. Modality-separated attention (Section 3.2): A custom attention mask enforces modality isolation. "Image tokens V can only attend to other image tokens, and text tokens (T, Q) can only attend to other text tokens. As a result, neither modality can directly access the other, and only the final prediction token p is permitted to attend to tokens from both modalities." This ensures unimodal predictions are well defined.

3. Cross-modality alignment (Section 4.2): An MSE loss pulls image and text representations of the same patient closer: L cross = (1/B) Σ i∈B z img i − z txt i22. This loss pulls the image and text representations of the same patient closer, encouraging both modalities to encode similar diagnostic information in a shared semantic space.

4. Within-modality alignment (Section 4.3): Supervised contrastive learning clusters same-diagnosis samples: L within = L within img + L within txt, using the SupCon loss with temperature τ = 0.07. This loss pulls samples with the same label together while pushing samples with different labels apart.

The total objective is L total = L cls + λ2L cross + λ3L within, with λ1 = 1 fixed and λ2, λ3 adaptively adjusted via GradNorm. The model uses LoRA for parameter-efficient fine-tuning (rank r = 32, α = 64) on InternVL2.5-8B.

Experiments. UniMod is evaluated on two datasets: Harvard-Glaucoma (fundus photographs paired with clinical notes for glaucoma detection) and CheXpert Plus (chest X-rays paired with radiology reports for pleural effusion detection). Text cleaning removes diagnosis-related keywords to prevent label leakage.

Main results (Table 1):

• On Harvard-Glaucoma: UniMod achieves 0.850 AUC, 0.768 accuracy, 0.786 F1, outperforming OGM-GE (0.835 AUC) and Gradient Blending (0.837 AUC) by 1.6–1.8% in AUC and 8.0–14.6% in F1.

• On CheXpert Plus: UniMod achieves 0.966 AUC, 0.919 accuracy, 0.920 F1, surpassing OGM-GE and G-Blend by over 5%, and surpassing even the Text-only baseline (0.961 AUC).

• Multimodal LoRA performs poorly (0.716 AUC on Harvard-Glaucoma, 0.908 on CheXpert Plus) due to shortcut learning.

Multi-label generalization (Table 5): UniMod extends to 5-class multi-label diagnosis on CheXpert Plus without architectural change, achieving 0.760 mAUC, 0.473 mAP, and 0.480 mF1 versus 0.664, 0.340, and 0.306 for CGGM, a gain of +0.097 mAUC and +0.174 mF1.

Shortcut reduction (Section 5.3): The paper directly measures the text shortcut. Multimodal LoRA drives CE mm to 0.026 while CE img remains at 0.774, i.e., the fused loss is satisfied without the image encoder becoming diagnostic. UniMod keeps all three terms balanced (CE mm = 0.415, CE img = 0.612, CE txt = 0.412 on Harvard-Glaucoma). Robustness tests show UniMod maintains much higher performance with missing modalities. On CheXpert Plus, Multimodal LoRA opens a 0.58 gap between text (0.91) and image (0.33) AUC, while UniMod reduces the gap to 0.06.

Ablation studies (Table 3): Removing within-modality alignment drops AUC from 0.850 to 0.838 (Harvard-Glaucoma) and 0.966 to 0.964 (CheXpert Plus). Removing cross-modality alignment drops AUC to 0.835 and 0.964 respectively. Independent feature extraction provides the largest improvement.

Case studies (Figure 4): UniMod correctly classifies cases where Multimodal LoRA fails, particularly when clinical text is ambiguous or indirect. For example, in a CheXpert Plus case where the report describes indirect findings without naming pleural effusion, Multimodal LoRA returns 0.487 while UniMod detects the effusion at 0.848.

Conclusion. The paper states: "We presented UniMod, a framework that mitigates shortcut learning in multi-modal medical diagnosis by requiring each modality to predict the diagnosis on its own, combining independent unimodal supervision with cross-modality and within-modality alignment under GradNorm weighting. On Harvard-Glaucoma and CheXpert Plus it improves over standard multi-modal fine-tuning and over modality-balancing baselines, and it transfers to 5-class multi-label diagnosis without architectural change. Our analysis indicates that gradient modulation alone does not remove the shortcut: explicit unimodal supervision is what drives each modality to develop diagnostic capability."

Improvements to AI systems based on UniMod:

1. Add independent unimodal supervision to multi-modal training. Instead of only training a fused model, add separate classification losses for image-only and text-only prediction paths. This forces each modality to learn diagnostic features on its own, preventing the model from ignoring one modality when the other is easier.

2. Enforce modality-separated attention in transformer backbones. Use an attention mask that blocks cross-modality token interactions until the final prediction token, so unimodal representations remain clean and cannot leak information from the other modality. This makes unimodal predictions well-defined and trainable.

3. Apply cross-modality alignment via representation matching. Pull image and text embeddings of the same patient closer using MSE loss, encouraging both modalities to encode shared diagnostic semantics rather than modality-specific shortcuts.

4. Add within-modality supervised contrastive learning. Cluster same-diagnosis samples and push apart different-diagnosis samples within each modality separately, improving feature discriminability and robustness to missing or ambiguous inputs.

5. Use adaptive loss weighting (e.g., GradNorm) instead of fixed weights. Dynamically balance classification, cross-modality, and within-modality losses to prevent any single objective from dominating, especially when dataset difficulty varies.

6. Monitor per-modality cross-entropy separately during training. Track CE img, CE txt, and CE mm to detect shortcut learning early (e.g., if CE mm is low but CE img is high, the image path is being neglected). Use this signal to adjust supervision or halt training.

7. Make the system robust to missing or incomplete modalities. Because each modality is independently supervised and aligned, the improved system can still produce reliable predictions when text is ambiguous, incomplete, or entirely absent—critical for real-world clinical deployment.

What the improved AI system can do:

• Achieve higher diagnostic accuracy (e.g., +1.6–1.8% AUC on glaucoma, +5% on pleural effusion) compared to standard multi-modal fine-tuning and modality-balancing baselines.

• Maintain balanced performance across modalities, reducing the gap between text-only and image-only AUC from 0.58 to 0.06.

• Generalize to multi-label diagnosis (e.g., 5-class chest X-ray findings) without architectural changes, improving mean AUC by +0.097 and F1 by +0.174.

• Correctly diagnose cases where clinical text is indirect or does not name the condition (e.g., detecting pleural effusion from indirect radiology report language).

• Provide reliable predictions even when one modality is missing or degraded, making it suitable for real-world clinical workflows where imaging or notes may be incomplete.

Multi-modal learning combining medical images and clinical text is promising for disease diagnosis. However, standard multi-modal training leads to shortcut learning: models exploit the easier modality (e.g., diagnostic cues in text) while neglecting harder-to-learn features (e.g., subtle visual patterns). We propose UniMod, a framework that mitigates shortcut learning by requiring each modality to predict the diagnosis on its own. It supervises image-only, text-only, and multi-modal classification simultaneously, so each modality must extract diagnostic features. We add cross-modality alignment for knowledge transfer and within-modality supervised contrastive alignment over same-diagnosis patients. On Harvard-Glaucoma, UniMod reaches 0.850 AUC, outperforming OGM-GE and Gradient Blending by 1.6-1.8%; on CheXpert Plus, it reaches 0.966 AUC, surpassing them by over 5%. UniMod also extends to 5-class multi-label diagnosis without architectural change, improving mean AUC by 0.097 over CGGM.