PyTorch tensors and autograd

A 3D Gaussian splatting trainer (train.py in graphdeco-inria/gaussian-splatting) is, at heart, eight lines repeated 30 000 times:

for iteration in range(1, opt.iterations + 1):
    render_pkg = render(viewpoint_cam, gaussians, pipe, background)
    image = render_pkg["render"]
    Ll1 = l1_loss(image, gt_image)
    loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
    loss.backward()
    with torch.no_grad():
        ...  # densify, prune, reset opacity
        gaussians.optimizer.step()
        gaussians.optimizer.zero_grad(set_to_none=True)

Everything interesting is hidden in three places: how loss.backward() finds the gradient of a scalar with respect to millions of raw parameters, how the rasterizer — a CUDA kernel autograd can't see into — takes part in that, and how the parameters are shaped so that an unconstrained optimiser can't push them into nonsense. This lesson covers all three in plain PyTorch. The C++ side (gfx-differentiable-rendering) derived the same gradients by hand; here you'll see how research code wires them into autograd.

1. Tensors and the graph they record

A torch.Tensor is an n-dimensional array like a NumPy ndarray — same broadcasting rules, same ... indexing — with three differences that matter here:

• the default floating dtype is float32 (NumPy's is float64); torch.from_numpy(a) keeps the array's dtype and shares its memory;
• it can live on a GPU (.to("cuda")); the sandbox here is CPU-only, and research code is written so that the same lines run on both;
• when any input has requires_grad=True, every operation records itself in a graph, so that the gradient can be computed later.

Tensors you create with requires_grad=True are leaves: the parameters. Every result computed from them has a grad_fn — the node that knows how to send a gradient back to its inputs.

import torch

logit = torch.tensor([0.0, 2.0], requires_grad=True)   # a raw parameter: a leaf
color = torch.tensor([1.0, 0.5])                        # data: no gradient needed
opacity = torch.sigmoid(logit)
weight = opacity * 0.5
loss = (weight * color).sum()

for name, t in [("logit", logit), ("opacity", opacity), ("weight", weight), ("loss", loss)]:
    print(f"{name:8} leaf={t.is_leaf!s:5} requires_grad={t.requires_grad!s:5} grad_fn={type(t.grad_fn).__name__}")
print("loss.grad_fn.next_functions:", [type(f).__name__ for f, _ in loss.grad_fn.next_functions])

2. backward() and gradient accumulation

loss.backward() computes ∂loss/∂leaf for every leaf that requires grad and adds it into leaf.grad. It never overwrites: a second backward adds a second gradient. That's deliberate — gradient accumulation over micro-batches, or one loss term per camera, needs no extra code — and it's why every training step clears the gradients first:

• optimizer.zero_grad(set_to_none=True) sets every .grad to None (the default since PyTorch 2.0; cheaper than filling with zeros);
• the graph is freed after backward() (its saved tensors are released), so calling it twice raises unless the first call passed retain_graph=True;
• backward() needs a scalar; for a tensor output y, y.backward(v) computes the vector-Jacobian product vᵀ ∂y/∂x — which is all reverse mode ever computes, and why one backward costs about as much as one forward whatever the number of parameters.

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

for step in range(2):
    loss = (x ** 2).sum()           # d loss / dx = 2x
    loss.backward()
    print(f"after backward {step}: x.grad = {x.grad.tolist()}")
x.grad = None                       # what zero_grad(set_to_none=True) does

y = x * torch.tensor([10.0, 20.0, 30.0])
y.backward(torch.tensor([1.0, 0.0, 1.0]))   # v^T J for a vector output
print("vector-Jacobian product:", x.grad.tolist())

# torch.autograd.grad returns gradients instead of accumulating them.
(g,) = torch.autograd.grad((x ** 3).sum(), x)
print("autograd.grad:", g.tolist(), "| x.grad untouched:", x.grad.tolist())

3. Switching recording off: no_grad, detach, item

Recording costs memory (every saved tensor stays alive until backward), and some computations must not be differentiated at all:

model.eval() is unrelated: it switches modules like dropout and batch norm to inference behaviour and records just as much. A hand-written optimiser shows the first row in action — the update itself must not be recorded, and the leaf must be modified in place so it stays the same leaf:

import torch

torch.manual_seed(0)
target = torch.tensor([0.25, -1.0])
w = torch.zeros(2, requires_grad=True)

for step in range(50):
    loss = ((w - target) ** 2).sum()
    loss.backward()
    with torch.no_grad():          # an SGD step: not part of any graph
        w -= 0.1 * w.grad
    w.grad = None
print("w:", [round(v, 4) for v in w.tolist()], "| still a leaf:", w.is_leaf)

d = w.detach()
print("detached: requires_grad =", d.requires_grad, "| same storage as w:", d.data_ptr() == w.data_ptr())

detach() never copies: write into d in place and you have changed w. Use w.detach().clone() for an independent copy — the idiom for "make trainable copies of these tensors".

4. In-place operations and the version counter

Methods with a trailing underscore (add_, mul_, sigmoid_, clamp_) and indexed assignment (t[mask] = 0) modify a tensor in place. Autograd allows it, with two rules:

1. A leaf that requires grad can't be modified in place outside no_grad — its value is the parameter, and its .grad refers to it.
2. A tensor saved for backward can't be modified afterwards. Every tensor has a version counter that in-place ops bump; saved tensors remember the version, and backward raises if it changed. sigmoid and exp save their output (σ' = σ(1 − σ), exp' = exp); x * y saves both inputs; x * 3 saves nothing of x.

import torch

leaf = torch.zeros(3, requires_grad=True)
try:
    leaf.add_(1)
except RuntimeError as e:
    print("leaf:", str(e).split(".")[0])

s = torch.sigmoid(leaf)
s.mul_(2)                                # overwrites what sigmoid saved
try:
    s.sum().backward()
except RuntimeError as e:
    print("saved tensor:", str(e).split(":")[0])

e = torch.exp(leaf)
e = e * 2                                # out of place: a new tensor
e.sum().backward()
print("out of place is fine:", leaf.grad.tolist())

A third trap is aliasing. expand, view, slicing and transpose return views that share memory. expand in particular gives stride-0 dimensions: writing into an expanded tensor writes the same element many times (PyTorch refuses with "unsupported operation: more than one element of the written-to tensor refers to a single memory location"). That's why 3DGS's SSIM window is built as _2D_window.expand(channel, 1, ws, ws).contiguous() — the .contiguous() makes a real copy before anyone writes to it.

5. Parameter activations: optimise free numbers

Adam treats every parameter as an unconstrained real number. Physical quantities aren't: a scale is positive, an opacity lies in (0, 1), a rotation quaternion has unit length. Research code stores a raw parameter and applies an activation in every forward pass — 3DGS's GaussianModel:

The activation shapes the optimisation, not just the range. With scale = exp(s), ∂scale/∂s = scale: an Adam step of roughly constant size in s is a constant factor on the scale, right for sizes spanning millimetres to metres. With opacity = σ(ℓ), ∂opacity/∂ℓ = σ(1 − σ), which vanishes as the opacity saturates: a Gaussian whose logit drifted to −10 is effectively frozen — one reason 3DGS prunes Gaussians below opacity 0.005 and resets all opacities to 0.01 every 3000 iterations.

The normalisation is part of the function too, and its gradient has a geometric meaning: q / ‖q‖ doesn't change when q is scaled, so the gradient on the raw quaternion is orthogonal to q — (I − q̂q̂ᵀ) g / ‖q‖. Detach the norm (a tempting "it's just a constant") and you lose that projection.

import torch

torch.manual_seed(0)
raw = torch.randn(4, 4, requires_grad=True)             # four raw quaternions
q = raw / raw.norm(dim=-1, keepdim=True)                 # keepdim: (4, 4) / (4, 1)
(q * torch.randn(4, 4)).sum().backward()
print("‖q‖:", [round(v, 6) for v in q.norm(dim=-1).tolist()])
print("grad · raw per quaternion:", [round(v, 6) + 0.0 for v in (raw.grad * raw).sum(-1).tolist()])

raw.grad = None
q_bad = raw / raw.norm(dim=-1, keepdim=True).detach()
(q_bad * torch.randn(4, 4)).sum().backward()
print("with a detached norm:     ", [round(v, 3) for v in (raw.grad * raw).sum(-1).tolist()])

6. A custom torch.autograd.Function

Autograd through the rasterizer's per-pixel loop would record a node for every multiply of every Gaussian at every pixel — gigabytes of saved tensors. So 3DGS (diff_gaussian_rasterization._RasterizeGaussians) and gsplat (_RasterizeToPixels) are custom Functions: forward launches a CUDA kernel and returns tensors; backward launches a hand-derived kernel. Inside forward nothing is recorded — the whole call is one node — and you decide what to keep with ctx.save_for_backward.

The core of that backward is front-to-back alpha compositing. For one pixel with samples sorted front to back,

rgb = Σ_k w_k c_k,   w_k = T_k α_k,   T_k = Π_{j<k} (1 − α_j)

T_k is the exclusive product — the light that reaches sample k, before k absorbs any. torch.cumprod is inclusive, so research code shifts it: nerf-pytorch writes torch.cumprod(torch.cat([torch.ones((N, 1)), 1. - alpha + 1e-10], -1), -1)[:, :-1].

The gradient has two parts. c_k appears only in its own term: ∂rgb/∂c_k = w_k. But α_k appears in w_k and in every later T_m, through the factor (1 − α_k):

∂rgb/∂α_k = T_k c_k − Σ_{m>k} w_m c_m / (1 − α_k)

The second term is how a front Gaussian learns that becoming more opaque hides what's behind it. Contracted with the incoming gradient g = ∂L/∂rgb, the suffix sum over m > k is a reversed cumulative sum:

import torch


class Composite(torch.autograd.Function):
    @staticmethod
    def forward(ctx, alpha, color):                  # (..., K), (..., K, C)
        trans = torch.cumprod(1 - alpha, dim=-1)
        trans = torch.cat([torch.ones_like(trans[..., :1]), trans[..., :-1]], dim=-1)
        ctx.save_for_backward(alpha, color, trans)
        return ((trans * alpha).unsqueeze(-1) * color).sum(-2)

    @staticmethod
    def backward(ctx, g):                            # g: (..., C)
        alpha, color, trans = ctx.saved_tensors
        w = trans * alpha
        gc = (color * g.unsqueeze(-2)).sum(-1)       # g · c_k
        behind = (w * gc).flip(-1).cumsum(-1).flip(-1) - w * gc
        return trans * gc - behind / (1 - alpha), w.unsqueeze(-1) * g.unsqueeze(-2)


torch.manual_seed(0)
alpha = (torch.rand(3, 5, dtype=torch.float64) * 0.99).requires_grad_(True)
color = torch.rand(3, 5, 3, dtype=torch.float64, requires_grad=True)
rgb = Composite.apply(alpha, color)
print("rgb shape:", tuple(rgb.shape), "| grad_fn:", type(rgb.grad_fn).__name__)
print("gradcheck:", torch.autograd.gradcheck(Composite.apply, (alpha, color)))

The CUDA rasterizer computes the same thing per pixel, walking its sorted list back to front and accumulating the colour behind — and instead of saving every T_k it recomputes them by dividing the final transmittance by (1 − α), trading a division for memory. That division is one reason 3DGS clamps α to 0.99:

// diff-gaussian-rasterization, renderCUDA backward (simplified)
float T = T_final;
for (int j = last; j >= 0; --j) {
    T = T / (1.f - alpha[j]);                       // T_j, recovered
    float dchannel_dcolor = alpha[j] * T;           // ∂rgb/∂c_j = w_j
    float dL_dalpha = 0.f;
    for (int ch = 0; ch < C; ++ch) {
        dL_dalpha += (color[j][ch] - accum_rec[ch]) * dL_dpixel[ch];
        accum_rec[ch] = alpha[j] * color[j][ch] + (1.f - alpha[j]) * accum_rec[ch];
    }
    dL_dalpha *= T;                                 // T_j (c_j − colour behind) · g
}

accum_rec there is Σ_{m>k} w_m c_m / T_{k+1} — the colour behind, renormalised — so T_k (c_k − accum_rec) is exactly the formula above. A few more rules for Functions: call them with .apply (never instantiate); return one gradient per forward argument, shaped like it, or None for non-tensors and inputs that need none (ctx.needs_input_grad[i] says which); use ctx.mark_dirty if forward modifies an input in place.

7. The optimiser loop: Adam and parameter groups

Adam (Kingma & Ba 2015) keeps per-element running means of the gradient m and of its square v, and steps by lr · m̂ / (√v̂ + ε) — about lr in magnitude whatever the gradient's scale. That's what makes the activations above behave: step sizes are set by the learning rate, in the raw parameter's units. Because positions, log-scales, logits, quaternions and SH coefficients live on very different scales, 3DGS's training_setup gives each its own parameter group with its own rate (positions 1.6e-4 × scene extent with an exponential decay, scales 5e-3, rotations 1e-3, SH DC 2.5e-3, the higher SH bands 20× less):

l = [
    {"params": [self._xyz], "lr": training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
    {"params": [self._features_dc], "lr": training_args.feature_lr, "name": "f_dc"},
    {"params": [self._opacity], "lr": training_args.opacity_lr, "name": "opacity"},
    {"params": [self._scaling], "lr": training_args.scaling_lr, "name": "scaling"},
    {"params": [self._rotation], "lr": training_args.rotation_lr, "name": "rotation"},
]
self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)

(The name keys let densification find each group's tensor and its Adam state when it concatenates new Gaussians or prunes old ones.)

Put together — fit two semi-transparent layers so that, composited, they show a target colour; every raw parameter is unconstrained, and copies keep the caller's tensors intact:

import torch

torch.manual_seed(0)
target = torch.tensor([0.8, 0.4, 0.2])
init = {"logit": torch.zeros(2), "color": torch.zeros(2, 3)}
p = {k: v.detach().clone().requires_grad_(True) for k, v in init.items()}
opt = torch.optim.Adam([{"params": [p["logit"]], "lr": 0.05}, {"params": [p["color"]], "lr": 0.1}])

for step in range(301):
    opt.zero_grad(set_to_none=True)
    alpha = torch.sigmoid(p["logit"]) * 0.99
    trans = torch.cat([torch.ones(1), torch.cumprod(1 - alpha, 0)[:-1]])
    rgb = ((trans * alpha)[:, None] * torch.sigmoid(p["color"])).sum(0)
    loss = ((rgb - target) ** 2).sum()
    loss.backward()
    opt.step()
    if step % 100 == 0:
        print(f"step {step:3}: loss {loss.item():.2e}  rgb {[round(v, 3) for v in rgb.tolist()]}")
print("init untouched:", bool((init["logit"] == 0).all()))

In an interview

Expect "what does backward() compute and why is it cheap?" (a vector-Jacobian product per node, reverse order: one backward ≈ one forward, whatever the number of parameters), "why zero_grad?" (accumulation), "detach vs no_grad?" (one tensor vs a region), and the in-place error messages. For a splatting or NeRF role: "how would you make a CUDA kernel differentiable?" — a torch.autograd.Function, saved tensors, a hand-derived backward, gradcheck in float64 against a slow reference — and "why are opacities stored as logits and scales as logs?" Be ready to derive the compositing gradient on a whiteboard: the term most people forget is the occlusion term − Σ_{m>k} w_m c_m / (1 − α_k).