Гайды · · 4 мин
Как исключить bias из weight decay через Optax
Optax 0.2.6: маска add_decayed_weights применяет decay к kernel и оставляет bias без изменения.
Задача и применимость
Актуальный запрос: «как исключить bias из weight decay Optax add_decayed_weights». Применить add_decayed_weights только к kernel. Маска имеет ту же структуру PyTree: True включает transform, False исключает лист. При нулевых gradients, weight_decay 0,1, kernel 2 и bias 3 выход должен содержать добавку 0,2 для kernel и 0 для bias. Сам transform ещё не меняет знак и не умножает на learning rate. Для production отдельно зафиксируйте структуру PyTree, dtype, момент вызова относительно JIT и порядок соседних transformations. Добавьте негативный тест на несовместимую маску, неверный счётчик или отсутствующий batch axis — в зависимости от задачи. Численный unit test должен выполняться при каждом обновлении Optax.
Материал проверен на Optax 0.2.6, JAX/jaxlib 0.11.2 и Python 3.12 на CPU. Используется актуальный API from flax import nnx, а не устаревший экспериментальный import и не Linen. Пример не загружает веса или внешние данные и решает одну изолированную задачу. Численный тест подтверждает контракт кода в указанной версии, но не качество нейросети на рабочем датасете.
Перед запуском установите совместимые версии пакетов в отдельное окружение и выведите flax.__version__, jax.__version__ и jaxlib.__version__. Не смешивайте разные способы накопления, маскирования и усреднения: похожие названия имеют разные контракты. При переносе сначала закрепите ожидаемую математику unit-тестом и следуйте документации установленной версии.
Пошаговое решение
- Создайте PyTree kernel/bias и нулевые gradients.
- Задайте boolean mask с True только для kernel.
- Вызовите add_decayed_weights и передайте params в update.
- Проверьте, что bias остался нулевым в updates.
Готовый проверяемый шаблон
import jax
import jax.numpy as jnp
import optax
params = {"kernel": jnp.array([2.0]), "bias": jnp.array([3.0])}
grads = jax.tree.map(jnp.zeros_like, params)
mask = {"kernel": True, "bias": False}
tx = optax.add_decayed_weights(0.1, mask=mask)
updates, _ = tx.update(grads, tx.init(params), params)
assert jnp.allclose(updates["kernel"], jnp.array([0.2]))
assert jnp.allclose(updates["bias"], jnp.array([0.0]))
print(updates)