K-MaT: Knowledge-Anchored Manifold Transport for Cross-Modal Prompt Learning in Medical Imaging

K-MaT is a novel prompt-learning framework that enables the zero-shot transfer of large-scale biomedical vision-language models from high-end to low-end imaging modalities by anchoring prompts to clinical text and aligning their decision manifolds via Fused Gromov-Wasserstein optimal transport, thereby achieving state-of-the-art performance while mitigating catastrophic forgetting.

The Gist

Imagine you are a brilliant medical detective who has spent years studying high-resolution, 3D MRI scans and CT scans. You've learned to spot diseases with incredible precision because these images are like high-definition movies: you can see every tiny detail, texture, and shadow.

Now, imagine you are sent to a rural clinic where they only have basic X-rays and ultrasound machines. These images are like grainy, black-and-white sketches compared to your high-def movies.

The Problem:
If you try to use your "MRI brain" to read these "sketches," you get confused. You might start looking for details that don't exist in the sketch, or you might get tricked by the grainy noise. In the world of AI, this is called "catastrophic forgetting." The AI learns the high-end images so well that it forgets how to handle the low-end ones, or worse, it learns "shortcuts" (like guessing "cancer" just because the X-ray is blurry) that only work for one type of machine.

The Solution: K-MaT
The paper introduces a new method called K-MaT (Knowledge-Anchored Manifold Transport). Think of it as a special training program that teaches your AI detective how to translate its high-definition knowledge into a language that works for the sketchy, low-end machines—without ever seeing a single low-end image during training.

Here is how K-MaT works, using three simple analogies:

1. The "Universal Translator" (Prompt Factorization)

Usually, AI learns a single set of instructions (prompts) to solve a problem. K-MaT splits these instructions into two parts:

• The Core Knowledge: The universal medical facts (e.g., "a tumor looks like a mass").
• The Machine-Specific Dialect: The specific way that fact looks on an MRI vs. an X-ray.
  By separating these, the AI keeps its "medical brain" intact while learning to speak the "X-ray dialect" without losing its original wisdom.

2. The "Anchored Compass" (Knowledge Anchoring)

When an AI tries to learn a new task, it can easily drift off course, inventing its own weird rules. K-MaT uses LLMs (Large Language Models) as a "Compass."

• Imagine the AI is a ship. The LLM generates detailed text descriptions of diseases (e.g., "a malignant mass with jagged edges").
• These text descriptions act as anchors tied to the seabed.
• As the AI learns, it is physically tethered to these anchors. No matter how much the image quality changes, the AI is forced to stay close to the true, medically accurate definition of the disease. It can't drift into "shortcut land."

3. The "Rubber Sheet" (Manifold Transport)

This is the most magical part. Imagine the high-end images (MRI) and low-end images (X-ray) are two different shapes made of rubber.

• The high-end shape has a perfect, complex structure where "cancer" is always a certain distance from "healthy tissue."
• The low-end shape is squashed and distorted.
• K-MaT uses a mathematical tool called Fused Gromov-Wasserstein (FGW) transport. Think of this as a smart rubber sheet stretcher. It doesn't just stretch the low-end shape; it stretches it exactly to match the internal geometry of the high-end shape.
• It forces the "relationships" between diseases in the low-end world to mirror the relationships in the high-end world. If "Tumor A" is close to "Tumor B" in the MRI world, K-MaT ensures they stay close in the X-ray world, even if the pictures look totally different.

The Result

When the researchers tested this:

• Old AI methods would get great scores on the high-end MRI but would crash and burn (dropping to 27% accuracy) when faced with the low-end X-rays. They forgot everything.
• K-MaT kept its high-end skills and learned to handle the low-end images, achieving the best results ever recorded for this type of task.

In a nutshell:
K-MaT is like teaching a master chef (the AI) to cook a gourmet meal using only a campfire (low-end modality) by giving them a recipe book (LLM anchors) and a special map (optimal transport) that shows them exactly how to translate their fancy kitchen techniques to the outdoors, without ever letting them practice on the campfire first. It ensures the food tastes just as good, no matter the equipment.

Technical Summary

1. Problem Statement

Medical imaging models often suffer from distribution shifts when transferring from high-end modalities (e.g., CT, MRI, Mammography) to low-end, frontline modalities (e.g., X-ray, Ultrasound, Clinical photos).

• The Core Issue: Standard adaptation methods, including Vision-Language Models (VLMs) and prompt learning techniques (like CoOp), tend to collapse into modality-specific shortcuts. When optimized on high-end data, they overfit to the specific acquisition physics of that modality and fail to generalize to low-end modalities.
• Catastrophic Forgetting: In cross-modal settings, learnable prompts optimized solely on high-end data often lose the essential general textual knowledge required for shared diagnostic semantics, leading to severe performance drops (e.g., accuracy dropping from ~75% to ~27%) when applied to low-end images in a zero-shot regime.
• The Challenge: How to transfer diagnostic knowledge from high-end visual data to low-end modalities without access to low-end training images (strict zero-shot transfer).

2. Methodology: K-MaT Framework

The authors propose K-MaT (Knowledge-Anchored Manifold Transport), a prompt-learning framework built upon the frozen BiomedCLIP backbone. The method consists of three core mechanisms:

A. Factorized Prompt Parameterization

