KuaiRand Autonomous ML Research Agent

Inspiration

Video recommendation ranking is hard. You can throw a fancy ML model at it, but production datasets are noisy, biased by what users saw, and full of edge cases. The KuaiRand-Pure benchmark is a real-world TikTok dataset with exactly those problems — a perfect testbed for autonomous research.

We wanted to build something that the competition explicitly asks for: an autonomous ML research agent that improves a recommendation model's score on KuaiRand through its own iteration loop, not just a hand-tuned model. Most teams would train one model, tune hyperparameters manually, and submit. We decided to build a system that proposes, trains, evaluates, and iterates automatically, making real decisions based on what it learns.

What it does

The agent is a self-driving ML pipeline for ranking optimization:

Core task: Within-user ranking — given a user's past watched videos, rank their own future interactions to predict which ones they'll actually watch long enough to finish.

The autonomous loop:

Read prior run results from the log
Propose a hypothesis (deterministic grid or Claude API reasoning)
Apply it as a config change (e.g., "use IPS weights" or "enable sequence encoder")
Train the model on KuaiRand-Pure
Evaluate on validation only (never peek at test)
Log the result with metrics (GAUC, nDCG@5, primary = mean of both)
Check convergence: if validation primary hasn't improved by ε=0.002 over 3 consecutive iterations, stop
Repeat

Evaluation: Two metrics with a clear definition:

GAUC (grouped AUC): per-user AUC, then averaged across users
nDCG@5 (normalized discounted cumulative gain): ranking quality at top-5, using 2^rel − 1 gain
Primary metric: mean(GAUC, nDCG@5) — both must improve together

