Был поздний вечер. Ужасно болела голова, простуда не давала нормально дышать, а градусник упрямо полз вверх. Что можно делать в таком состоянии, когда кажется, что сама Смерть уже натирает лыжи где‑то рядом?
Казалось бы, ничего значимого. Поэтому я просто решил разобрать груду скопившихся на столе бумаг — надо же иногда наводить порядок. И тут взгляд зацепился за старую тетрадь торчавшую из‑за шкафа.
Я аккуратно развернул ее, и взгляд зацепился за старый рисунок, который был сделан немногим позже ухода NVIDIA из России, где я работал ведущим инженером и в основном занимался написанием фреймворка для нейросетей, но уезжать не стал. А рисунок набросал, размышляя о том, как бы заполучить заветные ядра свертки Винограда без китайских теорем, универсально, для абсолютно любой формы ядра, включая те, которые не решаются аналитически.
Идея была простой: почему бы матрицы Винограда не получить методом градиентного спуска?
Я быстро озадачил свой локальный ИИ вопросом: а не сделал ли кто‑то за эти годы то же самое? Успокаивающее шуршание дросселей на старых V100, говорило о том что процесс пошел, а помаргивание коммутатора подтверждало что нейронки активно занимаются поиском ответа...
Выскочила пара ссылок на AlphaTensor от DeepMind и ещё несколько академических работ, довольно слабо относящихся к реальной «серебряной пуле». Но главное — ни у кого из них я не увидел двух ключевых вещей: открытого кода и универсальности подхода для построения точных алгоритмов под ЛЮБЫЕ свертки (с произвольным Stride, Dilation и асимметрией), но с минимальным числом умножений, как это делал тот крутой мужик, которого звали Шмуэль Виноград.
Гештальт, висевший в воздухе годами, потребовал срочной развязки....
Скажу честно: я пишу в основном на C++, CUDA и HDL (VHDL/SystemVerilog). Всё остальное без серьезной надобности не употребляю. Но сегодня спуститься до уровня Python, думаю, было вполне простительно — держать стандартную дисциплину разработки я был попросту не в состоянии.
Поэтому я просто попросил ИИ набросать мне каркас нейросети прямо по моему рисунку из тетради. Модель справилась легко — архитектура‑то концептуально очень простая.
Два полносвязных слоя (fc_D и fc_F) кодируют входные данные D и ядро F в промежуточное пространственное «горлышко» (inter_dim). В горлышке происходит поэлементное умножение. Третий полносвязный слой (fc_C) декодирует результат обратно в форму выхода.
class WinogradFromSketch(nn.Module): def __init__(self, in_shape, out_shape, kernel_shape, inter_dim): super().__init__() self.fc_D = nn.Linear(math.prod(in_shape), inter_dim, bias=False) self.fc_F = nn.Linear(math.prod(kernel_shape), inter_dim, bias=False) self.fc_C = nn.Linear(inter_dim, math.prod(out_shape), bias=False) def forward(self, D, F_tensor): out_D = self.fc_D(D.view(B, -1)) out_F = self.fc_F(F_tensor.view(B, -1)) return self.fc_C(out_D * out_F).view(B, 1, *out_shape)
Дальше я выстроил простейший цикл обучения, где источником «истины» и входных данных служила сама эталонная свертка (F.conv1d, F.conv2d или F.conv3d), которую мы хотим сжать. Генерируем случайные нормальные батчи D и F, считаем от них истинную свертку — вот вам и идеальный бесконечный датасет с нулевым шумом.
Чтобы заставить сеть обнулять лишние связи в матрицах (ведь мы ищем разреженность, как у Винограда), я повесил на веса обычный L1-штраф и небольшой прунинг по теоретически минимальной дроби:
Натравил Adam, настроил простейшую 1D‑свертку, нажал Run и совершенно спокойно приготовился ждать провала операции...
И когда на экране выкристаллизовалась совершенно ясная, чистая матрица! Моя реакция была: «Да ладно?!», как когда‑то у моего напарника по лабораторке, когда я написал сотни строк кода на ассемблере, и они заработали строго с первого раза.
--- Matrix fc_D (data) (shape: (4, 4)) --- Zero percentage (< 0.08333333333333333): 37.5% [[ 0.00 0.76 0.98 0.00] [ 0.00 0.89 -0.70 0.00] [ 0.66 0.33 -0.66 0.00] [ 0.00 0.63 0.31 -0.62]] --- Matrix fc_F (kernel) (shape: (4, 3)) --- Zero percentage (< 0.08333333333333333): 33.3% [[ 0.48 0.62 0.80] [ 0.84 -0.65 0.51] [ 1.24 0.00 0.00] [ 0.00 0.00 -1.18]] --- Matrix fc_C (output) (shape: (2, 4)) --- Zero percentage (< 0.08333333333333333): 25.0% [[ 0.80 -1.06 1.22 0.00] [ 1.02 0.83 0.00 1.36]]
Итак, вечер переставал быть стремным. А кто‑то там в темноте, кажется, тихо снял лыжи, отложил косу и подсел поближе...
Конечно, сразу захотелось попробовать геометрии побольше. Алгоритм очевидно сходился, но делал это довольно неторопливо. Модели при обучении явно требовались хорошие подзатыльники. Откуда их взять? Если подумать, у нас же прямо в процессе генерируются данные, которые хуже всего вписываются в текущую математику — те самые случайно сгенерированные примеры, на которых модель выдает максимальную ошибку.
Ага! Будем копить эти «тяжелые» примеры в специальный буфер (HardExampleBuffer) и периодически отдельно прогонять сеть по ним. И это сработало: по ощущениям, скорость сходимости подскочила примерно на порядок!
class HardExampleBuffer: def __init__(self, capacity=1024 * 8): self.capacity = capacity self.D_buffer = None self.F_buffer = None def add_examples(self, D, F_tensor, errors, top_k=32): # Отбираем самые "тяжелые" примеры с максимальным MSE _, top_idx = torch.topk(errors, k=min(top_k, D.size(0))) # ... и дообучаем модель на этом буфере каждые пару шагов
Но в голове, хоть и очень тяжелой от температуры, высвободившийся ресурс нейронов — за счет Python и дозволения себе поговнокодить совместно с ИИ — четко предсказывал еще одну проблему: проблему чисел, больших единицы по модулю (
).
Первые две матрицы (
и 
) большие значения в алгоритмах Винограда не содержат — они живут в идеальном диапазоне ![[-1, 1]](https://habrastorage.org/getpro/habr/formulas/7/7d/7de/7dec1d46e68831c4eca28b020fcb1604.svg)
(состоя из 
). Поэтому основная проблема сидит в выходной матрице 
и на самом выходе из свертки как таковой. Что с этим делать? Да все просто — вносим множители! И градиентный спуск отдаст нам все, что может.
Но есть важнейший нюанс: выдавать большой scale для весов выходной матрицы прямо сразу — категорически нельзя...
Если задрать scale с первого же шага, ленивая нейросеть тут же найдет «грязное» непрерывное решение и выдаст отличный результат с точки зрения погрешности, но абсолютно бестолковый с точки зрения плотности (нулей мы там не увидим).
Поэтому мы должны сначала заставить модель поголодать и провалиться в каньон разреженных решений. Для этого мы держим веса на строгой диете (
) и плавно повышаем scale от 
до теоретического максимума в процессе обучения!
# S-образный сглаженный разгон scale в процессе обучения new_scale = get_current_scale(step, steps, float(winograd_scale)) if new_scale != model.scale: with torch.no_grad(): # Инвариантный пересчет весов, чтобы MSE не скакал при росте scale model.fc_C.weight.mul_(model.scale / new_scale) model.scale = new_scale
Решения начали сходиться и для больших весовых коэффициентов выходной матрицы. Казалось бы, всё, конечно, очень здорово и даже имеет реальную практическую пользу для приближенных вычислений. Но нам нужны знания — абсолютные точные формулы с рациональными числами, которые мы можем редуцировать до применимых при построении в вычислительных ядрах выражений!
И тут я вспомнил один фокус. Давно я как‑то применял метод наименьших квадратов (МНК / Least Squares) для сопоставления сигналов при их фрактальном сжатии и подумал: а почему бы не взять эти неидеальные матрицы и не попытаться сопоставить их с идеальной функцией свертки?
# --- Ядро аналитического сольвера (SymPy + МНК) --- def prove_winograd_exact(model, target_sym, monomials, max_denom=36): WD_float, WF_float = model.fc_D.weight.numpy(), model.fc_F.weight.numpy() # 1. Перебираем знаменатели q и квантуем WD и WF к дробям Q for q in range(1, max_denom + 1): WD_exact = snap_to_sympy_rationals(WD_float, limit_denom=q) WF_exact = snap_to_sympy_rationals(WF_float, limit_denom=q) # 2. Аналитически находим идеальную WC через МНК (Least Squares) M_num = build_monomial_matrix(WD_exact, WF_exact, monomials) WC_num, *_ = np.linalg.lstsq(M_num, Target_num.T, rcond=None) WC_exact = snap_to_sympy_rationals(WC_num.T, limit_denom=max_denom) # 3. Символьная верификация в кольце полиномов SymPy M_sym = sp.Matrix(build_symbolic_monomials(WD_exact, WF_exact)) residual = sp.expand(sp.Matrix(WC_exact) * M_sym.T - target_sym) # 4. Если остаток матриц строго равен нулю — алгоритм найден! if residual == sp.zeros(*residual.shape): print(f"[УСПЕХ] Точное тождество доказано при q = {q} (ошибка = 0)!") return WD_exact, WF_exact, WC_exact, True return None, None, None, False
После некоторого колдунства над этой задумкой — уже на заре, запускаем и опять: «Да ну нафиг!». 100% копия матриц Винограда!
--- Symbolic Matrix fc_D (data) (shape: 4x4) --- [ 0 1 1 0 ] [ 0 1 -1 0 ] [ 1 0 -1 0 ] [ 0 1 0 -1 ] Zero percentage: 50.0% --- Symbolic Matrix fc_F (kernel) (shape: 4x3) --- [ 1 1 1 ] [ 1 -1 1 ] [ 1 0 0 ] [ 0 0 -1 ] Zero percentage: 33.3% --- Symbolic Matrix fc_C (output) (shape: 2x4) --- [ 1/2 -1/2 1 0 ] [ 1/2 1/2 0 1 ]
Я завариваю бессчетную чашку чая и начинаю прогонять ядра для которых я никогда не видел и не слышал решения в форме алгоритма Винограда и это работает! Метод работает!
================================================================= Experiment Parameters (Winograd Algorithm Mining) ================================================================= Device : cuda Data Type (dtype) : torch.float32 Input Shape (in_shape) : (5, 5) (25 elems) Kernel Shape (kernel) : (3, 3) (9 elems) Output Shape (out_shape) : (2, 2) (4 elems) Stride & Dilation : stride=(2, 2), dilation=(1, 1) ------------------------------------------------------------- Naive Multiplications : 36 Winograd Multiplications : 25 FLOPs Reduction : +30.6% ------------------------------------------------------------- Total Steps : 100000 Batch Size : 2048 Initial LR / L1 Lambda : 0.001 / 0.00015 Pruning Threshold : 0.000868 Max Denominator (q_max) : 576 Target Scale (max scale) : 36 ================================================================= step 1/100000 | lr: 1.00e-03 | scale: 1.00 | mse: 8.838e+00 | l1 sum: 112.77 step 10000/100000 | lr: 9.76e-04 | scale: 1.00 | mse: 4.320e-06 | l1 sum: 171.90 step 20000/100000 | lr: 9.05e-04 | scale: 1.71 | mse: 1.602e-06 | l1 sum: 102.98 step 30000/100000 | lr: 7.94e-04 | scale: 7.04 | mse: 7.585e-06 | l1 sum: 66.34 step 40000/100000 | lr: 6.55e-04 | scale: 16.01 | mse: 3.799e-07 | l1 sum: 52.72 step 50000/100000 | lr: 5.01e-04 | scale: 25.77 | mse: 1.839e-07 | l1 sum: 45.42 step 60000/100000 | lr: 3.46e-04 | scale: 33.22 | mse: 7.071e-05 | l1 sum: 41.33 step 70000/100000 | lr: 2.07e-04 | scale: 36.00 | mse: 4.990e-07 | l1 sum: 38.97 step 80000/100000 | lr: 9.64e-05 | scale: 36.00 | mse: 1.407e-05 | l1 sum: 37.74 step 90000/100000 | lr: 2.54e-05 | scale: 36.00 | mse: 7.816e-08 | l1 sum: 36.77 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! [SUCCESS]: Early stopping triggered at step 97000/100000! Exact Winograd algorithm proved via SymPy (Residual = 0)! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Что ж, похоже, я снова сэкономил для планеты гигаватты энергии в год. Но на этот раз - без закрытых корпоративных библиотек и проприетарных движков.
Пока я оформлял исходники в открытый репозиторий под MIT лицензией, кто‑то в темноте клацал по клавишам и что‑то бубнил себе под нос. Наконец, я залил все в репозиторий: https://github.com/MaxAkaAltmer/apex_winograd и спокойно уснул...
