MIRNet: Integrating Constrained Graph-Based Reasoning with Pre-training for Diagnostic Medical Imaging

AI summary

Overview

Research area: Computer vision for diagnostic medical imaging, specifically automated tongue image diagnosis (traditional Chinese medicine), combining self-supervised visual pretraining with graph neural network reasoning and clinical constraint optimization.

Technical level: Advanced. The paper assumes familiarity with masked autoencoders, Vision Transformers, graph attention networks, multi-label classification metrics, and Lagrange-relaxed constrained optimization.

Scope: The paper introduces MIRNet, an end-to-end framework for multi-label tongue diagnosis, together with a new expert-annotated benchmark (TongueAtlas-4K), and reports benchmark comparisons and ablation studies.

What This Paper Is About

Automated reading of tongue images is hard because annotated medical data is scarce, diagnostic labels are severely imbalanced, labels are statistically correlated (for example, "pale tongue" often co-occurs with "white tongue coating"), and unconstrained models can output clinically impossible combinations (such as a tongue being both "thin" and "enlarged"). The paper's goal is to build a single framework that learns visual features without labels, models label dependencies through an expert-informed graph, and penalizes clinically implausible predictions during training, then evaluate it against existing tongue-diagnosis and general classification baselines.

Key Contributions

1. MIRNet framework. An end-to-end architecture that couples self-supervised masked autoencoder visual pretraining with constrained graph-based reasoning, intended to address annotation scarcity while modeling diagnostic dependencies through clinical knowledge graphs.
2. TongueAtlas-4K benchmark. Described by the authors as the largest publicly available expert-curated benchmark for tongue analysis: 4,000 images annotated with 22 clinically validated labels spanning tongue color, tongue shape, property of tongue coating, and color of tongue coating (released via a Zenodo DOI).
3. Differentiable clinical constraint engine. A constraint-aware optimization scheme using KL divergence and domain-driven regularization losses to encode medical knowledge (such as physiological incompatibilities) as soft constraints, combined with asymmetric loss (ASL) to mitigate label imbalance.
4. Reported state-of-the-art results. The authors report that MIRNet and MIRNet-Boosting outperform all baselines across every reported metric, with Macro Recall improved by 77.8% and Macro-F1 by 33.2% over the strongest competing method.

Main Findings

