Reflectron

An ML research agent that runs the whole loop on its own; and never hides a bad result.

Comment

Project Story

Inspiration

We went into TikTok TechJam wanting two things at once: a real shot at a hard, open-ended challenge, and a genuine excuse to get our hands dirty with machine learning itself. Not just calling an API, but actually understanding how the frameworks differ, how a model's parameters move during training, and why an ML engineer's day looks the way it does.

Challenge 2 turned out to be exactly that excuse. The premise is almost recursive: build an agent that does what an ML research engineer does (read the problem, inspect the data, engineer features, train and tune a model, evaluate it, then reflect and revise) and have it run that loop on itself, autonomously, against a real recommendation benchmark: KuaiRand Pure, 1.4 million interactions from Kuaishou's short video feed. To build that agent well, we had to actually understand every stage of the loop we were automating. You can't design a diagnosis system that recognizes why a model regressed without understanding backpropagation, overfitting, and what a validation curve is supposed to look like when it's healthy.

We also drew directly on three papers the problem statement pointed us toward: AIDE, which showed that tree-structured search over ML solutions beats linear iteration by roughly four times on MLE Bench; AI Scientist v2, which separates the decision of what to try next from the LLM call that writes the code; and MLE Bench itself, the evaluation paradigm this whole challenge is modeled on. Reading those before writing a line of code changed the shape of what we built. We didn't want a chatbot that edits a script in a loop; we wanted a system with real memory of its own experiments.

What it does

Given the KuaiRand Pure benchmark and a fixed baseline, the agent autonomously reproduces that baseline, then iterates: proposing, training, evaluating, and diagnosing new candidates until the validation score converges by the organizer's own rule.

Formally, a run is converged once

$$ \text{best_valid_primary}i - \text{best_valid_primary}{i-N} \leq \epsilon, $$

where

$$ \epsilon = 0.002,\qquad N = 3. $$

In other words, the best score seen so far has not improved by more than ( \epsilon ) over the last ( N ) consecutive iterations.

Every iteration's hypothesis, code diff, metrics, and any error or recovery event gets logged automatically and rendered into a clickable Research Map.

The scored metric itself is a simple average of two ranking metrics:

$$ \text{primary} = \frac{1}{2}\left(\text{GAUC} + \text{nDCG@5}\right). $$

Every result is judged against the baseline as

$$ \Delta = \text{primary}{\text{agent}} - \text{primary}{\text{baseline}}. $$

The result that shipped, a multi-task DeepFM model we call deepfm_mtl_v1, achieves

$$ \Delta = +0.0028 $$

on the hidden test set, verified across 3 random seeds. That edge holds, and even grows, on TikTok's own unbiased randomized exposure log: real evidence it isn't an artifact of the platform's normal serving bias.

How we built it

We built this in layers, each one earning the right to add the next.

Foundation first

Before any modeling, we built the boring but critical plumbing: an evaluator wrapper that calls the organizer's own scoring code directly instead of reimplementing it (the single easiest way to silently score yourself wrong), and a convergence detector that reads the required threshold and iteration count live from the organizer's own published file rather than hardcoding it. Only once that harness self-checked correctly against a known reference score did we trust anything it reported afterward.

The models, hand-derived first

Our first real model was a Factorization Machine, ported faithfully from the official baseline. In plain terms, its prediction is three things added together: a starting bias, a simple weighted sum of the input features, and a learned interaction term that lets sparse categorical features like user ID, video ID, and author ID cross with each other, without an explosion of parameters, by representing each one as a small learned vector and taking the dot product of every pair:

$$

\hat{y}(x)

b + \sum_{i=1}^{n} w_i x_i + \sum_{i=1}^{n} \sum_{j=i+1}^{n} \langle v_i, v_j \rangle x_i x_j. $$

We hand-derived and hand-coded its backward pass,

$$ \frac{\partial \hat{y}}{\partial v_i} \qquad\text{and}\qquad \frac{\partial \hat{y}}{\partial w_i}, $$

in plain NumPy rather than reaching for a framework immediately. This turned out to be genuinely useful for understanding what autograd is actually doing under the hood once we did start using it.

DeepFM extended this with a small neural network component reading the same embeddings, still hand-rolled.

PyTorch, added deliberately, not by default

