Consider using fft for convolution to reduce cost of multiple convolutions with shared inputs

[ricardoV94]

Description

fft convolution is actually faster than direct convolution for large inputs. Scipy uses it for n>500 [citation needed]. But even if our inputs are not large, we may want to use it if we see the same inputs used in multiple places.

One useful but more involved case is the forward + grad of a convolution, where the reverse of the input is convolved with the output gradient. There seems to be a way to obtain the fft of the reversed input, for the fft of the input: https://dsp.stackexchange.com/questions/82725/recover-fourier-transform-of-flipped-signal-from-the-fft-of-orignal-signal

[victorgarcia98]

Hey @ricardoV94 !

This issue is very appealing to me 😄 Could you provide me some support, please?

I have some questions to start:

1. Where should the logic be implementd? Is it in class AbstractConv?
2. How would be keep the FFT of each matrix cached for to reuseit when we are applying a kernel multiple times?

Thanks a lot in advance!

[ricardoV94]

@victorgarcia98

I would suggest working on top of the new Convolve1d which has a much simpler API than the other convs. We want to replace those by simpler Convolve2d/3d/nd as well.

"caching" just means that if you compile a pytensor function that has the output of the convolution and the gradient wrt to one of the inputs (value and grad in jax parlance), using the fft approach ends up computing the fft of the input only once.

To get started with this work I suggest a proof of concept in a jupyter notebook that we can get meaningful speedups using this approach.

Then in terms of codebase:

1. the first step is to allow the user to specify they want the convolution using fft (like scipy convolve does) and make it work. Hopefully we already have fft Ops we can use.

2. Then make sure we have the clever gradient that uses the trick to avoid computing the fft twice. This may be possible at graph definition or only during compilation rewrites

3. The step which we have to be more careful is to optimize automatically convolution and gradient to use fft even if the user didn't specify they wanted to use it. We may need to use some heuristics based on input size lile scipy does when you call convolve (but done at compile time, not runtime) unless the fft dominates across all sizes. Maybe we only do the trick if there's both value and gradient. That's why step 0 of just exploring the performance in a notebook is useful.

Let me know if this makes sense.

[victorgarcia98]

Hi @ricardoV94 ,
Thanks for the guidance. I've run the benchmark with different lengths of the signal and window size. Not surprisingly, the time depends on the size of the inputs.

Small signal (N=10³)

Large signal (N=10⁶)

---

Code

import pytensor
import pytensor.tensor as pt
import numpy as np
import matplotlib.pyplot as plt
from pytensor.tensor import conv
from pytensor.tensor import fft
import time

rng = np.random.default_rng(1234)

# Parameters
T = 1000
window_size = 500

data = np.cumsum(rng.normal(size=T))
kernel = np.ones(window_size) / window_size

N = T
data_pad = np.pad(data, (0, N - len(data)))
kernel_pad = np.pad(kernel, (0, N - len(kernel)))

r1_numpy = np.fft.irfft(np.fft.rfft(kernel_pad) * np.fft.rfft(data_pad))
r2_numpy = np.convolve(data, kernel)

data_tensor = pt.vector(dtype='float64')
filters_tensor = pt.vector(dtype='float64')

# Padding the data and filter tensors
padded_data_tensor = pt.pad(data_tensor, (0, N - data_tensor.shape[0]), mode='constant', constant_values=0).reshape((1, 1, -1))
padded_filter = pt.pad(filters_tensor, (0, N - filters_tensor.shape[0]), mode='constant', constant_values=0).reshape((1, 1, -1))

# Perform FFT
fft_data = fft.rfft(padded_data_tensor)
fft_filter = fft.rfft(padded_filter)

cplx_multiply = pt.stack([
    fft_data[..., 0] * fft_filter[..., 0] - fft_data[..., 1] * fft_filter[..., 1],
    fft_data[..., 0] * fft_filter[..., 1] + fft_data[..., 1] * fft_filter[..., 0],
], axis=3)

r1 = fft.irfft(cplx_multiply)

r2 = conv.causal_conv1d(data_tensor.reshape((1, 1, -1)), filters_tensor.reshape((1, 1, -1)), filter_shape=(1, 1, window_size))

f_fft = pytensor.function([data_tensor, filters_tensor], [r1])
f_conv = pytensor.function([data_tensor, filters_tensor], [r2])


fft_times = []
conv_times = []

for i in range(100):
    print(i)
    data = np.cumsum(rng.normal(size=T))
    kernel = np.ones(window_size) / window_size

    start_time_r1 = time.time()  # Start time for r1
    r1_eval = f_fft(data, kernel)
    end_time_r1 = time.time()  # End time for r1
    r1_time = end_time_r1 - start_time_r1  # Calculate time taken for r1

    start_time_r2 = time.time()  # Start time for r2
    r2_eval = f_conv(data, kernel)  # Evaluate only r2
    end_time_r2 = time.time()  # End time for r2
    r2_time = end_time_r2 - start_time_r2  # Calculate time taken for r2

    fft_times.append(r1_time*1000)
    conv_times.append(r2_time*1000)


print("\n--- Summary ---")
print(f"FFT Mean Time: {np.mean(fft_times):.6f}ms, Std: {np.std(fft_times):.6f}ms")
print(f"Conv Mean Time: {np.mean(conv_times):.6f}ms, Std: {np.std(conv_times):.6f}ms")

plt.figure(figsize=(10, 6))
plt.hist(fft_times[1:], label='FFT-based Convolution')
plt.hist(conv_times[1:], label='Causal Convolution')
plt.xlabel("Time (ms)")
plt.title(f"Convolution Time Benchmark over 100 Iterations T={T}, window_size={window_size}")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig(f"bench-{T}-{window_size}.png")
plt.close()

[ricardoV94]

Seems like fft has much better worst-case scenario. Couple of questions:

1. Can you add the new tensor.signal.convolve1d to the mix? We did some small benchmarks and that seemed better than causal_conv1d

2. Can you try the backwards pass? The idea being that in the backwards pass we can reuse the fft with the expression in that stackoverflow page, for the input that is reused. Then even in those cases where fft looses in the forward pass it may still gain in the forward+backward because of only having to do the fft transform once.

Either way would be great to open a PR for extending the convolve1D with the method="fft" before we decide on being opinionated about it.