When an encoder is frozen but a linear head still trains, no_grad is usually the more suitable starting point than replacing the extraction region with inference_mode. Features need not receive gradients themselves for their values to be needed during backward. The distinction matters at the boundary between a frozen producer and a trainable consumer.

Why can inference_mode affect downstream training even after the feature extractor has been frozen? A two-dimensional head makes the dependency visible. We will derive the values backward needs, then use that dependency to select a mode and define acceptance checks. 中文版.

A weight gradient can require features without feature gradients

Let the encoder output h=(2,−1), the head weights be w=(3,4), prediction z=w·h=2 and target y=1. For L=(z−y)²/2, the loss is 0.5. We train only w and treat h as a fixed input for this step. The chain rule gives ∂L/∂w=(z−y)h=(2,−1).

There is no requirement here to compute ∂L/∂h. Nevertheless, h appears explicitly in the weight-gradient formula. Discarding it after forward generally leaves insufficient information to reconstruct both weight gradients from the scalar loss. Different features and residuals can produce the same loss. The derivative formula, rather than an input’s gradient flag alone, determines which values remain necessary.

For a concrete counterexample, replace the features with (0,0) while keeping the same target. The loss is again 0.5, but the weight gradient becomes (0,0). Equal losses and unequal gradients disprove the idea that retaining only the loss is sufficient. This argument does not depend on network size or any assumption about a particular GPU kernel.

The two modes make different boundary commitments

The PyTorch 2.8 autograd notes distinguish outputs from no-grad computations that can subsequently be used in grad mode from inference mode’s restrictions. Inference mode is intended for computations without later autograd interaction. Reducing tracking overhead does not establish unconditional compatibility with a differentiable consumer.

The inference_mode API documentation identifies differences involving view tracking, version counters and forward-mode automatic differentiation. This context manager does not infer the lifetime of every output. The caller must understand its consumers, especially whether a subsequent operation needs to save it for backward.

For the linear head above, excluding upstream differentiation is only one requirement. Training the downstream weights still needs the feature values. Write both requirements into the boundary contract: do not build the unwanted upstream graph, and preserve valid downstream differentiation. Seeing requires_grad=False on a feature does not test the second requirement. Nor should the restriction be summarized as every later operation always failing. What matters is whether the actual operation participates in backward recording and needs to save the relevant value. Validate the entire intended path rather than extrapolating from one successful arithmetic operation.

Choose the boundary before choosing the mode

Freezing parameters can itself be expressed through their requires_grad settings. Whether a computation should additionally disable gradient recording depends on its input dependencies. A blanket context around every frozen module loses that distinction. It can be correct for one pipeline and wrong for another using the same module weights.

If an encoder and all subsequent processing only produce final predictions used for display, ranking or persistence, their complete lifetime can be considered as a candidate inference region. Verify that returned objects are not reused by another training path. With a frozen encoder feeding a trainable head, no_grad is generally the more suitable extraction boundary.

A different arrangement is easy to misclassify. Suppose a trainable adapter precedes the encoder. Its parameters need gradients propagated through the encoder, even though the encoder’s own parameters are frozen. Wrapping the entire encoder call in no_grad merely because its parameters do not update can remove the signal needed by the adapter. Keeping parameters fixed and removing derivatives with respect to inputs are separate decisions.

Ask three questions: which parameters should update, which inputs need derivatives, and whether a later backward pass needs to read the outputs. These are data-flow questions. Folder names such as training or inference cannot answer them. Even an offline feature-generation script can produce data eventually consumed by a training stage; its name is not evidence of a compatible object boundary.

Nine scalar checks test dependencies, not framework errors

The following program uses only Python’s standard library. It computes prediction, residual and analytic head gradients, then checks each weight derivative with central differences. One update with step size 0.1 reduces the loss from 0.5 to 0.125. Finally, zero features provide the equal-loss, different-gradient counterexample.

# Scalar dependency check, not an implementation of PyTorch grad modes.
from math import isclose
h=(2.,-1.); w=(3.,4.); target=1.
def loss(weights,features=h):
    prediction=sum(a*b for a,b in zip(weights,features))
    return (prediction-target)**2/2