Instead of a single prompt, K-MaT introduces a factorized learnable prompt for each class $i$ and modality $m$:
$$t_{i,m} = [v^1_{i,m}, v^2_{i,m}, v^3_{i,m}, v^4_{i,m}, c_i]$$

• Class-Specific Context (CSC): Tokens specific to the pathology.
• Modality-Specific Context (MSC): Tokens specific to the imaging modality.
• This factorization allows the model to learn distinct representations for different modalities while sharing the core class knowledge, mitigating interference.

B. Space Anchoring (Semantic Constraint)

To prevent the learnable prompts from deviating into modality-specific statistics and forgetting general medical knowledge, the authors introduce a Space Anchoring constraint.

• Mechanism: An LLM (GPT-5) generates multiple clinical text descriptions ($p^j_{i,m}$) for each class. These are encoded to create fixed textual prototypes ($\bar{p}_{i,m}$).
• Loss Function ($L_{anc}$): A Euclidean distance loss minimizes the discrepancy between the learnable prompt embeddings ($w_{i,m}$) and the fixed LLM-generated prototypes.
• Goal: This acts as a semantic anchor, ensuring the learned prompts remain grounded in clinically meaningful semantics rather than just visual noise.

C. Cross-Modal Manifold Alignment via FGW

To enforce structural consistency between the high-end and low-end prompt spaces without low-end images, K-MaT uses Fused Gromov-Wasserstein (FGW) Optimal Transport.

• Concept: FGW aligns the relational geometry of the low-end prompt manifold to the high-end manifold. It treats the high-end textual embeddings as a fixed reference.
• Mechanism: It computes intra-modal distance matrices ($D_H, D_L$) and a cross-modal cost matrix ($M$). The loss ($L_{fgw}$) seeks an optimal coupling matrix $\Gamma$ that minimizes the difference between the feature distances and the structural distances.
• Goal: This forces the low-end prompt manifold to strictly mirror the relational structure of the high-end manifold, effectively "transporting" the decision boundaries from the high-end space to the low-end space.

D. Overall Objective

The total loss function combines standard cross-entropy (on high-end data), space anchoring, and manifold alignment:
$$L = L_{ce} + \lambda_{anc}L_{anc} + \lambda_{fgw}L_{fgw}$$
During inference, the model uses the learned low-end embeddings to classify unseen low-end images, with all backbone parameters frozen.

3. Key Contributions

1. Strict Zero-Shot Asymmetric Transfer: A novel strategy that transfers knowledge from high-end to low-end modalities using only high-end visual data and LLM-generated text, eliminating the need for low-end training images.
2. Mitigation of Catastrophic Forgetting: The introduction of Space Anchoring using LLM-generated prototypes prevents prompts from collapsing into modality-specific shortcuts.
3. FGW-Based Manifold Alignment: A structural regularizer that aligns the low-end prompt manifold to the high-end space, preserving shared diagnostic semantics and relational geometry.
4. Factorized Prompt Design: The use of CSC and MSC tokens effectively balances class-specific and modality-specific learning.

4. Experimental Results

The method was evaluated on four cross-modal benchmarks:

1. Dermoscopy $\to$ Clinical Images (Skin I)
2. Dermoscopy $\to$ Clinical Images (Skin II)
3. Mammography $\to$ Ultrasound (Breast)
4. CT $\to$ Chest X-ray (Chest)

Key Findings:

• State-of-the-Art Performance: K-MaT achieved an average Harmonic Mean (H) Accuracy of 44.1% and Macro-F1 of 36.2%, outperforming the previous best baseline (BiomedCoOp) by 2.1% in H-Accuracy.
• Robustness in Breast Imaging: In the challenging Mammography-to-Ultrasound task, standard methods like CoOp suffered catastrophic forgetting (dropping to 27.0% accuracy on low-end). K-MaT preserved robust performance, achieving 38.4% accuracy on the low-end modality and an H-score of 50.3%.
• Ablation Studies:
  • Removing FGW ($L_{fgw}$) caused the low-end manifold to fail to preserve structural relationships, leading to poor discrimination.
  • Removing Space Anchoring ($L_{anc}$) led to semantic divergence.
  • The combination of all components yielded the highest performance, with a 13.75% relative improvement in H-F1 over the baseline.

5. Significance and Limitations

Significance:

• K-MaT provides a highly effective route for zero-shot cross-modal deployment of medical VLMs, addressing a critical bottleneck where high-quality medical data (CT/MRI) is abundant but frontline data (X-ray/Ultrasound) is scarce or unavailable for training.
• It demonstrates that structural alignment via optimal transport is superior to simple feature alignment for preserving diagnostic semantics across modalities with vastly different acquisition physics.

Limitations:

• Absolute Low-End Performance: While K-MaT significantly improves over baselines, the absolute accuracy on low-end modalities still shows limited improvement over the raw zero-shot BiomedCLIP baseline in some cases.
• Visual-Textual Gaps: The framework relies heavily on text anchoring. If the visual discrepancy between modalities is too severe, purely text-anchored alignment may not fully bridge the gap without additional visual signals.

Conclusion:
K-MaT represents a significant advancement in medical AI by enabling the transfer of diagnostic intelligence from high-end to low-end imaging modalities without requiring target domain data, thereby enhancing the accessibility and reliability of AI diagnostics in resource-constrained settings.