[Перевод] Что спрятано в 2500 слоях: как по весам нейросети восстановили MD5

Во многих ML‑головоломках в духе CTF участникам дают нейросеть — «чёрный ящик» и предлагают выяснить, что она делает. Когда в начале прошлого года мы задумались о собственной головоломке по машинному обучению, нам захотелось немного изменить привычный формат.

Мы решили, что будет интересно предоставить участникам полное описание нейросети, включая все её веса. Чтобы восстановить принцип работы сети, им пришлось бы воспользоваться методами механистической интерпретируемости. С похожей задачей мы порой сталкиваемся и в собственных исследованиях, когда пытаемся интерпретировать внутренние признаки сложных моделей.

Мы опубликовали головоломку в феврале прошлого года. На тот момент мы даже не были уверены, что её вообще можно решить. Созданная нами нейросеть выдавала 0 почти для любых входных данных. Логично было предположить, что задача состоит в том, чтобы подобрать вход, при котором сеть вернёт 1 или любое другое ненулевое значение.

Но, как вы скоро увидите, мы специально устроили сеть так, чтобы до ответа нельзя было добраться традиционными методами полного перебора — например, выполнив обратное распространение от ненулевого выходного значения до входного слоя. Нужно было действительно понять, что делает сеть.

Реакция на головоломку нас поразила. Во многом благодаря удаче нам удалось почти идеально подобрать уровень сложности: задача оказалась не настолько трудной, чтобы её никто не решил, но и не настолько простой, чтобы нас завалили ответами. Более того, если вам по силам решить эту головоломку, велика вероятность, что вы хорошо впишетесь в команду Jane Street.

Ниже мы ещё раз сформулируем условие, но предупреждаем: дальше будет огромное количество спойлеров. Если хотите сначала решить головоломку самостоятельно, лучше пока не читайте. В оставшейся части статьи мы шаг за шагом разберём путь одного из участников — со всеми тупиками и неожиданными поворотами, через которые он прошёл, прежде чем наконец нашёл решение.

Условие

Сегодня я отправился в поход и нашёл под неолитическим курганом груду тензоров! Я отнёс их местному нейросетевому сантехнику, и ему удалось собрать на скорую руку вот это.

model.pt

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

Входные данные модели

vegetable dog

Выходное значение модели

0

Файл model.pt, по сути, представляет собой модель PyTorch в формате pickle.

Решение

С чего начать

Студент выпускного курса по имени Алекс сидел у себя в общежитии, когда сосед по комнате рассказал ему о головоломке, которую активно обсуждали в Twitter. Сосед и сам пытался её решить, но сдался после двух бессонных ночей. Алекс проводил в университете свою последнюю зиму, искал, чем бы заняться, и решил взглянуть на задачу.

Он скачал модель и начал разбираться, уделив особое внимание последнему слою:

import torch
import plotly.express as px
model = torch.load('./model.pt')
linears = [x for x in model if isinstance(x, torch.nn.Linear)]
px.imshow(linears[-1].weight.detach())

Сразу стало ясно, что это не обычная нейросеть. Её явно не обучали: все веса были целыми числами. Скорее всего, сеть спроектировали вручную для выполнения какого‑то вполне определённого вычисления.

Последний слой представлял собой матрицу размером 48 × 1, явно разделённую на три секции. Если посмотреть на значения активаций предыдущего слоя, там действительно всегда трижды повторялось одно и то же.

Предпоследний слой, судя по всему, содержал три копии одних и тех же весов, а в его смещении повторялись одни и те же 16 байт, каждый раз увеличенные на единицу. Как будто там были закодированы вектор v, затем v + 1 и v + 2.

Вот как выглядели веса предпоследнего слоя:

px.imshow(linears[-1].weight.detach())

А вот так выглядели смещения:

px.imshow(linears[-2].bias.detach().unsqueeze(0))

Немного поразмыслив и вспомнив, что последний слой выдаёт один бит, Алекс понял: предпоследний слой ReLU, вероятно, проверяет равенство двух 16-байтовых целых чисел, причём каждому байту соответствует отдельный нейрон.

Судя по всему, это работало так: слой создавал три копии входного вектора v, представляющего 16-байтовое число, и сравнивал их с эталонным числом x, заданным смещением предпоследнего слоя. В результате три копии представляли выражения:

Последний слой применял к ним веса 1, -2 и 1 соответственно.

Рассмотрим отдельное значение и разберём возможные случаи.

Вычисляется выражение:

Если v = x, результат равен 1. Остальные случаи мы здесь расписывать не будем, но во всех них получается 0.

Смещение последнего слоя равнялось -15, поэтому итоговый нейрон срабатывал только в том случае, когда для всех 16 байт.

