TL;DR: New speedrun here. Details below. I generally experiment in PyTorch on my local GPUs. Recently, I received some TPU credits from the TPU Research Cloud (TRC) program to do some research, which gave me a good reason to finally learn JAX. Near about all I knew was that you could often replace numpy with jax.numpy and get code to run on accelerators, and that JAX code is supposed to be functional. This was also the first time I wrote code for TPUs. With this, I set out on a small side project: porting the modded nanoGPT speedrun to pure JAX with the goal of achieving the best possible performance. ...
Table of Contents Introduction The Setting The Optimization Problem Subgradient of the Loss Function The Family of Geometric Optimizers A Convergence Theorem Convergence Bounds for Specific Scenarios Normalized Gradient Descent Muon Adam/SignSGD A Note on Bound Tightness and Rate Variation Analysis of Rate Variation Analysis of Bound Tightness Using This In Practice: A Comparison of Optimizers The Problem with Vanilla Gradient Descent Geometric Optimizers as the Solution Comparing the structure for Normalized Gradient Descent, Muon, and Adam/SignSGD The Nature of the Updates Building Intuition from the Convergence Bounds The Meaning of \(C_{\text{lower}}\) and \(C_{\text{upper}}\) Rate Stability and Comparing Optimizers Further reading Introduction We want to answer the question: how do optimizers fundamentally differ in their approach to finding a minimum? To explore this question in a controlled environment, we focus on a simple problem: minimizing the matrix error term \(\min \lVert AX - B \rVert\). ...
Figure 1: Attention score similarity with our 3D RoPE. These plots show how positional similarity changes across an axis. Introduction RoPE has become the de facto positional embedding for transformer models. Its popularity mainly stems from its performance, but the “derivation” in the paper is also quite elegant (but flawed). Implementing high dimensional RoPE also pushes us to think about generalizing the underlying ideas as much as possible (alongside using signal processing intuition) - there’s code at the end of the post that implements things based on the ideas we develop here. ...
So I’ve been solving IMO problems almost every year, ever since I was a math Olympiad contestant, either posting them to AoPS or just keeping the solutions to myself/discussing with other math Olympiad people. Usually, I used to solve the problems in a single sitting, but given the fact that I missed a couple of years at some point, and that I had a bit of time this time, I decided to do a proper IMO mock this time (spoiler: P4 left a bad enough taste that I didn’t want to really put in effort into P6). ...
Motivation Let’s start by doing some arithmetic about large language models (LLMs). These are neural networks with huge parameter counts, with state-of-the-art open-weights models (i.e., ones you can download) having parameter counts of the order of 100B (\(10^{11}\)) or so (and usable ones around one order of magnitude smaller). Take the latest SOTA release Qwen 3 235B-A22B, for instance, which has roughly 235B parameters. If all these parameters were to be stored in a naive array of 32-bit (4 byte) floating point numbers, this model would require around 940 GB of storage as well as memory for a usable speed. Running this model purely on CPU with dual channel DDR4 RAM (which is likely the kind of RAM you have on your computer) would take you multiple seconds to output a single token/word (and even this is quite fast for the total size of the model because the architecture is what is called a Mixture of Experts, more on that later, so don’t worry yet). ...
Background and context Some background on me just so that you have a rough idea of where this review is coming from: I did Olympiads as a kid and have been involved in fairly math-heavy fields ever since, through university at least, depending on your definition of math-heavy. I’ve also been completely self-taught when it comes to the non-elementary-school math I know. So my plan on using my Math Academy subscription was to mostly brush up on things. I also did only the university-level courses. ...
If you’re here for the game, visit this. Calibration Everyone knows that being overconfident can often lead you to making reckless, unnecessarily aggressive decisions, and being underconfident leads to not taking enough opportunities. This duality shows up everywhere in real life - from aspects like investing/business, where real money is at stake, to more personal matters like career progression and navigating interpersonal relationships. If we were always right about things, we could blindly believe in ourselves, and overconfidence would not exist. If we were completely clueless (for some definition of completely clueless), we would be better off asking a stochastic parrot to take our decisions for us. ...
If you have been on the internet for a while, you probably know by now that LLMs (large language models) are a thing, and talking to them feels pretty close to talking to a human. If you have been on the internet for a decade or so, you can probably guess what this blog is going to be about. A harmless query… or was it? The task at hand was downloading an interesting video from YouTube that I happened to come across in my old link archive using the ever-handy yt-dlp. ...
Introduction When we reason qualitatively, we often rely on our intuition. However, intuition is often loaded with certain meta-biases that cloud our judgment; one of these biases comes into play when we think about “local” and “global” statements. What is a local or a global statement? One way of distinguishing between them is how many conditions we must guarantee hold before we can talk about the statement - so this terminology is relative. A local statement is a statement that uses many more conditions (a higher amount of specificity) than a global statement. For example, any statement about a certain group of people in the past few years is a local statement relative to a statement that is about all species that have ever existed on the earth. ...
As one can easily guess, this blog is about a mental model of thought. I chose to write about it since I feel that introspection about ways of thinking (and consequently what “thinking before doing something” means) is greatly lacking among people, and that it is critical to make better decisions (and realizing when there is no “better” decision). I won’t bore you with the philosophical details, so I’ll approach it from a probabilistic perspective, which is closer to what I personally choose to think in terms of. A word of caution: I will sometimes oversimplify things in order to drive home the point, but sometimes it might seem to contradict with what I say in the later parts of this post. The key here is context, and if you keep track of it, things will make more sense. To understand probability concepts that I mention here in a bit more detail, I recommend reading my post on probability. If you don’t understand/don’t want to go through the math examples here, don’t worry - I’ll intersperse it with a general idea of what we are trying to do, so looking for those explanations should help. Wherever you find something that’s not something you already know, you should probably just make a note of it and go ahead (and read more about it later). If it is still not clear, you can let me know and I’ll try to clarify that part (for you and the other readers of this post). Or even by just looking up the thing on your favorite search engine/LLM, you’ll likely learn a lot, even if in a less pointed manner. ...