Training neural network to implement discrete Fourier transform (DFT/FFT)

readme.md

My third neural network experiment (second was FIR filter). DFT output is just a linear combination of inputs, so it should be implementable by a single layer with no activation function.

Animation of weights being trained:

Red are positive, blue are negative. The black squares (2336 out of 4096) are unused, and could be pruned out to save computation time (if I knew how to do that).

Even with pruning, it would be less efficient than an FFT, so if the FFT output is useful, probably best to do it externally and provide it as separate inputs?

This at least demonstrates that neural networks can figure out frequency content on their own, though, if it's useful to the problem.

The loss goes down for a while but then goes up. I don't know why:

DFT_ANN.py

"""
Train a neural network to implement the discrete Fourier transform
"""
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential

N = 32
batch = 10000

# Generate random input data and desired output data
sig = np.random.randn(batch, N) + 1j*np.random.randn(batch, N)
F = np.fft.fft(sig, axis=-1)

# First half of inputs/outputs is real part, second half is imaginary part
X = np.hstack([sig.real, sig.imag])
Y = np.hstack([F.real, F.imag])

# Create model with no hidden layers, same number of outputs as inputs.
# No bias needed.  No activation function, since DFT is linear.
model = Sequential([Dense(N*2, input_dim=N*2, use_bias=False)])
model.compile(loss='mean_squared_error', optimizer='adam')
model.fit(X, Y, epochs=100, batch_size=100)

# Confirm that it works
data = np.arange(N)


def ANN_DFT(x):
    if len(x) != N:
        raise ValueError(f'Input must be length {N}')
    pred = model.predict(np.hstack([x.real, x.imag])[np.newaxis])[0]
    result = pred[:N] + 1j*pred[N:]
    return result


ANN = ANN_DFT(data)
FFT = np.fft.fft(data)
print(f'ANN matches FFT: {np.allclose(ANN, FFT)}')

# Heat map of neuron weights
plt.imshow(model.get_weights()[0], vmin=-1, vmax=1, cmap='coolwarm')

@endolith That's one idea they explored in the paper, yes. They showed you can learn some underlying sparse structure of otherwise "fully connected" layers using their method. Using it like that is a new kind of matrix/model compression, and it contrasts a lot with the other approach that's popular in DL, pruning.

They also use the ideas of learning the weights of matrices with butterfly-sparsity-patterns and learning a permutation out a discrete set of choices to show that Adam does indeed find solutions that are the usual "fast" linear transforms from signal processing (FFT, fast DCT, fast DST, etc.).

Hello. I tried to use your code and tried to find the difference between in computing the FFT using numpy and neural network and there was a big difference. Can you provide some insights on how to solve that? I saw multiple papers that are being published in which the authors mention that using the neural network as FFT reduces the time complexity greatly.

@msaad1311

Yes, it's highly inefficient, as I said in the description and the comments. Even more so than a direct DFT because of all the zero weights being calculated unnecessarily. This isn't something that you should be doing. It was just an experiment while teaching myself neural nets.  If your neural net would benefit from frequency domain information, it's better to just do a numpy FFT and pass the output to the net (possibly adding a magnitude function afterward, since that nonlinearity was much harder to learn in my tests).

If you need the net to learn FFT-like transforms in general, look at the comment above about butterfly networks https://gist.github.com/endolith/98863221204541bf017b6cae71cb0a89#gistcomment-3461984

Can you shed some light on the motivation on this? Why do you want to approximate an existing algorithm with a (likely less accurate and less efficient) black box approach?

@sebraun-msr It was just an experiment. I thought that ANN nodes could be equivalent to signal processing nodes and tried implementing some things that way. I've said multiple times that you shouldn't actually do this. :)

For example, I learned that:

• ANNs can learn linear functions like this very easily
• They only learn at frequencies that are in the dataset (so need to be trained on white noise)
• Tried all the different optimizers
• (In a later experiment) learned that abs(fft(x)) is much harder to learn with any accuracy: https://stats.stackexchange.com/q/363352/11633 https://stats.stackexchange.com/q/379884/11633

Added the earlier FIR filter experiment https://gist.github.com/endolith/1de9a8700f72b97974a2e93b0fba316a (and single-neuron experiments https://gist.github.com/endolith/371f0945ab849aa111817fabc395869e)

This work is cited here,
https://youtu.be/mYy4Zv80tXQ
@20:21

@ShriAka That's cool, thanks for pointing it out :)

The demo is using random noise data ， if specified data,such as time series data or picture data is used , is it easy to learn the approprate weight?

@Chenhubget It only learns at the frequencies present in the data