ggml_opt_step_adamw: AdamW Optimizer Step (Graph Node)
In ggmlR: 'GGML' Tensor Operations for Machine Learning

ggml_opt_step_adamw

R Documentation

AdamW Optimizer Step (Graph Node)

Description

Builds a node that applies one AdamW step when the graph is computed:

m <- beta1 * m + (1 - beta1) * grad
v <- beta2 * v + (1 - beta2) * grad^2
a <- a * (1 - alpha * wd) - alpha * (m * beta1h) / (sqrt(v * beta2h) + eps)

a, m and v are updated in place; the returned tensor is a view of a. Because one kernel reads and writes each element, the update is free of the read-before-write ordering that a chain of separate ops would need. With wd = 0 the decay factor is exactly 1.

Usage

ggml_opt_step_adamw(ctx, a, grad, m, v, adamw_params)

Arguments

`ctx`	GGML context
`a`	Weight tensor, F32, contiguous. It need not be marked with `ggml_set_param`: the flag ggml asserts is set only while the node is built, then the tensor's own flags are restored (a lasting PARAM flag would turn the weight into a node of every later graph)
`grad`	Gradient, same shape as `a`, F32
`m`	First moment, same shape as `a`, F32
`v`	Second moment, same shape as `a`, F32
`adamw_params`	F32 tensor of 7 elements: `alpha, beta1, beta2, eps, wd, beta1h, beta2h`

Details

Precision: the kernel receives beta2 as a float and forms 1 - beta2 itself, so the weight of grad^2 in v is 1 - fl(beta2), not 1 - beta2. Its relative error is 9.5e-7 for 0.99, 1.29e-5 for 0.999 and 1.66e-4 for 0.9999 – the closer beta2 is to 1, the larger. The effect on the weight is at most about half of that. Expect agreement with a double-precision Adam to that level, not to the last bit.

Bias correction is passed in, not derived: for step t, beta1h = 1 / (1 - beta1^t) and beta2h = 1 / (1 - beta2^t).

Only F32 is supported (the CPU and Vulkan kernels are F32-only); any other type or a shape mismatch is an R error.