Test-Time Accuracy-Cost Control in Neural Simulators via Recurrent-Depth

Test-Time Accuracy-Cost Control for Neural Simulators

Scientific computing often requires explicit control of the accuracy-cost trade-off: fast, approximate predictions for exploration or real-time control, and accurate (but slow) predictions for final analysis or critical decisions. Numerical methods naturally provide this flexibility through discretization, order, and tolerances. Neural simulators do not: once trained, their expected accuracy and inference cost are fixed.

We introduce RecurrSim, an architecture-agnostic framework that enables test-time accuracy-cost control in neural simulators. RecurrSim models i) serve multiple computational budgets, ii) support anytime prediction, and iii) are highly parameter-efficient.

Why Test-Time Accuracy-Cost Control for Neural Simulators Matters

Simulations Are Fundamental to Science and Engineering

Simulations allow scientists to study complex systems and engineers to optimize designs without expensive or impractical experiments.

Numerical Methods Have Explicit Control of the Accuracy-Cost Trade-Off

Finer discretizations, higher-order methods, and stricter tolerances usually yield more accurate predictions, but require more compute.

Many Choices Affect the Accuracy-Cost Trade-Off Only During Training, Not Inference

During training, allocating more computational resources generally leads to more accurate predictions; for example, through more data, larger models, or longer training. But after training, typically models have a fixed expected accuracy and inference cost.

Existing Adaptive Inference Methods Have Limitations

Deep Equilibrium models can, in principle, adapt compute by iterating toward a fixed point, but in practice they often fail to converge; therefore, DEQs with additional iterations do not reliably yield more accurate predictions.

Diffusion-based methods offer a user-controlled knob, but their intermediate states are generally not directly usable as predictions; therefore, diffusion models do not provide anytime prediction.

What RecurrSim Enables

Recurrent-Depth Simulator

RecurrSim consists of an encoder that maps the input state x to a conditioning vector c; a recurrent-depth block that iteratively updates a latent state z_k conditioned on conditioning vector c; and a decoder that maps the final latent state to the output state.

RecurrSim is trained end-to-end. For each sample, the number of recurrent iterations K is drawn from a distribution, the recurrent-depth block is applied K times, and gradients are back-propagated. To bound memory, we use truncated backpropagation-through-depth.

At test-time, the user chooses the number of recurrent iterations K according to their desired accuracy and computational resources. The first few iterations make the largest adjustments, and subsequent iterations contribute progressively smaller adjustments.

During training, the number of recurrent iterations K is drawn from a Poisson log-normal distribution, exposing the model to a broad spectrum of compute budgets which encourages stability across both shallow and deep rollouts.

def forward(self, x):
  z = self.encoder(x)

  z = self.processor(z)












  x = self.decoder(z)
  return x

def forward(self, x, K=None):
  c = self.encoder(x)

  z = sample_noise()
  if self.training:
    with no_grad():
      K = sample_K()
      for _ in range(K - B):
        z = self.processor(cat([c, z], dim=1))
    for _ in range(B):
      z = self.processor(cat([c, z], dim=1))

  if not self.training:
    for _ in range(K):
      z = self.processor(cat([c, z], dim=1))

  x = self.decoder(z)
  return x

RecurrSim Enables Explicit Test-Time Accuracy-Cost Control

As the number of recurrent iterations K increases, trajectory error consistently converges—even beyond the maximum number of iterations seen during training (adaptive depth)—while predictions remain physically meaningful at every iteration (anytime prediction).

RecurrSim Outperforms Existing Adaptive Inference Methods

RecurrSim is the only method providing both adaptive depth and anytime prediction, while also achieving the best overall performance.

RecurrSim Transfers Across Architectures