Baseline to beat: Official FM (Factorization Machine) baseline scores 0.5946 primary on held-out test. Ceiling (perfect ranking): 0.8645 (nDCG's true ceiling is only 0.7289 because 27.1% of test users have no positive labels).

What we achieved: 5-seed DeepFMMTL ensemble reached 0.5983 primary on test — +0.0037 over baseline, real but modest, and reported honestly with multi-seed variance shown, not cherry-picked. The same pipeline, retargeted at the bonus KuaiRand-1K benchmark, scored 0.6408 primary on its test set (single-seed, no official baseline exists there to compare against).

How we built it

The orchestration layer

agent/orchestrator.py is the heartbeat. It:

Reads logs/run_log.jsonl (each line: {iteration, config, metrics, timestamp})
Calls either agent/hypothesis.py (deterministic grid) or agent/llm_hypothesis.py (Claude API)
Launches python src/train.py --config <proposed_config>
Runs scripts/eval_checkpoint.py to score on validation
Logs one structured entry
Checks if we've converged; if not, repeats

Ran to genuine convergence: 25 iterations, stopped because validation primary plateaued (not because it hit the 50-iteration cap).

Two hypothesis proposers

Deterministic (agent/hypothesis.py):

Grid over model variants: base DeepFM, +MultiTask, +IPS, +Sequences, +BPR, combinations
No API key needed, always available
Simple but covers the design space we cared about

LLM-driven (agent/llm_hypothesis.py):

Calls Claude Sonnet 4.5 API with full run history as context
Claude reasons: "What should we try next based on these results?"
Claude picks from a fixed, type-checked allowlist of config keys: use_mtl, use_ips, use_seq, use_bpr, learning_rate, batch_size, seed
Deliberately scoped: LLM can never write code, never touch src/data.py (splits), never touch src/evaluate.py (scoring), never access test data
If the API call fails → orchestrator logs the exception, falls back to grid, keeps going ✓

Techniques explored (branched teams)

Five team members, each owning a branch and file responsibility:

Branch	Owner	Technique	Files	Impact
`feat/harness`	Rishi	Baseline reproduction, integration	`agent/orchestrator.py`, `agent/convergence.py`, `src/data.py`	Coordinated the whole loop
`feat/multitask`	Min	Multi-task DeepFM	`src/models/deepfm_mtl.py`, `src/features/base.py`	+0.0017 primary vs. FM (small but real)
`feat/ips-debias`	Vidush	Inverse propensity scoring	`src/features/propensity.py`	Honest finding: flat or slightly negative on valid
`feat/sequential-ssl`	Nandit	GRU sequence encoder + SSL pretraining	`src/models/seq_encoder.py`, `src/features/sequential.py`	+0.0019 primary (n=2, needs more seeds)
`feat/eval-submission`	Sarthak	Rank-average ensembling	`src/models/ensemble.py`, `src/submit.py`	Final submission: 5-seed blend

KuaiRand-specific innovation: IPS + off-policy eval

KuaiRand-Pure uniquely ships two logs:

Production policy log (biased, what TikTok's algorithm showed)
Randomized-exposure log (uniformly random, unbiased labels)

We:

Built propensity estimates from the random log: P(exposed | video features)
Weighted training data by 1/propensity (IPS reweighting)
Built a genuinely unbiased evaluation set from random-log dates in the official validation window
Measured if IPS actually helps against ground truth

Honest finding: IPS cost ~0.002 primary on the (biased) validation set — a real, reproducible gap (~6× the combined std), not noise. Its counterfactual benefit remains unconfirmed either way once checked against a genuinely unbiased off-policy set (below).

Model stack

DeepFMMTL (Multi-task Deep Factorization Machine):

Shared embedding layer (5 feature domains: user_id, video_id, author_id, tab, duration bucket)
FM interaction matrix
Deep tower (2 hidden layers, 64 → 32 → 1)
Primary head: long_view (binary, BCE). Auxiliary heads (click/like/play_time) exist in the architecture and compute_loss already supports weighting them in (default aux weight 0.3) — but were never actually fed real labels in any run reported here, since src/data.py doesn't expose is_click/is_like/play_time_ms yet. Structurally ready, not yet multi-task in practice — see "What's next."

Sequence encoder (optional, +50s/epoch):

GRU over user's past k=30 videos
Output: sequence embedding, concatenated to static features
SSL pretraining: item2vec-style (video → video embeddings via skip-gram)

Loss functions tested:

Pointwise BCE (default, works best)
Pairwise BPR (Bayesian Personalized Ranking): per-user (pos, neg) pairs — underperforms BCE despite theoretically optimizing ranking, because it trains on fewer effective examples (~383K pairs/epoch vs. 1.14M rows)

Bonus: KuaiRand-1K

Same pipeline, retargeted at the bonus KuaiRand-1K benchmark (base DeepFMMTL only — IPS/BPR were tested and found not to transfer cleanly from Pure-tuned hyperparameters to 1K's larger scale). Single seed rather than Pure's 5-seed ensemble, given 1K's much larger per-run cost (~11.7M-row standard logs vs. Pure's ~1.4M).

	GAUC	nDCG@5	primary
1K submission (single-seed DeepFMMTL)	0.6738	0.6079	0.6408

4,132,081 test rows, passes submit.py --check. No official 1K baseline exists to compare against (unlike Pure's 0.5946), and the raw magnitude isn't comparable to Pure's anyway — 1K's much denser per-user interaction history makes within-user ranking structurally easier, not evidence of a better model. Only directional findings from 1K carry over.

The trained checkpoint (~199MB) and submission file (~121MB) both exceed GitHub's 100MB limit, so they're hosted on Drive instead of committed: https://drive.google.com/drive/folders/1gS479mnNW6D4zYvcWMASsrR10VqlXRRi — both are exactly reproducible from the pipeline anyway.

Challenges we ran into

Noise floor is high, signals are small

Baseline FM: 0.6016 (valid)
DeepFMMTL: 0.6033 ± 0.0003
Difference: 0.0017 ± 0.0004 (4× the gap!)

Had to run 5+ seeds per variant before confidently saying if something helped or hurt. A single run that looks great is just luck.

IPS doesn't work the way textbooks suggest

Inverse propensity scoring is theoretically beautiful for debiasing exposure bias. But:

On the biased validation set: IPS cost ~0.002 primary (a real, ~6× the combined std, gap at n=5)
We built an unbiased off-policy evaluation set from the random log to check if IPS helps against ground truth
Result: IPS still flat-to-slightly-worse, 0.0002 gap (~0.29× combined std, indistinguishable from noise)

Conclusion: We're honest. The counterfactual correction question is genuinely open.

Pairwise ranking loss backfired

BPR (Bayesian Personalized Ranking) directly optimizes what GAUC measures: per-user ranking of positive vs. negative items. Theoretically perfect. Practically:

BPR underperforms BCE by ~0.002 (4.4× combined std over 5 seeds)
Even loses on GAUC itself, the metric it's supposedly optimizing
Root cause: fewer training examples per epoch (only pos/neg pairs, ~383K vs. 1.14M rows)
Natural next step: hybrid loss (BCE + BPR term, not instead) — not yet tried

Sequence encoder is slow

GRU forward/backward over padded histories: ~50s/epoch
Plain DeepFM: ~1s/epoch, 50× faster
Made sequences the bottleneck; we only ran n=2 seeds before time ran out

Test-set discipline is hard

Easy to overfit if you touch test labels during model selection
We never did: model selection used valid only
Off-policy eval used its own random-log-based set from official valid dates only
Final test scoring done once, on the final 5-seed ensemble

Accomplishments we're proud of

Reproducible, honest wins

Beat official FM baseline (0.5946) with 0.5983 on genuinely held-out test
+0.0037 primary — not huge, but real (inside the noise floor would have been 0.5956-0.5966)
Multi-seed ensembling: reduces variance by ~0.0007 vs. any single seed
Same pipeline retargeted at bonus KuaiRand-1K, no code changes beyond config: 0.6408 primary on its held-out test set, single seed, passes submit.py --check on 4.1M rows
Reported with full methodology, all seeds, all failures included

A genuinely autonomous agent

The orchestrator ran to convergence (25 iterations) without human intervention. When the LLM API had no credentials, it:

Caught the exception
Logged it verbatim
Fell back to deterministic grid
Completed the iteration normally
Continued the loop

That's robustness. Not "never fails," but "fails gracefully and keeps going."

Rigorous experimental methodology

5+ seeds per variant before claiming wins/losses
Off-policy evaluation on random-exposure labels (not a simulation, real unbiased data)
Multi-variant ablation (tried 6 variants, documented all of them)
Baseline and ceiling metrics so you see actual headroom (0.8645, not 1.0)

Team structure that scales

Clear ownership:

One branch per technique (feat/ips-debias, feat/sequential-ssl, etc.)
One owner per branch (Vidush owns propensity.py, Nandit owns seq_encoder.py)
CI gating: .github/workflows/check.yml runs tests on every PR
main merge-only, reviewed by Rishi
Five people, minimal friction

Transparent failure reporting

Included a section "Already tried, no lift — don't repeat":

+13 feature domains from CWM dataset: no difference
Embedding dimensions k=8/16/32: barely moves needle
Documented what didn't work so future researchers don't retrace steps

What we learned

1. Seed variance is your enemy

One run: 0.6035, looks great! Five runs: 0.6033 ± 0.0003, turns out it was noise. Always use n≥5 for any binary claim. A single seed is not evidence.

2. Theory ≠ practice (but check both)

IPS sounds perfect on paper (correct for exposure bias via importance weighting). Didn't help empirically.

BPR directly optimizes ranking pairs. Lost to simpler BCE in practice.

Auxiliary tasks (clicks, likes, play_time_ms) should help multi-task learning. Still waiting for real gains.

Empirical evidence is the only truth. But we also check it rigorously — not just one seed.

3. Off-policy evaluation is hard to get right

We had random-exposure labels (genuinely unbiased). But the base rate of positives under random exposure is so different from the production-policy log that absolute scores become meaningless (0.37 vs. 0.60). Only relative ordering between variants matters — and even there, n=2 isn't enough to claim something works.

4. Ensembling repeated seeds is a free lunch (but not what you think)

5-seed DeepFMMTL averaged: 0.6040 5-seed DeepFMMTL + 2-seed sequences averaged: 0.6040 (no improvement!) All 19 checkpoints (all 5 techniques, mixed seeds): 0.6040

The lesson: The gain from ensembling is pure variance reduction, not complementary diversity. You're buying accuracy by averaging the same model trained 5 times, not by combining different techniques. Both are valuable, but the story matters for what to try next.

5. LLM agents need guardrails, not cleverness

Letting Claude propose any 1-line Python is dangerous. Letting Claude pick config keys from a validated allowlist is both safe and useful:

Safe: no code execution, no data access, no split-logic changes
Useful: Claude can still reason over run history and pick good experiments

Scope is the security boundary, not intelligence.

What's next for kuairand-agent

Short-term (within the scope of this challenge)

Multi-task labels: They're defined but not exposed in src/data.py. Wiring in is_click, is_like, play_time_ms as true auxiliary targets (not just dummy columns) could unlock the Priority 1 headroom.
Hybrid BCE+BPR loss: BPR alone loses to pointwise BCE; adding BPR as a regularization term on top (not instead) is theoretically motivated and hasn't been tried.
Sequence encoder scaling: Need 3+ more seeds (n=2 now) to confirm if sequences help or are just noise at scale.

Medium-term

Real LLM hypothesis generation at scale: With credentials available, run --proposer llm end-to-end and compare experiment selection against the deterministic grid. This is the proof-of-concept we couldn't demo here.
Expand to the remaining KuaiRand benchmark: KuaiRand-1K is done (0.6408 primary, see above); KuaiRand-27K (9.9GB) was attempted but its archive didn't finish downloading within this project's time budget, and even a completed download would mean training ~10x Pure's row count on CPU-only hardware — consistent with the playbook's own guidance to treat it as lowest priority if time is short.

Long-term

Package for the community: Container + CLI tool for any researcher working on recommendation ranking to run their own autonomous experiments
Open agent to more researchers: Make it easy to add new techniques (new model files, new loss functions) without rewriting the orchestrator
Production readiness: Move from validation-only evaluation to production metrics, handle cold-start users, streaming inference

Technical foundations

Development tools used

Claude Code (Anthropic) — agentic CLI, used throughout for implementation, debugging, and running experiments across team branches
VS Code — primary editor
Git / GitHub — version control, branch-per-role, PR-gated merges with CI checks
Docker — Dockerfile pins python:3.12-slim + requirements.txt as the cross-platform tiebreaker when results don't reproduce locally

APIs and libraries

Anthropic Claude API (sonnet-4-5-20250929) — powers optional LLM hypothesis proposer only
PyTorch (2.9.1) — DeepFMMTL, GRU encoder, all training loops (pointwise BCE and pairwise BPR)
NumPy (2.1.3) — data loading, SSL pretraining, evaluation math
PyYAML (6.0.2) — config file loading
pytest (8.3.4) — test suite (11 files, ~60 tests) gating every push
anthropic (1.2.0) — official Python SDK for Claude API

Data and evaluation

KuaiRand-Pure (required benchmark from Zenodo 10439422)
- Production-policy log (biased) + randomized-exposure log (unbiased)
- Splits: train 20220408–20220421 / valid 20220422–20220428 / test 20220429–20220508
- Official evaluation code (src/evaluate.py, frozen, not modified)
KuaiRand-1K (bonus, attempted and submitted)
KuaiRand-27K (bonus, attempted but incomplete within time budget)
No external data, no synthetic labels, no manual annotation

How our solution addresses the problem statement

The challenge asks for: an autonomous ML research agent that improves a recommendation model through its own iteration loop.

We delivered:

✓ Real code, not narrative — agent/orchestrator.py implements the full cycle, ran to genuine convergence

✓ Two swappable proposers — Deterministic grid (always available) + Claude API-driven (when credentials present), both type-checked and sandboxed

✓ Failure resilience — Caught a real missing-credentials error live, fell back gracefully, kept going

✓ KuaiRand-specific techniques — Propensity weighting + unbiased off-policy evaluation using the random-exposure log

✓ Honest reporting — Every technique with multi-seed std, including ones that didn't help, not cherry-picked wins

✓ Measurable improvement — Beat the baseline (0.5946 → 0.5983) on held-out test, real but modest, reported with full methodology

Built With

api-integration
autonomous-ml
ci-cd
deepfm
docker
ensemble-methods
kuairand-benchmark
llm-reasoning
pytorch
ranking-models
recommendation-systems