Comment by leecarraher
2 months ago
this seems to follow the similar FNO work by nvidia, and switching to frequency domain is usually in any computer scientist's toolbox at this point, however, I'm curious if this translates to real gains for real architectures. FFT makes use of imaginary numbers to encode the harmonics of the signal, these are generally not amenable to gpu architectures. Would fast walsh hadamard suffice? Sometimes the 'signal mixing' is more important than the harmonics of a compositions of sines. Or do we go further down the rabbit hole of trained transformation and try out wavelets? I am an avid FFT fan, (love fast johnson lindenstrauss transform using the embedded uncertainty principle for RIP), but sometimes real hardware and good theory dont always align (eg there are sub ternary matrix multiplies, but they are rarely used in DL)
Complex numbers work just fine on a GPU. You just represent the data as a real-part matrix and an imaginary-part matrix and do the (ac-bd)+(ad+bc)i stuff in matrix land instead of on complex scalars.
What's wrong with complex numbers on GPUs? You don't have to do anything special. It's obviously faster if you can make simplifying assumptions like "the input signal is purely real" but otherwise at worst you're dealing with pairs of reals (or floats) and don't have to think about philosophical implications.
https://docs.nvidia.com/cuda/cufft/
gpus dont implement complex number fp math, you have to bolt it on as extra logic. cufft works because you can recursively predict the imaginary and real component paths in the butterfly network. between layers you have fft->ifft , is this cost memory locality-wise worth it, or is it better to find ways to tamp down n in n^2 self attention by windowing, batching, gating, many other solutions. im not saying this work isn't cool, FNOs are really cool especially for solving PINNs and related continuous problems, are llms continuous problems, does n have to span the entire context window? I'll probably end up experimenting with this as theyve made the code available, but sometimes good theory is good theory, but not necessarily practical.