Gradient Safety and Differentiability

NeuralPOM is designed for end-to-end differentiable simulation -- the gradient of any scalar loss with respect to model parameters, initial conditions, or embedded neural network weights must flow through every numerical operation in the time loop. This page documents the strategies used to maintain gradient safety across hundreds of thousands of time steps.

The Core Challenge

The Princeton Ocean Model (like most legacy Fortran ocean codes) was originally written around in-place array updates -- the same memory locations are overwritten at each time step. This is fundamentally incompatible with PyTorch's autograd engine, which builds a directed acyclic computation graph and prohibits in-place mutations of tensors that are needed for gradient computation.

Consider this simplified example of the problem:

# Problematic pattern: in-place update of a leaf tensor
x = torch.zeros(10, requires_grad=True)
for t in range(100):
    x[2:5] = x[2:5] + 0.1  # RuntimeError: a leaf Variable that
                            # requires grad is being used in an
                            # in-place operation.

The error occurs because x is a leaf tensor (created directly, not from an operation), and leaf tensors with requires_grad=True cannot be modified in-place since autograd needs to retain their original values for backward computation.

Strategy 1: Leaf Variable Cloning

The primary defense mechanism in NeuralPOM is to clone all prognostic leaf variables into intermediate tensors before entering the main time loop. This converts them from leaf tensors (which autograd protects) into intermediate tensors (which can be safely updated in-place as part of the computational graph).

The cloning block at the entry point of both time loops:

# From src/neuralpom/cores/explicit/time_loop.py
# CRITICAL FIX FOR AUTOGRAD: Clone Leaf Variables
# Convert leaf tensors from the initial condition into
# intermediate tensors so later in-place updates are valid.
self.u = self.u.clone()
self.v = self.v.clone()
self.w = self.w.clone()
self.t = self.t.clone()
self.s = self.s.clone()
self.rho = self.rho.clone()

self.ua = self.ua.clone()
self.va = self.va.clone()
self.el = self.el.clone()
self.et = self.et.clone()
self.elb = self.elb.clone()
self.etb = self.etb.clone()

self.q2 = self.q2.clone()
self.q2l = self.q2l.clone()
self.km = self.km.clone()
self.kh = self.kh.clone()
self.kq = self.kq.clone()

self.uab = self.uab.clone()
self.vab = self.vab.clone()

# Surface forcing variables
self.wusurf = self.wusurf.clone()
self.wvsurf = self.wvsurf.clone()
self.wtsurf = self.wtsurf.clone()
self.wssurf = self.wssurf.clone()
self.swrad = self.swrad.clone()
self.vfluxf = self.vfluxf.clone()

After cloning, self.u, self.v, etc. are no longer leaf tensors -- they are intermediate results of the clone() operation. PyTorch's autograd tracks subsequent in-place operations on intermediate tensors by recording them as graph mutations in the backward pass.

Clone Before Any Mutation

The cloning must occur before any in-place tensor updates. If a tensor has already been mutated in-place, cloning it afterward will not recover the original value needed for the backward pass. The clone block is placed at the very beginning of run_time_loop, immediately after slice definitions.

Strategy 2: Clone-on-Write for Slice Assignments

Even with the initial cloning, direct slice assignments like self.u[1:, :, :kbm1] = ... can cause issues if the tensor already participates in a computation graph. The time loop uses a clone-on-write pattern: create a fresh clone of the tensor, perform the slice assignment on the clone, then rebind the attribute to the clone:

# From src/neuralpom/cores/explicit/time_loop.py
# FIX: Clone u to avoid in-place modification of graph leaf/node
u_new = self.u.clone()
u_new[1:self.im, :, :self.kbm1] = (
    u_slice - tps_slice
) + corr_slice
self.u = u_new  # rebind to new tensor

# Same for v
v_new = self.v.clone()
v_new[:, 1:self.jm, :self.kbm1] = (
    v_slice - tps_v_slice
) + corr_v_slice
self.v = v_new

This pattern is used consistently throughout the baroclinic correction step, the Smagorinsky viscosity update, and the Robert-Asselin filter step:

# Smagorinsky viscosity update with clone-on-write
aam_new = self.aam.clone()
aam_new[s_i, s_j, :self.kbm1] = smag_visc
self.aam = aam_new

# Robert-Asselin filter with clone-on-write
u_filtered = self.u.clone()
u_filtered[:,:,:self.kbm1] = (
    self.u[:,:,:self.kbm1]
    + 0.5 * self.smoth * (term_u - tps_unsqueezed)
)
self.ub = u_filtered
self.u = self.uf.clone()

