Поиск алгоритмов Винограда через градиентный спуск: ночная разработка

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

В основе лежит автоматический поиск матриц Винограда с помощью нейросетей, объединяющий глубокое обучение и символьную математику для поиска оптимальных вычислительных путей. Идея зародилась еще во времена ухода NVIDIA из России, где автор работал ведущим инженером и занимался написанием фреймворка для нейросетей, размышляя о том, как получить заветные ядра свертки без китайских теорем — универсально, для абсолютно любой формы ядра, включая нерешаемые аналитически.

Идея метода и архитектура нейросети

Поисковый запрос через локальный ИИ выдал пару ссылок на AlphaTensor от DeepMind и несколько академических работ, однако открытого кода и универсального подхода для построения точных алгоритмов под ЛЮБЫЕ свертки (с произвольным Stride, Dilation и асимметрией) с минимальным числом умножений, как у Шмуэля Винограда, обнаружено не было. Было принято решение задействовать Python и обучить простую нейросетевую модель, используя для каркаса концептуально простую архитектуру.

Архитектура включает три полносвязных слоя. Входные данные и ядро кодируются в промежуточное пространственное горлышко через слои fc_D и fc_F, где выполняется поэлементное умножение, после чего третий слой fc_C декодирует результат обратно в форму выхода.

Обучение модели и достижение разреженности

В качестве источника обучающих данных и «истины» выступала эталонная свертка (F.conv1d, F.conv2d или F.conv3d), генерирующая случайные нормальные батчи D и F, создавая идеальный бесконечный датасет с нулевым шумом. Для получения разреженных матриц, характерных для алгоритмов Винограда, на веса накладывался обычный L1-штраф совместно с небольшой процедурой прунинга по теоретически минимальной дроби.

Первый запуск с оптимизатором Adam позволил выявить четкую и чистую структуру матриц с нужным процентом нулевых элементов. Для ускорения сходимости модели (примерно на порядок) был добавлен специальный буфер сложных примеров HardExampleBuffer, куда сохранялись итерации с максимальной ошибкой MSE, после чего модель периодически дообучалась на этом буфере.

Работа с большими значениями и масштабирование

Для преодоления проблемы чисел, больших единицы по модулю (так как первые две матрицы живут в идеальном диапазоне, а основные трудности приходятся на выходную матрицу и выход свертки), применяются специальные масштабные коэффициенты (scale) и множители. Резкое повышение scale с первого же шага приводит к тому, что ленивая нейросеть находит «грязное» непрерывное решение с отличной погрешностью, но абсолютно бестолковое с точки зрения плотности и нулей.

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

Символьная верификация и точные формулы

Для получения абсолютных точных формул с рациональными числами, применимых в вычислительных ядрах, неидеальные матрицы были сопоставлены с идеальной функцией свертки с помощью метода наименьших квадратов (МНК) и библиотеки SymPy. Перебор знаменателей $q$, квантование весов к дробям $Q$, аналитический поиск идеальной матричной составляющей и проверка остатков в кольце полиномов позволили доказать точные тождества.

Финальная проверка подтвердила совпадение с классическими матрицами Винограда, а также работоспособность метода для ядер произвольной конфигурации (например, для входной формы (5, 5), ядра (3, 3) и шага stride=(2, 2) с сокращением FLOPs на 30.6%). Разработанный инструмент опубликован в открытом репозитории под лицензией MIT.