Storage and Gradient Snapshots

Reverse-mode differentiation needs information from the forward trajectory. TIDE stores selected forward-time states on the compute device so the adjoint pass can correlate them with reverse fields and accumulate material gradients. StorageOptions controls whether snapshots are retained, how they are represented, and whether a configured device-capacity limit permits the call.

Why storage matters

Snapshot demand grows with spatial volume, shot batch size, number of stored field components, and number of sampled time levels. A forward-only call may need no derivative snapshots. A VJP or full Hessian-vector product can require substantial state, even when the final receiver tensor is small.

The receiver data shape is therefore not a useful estimate of gradient memory.

Storage modes

Mode	Behavior	Use
device	Keep snapshots on the compute device	Lowest access latency
none	Do not allocate snapshots	Forward-only work or compatible reference paths
auto	Use device when the capacity limit permits it	Guard device memory explicitly

CPU and disk snapshot modes are not supported. TIDE does not silently spill snapshots away from the compute device.

Configuring storage

storage = tide.StorageOptions(
    mode=tide.StorageMode.AUTO,
    compression="bf16",
    bytes_limit_device=6 * 1024**3,
)

operator = tide.MaxwellTM(
    discretization,
    experiment,
    storage=storage,
)

Capacity limits should leave headroom for material tensors, live wavefields, receiver data, gradients, optimizer state, CUDA context, and unrelated processes. auto raises MemoryError when the required snapshot capacity exceeds bytes_limit_device; it never falls back to host memory or disk. Do not set the device limit equal to total advertised VRAM.

Compression

compression="bf16" stores eligible float32 snapshots in bfloat16 while propagation and model tensors remain float32. This reduces storage volume and transfer traffic, but quantizes stored states.

For non-dispersive native Maxwell VJPs, TIDE stores already-combined physical-parameter statistics rather than coefficient-space gradients. Snapshot storage is allocated only for the requested material parameters, so requesting additional material gradients increases capacity proportionally to the state components needed by those gradients.

Validate compression on the quantity the workflow uses:

1. Run the same small objective with and without compression.
2. Compare receiver data to confirm the primal path is unchanged.
3. Compare gradient norms and a directional derivative.
4. Compare one or more optimizer steps, not only one gradient image.

A visually similar gradient can still change line-search behavior.

Temporal sampling

model_gradient_sampling_interval reduces the cadence used for model-gradient accumulation. It is separate from storage mode and compression. Increasing the interval can reduce stored time levels and correlation work, but it changes the derivative approximation. Establish an acceptable interval against the default value of 1.

Shot batching and storage

Reducing shot batch size lowers simultaneous wavefield and snapshot demand. The total computation remains similar, while peak memory falls. Gradient accumulation must preserve the desired dataset normalization across batches.

A useful tuning order is:

1. Measure one shot with device storage.
2. Increase batch size until memory or throughput stops improving.
3. If one shot does not fit, try BF16 snapshots.
4. Increase gradient sampling interval only after a derivative comparison.
5. If snapshots still do not fit, reduce the spatial problem or use an explicitly supported recomputation/reference strategy.

Diagnosing storage failures

Device out of memory before propagation

Model, optimizer state, or batch size is already too large. Snapshot placement cannot solve allocations that occur before the forward trajectory begins.

Device out of memory during forward

Reduce shot or model batch size, enable BF16 snapshots, or increase the configured device limit when capacity is available.

Reverse pass is much slower than forward

Compression conversion or sparse access to stored snapshots may dominate. Profile the full forward-backward operation.

See configuration for storage fields and performance for measurement strategy.