optax_adamw

phasic.optax_wrapper.optax_adamw(
    learning_rate=0.001,
    b1=0.9,
    b2=0.999,
    eps=1e-08,
    weight_decay=0.01,
)

Create Optax AdamW optimizer wrapped for phasic.

AdamW implements Adam with decoupled weight decay regularization, which can help prevent overfitting.

Parameters

learning_rate : float or optax.Schedule = 0.001

Learning rate.

b1 : float = 0.9

Exponential decay rate for first moment.

b2 : float = 0.999

Exponential decay rate for second moment.

eps : float = 1e-8

Small constant for numerical stability.

weight_decay : float = 0.01

Weight decay coefficient.

Returns

: OptaxOptimizer

Wrapped AdamW optimizer compatible with phasic SVGD.