LoGra-Med: Long-Context Multi-Graph Alignment for Medical Visual-Language Models

Dear Reviewer cEx3, we appreciate your positive feedback on our proposed algorithm. Below, we address your remaining concerns regarding the stability of LoGra-Med with respect to paraphrased text generated by other LLM models, as well as other questions you raised.

Q1. The potential sensitivity LoGra-Med on GPT-4, how does it work with other language models or simpler context extension?

To validate the robustness of LoGra-Med w.r.t the LLMs used to generate extended captions, we conducted additional experiments. (i) First, we perform simple paraphrasing during instruction tuning by replacing keywords with their synonyms, and (ii) We use the Gemini API from Google to generate paraphrase captions.

We test (i) on weights pre-training with 40% data and (ii) for both 10% and 40% pre-training data. Model performance are evaluated on VQA-RAD and SLAKE, with average scores computed for open- and closed-based questions in these datasets. These new experiments are reported in Table 7 in the attached link https://anonymous.4open.science/r/logra-med-386F/readme.md.

We conclude that LoGra-Med continues to perform well when using Gemini LLM models, achieving significantly better results than the LLaVa-Med baseline, especially on the 10% and 40% data subsets of VQA-RAD. In contrast, a simpler model version does not produce similar improvements. This difference can be attributed to the lack of richer, more contextually aware descriptions incorporated into LoGra-Med, which helps the model better use the image-caption relationships.

Q2. Variable performance on fine-tuned PathVQA and Zero-shot classification task.

We believe that the variations in LoGra-Med’s performance can be attributed to task-specific adaptation. In the case of PathVQA, the model was fine-tuned with auto-regressive learning on the provided training samples, allowing it to specialize its general knowledge for the domain-specific task. In contrast, the zero-shot image classification task does not benefit from fine-tuning, which makes it more difficult for the model to generate accurate predictions, particularly when it hasn’t seen sufficient examples or when visual features aren’t strongly tied to the classes in the model’s knowledge base.

In conclusion, we argue that the zero-shot performance of LoGra-Med and other baselines could be significantly improved by fine-tuning task-specific training data.

Q3. Analysis of graph construction step with different K values in K-nearest neighbors.

