KANs and AlphaFold 3

Taking a break from my posts and reading on thermodynamic computing just earmarking here a lot of interesting things that are going on in the ML/AI world.

Kolmogorov-Arnold networks (KAN)

Last week there was a preprint about a new deep learning architecture called Kolmogorov-Arnold networks(KAN). These are in contrast to multi-layer perceptron (MLP). What is really interesting is that instead of learning the linear weights on edges, $\mathbf{W}_i$¹, of a layer:

$$ \sigma(\mathbf{W}_i\cdot\mathbf{x}_{i-1} + \mathbf{b}_i) $$,

that are passed to a fix non-linear activation function, $\sigma$, the network learns the parameterized non-linear activation function themselves. The learning of the non-linear functions is done on the edges and then a summation operation on the nodes (i.e., where the activation function in MLPs would normally be done) is performed. Mathematically for the interlayer (i.e., edge operation) you would have:

$$ \begin{equation} \Phi(\sum_i^n \phi_i(x,\mathbf{\alpha})) \label{eq:KA} \end{equation} $$

here the $i$'s are the edges the functions live-on that link node in between layers. The key point is $\phi_i$ are the learned non-linear functions with parameters $\mathbf{\alpha}$². In the paper they use b-splines to learn the functions $\phi_i$. My guess is they choose these over other basis functions because b-splines can be locally controlled and are smooth, but there should be no reason one can't use other functions.

I didn't state why this would work, it turns out that KAN are based on the Kolmogorov-Arnold representation theorem which is a connivent way to represent real-valued multivariate functions as a sum of univariate functions over the domain $[0,1]$. I did not know about this!

This is my best understanding of the KAN paper, but I've barely scratched the surface of the 48 pages. The other thing is that the function composition in eq. $\eqref{eq:KA}$ is represented in KANs as a matrix operations, i.e., $\text{KAN}(\mathbf{x})= (\Phi_{N} \;\circ \Phi_{N-1} \; \circ \;\dots \;\circ \Phi_1)(\mathbf{x})$, so I think its just a linear function matrix (see eq. 2.6 in [1]). I need to understand better.

If I am getting this wrong above please drop a comment. The graphical abstract covers all what I tried to convey above better:

Abstract graphic for KANs from [1].

The thing that interests me is that KANs actually seem more intuitive. The network learns the representative non-linear functions that capture characteristics of the data. I guess a MLPs do this as well but less effective? Well at least based on the paper which shows a much better scaling law compared to MLPs.

I don't think I'll have much time to focus more on this but I am very curious about the use of KANs over Gated MLPs in things like Graph Neural Networks. The authors of the preprint put together a very nice and usable python package, pykan.

Update 16 May, 2024

The author of ref. [1] posted a review video of the paper.

AlphaFold 3

Now onto the field of computational biochemistry and molecular biology. Google's Deepmind just published there research results for AlphaFold 3 [2]. Basically, it appears they have the ability to develop and discover new drugs using computational means as well as address other life science challenges. Google is probably sitting on a very lucrative resource, so kudos them for investing in it.

The thing thats a little bit off, is how all these big tech companies are doing research and publishing in science journals yet no one has access or knowledge into exactly how these models work and are trained. Its not bad that they want to keep the information internally but they probably shouldn't publish theb. I should make a blog post on this topic about how big tech scientist are now in the business of publishing papers; just a new kind era I guess.

I should mention that they did make available a compute server where non-commercial users can make queries with AlphaFold 3.

Footnotes

---

1. I'll use $i$ to represent the layer edges and therefore the inputs $\mathbf{x}$ will be $i-1$, but this isn't standard notation. ↩

2. This is totally different from how its written in ref. [1], but it helps me understand; at least I think 😏. ↩

---

References

[1] Z. Liu, Y. Wang, S. Vaidya, F. Ruehle, J. Halverson, M. Soljačić, T.Y. Hou, M. Tegmark, KAN: Kolmogorov-Arnold Networks, (2024). arXiv.

