Revisiting an old research project with Astra

Background

In 2025, I coauthored SelfDB, a self-supervised diffusion bridge for MRI reconstruction. MRI scanners collect frequency-space measurements, called k-space. Requiring fewer measurements can shorten how long an MRI scan takes. But the reconstruction model must then recover an image from less information.

Diffusion models learn to undo corruption over a sequence of steps. SelfDB starts from a degraded image built from incomplete MRI measurements. A neural network repeatedly predicts a reconstruction, using each prediction to construct the next input. During training, the diffusion bridge connects more- and less-degraded versions of a scan, with added noise, so the network learns to recover image structure.

The model learned this without clean images as training targets. We hid additional measurements from an already incomplete scan, then compared the predicted image’s simulated scanner readings with the measurements we had acquired. The withheld readings provided supervision even though a complete image was unavailable.

OpenAI recently released GPT-6 Astra, which scored 64.6% on Terminal-Bench Science 0.1, about 23% higher than the next-best model. I revisited this project to see whether that jump could translate into a substantial advance beyond the original paper on a $200 compute budget. Across a 21-hour research session, Astra used approximately 28 GPU-hours and $117.04 of that budget on Modal, reducing reconstruction error by 19.2%.

Method

The original SelfDB objective trains one prediction at a time, using a synthetic intermediate state at a randomly chosen time. At inference, however, each prediction helps construct the next input. Later steps must therefore handle errors introduced by the model itself—inputs that can differ from its training examples.

Astra adapted rollout training to SelfDB: run the reconstruction sequence during training and backpropagate through all its steps. An error in the final image can then change what the network does at the beginning. Every step uses the same U-Net weights.

Astra also incorporated measurement consistency correction into the training sequence. After each prediction, a numerical solver adjusts the image so its simulated scanner readings better match the actual input measurements, while staying close to the network’s prediction. Training differentiates through this correction too. Measurements hidden for the training loss remain hidden from the solver.

For a fair comparison, we also give the SelfDB baseline the same measurement correction at evaluation. The measured improvement therefore comes from training through the corrected sequence, with the reconstruction procedure held fixed.

Experiments

Astra and I started with a small model to test the approach, then used the original repository’s 149.6M-parameter U-Net. We used Calgary–Campinas, a public dataset of multi-coil brain MRI measurements for the same reconstruction task. The study used 5,032 training slices from 37 acquisitions, 8 validation scans, and 16 test scans containing 192 slices.

For each of three random seeds, we trained an ordinary baseline and separate two-, four-, and eight-step rollout models from the same pretrained checkpoint: 12 models total. Each branch received 24,000 additional network calls. This matches logical calls, but not FLOPs, optimizer updates, or elapsed training time.

Results

The graphs compare rollout training (new) with a SelfDB baseline trained for this Calgary study. Rollout training reduced image error by 19.2% and improved the other two main metrics, using the same network and four inference steps. Each method is evaluated at its best validation checkpoint.

NRMSE measures pixel error, LPIPS measures perceptual distance, and SSIM measures local structural similarity. Foreground scoring excludes dark background from NRMSE and SSIM; LPIPS uses the full image. Paired 95% bootstrap intervals supported improvement on all three metrics.

Reconstruction examples

The same test image reconstructed by the SelfDB baseline and rollout training (new), alongside the reference. Both use four inference steps and measurement correction. These examples use the final training checkpoints; the 19.2% headline uses the best validation checkpoints.

Example 1Example 2Example 3

SelfDB baseline

Rollout training (new)

Reference

Error maps

SelfDB baseline: absolute error

Rollout training (new): absolute error

00.15

Absolute error divided by the reference display maximum. The same scale is used for both methods and all three examples; darker means less error. Background is included.

Data: Calgary–Campinas (CC BY-ND 4.0). Model reconstructions and error maps from this study.

The fixed-budget endpoint showed a larger 29.2% error reduction, but ordinary training deteriorated as training continued. I lead with the smaller best-validation comparison because it gives the baseline its strongest validation-selected checkpoint.

Two-step-specific training helped at two-step inference. Eight-step-specific training did worse than running the four-step model for eight steps. An earlier external test on low-field MRI improved error and perceptual distance but did not establish an SSIM improvement. The final Calgary test set was held out from training, but had already been inspected during development.

Working with Astra

Astra handled research, ideation, implementation, training, debugging, and analysis autonomously, with periodic steering from me on the research direction and evaluation. It prepared data, launched experiments, recovered interrupted jobs, and compared the results.

Reflections