Differentiating transport when photon paths are random
Changing a transport parameter can alter both a photon history and the probability of drawing it. Gradients of the expected detector signal must account for both effects.
Move a bone slightly across the beam. The primary image changes because rays cross different thicknesses of tissue. In the scattered image, moving the bone changes where photons can collide and which outgoing directions lead to the detector. A history that missed the detector before the move may now contribute. An absorbed photon might instead scatter. The derivative has to account for those changes in probability, although the simulation only records the history that actually happened.
A transport program can produce the right average image while its automatically differentiated version produces the wrong average gradient. Every multiplication may have been differentiated correctly. The missing contribution lives in the decisions between those multiplications: which interaction occurred, which boundary was crossed, which detector pixel received the photon.
We need a derivative of the expected measurement to decide how to change the object or the source. That requires choosing a representation of the expectation, deriving an estimator for its derivative, and arranging the GPU calculation around that estimator. The probability calculation comes first. Giving a photon history a tape does not make it confess to the paths it did not take.
Build the derivative by extending Chapter 9’s history walk. For material-density changes in fixed cells, accumulate the likelihood contribution along each history. For a source-amplitude change, differentiate the multiplier on its score. These estimates will feed a stochastic inverse solver whose signal and derivative calculations use independent histories.
10.1 Differentiate an expected measurement
Take one detector pixel and one scalar parameter . It might be a material coefficient, a translation in mm, an angle in radians or a source-spectrum parameter. Other parameters are held fixed during this partial derivative. Write for the physical history of one emitted photon, including its terminal event. Its normalised physical probability density is , and is its detector score. The integral over histories includes sums over discrete choices and history lengths as well as integrals over continuous states.
Let be the expected number of photons emitted in the exposure. For now, hold it fixed. The expected signal is
In equation (10.1), the detector score carries the measurement units. A score of one for an accepted arrival gives expected counts, while an energy-dependent detector response gives that response’s output units. The source and score conventions are those of Chapter 9. The simulated history count controls the precision of an estimate of this integral. It does not replace or change the exposure.
Both factors inside the integral can depend on the parameter. Differentiating a detector gain changes the score. Differentiating an interaction coefficient changes the history law. For a fixed history domain, suppose the integrand is differentiable in a neighbourhood of the parameter and its derivative has an integrable bound independent of the parameter in that neighbourhood. Differentiation can then pass under the integral. Where is positive, define its logarithmic derivative, or likelihood score, and apply the product rule:
The two terms in equation (10.2) change the value assigned to a fixed history and the probability of drawing it. Here holds the history coordinates fixed. The likelihood score has units reciprocal to those of . It can be negative even when every photon contribution is non-negative. It is distinct from the detector score .
Changing support or moving a jump in requires further treatment in §10.4. Those cases violate the simple smooth-integrand argument above. If the expected emitted population also varies, the product rule adds . Keeping that prefactor explicit prevents a source-intensity gradient from disappearing into a normalised probability density.
10.2 Pathwise and score-function estimators
A random number generator supplies variables with a parameter-independent distribution. For example, a uniform variate remains uniform when the material density changes, while the sampled flight distance changes through the transformation applied to that variate. Suppose a history can be generated as from base randomness with law . If the transformed score is differentiable and its difference quotients admit an integrable bound, the pathwise construction gives
In equation (10.3), the random numbers stay fixed while the sampled trajectory moves. The total derivative of includes the changing positions, energies and detector response along that trajectory. This is where automatic differentiation can carry a large amount of useful arithmetic.
Mohamed and colleagues organise gradient estimation around pathwise, likelihood-score and measure-valued constructions, with different regularity and support requirements. For our walker, a likelihood score accounts for the changing probability of collision events without differentiating a categorical branch. Smooth transformations can use pathwise derivatives, while moving discontinuities need the treatment in §10.4. The choice follows the dependence of the measurement on the parameter. [28]
The pathwise and likelihood-score constructions represent the same derivative under their respective conditions. Adding their complete answers would count it twice. They can, however, handle different choices within one estimator. Consider a discrete choice with probabilities , followed by a continuous sample generated from independent base randomness. Let be the resulting score. Differentiating the sum over and then its conditional integral gives
The first term in equation (10.4) differentiates the continuous calculation with fixed. The second accounts for the changing probability of choosing . A full history applies this reasoning conditionally at successive decisions. Each probability is conditioned on the preceding state, and its derivative includes the parameter dependence of that state under the chosen representation.
This division of labour is especially useful for transport. Flight lengths can sometimes be reparameterised smoothly. Absorption versus scattering is a categorical choice. A detector hit can introduce a jump even after every preceding sample was continuous. For each dependency, identify which estimator term covers it before building the backward calculation. The existence of a differentiable inverse sampler does not establish the smoothness of the final detector score.
Figure 10.1 traces these dependencies from the source sample to the detector score.
- Source sampleCP
Move a continuous sample or change the probability of a discrete source choice.
- Free flightC
In a homogeneous medium, holding U fixed gives a smooth inverse-sampled distance.
- Interaction choiceP
Absorption and scattering branches have parameter-dependent probabilities.
- Scattering sampleC
A smooth angular sampler changes the outgoing direction and energy.
- Detector scoreCB
Response changes smoothly, while a moving acceptance edge changes which histories count.
One possible division of derivative terms, with fixed source amplitude. C and P cover different choices within the same estimator. Complete pathwise and likelihood estimates must not be added together.
Estimator scope
The likelihood estimator in the accompanying source covers fixed-cell densities and source amplitude. Moving geometry and score boundaries require additional derivative terms.
10.3 Differentiate the sampling law
Start with a homogeneous medium whose total interaction coefficient is , in mm. A uniform variate produces an unbounded free-flight distance in mm:
The negative derivative in equation (10.5) has a physical meaning: increasing the collision rate shortens the flight generated by the same random number. Its units are mm, because we differentiate a length with respect to an inverse length. The limit is the separate vacuum branch.
Now put the end of the medium at a fixed distance . Record if a collision occurs before that end and if the photon escapes uncollided. The travelled distance is . The collision branch has density , while the escape branch has probability . Their likelihood scores fit into one expression:
Equation (10.6) includes the escape probability. Treating every recorded segment as though it ended in a collision would add a spurious on escaping histories. For a primary-count score , only escapes contribute, so the expected per-emitted-photon derivative is
This recovers Chapter 2’s homogeneous transmission derivative in equation (10.7) directly from a random-history estimator. At and , and . Differentiating the realised binary escape score through the comparison instead would return zero almost surely. The expectation changes because the threshold selects a different fraction of random numbers.
In a heterogeneous medium, an uncollided segment carries the survival factor for its whole traversed length. Hold the segment geometry fixed for a material derivative. Write for the coefficient at along that segment, and define
Equation (10.8) contributes on every segment, including the last segment to the detector or domain boundary. A collision adds the logarithmic derivative of its local rate. For event type , it is often convenient to factor this into a total collision rate and a conditional event probability:
Combining the two rate terms in equation (10.9) leaves , as it should. A scattering event also contributes the derivative of its normalised angular and energy distribution. The source distribution contributes its own term. Probabilities or rates equal to zero require an alternative representation or a suitable one-sided limit. Dividing by zero is not a derivative estimator.
The supported material parameter is the logarithm of a positive density multiplier. That multiplier scales both absorption and scattering coefficients for one selected material, leaving their ratio and the conditional scattering law unchanged. Under analogue sampling, a collision in that material therefore contributes one to the log-density likelihood score, while each traversed segment contributes minus its total optical thickness. The accumulator in Listing 9.2 includes the final escape segment as well as collision-ending segments. This is the material-wise version of equations (10.6) and (10.8), with the logarithmic chart’s chain factor already included.
The derivative kernel replays that walk and combines its density accumulator with the complete detector contribution. The same entry point also handles direct and logarithmic source-amplitude derivatives under a fixed source law.
@wp.kernel
def derivative_histories(
positions: wp.array(dtype=wp.vec3d),
directions: wp.array(dtype=wp.vec3d),
weights: wp.array(dtype=wp.float64),
density: wp.array(dtype=wp.float64),
material_ids: wp.array(dtype=wp.int32),
energies: wp.array(dtype=wp.float64),
absorption: wp.array(dtype=wp.float64),
scattering: wp.array(dtype=wp.float64),
parameters: Parameters,
seed: wp.uint64,
first_history: wp.uint64,
active_material: int,
source_amplitude: wp.float64,
out_pixel: wp.array(dtype=wp.int32),
out_derivative: wp.array(dtype=wp.float64),
out_status: wp.array(dtype=wp.int32),
status: wp.array(dtype=wp.int32),
):
index = wp.tid()
amplitude = source_amplitude
base_weight = weights[index]
if active_material >= 0:
# A density derivative is formed directly below. An unused primal
# overflow/underflow must not determine its representability.
base_weight = wp.float64(0.0)
if active_material == SOURCE_AMPLITUDE:
amplitude = wp.float64(1.0)
result = walk(
positions[index],
directions[index],
base_weight,
amplitude,
seed,
first_history + wp.uint64(index),
active_material,
material_ids,
energies,
absorption,
scattering,
density,
parameters,
)
out_pixel[index] = result.pixel
derivative = wp.float64(0.0)
if active_material >= 0 and result.pixel >= 0:
score_energy = wp.float64(1.0)
if parameters.energy_score != 0:
score_energy = result.energy
derivative = product4(weights[index], source_amplitude, score_energy, result.density_score)
if parameters.continuous_absorption != 0:
derivative = attenuated_product4(
result.absorption_depth,
weights[index],
source_amplitude,
score_energy,
result.density_score,
)
if active_material == SOURCE_AMPLITUDE:
# d(a * base_score)/da, including at a=0; never divide by amplitude.
derivative = result.score
if active_material == LOG_SOURCE_AMPLITUDE:
derivative = result.score
out_derivative[index] = derivative
out_status[index] = result.status
if result.status >= 2:
wp.atomic_or(status, 0, 1 << result.status)
if not wp.isfinite(derivative):
wp.atomic_or(status, 0, 16)
For a density derivative, base_weight is set to zero during the walk because the forward score is unused. The derivative is then formed directly from source weight, amplitude, optional energy and density_score, using mantissa/exponent factorisation. The analogue branch calls product4, while continuous absorption weighting calls attenuated_product4 with the accumulated absorption depth. An underflowed forward tally cannot erase a representable derivative, and an unused primal overflow cannot invalidate one. The probability calculation still uses the actual density. Zeroing that temporary score does not detach the physical law.
The direct amplitude branch instead evaluates the base score with amplitude one. This gives the derivative even at amplitude zero without dividing by zero. The log-amplitude branch requires positive amplitude and retains its chain factor in the score. These are derivatives of an intensity multiplier. Changing the source’s angular distribution, spectrum or sampling support would introduce different terms and is outside this interface.
Under continuous absorption weighting, the density derivative has two sources. The scattering-path likelihood contributes , where counts scattering events in the active material. Differentiating the explicit absorption weight contributes . Thus a complete detector contribution has derivative estimator . The familiar total optical depth reappears, but its absorption part now comes from differentiating a weight. Treating every subtraction as a likelihood term would obscure which probability law we are differentiating.
This argument holds for positive density multipliers, fixed coefficient zeros, geometry and source law, and a density-independent conditional scattering law. Interchanging the derivative and expectation still needs an integrable derivative. In the free-electron Compton model, density changes collision rates while the conditional energy-angle law stays fixed. Each segment uses the coefficients at its own energy. To check that calculation independently, differentiate the primary-plus-single-scatter integral and bound the derivative of the omitted scattering orders. A bound on their detector score alone is insufficient.
A sampling proposal can simplify these terms. At a reference parameter , choose a fixed history proposal that covers the physical integrand and its derivative throughout a neighbourhood. On a fixed domain with the same interchange conditions, importance sampling gives
In equation (10.10), is held fixed while differentiating. The physical density still changes inside the numerator. Detaching a proposal means holding that sampling law fixed in the derivative. Detaching the physical attenuation or scattering factor would delete the sensitivity we are trying to estimate.
For null-collision tracking, one option is to hold a valid majorant fixed locally. The candidate process is then unchanged, but acceptance and rejection probabilities still depend on the material. A likelihood treatment includes both outcomes, or a weighted formulation derives their replacement explicitly. The majorant must bound the physical rate over the parameter neighbourhood being used. Changing a majorant or a queue rule is an algorithmic change whose effects must cancel in the expected physical measurement.
10.4 Moving boundaries and discontinuities
A single detector edge exposes what ordinary trace differentiation misses. Let a crossing coordinate be uniform over a strip of width in mm. An aperture accepts , with and fixed . Its expected count is
For every sampled , the local derivative of the indicator in equation (10.11) is zero. Yet moving the edge by a small positive amount admits an interval of previously rejected crossings. Its width shrinks with the perturbation, while the change in each admitted score remains one. Dividing by the perturbation leaves a finite contribution. The derivative resides on the moving edge.
The same argument works when the two sides carry different smooth contributions. Let and be signal densities per mm to the left and right of a moving boundary in a fixed strip. Integrating each side and differentiating its endpoint gives
The final term in equation (10.12) is the jump in contribution multiplied by boundary motion. It vanishes when the contributions agree across the boundary. Otherwise, differentiating the smooth pieces leaves it out. In more dimensions, the corresponding term integrates over the boundary and uses its normal velocity.
Transport contains several such boundaries: an object silhouette, the edge of a detector pixel, a collimator opening and the interface at which one material becomes another. A geometric interface need not create a jump in every chosen integrand, so identify the actual contribution on both sides. Fixed random numbers do not remove a jump. Nor does assigning a zero derivative to a branch comparison account for it.
Chapter 5 encountered this problem in a sampled line integral through a sharp interface. In a transport history, the discontinuity can occur at any interaction or detector decision.
There are several ways forward. Integrate a simple acceptance probability analytically, as in the strip example. Sample the boundary contribution alongside the smooth contribution. Or transform the integration variables so that the problematic boundary becomes fixed and differentiate the transformed integrand and its Jacobian. Each method needs a derivation for the source, geometry and detector score being used. Replacing a sharp detector edge with a smooth filter defines a different measurement unless that filter belongs to the physical response model.
Li and colleagues derive an edge-sampling estimator for differentiable rendering that evaluates the jump across visibility boundaries in addition to the smooth contribution. Their construction makes the missing term explicit: ordinary area samples almost surely miss the lower-dimensional boundary integral. The same mathematical issue appears in photon transport, although the material and detector model determine the boundary contribution that must be sampled. [29]
We will keep material interfaces, detector faces and the source distribution fixed when differentiating density and amplitude. Adding moving geometry requires the corresponding interface contribution derived here; changing the source or conditional scattering law requires its probability derivative. These are additions to the estimator before they are additions to the parameter vector.
The swept interval in Figure 10.2 is the source of the boundary contribution.
Uniform crossings over 10 mm
The admitted fraction rises by 20 percentage points. Its derivative is .
10.5 Replay, random numbers and memory
A long history can visit many materials before contributing to a detector pixel. Reverse differentiation needs the states at which its local factors were evaluated. Storing every position, energy, probability and weight for every history makes the memory cost grow with total path length. Recomputing the first event, then the first two, then the first three soon repeats a great deal of work: regenerating every prefix takes quadratic time in history length. A replay scheme should specify which states it reconstructs and how often.
Path replay backpropagation, developed by Vicini and colleagues, reconstructs local quantities for reverse differentiation using invertible local scattering Jacobians. Their formulation takes work linear in path length and constant path-local memory. We also regenerate histories to save storage, but accumulate a likelihood score during each forward replay for one selected parameter. This simpler organisation suits a small parameter set; their reverse-mode construction targets gradients with respect to many scene parameters. [30]
Reproducibility starts with the identity of a history. Assign random draws to a source-history identifier, event index and sampling purpose, with a further draw index for rejection samplers. Splitting requires persistent descendant identifiers as well. Then compaction or a different GPU scheduling order can move work between threads without changing which variates belong to it. A global seed alone does not provide this property if thread scheduling determines who consumes the next number.
During a backward replay at unchanged parameters, reproduce the same physical and sampling decisions. Retain enough information to reconstruct proposal densities, roulette probabilities, detector contributions and local states. A replay that encounters different random decisions is differentiating another realised computation. Within a parameter perturbation, however, the same variates are allowed to produce different decisions: that change is part of the finite difference.
To assemble derivative_histories, start from resident source states and a HistoryBatch, then run the walker with the selected parameter’s likelihood accumulation. Reusing a batch at unchanged inputs reconstructs its histories, using the counter implementation in Listing 9.6. No preceding forward image is needed to evaluate this derivative estimator.
def derivative_histories(
positions: Any,
directions: Any,
weights: Any,
density: Any,
*,
parameter: TransportParameter,
batch: HistoryBatch,
workspace: TransportWorkspace,
out_pixel: Any,
out_derivative: Any,
out_status: Any,
source_amplitude: float = 1.0,
stream: Any = None,
validate: bool = True,
) -> None:
"""Replay complete histories and write their sparse expected-score derivatives.
One selected parameter costs one complete-history launch and O(H) caller
outputs; no history-by-event tape or material-by-history scratch is retained.
Several parameters can use sequential calls and reuse these buffers. This
trades recomputation for an explicit bounded memory footprint; its useful
parameter-count range requires profiling. Binary64 signed scores are kept
before reduction so uncertainty remains defined at the original-history level.
Forward execution is not a prerequisite. Supplying an earlier batch identity
and unchanged inputs reproduces its paths; independent score and derivative
estimates for nonlinear losses must instead use disjoint batches.
The prepared estimator selects both the flight law and derivative measure.
Each output combines original factors and absorption log weight independently
so forward-score underflow cannot erase a representable derivative.
"""
material = SOURCE_AMPLITUDE if parameter.material is None else parameter.material
if parameter.kind == "log-source-amplitude":
if source_amplitude <= 0:
raise TransportError("log-amplitude derivatives require strictly positive amplitude")
material = LOG_SOURCE_AMPLITUDE
if material >= workspace.spec.grid.materials:
raise TransportError("active material is outside this workspace")
wp = workspace.context.wp
_validate_call(
workspace,
batch,
positions,
directions,
weights,
density,
[
("out_pixel", out_pixel, wp.int32),
("out_derivative", out_derivative, wp.float64),
("out_status", out_status, wp.int32),
],
source_amplitude,
stream,
validate,
)
workspace._launch(
workspace._kernels.derivative_histories,
batch.count,
[
positions,
directions,
weights,
density,
*workspace._model_inputs(),
wp.uint64(batch.seed),
wp.uint64(batch.first_history),
material,
source_amplitude,
out_pixel,
out_derivative,
out_status,
workspace._status,
],
)
if validate:
workspace.check_status()
Each call writes pixel, signed binary64 derivative and terminal status per original history. Several parameters can reuse those buffers sequentially. Memory is therefore proportional to histories plus detector outputs, with no history-by-event tape or retained material-by-history array, while work grows by a full replay for each active parameter. This is an intentional tradeoff for a small declared parameter set. It should be remeasured before extending the inverse problem to hundreds of material parameters.
Keep source and model arrays immutable until replay finishes, and reject a derivative batch containing a failed or truncated history. This manual first-order calculation reconstructs the histories needed by §10.3’s estimator without recording a Warp tape.
For a central difference, let and be two simulation estimates made with common random numbers at and . Write their variances as and and their covariance as . The central difference satisfies
The covariance in equation (10.13) explains why coupling can make a finite difference far less noisy. Positive covariance cancels shared fluctuations. It is a property to check: different rejection counts or highly changed paths can weaken the coupling. Both marginal simulations must still have the intended law.
Use independent groups of coupled simulations to estimate uncertainty. A single fixed set of histories can be excellent for comparing nearby parameters, but it is only one realisation. Repeating its replay adds no independent evidence about the expected derivative. Table 10.1 lists the history state needed to reproduce and differentiate those simulations.
| Quantity | Why the derivative calculation needs it |
|---|---|
| Source and descendant identifiers | Reconstruct the same random streams after scheduling or compaction |
| State at a retained checkpoint | Restart a path segment without regenerating the entire prefix |
| Physical and proposal factors | Reconstruct the contribution and its likelihood or importance terms |
| Roulette and splitting decisions | Preserve weights, descendants and source-history grouping |
| Detector assignment and response | Reproduce the scored output and any boundary treatment |
| Estimator and model versions | Replay the same probability law and numerical conventions |
Figure 10.3 shows how retained checkpoints trade memory for repeated path evaluation.
Full state storage
Derivative traversal ←
Read each retained state when its derivative contribution is needed.
Reconstruct each prefix
Start from the source state whenever an earlier state is needed. Prefixes repeat.
Checkpointed replay
Rebuild 3 → 6, use its intermediate states, then rebuild 0 → 3. Discard each temporary segment after use.
Replay preserves the source, history ID, draw mapping, physical state and model conventions. Reconstructing the same history leaves the number of independent samples unchanged.
10.6 Variance of the gradient estimator
Let be the complete scalar derivative contribution from source history , including the exposure factor and any associated boundary samples. For independent, identically distributed source histories with finite second moment, the gradient estimate and its estimated variance are
Equation (10.14) groups correlated descendants with their source history. If boundary samples use a separate independent estimator, estimate its variance separately and add it. Shared randomness introduces a covariance term. Approximate normal intervals require enough effectively independent contributions for the mean’s sampling distribution to be well behaved. Rare, heavily weighted histories can make a small empirical standard error misleading.
A likelihood estimator often multiplies a noisy contribution by another noisy, signed quantity. Under the fixed-support regularity conditions, normalisation implies . Subtracting a constant baseline from the detector score therefore preserves the likelihood term’s expectation:
The baseline in equation (10.15) changes variance without changing the mean. Expanding the second moment as a quadratic in gives the scalar baseline that minimises the variance of this likelihood term, when the denominator is positive:
Equation (10.16) weights the detector score by the squared likelihood score. Simply using the mean detector signal need not minimise variance. Estimate a baseline from independent pilot histories, or use a conditional baseline independent of the current random choice given its preceding state. Fitting a baseline to the same contribution without the appropriate correction can alter the expected gradient. In a custom derivative rule, the baseline supplies a control variate. Its own parameter dependence is not an extra physical score term.
Other variance reductions can exploit the transport structure. Evaluate a primary contribution and its derivative deterministically, then estimate the remaining scattered contribution with matching normalisation. Replace a binary detector decision by its conditional expected response when that integral is available. Such conditioning also removes the variance contributed by that decision. It may eliminate a discontinuity that made direct pathwise differentiation fail.
Compare uncertainty in physically scaled parameter coordinates. If , where is dimensionless and is a chosen physical scale, then . Both the gradient and its standard error acquire the same factor. A rotation derivative per radian and a translation derivative per mm cannot be compared by their raw magnitudes to decide which parameter is well constrained.
10.7 Independent gradient validation
The homogeneous escape example tests the sampling-law derivative without complicated geometry. The moving aperture tests the missing boundary term. A smoothly varying detector response tests continuous pathwise arithmetic. These examples diagnose different failures, so an agreement on one is insufficient grounds to skip the others. Table 10.2 pairs the controlled problems with their expected derivatives and tested dependencies.
| Controlled problem | Expected derivative | Dependency tested |
|---|---|---|
| Homogeneous primary survival, fixed | Collision and escape probabilities | |
| Aperture edge inside a uniform strip | Moving score discontinuity | |
| Detector gain , fixed transport | Explicit score dependence | |
| Source population, fixed normalised law | Exposure normalisation | |
| Changing a valid proposal only | Zero change in the physical expectation | Importance-weight cancellation |
For a general parameter direction, compare the estimator with central differences of expected measurements over a range of perturbations. Estimate those differences with common random numbers and independent replications. On a sufficiently smooth expectation, central-difference truncation error is quadratic in the perturbation. Its sampling error also depends on the coupling and the discontinuities of individual histories.
The aperture makes that dependence explicit. For , one history contributes to the central difference only when lies in the interval of width swept by the edge. Its derivative contribution is then . The contribution’s mean is exactly , and the variance of an -history average is
Equation (10.17) grows as when becomes small. Common random numbers help, but they do not make an arbitrarily small perturbation useful. With too few histories, none may land in the swept strip. The finite difference is then zero even though the expected derivative is positive. Increasing at a fixed step and varying the step at a fixed distinguish that sampling failure from a missing estimator term.
Figure 10.4 expresses this tradeoff through the exact standard error of the aperture estimator.
A narrower strip gives a noisier derivative
- N = 1,000
- N = 10,000
- N = 100,000
Many draws miss the strip entirely
- N = 1,000
- N = 10,000
- N = 100,000
Uniform aperture: w = 10 mm, a = 5 mm. Both curves use the exact coupled-strip moments.
Report the parameter and its units, perturbation, history count, seed grouping, analytic target where available, and uncertainty of the discrepancy. If the derivative estimate and finite difference share histories, estimate their paired difference directly so that their covariance is included. Use an independently derived formula or estimator as the reference: a finite difference of a fixed, discontinuous sampled trace can share the very omission being tested.
For a vector parameter, also use Chapter 5’s directional checks and the pairing between Jacobian-vector and transpose-Jacobian products. These algebraic checks catch component ordering, frame and accumulation mistakes. They complement the expectation checks above, which catch missing probability and boundary contributions even when the implemented linear algebra is internally consistent.
Begin the implementation checks with slab log density, where analytic survival gives an independent loss derivative. Repeat the estimate with independent batches and compare the discrepancy with its sampling uncertainty. Use the aperture example separately to check any boundary estimator you add.
10.8 A differentiable transport implementation
The transport interface needs to state what its derivative returns. Given physical parameters , an exposure definition and a random-stream specification, the forward calculation estimates the expected detector image . The derivative calculation should accept a fixed detector weight vector and estimate the corresponding parameter sensitivity:
Equation (10.18) allows the implementation to accumulate a transpose-Jacobian product without materialising every pixel-by-parameter entry. The fixed weights combine detector scores before their history derivatives are accumulated. A random weight computed from the same noisy image requires the additional expectation analysis in §10.9.
Construct the contraction one parameter at a time: average its sparse history derivatives into a reusable detector buffer, apply the fixed pixel weights and reduce on the device. This stores one derivative image, which is useful for inspecting spatial sensitivity, and avoids retaining the entire pixel-by-parameter Jacobian. If memory traffic warrants fusing the weighted contraction into the history tally, preserve the same per-history estimator and original-history denominator.
The parameter set determines where contributions accumulate. A spatial material-field derivative can touch many volume entries; a pose derivative has six components summed over histories, with the interface terms from §10.4. Choose output buffers after deriving those contributions. For pose, retain Chapter 3’s frame and translation-before-rotation ordering.
Warp provides array gradients and custom replay and gradient functions for the continuous arithmetic in an estimator. Dynamic loops can overwrite intermediate values needed by the reverse calculation, so explicitly store or reconstruct them. The Warp differentiation documentation explains those mechanisms and their limitations. Combine them with the probability and boundary terms derived for the measurement; allocate retained state according to what that backward calculation needs.
Allocate volume data, cross-section tables, source state, queues, tallies and gradient buffers before the repeated solve. Record the dtype of each buffer and of reductions, because signed contributions can cancel strongly. Clear accumulated gradients between independent evaluations. Keep random-stream identities attached to work items during compaction. Queue overflow and a maximum interaction count require an explicit continuation or unbiased termination scheme: dropping the remaining histories changes both the expected image and its gradient.
When choosing between storage and replay, measure state traffic, launches and host synchronisation alongside elapsed time. Inspect event divergence and detector-tally contention. Vary live histories and active parameters: a small path-local state can still accompany large source, detector and gradient arrays.
A revised sampler needs derivative weights for its revised proposal. A faster reduction may change the accuracy of signed sums through its accumulation order. Caching a physics query must retain the parameter dependence of its result. The forward and derivative calculations must continue to describe the same physical expectation.
10.9 Inverse problems with stochastic gradients
An unbiased image estimate and an unbiased image derivative do not automatically give an unbiased inverse-problem gradient. Consider one pixel with expected signal , observed value and squared-error objective . Let and be unbiased estimates of and . Multiplying the estimated residual by the estimated derivative gives
The covariance in equation (10.19) is a gradient bias when the target is squared error against the expected image. Reusing the same photon histories generally correlates the two estimates. More histories may reduce this term, but its sign and size follow from the joint estimator. They cannot be inferred from two separate unbiasedness claims.
For this quadratic objective, independent batches provide a direct correction. Use batch for the signal estimate and batch for its derivative:
Independence factorises the expectation in equation (10.20). Within each batch, variance reduction and replay can still share random numbers according to their own estimator design. The two batches must remain independent conditional on the current parameters and any fixed pilot information.
Form the weighted residual from image A and multiply it by derivative image B in independent_squared_gradient. Its binary64 pixel contributions are then summed on CUDA. Give A and B non-overlapping history streams. Random source samples must also be independent between them; a fixed deterministic source may be shared. The call checks declared history ranges and source namespaces, so construct the source samples consistently with those declarations.
def independent_squared_gradient(
mean_a: Any,
derivative_mean_b: Any,
observed: Any,
weights: Any,
*,
batches: tuple[HistoryBatch, HistoryBatch],
workspace: TransportWorkspace,
out_components: Any,
stream: Any = None,
validate: bool = True,
) -> None:
"""Write unbiased per-pixel contributions to the loss-of-expected-score gradient.
For fixed finite observation y and nonnegative fixed pixel weight w, the
target is 1/2 sum_p w_p (E[S_p]-y_p)^2. Independent unbiased estimates A of
E[S] and B of its derivative give E[w(A-y)B] = w(E[S]-y)dE[S]. Using the same
histories in both factors generally adds their covariance and is rejected.
Sum components on device using the shared reduction operator. Estimate the
scalar gradient's uncertainty across independent product replicates, because
pixel covariances generally do not vanish.
"""
_product(
mean_a,
derivative_mean_b,
observed,
weights,
batches=batches,
workspace=workspace,
out_components=out_components,
loss=False,
stream=stream,
validate=validate,
)
The companion independent_squared_loss multiplies residual means from two independent batches and applies the one-half factor. A particular estimate can be negative although the true squared-error loss is non-negative. Clipping it would bias the estimator. Candidate and incumbent can share random numbers within factor A and within factor B to reduce the variance of their difference, while A and B remain independent of each other. Pixel correlations are retained by reducing each complete replicate to a scalar before estimating the replicate uncertainty.
There is also an objective-level explanation. The expected squared residual of the Monte Carlo image is
Equation (10.21) contains a penalty for simulation variance. If that variance depends on the parameters, minimising this expectation can favour parameters that are easier for the sampler. Increasing the photon budget changes the unwanted term. A solver should specify whether it targets a loss of the physical expectation or an expectation of a noisy loss.
Assemble TransportSquaredOracle from a fixed pencil or uniformly sampled rectangular ParallelBeam, immutable observations and coefficient tables, and reusable source, history and tally buffers. Express selected density scales and source amplitude in absolute log coordinates. Each trial uploads that small vector, retains image A while replaying B for each derivative, and reduces the resulting loss gradient. For a candidate loss change, reuse each independent factor at both the incumbent and candidate.
def _gradient_replicate(
self,
parameters: Vector,
mean_batch: HistoryBatch,
derivative_batch: HistoryBatch,
capture_model: bool = False,
) -> Vector:
"""Use disjoint source/transport samples for the two nonlinear factors."""
require_no_tape()
require_independent(mean_batch, derivative_batch)
amplitude = self._chart(parameters)
self._mean(mean_batch, amplitude, self._arrays["mean_a"])
inputs = self._source_arrays(derivative_batch)
pixel = self._arrays["pixel"][: derivative_batch.count]
derivative = self._arrays["score"][: derivative_batch.count]
status = self._arrays["status"][: derivative_batch.count]
result: list[float] = []
for column, parameter in enumerate(self.parameters):
derivative_histories(
*inputs,
parameter=parameter,
batch=derivative_batch,
workspace=self.workspace,
out_pixel=pixel,
out_derivative=derivative,
out_status=status,
source_amplitude=amplitude,
stream=self.workspace.context.stream,
validate=False,
)
self.histories_traced += derivative_batch.count
history_mean(
pixel,
derivative,
status,
batch=derivative_batch,
workspace=self._moments,
out_mean=self._arrays["mean_b"],
stream=self.workspace.context.stream,
validate=False,
)
if capture_model:
self.workspace._launch(
self._model_kernels.store_column,
self.workspace.spec.detector.pixels,
[
self._arrays["mean_b"],
column,
self.workspace.spec.detector.pixels,
self._arrays["jacobian"],
],
)
independent_squared_gradient(
self._arrays["mean_a"],
self._arrays["mean_b"],
self.observation,
self.pixel_weights,
batches=(mean_batch, derivative_batch),
workspace=self.workspace,
out_components=self._arrays["components"],
stream=self.workspace.context.stream,
validate=False,
)
value = self._scalar()
if not math.isfinite(value):
raise NumericalError("inverse chart derivative overflow")
result.append(value)
return tuple(result)
def model_replicate(
self,
parameters: Vector,
mean_batch: HistoryBatch,
derivative_batch: HistoryBatch,
) -> tuple[Vector, tuple[Vector, ...]]:
"""Return an unbiased gradient and a PSD *proposal* metric, not an unbiased Hessian.
Jacobian columns and all pixel contractions stay on CUDA. Only the small
dense metric crosses to prepared pinned staging. Sampling variance biases
its diagonal upwards; fresh independent acceptance decides whether to move.
"""
if not self._model_partials:
raise TransportError("prepare the inverse oracle with local_model=True")
context = self.workspace.context
wp = context.wp
try:
gradient = self._gradient_replicate(parameters, mean_batch, derivative_batch, True)
dimension = len(self.parameters)
pixels = self.workspace.spec.detector.pixels
count = (pixels + 255) // 256
wp.launch_tiled(
self._model_kernels.gram_tiles,
dim=dimension * dimension * count,
block_dim=256,
inputs=[
self._arrays["jacobian"],
self.pixel_weights,
pixels,
dimension,
count,
self._model_partials[0],
self.workspace._status,
],
device=context.device,
stream=context.stream,
record_tape=False,
)
previous = self._model_partials[0]
for destination in self._model_partials[1:]:
next_count = (count + 255) // 256
wp.launch_tiled(
self._model_kernels.sum_gram_tiles,
dim=dimension * dimension * next_count,
block_dim=256,
inputs=[previous, count, next_count, destination],
device=context.device,
stream=context.stream,
record_tape=False,
)
previous, count = destination, next_count
wp.copy(self._model_host, previous, stream=context.stream)
self.workspace.check_status()
self.scalar_download_bytes += 8 * dimension * dimension
curvature = tuple(
tuple(float(self._model_view[i * dimension + j]) for j in range(dimension))
for i in range(dimension)
)
if not all(math.isfinite(x) for row in curvature for x in row):
raise NumericalError("proposal curvature exceeds the finite chart range")
return gradient, curvature
finally:
wp.synchronize_stream(context.stream)
def _change_replicate(
self,
before: Vector,
after: Vector,
first: HistoryBatch,
second: HistoryBatch,
) -> float:
"""Independent product factors; common random numbers across parameter points."""
require_no_tape()
require_independent(first, second)
losses: list[float] = []
for parameters in (before, after):
amplitude = self._chart(parameters)
self._mean(first, amplitude, self._arrays["mean_a"])
self._mean(second, amplitude, self._arrays["mean_b"])
independent_squared_loss(
self._arrays["mean_a"],
self._arrays["mean_b"],
self.observation,
self.pixel_weights,
batches=(first, second),
workspace=self.workspace,
out_components=self._arrays["components"],
stream=self.workspace.context.stream,
validate=False,
)
losses.append(self._scalar())
difference = losses[1] - losses[0]
if not math.isfinite(difference):
raise NumericalError("objective-change estimate overflow")
return differenceKeep the image arrays on CUDA and return scalar reductions and status to the controller. Complete pending work before reusing a pinned staging buffer, including on failure. Replaying B across parameters can correlate components within one gradient vector; independent complete replicates supply its uncertainty. Count repeated trace work separately from unique histories when comparing the cost of this design.
Nonlinear likelihoods require their own analysis. For an ideal Poisson count with positive expected count , the negative log-likelihood up to a data-only constant has derivative
Even with independent batches, replacing the reciprocal in equation (10.22) by generally gives bias: an unbiased estimate of does not have an unbiased reciprocal. An estimate equal to zero also makes that substitution undefined. Options include evaluating the mean accurately enough to control the resulting bias or deriving an estimator for the specific nonlinear objective. Flooring the noisy reciprocal changes the estimator and its bias. If the physical mean is floored in the objective itself, differentiate that explicitly modified objective. An energy-integrating detector needs the observation model developed in Chapter 8, rather than inheriting the Poisson formula from a photon counter.
Work in scaled parameter coordinates and choose optimisation steps in relation to the gradient uncertainty. Large steps cannot be justified merely because one noisy gradient has a large magnitude. Reusing histories over several nearby iterations can make the sampled objective easier to optimise, but fresh independent evaluations are needed to find out whether the expected objective improved. Increase the simulation budget when uncertainty prevents resolving a useful update, and keep acquisition noise separate from that numerical decision.
The deterministic strong-Wolfe search in Chapter 6 assumes repeated evaluations of the same objective and derivative. For stochastic transport, separate the histories used to propose a step, assess its loss change and check the returned point. recover_expected_signal uses a trust region to bound the proposed move. A linear model moves to its boundary; a quadratic model can use curvature to choose a smaller correction near a solution.
With proposal="quadratic", minimise the local model inside the Euclidean ball . In one dimension with positive curvature , an interior minimiser is , provided . Thus a smaller estimated slope produces a smaller step without first shrinking the trust radius. Keep the coordinates in the declared log-parameter chart. Estimate from the resident derivative images and reduce the Gram matrix on CUDA; only the small gradient and matrix need to reach the host solve.
Squaring a noisy Jacobian generally biases this curvature estimate. Use it to propose the step, retain independent residual and derivative factors in the gradient, and assess the candidate’s loss change with separate histories. Inspect rank before damping: stabilising a singular matrix does not add information to the acquisition. Chen and colleagues analyse this separation of random models and step assessment under probabilistic accuracy conditions. The empirical error bands used here are a practical decision rule, without those convergence guarantees. [41]
An interior step can be much smaller than the radius. If a candidate is rejected, reuse the incumbent’s local model while contracting the radius, then draw fresh acceptance histories for the new candidate. If acceptance is inconclusive, hold the candidate fixed and increase the batch. Build a new model after an accepted move. Grow the radius only after a boundary-constrained step with sufficient agreement between predicted and estimated decrease.
def recover_expected_signal(
oracle: IndependentSquaredOracle,
initial: Vector,
*,
seed: int,
policy: StochasticPolicy | None = None,
) -> StochasticRecoveryResult:
"""Fresh proposal/acceptance pools with an optional held-out final checkpoint.
Quadratic proposals cache a PSD metric at the unchanged incumbent; rejected
radii use fresh acceptance identities. Final sample size is fixed here before
observing any samples. A failed final checkpoint terminates unresolved and
never feeds a new proposal. ``linear`` with no explicit final batch retains
the original trajectory for ablation, including its heuristic stopping rule.
The explicit ``deterministic_sampling`` oracle flag must be established from
immutable physical/source properties. Empirical variance never establishes
that classification. Its numerical allowance requires independent validation
by the caller. Stochastic norm bands are heuristic, not simultaneous-vector
or repeated-look confidence bounds. Unique identities are not replay costs.
"""
selected: StochasticPolicy = policy or StochasticPolicy()
quadratic = selected.proposal == "quadratic"
integer(seed, "seed", maximum=2**64 - 1)
parameters: Vector = tuple(finite_scalar(value, "initial parameter") for value in initial)
if not parameters:
raise ContractError("a stochastic inverse problem needs active parameters")
if quadratic and len(parameters) > 16:
raise ContractError("dense stochastic models support at most 16 active parameters")
if quadratic and not callable(getattr(oracle, "model_replicate", None)):
raise ContractError("quadratic proposals require an oracle model_replicate operation")
classification = getattr(oracle, "deterministic_sampling", False)
if type(classification) is not bool:
raise ContractError("deterministic_sampling must be an explicit boolean contract")
deterministic = classification
final_size = selected.final_validation_batch
if quadratic and final_size is None:
final_size = selected.initial_batch
final_cost = 0 if final_size is None else 2 * final_size * selected.replicates
first_history = 0
radius = selected.initial_radius
batch_size = selected.initial_batch
damping_relative = selected.damping_relative
rank_tolerance = selected.rank_tolerance
steps: list[StochasticStep] = []
domain_rejections: list[DomainRejection] = []
gradient_attempts: list[GradientAttempt] = []
acceptance_attempts: list[AcceptanceAttempt] = []
reservations: list[WorkReservation] = []
local_models: list[LocalModel] = []
validation_attempts: list[GradientAttempt] = []
validation_trigger: str | None = None
cached: list[LocalModel] = []
band_method = (
"deterministic_numerical_allowance" if deterministic else "heuristic_euclidean_marginal_se"
)
def reserve(size: int, operation: str, iteration: int, holdback: int = 0) -> int | None:
nonlocal first_history
required = 2 * size * selected.replicates
start = first_history
if first_history + required + holdback > selected.unique_history_budget:
reservations.append(
WorkReservation(
operation, iteration, start, required, 0, holdback, (), "budget_denied"
)
)
return None
pairs: list[tuple[HistoryBatch, HistoryBatch]] = []
for _ in range(selected.replicates):
left = HistoryBatch(seed, first_history, size, f"inverse-source-{first_history}")
first_history += size
right = HistoryBatch(seed, first_history, size, f"inverse-source-{first_history}")
first_history += size
pairs.append((left, right))
reservations.append(
WorkReservation(
operation, iteration, start, required, required, holdback, tuple(pairs), "reserved"
)
)
return len(reservations) - 1
def execute(
index: int, operation: Callable[[HistoryBatch, HistoryBatch], _Replicate]
) -> tuple[_Replicate, ...]:
outputs: list[_Replicate] = []
for left, right in reservations[index].pairs:
reservations[index] = replace(
reservations[index], attempted_replicates=len(outputs) + 1
)
try:
outputs.append(operation(left, right))
except Exception as error:
reservations[index] = replace(
reservations[index],
status="domain_error"
if isinstance(error, TrialDomainError)
else "oracle_error",
completed_replicates=len(outputs),
)
raise
reservations[index] = replace(reservations[index], completed_replicates=len(outputs))
reservations[index] = replace(reservations[index], status="complete")
return tuple(outputs)
def result(
reason: Literal[
"gradient_band",
"sampling_unresolved",
"history_budget",
"radius_limit",
"iteration_budget",
],
detail: str = "",
) -> StochasticRecoveryResult:
return StochasticRecoveryResult(
parameters,
reason,
first_history,
tuple(steps),
tuple(domain_rejections),
tuple(gradient_attempts),
detail or reason,
acceptance_attempts=tuple(acceptance_attempts),
reservations=tuple(reservations),
local_models=tuple(local_models),
final_validation_attempts=tuple(validation_attempts),
sampling_classification="verified_deterministic_expectation"
if deterministic
else "stochastic",
proposal_policy=selected.proposal,
uncertainty_policy=band_method,
numerical_gradient_allowance=selected.numerical_gradient_allowance,
final_validation_batch=final_size,
final_validation_trigger=validation_trigger,
)
def record_exception(error: Exception) -> None:
# Preserve the normal numerical exception contract, with a serialisable
# snapshot for callers recording a failed run. Partial device work is
# deliberately unknown here; the production operator records its cost.
error.__dict__["recovery_diagnostics"] = result("sampling_unresolved", "oracle_exception")
error.add_note("Recovery reservation/attempt snapshot is in recovery_diagnostics.")
def log_acceptance(
base: AcceptanceAttempt,
changes: list[float],
outcome: str,
change: float | None = None,
error: float | None = None,
) -> None:
ratio = None if change is None else -change / base.predicted_decrease
acceptance_attempts.append(
replace(
base,
outcome=outcome,
replicate_changes=tuple(changes),
mean_change=change,
standard_error=error,
agreement=ratio if ratio is None or math.isfinite(ratio) else None,
)
)
def finish_stationarity(
iteration: int, trigger: str = "proposal_gradient_band"
) -> StochasticRecoveryResult:
nonlocal validation_trigger
if final_size is None:
return result("gradient_band")
validation_trigger = trigger
index = reserve(final_size, "final_validation", iteration)
if index is None:
return result("history_budget", "final_validation_reservation_failed")
try:
values = execute(index, partial(oracle.gradient_replicate, parameters))
attempt = _gradient_attempt(
values,
parameters,
selected,
deterministic,
iteration=iteration,
batch_size=final_size,
first_history=reservations[index].first_history,
histories_used=reservations[index].reserved_histories,
model_id=None,
pool="final_validation",
)
except Exception as error:
record_exception(error)
raise
validation_attempts.append(attempt)
if attempt.decision == "gradient_band":
return result("gradient_band", "held_out_gradient_band")
return result("sampling_unresolved", "final_validation_failed")
def finish_normal(
reason: Literal["sampling_unresolved", "radius_limit"], detail: str, iteration: int
) -> StochasticRecoveryResult:
# A development decision can be unresolved even when the fixed incumbent
# is stationary. Spend its already reserved, independent checkpoint once.
# Final failure terminates; none of these samples select another point.
if final_size is not None and first_history + final_cost <= selected.unique_history_budget:
return finish_stationarity(iteration, detail or reason)
return result(reason, detail)
for iteration in range(selected.iterations):
if not cached:
# Fresh-batch growth is deliberately retained: there is no pooling
# across adaptive epochs or different parameter values.
while True:
checkpoint = final_cost
if final_size is not None:
checkpoint += 2 * batch_size * selected.replicates
index = reserve(
batch_size,
"model" if quadratic else "gradient",
iteration,
checkpoint,
)
if index is None:
if final_size is not None and any(value.accepted for value in steps):
return finish_stationarity(
iteration, "proposal_checkpoint_reservation_failed"
)
return result(
"history_budget",
"proposal_checkpoint_reservation_failed"
if checkpoint
else "history_budget",
)
model_id = len(local_models) if quadratic else None
try:
curvature: Matrix = ()
if quadratic:
model_oracle = cast(QuadraticSquaredOracle, oracle)
models = execute(index, partial(model_oracle.model_replicate, parameters))
replicates = tuple(value[0] for value in models)
size = len(parameters)
if any(
len(value[1]) != size or any(len(row) != size for row in value[1])
for value in models
):
raise ContractError(
"model curvature dimension differs from the active chart"
)
curvature = tuple(
tuple(
mean_standard_error(tuple(value[1][i][j] for value in models))[0]
for j in range(size)
)
for i in range(size)
)
validate_curvature(curvature, size)
else:
replicates = execute(index, partial(oracle.gradient_replicate, parameters))
attempt = _gradient_attempt(
replicates,
parameters,
selected,
deterministic,
iteration=iteration,
batch_size=batch_size,
first_history=reservations[index].first_history,
histories_used=reservations[index].reserved_histories,
model_id=model_id,
)
except Exception as error:
record_exception(error)
raise
gradient_attempts.append(attempt)
if model_id is not None:
local_models.append(
LocalModel(
model_id,
parameters,
attempt.gradient,
curvature,
attempt.first_history,
attempt.histories_used,
batch_size,
)
)
if attempt.decision == "zero_sample":
if batch_size == selected.maximum_batch:
return finish_normal(
"sampling_unresolved",
"zero_gradient_and_variance_at_maximum_batch",
iteration,
)
batch_size = min(2 * batch_size, selected.maximum_batch)
continue
if attempt.decision == "gradient_band":
return finish_stationarity(iteration)
if attempt.decision == "resolved":
gradient, norm = attempt.gradient, attempt.gradient_norm
if model_id is not None:
cached[:] = [local_models[-1]]
break
if batch_size == selected.maximum_batch:
return finish_normal(
"sampling_unresolved", "gradient_uncertainty_at_maximum_batch", iteration
)
batch_size = min(2 * batch_size, selected.maximum_batch)
else:
gradient, norm = cached[0].gradient, math.hypot(*cached[0].gradient)
proposal: QuadraticProposal | None = None
active_curvature: Matrix = ()
if quadratic:
assert cached
active_curvature = cached[0].curvature
proposal = quadratic_proposal(
gradient,
active_curvature,
radius,
damping_relative=damping_relative,
rank_tolerance=rank_tolerance,
)
step, predicted, boundary = (
proposal.step,
proposal.predicted_decrease,
proposal.boundary,
)
else:
step = tuple(-radius * (value / norm) for value in gradient)
predicted, boundary = radius * norm, True
candidate: Vector = tuple(a + b for a, b in zip(parameters, step, strict=True))
if not math.isfinite(predicted) or predicted <= 0 or not all(map(math.isfinite, candidate)):
raise NumericalError("stochastic proposal exceeds the finite chart range")
if candidate == parameters:
return finish_normal(
"sampling_unresolved", "proposal_below_chart_resolution", iteration
)
# Addition in a large finite chart can round the requested displacement.
# Log the realised move and assess its actual quadratic prediction. The
# legacy linear prediction/decisions remain unchanged for the ablation.
step = tuple(b - a for a, b in zip(parameters, candidate, strict=True))
if proposal is not None:
actual_norm = math.hypot(*step)
if actual_norm > radius * (1 + 1e-12):
return finish_normal(
"sampling_unresolved", "chart_rounding_exceeds_trust_radius", iteration
)
predicted = quadratic_reduction(gradient, active_curvature, step)
boundary = proposal.boundary and actual_norm >= radius * (1 - 1e-12)
proposal = replace(proposal, step=step, predicted_decrease=predicted, boundary=boundary)
look = 0
while True:
look += 1
start = first_history
index = reserve(batch_size, "acceptance", iteration, final_cost)
changes: list[float] = []
threshold = -selected.acceptance_fraction * predicted
base = AcceptanceAttempt(
iteration,
parameters,
candidate,
None if not cached else cached[0].model_id,
step,
predicted,
radius,
batch_size,
start,
first_history - start,
2 * batch_size * selected.replicates,
(),
None,
None,
band_method,
look,
threshold,
"reserved",
None,
boundary,
proposal,
)
if index is None:
log_acceptance(base, changes, "budget_denied")
return result(
"history_budget",
"acceptance_checkpoint_reservation_failed" if final_cost else "history_budget",
)
def change_operation(
a: HistoryBatch,
b: HistoryBatch,
*,
before: Vector = parameters,
after: Vector = candidate,
outputs: list[float] = changes,
) -> float:
value = oracle.change_replicate(before, after, a, b)
if not math.isfinite(value):
raise NumericalError("objective-change replicate is nonfinite")
outputs.append(value)
return value
try:
execute(index, change_operation)
change, standard_error = mean_standard_error(changes)
band = selected.standard_error_multiplier * standard_error
if not math.isfinite(band):
raise NumericalError("objective-change uncertainty band exceeds finite range")
except TrialDomainError as error:
log_acceptance(base, changes, "domain_error")
domain_rejections.append(DomainRejection(iteration, radius, str(error)))
radius *= 0.5
if radius < selected.minimum_radius:
return finish_normal("radius_limit", "radius_limit", iteration)
break
except Exception as error:
log_acceptance(base, changes, "oracle_error")
record_exception(error)
raise
if change + band < threshold:
log_acceptance(base, changes, "accepted", change, standard_error)
steps.append(
StochasticStep(
iteration,
True,
change,
standard_error,
predicted,
radius,
batch_size,
start,
first_history - start,
)
)
parameters = candidate
cached.clear()
if not quadratic or (
boundary and -change / predicted >= selected.radius_growth_agreement
):
radius = min(2 * radius, selected.maximum_radius)
break
if change - band >= threshold:
log_acceptance(base, changes, "rejected", change, standard_error)
steps.append(
StochasticStep(
iteration,
False,
change,
standard_error,
predicted,
radius,
batch_size,
start,
first_history - start,
)
)
radius *= 0.5
if radius < selected.minimum_radius:
return finish_normal("radius_limit", "radius_limit", iteration)
break
log_acceptance(base, changes, "ambiguous", change, standard_error)
if batch_size == selected.maximum_batch:
return finish_normal(
"sampling_unresolved", "acceptance_uncertainty_at_maximum_batch", iteration
)
batch_size = min(2 * batch_size, selected.maximum_batch)
if final_size is not None:
return finish_stationarity(selected.iterations, "iteration_budget")
return result("iteration_budget")
Assess loss changes using a chosen multiple of the empirical standard error across independent scalar replicates. These bands are heuristic: repeated consultation does not give them simultaneous coverage. Under stochastic sampling, an all-zero gradient and error estimate can mean that no useful histories arrived. Increase the batch and return sampling_unresolved if the available budget cannot resolve the estimate. The deterministic absorption limit instead uses a separately declared numerical gradient allowance.
Reserve enough sampling budget to assess a candidate and check the returned point before spending on another model. Test provisional stationarity with fresh held-out gradients that never select a step. If that check fails, return an unresolved result. Keep this numerical stopping decision separate from the parameter-identifiability analysis in Chapter 6.
10.10 Recover a source amplitude
The source-amplitude example connects the whole calculation: construct the scattering problem, estimate its signal and derivative independently, choose a step and assess the returned point with fresh histories. Its accepted trajectory shows two radius-limited moves followed by smaller corrections from the quadratic model.
Keep the geometry and material fixed so that we can see what the stochastic optimiser is doing. An 80 keV pencil beam enters a 1 mm cube with absorption coefficient and isotropic elastic-scattering coefficient . A single detector pixel spans at . The source is at and points along positive . Continuous absorption weighting integrates out absorption events, while scattering paths and detector hits remain random.
_ABSORPTION = 0.3
_SCATTERING = 0.1
_OBSERVATION = 0.08
_INITIAL_AMPLITUDE = 0.08
def _scattering_problem() -> tuple[TransportSpec, ParallelBeam]:
return (
TransportSpec(
MaterialGrid((-0.5, -0.5, 0.0), (1.0, 1.0, 1.0), (1, 1, 1), 1),
PlanarDetector((-2.0, -2.0), (4.0, 4.0), (1, 1), 2.0),
80.0,
"declared isotropic scattering cube for recovery validation",
estimator="continuous-absorption",
),
ParallelBeam((0.0, 0.0, -1.0), (0.0, 0.0)),
)
Let be the expected detector score at unit amplitude and write . Prescribe a target score with unit squared-error weight. The expected signal is , so the objective is and its derivative is . This simple expression gives us an independent way to assess the result. The optimiser still has to estimate its residual and derivative with independent histories: it receives neither nor the implied optimum . The observation is fixed throughout; drawing more simulated photons improves numerical precision without changing the physical problem.
Assemble the sampling batches
Combine the fixed scattering problem above with the recovery loop from §10.9. Generate an independent estimate of the unit-amplitude response for checking the result, keeping it outside the optimiser. The following discussion follows one recorded trajectory.
Give proposal, acceptance and final checking separate history batches, with an independent stream for the reference. When uncertainty prevents useful decisions, increase sampling throughout the solve. Increasing only the final check leaves earlier decisions just as noisy.
Follow the accepted updates
The first two steps reach the trust-radius boundary and increase the amplitude substantially. The third is an interior quadratic-model step. After it, the estimated derivative changes sign, so the fourth makes a small correction in the opposite direction. As the gradient shrinks, the proposed move shrinks without repeatedly halving the radius.
Figure 10.5 follows the amplitude through these updates.
Accepted amplitude updates
- Accepted source amplitude
Recorded plot values
Simulation records for this example. Download the
complete first trajectory
,
all 32 independent assessments
,
reference batches
and
data sources and computation details
.
After the fourth accepted update, assess the gradient with fresh histories. The stopping rule requires its absolute estimate plus twice its empirical standard error to lie below the declared tolerance. Here the independent reference also places the derivative inside that tolerance. Use the reference to check the returned point, keeping it out of step selection.
Sampling uncertainty and parameter error
If the gradient band crosses the tolerance, the sampling precision cannot resolve the stopping decision. Increasing the history count can narrow that band without changing the objective or relaxing the tolerance.
A small gradient can also coexist with appreciable density error when strong attenuation makes the signal insensitive to density. The measurement is weakly sensitive to density in this regime. Curvature can improve the proposed step, but a precise density estimate may still require a more informative acquisition.
Check predictions on withheld views or exposures, using fresh simulation histories. Fresh histories assess numerical error; withheld observations test whether the fitted model predicts another measurement. Inspect residuals and parameter sensitivities together. If source intensity can trade against attenuation, constrain one of them or add an observation that separates their effects.
References
- Mohamed, Shakir, Rosca, Mihaela, Figurnov, Michael and Mnih, Andriy (2020). Monte Carlo Gradient Estimation in Machine Learning. Journal of Machine Learning Research, 21(132), 1-62. https://jmlr.org/papers/v21/19-346.html
- Li, Tzu-Mao, Aittala, Miika, Durand, Frédo and Lehtinen, Jaakko (2018). Differentiable Monte Carlo Ray Tracing through Edge Sampling. ACM Transactions on Graphics, 37(6), 222:1-222:11. https://doi.org/10.1145/3272127.3275109
- Vicini, Delio, Speierer, Sébastien and Jakob, Wenzel (2021). Path Replay Backpropagation: Differentiating Light Paths using Constant Memory and Linear Time. ACM Transactions on Graphics, 40(4), 108:1-108:14. https://doi.org/10.1145/3450626.3459804