Neural volatility surrogates in Rust

Neural-network volatility surrogates — Heston, one-factor Bergomi, rough Bergomi. Trained offline, inference at sub-millisecond speeds via candle.

The stochastic-rs-ai crate ships neural-network volatility surrogates for fast pricing / calibration. The crate is feature-gated behind ai; it pulls in candle-core for the inference runtime.

Status: experimental. The surrogates are useful for warm-starting a calibrator or for sub-millisecond price approximations in a feedback loop, not as the primary pricing surface for production trading.

Models shipped

Model	Spec module	Training-set archive
Heston	`volatility::heston::HestonNn`	`HestonTrainSet.txt.gz`
One-factor Bergomi	`volatility::one_factor::OneFactorNn`	`Bergomi1FactorTrainSet.txt.gz`
Rough Bergomi	`volatility::rbergomi::RBergomiNn`	`rBergomiTrainSet.txt.gz`

Each surrogate maps (model parameters, strike grid, expiry grid) → implied-vol surface. The surrogate is a small MLP (3–5 hidden layers) trained offline on a Fourier-pricer-generated dataset. At inference time the pipeline is:

Normalise the inputs via BoundedScaler / StandardScaler.
Forward-pass through the MLP.
Denormalise the output to a flat IV grid.
Hand the grid to ImpliedVolSurface::from_flat_iv_grid for downstream pricing or smile / skew analytics.

Example — Heston surrogate inference

