Потерянная тетрадь — разработка длинною в ночь

Был поздний вечер. Ужасно болела голова, простуда не давала нормально дышать, а градусник упрямо полз вверх. Что можно делать в таком состоянии, когда кажется, что сама Смерть уже натирает лыжи где-то рядом?

Казалось бы, ничего значимого. Поэтому я просто решил разобрать груду скопившихся на столе бумаг — надо же иногда наводить порядок. И тут взгляд зацепился за старую тетрадь торчавшую из-за шкафа.

Я аккуратно развернул ее, и взгляд зацепился за старый рисунок, который был сделан немногим позже ухода 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 и дозволения себе поговнокодить совместно с ИИ — четко предсказывал еще одну проблему: проблему чисел, больших единицы по модулю ().

Первые две матрицы ( и ) большие значения в алгоритмах Винограда не содержат — они живут в идеальном диапазоне (состоя из ). Поэтому основная проблема сидит в выходной матрице и на самом выходе из свертки как таковой. Что с этим делать? Да все просто — вносим множители! И градиентный спуск отдаст нам все, что может.

Но есть важнейший нюанс: выдавать большой 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.77step   10000/100000 | lr: 9.76e-04 | scale:  1.00 | mse: 4.320e-06 | l1 sum: 171.90step   20000/100000 | lr: 9.05e-04 | scale:  1.71 | mse: 1.602e-06 | l1 sum: 102.98step   30000/100000 | lr: 7.94e-04 | scale:  7.04 | mse: 7.585e-06 | l1 sum: 66.34step   40000/100000 | lr: 6.55e-04 | scale: 16.01 | mse: 3.799e-07 | l1 sum: 52.72step   50000/100000 | lr: 5.01e-04 | scale: 25.77 | mse: 1.839e-07 | l1 sum: 45.42step   60000/100000 | lr: 3.46e-04 | scale: 33.22 | mse: 7.071e-05 | l1 sum: 41.33step   70000/100000 | lr: 2.07e-04 | scale: 36.00 | mse: 4.990e-07 | l1 sum: 38.97step   80000/100000 | lr: 9.64e-05 | scale: 36.00 | mse: 1.407e-05 | l1 sum: 37.74step   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)!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

Что ж, похоже, я сэкономил для планеты гигаватты энергии в год. Снова — как когда-то с обычными и слитыми свертками, которые так и остались в недрах проприетарного кода NVIDIA. Но не сегодня. Сегодня всё будет иначе.

Пока я оформлял исходники в открытый репозиторий под MIT лицензией, кто-то в темноте клацал по клавишам и что-то бубнил себе под нос. Наконец, я залил все в репозиторий: https://github.com/MaxAkaAltmer/apex_winograd и спокойно уснул…