google-deepmind
Geometric algebra
Ive recently authored a geometric algebra package for JAX, and it occurred to me that a project like alphafold might benefit quite a bit from adopting it.
If nothing else, one feature of using numga I have found very useful in my own work, is to easily swap out the execution strategy for your operations. In my experience, this can matter up to 10x, both in compilation and runtime performance. Right now alphafold appears to use a mix of dense and sparse product execution styles; eg (quat_multiply) and (vecs_cross_vecs).
Numga completely seperates the mathematical logic from the execution logic, and it is easy to swap out execution styles. If you toggle this line you can get an impression of how much it matters on your hardware, on a fairly representative piece of code being executed. In my benchmarking, dense products can compile up to 10x faster, while sparse ones tend to execute 10x faster (at least on CPU).
While it is super elegant, it is not obvious to me that going all-in on the GA 'motor' paradigm is the right call for a performance sensitive application like allphafold; and i have some open questions with respect to numerical integration errors myself. But one can perfectly well use a geometric algebra library in conjunction with a 'classical' translation + quaternion state encoding. For my own work, my preferred approach is to start prototyping my code in the most mathematically elegant manner, and to drop down to lower level constructs only where I can show it provides a benefit; and to keep both implementations around, so I can unit test them against one another. But maybe encoding all states with motors would work out really well for alphafold; I wouldnt dare say without having tried.
I dont have the bandwidth to dive into alphafold to try and make a branch myself, but id be happy to see where I can help if anyone else is interested in giving it a go. If nothing else, playing around a little with geometric algebra is a lot of fun, I can promise that much.
Just replacing the /geometry/, r3.py and quat_affine.py math functions with GA equivalents should not have to be a lot of work, come to think of it, as long as the state representations are kept the same..
Note that the current JAX backend only implements array-of-struct memory layout multivectors. Adding support for struct-of-array should not be hard at all; the thing im still trying to wrap my head around is what should happen if SOA and AOS multivectors interact? Is there a general rule for this; or would it make most sense to implement a user-overridable lookup table to specify what type you are supposed to allocate when you multiply a SOA-quat and a AOS-vector, and so on?
I do not quite understand the mix of SOA and OAS in alphafold in the first place. My first guess would be that SOA is always preferred, when working on GPU/TPU, with large batches of identical objects; as is the case for alphafold. Why mix in AOS layouts in the first place? Is that a deliberate choice; or more a matter of it being too annoying to rewrite all code to SOA?
@Augustin-Zidek I see youve been working on a lot of this part of the codebase; would you be able to comment on the above questions concerning SOA-vs-AOS?
Hi @EelcoHoogendoorn, sorry for late reply. We stick to the AlphaFold development plan; unfortunately we does not plan to substitute our current libraries with the library you have kindly mentioned, as this may have negative effect on the future codebase support and affect other projects.