We kept the NumPy-only philosophy for as long as it made sense, and only reached for PyTorch when automatic differentiation was a genuine win rather than a shortcut.

Specifically, for a five-headed multi-task network trained on a combined objective,

$$

L

L_{\text{main}} + \lambda \sum_{k=1}^{4} L_{\text{aux}}^{(k)}, $$

we had one main loss plus four auxiliary engagement signals, all sharing one embedding table, whose backward pass has five different gradients merging back into a shared trunk. Hand-deriving that would have been real effort spent on the wrong problem.

That one model, deepfm_mtl_v1, turned out to be the actual project best, and it landed there on its very first attempt.

A memory, not a log

Rather than a flat list of what we tried and what happened, we built a persistent, tree-structured Research Map, directly inspired by AIDE, where every experiment is a node with a parent, a type of edge (a fresh idea, an improvement, or a bug fix), and a diagnosis.

A metric-aware diagnosis engine reads the pattern across both scored metrics, not just whether the headline number went up. So a result where broad ranking quality improved but top-of-list precision got worse is correctly tagged as a genuine trade-off, not lumped in with a clean win or a clean loss. That distinction mattered more than once.

Then we let an LLM drive

With that memory and diagnosis layer in place, we wired in Google Gemini as a genuine research strategist: it reads the live Research Map, every prior hypothesis and result, and proposes the next experiment on its own, validated against our model registry before anything runs.

On its first real round, it independently noticed an overfitting pattern our own hand-authored logic had missed, and proposed exactly the right fix, which became a real, verified improvement.

Then we searched, honestly, for more

Once we had one real win, we spent the rest of the project trying to beat it, mostly failing on purpose while testing everything we could:

  • Five variants of pairwise ranking loss
  • A technique that weights each ranking mistake by exactly how much fixing it would move our target metric
  • Focal loss, which down-weights training examples the model already gets confidently right
  • Listwise softmax
  • A new architecture called DCNv2
  • Four different ensembling methods
  • Checkpoint averaging
  • Our one genuinely non-standard idea: initializing embeddings by propagating signal across the user-video interaction graph instead of starting from random noise

Almost all of it came back negative or a tie. We reported every single one of those results exactly as measured.

Challenges we ran into

The task definition itself was ambiguous, at first

