Хабр·Russia·Monday, October 12, 2026 at 09:52 PM UTC

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

Translate

MaxAkaAltmer12 минут назад

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

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

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

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

Но есть важнейший нюанс: выдавать большой 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)!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

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

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

View the original on Хабр →

KioskNews shows a cleaned-up reading view extracted from the publisher’s page — the original always lives on their site, not ours.