Два вызова JAX дают один шум: как разделить ключи и сохранить повторяемость

В JAX повторяемость удобно строить явно: случайное число зависит от переданного ключа. Это помогает повторить эксперимент и разнести вычисления по подзадачам. Но привычка из библиотек с автоматически меняющимся генератором может привести к ошибке: два вызова с одним ключом не создают две независимые выборки.

Ключ не делает шаг сам

В руководстве по случайным числам JAX ключ описан как неизменяемое значение. Вызов функции выборки не продвигает скрытый счётчик. Если снова передать тот же ключ в ту же операцию с теми же параметрами, получится тот же результат. Это ожидаемое поведение, полезное для проверки.

Представьте учебный эксперимент: нужно добавить шум к двум копиям набора признаков. Если использовать один ключ дважды, маски случайности могут совпасть вместо ожидаемого независимого возмущения. Качество модели этим примером не измеряется; он только помогает проверить, правильно ли организован поток случайных данных.

from jax import random

root = random.key(42)
a = random.normal(root, shape=(4,))
b = random.normal(root, shape=(4,))
# a и b совпадают: ключ не изменился

left, right = random.split(root, 2)
a = random.normal(left, shape=(4,))
b = random.normal(right, shape=(4,))

Назначьте ключ каждой подзадаче

split детерминированно выводит новые ключи из родительского. В примере каждая копия данных получает отдельную ветку. Для повторения всего опыта нужно восстановить исходный seed и тот же порядок назначения веток, а не случайно переиспользовать один дочерний ключ во всех местах.

Проверяйте распределение ключей по смыслу задачи: отдельная ветка для инициализации, другая для шума, ещё одна для следующего независимого опыта. Запишите это соответствие рядом с параметрами эксперимента. Тогда добавление вспомогательного случайного действия проще заметить при сравнении двух запусков.

В справочнике jax.random рекомендуются типизированные ключи random.key; старый PRNGKey тоже поддерживается, но имеет другой формат. Там же описаны альтернативные реализации генератора. Не переносите их особенности в эксперимент незаметно: выбранная реализация — часть условий сравнения.

Что считать успешной проверкой

Начните с малого массива без обучения большой модели. Два одинаковых вызова с одним ключом должны совпасть; повторное создание корня и тех же веток должно воспроизвести соответствующие выборки. Две разные ветки предназначены для независимых выборок, но не обещают отсутствия любого случайного совпадения элементов.

Сохраните seed, версии библиотек, выбранный генератор, формы массивов и способ назначения ключей. При изменении размера выборки не требуйте автоматически той же последовательности, что получалась при поэлементных вызовах: JAX не обещает такую последовательную эквивалентность.

Практическое правило: повторно используйте ключ, когда намеренно хотите повторить опыт; для независимой случайности выделяйте отдельные ветки. Так повторяемость остаётся управляемым свойством эксперимента, а не случайной одинаковостью шума.