Гайд · TNWS AI
Почему Optax l2_loss включает множитель 0,5
Optax 0.2.6: l2_loss возвращает 0,5·(prediction-target)² поэлементно; проверяем значения и отличие от MSE.
Задача и применимость
Актуальный запрос: «как посчитать L2 loss Optax». Показать точный контракт l2_loss: это половина квадрата ошибки для каждого элемента, а не готовое mean squared error. Для errors 1 и 2 результат равен 0,5 и 2. Множитель 0,5 упрощает производную, но означает, что численное значение вдвое меньше обычного squared error при одинаковой агрегации. При использовании sample weights согласуйте, умножаются ли поэлементные значения до или после reduction. Для неравных batch sizes на последнем шаге усредняйте метрику по числу объектов, а не по числу batch.
Материал проверен на 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__. Не смешивайте loss-функции с похожими названиями: они различаются ожидаемыми labels, формой результата и численной нормировкой. Сначала закрепите logits/targets/reduction контракт тестом и следуйте документации установленной версии.
Пошаговое решение
- Создайте predictions и targets одной или broadcast-compatible формы.
- Вызовите l2_loss без reduction.
- Сравните с 0.5*(predictions-targets)**2.
- Выберите mean или sum согласно контракту обучения.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
predictions = jnp.array([1., 3.])
targets = jnp.array([0., 1.])
losses = optax.l2_loss(predictions, targets)
manual = 0.5 * (predictions - targets) ** 2
assert jnp.allclose(losses, jnp.array([0.5, 2.0]))
assert jnp.allclose(losses, manual)
print(losses, losses.mean())
Вход и ожидаемый результат
Ошибки 1 и 2 превратились в [0.5,2.0]; mean равен 1,25. Обычный MSE на этих данных был бы 2,5.
Как проверить результат
Запишите reduction в названии метрики и тесте. Не сравнивайте отчёты разных библиотек без проверки множителя и mean/sum. Для multi-output определите, какие оси усредняются первыми.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт 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, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не называть element-wise l2_loss готовым MSE.
- Не терять множитель 0,5 при сравнении метрик.
- Не позволять broadcasting скрыть ошибочную форму target.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Почему есть 0,5?
Это стандартная нормировка, упрощающая производную.
Есть ли встроенный mean?
Нет, loss поэлементный.
Это L2-регуляризация весов?
Нет, здесь сравниваются predictions и targets.
Официальные источники
Частые вопросы
Почему есть 0,5?
Это стандартная нормировка, упрощающая производную.
Есть ли встроенный mean?
Нет, loss поэлементный.
Это L2-регуляризация весов?
Нет, здесь сравниваются predictions и targets.
Читайте также
Как использовать Huber loss в Optax для регрессии с выбросами
Optax 0.2.6: Huber loss квадратичен около нуля и линеен после delta; проверяем ошибки 0,5 и 2.
Как использовать sigmoid focal loss в Optax
Optax 0.2.6: focal loss для logits 0 и labels 0/1 при gamma=2; проверка симметричного базового значения.
Как посчитать binary cross-entropy по logits в Optax
Optax 0.2.6: sigmoid_binary_cross_entropy принимает logits и метки 0/1; проверяем базовый уровень ln(2).
Комментарии
Пока тихо. Скажите первое слово