Losses and optimizers

Applicable version · ArcQML 0.1.0

On this page

Losses: input constraints and forward formulas

Function Key constraints
mse_loss Nonempty prediction and target, same dtype (F32/F64), broadcast-compatible shapes.
l1_loss Nonempty prediction and target, same dtype (F32/F64), broadcast-compatible shapes.
binary_nll_loss One-dimensional prediction and labels of equal length; labels must be 0 or 1.
BCE with logits logits and targets have identical shape and dtype; target values lie between 0 and 1.
cross_entropy_loss A two-dimensional logits Tensor; the labels count equals batch size, and every label is a valid class index.

These losses use the current CPU Dense numerical operators. Let N be the element count of the broadcasted error Tensor, with the i th broadcasted elements of prediction and target denoting predicted and target values. For a [B, C] error tensor, for example, N = B × C. Mean squared error is:

LMSE=12Ni=1N(y^iyi)2.L_{\mathrm{MSE}} =\frac{1}{2N}\sum_{i=1}^{N} \left(\hat{y}_i-y_i\right)^2.

The L1 loss is:

LL1=1Ni=1Ny^iyi.L_{\mathrm{L1}} =\frac{1}{N}\sum_{i=1}^{N} \left\lvert\hat{y}_i-y_i\right\rvert.

In binary_nll_loss, prediction is a Pauli-Z expectation, not an ordinary probability or logit. The framework first clamps the expectation:

ziclip=clamp ⁣(zi,1+2ε,12ε),ε=1012.z_i^{\mathrm{clip}} =\operatorname{clamp}\!\left( z_i,-1+2\varepsilon,1-2\varepsilon \right), \qquad \varepsilon=10^{-12}.

The clamped expectation is then converted to probabilities for two classes:

pi(0)=1+ziclip2,pi(1)=1ziclip2.p_i(0)=\frac{1+z_i^{\mathrm{clip}}}{2}, \qquad p_i(1)=\frac{1-z_i^{\mathrm{clip}}}{2}.

The corresponding binary negative log-likelihood is:

Lbinary NLL=1Ni=1N[(1yi)logpi(0)+yilogpi(1)],yi{0,1}.L_{\mathrm{binary\ NLL}} =-\frac{1}{N}\sum_{i=1}^{N} \left[ (1-y_i)\log p_i(0)+y_i\log p_i(1) \right], \qquad y_i\in\{0,1\}.

Binary cross entropy taking logits directly uses a numerically stable form, where N is the total number of logits:

LBCE=1Ni=1N[max(xi,0)+log ⁣(1+exi)yixi].L_{\mathrm{BCE}} =\frac{1}{N}\sum_{i=1}^{N} \left[ \max(x_i,0)+\log\!\left(1+e^{-\lvert x_i\rvert}\right)-y_ix_i \right].

For multiclass logits with B samples and C classes, cross entropy is:

LCE=1Bb=1Blog ⁣[softmax(Xb)yb].L_{\mathrm{CE}} =-\frac{1}{B}\sum_{b=1}^{B} \log\!\left[ \operatorname{softmax}(X_b)_{y_b} \right].

Optimizers

SGD: updates with coupled L2 regularization

Sgd::step updates each Parameter that is trainable, has requires_grad enabled, and has a grad. For F32 or F64, each element is updated as:

θθη(g+λθ).\theta\leftarrow\theta-\eta\left(g+\lambda\theta\right).
η=learning rate,g=current gradient,λ=weight decay coefficient.\eta=\text{learning rate}, \qquad g=\text{current gradient}, \qquad \lambda=\text{weight decay coefficient}.

Parameters without gradients count toward skipped_no_grad; frozen parameters or those with requires_grad=false count toward skipped_frozen.

Adam: bias correction and weight decay placement

At the first step, Adam creates state slots by position in the parameter slice. Later calls must preserve parameter count and order. It stores the step count, first moment, and second moment. Each update first adds the coupled L2 term to the current gradient:

g~t=gt+λθt1.\widetilde{g}_t =g_t+\lambda\theta_{t-1}.

Then it updates the first and second moments:

mt=β1mt1+(1β1)g~t,vt=β2vt1+(1β2)g~t2.\begin{aligned} m_t&=\beta_1m_{t-1}+(1-\beta_1)\widetilde{g}_t,\\ v_t&=\beta_2v_{t-1}+(1-\beta_2)\widetilde{g}_t^2. \end{aligned}

After bias correction, it updates the parameters:

m^t=mt1β1t,v^t=vt1β2t,θt=θt1ηm^tv^t+ε.\begin{aligned} \widehat{m}_t&=\frac{m_t}{1-\beta_1^t},\\ \widehat{v}_t&=\frac{v_t}{1-\beta_2^t},\\ \theta_t&=\theta_{t-1} -\eta\frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}+\varepsilon}. \end{aligned}

Here weight_decay therefore enters the moment estimates and is not the decoupled weight decay used by AdamW.

Constructors require finite, nonnegative learning_rate and weight_decay, beta1 and beta2 in [0, 1), and finite positive epsilon. Reuse the same Adam instance during training; recreating it loses momentum and step counts.

Recommended training loop

rust
let mut optimizer = Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0)?;
for _ in 0..steps {
    let prediction = simulator.run(&circuit, &observable)?;
    let loss = mse_loss(&prediction, &target)?;
    loss.backward()?;
    optimizer.step(circuit.parameters())?;
    optimizer.zero_grad(circuit.parameters());
}

Call zero_grad after each optimizer.step to clear current gradients and prevent accidental accumulation in the next iteration. Do not replace or modify a forward-pass Parameter Tensor in place before backward, or version checks may report modified forward data. Let optimizer.step perform updates.

Custom losses

Rust users can connect quantum outputs to differentiable operators such as sub, square, mul, and mean to build scalar losses. Ordinary operator composition requires no handwritten backward pass. To specify derivatives for out-of-graph computation or a complete expression, use CustomOp and apply_custom_op. Autograd follows local backward rules in the actual graph; it does not symbolically simplify arbitrary equivalent mathematical expressions.

See the Custom loss tutorial (docs/tutorial/custom_loss.md in the release package) and Custom loss example (rust/custom_loss.rs in the release package) for full input contracts, weighted MSE, batch circuit training, analytical/numerical gradient checks, and custom backward examples. Python does not currently export the corresponding general operator-composition or CustomOp interfaces.