We conducted new experiments, detailed in Table 6 (attached link: https://anonymous.4open.science/r/logra-med-386F/readme.md), to evaluate different K values used in the graph construction step. The table reports the model performance on the VQA-RAD dataset as well as the training time (Step-2 in pre-training) for each value of K. Our results show that K = 5 is the optimal choice. These findings will be included in the revised version of the paper.

Q4. Analysis of interpolation between the traditional autoregressive loss and the graph alignment objective.

In our default configuration, we choose alpha = 1 as the coefficients to combine auto-regressive and graph alignment objectives. To verify this choice, we ran further another ablation study with alpha = 0.1, 0.5 and trained LoGra-Med on 10% pre-training data. We present results on the VQA-RAD and Slake datasets in Table 8, Section F (attached link: https://anonymous.4open.science/r/logra-med-386F/readme.md).

The records demonstrate that alpha = 1.0 is the optimal choice; when reducing the contribution of multi-graph alignment, the model tends to drop their results.

Q5. Analysis of Extra Running Time of LoGra-Med (1 hour)

We attribute the increased training time to two main factors: (a) solving three graph alignment problems simultaneously and (b) the need to compute second-order graph alignments using heuristic solvers and perform backward gradient calculations with black-box gradient estimations.

Despite the additional training time required by LoGra-Med, this extra cost leads to better performance on downstream tasks. Furthermore, an ablation study showed that even when LLaVa-Med is trained with or without extended contexts and for a longer duration, our model consistently outperforms LLaVa-Med.

Dear Reviewer cEx3,

Thank you for your response. While we respect your evaluation, we kindly request further clarification on whether there are remaining issues or if additional improvements are needed, which makes you keep the unchanged scores without further justifications.

To date, we have thoroughly addressed all your concerns through several additional experiments, as detailed in Tables 6, 7, and 8 in this link (https://anonymous.4open.science/r/logra-med-386F/readme.md). Furthermore, we extended our evaluation by generating new results for LoGra-Med using the Gemini API, creating extended captions for 24K conversations. This effort demonstrates that LoGra-Med performs robustly with other LLM models. Notably, this process required substantial computational resources and incurred significant costs (to call API) for pre-training and fine-tuning downstream tasks, underscoring our commitment to addressing your feedback comprehensively.

We hope to hear your feedback

Regards

Authors

Dear Reviewer cEx3,

We appreciate your efforts in reviewing our paper. We hope that our response and revision can address your concerns. There are only two days left until the end of the discussion period. We look forward to hearing from you. Please feel free to leave any additional questions or comments. Otherwise, we would greatly appreciate it if you could raise your rating.

Thank you very much for your time and consideration!

Best regards,

Authors

Dear Reviewer jqRH, we appreciate your taking the time to give us feedback. Below, we address your concern ​​individually.

Q1. Explaining the motivation for adding graph alignment task in the pre-training stage.

We address this question through the following points:

First, we explain why doing a contrastive alignment in the pre-training step can help LLM models. Specifically, by applying triplet contrastive learning to input images, original captions, and extended captions, the model learns to identify relevant similarities and differences within the data rather than memorizing individual samples. This approach thus enhances robustness, particularly when training on small pre-training datasets. Additionally, including extended captions enriches the model’s contextual and hierarchical understanding, which expands the latent space and reduces representation collapse, enabling the model to better distinguish subtle variations in medical images and descriptions.

Second, among various contrastive learning approaches, we adopt a graph-based formulation inspired by prior work (LVM-Med, Duy et al., 2024). This formulation extends beyond pairwise contrastive learning by utilizing combinatorial graph matching, which captures stronger structural constraints within feature embeddings, leading to improved generalization performance. While LVM-Med is designed for aligning two graphs - vision and augmented vision - we extend this approach to three graph alignments: one for the vision graph, one for the original caption-based graph, and one for the extended caption-based graph. Our intuition is that image embeddings should be aligned with both their short and extended captions, and the LLM text embeddings for both versions should exhibit a high correlation.

Based on these motivations, we designed the multi-graph alignment for pre-training. If the Reviewer is still unclear at some points, we are happy for further clarification.

---

MH Nguyen, Duy, et al. "Lvm-med: Learning large-scale self-supervised vision models for medical imaging via second-order graph matching." NeurIPS 2024

Q2. Not Including MedTrinity results in comparison

There are two primary reasons why we did not include MedTrinity in our comparison:

Reproducibility Issues: Using the pre-trained models provided by MedTrinity, we could not reproduce results comparable to the best version reported in their paper (denoted as LLaVA-Med++ (Ours, w/)), despite following the authors’ instructions.

Incompatible Pre-training Strategy: MedTrinity does not offer a single unified pre-trained LLM model. Instead, separate pre-trained weights are provided for each downstream task (e.g., VQA-RAD, SLAKE, and PathVQA). This is because their approach involves combining a subset of training data from each downstream task with their collected data to create task-specific pre-trained models. For instance, to train the VQA-RAD model, they combine a subset of VQA-RAD training data with their collected data, extended with region of interest (ROI) annotations. A similar strategy is applied for SLAKE and PathVQA.

In contrast, our approach (and those of other methods we compare against) involves producing a single pre-trained model that can be fine-tuned across multiple downstream tasks. We argue that this fundamental difference in methodology makes MedTrinity unsuitable for direct comparison with our framework.

Q3. (i) Clarify the “10% data” setting. Is it "10% data size for the Stage 2 instruction-tuning stage, but the data size for the Stage 1 pre-training is still full-size; (ii) Do we include extended context in fine-tuning the LLaVaMed baseline with 10%?

We address two concerns together due to their connection to our pre-training and fine-tuning setup. Below, we recap the details of our data usage and training stages:

Pre-training: This consists of two stages:

• Stage 1: We train only the MLP to align the output of vision models, keeping all other components frozen. This follows the same strategy as LLaVa-Med, using 600k image-text pairs and an auto-regressive training loss.

• Stage 2: Both the MLP and LLM are trained using auto-regressive learning combined with our multi-graph alignment approach. Specifically, we use 60k conversations, as in LLaVa-Med, for auto-regressive training, while the multi-graph alignment is trained with our extended captions derived from these 60k conversations.

Fine-tuning for downstream tasks: Using the pre-trained model from the pre-training phase, we fine-tuned the LLM weights with only auto-regressive learning, leveraging the training samples provided for each dataset.

Dear Reviewer jqRH,

We would like to thank you very much for your insightful review, and we hope that our response addresses your concerns regarding our paper. However, as the discussion period is expected to end in the next few days, please feel free to let us know if you have any further comments on our work. We would be willing to address any additional concerns you may have.

Thank you again for spending time on the paper, we really appreciate it!

Best regards,

Authors

Dear Reviewer jqRH,

We appreciate your efforts in reviewing our paper. We hope that our response and revision can address your concerns. There are only two days left until the end of the discussion period. We look forward to hearing from you. Please feel free to leave any additional questions or comments. Otherwise, we would greatly appreciate it if you could raise your rating.

Thank you very much for your time and consideration!

Best regards,

Authors

Dear Reviewer jqRH,

We would like to kindly remind you that today is the final day of the discussion period, and we have not yet to receive your feedback since November 23rd. We have put a large effort into addressing your concerns, such as developing an interface and conducting a user study with practicing doctors to validate the context generated by GPT-4 models (see Sections H and I in [https://anonymous.4open.science/r/logra-med-386F/readme.md]). Given the importance of your feedback to us, we would greatly appreciate it if you could take the time to share your thoughts in the remaining hours.

Thank you for your time and consideration, and we look forward to hearing from you.

Best regards,

The Authors

Authors

Dear Reviewer jqRH,

We thank you for your response. It's very regrettable that the Reviewer made comments just before the discussion deadline a few hours and raised a new concern about the novelty of LoGra-Med compared with prior work on a single domain (LVM-Med).

Below, we address your concern with a comprehensive comparison of the two approaches across several key aspects. Our formulation distinguishes itself from LVM-Med in terms of its original motivation, the design of efficient solvers for multi-domain tasks, the method of graph construction, achieved performance, and the theoretical insights contributing to a deeper understanding of the algorithm. We hope this table helps you better understand our contributions.

Dear Reviewer Kirs, thank you for giving us detailed feedback, especially on our methodology section, such as chosen solvers and computing features, comparing them with optimal transport, etc. In the following, we address each of your concerns with several figures and tables referred to in the anonymous link.

Q1. Testing with simpler solutions (two graph alignment and solver for multi-graph alignments).

We address this question by doing two different experiments.

(i) Contrastive loss between two graphs: one from visual features and the other is original captioning.

(ii) Solving three pair-wise graph alignments in multi-graph alignment rather than solving through a barycenter graph.

We provide experiments for (i) with 40% pre-training data and (ii) with 100% pre-training data. Performance is reported with Average scores on two downstream tasks (VQA-RAD and SLAKE) after fine-tuning.

The results demonstrate that incorporating extended context through multi-graph alignment significantly improves performance compared to using only two graphs, as in setting (i). Setting (ii) also highlights that solving independent pairwise graph alignments leads to lower performance than our proposed solution. This is because LoGra-Med accounts for all mutual relationships, enforcing consistency across the triplet and avoiding conflicts that may arise when pairwise alignments between subsets (e.g., A and B, B and C) contradict the relationship between A and C.

To provide further insights, we also measured the running time of setting (ii) compared to the default LoGra-Med configuration. The statistics indicate that both approaches require comparable training time, with approximately seven hours in Stage 1 and 7.5 hours in Stage 2.

In conclusion, this evidence validates the effectiveness of LoGra-Med’s configuration. Moreover, it is worth noting that our solver’s efficiency makes LoGra-Med scalable to additional domains, such as health records or surgery videos, further enhancing its applicability in building multi-modal LLMs.

Q2. Analysis of feature representation computing as average pooling for visual and language tokens.

We address this concern with two experiments:

• We trained LoGra-Med on 70% of the pre-training data, randomly sampling 1,000 unseen image-text pairs. The trained model extracted features using average pooling, and a box plot (Figure in Section G) visualized the central tendency, spread, and skewness of 1,000 positive and negative pairs. The results show: (i) the median similarity for positive pairs is significantly higher than for negative pairs, indicating clear separation; (ii) while some overlap exists in the interquartile ranges (IQRs), the shift in central tendency confirms the distinction; and (iii) outliers are present, particularly among negative pairs, but they minimally overlap with the core distribution of positive pairs.

• Another example can be found in Table 2 of our paper (we added again below), where we compare two baselines, InfoNCE and PLOT. Both utilize the same contrastive loss, but InfoNCE relies on cosine distance with averaged features, while PLOT directly computes optimal transport over sets of visual and text tokens. The results for these baselines are summarized below.

We see that using a more sophisticated distance metric, such as optimal transport, provides a slight improvement (around 1%) over the averaging approach. However, the performance gain is relatively modest. Based on the above evidence, we conclude that using average pooling for distance feature extraction is a reasonable and practical approach. Notably, a similar method was employed in VLAP (Park, Jungin et al., 2024), where its simplicity and consistency were validated.

However, we still acknowledge the Reviewer’s point that alternative methods may offer improvements, and we suggest this as a valuable direction for future exploration.

Park, Jungin, Jiyoung Lee, and Kwanghoon Sohn. "Bridging Vision and Language Spaces with Assignment Prediction." ICLR 2024

Q3. Compute macro-F1 scores for classification/VQA because medical datasets might be imbalanced.

Due to time constraints during the rebuttal phase, we focused on the VQA task to report additional macro-F1 scores. These scores were not included in the main paper since other baselines did not provide F1 metrics for VQA, reporting only Recall and Accuracy. To ensure consistency, we computed the F1 scores for both our model and LLaVa-Med using the checkpoints reproduced from their GitHub repository.

We evaluate F1 scores for VQA-RAD and SLAKE while skipping the PathQA dataset due to its mixed domains, where visual-question answering data cannot be clearly separated by image modalities (e.g., MRI, CT, X-ray) or specific organs (e.g., liver, lung, chest).

• VQA-RAD: questions in this data can be divided into categories of “Chest, Head, Abdomen.” We compute the F1 score for multiple-choice questions for each data body organ. We present the results in Table 1 in the attached link https://anonymous.4open.science/r/logra-med-386F/readme.md.

• SLAKE: Questions in this data can be divided into different data modalities, such as CT, X-ray, and MRI. We compute the F1 score for multiple-choice questions for each data modality and for all samples. Results are summarized in Table 2.

Based on the F1 scores, we conclude that LoGra-Med demonstrates consistent performance, aligning with its accuracy results, and continues to outperform LLaVa-Med on this metric.

Q4. Confusing about using the “long-context” term.

We thank the Reviewer for pointing out the potential confusion around the term “long context.” In our work, “long context” does not refer to the length of input tokens used in training LLMs or the window size of attention mechanisms, as is commonly understood in large-scale LLM training. Instead, “long-context” refers to the enriched captions generated by extending the original medical captions with paraphrased descriptions and additional relevant context using GPT-4. These enriched captions provide more detailed semantic information about the input image. We chose this term to emphasize the richer contextual information contained within the extended captions rather than any technical aspect related to token length or attention span. To address potential ambiguity, we will revise the terminology in the Abstract, Introduction, and Related Works sections to clearly specify its meaning within the context of our medical visual question-answering task.

Given these arguments, we prefer to retain the current term to minimize significant revisions across multiple sections, including method names and descriptions. However, if the Reviewer feels strongly about this, we are open to reconsidering and adopting an alternative term. Please let us know if this proposed solution is acceptable.

Q5.Why is graph perspective used? How is it different from optimal transport?

Our graph alignment solving the node-to-node correspondences under edge constraints indeed can be formulated using optimal transport, namely fused Gromov-wasserstein optimal transport (FGW-OT) (Nguyen, Duy MH, et al. 2024; Ma, Xinyu, et al. 2024). However, two main challenges hinder us from using optimal transport in LoGra-Med:

• First, performing the forward pass to compute alignment between graph pairs using optimal transport is computationally expensive (Nguyen, Duy MH, et al. 2024), making it impractical for scaling to large-scale LLM training with hundreds of thousands of samples. This obstacle is further challenged in LoGra-Med, where three separate graph alignment problems have to be solved, significantly incurring computational costs. In contrast, adopting a graph-based formulation enables the use of heuristic solvers specifically designed for combinatorial graph matching (Swoboda et al., 2017; Rolínek et al., 2020), providing an efficient solution to address the high computational complexity of graph alignment.

• Second, our training loss (Hamming loss, Eq.(5)) requires gradients from graph alignments to learn feature representations. Using optimal transport would necessitate backpropagation through its Sinkhorn iterations (50-100), adding substantial computational cost and GPU memory usage for storing intermediate variables.

Our graph-based formulation addresses the second challenge by leveraging modern gradient estimation techniques for black-box optimization (Niepert et al. 2021), making the backward step efficient for LLM training.

---

Nguyen, Duy MH, et al. "Structure-Aware E (3)-Invariant Molecular Conformer Aggregation Networks." ICML 2024 Ma, Xinyu, et al. "Fused Gromov-wasserstein graph mixup for graph-level classifications." NeurIPS 2024. Paul Swoboda et al. “A study of lagrangian decompositions and dual ascent solvers for graph matching.”, CVPR 2017 Niepert et al. “Implicit MLE: backpropagating through discrete exponential family distributions”, NeurIPS 2021

Q6. Without using extended contexts, how LoGra-Med is different from PLOT paper

We thank the Reviewer for your question. Without using the third graph from the extended caption, LoGra-Med reverts to a graph alignment between the image and its caption. While LoGra-Med and PLOT paper both works with two domain information, there are several significant differences between the two methods that we list below:

Q7. Providing some explanation for mixed results with DCI features.

We aim to integrate DCI into LoGra-Med to capture better multi-scale visual features from vision encoders, potentially benefiting medical image analysis by considering both local (detailed) and global (contextual) features.

To address the Reviewer’s concerns, we retrained LoGra-Med with the DCI component using more carefully optimized hyperparameters. Previously, we had applied the default LoGra-Med settings, including the learning rate, without tailoring them for the DCI-enhanced version. The updated results (Table 3 in the attached link https://anonymous.4open.science/r/logra-med-386F/readme.md) demonstrate that DCI achieves comparable performance to the standard LoGra-Med on SLAKE, with a slight decrease, while showing improved average performance on VQA-RAD and a particularly notable gain on PathVQA.

We explain this phenomenon because the impact of the DCI component across datasets can be attributed to differences in dataset characteristics and task requirements. In particular, PathVQA likely benefits more from multi-scale features due to its reliance on both local and global visual contexts, making DCI’s integration more effective. On the other hand, VQA-RAD and SLAKE may rely less on multi-scale visual representations and more on different factors like textual reasoning or specific feature interactions. In terms of practicality, we would suggest providing both variations where the better model can be determined based on specific applications.

Dear Reviewer 9ZPV,

First, we sincerely thank you for taking the time to provide us with your valuable feedback. Below, we address your concerns in detail.

Q1. Performance comparison with the LLaVa-Med trained with extended context and longer training time.

We composed this question based on two first concerns and the second question of the Reviewer (Weaknesses and Question section). To address this, we conducted the following experiments:

• We train LLaVa-Med baseline with 10% pre-training data (both Stages) using extended captions in their auto-regressive training loss.

• We train LLaVa-Med baseline with 40% pre-training data (both Stages) using extended captions in their auto-regressive training loss.

Both settings were trained for 3 epochs, with training times of 0.47 and 1.9 hours, respectively. In comparison, we trained LoGra-Med for 2 epochs, ensuring that LLaVa-Med had a longer training duration than ours (our training times were 0.45 and 1.8 hours). Afterward, we fine-tuned the pre-trained models on downstream tasks. The results are summarized in Tables 4 and 5 in the attachment link https://anonymous.4open.science/r/logra-med-386F/readme.md (highlighted in grey color).

We observe that incorporating extended captions during the pre-training phase of LLaVa-Med enhances its performance compared to the default settings. However, both of these new versions still fall short of LoGra-Med’s performance. This suggests that our graph alignment algorithm plays a significant role in enhancing LoGra-Med's performance.

Q2. Provide more details for hyper-parameters like batch size B, and values of k in graph construction.

We apologize for the missing report on these parameters. We used a batch size of B = 16 x number of GPUs, and set k = 5 for graph construction. To validate this choice of k , we present an additional ablation study on the VQA-RAD dataset, as shown in Table 6 in the attached link https://anonymous.4open.science/r/logra-med-386F/readme.md. The results indicate that k = 5 strikes the optimal balance between performance and complexity. We will add these new results and provide more details for hyper-parameters in the revised version.

Dear Reviewer 9ZPV,

Thank you for providing your feedback. We respectfully seek further clarification regarding your decision to maintain the same score, as this was done without additional explanation despite our comprehensive efforts to address your concerns. Could you kindly indicate whether additional studies or improvements are required, or if the results provided still fall short of your expectations?

We want to summarize again your concerns addressed during the rebuttal:

• We provided the hyperparameters used in our method and conducted extra experiments to validate these choices (Tables 4 and 5 in https://anonymous.4open.science/r/logra-med-386F/readme.md).
• We performed a direct comparison with LLaVa-Med under identical extended-context conditions and trained our method for longer durations to ensure a fair and thorough evaluation (Table 6).

We look forward to hearing your thoughts.

Best regards,

The Authors