CAFF

--- license: mit language: - en library_name: pytorch pipeline_tag: feature-extraction base_model: michiyasunaga/BioLinkBERT-large tags: - biomedical - knowledge-graph - multi-hop - retrieval - evidence-selection - kg-rag - rare-disease - orphanet - triple-filtering metrics: - f1 - precision - recall --- CAFF: Context-Aware Feedback Filtering for Multi-Hop Biomedical Knowledge Graph Evidence Selection Marwan Dhifallah * . Yu Liu Dalian University of Technology, Dalian, China marwan@mail.dlut.edu.cn . yuliu@dlut.edu.cn --- ## Table of Contents 1. [TL;DR](#tldr) 2. [The Context Blindness Error](#the-context-blindness-error) 3. [Approach](#approach) 4. [Repository Structure](#repository-structure) 5. [Installation](#installation) 6. [Data](#data) 7. [Training](#training) 8. [Evaluation](#evaluation) 9. [Results](#results) 10. [Ablation Study](#ablation-study) 11. [Configurations](#configurations) 12. [Hyperparameters](#hyperparameters) 13. [Reproducibility](#reproducibility) 14. [Hardware](#hardware) 15. [Scope and Future Work](#scope-and-future-work) 16. [Citation](#citation) 17. [License](#license) 18. [Acknowledgements](#acknowledgements) 19. [Contact](#contact) --- ## TL;DR CAFF is a triple filter for multi-hop biomedical knowledge graph retrieval-augmented generation. It addresses the **Context Blindness Error (CBE)**: filters that score each candidate triple from `(Query, relation, BFS_depth)` alone cannot distinguish whether the same triple is relevant or irrelevant under different upstream retained sets. CAFF fixes this with two coupled components: - **CSV** -- a parameter-free, permutation-invariant summary of the previous hop's retained set. - **DBM** -- a low-rank, sigmoid-gated dynamic perturbation of the bilinear scoring matrix, generated from the CSV. **Headline results** (Orphanet biomedical KG, 3,000 held-out test queries, 3 random seeds, paired bootstrap): | Metric | Mean +/- std | |-------------------------------|--------------------| | Test F1 (per-hop thresholds) | **0.5764 +/- 0.0022** | | Test F1 (autoregressive) | 0.5477 +/- 0.0006 | | Test MAP | 0.6741 +/- 0.0003 | | Test NDCG@10 | 0.7090 +/- 0.0003 | A complete leave-one-out ablation establishes CSV and DBM as the essential components; two additional losses that were considered during development (a depth-contrastive loss and an HC3-style contrastive loss) were measured and removed because they did not improve held-out F1 on this knowledge graph. See [Ablation Study](#ablation-study). --- ## The Context Blindness Error Consider the clinical query: > *"What drug targets the pathway of the causal gene of Fanconi anemia complementation group D1?"* The same hop-2 triple ` ` is: - **Relevant** when the hop-1 retained set is `{ }`. - **Irrelevant** when the hop-1 retained set is `{ }`. A filter that sees only `(Query, relation, hop)` cannot tell these two situations apart and therefore must score the triple identically in both. We call this the **Context Blindness Error** (CBE). By the Data Processing Inequality, any filter that ignores the previous hop's retained set has expected loss bounded below by `I(Y; S_{ell-1} | Q, r, ell)`, which is strictly positive whenever the retained set carries information about the gold label. The architectural fix is to feed a summary of `S_{ell-1}` into the scoring function for hop `ell`. CAFF does this with CSV (the summary) and DBM (a gating mechanism on the bilinear scorer). --- ## Approach CAFF operates in four stages during multi-hop evidence retrieval. ### Stage 1 -- BFS candidate stratification A BFS from the query's seed entities collects all triples reachable within `L=3` hops, stratified by hop depth. A per-relation frequency cap (`K_r=20`) prevents any one relation from dominating a candidate set on highly connected entities. ### Stage 2 -- Contextual Summary Vector (CSV) For each hop `ell > 1`, the retained set from the previous hop is summarized by a parameter-free, permutation-invariant pool over the relation embeddings of its triples: ``` z_{ell-1} = pool({ E[r] : (h, r, t) in S_{ell-1} }) ``` The default pool is `mean`. At `ell=1` the retained set is empty by convention, so `z_0 = 0` and the scorer reduces to its base form. ### Stage 3 -- Dynamic Bilinear Modulation (DBM) The base hop-conditioned bilinear scorer ``` s_base(Q, r, ell) = Q^T W_ell E[r] ``` is augmented with a low-rank, context-dependent perturbation generated from `z_{ell-1}`: ``` Delta_ell(z) = sigmoid(U z) * (A z) (B z)^T, with rank rho = 3.10 - CUDA 11.8 (an 8 GB consumer GPU is sufficient) - Git ### Setup ```bash # 1. Clone the repository git clone https://github.com/ /caff.git cd caff # 2. Create a clean virtual environment python -m venv .venv source .venv/bin/activate # Windows: .venv\Scripts\activate # 3. Install PyTorch pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cu118 # 4. Install remaining dependencies pip install -r requirements.txt # 5. Pre-download the BioLinkBERT-Large encoder (optional, also auto-downloads) python -c "from transformers import AutoModel; AutoModel.from_pretrained('michiyasunaga/BioLinkBERT-large')" ``` ### Core dependencies ``` torch>=2.0 transformers>=4.30 networkx>=3.0 numpy, scipy, scikit-learn, pandas tqdm, pyyaml ``` --- ## Data CAFF operates on a merged biomedical knowledge graph built from three public sources: **Orphanet** (rare-disease ontology, gene-disease links), **HPO** (phenotype ontology), and **OMIM** annotations (Mendelian inheritance, gene-phenotype). All three are publicly available and non-credentialed. ### Build the KG ```bash # 1. Convert raw ontologies to TSV python scripts/convert_orphanet_xml_to_tsv.py --in data/raw/orphanet/ --out data/processed/orphanet.tsv python scripts/convert_hpo_to_tsv.py --in data/raw/hpo/hp.obo --out data/processed/hpo.tsv # 2. Build base KG and merge in HPO/OMIM python scripts/build_kg.py --orphanet data/processed/orphanet.tsv --out data/processed/merged_kg.tsv python scripts/merge_hpo_into_kg.py --in data/processed/merged_kg.tsv --hpo data/processed/hpo.tsv --out data/processed/merged_kg_v2.tsv # 3. Sample QA records from the KG python scripts/build_orphanet_qa.py --kg data/processed/merged_kg_v2.tsv --n 20000 --out data/processed/ # 4. Pre-compute BFS candidates and gold annotations python scripts/extract_bfs.py --kg data/processed/merged_kg_v2.tsv --L 3 --K_r 20 python scripts/annotate_triples.py --kg data/processed/merged_kg_v2.tsv --qa data/processed/ ``` ### KG statistics | Property | Value | |---------------------------------|-----------| | Entities `|V|` | 38,456 | | Triples `|E|` | 291,335 | | Relation types `|R|` | 11 (after `min_relation_freq=50`) | | Maximum BFS hop depth `L` | 3 | | QA records (train/dev/test) | 14,000 / 3,000 / 3,000 | | Triple instances (train) | 473,471 (6.24% positive) | | Triple instances (test) | 102,317 (6.27% positive) | Triples on any shortest path from a seed entity to the gold answer entity receive label `y=1`; all others `y=0`. The test set is held out completely from training and threshold tuning. --- ## Training The default training configuration is `configs/no_dc.yaml`. To reproduce the headline numbers on three seeds: ```bash for s in 42 1337 2024; do python train.py --config configs/no_dc.yaml --seed $s done ``` Each seed takes about 40 minutes on an 8 GB consumer GPU (RTX 4060). `train.py` auto-detects CUDA and applies hardware-appropriate overrides (`micro_batch_size=4, grad_accum_steps=64, mixed_precision=fp16`, effective batch 256). Training is deterministic: the same seed produces bit-identical results across runs on the same hardware. ### Reproduce the full ablation suite ```bash # Train each variant on 3 seeds for cfg in no_dc caff_orphanet no_csv no_dbm no_freqcap caff_no_hc3 depthbilinear; do for s in 42 1337 2024; do python train.py --config configs/$cfg.yaml --seed $s done done # Per-hop threshold tuning on each variant for cfg in no_dc caff_orphanet no_csv no_dbm no_freqcap depthbilinear; do for s in 42 1337 2024; do python scripts/per_hop_threshold_sweep.py \ --config configs/$cfg.yaml \ --checkpoint runs/$cfg/seed_$s/best.pt done done ``` ### Key hyperparameters (`no_dc.yaml`) | Hyperparameter | Value | |-----------------------|-------:| | Optimizer | AdamW (weight decay 1e-2) | | Base learning rate | 3e-4 (cosine to 1e-5, 1-epoch warmup) | | Effective batch size | 256 | | Epochs | 10 (early stopping on dev F1, patience 5) | | Gradient clip | 1.0 | | Encoder | BioLinkBERT-large (340 M, frozen, d=1024) | | Random seeds | {42, 1337, 2024} | --- ## Evaluation ### Per-hop threshold sweep (the headline metric) ```bash python scripts/per_hop_threshold_sweep.py \ --config configs/no_dc.yaml \ --checkpoint runs/no_dc/seed_42/best.pt ``` Tunes per-hop retention thresholds on the dev set, then reports precision, recall, and F1 on the held-out test set under three regimes: global theta=0.50, global theta=0.80, and per-hop tuned thresholds. ### Paired bootstrap significance test ```bash python evaluate.py \ --checkpoint runs/no_dc/seed_42/best.pt \ --report-bootstrap-vs runs/caff_orphanet/seed_42/best.pt \ --mode autoregressive \ --output-json results/bench_no_dc_vs_full_seed_42.json ``` Reports test metrics (F1, MAP, NDCG@10, per-hop precision) at theta=0.80 in autoregressive inference mode (no leakage of gold relations from prior hops), plus a paired bootstrap on per-query AP versus the baseline checkpoint (10,000 resamples). --- ## Results All numbers below are measured on the held-out Orphanet QA test set (3,000 queries, 102,317 candidate triples), 3 seeds, deterministic. The full per-seed outputs are in `results/`. ### Default configuration (no_dc.yaml) | Metric | Mean +/- std | Inference mode | |------------------------------|--------------------|----------------| | Test F1 (per-hop) | **0.5764 +/- 0.0022** | teacher-forced | | Test F1 (autoregressive) | 0.5477 +/- 0.0006 | autoregressive | | Test MAP | 0.6741 +/- 0.0003 | autoregressive | | Test NDCG@10 | 0.7090 +/- 0.0003 | autoregressive | | Hop-1 precision | 0.8234 +/- 0.0047 | autoregressive | | Hop-2 precision | 0.4378 +/- 0.0006 | autoregressive | | Hop-3 precision | 0.2426 +/- 0.0026 | autoregressive | Per-hop is the headline metric: thresholds are tuned per hop on the dev set, then applied unchanged on test. Autoregressive is reported separately because it does not leak gold relations from prior hops during inference; the F1 gap of about 0.029 between the two modes is the cost of realistic deployment. ### Statistical significance versus alternative configurations Paired bootstrap on per-query AP, 10,000 resamples, computed per seed: | Comparison | delta_AP | 95% CI | p-value | |----------------------------------|----------------:|-----------------------|--------:| | no_dc vs caff_orphanet, seed 42 | +0.0295 | [+0.0251, +0.0340] | 0.0000 | | no_dc vs caff_orphanet, seed 1337| +0.0227 | [+0.0188, +0.0267] | 0.0000 | | no_dc vs caff_orphanet, seed 2024| +0.0227 | [+0.0190, +0.0267] | 0.0000 | | **mean** | **+0.0250** | (each CI excludes 0) | /seed_ /`. --- ## Hyperparameters The default configuration (`no_dc.yaml`) uses: | Symbol | Meaning | Value | |----------------------|--------------------------------------------------|------:| | `d` | Embedding dimension (BioLinkBERT-Large output) | 1024 | | `L` | Maximum BFS hop depth | 3 | | `rho` | DBM rank | 16 | | `theta` | Retention threshold (global default) | 0.80 | | `K_r` | Frequency cap per relation per head | 20 | | `lambda_C` | HC3 loss weight | 0.35 (inert) | | `lambda_D` | Depth-contrastive loss weight | **0.0** (default; 0.40 disabled) | | `min_relation_freq` | Drop singleton relations at KG load | 50 | | `gamma_C` | HC3 margin | 0.25 | | `gamma_D` | Depth-contrastive margin | 0.20 | A theta sensitivity analysis and a lambda_D sweep are listed under future work. --- ## Reproducibility - All results are **mean across three seeds**...