Medical Image Captioning: ImageHRM and ImageTRM Architectures
The research focuses on developing advanced models for generating accurate descriptions of medical images, using hierarchical and recursive structures.
ImageHRM (Hierarchical Reasoning Model)
- Integration: Connects a vision encoder directly into the model’s HRM reasoning process.
- System modules:
- Vision Module: Extracts image features.
- L-Module: Performs local computations and reasoning.
- M-Module: Integrates the results and updates the latent memory through a recurrent loop.
- H-Module: Generates the high-level decision or prediction toward the output head.
- Semantic Clustering: Includes an intermediate loop to enable the transition between high-level (HL) and mid-high-level (HML) representations.
- Tested Backbones: ResNet18, Swin Transformer, and FuseLIP. The pre-trained visual modules are kept "frozen" to preserve visual knowledge, while the reasoning core learns the structuring of clinical text.
ImageTRM (Tiny Recursive Model)
- Architecture: A 2-layer Transformer that runs recursively $N$ times with shared weights.
- Efficiency: Reduced model footprint, with only approximately 7 million parameters.
- Stability: Uses Exponential Moving Average (EMA) weight smoothing to maintain stability of deep recursion during training.
- SOTA Performance: Achieves state-of-the-art results compared to ImageHRM variants, reaching a CIDEr score of 0.449 and a ROUGE-L score of 0.199 using a Swin backbone.
Transferable Knowledge Base Augmented Language Models (KBLaM)
The portability of knowledge base-augmented language models across different medical domains and ontologies is explored.
Methodology and Observations
- Training: The KBLaM model is trained exclusively on the DOID dataset (BioML).
- Cross-Domain Inference: The model is applied without retraining to other ontologies and medical domains, such as NCIT, ORDO, or OMIM.
- Transferability: It was observed that the attention mechanism trained on one ontology performs effectively on different ontologies, demonstrating that the model learns attention structure independently of the dataset.
Inference Results (Accuracy)
Performance was evaluated across different ontologies with varying knowledge base (KB) sizes.
