ggml_ssm_scan_back: Backward Pass of the Selective State-Space Scan
In ggmlR: 'GGML' Tensor Operations for Machine Learning

ggml_ssm_scan_back

R Documentation

Backward Pass of the Selective State-Space Scan

Description

Gradients of ggml_ssm_scan with respect to all six of its differentiable inputs, given grad, the gradient of the loss with respect to the scan's packed forward result.

Usage

ggml_ssm_scan_back(ctx, s, x, dt, A, B, C, ids, grad)

Arguments

`ctx`	GGML context
`s`	Initial states, as passed to the forward scan
`x`	Inputs, as passed to the forward scan
`dt`	Per-token timestep, as passed to the forward scan
`A`	State decay, as passed to the forward scan
`B`	Input projection, as passed to the forward scan
`C`	Output projection, as passed to the forward scan
`ids`	Sequence-to-state mapping, as passed to the forward scan
`grad`	Gradient w.r.t. the packed forward result

Details

Training reaches this automatically through the backward graph. It is exported so the gradient can be checked directly against a numeric one.

Value

A flat tensor holding the six gradients, concatenated

Packed result

All six gradients come back in one flat tensor, in the same order the scan takes its inputs: [d_s | d_x | d_dt | d_A | d_B | d_C]. ids is integer and has no gradient. A single op returns a single tensor, and packing them keeps the recurrence from being replayed once per input.

Cost

The forward pass does not keep its intermediate states, so this replays the recurrence before walking it back – roughly twice the forward work, and no extra memory beyond one state buffer.