The problem statement's own prose said one thing, and the Starter Kit's actual pinned code said another. We had to decide, and defend, which one to trust (the code, since it's what actually gets run) before writing a single model.

The organizers later updated the problem statement to match the code directly, which was a genuinely good feeling to see confirmed.

Windows fought us more than the ML did

Python's multiprocessing needs a different startup mode on Windows than on Mac or Linux, which meant real subprocess isolation code for our failure recovery system, not a simple error catch.

The organizer's own submission validator crashes on a default Windows console because of a single special character in its success message. It was a genuinely confusing "is my file broken?" moment the first time we hit it, resolved by realizing the crash happened only after every real check had already passed.

LightGBM's native library outright refused to load under a Windows security policy on our primary machine; a teammate on macOS didn't have that problem at all.

Silent bugs are the dangerous kind

Twice we found bugs that had been silently wrong for a while without ever throwing an error.

The first was a report generator whose import ordering meant two of its own sections had never rendered correctly.

The second, and more serious, was a "current best" headline that had been silently stuck reporting an old result for days because nothing cross-checked it against our own persistent memory.

Neither crashed anything, which is exactly what made them dangerous. We only found them by actually reading the generated output critically, instead of trusting that no crash meant correct.

Knowing when to stop looking for a bug and accept a real result

After roughly twenty further attempts beyond our first win all failed to beat it, the hard part wasn't building the twenty-first attempt. It was recognizing that the pattern itself was the finding.

This benchmark's learnable signal, at this data volume, looks substantially extracted by a multi-task training objective, and no amount of architecture or loss-function tinkering was going to change that.

Reporting a wall of negative results honestly, instead of quietly going back for one more attempt and hoping for a better story, was its own kind of discipline.

Accomplishments that we're proud of

A real, verified improvement, not a lucky seed

A gain of 0.0028 on the primary metric sounds small until you look at the ceiling.

The metrics used here don't span a full 0 to 1 range; a perfect ranking only reaches about 0.86, since over a quarter of users have no positive label at all, which forces their score to zero no matter how good the model is.

Define the attainable headroom as the gap between a perfect ranking and pure random scoring:

$$

\text{headroom}

\text{primary}_{\text{oracle}}

\text{primary}_{\text{random}}

0.8645 - 0.4753

0.3892. $$

The official baseline already sits at

$$ \frac{0.5946 - 0.4753}{0.3892} \approx 31\% $$

of the way up that headroom.

Our own gain adds another

$$ \frac{0.0028}{0.3892} \approx 0.7\% $$

of that same headroom on top: a real, meaningful bite, not a rounding error.

We didn't trust a single lucky run for it either. Every number we call "the result" is verified across 3 seeds, and we independently confirmed it holds, and even grows, on a genuinely different, unbiased data distribution than the one we trained on.

Genuine autonomy, not a rubber stamp

Watching our LLM research strategist read its own experiment history and independently flag an overfitting pattern our own hand-authored logic had missed, then propose exactly the right fix, which became a real, verified win, was the moment this stopped feeling like a script and started feeling like a colleague.

Catching real bugs before they mattered

We found and fixed two genuinely silent bugs: a report generator whose sections had never once rendered correctly, and a "current best" headline stuck on a stale result for days.

We caught both purely by reading our own output critically instead of trusting that no crash meant correct. Neither would have shown up in a casual demo.

Robustness we actually tested, not just claimed

Our failure recovery system survived a genuine out-of-memory crash, a real allocation request of nearly 300 terabytes, not a simulated one, and a real training timeout mid-run, recovering both times without ever taking down the whole pipeline.

That is exactly the kind of thing that is easy to assert and easy to never actually verify.

An honest, thorough search, even when it kept saying no

Thirty-three real, logged experiments, only three of them a clean win, all reported exactly as measured instead of quietly filed away.

We're proud of the discipline that took as much as we're proud of the one result that worked.

What we learned

Concretely, on the machine learning side: what a hand-derived backward pass actually looks like for a bilinear interaction term, when automatic differentiation genuinely earns its complexity cost versus when it's just convenience, and why pairwise ranking losses are theoretically appealing but empirically harder to train stably than plain pointwise classification on a dataset this size.

The biggest surprise was that on a well-specified benchmark like this one, the training objective, meaning what you are asking the model to optimize for, can matter far more than the model's capacity, meaning how big it is.

Every architecture change we tried landed close to flat. The one lever that worked was multi-task learning: a change to what the model was being asked to predict, not how expressive it was allowed to be.

More broadly, a genuinely autonomous research loop is mostly an exercise in restraint and record keeping, not cleverness. The value was in never letting a result go unrecorded, never trusting a number we hadn't independently verified against the organizer's own pinned scoring code, and being honest, every single time, about which parts of the process were the agent's own reasoning and which were ours.

What's next for our project

Turn the graph from a feature into an architecture

Our one genuinely non-standard experiment, initializing embeddings from a graph propagation over the user-video interaction graph, came back a clean null, but only as a one-time initialization.

We never tried making that propagation a live part of the forward pass, recomputed against the model's current, evolving embeddings every step instead of a frozen snapshot from before training started.

That is a real architecture change, not a data-prep step, and it's the most promising untested idea we have.

Let the LLM run longer, with a real budget to manage

Our research strategist has only ever run a handful of iterations at a time.

With budget-aware stopping already built, the natural next step is a long, unattended run, hours rather than minutes, to see whether it converges on something none of us thought to try by hand.

Generalize the research critic gate

Right now it catches one specific repeated pattern: a pure capacity hyperparameter change tried twice.

A version that recognizes a fourth pairwise loss variant as the same kind of dead end as a fourth capacity-only change would make the agent's own judgment sharper, not just its search wider.

Scale up, properly

Two larger versions of the same benchmark are sitting right there as bonus targets.

We deliberately didn't force a shortcut through the Starter Kit's hardcoded loader just to claim the bonus points. A real, config-driven, validated loader for the larger benchmarks is worth building properly next, not rushed.

Take the counterfactual angle further

KuaiRand's randomized exposure log already let us confirm our result generalizes to unbiased data.

The same log is a genuine foundation for real off-policy evaluation and debiased training: using the unbiased subset not just to check a model trained on biased data, but to train one that corrects for that bias directly.

Built With

Share this project:

Updates