Extended Context Graph Alignment for Medical Vision-Language Models
A medical vision-language model must line up three views of the same training batch, one vector per sample: the pooled visual patch features Zv, the LLM features of the original answers Za, and the LLM features of GPT-4 extended answers Zae. Aligning every pair of views costs K(K−1)/2 comparisons. ExGra-Med instead forms one barycenter (the elementwise mean of the three node-feature matrices) and aligns each view to that single hub, K comparisons, reusing the source view's kNN edges for the barycenter. A second-order QAP scores both a node cosine term and an edge-difference inner-product term, and the training signal is the bidirectional Hamming distance between the solver's correspondence V̂ and the identity ground truth V*. Toggle first vs second order and watch the crossed red matches snap back to green.
Each batch has B = 6 samples with a hidden identity 1..B (the released code uses per-device B = 16). From that identity the demo forms the paper's three feature sets, one 2D vector per sample: Zv (image patches, teal), Za (answers, amber) and Zae (extended answers, violet), each turned into a kNN graph (k = 3, Euclidean; code k = 5) over the batch. The gold barycenter is the elementwise mean of the three node-feature matrices, barycenteri = (Zv,i + Za,i + Zae,i)/3, recomputed live, and it reuses the source view's kNN edges (code: edge_indice2 = edge_indice1). Each source node is matched to a barycenter node by a genuine QAP solve: first order scores the node cosine cost only, so look-alike samples get crossed red matches (a Hamming error against the identity V*); second order (SGA) adds the edge-difference inner-product cost (endpoint difference vectors on the reused kNN structure) and snaps the matches back to green. Exact code cost transforms: unary = −(cosine + Gumbel(0, 0.05) − 0.5), quadratic = −0.5·(inner-product + Gumbel(0, 0.05) − 0.5). The big number is the bidirectional Hamming loss, the mean over the three views of [Hamming(view→bary) + Hamming(bary→view)]/2 with target V* = identity, with a scrolling history and the per-view counts. Optimize runs the real I-MLE step: perturb the cost tensors with Gumbel(0, 0.05), solve for a MAP sample, take the blackbox finite-difference gradient g = (1/lambda)[solve(c + lambda·dL/dV) − solve(c)] with lambda = 80, and nudge the points (violet arrows) so the loss counts down toward 0, the Stage-2 training story; Scramble Zae distorts only the extended-answer view so its graph diverges and the loss jumps. On the right, the star vs the complete graph make the efficiency claim exact: aligning to one barycenter is K operations against K(K−1)/2 for all pairs (equal at K = 3; K = 8: 8 vs 28). Reported experiments: on medical VQA (VQA-RAD, SLAKE, PathVQA, open/closed accuracy), ExGra-Med matches LLaVA-Med using only 10% of the pretraining data, a +20.13% VQA-RAD gain at the 10% budget (VQA-RAD average: LLaVA-Med 10% = 52.39 to ExGra-Med 10% = 72.52, versus LLaVA-Med 100% = 72.64); full-data ExGra-Med is approx VQA-RAD 74.9 / SLAKE 85.5 / PathVQA 63.9. Config: vision CLIP-ViT-L/14, LLM LLaMA-7B, 2-layer GraphConv GNN, kNN k = 5, Gumbel beta = 0.05, BBGM lambda = 80, LPMP maxIter = 100, alpha = 1.0, per-device B = 16, Stage-2 lr 2e-5, 3 epochs, 4x A100-80GB. (No neural network runs here: there is no ViT, no LLM and no text encoder. The "embeddings" are hand-placed 2D points, not real Zv/Za/Zae; the LPMP (BBGM) message-passing solver is emulated by an exact brute-force QAP over the B nodes; the GNN message passing is a light denoising average, and the second-order edge term reads that GNN-denoised coordinate so identity is the structural optimum. Faithful and computed live: nodes = the batch samples, three panels = the three feature sets, the hub = the barycenter feature mean with the source view's reused kNN edges, unary = node cosine and quadratic = edge-difference inner-product with Gumbel(0, 0.05), each source-to-hub line = V̂ against V* = identity, the number = the bidirectional Hamming loss, Optimize = a real perturb-and-MAP blackbox gradient with lambda = 80, and the K vs K(K-1)/2 count is exact. Honest gaps vs the paper: the code uses a fixed lambda = 80 (the paper text describes a learnable step size) and Gumbel(0, 0.05) (the paper says Gumbel(0, 1)); the repo default generates extended answers with gpt-4o-mini (the paper reports GPT-4). It conveys the mechanism, not any medical-VQA accuracy.)
Run the experiments
Every animation runs live in your browser. Click a button to play that experiment on the demo (it scrolls up and starts); drag the slider to take over. Nothing is downloaded.
Align the graphs (Optimize)
Run the I-MLE optimisation: the bidirectional Hamming loss drives the three views into agreement with the barycenter.
First vs second order
Toggle the alignment order: the second-order Gromov-Wasserstein term snaps the crossed look-alike matches back to the identity.
Number of graphs K
Change the batch size and compare barycenter alignment against all-pairs cost.
The idea in three steps
A medical multimodal model learns better when the visual features, the answer text and a richer GPT-4 extended answer are pushed into agreement. ExGra-Med does that agreement as a graph-alignment problem against a shared barycenter, not as an expensive all-pairs comparison, and folds it into Stage-2 pretraining as L = Lautoregressive + alpha·Lalign with alpha = 1.0.
Three views, one vector per sample
Per batch, mean-pool the CLIP-ViT-L/14 patches into Zv, and push the original answers and the GPT-4 extended answers through the LLM (LLaMA-7B), mean-pooling each into Za and Zae. The nodes are the batch samples; a 2-layer GraphConv GNN and a kNN graph (k = 5) give each view its edge structure.
Align to one feature mean, not every pair
Average the three node-feature matrices into a barycenter and reuse the source view's kNN edges for it, then solve a second-order QAP (BBGM, LPMP backend, lambda = 80) with a node-cosine term and an edge-difference inner-product term. That is K alignments instead of K(K-1)/2 pairwise ones.
Bidirectional Hamming, backpropped by I-MLE
The loss is the bidirectional Hamming distance between the solver's correspondence V̂ and the identity V*, averaged over the three views and both directions. The solver is not differentiable, so I-MLE takes a Gumbel-perturbed blackbox gradient (perturb-and-MAP, lambda = 80) and folds it into pretraining.
For the extended-context construction, the second-order graph-alignment objective, the I-MLE gradient estimator and the medical visual-question-answering results (matching LLaVA-Med with only 10% of the pretraining data, and +20.13% on VQA-RAD at the 10% budget), see ExGra-Med: Extended Context Graph Alignment for Medical Vision-Language Models (Duy M. H. Nguyen, Nghiem T. Diep, TrungTin Nguyen et al., NeurIPS 2025).