NeurIPS 2026Balanced Fine-Tuning

Aligning LLMs with
Biomedical Knowledge
using Balanced
Fine-Tuning

Preserving the learning signal in difficult biomedical contexts.

A dual-scale post-training objective guided by local confidence, designed for knowledge-dense biomedical text.

Zhenchao Tang^*1, Fang Wang^*†1, Haohuai He¹, Jiale Zhou¹, Tianxu Lv¹, Jun Zhu¹, Shouzhi Chen¹, Minghao Yang¹, Yu Wang¹, Jiayang Wu¹, Yidong Song², Yaokun Li³, Jiehui Huang⁴, Jun Zhou⁵, Bing He¹, Jianhua Yao^†1

¹Tencent AI for Life Sciences Lab ²Nankai University ³Sun Yat-sen University
⁴The Hong Kong University of Science and Technology ⁵Tsinghua University

* Equal contribution † Corresponding authors

Read the paper ↗Code ↗NeurIPS poster ↗Poster PDF ↓

01 / The idea

Low confidence has a local structure.

In biomedical text, contiguous runs of low-confidence tokens often contain rare entities and mechanistic steps. Treating every uncertain token alike can suppress the knowledge we want the model to learn.

BFT compares each token’s confidence with its surrounding context and gives more weight to sequences containing difficult regions. The two signals allocate learning effort within and across sequences.

60.9%

Biomedical entities among low-confidence tokens in dense PMC windows, versus 12.5% in sparse windows.

+13.4pp

VCWorld accuracy over SFT for the aligned 70B backbone under a shared recipe.

+4.0pp

VCWorld improvement after subsequent GRPO, from 0.696 to 0.736.

Source: paper Introduction, VCWorld results, and Appendix diagnostic. Concept enrichment supports an interpretation of knowledge gaps; it is not direct ground truth for epistemic uncertainty.

02 / Method

One window. Two scales.

BFT uses a centered local window of valid target tokens, with a default size of g = 256. Prompt and padding positions do not contribute to the loss. Confidence-derived weights use stop-gradient.

BFT overview: sparse and dense low-confidence windows, token-level local normalization, and sequence-level reweighting — Figure 1. Biomedical uncertainty structure and the BFT objective. Click any figure to enlarge.

Within a sequence

ρ_b,t = clip(c_b,t / (C_b,t^loc + ε), 0, 1)

Compare token confidence with the local mean. A difficult token retains more signal when its entire context is difficult.

Across sequences

s_b = sg(1 − min_t C_b,t^loc)

Use the weakest local region to assign a bounded sample weight. Sequences with difficult regions receive greater relative emphasis.

Combined objective

L_BFT = (1/B) ∑_b s_b [∑_t m_b,t sg(ρ_b,t) ℓ_b,t] / [∑_t m_b,t + ε]

Explore the token weight

Same token. Different context.

Fix the target-token confidence at c = 0.10 and vary the local mean. This illustrates the token weight before sample-level scaling.

Local mean confidence 0.20

Difficult contextConfident context

SFT

1.000

DFT

0.100

BFT token

0.500

03 / Evidence

Biomedical reasoning and RL compatibility.

We independently fine-tune DeepSeek-R1-Distill backbones at 14B, 32B, and 70B using a shared optimization recipe. The medical and biological mixture contains 51,056 filtered samples, with 48,503 in the supervised training split.

Chemical perturbation reasoning

VCWorld · 70B · five held-out cell lines

Average Up / Down / No classification accuracy
Initialization	Before GRPO	After GRPO	Change
DFT	0.490	0.420	−0.070
SFT	0.562	0.486	−0.076
BFT-sample	0.596	0.624	+0.028
BFT-token	0.660	0.694	+0.034
BFT	0.696	0.736	+0.040

Three-decimal averages from the Results section. The paper’s main table rounds these to two decimals. Changes here use the displayed three-decimal values.

Before GRPO, BFT is numerically higher than the Gemini-2.5-Flash backbone (+0.018), but the difference is not statistically significant across five cell-line means (95% CI [−0.004, 0.040], p = 0.088). Reported inference variability and training-seed variability are separate. The RL pattern applies to the evaluated VCWorld setup.

Medical competence across HealthBench Hard axes and themes and biological process reasoning across GeneAgent benchmarks — Figure 2. Medical competence on HealthBench Hard and biological process reasoning with GeneAgent. BFT improves more consistently than the supervised baselines under the shared recipe.

General capability retention and window sensitivity

MMLU and CMMLU capability retention, 14B window-size sensitivity, and injected owl preference transmission — Figure 3. DFT forgets least, followed by BFT-token and BFT. The window-size study uses the 14B model and is not directly comparable to 70B VCWorld results. The injected-preference test covers one preference and is not a general safety evaluation.

04 / Beyond generation

Better profiles support biological representations.

BFT-70B generates biomedical entity profiles. Encoding those profiles with text-embedding-ada-002 yields gene representations, which also support expression-weighted cell embeddings.

Gene embedding UMAP, gene property and interaction prediction, and cell clustering results — Figure 4. BFT gives the best average gene-level performance and improves the evaluated cell-clustering tasks. GenePT remains higher on Lys4.

Perturbation response prediction across four datasets using BFT text embeddings with the STATE decoder — Figure 5. With the STATE decoder retained, BFT embeddings perform comparably to standard inputs across four perturbation datasets. Profile generation is zero-shot with respect to these datasets; decoder training is not.

For multimodal integration, BFT reaches 0.57 overall scIB, above Harmony and BBKNN but below paired-supervision scMODAL (0.65). Improvements are task-specific and do not establish a general factuality or deployment guarantee.

05 / Use & cite

Resources

Paper on arXiv ↗Training code ↗NeurIPS 2026 presentation ↗Editable poster (PPTX) ↓Print poster (PDF) ↓Conference poster (PNG) ↓Paper thumbnail (PNG) ↓

Quick start

Enable BFT in your LLaMA-Factory training configuration. See the repository for setup and the complete recipe.

stage: sft
do_train: true
use_bft_loss: true
# Default local confidence window: BFT_GROUP_SIZE=256

BibTeX

@article{tang2025aligning,
  title={Aligning LLMs with Biomedical Knowledge using Balanced Fine-Tuning},
  author={Tang, Zhenchao and Wang, Fang and He, Haohuai and Zhou, Jiale and
          Lv, Tianxu and Zhu, Jun and Chen, Shouzhi and Yang, Minghao and
          Wang, Yu and Wu, Jiayang and Song, Yidong and Li, Yaokun and
          Huang, Jiehui and Zhou, Jun and He, Bing and Yao, Jianhua},
  journal={arXiv preprint arXiv:2511.21075},
  year={2025},
  note={Accepted at NeurIPS 2026},
  doi={10.48550/arXiv.2511.21075},
  url={https://arxiv.org/abs/2511.21075}
}

Correspondence: Fang Wang and Jianhua Yao.

Aligning LLMs withBiomedical Knowledgeusing BalancedFine-Tuning