• Overall benchmark performance: On TongueAtlas-4K, MIRNet-Boosting reports Macro-F1 0.537061, Macro Recall 0.655388, and Macro PR-AUC 0.543415, which the authors describe as state of the art across all baselines. MIRNet without boosting reports Example-F1 0.680389, Micro-F1 0.683048, Macro-F1 0.525425, Macro Precision 0.507837, Macro Recall 0.599019, and Macro PR-AUC 0.527103.
• Baseline comparison: LGAN reports Macro-F1 0.397091 and Macro Recall 0.368678; Faster R-CNN 0.380543 and 0.338691; DenseNet121 0.403075 and 0.363772; C-GMVAE 0.346378 and 0.304918; YOLO12-CLS 0.290485 and 0.275153; IFRCNet 0.245823 and 0.245877. Every experiment was run five times and the mean is reported.
• Improvement margins: MIRNet-Boosting improves Macro Recall by 77.8% and Macro-F1 by 33.2% over the strongest competing method. MIRNet alone reports 62.5% higher Macro Recall and 30.4% higher Macro-F1 than the strongest baseline.
• Boosting trades some metrics for recall: MIRNet-Boosting has a higher Macro-F1 (0.537061 vs 0.525425), Macro Recall (0.655388 vs 0.599019), and Macro PR-AUC (0.543415 vs 0.527103) than MIRNet, but lower Example-F1 (0.674805 vs 0.680389), Micro-F1 (0.677620 vs 0.683048), and Macro Precision (0.499404 vs 0.507837).
• Rare-label behavior: The paper reports that all baselines show F1 < 0.25 on dark-red tongue (2.15% prevalence) and that gray-black coating (3.35% prevalence) is also poorly handled; MIRNet-Boosting reportedly raises these to 0.68 and 0.71 F1 respectively.
• Dimension-level averages: MIRNet maintains average F1 of 0.81 for tongue color, 0.77 for tongue shape, 0.76 for coating property, and 0.84 for coating color, compared with baseline averages of 0.59, 0.43, 0.51, and 0.68.
• Missed detections: Baseline maximum misses in the tongue shape dimension reach 209 cases; MIRNet-Boosting reduces this to 14 misses, a 93.3% reduction. Tongue color misses fall from a baseline range of 36-98 cases to 18 cases in MIRNet-Boosting, while coating property misses approach zero in MIRNet variants versus 3-33 in baselines.
• Ablation - constraints (MIRNet-C): Removing clinical knowledge integration reduced Example-F1 and Micro-F1 by 3.2% each and Macro-F1 by 4.4%.
• Ablation - graph decoder (MIRNet-G): Replacing the GAT with an MLP classifier cost 3.2% Macro-F1, an 8.1% recall drop, and a 3.1% precision decline, the steepest precision decline among the ablated variants.
• Ablation - pretraining (MIRNet-P): Skipping MAE pretraining caused the largest degradation, with Macro-F1 falling 23.0% and Macro Recall collapsing 29.0%.
• Pretraining data count is reported inconsistently: The Datasets section states that 15,905 unlabeled images were curated for large-scale pretraining, while the Model Architecture section states the ViT backbone was pretrained on 19,505 unlabeled images. The paper does not reconcile these two figures.
• Label distribution skew: The most frequent labels include white tongue coating (78.38%), thin tongue coating (67.58%), tongue with teeth marks (53.67%), and light-red tongue (52.80%); the rarest include dark-red tongue (2.15%), gray-black tongue coating (3.35%), and peeled tongue coating (3.52%).

Methodology in Plain English

The framework has four connected parts.

1. Learn vision without labels. A Vision Transformer (ViT-Base-Patch16-224) is pretrained as a masked autoencoder on a large pool of unlabeled tongue images. Patches are hidden with a 75% masking ratio and the model is trained to reconstruct the hidden pixels using a pixel-wise mean squared error loss. This gives the encoder generalizable visual features even though labeled data is scarce.
2. Model label relationships as a graph. Diagnostic labels become nodes and statistically significant co-occurrences become edges. The graph is built by counting how often labels appear together in the training data and thresholding at the 25th percentile of non-zero co-occurrences, producing a sparse graph. A two-layer GATv2Conv network then propagates information between labels, so a prediction for one label is informed by related labels. Two enhancements are added: rare-label boosting, which rescales attention by log(1/P(y_k = 1)), and correlation confidence weighting, which scales attention by normalized co-occurrence frequency.
3. Encode clinical rules as a soft penalty. The training objective combines a diagnosis loss with constraint and prior terms. The constraint term covers mutual exclusion (two diagnoses cannot co-occur), co-appearance (two diagnoses must both be present or both absent), and implication (one diagnosis requires another), each penalized only when violated via a max(0, .) form. The prior term uses KL divergence to align average predictions with empirical class frequencies. The diagnosis loss is asymmetric loss (ASL), which down-weights the negative gradients for frequent classes and uses asymmetric focusing parameters.
4. Recover rare classes with boosting. A base model is trained on the full dataset; a second model is fine-tuned only on underperforming classes using RandAugment, random erasing, and normalization. The final prediction replaces the five lowest-performing labels with the second model's outputs and keeps the base model's outputs for all other labels.

Data curation and training details: 4,000 annotated images came from two independent clinical sources. Ten systematically trained experts labeled samples independently, followed by mutually blinded cross-review, dual expert audits for disagreements, and adjudication by a senior traditional medicine practitioner. Data was split 80% training, 10% validation, 10% test. Images underwent color correction, haze removal, and reflectance normalization; tongue regions were segmented with DeepLabV3+ and manually refined in ITK-Snap. The GATv2Conv module uses in_dim=768, hidden_dim=64, out_dim=21, and num_head=8, and a shared two-layer MLP classifier with dimensions 640 x 320 and 320 x 1 with ReLU activation produces per-label predictions. Training uses the AdamW optimizer, learning rate 1e-3, batch size 200, layer-wise decay 0.75, over 200 epochs with λ1 = 0.1 and λ2 = 0.05, on an NVIDIA A800 GPU.

