Гайд · TNWS AI
Как ограничить gradient каждого примера через Optax
Optax 0.2.6: per_example_global_norm_clip отдельно ограничивает batch-градиенты и возвращает их сумму и число clipping-событий.
Задача и применимость
Актуальный запрос: «как ограничить норму градиента каждого примера Optax per_example_global_norm_clip». Ограничить L2-норму каждого примера до суммирования. Входные листья обязаны иметь batch в оси 0. Для gradients [3,4] с norm 5 и [0,6,0,8] с norm 1 при threshold 1 первый пример масштабируется до [0,6,0,8], второй остаётся прежним; сумма равна [1,2,1,6], clipped count равен 1. Для 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-тестом и следуйте документации установленной версии.
Пошаговое решение
- Подготовьте per-example gradients с batch axis 0.
- Вызовите функцию с l2_norm_clip=1.
- Прочитайте сумму clipped gradients и счётчик.
- При необходимости разделите сумму на batch size отдельно и явно.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
grads = [jnp.array([[3.0, 4.0], [0.6, 0.8]])]
summed, clipped = optax.per_example_global_norm_clip(grads, l2_norm_clip=1.0)
assert jnp.allclose(summed[0], jnp.array([1.2, 1.6]))
assert int(clipped) == 1
mean = summed[0] / grads[0].shape[0]
assert jnp.allclose(mean, jnp.array([0.6, 0.8]))
print(summed, clipped, mean)
Вход и ожидаемый результат
Функция вернула сумму [1,2,1,6], clipped count 1 и после явного деления среднее [0,6,0,8]. Она не усредняет batch автоматически.
Как проверить результат
Сначала сравните per-example norms с threshold. Проверьте, что все листья имеют одинаковый batch size. Логируйте долю clipped examples; резкий рост может означать сдвиг данных или ошибку loss scaling.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт batch: порядок осей, dtype, диапазон признаков и семантику target. Инициализацию, dropout и генерацию данных разделяйте по именованным RNG streams. Обучаемые nnx.Param, статистическое состояние и optimizer state сохраняйте и проверяйте раздельно. Перед обновлением Flax/JAX прогоняйте небольшой обезличенный regression batch и сравнивайте формы, loss, конечность значений и допустимое численное отклонение.
Критерии приёмки
- Версии Flax, JAX и jaxlib записаны в отчёте.
- Формы входа, параметров и результата совпадают с контрактом.
- Все численные outputs и gradients конечны.
- Train/eval режим задан явно для стохастических и stateful слоёв.
- Результат проверен независимой формулой, состоянием или контрольным вызовом.
- Негативный случай неверной формы или режима включён в тесты проекта.
Регрессионный протокол
Сохраните минимальные входы: один объект, небольшой batch, нулевые значения, несимметричные размеры осей и заведомо несовместимую форму. Сравнивайте массивы через allclose с документированным допуском. Воспроизводимость требует одинаковых seeds и одинакового порядка запросов RNG. Производительность оценивайте после JIT-прогрева и block_until_ready, отдельно от compile time. Один удачный запуск не заменяет validation на независимых данных.
Диагностика расхождений
При расхождении сначала проверьте layout и preprocessing, затем RNG streams, train/eval mode, типы переменных State, forward, loss и optimizer update. Упростите модель до одного слоя, отключите JIT и стохастику, сравните промежуточные shapes. Возвращайте компоненты по одной. Не увеличивайте допуск и не заменяйте NaN нулями до установления причины.
Контроль после внедрения
В эксплуатации ограничьте допустимые shapes и размер batch, чтобы запрос не исчерпал память и не вызвал поток перекомпиляций. Не логируйте исходные пользовательские тензоры, ключи и содержимое обучающих примеров. Храните предыдущий проверенный checkpoint и схему State для отката. После восстановления проверяйте контрольный forward до включения модели в трафик.
Границы применимости
CPU-проверка не доказывает скорость или точную численную идентичность GPU/TPU. Случайная инициализация не является обучением, снижение train loss не гарантирует обобщение, а правильная форма не подтверждает корректность данных. Для важных решений нужны validation/test, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не передавать уже усреднённый gradient без batch axis.
- Не считать возвращённый tree средним.
- Не путать число clipped examples с числом clipped элементов.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Что возвращает функция?
Сумму clipped gradients и число затронутых примеров.
Где batch axis?
В нулевой оси каждого массива.
Как получить среднее?
Явно разделить сумму на размер batch.
Официальные источники
Частые вопросы
Что возвращает функция?
Сумму clipped gradients и число затронутых примеров.
Где batch axis?
В нулевой оси каждого массива.
Как получить среднее?
Явно разделить сумму на размер batch.
Читайте также
Как исключить bias из weight decay через Optax
Optax 0.2.6: маска add_decayed_weights применяет decay к kernel и оставляет bias без изменения.
Как накапливать gradients через Optax MultiSteps
Optax 0.2.6: MultiSteps усредняет два microbatch-gradient и обновляет SGD только на каждом втором вызове.
Как обновлять target network через Optax periodic_update
Optax 0.2.6: periodic_update копирует online parameters на шагах, кратных update_period, и сохраняет target между ними.
Комментарии
Пока тихо. Скажите первое слово