Strategy 3: Time-Level Shifts via Clone

The Robert-Asselin filter shifts time levels at the end of each internal step. These shifts must use .clone() to break reference sharing, preventing autograd from attempting to trace cycles:

# From src/neuralpom/cores/explicit/time_loop.py
# Shift time levels (correct pattern)
self.elb = self.el.clone()
self.el = self.elf.clone()
self.d = self.h + self.el
self.uab = self.ua.clone()
self.ua = self.uaf.clone()
self.vab = self.va.clone()
self.va = self.vaf.clone()

Without .clone(), writing self.elb = self.el would make elb an alias for el, and the subsequent time loop iteration would overwrite the gradient history.

Strategy 4: the generation_fast_path Flag

The leaf variable cloning adds overhead to every run_time_loop call. For data generation runs (where gradients are not needed), this overhead is unnecessary. The implicit core provides a generation_fast_path flag that conditionally skips the cloning block:

# From src/neuralpom/cores/implicit/time_loop.py
generation_fast_path = getattr(self, "generation_fast_path", False)

if not generation_fast_path:
    # CRITICAL FIX FOR AUTOGRAD: Clone Leaf Variables
    self.u = self.u.clone()
    self.v = self.v.clone()
    # ... all clones ...

The flag is set by run_generation_day in the accelerated implicit core:

# From src/neuralpom/cores/implicit/core.py
def run_generation_day(self, steps_per_day, ...):
    self.generation_fast_path = True
    try:
        self.run_time_loop()
        return list(self._generation_reference_records)
    finally:
        self.generation_fast_path = False

When to use generation_fast_path

• ON (True): Offline training data generation, pure inference, benchmarking throughput. State snapshots use .detach() to ensure they do not participate in any computation graph.
• OFF (False): Training runs with backpropagation, any workflow that requires gradient flow through the time loop.

When generation_fast_path is active, the time loop also captures snapshots:

if generation_fast_path:
    capture_interval = getattr(
        self, "_generation_capture_interval_steps", 0
    )
    if capture_interval > 0 and \
       self.iint <= capture_limit and \
       self.iint % capture_interval == 0:
        self._generation_reference_records.append(
            self._capture_generation_reference_snapshot()
        )

Each snapshot uses .detach().clone() to create gradient-free copies:

def _capture_generation_reference_snapshot(self):
    return {
        "time": float(self.time),
        "u": self.u.detach().clone(),
        "v": self.v.detach().clone(),
        "t": self.t.detach().clone(),
        "s": self.s.detach().clone(),
        # ...
    }

Safe Numerical Utilities

Several numerical operations common in ocean models are numerically unstable for autograd in their naive form. The neuralpom.utils.safety module provides gradient-safe alternatives:

# From src/neuralpom/utils/safety.py
SAFE_DIV_EPS = 1e-20
SAFE_SQRT_EPS = 1e-8
SAFE_LOG_EPS = 1e-12

def safe_divide(numerator, denominator, eps=SAFE_DIV_EPS):
    return numerator / (denominator + eps)

def safe_sqrt(x, eps=SAFE_SQRT_EPS):
    return torch.sqrt(torch.clamp(x, min=eps))

def safe_abs_pow(x, exponent, eps=SAFE_SQRT_EPS):
    return torch.pow(
        torch.clamp(torch.abs(x), min=eps), exponent
    )

def safe_log(x, eps=SAFE_LOG_EPS):
    return torch.log(torch.clamp(x, min=eps))

def safe_reciprocal(denominator, eps=SAFE_DIV_EPS):
    return 1.0 / (denominator + eps)

These utilities prevent NaN and inf values that would otherwise break gradient propagation. For instance, safe_sqrt is used in the Smagorinsky viscosity computation and the bottom drag formulation:

# Smagorinsky viscosity (explicit time_loop.py)
smag_visc = self.horcon * dx_sl * dy_sl \
    * safe_sqrt(term1 + term2 + term3)

# Bottom drag velocity magnitude (explicit momentum.py)
mag_vel = safe_sqrt(u_sq + v_avg_sq)

The equation of state uses safe_abs_pow for the $S^{1.5}$ term, which would produce NaN gradients for negative salinity values:

# From explicit momentum.py, dens()
rhor += (-5.72466e-3 + 1.0227e-4 * tr - 1.6546e-6 * tr2) \
    * safe_abs_pow(sr, 1.5)

