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