Neural Methods

This section covers the neural_solver Python package, the mesh-free, differentiable companion to the Rust solver crate. It implements neural methods for backward stochastic differential equations (BSDE) and Hamilton-Jacobi-Bellman (HJB) equations in JAX.

Neural methods are a separate package, not part of the Rust workspace. They are coupled to the Rust solver only through committed reference values: the Rust solver emits exact values, and the JAX tests assert agreement against them. There is no live call graph between the two codebases.

Theory and empirical results are separate

This folder keeps the mathematical formulation and the measured results in distinct files so the two are not conflated:

The Rust numerical methods (finite difference, least-squares Monte Carlo BSDE, analytical solutions) are documented separately in the sections above; they are the rigorous reference against which the neural methods are validated.

Network architectures

Three architectures are implemented. All are built on a single feedforward block, Mlp (Xavier init, selectable activation tanh (default), softplus, relu).

Deep BSDE

Forward-backward SDE: learns a scalar initial value Y0 plus an integrand network on \((t, x)\).

• Z net (Mlp) — the diffusion integrand.
• \Gamma net (Mlp) — the predictable jump integrand, added by the jump extension.

Applied to Black-Scholes and Merton (diffusion), and to jump-LQ and Merton-jump (jump).

• Random-jump-size sampling — the \Gamma path currently samples fixed-amplitude jumps; the log-normal MertonJumpLognormal is reference-only with no bsde_problem. (planned)

Deep HJB (DGM)

Residual minimization over the value function, parameterized directly.

• value net (Mlp) on \((t, x)\).

Applied to Black-Scholes and jump-LQ.

• Full DGM residual for arbitrary problems — only the Black-Scholes and jump-LQ residuals are implemented; the general HJB residual is not. (planned)
• Adaptive activation — a tanh with a learnable per-unit scale and shift, to sharpen near kinks while staying differentiable. (planned)
• Partition-of-unity / local-basis networks — piecewise-smooth global functions for kink capture, discussed on the jump page. (planned)

Neural operator

Learns the solution operator over a family of instances (amortized), rather than one instance at a time.

• FNO — FourierOperator, a stack of spectral_layer Fourier-multiplier layers (a learned symbol).
• DeepONet — a branch Mlp + trunk Mlp factorization.

Applied to the jump-LQ family amortized over the jump amplitude.

• Exact-reduction match — the FNO/DeepONet primitives and the amortized training loop exist, but the stage validating the learned symbol against the exact jump-generator symbol is still in progress (07-neural-operators). (planned)

The architecture and activation details are on the BSDE/DGM and jump theory pages, and the operator architectures on the operators page.

Other planned architectures

Worked out in the neural-architectures research problem:

• Hamiltonian saddle network — value/control as a min-max of a learned Lagrangian, so the control supremum is exact by construction. (planned)
• Monotone neural operator — input-convex-plus-affine layers carrying the viscosity comparison principle. (planned)
• Measure transformer / Wasserstein attention — permutation-invariant attention over the BSDE particle cloud (dimension-free). (planned)
• Backward-causal operator — time-masked attention over reversed-time slices. (planned)
• Characteristic / Hamiltonian-flow transformer — parameterizes along the Hamiltonian characteristics. (planned)
• Actor-critic / learned policy — a learned control network replacing the analytic controls, for problems with no closed-form policy. (planned)
• Graph / equivariant network — a permutation-equivariant head for multi-asset inventory. (planned)
• Physics-informed neural network (PINN) — strong-form residual with explicit boundary/terminal conditioning. (planned)
• Score-based / diffusion generative model — for sampling the forward measure. (planned)
• Kolmogorov-Arnold network (KAN) — an MLP alternative. (planned)
• RNN / GRU sequence network — the cheap alternative to masked attention for the backward recursion. (planned)

Status

Implementation state is tracked in the neural_solver workstream. In short: the package, forward SDE simulation, closed-form reference models, the DGM diffusion-term helpers, the full deep BSDE (Z process) method (validated against Black-Scholes and Merton), and the jump-aware deep BSDE ($\Gamma$ process, validated against the jump-LQ and Merton jump problems) exist. The remaining items are the "(planned)" entries under each architecture in the Network architectures section above.