Skip to main content

ADAM

Constant ADAM 

Source
pub const ADAM: &str = r#"
#include <metal_stdlib>
using namespace metal;

kernel void optirs_adam(
    device float* x [[buffer(0)]],
    device const float* y [[buffer(1)]],
    device float* a [[buffer(2)]],
    device float* b [[buffer(3)]],
    device const float* result [[buffer(4)]],
    uint idx [[thread_position_in_grid]])
{
    uint n = as_type<uint>(result[7]);
    if (idx >= n) { return; }

    float lr = result[0];
    float beta1 = result[1];
    float beta2 = result[2];
    float eps = result[3];
    float wd = result[4];
    float bc1 = result[5];
    float bc2 = result[6];

    float p = x[idx];
    float g = y[idx];
    if (wd > 0.0f) { g = g + wd * p; }

    float mi = beta1 * a[idx] + (1.0f - beta1) * g;
    float vi = beta2 * b[idx] + (1.0f - beta2) * g * g;
    a[idx] = mi;
    b[idx] = vi;

    float m_hat = mi / bc1;
    float v_hat = vi / bc2;
    x[idx] = p - lr * m_hat / (sqrt(v_hat) + eps);
}
"#;
Expand description

Adam with coupled L2 weight decay. See super::wgsl::ADAM for bindings.