{"id":"ad84f52c-1d26-4e7c-a089-ebc65348e021","arxiv_id":"2509.07588","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"Aligning language model entity embeddings with UMLS knowledge graph subgraph embeddings during pre-training improves biomedical QA and entity linking across PubMedBERT and BioLinkBERT.","lead":"BALI is a pre-training method that aligns a biomedical language model's entity representations with knowledge graph subgraph representations using a contrastive loss, then keeps only the language model for downstream tasks. The paper reports consistent gains on question answering and entity linking benchmarks across PubMedBERT and BioLinkBERT, suggesting a cheap route to injecting structured biomedical knowledge without inference-time graph retrieval.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Zero-shot EL gains largely measure the training objective; without an MLM-only control on EL or a disjoint-concept test, the headline entity-linking improvements do not cleanly isolate the KG-alignment effect.","rationale":"The paper is honestly written and the method is well specified. The key positive evidence—large zero-shot EL gains—is real but partially circular because the evaluation metric is nearly identical to the InfoNCE objective optimized during pre-training. The QA and RE gains are more independent, and the Table 6 ablation already provides an MLM-only control showing that alignment adds beyond continued MLM on QA. However, no such control is given for EL, and EL is where the largest improvements appear. The central claim that BALI's cross-modal KG alignment drives the improvements therefore remains conditional on closing this gap. A single MLM-only EL evaluation would settle whether the headline EL gains are attributable to the alignment signal or to continued pre-training. I do not see grounds to reject; the reader's CONDITIONAL verdict is appropriate.","tokens_in":25107,"tokens_out":6118,"duration_ms":75928,"concrete_test":"Pre-train PubMedBERT with the same BALI pre-training data, batch size, and 65K steps (Table 1) but with L_align removed, keeping only MLM. Evaluate this model with the same zero-shot retrieval protocol on NCBI-Disease, BC5CDR-D, and BC5CDR-C. If the MLM-only model retains a substantial fraction of BALI's Accuracy@1 gains (e.g., NCBI >58.8, vs PubMedBERT 49.51 and BALI 68.14), the alignment loss is not the main driver of the headline EL improvements.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central empirical evidence for improved entity representations is Table 3's zero-shot entity linking. But this evaluation largely re-measures the BALI training objective. In Section 4.2, L_align maximizes cosine similarity between a pooled mention representation e_v and a graph representation g_v whose initial node features are LM-encoded concept names (Section 4.1). The zero-shot protocol in Section 5.3.2 retrieves concepts by cosine similarity between mention and concept-name representations. Thus a model trained to maximize exactly this similarity is expected to show large retrieval gains; those gains do not independently demonstrate that external KG structure was internalized. The QA results are more independent, but Table 6 shows that removing L_align (MLM-only) raises PubMedQA from 63.1 to 63.78 and BioASQ from 67.8 to 70.58, so a meaningful part of the QA improvement comes from continued pre-training. No MLM-only control is reported for entity linking, so the largest claimed gains cannot be attributed to the alignment term versus continued MLM on the same entity-linked corpus. This is a missing control, not an internal contradiction; the paper's own ablations are a good start.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes BALI, a pre-training method that augments biomedical language models (PubMedBERT, BioLinkBERT base/large) by aligning pooled entity-mention representations with UMLS knowledge-graph subgraph representations. A GAT encoder encodes local 1-hop subgraphs, whose initial node features are LM-encoded concept names; an InfoNCE contrastive loss pulls mention and graph representations of the same concept together, while MLM is retained as a joint objective. After pre-training on 1.67M PubMed sentences with BERN2-based entity linking, the GNN is discarded, and the resulting LM is evaluated on biomedical question answering (PubMedQA, MedQA, BioASQ), entity linking (zero-shot and supervised), and relation extraction (ChemProt, DDI, GAD). The paper reports consistent QA gains over base models, large zero-shot entity-linking gains for general biomedical LMs, small relation-extraction gains, and ablations showing that both objectives and attention-based graph aggregation matter.","tokens_in":25352,"tokens_out":6424,"duration_ms":70234,"significance":"If the claims hold, BALI would be a valuable and relatively cheap recipe for infusing structured biomedical knowledge into lightweight encoders without any inference-time KG access. The paper has concrete strengths: it releases code and pre-trained models, reports hyperparameters in Table 1, averages repeated fine-tuning runs for PubMedQA and BioASQ, and includes a useful ablation suite (Table 6) covering loss removal, alternative alignment losses, aggregation methods, and GNN depth. The core idea of explicit cross-modal alignment through entity anchors is clearly presented and is distinct from prior interaction-token or retrieval-augmented approaches. However, the current evidence for the entity-representation claim is weakened by the absence of an MLM-only control in the entity-linking experiments and by the close similarity between the zero-shot entity-linking protocol and the training objective. The QA results are more independent and, with the existing MLM-only ablation, already show that the alignment term contributes beyond continued pre-training, which is a useful partial control.","major_comments":[{"comment":"The zero-shot entity-linking evaluation largely re-measures the training objective. L_align (Eq. 5) maximizes cosine similarity between a pooled mention representation and a GNN subgraph representation whose initial node features are LM-encoded concept names (§4.1). The zero-shot retrieval protocol (§5.3.2) retrieves by cosine similarity between mention and concept-name representations. This is close to the positive-pair construction in training. More importantly, Table 3 reports no MLM-only control for entity linking, even though Table 6 shows that MLM-only accounts for a substantial portion of the QA gains (PubMedQA 63.1→63.78; BioASQ 67.8→70.58). Without an equivalent EL control or a disjoint-concept test, the headline EL gains (e.g., PubMedBERT +13.1 Accuracy@1) cannot be cleanly attributed to the KG-alignment term rather than to continued MLM on the same entity-linked corpus.","section":"§4.2, §5.3.2, Table 3"},{"comment":"The text claims that BioLinkBERT_large 'after BALI pretraining performs on par or better than the task-specific QA-GNN and GreaseLM methods.' This is contradicted by Table 2: BioLinkBERT_large+BALI(GNN) scores 68.7 on PubMedQA and 45.0 on MedQA, while QA-GNN scores 72.1/45.0 and GreaseLM 72.4/45.1. Even the linear-graph variant (70.9 PubMedQA) is below QA-GNN and GreaseLM. The comparison should be corrected, or the claim narrowed to specific datasets where it actually holds (e.g., MedQA).","section":"§5.3.1, Table 2"},{"comment":"MedQA results are reported as single means with no standard deviations or error bars, whereas PubMedQA and BioASQ have ± values averaged over 10 runs. Since the reported MedQA gains are small (e.g., PubMedBERT 38.1→39.8; BioLinkBERT_large 44.6→45.0), it is impossible to assess whether these differences are significant. Please report variance or at least the number of seeds used for MedQA.","section":"§5.1 / Table 2"},{"comment":"The pre-training corpus is built with BERN2 for named entity recognition and UMLS normalization, but the paper reports no quality check on this linking. If a substantial fraction of mentions are mapped to incorrect UMLS concepts, the contrastive loss will pull text representations toward wrong graph nodes, and the observed gains would reflect a different effect. The manuscript should at least report a sample-based precision estimate for the BERN2 linking on the pre-training data, or provide a robustness analysis (e.g., training on a filtered subset with high-confidence links).","section":"§5, 'Pretraining Data'"}],"minor_comments":[{"comment":"The evaluation-task list has a duplicate: item (iii) is listed as 'BC5CDR-D' but should presumably be the BC5CDR-Chemical corpus, matching Table 3's 'BC5CDR-C' column.","section":"§5.1"},{"comment":"The ablation text says 'larger (5 layers)' but Table 6 reports L=7 for the larger GNN. The text and table should be harmonized.","section":"§5.5.3 / Table 6"},{"comment":"The pre-training corpus size is given inconsistently as '1.5M sentences' (Introduction), '1.67M sentences' (Pretraining Data), and '1.7M sentences' (Conclusion). Please use one consistent figure.","section":"Abstract / §1 / §5 / Conclusion"},{"comment":"The underline convention ('best of two scores') is defined only for original vs. BALI-pretrained models. For the SapBERT and GEBERT rows, the notational significance of underlining is less clear; please clarify or add bolding for the overall best per column.","section":"Table 3"}],"recommendation":"major_revision","confidential_remarks":"The central methodological idea is promising and the QA evidence, together with the paper's own MLM-only ablation, gives some support to the alignment term. The main gap is the entity-linking evaluation: the near-objective nature of zero-shot EL and the lack of an MLM-only control make the strongest claim ('improved entity representations') currently under-supported. The overclaim regarding QA-GNN/GreaseLM in §5.3.1 should also be corrected before publication. These are addressable with additional experiments and textual revisions, so I recommend major revision rather than rejection."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"BALI is a solid, well-ablated pre-training idea: align a biomedical LM's entity embeddings with GNN-encoded UMLS subgraphs via InfoNCE, then discard the GNN at inference. No inference-time KG retrieval, cheap pre-training (65K steps, ~9 hours on 4 V100s), and code plus weights are public. The QA results are genuinely positive: for PubMedBERT, BALI lifts PubMedQA, MedQA, and BioASQ by 2.1, 1.7, and 6.2 points, and the ablation shows that removing L_align (MLM-only on the same corpus) still gains 0.68 on PubMedQA and 2.78 on BioASQ over the base. So continued pre-training explains part of it, but the full objective adds a further 1.4 and 3.4 points. That is a real signal.\n\nThe entity-linking story is less clean. The zero-shot protocol retrieves concepts by cosine similarity between pooled mention and concept-name embeddings, and the training objective optimizes essentially that similarity (via the GNN-initialized-from-concept-names). So the large zero-shot EL gains are partly the model conforming to the evaluation metric. The paper does not report an MLM-only control for EL, so we cannot attribute the +13-24 Accuracy@1 gains to the alignment term over continued pre-training on the same corpus. Given the QA ablations, this is a missing control rather than a fatal flaw.\n\nOther soft spots: MedQA lacks error bars; the InfoNCE temperature is not stated; and the base-model comparisons use original checkpoints rather than same-corpus MLM-only runs for BioLinkBERT. The PubMedBERT ablation partially covers this, but not fully.\n\nCredit where due: the node-representation ablations (GAT vs GraphSAGE vs linearized vs TransE/DistMult) are informative, the authors honestly report where BALI hurts (e.g., SMM4H zero-shot for SapBERT), and the paper is clearly written. The central claim—explicit cross-modal alignment adds something beyond continued MLM—holds up for QA.\n\nThis deserves a serious referee. I'd recommend acceptance with minor revisions: add the MLM-only EL control, report tau, and give MedQA variance. Worth engaging with.","headline":"BALI is a solid pre-training method for biomedical LMs, with real QA gains; the entity-linking improvements are partly circular and need an MLM-only control.","tokens_in":25859,"tokens_out":4117,"would_cite":true,"duration_ms":40519,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Aligning LMs with UMLS subgraphs boosts biomedical QA and linking.","keywords":["biomedical knowledge graph","UMLS","contrastive learning","language model pre-training","entity linking","biomedical question answering","graph neural network","InfoNCE"],"falsifier":"Pre-train BALI twice with identical data and compute, once with true mention-to-UMLS links and once with randomly shuffled links; if the shuffled model keeps the QA and entity-linking gains, then the knowledge-graph alignment itself is not the driver and the effect comes from the extra pre-training exposure or MLM.","tokens_in":1571,"feed_emoji":"🧬","tokens_out":2410,"duration_ms":59599,"temperature":0.7,"pith_summary":"The paper tries to establish that a single additional pre-training stage can inject biomedical factual knowledge into language models by aligning their entity representations with knowledge-graph subgraph representations. Concretely, BALI trains a graph encoder over local UMLS subgraphs and pulls each pooled textual entity mention toward its concept's graph embedding using an InfoNCE contrastive loss, while continuing masked-language modeling. The authors claim that this stage, using only 1.67 million PubMed sentences and about 600,000 UMLS concepts, improves PubMedBERT and BioLinkBERT on question answering and entity linking, and that the graph encoder can then be discarded with no knowledge-graph access needed at inference. If true, this offers a cheap, task-agnostic way to add structured biomedical knowledge to existing encoders.","feed_headline":"Aligning LMs with UMLS subgraphs boosts biomedical QA and linking","feed_subtitle":"A 65K-step contrastive stage on 1.67M PubMed sentences lifts PubMedBERT and BioLinkBERT across benchmarks.","key_machinery":"Cross-modal anchoring: a single biomedical concept is represented in two complementary modalities — a textual representation (mean-pooled token embeddings of its mention in context) and a structural representation (a GAT-encoded local UMLS subgraph, with initial node features obtained by encoding randomly sampled concept names with the LM). The InfoNCE contrastive loss over paired text-graph anchors is the alignment mechanism, MLM is retained to preserve language ability, and the graph encoder is discarded after pre-training.","core_discovery":"BALI's central claim is that explicit cross-modal alignment between a language model's entity representations and a graph neural network's concept representations transfers usable biomedical knowledge into the LM. For each mention in a masked sentence, the mention's pooled token embeddings serve as the textual anchor, and a GAT encoder produces a structural anchor from a 1-hop UMLS subgraph seeded with concept-name embeddings. An InfoNCE loss pulls paired anchors together while MLM preserves language ability. After this short pre-training, the graph encoder is discarded, yet the LM alone shows mean accuracy gains of 2.1, 1.7, and 6.2 points for PubMedBERT on PubMedQA, MedQA, and BioASQ, larg","pith_inferences":["The paper's own ablation shows that removing the alignment loss but keeping MLM already yields 63.78 on PubMedQA versus the 63.1 raw PubMedBERT baseline, so part of the gain is plausibly from continued masked-language pre-training rather than graph alignment; no fully matched equal-compute control is reported.","Because the training data relies on automatic entity recognition and normalization, incorrect mention-to-concept links would push text embeddings toward wrong graph nodes; a gold-annotated or confidence-filtered training set would likely sharpen the measured alignment signal.","The same anchor-based recipe should transfer to other text-attributed knowledge graphs, with the effective graph encoder choice depending on model capacity — small encoders appear to need the external GAT, larger ones can use linearized subgraphs.","A direct negative-control experiment, pairing sentences with randomly shuffled subgraphs instead of their true linked subgraphs, would isolate whether the contrastive alignment itself or the extra pre-training data drives the reported gains."],"forward_implications":["Knowledge can be injected during pre-training alone, so downstream tasks need no retrieved subgraphs or entity linking at inference time.","A small, balanced alignment corpus (1.67M sentences, ~600K concepts, 65K steps, roughly 9 GPU-hours for base models) is sufficient for consistent QA improvements.","Zero-shot entity linking improves sharply for general biomedical LMs, with Accuracy@1 rising by about 13 points for PubMedBERT and 24 points for BioLinkBERT-base on average across five corpora.","A BALI-adapted BioLinkBERT-base can match or slightly beat SapBERT, a task-specific model pre-trained on the full UMLS synonym space, on supervised entity linking.","GAT-based subgraph aggregation outperforms mean-pooling graph encoders, while larger LMs can benefit more from a single-encoder linearized-graph variant."],"supporting_citations":[{"why":"Supplies PubMedBERT, the primary base biomedical LM that BALI pre-trains and compares against.","marker":"[17]"},{"why":"Supplies BioLinkBERT, the second base LM, in base and large variants, and the fine-tuning hyperparameters used for evaluation.","marker":"[84]"},{"why":"BERN2 performs the named-entity recognition and normalization that links PubMed sentences to UMLS concepts, producing the BALI pre-training data.","marker":"[69]"},{"why":"Defines the UMLS knowledge graph, the source of concepts, relations, and local subgraphs used for alignment.","marker":"[5]"},{"why":"Provides the InfoNCE contrastive objective that pulls textual and graph representations of the same concept together.","marker":"[76]"},{"why":"Provides the graph attention network used as the external graph encoder for local subgraphs.","marker":"[78]"},{"why":"Defines SapBERT, the task-specific UMLS-pretrained entity-linking baseline that BALI compares against in supervised and zero-shot settings.","marker":"[35]"},{"why":"Provides DRAGON, the joint LM-KG pre-training baseline representing the alternative implicit-interaction approach that BALI is contrasted with.","marker":"[83]"}],"fun_headline_variants":["UMLS subgraph alignment sharpens biomedical LMs and entity vectors","BALI: Align LM tokens with KG subgraphs to lift biomedical accuracy","Graph-encoder pre-training then discard: LM alone gains on BioQA","Short contrastive KM-LM alignment improves PubMedBERT on three benchmarks","Link medical mentions to UMLS subgraphs for better biomedical language models"],"cache_read_input_tokens":27520,"weakest_assumption_plain":"The alignment signal is only as good as the automatic linking that matches mentions in the training sentences to UMLS concepts; if many links are wrong, the contrastive objective pushes text embeddings toward the wrong graph nodes, and the reported gains may come mostly from continued masked-language training instead.","fun_headline_variants_meta":{"raw":{"variants":["UMLS subgraph alignment sharpens biomedical LMs and entity vectors","BALI: Align LM tokens with KG subgraphs to lift biomedical accuracy","Graph-encoder pre-training then discard: LM alone gains on BioQA","Short contrastive KM-LM alignment improves PubMedBERT on three benchmarks","Link medical mentions to UMLS subgraphs for better biomedical language models"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000325,"raw_usage":{"total_tokens":1651,"prompt_tokens":728,"completion_tokens":923,"prompt_tokens_details":{"cached_tokens":256},"prompt_cache_hit_tokens":256,"prompt_cache_miss_tokens":472,"completion_tokens_details":{"reasoning_tokens":827}},"tokens_in":472,"tokens_out":923,"duration_ms":10542,"temperature":1.0,"reasoning_tokens":827,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-04T21:55:01.440229+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Pre-train BALI twice with identical data and compute, once with true mention-to-UMLS links and once with randomly shuffled links; if the shuffled model keeps the QA and entity-linking gains, then the knowledge-graph alignment itself is not the driver and the effect comes from the extra pre-training exposure or MLM.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies BioLinkBERT, the second base LM, in base and large variants, and the fine-tuning hyperparameters used for evaluation."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Provides the graph attention network used as the external graph encoder for local subgraphs."},{"cited_title":"Manning, Percy Liang, and Jure Leskovec","cited_arxiv_id":null,"evidence_quote":"Provides DRAGON, the joint LM-KG pre-training baseline representing the alternative implicit-interaction approach that BALI is contrasted with."}],"review_version":1}