File size: 14,570 Bytes
bca9583 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 | # ResearchIT Phase 6 β LightGBM Reranker
> **Status:** Pipeline tested on synthetic data. Awaiting real citation edges from Semantic Scholar API to produce the production model.
> **Parent project:** [siddhm11/ResearchIT](https://huggingface.co/spaces/siddhm11/ResearchIT)
> **Replaces:** Hand-tuned heuristic scorer in `app/recommend/reranker.py`
---
## What This Is
A LightGBM lambdarank model that replaces the hand-tuned heuristic scorer in ResearchIT's recommendation pipeline. The heuristic uses 5 features with fixed weights. This model uses 37 features and learns optimal weights from data.
**This repo contains:**
- `scripts/` β 3 Python scripts that build the full pipeline (data β train β evaluate)
- `synthetic_model/` β A model trained on synthetic data (proof the pipeline works)
- `test_results.json` β Full benchmark results from the test suite
**This repo does NOT yet contain:**
- A production model trained on real citation data (that's the next step β run Script 1 with your S2 API key)
---
## How It Works
### The Problem
ResearchIT recommends arXiv papers using a multi-stage pipeline:
```
Qdrant ANN retrieval β Quota fusion β Reranking β MMR diversity β Feed
```
The **reranking** step currently uses a heuristic scorer:
```python
score = 0.40 Γ cos(paper, long_term_profile)
+ 0.25 Γ cos(paper, short_term_profile)
+ 0.15 Γ recency_decay
+ 0.10 Γ retrieval_rank_confidence
- 0.15 Γ cos(paper, negative_profile)
```
These weights are hand-tuned guesses. The model can't use citation count, co-citation, category match, or any feature interactions.
### The Solution
Train a LightGBM lambdarank model on citation-graph pseudo-labels:
- **Each arXiv paper acts as a "pseudo-user"** β its reference list simulates what that researcher would "save"
- **Direct citations β label 2** (strong positive)
- **Co-citations β label 1** (weak positive β papers that share references)
- **ANN-retrieved but not cited β label 0** (negative)
The model learns: given 37 features about a (user, paper) pair, which papers should rank higher?
---
## The 3-Script Pipeline
### Script 1: `01_fetch_citation_edges.py`
**What:** Fetches the citation graph from Semantic Scholar API.
**Input:** `arxiv_ids.txt` (your 1.6M arXiv IDs, one per line)
**Output:** `citations.parquet` β (citing_arxiv_id, cited_arxiv_id, is_influential)
```bash
python scripts/01_fetch_citation_edges.py \
--corpus-file arxiv_ids.txt \
--output citations.parquet \
--api-key YOUR_S2_API_KEY # optional, faster with key
```
- Supports checkpoint/resume β safe to interrupt
- Supports both batch API (simple) and bulk download (faster for 1.6M papers)
- Filters to in-corpus edges only (both papers must be in your Qdrant collection)
- Expected output: ~20-40M citation edges
### Script 2: `02_generate_training_triples.py`
**What:** Generates labeled training data by running ANN searches against Qdrant.
**Input:** `citations.parquet` + Qdrant access + Turso access
**Output:** `ltr_dataset/train.parquet` + `ltr_dataset/eval.parquet` + `feature_schema.json`
```bash
python scripts/02_generate_training_triples.py \
--citations citations.parquet \
--corpus-file arxiv_ids.txt \
--qdrant-url "https://YOUR_QDRANT_URL" \
--qdrant-api-key "YOUR_KEY" \
--qdrant-collection arxiv_bgem3_dense \
--turso-url "https://YOUR_TURSO_URL" \
--turso-token "YOUR_TOKEN" \
--output-dir ./ltr_dataset \
--num-queries 100000 \
--candidates-per-query 50
```
- Enforces time-split: train on papers before 2023, eval on 2023+
- Asserts no temporal leakage: `max(train.year) < min(eval.year)`
- Uses `scroll()` + `query_points()` for Qdrant (correct API for qdrant-client 1.17+)
- Expected output: ~5M train rows, ~500K eval rows
### Script 3: `03_train_lightgbm.py`
**What:** Trains LightGBM and compares against the heuristic baseline.
**Input:** `train.parquet` + `eval.parquet`
**Output:** `reranker_v1.txt` (~100-300KB model file) + metrics
```bash
python scripts/03_train_lightgbm.py \
--train-file ltr_dataset/train.parquet \
--eval-file ltr_dataset/eval.parquet \
--output-dir ./model_output \
--num-boost-round 500 \
--learning-rate 0.05
```
- Evaluates: nDCG@5, nDCG@10, nDCG@20, Recall@10, Recall@50, HR@10, MRR
- Compares LightGBM vs the exact heuristic from `reranker.py`
- Reports feature importance, latency benchmark, per-query win rates
- Outputs verdict: deploy / marginal / no improvement
---
## Feature Schema (37 features)
The model uses 37 features organized in 4 groups:
### Content/Retrieval Features (0-19) β Populated during pseudo-label training
| # | Feature | Description | Source |
|---|---------|-------------|--------|
| 0 | `qdrant_cosine_score` | BGE-M3 cosine similarity from ANN search | Qdrant |
| 1 | `candidate_position` | Rank position in ANN results (0-indexed) | Qdrant |
| 2 | `candidate_citation_count` | Total citation count | Turso |
| 3 | `candidate_log_citations` | log(citation_count + 1) | Computed |
| 4 | `candidate_influential_citations` | Influential citation count (Semantic Scholar) | Turso |
| 5 | `candidate_age_days` | Days since publication | Turso |
| 6 | `candidate_recency_score` | exp(-0.002 Γ age_days) β matches heuristic formula | Computed |
| 7 | `query_citation_count` | Citation count of the query/user paper | Turso |
| 8 | `query_age_days` | Days since query paper was published | Turso |
| 9 | `year_diff` | |query_year - candidate_year| | Computed |
| 10 | `same_primary_category` | 1 if same primary arXiv category, else 0 | Turso |
| 11 | `co_citation_count` | Papers that cite BOTH query and candidate | Citation graph |
| 12 | `shared_author_count` | Number of shared authors (case-insensitive) | Turso |
| 13 | `candidate_is_newer` | 1 if candidate published after query | Computed |
| 14 | `query_log_citations` | log(query_citation_count + 1) | Computed |
| 15 | `citation_count_ratio` | candidate_citations / (query_citations + 1) | Computed |
| 16 | `age_ratio` | candidate_age / (query_age + 1) | Computed |
| 17 | `candidate_citations_per_year` | citation_count / max(age_years, 0.5) | Computed |
| 18 | `query_num_references` | How many papers the query paper cites (in-corpus) | Citation graph |
| 19 | `candidate_num_cited_by` | How many corpus papers cite the candidate | Citation graph |
### User Behavior Features (20-30) β Zero-filled for pseudo-labels, active for real users
| # | Feature | Description | Source in ResearchIT |
|---|---------|-------------|---------------------|
| 20 | `ewma_longterm_similarity` | cos(candidate, user long-term EWMA profile) | `profiles.py` Ξ±=0.03 |
| 21 | `ewma_shortterm_similarity` | cos(candidate, user short-term EWMA profile) | `profiles.py` Ξ±=0.40 |
| 22 | `ewma_negative_similarity` | cos(candidate, user negative EWMA profile) | `profiles.py` Ξ±=0.15 |
| 23 | `cluster_importance` | Importance weight of the serving cluster | `clustering.py` |
| 24 | `cluster_distance_to_medoid` | cos(candidate, cluster medoid embedding) | `clustering.py` |
| 25 | `is_suppressed_category` | 1 if candidate's category is suppressed (β₯3 dismissals in 14 days) | `db.py` |
| 26 | `onboarding_category_match` | 1 if candidate matches user's onboarding selections | `db.py` |
| 27 | `user_total_saves` | Total papers user has saved | `interactions` table |
| 28 | `user_total_dismissals` | Total papers user has dismissed | `interactions` table |
| 29 | `user_days_since_last_save` | Days since user's most recent save | `interactions` table |
| 30 | `user_session_save_count` | Saves in current session | In-memory state |
### Cross Features (31-36) β Feature interactions
| # | Feature | Formula |
|---|---------|---------|
| 31 | `cosine_x_recency` | qdrant_cosine_score Γ candidate_recency_score |
| 32 | `cosine_x_citations` | qdrant_cosine_score Γ candidate_log_citations |
| 33 | `category_x_recency` | same_primary_category Γ candidate_recency_score |
| 34 | `cosine_x_cocitation` | qdrant_cosine_score Γ log(co_citation_count + 1) |
| 35 | `position_inverse` | 1 / (candidate_position + 1) |
| 36 | `citations_x_recency` | candidate_log_citations Γ candidate_recency_score |
### Why 37 Features Instead of 5?
The heuristic uses features 0, 1, 5 (via recency), 20, 21 β that's it. LightGBM additionally uses:
- **Citation count** (#2, #3, #17) β papers with more citations are generally higher quality
- **Co-citation count** (#11) β the strongest signal in our tests. Papers cited by the same community as your interests are highly relevant
- **Category match** (#10) β same arXiv category = topically aligned
- **Shared authors** (#12) β self-citation / collaborator network signal
- **Feature interactions** (#31-36) β e.g., a paper that is both semantically similar AND frequently co-cited is much more relevant than either signal alone
---
## Benchmark Results (Synthetic Data)
> These results are from realistic synthetic data that mimics citation graph patterns. Real data results will differ.
### LightGBM vs Heuristic Baseline
| Metric | Heuristic | LightGBM | Improvement |
|--------|:---------:|:--------:|:-----------:|
| nDCG@3 | 0.9305 | **0.9979** | +7.2% |
| nDCG@5 | 0.9138 | **0.9980** | +9.2% |
| **nDCG@10** | **0.9111** | **0.9985** | **+9.6%** |
| nDCG@20 | 0.9472 | **0.9988** | +5.5% |
### Per-Query Win Rate (500 eval queries)
- LightGBM wins: **91.4%**
- Heuristic wins: 0.4%
- Ties: 8.2%
### Production Metrics
| Metric | Value | Target | Status |
|--------|-------|--------|--------|
| Latency (100 candidates) | 0.088ms | <1ms | β
11Γ under budget |
| Model size | 286 KB | <500 KB | β
|
| Model reload | Identical predictions | β | β
|
| Handles NaN input | Graceful | β | β
|
| Handles extreme values | No crash | β | β
|
| Train-eval gap | 0.0008 | <0.05 | β
No overfitting |
### Top 10 Features by Importance
| Rank | Feature | Importance |
|------|---------|:----------:|
| 1 | cosine_x_cocitation | 95,386 |
| 2 | qdrant_cosine_score | 37,300 |
| 3 | co_citation_count | 14,124 |
| 4 | candidate_citation_count | 3,399 |
| 5 | candidate_position | 2,615 |
| 6 | cosine_x_citations | 1,623 |
| 7 | year_diff | 1,229 |
| 8 | candidate_log_citations | 894 |
| 9 | age_ratio | 564 |
| 10 | position_inverse | 427 |
The #1 feature is `cosine_x_cocitation` β the interaction between embedding similarity and co-citation count. This is a signal the heuristic cannot access.
---
## How to Produce the Production Model
### Prerequisites
```bash
pip install httpx pyarrow tqdm numpy qdrant-client lightgbm
```
### Step 1: Export corpus IDs
Export your 1.6M arXiv IDs from Turso to a text file:
```sql
SELECT arxiv_id FROM papers;
```
Save as `arxiv_ids.txt`, one ID per line.
### Step 2: Run the pipeline
```bash
# Fetch citation edges (~30-90 min)
python scripts/01_fetch_citation_edges.py \
--corpus-file arxiv_ids.txt \
--output citations.parquet
# Generate training data (~1-2 hours, needs Qdrant + Turso access)
python scripts/02_generate_training_triples.py \
--citations citations.parquet \
--corpus-file arxiv_ids.txt \
--qdrant-url "$QDRANT_URL" \
--qdrant-api-key "$QDRANT_API_KEY" \
--turso-url "$TURSO_URL" \
--turso-token "$TURSO_DB_TOKEN" \
--output-dir ./ltr_dataset
# Train (~5 min)
python scripts/03_train_lightgbm.py \
--train-file ltr_dataset/train.parquet \
--eval-file ltr_dataset/eval.parquet \
--output-dir ./model_output
```
### Step 3: Deploy
Copy `model_output/reranker_v1.txt` to your ResearchIT Space.
---
## Integration Into ResearchIT
The model integrates into `app/recommend/reranker.py` with minimal changes:
```python
# Load once at startup
import lightgbm as lgb
_lgb_model = None
try:
_lgb_model = lgb.Booster(model_file="reranker_v1.txt")
print("[reranker] LightGBM model loaded")
except Exception:
print("[reranker] Using heuristic fallback")
# In rerank_candidates():
features = compute_features_v2(...) # expanded to 37 features
if _lgb_model is not None:
scores = _lgb_model.predict(features)
else:
scores = heuristic_score(features) # existing fallback β always works
```
The heuristic scorer remains as a fallback. If the model file is missing or fails to load, the system silently uses the heuristic. No user-facing impact.
This is **LightGBM-1** in the Doc 07 multi-stage architecture:
```
Qdrant ANN β LightGBM-1 (this model) β [TinyBERT β LightGBM-2] (Phase 8b, future)
```
---
## Time-Split Evaluation
Training data uses a strict temporal split to prevent leakage:
- **Train:** Query papers published before 2023-01-01
- **Eval:** Query papers published on or after 2023-01-01
- **Assertion:** `max(train.year) < min(eval.year)` β verified in Script 2
This means the model never sees future papers citing past papers during training β matching how it would perform in production on newly published papers.
---
## Known Limitations
### Citation β User Interest
Citation pseudo-labels ("this author cited that paper in their bibliography") are different from real user signals ("this user saved that paper in their feed"). A foundational paper like "Attention Is All You Need" gets label=2 in citation data but might be dismissed by users who've already read it.
**Mitigation:** The `candidate_log_citations` and `candidate_citations_per_year` features help LightGBM learn a popularity penalty. When 500+ real user interactions accumulate, retrain on actual save/dismiss data β the 11 user behavior features (20-30) activate and the model learns real preferences.
### Synthetic vs Real
The benchmark numbers above are from synthetic data. Real citation data will have:
- Noisier labels (not all citations are "relevant" β some are obligatory references)
- Lower nDCG improvement over heuristic (expect +3-5% on real data vs +9.6% on synthetic)
- Different feature importance rankings
---
## Dependencies
```
lightgbm>=4.0
httpx>=0.24
pyarrow>=12.0
numpy>=1.24
qdrant-client>=1.17
tqdm>=4.65
```
---
## References
- ResearchIT Doc 06 Β§3.1 β LightGBM lambdarank architecture decision
- ResearchIT Doc 07 Β§A6 β Time-split evaluation protocol
- ResearchIT Doc 07 Β§B.4 β Multi-stage reranker architecture
- PinnerSage (Pal et al., KDD 2020) β Ward clustering + importance-weighted retrieval
- Taobao ULIM (Meng et al., RecSys 2025) β Quota allocation, +5.54% CTR
- YouTube (Xia et al., 2023) β 3Γ gain from negative signals in reranking
- Bruch et al. (SIGIR 2022) β RRF optimizes Recall not nDCG
|