Gradient Clipping in Training

During neural network training, the accumulated gradient through many timesteps of the ocean model can become large or explosive. Gradient clipping is applied in the training loop:

# From src/neuralpom/training/trainer.py
if (self.args.grad_clip_norm is not None
    and float(self.args.grad_clip_norm) > 0
    and self.trainable_params):
    torch.nn.utils.clip_grad_norm_(
        self.trainable_params,
        max_norm=float(self.args.grad_clip_norm)
    )

The default clipping norm is configured in config.py:

# From src/neuralpom/config.py
grad_clip_norm: float = 1.0

This clips the global norm of all trainable parameter gradients to 1.0, preventing gradient explosion in long backpropagation-through-time sequences.

Deep Freeze Pattern (Momentum Module)

The profu and profv vertical friction solvers use a deep freeze pattern to protect the tridiagonal solve from autograd graph entanglement. During the forward sweep, coefficients a and c are frozen with .clone() to prevent gradient leakage through recursive dependencies, then ee and gg are similarly frozen before the backward substitution:

# From src/neuralpom/cores/explicit/momentum.py, profu()
# DEEP FREEZE 1: Freeze 'a' and 'c'
a_frozen = a.clone()
c_frozen = c.clone()

# Forward sweep using frozen a, c ...
# (ee and gg computed)

# DEEP FREEZE 2: Freeze 'ee' and 'gg'
ee_frozen = ee.clone()
gg_frozen = gg_store.clone()

# Backward substitution using frozen ee, gg ...

This pattern ensures the tridiagonal solution remains numerically stable while remaining fully differentiable -- the gradient flows through the clone operations without the numerical instability that would arise from directly autograd-differentiating the recursive sweep.

Common Pitfalls and Debugging

Pitfall 1: "a leaf Variable that requires grad is being used in an in-place operation"

Symptom: PyTorch raises a RuntimeError during the time loop.

Cause: A prognostic tensor (e.g., self.u, self.v, self.el) was not cloned before an in-place update.

Fix: Ensure all prognostic variables are cloned at the top of run_time_loop. If adding a new prognostic variable, add it to the clone block.

Pitfall 2: Gradient scale diverges exponentially with time steps

Symptom: loss.backward() produces NaN or inf gradients after many time steps.

Cause: Unbounded gradient accumulation through the recurrent time loop.

Fix:

1. Verify grad_clip_norm is set to a reasonable value (default 1.0).
2. Check that the Robert-Asselin filter is being applied (it provides numerical damping).
3. Reduce the number of time steps in the differentiable window.

Pitfall 3: Zero gradients after the CG solver

Symptom: Gradients are zero after backward through the implicit barotropic step.

Cause: The accelerated PCG solver's .item() calls and while-loop control flow break gradient propagation.

Fix: The accelerated core automatically falls back to the reference CG implementation when torch.is_grad_enabled() is True. If you have modified the solver, ensure no .item() or Python-native control flow appears in the gradient-enabled path.

Pitfall 4: Memory explosion during backpropagation

Symptom: CUDA out-of-memory errors during loss.backward().

Cause: The autograd graph retains intermediate activations for the entire time loop history.

Fix:

1. Use torch.utils.checkpoint.checkpoint on the heaviest subroutines (turbulence, tracers, momentum).
2. Reduce the number of differentiable time steps.
3. Use mixed-precision training (torch.cuda.amp).

Pitfall 5: Inconsistent results between generation_fast_path and normal mode

Symptom: Data generated with generation_fast_path=True produces different physical fields than a normal run.

Cause: With generation_fast_path, no leaf cloning occurs, so the initial condition tensors are mutated directly. Reusing the same core object for both modes without reinitialization can lead to state corruption.

Fix: Always call run_initialization() before switching between generation and gradient modes, or use separate core instances.

Best Practices Summary

1. Always clone before mutating: Use .clone() before any slice assignment on tensors that are part of the autograd graph.
2. Use .clone() for time-level shifts: Never do self.elb = self.el; always do self.elb = self.el.clone().
3. Use safe_* utilities: Avoid raw torch.sqrt, /, and torch.log in numerical code; use the wrapped versions from neuralpom.utils.safety.
4. Detach generation snapshots: Always use .detach().clone() when capturing model state for offline use.
5. Separate generation from training: Use generation_fast_path during data generation, but always reinitialize before entering a gradient-tracked run.
6. Monitor gradient norms: Log gradient statistics during training to detect unhealthy gradient flow early.