[2] Google DeepMind and Isomorphic Labs, Accurate structure prediction of biomolecular interactions with AlphaFold 3, Nature (2024) 1–3. https://doi.org/10.1038/s41586-024-07487-w.

---

Reuse and Attribution

Prime, Compress, Author

When I was in graduate school studying for my qualifying/comprehensive exams, I had to devise a way I could learn as much information as I could. For one, I'm a terrible timed-test taker, I do not compete well in this setting. Remove the time constraint and I'm much more competitive; probably the case with others as well. However, that is a luxury that is rarely afforded to students since the goal of timed examination is to determine how well an individual can recall and utilize concepts in comparison to their peers. Because of my poor test taking abilities I was obsessed with preparing the best I could and formulated a technique that works for me. I'm positive elements of this approach exist in more well known standard methods.

My reason for writing about this approach now rather than 10 years ago is that I've been spending a lot of my time learning new things in the fields of AI, neuroscience, and thermodynamics computing. Because I have no formal background in these topics, I needed to leverage my previous approaches that worked well for me, hence my method used for my comprehensive exams.

Overview

My method is primarily designed for learning from books, but could potentially be adapted for audio or video content. It involves three key steps:

1. Prime: Prepare your mind by uninterrupted reading.
2. Compress: Annotate key ideas in your own concise style.
3. Author: Write as if you were the author of a textbook.

This is what I refer to as the PCA method, which is somewhat ironic given that principal component analysis is a method of compressing data, yet retains a richness of information.

🧨 Prime

The goal here is to familiarize yourself with the concepts and ideas in the learning material without interruptions. During this initial exposure, do not take notes; focus solely on understanding the content. Additionally, avoid consulting supplementary materials, even if parts of the text are unclear. This maintains the focus and efficiency of the session.

For effective priming, shorter segments are generally more manageable than entire chapters, as chapters often cover broad topics with complex details. Engaging with smaller content blocks allows for a more focused approach.

At the end of reading you should feel like you have a general sense of the content and what is important.

📦 Compress

I find this step crucial as it solidifies the understanding of the material. Re-read the text from the Prime step and take notes using paper and pen in real-time. The format of your notes should reflect your personal preferences, whether as text chunks or bullet points, etc. Include diagrams and summaries of important figures in your own style. Add footnotes for references or important details to revisit later.

The aim is to distill the content into a personalized, dense, and informative resource. These notes are meant for your use, tailored to your understanding. Don't write them thinking you'll share them with others, so "being messy" is okay!

📚 Author

Now, adopt the mindset of an "Author" or "Teacher." Use the notes from the Compress step as a basis to expand into a clear, detailed, and presentable format. This step involves refining your notes into a resource useful to others. Aim for thoroughness and professionalism.

Initially, try to rely on your memory and notes to guide your writing, ensuring the content is conveyed effectively. However, feel free to refer back to the original material and your notes, research further as needed, and enhance your explanations with additional resources.

Sometimes I call this step the "Being Feynman" step, in reference to the Richard Feynman technique. You could say PCA is a form of the Feynman technique, but as Dirac's self-appointed student it would be blasphemous to do so, 😉.

Re-read & Polish (Optional)

If time allows, revisiting the original text can enhance the accuracy of your work. Review your notes and the final document, possibly consulting with peers to identify and correct misunderstandings. Recognizing errors independently can be challenging, making peer feedback invaluable.

Tools You May Want to Use

1. Prime: Utilize the primary text, audio, or video only.
2. Compress: Choose quality writing tools and your preferred type of paper; I favor plain white printer paper.
3. Authoring: Consider various tools based on your needs:

App	Scope
$\LaTeX$	Ideal for non-computational writing, with Overleaf providing a quick start.
Obsidian	Excellent for non-computational, markdown-based documentation.
Jupyter Notebooks	Best for including computations and visualizations.
Quarto	Similar to Jupyter Notebooks but with expanded document format options.

Footnotes

---

1. Dividing the Prime step into smaller segments simplifies the Compress step by reducing the volume of notes and easing cross-reference with previous annotations. ↩

Reuse and Attribution