A New Chapter?
How can one resist the temptation to understand intelligence?
This is my very first blog post, so I want to briefly explain my interests. We are building human-level (perhaps soon beyond-human-level) intelligence in the lab. This is mind-blowing, isn’t it? Anyone with the slightest curiosity about how the universe works and how we come about perceiving it cannot possibly resist the temptation to understand it. The deeper I dig, the more I realize this would be a full-time effort.
However, it is very scary and identity-crisis-inducing to imagine transitioning away from a career you spent all adulthood on! To alleviate this feeling, I decided to write this post, or more like a journal for myself, and walk myself back through those moments when I felt particularly persuaded into thinking about AI. Hopefully, I will get a fresh eye on what I really want to do!
On the more practical side, I will explain why I find pretraining science interesting and some research efforts I have made.
My first astonishment
Prior to the summer of 2024, I was pretty underwhelmed by chatbots such as GPT-4o or Claude 3. They were useful for email revision or brainstorming dinner ideas, but I did not get any insights about physics or math out of them. Then OpenAI’s o1 came out. I asked it to solve integrals of special functions, and it did okay on elliptic integrals but struggled more with hypergeometric functions (I was obsessed with the quantum 6j symbols at the time). But it was very surprising to me that the harder the task I gave it, the longer it thought. How did that work? To me, it requires expertise to gauge the difficulty of these questions.
And then o3 came out, and it was very impressive. I gave it the question about hypergeometric integrals again, and, this time, it solved it in like 30 seconds and even correctly deduced that this is a calculation about out-of-time-order correlators (OTOCs) in JT gravity or the SYK model, where a pile of Gamma functions is supposed to collapse into a hypergeometric function. Doing it by hand had cost me an afternoon (after spending weeks on learning special functions). Of course, in retrospect, an AI model can probably just memorize the tables of special-function identities or a paper from Zhenbin Yang, but it is unclear if that is what the model did. It seemed to not just memorize training data but actually reason through all the steps (well, maybe it memorized all the steps too, but that’s already better than most first-year PhDs). But anyways, o3 nailed it, and I was very impressed.
I then escalated to a research question that Geoff and I have been thinking about since my first year in PhD: in what sense can a limit of a von Neumann algebra become a Poisson algebra, while the modular flows, which are unitary and analytic in modular time, become non-analytic? It thought for three (or four) minutes and gave me circular hand-waving. It restated my question in fancier language, decorated the restatement with the right vocabulary, and presented it back to me as an answer. It failed terribly. But again, it somehow knew this question was much harder than the SYK OTOCs, and allocated its effort accordingly. It seemed to understand and appreciate the varying difficulty of the questions.
How does it work
I was desperate to understand how this works, so I went on a reading spree: chain-of-thought, DeepSeek R1, test-time compute, Stanford RL open course, and so on. It is amusing to me that most algorithms in reinforcement learning are essentially about the variance-bias tradeoff, just like good-old-days, classic machine learning. GRPO (the algorithm of R1) is simply a Monte Carlo algorithm, i.e., no value function. GRPO works in 2025 and not before because (not for lack of trying) the base models became good enough, and thereby the variance small enough, for Monte Carlo to work.
It is amazing how simple the core algorithm is (of course, it must be finicky to get it to work in production). But why does RL encourage generalization and reasoning? In essence, it just boosts the probability of successful rollouts, each of which previously had a low probability. But those rollouts were in the model, in the sense that they had non-zero probability of being sampled. I think Nathan Lambert would call this the elicitation theory of post-training? But I personally find this perspective quite confusing. If a model always outputs a uniform distribution over the vocabulary, does it contain all the knowledge in the world? This smells like the common but wrong argument for the power of quantum computers: it searches over parallel outcomes in superposition, so it should solve every problem in NP in polynomial time, i.e., NP ⊆ BQP. This is wrong because in the end you don’t get to see any outcome freely but must observe according to some probability distribution that can be calculated.
So, elicitation theory feels a bit vacuous (anything can be elicited from random initialization) and doesn’t explain why more capable models seem to generalize better after RL. But I do agree with the statement that the magic is in the pretraining. Thinking about this for a while, I came to the intuition (perhaps absurd) that this is about how well the model can come up with accurate analogies of the rollouts. More precisely, I think the model is, either internally in weights or verbally in chain of thoughts, solving a classification problem. The model first augments the sets of successful and failed rollouts by its own heuristics, e.g., recalling that “this proof resembles a standard induction argument.” But how does the model learn these heuristics? From imitating the pretraining corpus! So I imagine the model first recites samples (internally) that it finds analogous to each rollout, augments the positive and negative sets, and performs the classification. This would explain why more pretraining helps post-training in a concrete, classic-ML way: a more pretrained model can recall more examples and thereby forms a richer dataset for a better classification result. This would also explain why a certain amount of memorization is necessary for reasoning. This would also (more?!) explain why SFT on in-distribution or even on-policy samples generalizes better: the model can come up with more analogies for in-distribution samples, and thereby a richer classification dataset again. One thing this perspective cannot explain is in-context learning, and I will just attribute that to a different internal mechanism inside LLM. Of course, I don’t claim that we can actually find a classification circuit inside an LLM (I will say more about my bitter experience in mech interp at the end), so this is more just for intuition; one could also argue that my intuition is vacuous. I am sure experts in the field must already have ways to reason about these things. I am just exposing my understanding out here for critiques or feedback. Excited about learning more!
Anyways, long story short, I convinced myself that if I want to understand where the intelligence comes from, I have to start at pretraining.
Also, I later found that pretraining research is more accessible than it first looks. A good architecture or optimizer intervention often shows a clear signal in a scaling-law experiment at the ~100M-parameter scale, doable on a few GPUs rented from, for example, Vast.ai or VESSL AI. I will report some results in a later post.
What did I do?
A few months ago, I started thinking more seriously about deep learning research and decided to take some courses at Berkeley: CS288 Advanced NLP with Sewon Min and CS281B Seminar in Deep Learning Theory with Jason Lee. They were more fun and useful than I expected! And this was mostly because of the people. Sewon advised my diffusion language model project (which I will share in the next post) and helped me get compute support from VESSL AI. Through discussions with her, I learned a lot of practical tips to stabilize MoE training and came to appreciate topics such as data efficiency and modularity. One question I have been thinking about recently is, morally speaking, MoE under infinite compute, obviously inspired by Kim et al. They found that tuning weight decay or ensembling leads to much better data efficiency for dense transformers. What about MoEs in practice? Say we have a limited amount of high quality code data for midtraining that we have enough compute to epoch multiple times, what should we do? Unfortunately I don’t have any good idea that could work within the budget of my PhD paycheck 😭
In Jason’s class, I was exposed to many recent ideas about deep learning. It is a seminar class, so each week 2 to 3 people presented papers. What struck me the most were Muon and μP. In fact, I presented μP (the spectral-norm version) and a recent paper about fast hyperparameter transfer that emphasizes that μP is only useful if the scale-invariant part of the hyperparameters converges fast enough. Otherwise, the improvement on optimal loss at scale would be dominated by the suboptimality error. I will write about it in a future post as well. I am also amazed by Muon. Juno talked about a paper he wrote with Jason about Muon’s superior performance on associative memory at the first step. It is very cool! Since optimizers seem like a hot topic nowadays, I will probably try to think more about hyperparameter transfer, even though they are not unrelated (if not anti-related, if you know what I mean).
Anyways, there is really no substitute for talking to experts at the frontier of their fields!
I have actually tried to turn my interest in diffusion LMs and pretraining into a real project, which I will write about in the next post. My hypothesis is that a certain amount of auto-regressive pretraining is helpful even for diffusion, but not too much. Right now, the open-sourced diffusion LLMs all have more than 90% of their compute invested in the AR phase. I tested this at 1M-parameter scale on TinyShakespeare, with no funding at all, and then, once Sewon helped me secure compute from VESSL AI, I scaled it up to 70M MoE models. The same optimal 30% autoregressive fraction showed up at both scales!! I still remember how surprising that was to see. Pure dopamine drive (although quite a small result 😂)
Final thoughts
My PhD taught me one thing: don’t set lofty goals, lower my standards. Fantasizing about progress is counterproductive. Just be honest with myself about whether I understand something, explain it until it sounds trivial, and then iterate on that as fast as possible.
My clock paper is a good example. A famous paper had shown that adding an observer to an idealized universe yields a sensible entropy and a nice algebra in 2022, and the community was excited. Many follow-ups claimed that the same procedure (the modular crossed product) should work in general, despite nobody being able to pin down the exact details. They mostly assume (or hand-wave) the edge-mode structure of gravitational Hilbert space, which is actually subtle in more than three dimensions. On the other hand, what Geoff and I did was very pedantic: we followed our noses and did the simplest calculation in the simplest setting (but slightly more complicated than the idealized setting above) we could fully understand. It turned out the procedures we needed were mostly new (not a crossed product, at least not by construction). Two years later, I still get invitations to workshops and seminars for this paper, and people keep telling me they appreciate a calculation we did (the Euclidean path integral for preparing the canonical state in slow-roll inflation and the boost Hamiltonian). But at that time, I thought that calculation was obvious and boring. If it were not for Geoff, who forced me to write it down as explicitly as possible, I might have ended up spending too much time on speculations about evaporating black holes and more general spacetimes, and the paper would not be as complete and rigorous on slow-roll inflation. Sure, speculations are usually more exciting to talk about, but in retrospect, the main contributions of our paper to the field, at least according to people I have met at the workshops, come from the rigorous, more settled parts rather than the speculative, seemingly exciting ones.
Sorry for the sudden physics drop! Always enjoy chatting about physics🥳
For AI research, I expect myself to apply the same research taste. Lower my standards. Don’t get too fixated on the fantasy. Validate simple things, and explain them until they sound trivial. Then, iterate fast. With that being said, if I am forced to state an overarching goal, summarizing everything above, it probably would be this:
How do we make pretraining more predictable, more principled, and more (especially data-) efficient?
Finally, I have to say, it feels fresh to be this ignorant again. Exposing my ignorance and curiosity to the world like this is a bit scary — and so, so exciting.
P.S. Why not mech interp?
One might ask: “if you want to understand how the model works, why don’t you do mechanistic interpretability?” Yes, that is exactly where I started. I was very impressed by Anthropic’s papers on superposition and induction heads. In fact, a lot of my mental model of the transformer’s residual stream comes from the superposition post.
I then decided to spend some time working on a mech interp project. The goal was to first create model organisms from open-sourced models for misaligned outputs and then train a lie detector by linear regression on the model’s activations on these outputs. I was curious whether a model behaves differently on its own model organism’s outputs versus other models’, i.e., “on- vs off-policy misalignment.” Through the project, I learned useful things, but mostly from the first part: how to set up an inference pipeline for alignment research and data collection. I don’t think I learned much from the mech interp part. The probe I trained on Llama-3 did not generalize well. Of course, I am sure this is mostly a skill issue on my side. There are many fantastic mech interp researchers making real progress on how LLMs work. I will just speak for myself that I have learned more about why LLMs work from, say, μP and attention sinks than from all the mech interp papers after Superposition and Induction Heads combined.