prediction=sum(a*b for a,b in zip(w,h))
residual=prediction-target
grad=tuple(residual*x for x in h)
assert prediction == 2.
assert residual == 1.
assert grad == (2.,-1.)
for j in range(2):
    plus=list(w); minus=list(w)
    plus[j]+=1e-5; minus[j]-=1e-5
    numerical=(loss(plus)-loss(minus))/(2e-5)
    assert isclose(numerical,grad[j],rel_tol=1e-8)
assert loss(w)==0.5
new_w=tuple(a-0.1*g for a,g in zip(w,grad))
assert isclose(loss(new_w),0.125)
# No gradient for features does not make their values irrelevant.
assert loss(w,(0.,0.))==0.5
assert (residual*0.,residual*0.) != grad
print('9 scalar dependency checks passed')

These nine assertions verify mathematical dependencies. PyTorch is unavailable in the execution environment, so no inference tensor, model backward or device test was run. An expected framework error is therefore not presented as an observed result. Versioned official documentation supports the mode semantics; executed scalar calculations support the numbers. These are different forms of evidence for different claims.

The finite differences perturb w while holding h and y fixed. Recomputing a feature that itself depends on w at every perturbation would test another function. For a head trained on frozen features, first establish the mathematical objective, then check whether the framework path implements it.

Another boundary mistake is leaving the head forward inside the no_grad context too. That disables the intended recording for the trainable consumer, not merely for feature extraction. Draw the region at the call level: leave the context after the encoder call, then perform the head forward, loss and backward in the normal gradient region. Check project decorators and nested contexts so that they do not silently widen the boundary. A correct-looking indentation example in isolation cannot validate the actual call stack.

An absence of errors is an incomplete acceptance test

A project-level test should include an actual head forward pass, loss, backward pass and update, instead of only running the encoder and inspecting feature shapes. Check that the intended head parameters receive finite gradients, then compare an interpretable example against an analytic result or numerical differences. Separately verify that frozen encoder parameters do not receive gradients the design excludes.

For an input-side adapter trained through a frozen module, explicitly inspect the adapter’s gradients too. Checking only that the encoder has no gradients can miss the case where the entire upstream signal has been accidentally cut. Specify which gradients should be zero and which should be absent; treating every zero as success hides legitimate failures.

If an existing pipeline uses inference_mode, first narrow its context or switch the extraction region to no_grad and validate the complete training path. Do not scatter tensor copies through the program solely to suppress an error. Copies consume storage, and their placement inside or outside a mode matters. Establish a correct, explainable baseline before optimizing copying and caching.

Module state, cache ownership and performance need separate checks

A model’s eval state controls behavior that depends on training state; it does not replace a gradient-mode decision. The earlier SDPA Dropout investigation also shows why custom call parameters require inspection. The new question here is whether an output crosses into a training region. Each question needs its own acceptance criteria.

Likewise, no_grad does not automatically preserve features as an independent numerical snapshot. If a producer later reuses its buffer, check storage ownership separately; the detach and clone storage comparison develops that issue. Separating graph relationships, module state and storage lifetime prevents one mechanism from becoming a supposed remedy for every failure.

A performance comparison must hold inputs, device, precision, module state and the measured call boundary constant. Measure the metrics actually required. No timing was collected here, so no speedup percentage is promised. A pipeline that copies outputs, preserves them across stages or resumes training cannot be evaluated by timing only its encoder. The result offered here is a decision method, not an unexecuted throughput claim.

If the two-dimensional derivation and boundary tests help diagnose gradients around a frozen model, voluntary support through the section below is welcome. The reusable procedure is to identify the values backward needs before choosing where inference ends.

Sources and backfill note

Inference mode API · Autograd mechanics · no_grad API

Checked against PyTorch 2.8 documentation available before the planned slot; the autograd notes were updated June 16, 2025. Diagrams and scalar examples are original. Execution covered standard-library numerical checks only. This backfill retains October 21, 2026, 14:00 Beijing as its article date; actual publication time is recorded separately.