Гайд · TNWS AI
Почему cross-entropy даёт NaN при −inf и как помогает Optax safe
Optax 0.2.6: safe_softmax_cross_entropy трактует 0·log(0) как 0 и избегает NaN для нецелевого −inf logit.
Задача и применимость
Актуальный запрос: «почему softmax cross entropy дает NaN при minus inf Optax safe». Воспроизвести особый случай logits [0,-inf] и one-hot labels [1,0]. Обычная softmax_cross_entropy формирует NaN из произведения нулевой метки и -inf; safe_softmax_cross_entropy следует соглашению 0·log(0)=0 и возвращает конечный loss 0. Это не делает корректной строку, где все logits равны -inf. Такой случай часто появляется после masking attention или недопустимых классов. Валидируйте, что у каждой строки остаётся хотя бы один разрешённый класс и целевой класс никогда не замаскирован.
Материал проверен на 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 контракт тестом и следуйте документации установленной версии.
Пошаговое решение
- Создайте один конечный целевой logit и
-infтолько у нулевой метки. - Вычислите обычную и safe cross-entropy.
- Подтвердите NaN первой и конечный ноль второй.
- Отдельно отклоняйте полностью masked строки до loss.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
logits = jnp.array([[0., -jnp.inf]])
labels = jnp.array([[1., 0.]])
regular = optax.softmax_cross_entropy(logits, labels)
safe = optax.safe_softmax_cross_entropy(logits, labels)
assert jnp.isnan(regular[0])
assert jnp.allclose(safe, jnp.array([0.]))
assert jnp.isfinite(safe).all()
print(regular, safe)
Вход и ожидаемый результат
Обычная функция вернула [nan], safe-вариант — [0]. Целевой класс имел probability 1, поэтому нулевая потеря математически ожидаема.
Как проверить результат
Проверьте, что labels неотрицательны и суммируются в 1. Перед loss обнаруживайте строки без единого конечного logit. Не заменяйте все NaN safe-функцией: источник может быть в параметрах, данных или переполнении.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт 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, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не считать safe-вариант лечением произвольных NaN.
- Не допускать все logits=-inf.
- Не передавать integer labels: здесь нужны распределения формы classes.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Какой случай исправляется?
Нулевая label умножается на log-probability -inf.
Можно ли все logits сделать -inf?
Нет, такая строка не задаёт распределение.
Какие labels нужны?
Вероятностные распределения, например one-hot.
Официальные источники
Частые вопросы
Какой случай исправляется?
Нулевая label умножается на log-probability `-inf`.
Можно ли все logits сделать -inf?
Нет, такая строка не задаёт распределение.
Какие labels нужны?
Вероятностные распределения, например one-hot.
Читайте также
Как использовать 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).
Комментарии
Пока тихо. Скажите первое слово