Гайд · TNWS AI
Как посчитать gradient для каждого примера через jax.vmap и grad
Per-example gradients без Python-loop в JAX 0.11.2: vmap(grad), PyTree batch-axis и проверка среднего против batch loss.
Задача и применимость
Актуальный запрос: «как получить per example gradients JAX vmap grad». Получить отдельный gradient параметров для каждой строки batch, а не только gradient усреднённой loss. Композиция jax.vmap(jax.grad(one_loss), in_axes=(None,0,0)) оставляет параметры общими и векторизует пример и target по ведущей оси. Каждый gradient leaf получает дополнительную batch-ось.
Материал применим к JAX 0.11.2 и jaxlib 0.11.2. JAX предоставляет массивы, автоматическое дифференцирование и программные преобразования; он не является готовой библиотекой слоёв, датасетов или полного тренировочного цикла. Пример намеренно мал и выполняется локально на доступном backend без загрузки весов и внешних данных. Численный критерий доказывает поведение конкретного кода, но не качество рабочей нейросети.
Пошаговое решение
- Опишите scalar loss одного примера.
- Примените grad по параметрам первого аргумента.
- Оберните функцию в vmap с общими params и batch-осями данных.
- Усредните per-example gradients и сравните с gradient mean batch loss.
Готовый проверяемый шаблон
import jax
import jax.numpy as jnp
params = {"w": jnp.array([0.5, -1.]), "b": jnp.array(0.2)}
X = jnp.array([[1., 0.], [0., 1.], [1., 1.]])
y = jnp.array([1., -1., 0.])
def one_loss(p, x, target):
return (x @ p["w"] + p["b"] - target) ** 2
per_example = jax.vmap(jax.grad(one_loss), in_axes=(None, 0, 0))(params, X, y)
batch_grad = jax.grad(lambda p: jnp.mean(
jax.vmap(one_loss, in_axes=(None, 0, 0))(p, X, y)
))(params)
for per_leaf, batch_leaf in zip(jax.tree.leaves(per_example), jax.tree.leaves(batch_grad)):
assert jnp.allclose(jnp.mean(per_leaf, axis=0), batch_leaf)
assert per_example["w"].shape == (3, 2)
assert per_example["b"].shape == (3,)
print(jax.tree.map(lambda z: z.shape, per_example))
Вход и ожидаемый результат
Batch из трёх строк дал gradient weights формы (3,2) и bias формы (3,). Среднее каждого per-example leaf совпало с gradient loss, которая усредняет те же три ошибки. Это также проверяет выбранную ось агрегации.
Как проверить результат
Проверьте leaf shapes, порядок примеров и равенство среднего с batch gradient. Для loss с sum вместо mean ожидаемый коэффициент другой. Перед использованием для differential privacy или influence analysis отдельно проверьте математический контракт и память.
Перенос на настоящую нейросеть
Для рабочей модели сначала зафиксируйте контракт параметров и batch: pytree structure, формы, dtype, порядок признаков и семантику выходов. Генерацию случайностей оформите как явный поток keys. Разделите чистые функции forward, loss и update, чтобы их можно было независимо тестировать, дифференцировать и компилировать. Данные train, validation и test не должны смешиваться. Любое изменение формы, precision, backend или версии JAX проверяйте на сохранённом обезличенном regression-наборе до развёртывания.
Критерии приёмки
- Все формы и dtype совпадают с контрактом.
- Каждый PRNG subkey используется в одной логической случайной операции.
- Loss и gradients конечны на положительном тесте.
- Негативный вход вызывает ожидаемую ошибку, а не случайный output.
- Обычный и преобразованный через jit/vmap путь численно согласованы, когда это применимо.
- Устройство, версии JAX/jaxlib и численный допуск записаны в отчёте.
Регрессионный протокол
Храните несколько малых входов: стандартный batch, один объект, нули, граничные значения и заведомо невалидную форму. Сравнивайте outputs и loss с эталоном через allclose с явно выбранным допуском. После смены backend проверяйте допустимую численную погрешность, а не требуйте безусловной битовой идентичности. Производительность измеряйте после компиляционного прогрева и завершения асинхронной работы. Не выбирайте лучший единичный замер: сообщайте распределение задержки и условия машины.
Диагностика расхождений
Если результат изменился, локализуйте слой: данные и preprocessing, key-поток, shapes/dtype, forward, loss, gradients, update, затем jit и backend. Сначала выполните функцию без jit на минимальном входе, затем верните transformations по одной. Сравните tree structure и каждый leaf. Не увеличивайте допуск и не применяйте nan_to_num, пока не выяснено, не скрывает ли это неверную форму, переполнение или повтор ключа.
Контроль после внедрения
В эксплуатации ограничьте допустимые batch shapes и объём входа, чтобы запрос не исчерпал память и не создал бесконечный поток перекомпиляций. Собирайте агрегированные метрики latency, compile count, долю невалидных входов и численные ошибки без записи исходных чувствительных данных. Сохраняйте предыдущий проверенный артефакт и код параметров для отката. Автоматическая замена весов требует той же независимой проверки, что и обычный выпуск.
Границы применимости
Пример на CPU не доказывает свойства GPU или TPU, а маленький массив не предсказывает производительность большой модели. JIT может ускорить повторяющиеся вычисления, но стоимость компиляции, передача данных и новые формы способны ухудшить end-to-end время. Autodiff возвращает производную написанной функции, а не подтверждает правильность математической постановки. Для критичных решений нужны предметные метрики, мониторинг и человеческий контроль.
Типичные ошибки и чего не делать
- Не ставьте batch-axis на общие параметры вместо None.
- Не путайте среднее per-example gradients с gradient суммы.
- Не храните огромный tensor всех градиентов без оценки памяти.
Не публикуйте только удачный вывод без негативного теста и версий среды. Не используйте тестовую выборку для настройки модели. Не загружайте непроверенные параметры и не передавайте в диагностику персональные данные. Для важных решений предусмотрите безопасный отказ.
FAQ
Где появляется batch-ось?
В начале каждого gradient leaf.
Почему params получают in_axes=None?
Один набор параметров общий для всех примеров.
Совпадает ли среднее с batch grad?
Да, если batch loss является средним тех же scalar losses.
Официальные источники
Частые вопросы
Где появляется batch-ось?
В начале каждого gradient leaf.
Почему params получают in_axes=None?
Один набор параметров общий для всех примеров.
Совпадает ли среднее с batch grad?
Да, если batch loss является средним тех же scalar losses.
Читайте также
Как инициализировать двухслойную нейросеть в JAX
Двухслойный MLP без фреймворка: random.split, масштаб весов, ReLU, PyTree параметров, одиночный и пакетный входы.
Как найти NaN в JAX-нейросети через debug_nans
Локальная диагностика NaN в eager и jit: контекст jax.debug_nans, FloatingPointError, безопасная функция и запрет маскировки проблемы nan_to_num.
Как обновить веса нейросети через value_and_grad в JAX
Один проверяемый шаг gradient descent для PyTree параметров: MSE, value_and_grad, tree.map, формы градиентов и снижение loss.
Комментарии
Пока тихо. Скажите первое слово