Dark Matter by Feedback: Uncertainty-Driven Simulation for Strong-Lensing Detectors

DeepLenseLoop is an agent that plans and runs strong gravitational lensing simulation batches, trains classifiers and anomaly detectors on the results, and uses what those models are uncertain about to decide what to simulate next. A human approves every batch before it runs.

It was built in six phases over a week on Pydantic AI, Lenstronomy and Modal, with the machine-learning architectures ported from two papers by the same research group.

A massive galaxy bends the light of one behind it into arcs. Small clumps of dark matter in the lens perturb those arcs by about one percent, and what the clumps look like depends on what dark matter is. Cold dark matter predicts point-like subhalos; ultralight axion dark matter predicts vortex lines.

Both signatures are subtle and real strong lenses are rare, so the only way to train a detector is to simulate: thousands of images per hypothesis, across every plausible halo mass, redshift, mass fraction and axion mass. Alexander et al. (2019) classified such simulations with ResNet-18 and AlexNet; Alexander et al. (2021) did it without labels using autoencoders. Both spent most of their effort generating and curating simulations by hand.

HEPTAPOD (Menzo et al., 2026) showed the pattern for high-energy physics: wrap the simulation engines in schema-validated tools, coordinate them with run cards a human signs off on, and let an agent orchestrate. This project applies that pattern to lensing and adds the piece it was missing, a feedback signal.

The claim the project tests is simple: a batch of simulations chosen by where the current model is uncertain should help more than a batch of the same size chosen uniformly. If it does, the agent is planning under a real feedback signal rather than running a script. Everything else exists to make that comparison cheap enough to run repeatedly.

The same engine the agent drives. Switch to the vortex class and slide the axion mass down: the vortex line stretches across the image and its perturbation of the arcs fades. That region is where the classifier turned out to be weakest.

The design went through three versions before any code existed. The first built everything by hand, including the agent loop. The second sat on Orchestral AI, the framework HEPTAPOD itself uses. The third, and the one that shipped, uses Pydantic AI, mostly because the run cards and dataset records were already Pydantic models, so tool arguments and structured outputs came for free. It also turned out to support approval-gated tools natively, which collapsed a planned two-agent propose-then-execute design into a single agent with one gated tool and an approval callback in our own code. That callback is a plain function returning approve, approve with edits, or deny, so it is unit-tested without an LLM, and the same function drives the terminal prompt, the benchmark's auto-approver, and the loop driver.

Compute and data went to Modal rather than a rented VM or the laptop. The pipeline is function-shaped, so each tool became a remote function behind the same submit-and-status interface as the local backend; batches shard across containers and training gets an L4 on demand. The images live on a Modal volume that mounts straight into the GPU containers, after a brief detour considering a vector database. The agent runs on Kimi K3 through Fireworks, though any Pydantic AI provider is a one-line change, and the ML models were ported layer for layer from the group's 2019 and 2021 papers rather than redesigned, since the contribution here is the loop around them, not the networks.

Tools never talk to the simulator or to PyTorch directly. They talk to a compute backend, so a command-line flag moves the whole pipeline from the laptop to Modal.

The agent never sees pixels. It reasons over one JSON record per image with cheap statistics computed at generation time, and over per-image model scores produced at evaluation time. That is what makes the uncertainty step a small numpy function rather than another model.

Simulation batches and training jobs cost money, so those tools require approval. Pydantic AI pauses the run and returns the pending calls; a callback in our own code decides; the run resumes. The same callback drives the terminal prompt, the auto-approver used by benchmarks, and the loop driver.

A frozen Pydantic model with physics validation: the source must sit behind the lens, the mass fraction must be positive for substructure classes, and a card whose field of view cannot contain the Einstein radius is rejected before any pixel is rendered.

The models were trained on 15,000 simulated images. Every number below was measured on a separate set of 3,000 images, 1,000 per class, generated with different seeds and kept out of training, so the scores reflect what the models learned rather than what they memorised.

5,000 training and 1,000 test images per class at 64 pixels in a Euclid-like setup, plus the batches added by the loop. 18,000 images rendered in about a minute of wall time, sharded across Modal containers.

ResNet-18 with a single-channel stem, trained from scratch for ten epochs on the 15,000 training images. It is essentially perfect on smooth lenses and misses about one vortex in five, and the uncertainty analysis says exactly which ones.

The same analysis, live on the 3,000 held-out scores. Pick a class and an axis, change the binning, and click a bar to see the most uncertain images in that cell.

The agent simulated 2,000 more vortex images at the axion mass of the weakest cell and retrained under an identical recipe. As a control, 2,000 images drawn uniformly across the three classes were added instead.

The uniform batch is indistinguishable from adding nothing. The targeted batch doubles accuracy in the weak region and lifts vortex accuracy from 0.82 to 0.91, at the cost of subhalos, which the added light-axion vortices most resemble and which are now outnumbered. A second round would flag subhalos as weakest and propose the balancing batch. That is the loop working as intended: the uncertainty signal finds a region uniform sampling cannot.

Same model, same prompt, same sandbox. One arm has the domain tools; the other has only a shell and file tools, Python with lenstronomy and torch on the path, and the required output layout spelled out. Graded by a cumulative rubric read off the sandbox: simulate, label, train on training data only, report the held-out AUC correctly.

A strong model builds the whole pipeline from scratch and even reaches a higher AUC by engineering its own preprocessing: an asinh stretch, dihedral augmentation and a small-image ResNet stem. The domain tools trade that ceiling for 24 times fewer tokens, nine times less wall time, and an identical result on every trial.

The same lens systems rendered through PyAutoLens agree with the lenstronomy port to about one percent of peak brightness. The residual is the packages' slightly different elliptical-radius conventions, and the point-mass substructure agrees as well as the smooth lens does.

• Three trials per arm. The pass@k estimators are exact for the sample, but the sample is small. Ten or more trials and a second model are the obvious next benchmark.
• The anomaly detector is at chance. The AAE reconstructs well but separates poorly, AUC 0.54, far from the 2021 paper's 0.93. At this resolution and mass fraction the substructure signal is one to three percent of peak, and ten epochs is a fraction of what the paper trained for.
• Targeting has a cost. The loop's gain in the weak region came with a subhalo regression from class imbalance. Class-balanced batches or re-weighting is the natural fix and the natural second round.
• The core arm out-engineered the fixed recipe. Its asinh stretch and augmentation are data-pipeline choices the ML tier should adopt. Its architecture change stays out, because the port is meant to match the papers.

Everything on this page came from the repository.