Теперь вопрос звучал так: как добиться, чтобы значения активаций предпоследнего слоя совпали с x?

Восстанавливаем программу в ядре сети

Алекс рассудил: если в самом конце сеть сравнивает результат с некоторым числом, значит, вся её остальная часть должна вычислять какое‑то большое выражение. И структура у сети действительно явно просматривалась — это видно даже по графику размеров 2500 линейных слоёв, составляющих примерно половину всей сети:

px.line([l.out_features for l in linears])

Алекс начал исследовать разные фрагменты сети, прослеживая зависимости между ними. Для этого ему пришлось подолгу всматриваться в графовые структуры:

Но после нескольких часов поисков Алекс так и не нашёл достаточно понятных вычислительных схем. Сложность сети пока явно не позволяла проследить всё вручную.

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

Разумеется, из‑за множества слоёв ReLU напрямую это сделать нельзя: ReLU нелинейна. Но её можно смоделировать, добавив дополнительную целочисленную переменную, которая соответствует утверждению «это значение активации отрицательно».

Таким образом сеть можно представить как задачу целочисленного линейного программирования и передать решателю ограничений, поддерживающему целочисленное программирование.

Алекс так и поступил: добросовестно написал код, преобразующий слои нейросети в огромную задачу линейного программирования, и запустил решатель.

И оставил его работать.

Дело явно не двигалось, поэтому Алекс решил сократить число переменных в задаче. Возможно, сеть можно было как‑то упростить?

Он заметил, что многие слои в основном похожи на единичные матрицы. Более того, примерно в 1500 слоях 80% узлов просто выполняли тождественное преобразование.

Алекс представил каждый нейрон сети как узел в ориентированном ациклическом графе, где каждый узел с определёнными весами связан с узлами следующего слоя.

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

Были и более хитрые способы упрощения:

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

• если у двух нейронов одного слоя полностью совпадают входные векторы, их можно объединить, а все связи с узлами‑потомками перенаправить на новый общий нейрон;

• этот процесс можно повторять много раз.

К этому моменту Алекс уже потратил на анализ немало часов. Он обнаружил вычислительные схемы, которые, судя по всему, повторялись во множестве слоёв.

Алекс выводил разные классы эквивалентности узлов и сравнивал последовательности входных весов для каждого из них. Оказалось, что разновидностей узлов было совсем немного.

Например, существовал класс узлов, которые фактически просто передавали значение из слоя двумя уровнями выше.

Сворачивание таких узлов и другие похожие упрощения сократили задачу линейного программирования примерно с двух миллионов узлов до 75 тысяч.

Но даже после этого Алекс снова запустил решатель, и тот продолжал вычисления, так и не завершаясь.

Последние упрощения

А что, если распространить границы значений по всей сети?

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

Оказалось, что даже при довольно консервативных оценках у многих узлов получаются очень узкие границы, например от 0 до 1. Возможно, этого хватит, чтобы сделать задачу практически решаемой?

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

В SAT‑представлении для каждой пары «узел — возможное значение из его диапазона» создавалась отдельная булева переменная. После всех упрощений получилось около 200 тысяч переменных. За сутки работы SAT‑решатель сократил задачу до 20 тысяч переменных.

Дальше процесс, похоже, не продвигался.

По сути, Алекс обнаружил внутри нейросети вычислительное ядро, которое уже не поддавалось дальнейшему упрощению и, к его большому разочарованию, всё ещё было слишком велико для полного перебора.

Через несколько дней ему пришлось отступить и пересмотреть подход: фактически он так никуда и не продвинулся.

В шаге от решения

Алекс решил взглянуть на задачу с другой стороны. Головоломка ведь должна иметь решение, верно? Как человек мог бы построить такую задачу, чтобы её было интересно решать?

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

Необратимая функция. Какие необратимые функции первыми приходят на ум?

Алекс попросил ChatGPT перечислить распространённые хеш‑функции и сопоставил их с простыми графиками ширины слоёв, у которых обнаружилась периодическая структура.

На графиках было 32 одинаковых периода длиной 48.

Возможно, сеть 32 раза выполняла один и тот же вычислительный блок?

Алекс снова обратился к ChatGPT: есть ли распространённые хеш‑функции, использующие 32 вычислительных блока?

Вот оно.

Оказалось, что примерно так устроены почти все из них.

Чтобы понять, какая именно хеш‑функция скрывалась в сети, Алекс начал экспериментировать вручную. Он подавал на вход сети какую‑нибудь строку, отдельно вычислял для неё разные хеши, а затем смотрел на предпоследний слой.

MD5 совпал, а остальные распространённые хеш‑функции — нет.