Why This Matters

Impact on research. The paper argues that diagnostic medical imaging has underexplored the paradigm of combining domain knowledge with data-driven learning that has proven useful in other scientific domains (the authors cite Deep Reasoning Networks and Physics-Informed Neural Networks as precedents). It offers a concrete recipe for injecting expert rules as differentiable penalties, and contributes a public 4,000-image, 22-label benchmark for a specialty domain where prior large datasets, such as the 8,676-image, seven-category set from Jiang et al. (2022), were not public.

Real-world applications:

• Point-of-care tongue image screening in traditional Chinese medicine clinics, where rare but clinically significant findings such as dark-red tongue or gray-black coating are frequently missed by existing models.
• Triage and decision-support tools that flag possible diagnoses for a practitioner rather than replacing clinical judgment, given the reported reductions in dimension-level missed detections.
• Standardized annotation and quality-control pipelines for building other specialized medical imaging datasets, using the described consensus labeling, cross-review, and audit procedure.
• Transfer of the same pipeline (self-supervised pretraining plus constrained label graph) to other multi-label diagnostic imaging tasks where labels are correlated and rules are known.

Industry relevance. The framework targets the common industrial situation of abundant unlabeled images and scarce expert labels, and the reported ablation shows pretraining removal costs 23.0% Macro-F1 and 29.0% Macro Recall, which quantifies the value of investing in unlabeled-data pipelines. Constraint-aware training also offers a way to keep deployed models from emitting clinically impossible label combinations.

Future Directions

• Multi-modal extension. The conclusion states that future work will incorporate multi-modal signals, without specifying which modalities.
• Clinical deployment evaluation. The authors plan to evaluate deployment within clinical workflows to improve robustness, reliability, and interpretability; no deployment study is reported in this paper.
• Generalization beyond tongue diagnosis. The paper claims the framework "readily generalizes to broader diagnostic medical imaging tasks," but the experiments reported here cover only tongue images, so external validation remains open.
• Resolving reporting details. The discrepancy between 15,905 and 19,505 unlabeled pretraining images, and the GATv2Conv out_dim of 21 versus the 22 described diagnostic labels, are details the paper does not explain and that would need clarification for replication.

Target Audience

Researchers and practitioners in medical image analysis and computer vision who work on multi-label diagnosis with limited annotations; developers building clinical decision support for traditional Chinese medicine or other rule-constrained diagnostic domains; and methodologists interested in combining self-supervised pretraining, graph neural networks, and constraint-aware loss design. Readers need a working knowledge of vision transformers, graph attention networks, and multi-label metrics to follow the methodology and ablation results.

Authors’ abstract

Automated interpretation of medical images demands robust modeling of complex visual-semantic relationships while addressing annotation scarcity, label imbalance, and clinical plausibility constraints. We introduce MIRNet (Medical Image Reasoner Network), a novel framework that integrates self-supervised pre-training with constrained graph-based reasoning. Tongue image diagnosis is a particularly challenging domain that requires fine-grained visual and semantic understanding. Our approach leverages self-supervised masked autoencoder (MAE) to learn transferable visual representations from unlabeled data; employs graph attention networks (GAT) to model label correlations through expert-defined structured graphs; enforces clinical priors via constraint-aware optimization using KL divergence and regularization losses; and mitigates imbalance using asymmetric loss (ASL) and boosting ensembles. To address annotation scarcity, we also introduce TongueAtlas-4K, a comprehensive expert-curated benchmark comprising 4,000 images annotated with 22 diagnostic labels--representing the largest public dataset in tongue analysis. Validation shows our method achieves state-of-the-art performance. While optimized for tongue diagnosis, the framework readily generalizes to broader diagnostic medical imaging tasks.