Grokking

Grokking — это отсроченное обобщение: нейронная сеть быстро достигает идеальной точности на обучающих данных, долго остается на случайном уровне на валидационных данных, а затем, уже после того как кажется, что модель переобучилась, внезапно начинает обобщать. Power и др. (2022) из OpenAI назвали эффект по результатам обучения небольших трансформеров алгоритмическим задачам, например модульной арифметике: точность на валидации росла от случайного уровня до идеального уже после того, как точность на обучении достигла максимума. Термин происходит из романа Роберта Хайнлайна «Чужак в чужой стране», где to grok означает полностью понять что-либо.

Как развивается grokking#

Классическая задача для изучения grokking — сложение по модулю: сеть получает пары (a, b) и должна выдать (a + b) mod p, причем для обучения доступна лишь часть всех пар. Приведенный ниже запуск обучает небольшой многослойный перцептрон с оптимизатором AdamW из PyTorch на 30% из 9 409 пар при p = 97; на CPU обучение занимает меньше минуты:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
p = 97  # task: predict (a + b) mod 97 from one-hot encodings of a and b
a, b = torch.cartesian_prod(torch.arange(p), torch.arange(p)).T
x, y = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float(), (a + b) % p
idx = torch.randperm(p * p)
train, val = idx[: p * p * 3 // 10], idx[p * p * 3 // 10 :]  # train on 30% of all pairs

model = torch.nn.Sequential(torch.nn.Linear(2 * p, 256), torch.nn.ReLU(), torch.nn.Linear(256, p))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)  # strong weight decay
for step in range(1, 20001):
    loss = F.cross_entropy(model(x[train]), y[train])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 4000 == 0:
        train_acc, val_acc = [(model(x[i]).argmax(1) == y[i]).float().mean().item() for i in (train, val)]
        print(f"step {step:5d}: train accuracy {train_acc:.2f}, validation accuracy {val_acc:.2f}")

На шаге 4 000 точность на обучении уже равна 1,00, тогда как на валидации она составляет 0,00. Затем точность на валидации повышается до 0,16 на шаге 8 000, 0,54 на шаге 12 000, 0,84 на шаге 16 000 и 0,97 на шаге 20 000. С weight_decay=0.0 точность на валидации в том же запуске остается равной 0,00 все 20 000 шагов.

Nanda и др. (2023) объяснили плато, выполнив обратный анализ однослойного трансформера, обученного сложению по модулю. Сеть учится преобразовывать сложение во вращение по кругу, используя дискретные преобразования Фурье и тригонометрические тождества. Обучение проходит три непрерывные фазы: запоминание, формирование схемы и очистка. Обобщающая схема постепенно усиливается, пока точность на валидации остается неизменной; скачок происходит, когда затухание весов устраняет заученное решение, которое по-прежнему доминировало в выходных данных.

Что вызывает grokking#

• Регуляризация: В исходных экспериментах наиболее эффективным приемом было затухание весов; без него приведенный выше запуск вообще не обобщает.
• Размер набора данных: Чем меньше доля обучающих данных, тем больше оптимизации требуется до начала обобщения; ниже определенного порога сеть не обобщает в пределах практичного бюджета.
• Норма весов: Omnigrok связывает grokking с величиной весов: у обобщающих решений норма весов меньше, чем у заученных. Авторы добились grokking в задачах с изображениями, текстом и молекулами, начав с больших начальных весов, а в алгоритмических задачах устранили его, ограничив норму весов.
• Численная стабильность: Prieto и др. (2025) обнаружили, что без регуляризации обучение с перекрестной энтропией продолжает увеличивать логиты, пока ошибки округления с плавающей точкой при вычислении softmax не остановят обучение. Это явление они называют коллапсом Softmax. Их функция активации StableMax и оптимизатор ⊥Grad позволяют получить grokking без затухания весов.

Grokking и связанные понятия#

• Переобучение: Разрыв между качеством на обучении и валидации, который сохраняется или увеличивается. При grokking разрыв закрывается сам по себе при продолжении обучения.
• Двойной спуск: Тестовая ошибка сначала падает, затем растет и снова падает по мере увеличения размера модели или длительности обучения. Двойной спуск по эпохам — еще одна форма отсроченного улучшения: во время обучения тестовая ошибка некоторое время растет, прежде чем снова снизиться.
• Механистическая интерпретируемость: Изучение алгоритмов внутри обученных сетей. Grokking на примере модульной арифметики стал одним из основных тематических примеров, поскольку изученный алгоритм удалось полностью восстановить.

Grokking на практике#

Grokking особенно заметен на небольших алгоритмических наборах данных, управляемых одним точным правилом. Он важен и для планирования экспериментов, поскольку ранняя остановка завершает запуск, когда метрика валидации перестает улучшаться, и может оборвать запуск с grokking во время плато. Обучение Ultralytics YOLO прекращается после patience эпох без улучшения показателя качества на валидации; по умолчанию это 100 эпох. При значении patience=0 обучение выполняется каждую эпоху.