Это была хорошая новость: по смещениям предпоследнего слоя Алекс уже знал, каким должен быть искомый хеш. Теперь задача сводилась к поиску входной строки с заданным MD5-хешем. Но как это сделать, было совершенно непонятно, особенно учитывая, что у него не было строгого доказательства, что сеть всегда вычисляет именно MD5.

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

Сбой в Матрице

Алекс заметил в сети одну странность. Похоже, в ней была ошибка: если длина входных данных превышала 32, сеть переставала вычислять правильный MD5-хеш. Возможно, именно эта ошибка позволяла обратить зашитое в сеть значение хеша?

Следующие два дня он потратил на обратную разработку ошибки.

Для начала Алекс попросил Gemini написать реализацию хеш‑функции MD5. Затем сопоставил каждый нейрон сети с соответствующей переменной алгоритма MD5. Он написал код, который сохранял последовательность значений заданной промежуточной переменной, а затем искал эти значения в каждом из 32 блоков сети.

Так удалось определить, какие диапазоны нейронов соответствуют битам отдельных переменных. Оказалось, что одни диапазоны битов точно соответствовали переменным алгоритма, а другие хранили промежуточные результаты вычислений.

После этого Алекс мог подавать входные данные длиной больше 32 и кропотливо прослеживать вычисления по всем блокам, пока не находил точное место, где сеть начинала расходиться с правильным алгоритмом.

Где была ошибка

Проблема скрывалась в первых семи слоях.

Там находилась вычислительная схема, которая определяла длину входных данных и пыталась сохранить её в четырёх байтах в порядке от младшего байта к старшему (little‑endian). Но если длина составляла 256 бит или больше, переменная длины содержала само значение 256 вместо его правильного представления.

Например, при длине больше 384 бит байты длины должны были выглядеть как:

128 1 0 0

но сеть вместо этого кодировала их как:

384 0 0 0.

Теперь предстояло понять, можно ли воспользоваться этой ошибкой, специально сформировав сообщение длиной не менее 256 бит.

Дальнейший кропотливый анализ позволил сделать несколько наблюдений.

Во‑первых, возможных значений длины было не так много. Входов было всего 55, поэтому Алекс мог выполнить полный перебор и посмотреть, как сеть ведёт себя с этими странными значениями.

Во‑вторых, некорректное значение длины преобразовывалось в двоичное представление, а затем проходило через все слои сети. В этом представлении все биты принимали значение 1, а оставшаяся часть числа оказывалась сосредоточена в младшем бите. Поэтому 384 кодировалось как:

130, 1, 1, 1, 1, 1, 1, 1.

В‑третьих, некорректные байты длины сообщения использовались лишь в нескольких блоках вычисления MD5, который всегда считывал байты входных данных в одном и том же порядке.

Опираясь на эти наблюдения, можно было записать модифицированную версию алгоритма MD5, которая в нужных блоках корректировала свои вычисления, чтобы они совпадали с поведением нейросети.

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

Ответ его удивил: ошибка была непреднамеренной.

Зная это, почему бы не попробовать решить головоломку ещё раз — в последний раз?

Возвращение полного перебора

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

Автор намеренно подобрал хеш, который легко найти полным перебором, и оставил в описании задачи и Python‑коде несколько небольших подсказок:

ответ состоял из двух английских слов в нижнем регистре, разделённых пробелом.

На самом деле Алекс уже пытался подобрать хеш раньше, но использовал список из 10 тысяч самых распространённых слов. Его оказалось недостаточно.

Когда Алекс нашёл список побольше, он получил ответ.

Ещё одна головоломка

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

MD5 показался нам удачным компромиссом, хотя простой эту задачу точно не назовёшь.

Поскольку MD5 использует сложение по модулю, при создании головоломки пришлось реализовать параллельный сумматор с переносом примерно в 20 слоях нейросети.

Задача не из лёгких!

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

Опыт создания и публикации головоломки, а также общения с теми, кто её решил, оказался настолько удачным, что мы решили повторить эксперимент.

Здесь вы найдёте нашу новую задачу. В ней слои нейросети перемешаны, и их нужно вернуть в правильный порядок… Справитесь?

Когда хочется не просто получить ответ от модели, а понять, почему она ведёт себя именно так, быстро упираешься в устройство алгоритмов, интерпретируемость и диагностику сложных систем.

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

• 8 сентября, 20:00. «Что надо знать про работу LLM моделей». Записаться

• 23 сентября, 18:00. «Дерево решений — простой и интерпретируемый ML‑алгоритм». Записаться

• 8 сентября, 20:00. «AI против бага: как разобрать инцидент в Python‑проекте от логов до исправления». Записаться

Полный список бесплатных уроков августа смотрите в дайджесте.