RankAgent

An autonomous multi-agent AI research team that invents, writes, tests, and self-heals recommendation system models directly in PyTorch code.

Comment

RankAgent

1. Project Description

Problem Statement

The KuaiRand-Pure benchmark evaluates within-user ranking of micro-video impressions for recommender systems. The primary target is predicting long_view (watch-time completion threshold) scored by the arithmetic mean of Group AUC (GAUC) and Normalized Discounted Cumulative Gain at rank 5 (nDCG@5):

$$\text{Primary Score} = \frac{\text{GAUC} + \text{nDCG@5}}{2}$$

The challenge presents several real-world ML difficulties:

  • Extreme Categorical Sparsity & Cold-Start: High-cardinality user/video ID spaces with sparse positive interactions.
  • Duration Confounding & Label Leakage: long_view is strongly correlated with raw watch duration. Exposing post-impression duration signals on test sets creates severe label leakage.
  • Ragged User Groupings: Impressions vary across users, making standard batch pointwise objectives sub-optimal compared to true within-user listwise ranking.
  • Autonomous Exploration Bottlenecks: Exploring 7 distinct research dimensions (loss functions, neural architectures, feature transformations, capacity, multi-task learning, sequence modeling, and optimization schedules) requires structured multi-agent collaboration rather than simple hyperparameter sweeps.

Solution Approach

RankAgent is an autonomous, headless machine learning research agent that operates over the action space of direct source code generation, transactional patching, and empirical validation:

  1. Multi-Agent Research Team: Rather than relying on a single monolithic prompt, RankAgent coordinates 5 specialized agents over typed Pydantic contracts:
    • Product Manager: Directs exploration and exploitation across 7 research dimensions and rotates strategies when improvements plateau.
    • ML Researcher: Formulates scientific hypotheses citing mechanisms from an immutable recommendation systems knowledge base.
    • Engineer: Authors SEARCH/REPLACE AST-valid code diffs in mutable files (pipeline/models.py, pipeline/features.py, pipeline/train.py), registers new models/losses in CLI registries, and executes self-healing error repairs (up to 2 attempts).
    • Feature Steward: Explores and validates tabular and sequential feature transformations.
    • QA / Verifier: Enforces pre-flight static leak audits, import allowlists, and metric score validity.
  2. Safety & Leakage Verification Boundary:
    • Sealed Test Set: Post-impression outcomes (play_time_ms, auxiliary feedback) on the test split are redacted with sentinel -1 values in pipeline/data.py.
    • Dynamic Mutation Audits: pipeline/feature_agent.py programmatically mutates labels to verify that computed features remain strictly invariant.
    • Process Isolation: Each experiment executes in an isolated child process with filtered environments (API keys stripped to prevent exfiltration) and OpenMP thread safety.
  3. Empirical Results:
    • Automatically reproduces the official Factorization Machine baseline (0.6015 vs. published 0.6016).
    • Autonomously discovers improved neural architectures (such as cross_attention achieving 0.6041, $+0.0026$ / $+3.2\sigma$ over baseline).

2. Development Tools Used

  • Visual Studio Code (VSCode): Primary Integrated Development Environment (IDE) used for codebase development, refactoring, multi-file navigation, and interactive debugging.
  • Git & GitHub: Version control system used for tracking experiment branches, pull requests, and commit history auditing across multi-agent experiments.
  • Windows PowerShell & Unix Bash: Command-line environments used for cross-platform subprocess orchestration, script execution, and test running.
  • Pytest: Automated testing framework used for unit testing, prompt verification, patch transaction tests, and safety gate regression testing (86 passing test cases).

3. APIs Used

  • OpenAI API (gpt-4o, gpt-4o-mini): Primary reasoning engine for iterative hypothesis generation, RecSys literature citation alignment, code patch authoring, and self-healing traceback analysis.
  • Anthropic API (claude-3-5-sonnet): Supported alternative Large Language Model provider for autonomous reasoning and source code synthesis.

4. Libraries and Frameworks Used

  • PyTorch (torch, torch.nn, torch.optim): Deep neural ranking models (Factorization Machines, DeepFM, DIN, MMoE, DCN-v2, PLE, Cross-Attention), tensor computations, GPU/CPU acceleration, dynamic autograd computation graphs, and custom ranking loss implementations.
  • LightGBM: Gradient Boosted Decision Tree (GBDT) engine utilizing the LambdaMART (lambdarank) ranking objective with top-5 truncation over causal engagement statistics.
  • NumPy & SciPy: Vectorized data manipulation, matrix operations, numerical stability guards, and official NumPy FM baseline reproduction.
  • Scikit-Learn: Categorical encoding (ordinal/one-hot), data scaling utilities, and baseline evaluation metrics.
  • Pandas: Tabular data manipulation, log analysis, and submission CSV export.
  • Pydantic (v2): Strict data modeling, runtime type validation, agent communication contracts (ResearchContext, TrialSpec), and structured JSON parsing.
  • PyYAML: Parsing configuration files for agent policies and benchmark constraints (configs/).

5. Datasets and Assets Used

  • KuaiRand-Pure Dataset (Zenodo Record 10439422): Micro-video sequential log containing 1.14M training interactions (04/08–04/21), 125K validation impressions (04/22–04/28), and 171K sealed test impressions (04/29–05/08) with 12 auxiliary user/video features.
  • KuaiRand Starter Kit: Official benchmark evaluation scripts for GAUC and nDCG@5.
  • RecSys Domain Knowledge Base (agents/knowledge.py): Curated repository of 12 citation-backed recommendation system methods (DCN-v2, PLE, BST, NeuralNDCG, BPR, ListMLE, MMoE, DIN, etc.).

Built With

Share this project:

Updates