HestonNn is not generic and is not constructed from a bare path string — load takes a model directory plus a candle_core::Device, and predict_surface takes a fixed-size [f32; INPUT_DIM] parameter array (INPUT_DIM = 5) rather than separate strike/expiry vectors (the strike and maturity grid is baked into the training set; use predict_implied_vol_surface, gated behind this crate's own quant feature, to get a full ImpliedVolSurface back). That grid is not arbitrary: strikes.len() * maturities.len() must equal the model's OUTPUT_DIM, which for HestonNn is fixed at 88 — any other shape returns Err("model output_dim 88 does not match strikes×maturities …"). No trained .safetensors model ships in this repository — only the gzip-npy training sets the round-trip test consumes — so this block is illustrative rather than compiled:

use candle_core::Device;
use stochastic_rs::ai::volatility::heston;
use stochastic_rs::ai::volatility::heston::HestonNn;

let device = Device::Cpu;
let nn = HestonNn::load("path/to/trained/heston", &device)?;

// Five parameters, each within `heston::PARAM_LB` / `heston::PARAM_UB`
// (module constants, not associated ones):
//   LB = [0.0001, -0.95, 0.01, 0.01, 1.0]
//   UB = [0.04,   -0.1,  1.0,  0.2,  10.0]
// `BoundedScaler` does not clamp, so out-of-range values are silently
// mapped outside the training domain.
let params: [f32; heston::INPUT_DIM] = [0.03, -0.5, 0.3, 0.04, 2.0];

// 11 strikes x 8 maturities = 88 = heston::OUTPUT_DIM.
let strikes: Vec<f64> = (0..11).map(|i| 70.0 + 6.0 * i as f64).collect();
let maturities = vec![0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0];
let forwards = vec![100.0; maturities.len()];

// Inference returns a full ImpliedVolSurface in <1ms
let surface = nn.predict_implied_vol_surface(&params, strikes, maturities, forwards)?;
// `ivs` is the raw (maturities x strikes) grid; there is no `iv_at(k, t)`
// point-query helper today, so read it directly or use `smile_slice`.
println!("1y smile: {:?}", surface.smile_slice(/* maturity_idx */ 4));

Training round-trip

Every surrogate ships a train_save_load_<model> integration test that:

Loads a known training-set archive
Trains the MLP for a few epochs
Saves to safetensors
Reloads and asserts inference parity

This guards the on-disk format from silent regressions.

Calibrating with a surrogate

The second step of Horvath, Muguruza & Tomas (2021): once a network $\tilde F$ reproduces the implied-volatility surface on its fixed strike–maturity grid, a calibration is the deterministic least-squares problem

\hat\theta = \arg\min_{\theta}\ \sum_{ij} w_{ij}^2\bigl(\tilde F(\theta)_{ij} - \sigma^{\text{mkt}}_{ij}\bigr)^2 ,

solved with Levenberg–Marquardt on the surrogate's exact Jacobian (reverse-mode differentiation through the network, one backward pass per grid point) — the pricer never runs during the calibration. SurrogateCalibrator works on any StochVolNn; HestonSurrogateCalibrator and RBergomiSurrogateCalibrator add the parameter mapping to quant's HestonParams / RBergomiParams and implement Calibrator + ToModel, so the fitted surrogate drops straight into the vol-surface pipeline. The result reports whether $\hat\theta$ stayed inside the training box: the network is untrained outside it, and a solution on the boundary is a warning, not an answer.

tests/doctest_ai_surrogate_calibration.rs

// docs: ai#calibrating-with-a-surrogate
//! Backs the surrogate-calibration example on the AI page.
#![cfg(feature = "ai")]

use ndarray::Array2;
use stochastic_rs::ai::Device;
use stochastic_rs::ai::calibration::HestonSurrogateCalibrator;
use stochastic_rs::ai::volatility::common::TrainConfig;
use stochastic_rs::ai::volatility::heston::HestonNn;
use stochastic_rs::ai::volatility::heston::INPUT_DIM;
use stochastic_rs::ai::volatility::heston::OUTPUT_DIM;
use stochastic_rs::ai::volatility::heston::PARAM_LB;
use stochastic_rs::ai::volatility::heston::PARAM_UB;
use stochastic_rs::quant::calibration::heston::HestonParams;
use stochastic_rs::traits::CalibrationResult;
use stochastic_rs::traits::Calibrator;
use stochastic_rs::traits::ModelPricer;
use stochastic_rs::traits::ToModel;

/// A smooth stand-in for a Heston surface generator, so the example trains
/// in seconds; a production surrogate is trained on pricer output instead.
fn synthetic_surface(params: &[f32]) -> Vec<f32> {
  (0..OUTPUT_DIM)
    .map(|k| {
      let mut v = 0.2 + 0.03 * k as f32 / OUTPUT_DIM as f32;
      for (j, &p) in params.iter().enumerate() {
        let x = (p - 0.5 * (PARAM_LB[j] + PARAM_UB[j])) / (0.5 * (PARAM_UB[j] - PARAM_LB[j]));
        v +=
          (0.08 + 0.02 * (j + 1) as f32) * x * ((k as f32 + 1.0) * (j as f32 + 1.0) * 0.11).sin();
      }
      v
    })
    .collect()
}

#[test]
fn calibrate_heston_on_a_surrogate() {
  // 1. Train a small surrogate on (parameters → surface) pairs.
  let rows = 256;
  let mut params = Array2::<f32>::zeros((rows, INPUT_DIM));
  let mut surfaces = Array2::<f32>::zeros((rows, OUTPUT_DIM));
  for i in 0..rows {
    let theta: Vec<f32> = (0..INPUT_DIM)
      .map(|j| {
        let u = ((i * 7 + j * 13) % 97) as f32 / 96.0;
        PARAM_LB[j] + u * (PARAM_UB[j] - PARAM_LB[j])
      })
      .collect();
    params
      .row_mut(i)
      .assign(&ndarray::Array1::from_vec(theta.clone()));
    surfaces
      .row_mut(i)
      .assign(&ndarray::Array1::from_vec(synthetic_surface(&theta)));
  }
  let mut model = HestonNn::new(&Device::Cpu).unwrap();
  let cfg = TrainConfig {
    epochs: 30,
    ..TrainConfig::default()
  };
  model.train(&params, &surfaces, &cfg).unwrap();

  // 2. Calibrate to a market surface: LM on the network's exact Jacobian.
  let market: Vec<f64> = synthetic_surface(&[0.02, -0.6, 0.4, 0.05, 3.0])
    .iter()
    .map(|&v| v as f64)
    .collect();
  let calibrator = HestonSurrogateCalibrator::new(&model, market).unwrap();
  let result = calibrator.calibrate(None).unwrap();
  assert!(result.converged(), "{}", result.message().unwrap_or(""));
  assert!(result.fit.in_bounds);
  assert!(result.rmse() < 0.05, "rmse {}", result.rmse());
  let fitted: HestonParams = result.params();
  assert!(fitted.kappa > 0.0 && fitted.rho < 0.0);

  // 3. The result is a pricer: price a vanilla off the calibrated parameters.
  let pricer = result.to_model(0.01, 0.0);
  let call = pricer.price_call(100.0, 100.0, 0.01, 0.0, 1.0);
  assert!(call > 0.0 && call < 100.0);
}

# maturin develop --features ai   (the wheels on PyPI leave candle out)
import numpy as np
import stochastic_rs as srs

model = srs.HestonNn()                       # best available device
report = model.train(params, surfaces, epochs=200)   # rows: [v0, rho, sigma, theta, kappa] → 88 IVs
fit = srs.calibrate_surrogate(model, market_surface)  # LM on the exact network Jacobian
print(fit["params"], fit["rmse"], fit["converged"], fit["in_bounds"])
surface, jacobian = model.predict_surface_with_jacobian(fit["params"])

Requires the ai feature. The example trains a tiny network on a synthetic surface so it runs without the bundled train sets; with a production surrogate, load it with HestonNn::load instead.

Training on the GPU

candle runs on the CPU by default. The umbrella's ai-metal feature (macOS) compiles candle's Metal back-end; for CUDA enable candle-core/cuda in your own manifest (candle's kernels need a CUDA toolkit to build, so the crate cannot carry that feature). ai::device::best_available() detects the fastest device at run time — pass it to any surrogate constructor instead of Device::Cpu. Weights are stored device-independently, so a model trained on a GPU loads on the CPU and vice versa.

use stochastic_rs::ai::device::best_available;
use stochastic_rs::ai::volatility::heston::HestonNn;

let device = best_available()?;               // Cuda / Metal / Cpu
let mut model = HestonNn::new(&device)?;      // trains and predicts on that device

Adding a surrogate

See the vol-surrogate-nn SKILL. It covers StochVolModelSpec, BoundedScaler / StandardScaler conventions, gzip-npy training-set loading, the round-trip test, and predict_surface integration with ImpliedVolSurface::from_flat_iv_grid.

AI surrogates