It was late in the evening. I had a splitting headache, a nasty cold made it hard to breathe, and the thermometer was stubbornly creeping upward. What can you even do in a state like that, when it feels like the Grim Reaper is right around the corner, already waxing its skis?
Naturally, nothing productive. So I simply decided to sort through the pile of papers cluttered on my desk - you’ve got to clean up every now and then, right? And that’s when my eyes caught an old notebook sticking out from behind the cabinet.
I carefully opened it, and my eyes landed on an old sketch drawn shortly after NVIDIA exited Russia. Back then, I was working there as a Lead Engineer, primarily developing a deep learning framework, but I chose not to relocate.
I had scribbled that drawing while pondering a question: how could we derive those coveted Winograd convolution kernels without relying on the Chinese Remainder Theorem - universally, for absolutely any kernel geometry, including those impossible to solve analytically?
The idea was simple: why not derive Winograd matrices using gradient descent?
I quickly tasked my local AI with a question: had anyone done the exact same thing over all these years? The soothing coil whine of my old V100s signaled that the search was underway, while the blinking network switch confirmed the models were actively hunting for answers...
A couple of links popped up pointing to DeepMind’s AlphaTensor and a few academic papers, though none of them really felt like a practical "silver bullet." Most importantly, none of them provided two crucial things: open-source code and a universal approach for building exact algorithms for ANY convolution geometry (with arbitrary stride, dilation, and asymmetry) while minimizing multiplications - just like that legendary guy named Shmuel Winograd did.
This unfinished business, lingering in the air for years, demanded an immediate closure...
To be honest, I mostly write C++, CUDA, and HDL (VHDL/SystemVerilog). I don't touch anything else unless absolutely necessary. But dropping down to Python tonight felt entirely forgivable - in my feverish condition, maintaining standard software engineering discipline was out of the question.
So I simply asked the AI to sketch out a neural network scaffold directly from my notebook drawing. The model handled it effortlessly - conceptually, the architecture is dead simple.
Two fully connected layers (fc_D and fc_F) encode the input data D and the kernel F into an intermediate spatial "bottleneck" (inter_dim). Inside this bottleneck, an element-wise multiplication takes place. A third fully connected layer (fc_C) then decodes the result back into the required output shape.
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)
Next, I built a simple training loop where the reference convolution itself (F.conv1d, F.conv2d, or F.conv3d) served as the ground truth. We generate random, normally distributed batches of D and F and compute the true convolution on the fly - voilà, an ideal, infinite, zero-noise dataset.
To force the network to zero out redundant connections in the matrices (after all, we’re hunting for Winograd-style sparsity), I applied a standard L1 penalty to the weights along with a subtle pruning threshold based on the theoretical minimum fraction:
I unleashed Adam, configured a basic 1D convolution, hit Run, and calmly braced myself for the whole thing to fail...
But then, a crystal-clear, remarkably clean matrix materialized right on the screen! My reaction was a genuine "No way!" - much like my lab partner's face back in college when I wrote hundreds of lines of Assembly code and it actually worked on the very first try.
--- 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]]
So, the evening was turning out to be far less grim. And someone in the dark seemed to have quietly unstrapped their skis, set aside the scythe, and pulled up a chair...
Naturally, I immediately wanted to try larger geometries. The algorithm was clearly converging, but it was doing so quite sluggishly. The model evidently needed a kick in the pants during training. But where do you get it? Come to think of it, during the process we were generating data that fit the current math the worst - those exact randomly generated samples where the model yielded the highest error.
Aha! We’ll cache these "hard" examples in a dedicated buffer (HardExampleBuffer) and periodically run separate passes on them. And it worked: convergence speed felt like it shot up by an order of magnitude!
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): # Select the "hardest" samples with the highest MSE _, top_idx = torch.topk(errors, k=min(top_k, D.size(0))) # ... and fine-tune the model on this buffer every couple of steps
Yet in my fever-heavy head, the neural capacity freed up by switching to Python and allowing myself to hack together code alongside AI clearly predicted another hurdle: values with a magnitude greater than one (
).
The first two matrices (
и 
) rarely contain large numbers in Winograd algorithms — they comfortably reside in the ideal ![[-1, 1]](https://habrastorage.org/getpro/habr/formulas/7/7d/7de/7dec1d46e68831c4eca28b020fcb1604.svg)
interval (consisting of
). Therefore, the main bottleneck sits in the output matrix 
and the final convolution output itself. How to fix it? Simple - introduce scaling factors! And gradient descent will hand over everything it's got.
However, there’s a crucial catch: you absolutely cannot crank up the scale for the output matrix weights right out of the gate...
If you jack up the scale from step one, the lazy neural network will immediately settle into a "dirty" continuous solution that looks great on loss, but is completely useless in terms of sparsity - you won't see any zeros there.
So, we must first starve the model and force it to drop into the canyon of sparse solutions. To do this, we keep the weights on a strict diet (scale = 1.0) and smoothly ramp the scale from 1.0 to its theoretical maximum over the course of training!
# Smooth S-curve scale warmup during training new_scale = get_current_scale(step, steps, float(winograd_scale)) if new_scale != model.scale: with torch.no_grad(): # Invariant weight rescaling to prevent MSE spikes as scale grows model.fc_C.weight.mul_(model.scale / new_scale) model.scale = new_scale
Solutions began to converge even for the larger weight coefficients in the output matrix. On the surface, this was great and offered tangible value for approximate computing. But we wanted absolute truth - exact, closed-form expressions with rational numbers Q that can be reduced to deterministic equations for compute kernel generation!
And that’s when I remembered an old trick. Years ago, I used Ordinary Least Squares (OLS) for signal matching in fractal compression, and a thought struck me: why not take these noisy, continuous matrices and fit them against the exact, symbolic target convolution?
# --- Analytical Solver Core (SymPy + OLS) --- 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. Iterate over denominators q and snap WD/WF to rationals 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. Analytically solve for ideal WC using Least Squares (OLS) 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. Symbolic verification in SymPy polynomial ring 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 matrix residual is strictly zero, the exact algorithm is proven! if residual == sp.zeros(*residual.shape): print(f"[SUCCESS] Exact identity proved at q = {q} (residual = 0)!") return WD_exact, WF_exact, WC_exact, True return None, None, None, False
After some black magic with this concept - right at the break of dawn - I hit run, and once again: "No freaking way!" A 100% exact match for Winograd matrices!
--- 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 ]
I brewed yet another cup of tea and started running kernels for which I had never seen or heard of a Winograd formulation - and it works! The method actually works!
================================================================= 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)! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Well, it seems I’ve saved the planet gigawatts of energy per year once again. But this time - without closed corporate libraries or proprietary engines.
While I was packaging the source code into an open-source repository under the MIT license, someone in the dark was clicking away at a keyboard, muttering something under their breath. Finally, I pushed everything to the repo: https://github.com/MaxAkaAltmer/apex_winograd and peacefully fell asleep...
