Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

You know how there are a bunch of ways to get better results out of a fixed LLM by throwing more compute at the inference (e.g., majority voting across multiple generations)? This paper asks the question, “How well can you get your output quality to scale with various approaches to test-time computation?” And in particular, can we beat big models with small models and more complex inference techniques?

You can’t test every possible approach to test-time compute, so they split them into two buckets and sample representative approaches from each. The buckets are:

1. improving the proposal distribution (sampling from the model) and

2. improving the “verifier” that aggregates different outputs.

The representative approaches they choose for improving the proposal distribution are sampling multiple responses in parallel vs training a model to iteratively refine the previous output. They also try combinations of both.

The representative approaches for aggregating outputs are weighted voting, beam search, and lookahead search (which subsumes beam search as a special case and resembles derandomized MCTS). All of these lean on a trained reward model to score outputs; the voting uses it to weight votes while the latter two use it to score nodes in the search tree.

When trying to improve the proposal distribution, you want to optimize your mix of sequential output revisions and parallel generations (local maxima within a line or group of bars, below). But even if you don’t, throwing more compute at the problem consistently helps.

This need for a balance makes sense because they find that you eventually get extreme diminishing returns from more steps of refinement, with correct answers getting broken about as often as incorrect ones get fixed.

When trying to aggregate responses, you usually do the best with a simple beam search. Although at high numbers of test-time generations, weighted voting starts winning. It matters less what you use for easy math problems than for harder ones (right subplot).

If you choose the best test-time strategy based on your generation budget and a learned estimate of how hard the problem is, you can do better than any individual technique.

Also, using a Process Reward Model (PRM) works better than a simple Outcome Reward Model (ORM)—which makes sense given that it’s being used to score intermediate nodes in a search tree.

So now that we’ve learned a lot about what works to scale up test-time compute, how does their most effective test-time scaling compare to scaling up pretraining?

Well…it’s a little complicated. In the figures below, the x-axis is always compute, with 3 pretraining FLOPs exchangeable for 1 inference FLOP (because of the backward pass). The dotted lines are the combined FLOPs for just using a bigger model, with different lines for different ratios of pretraining to inference tokens. The stars on the dotted lines are the bigger model’s accuracy for a given difficulty bucket and token ratio. So basically, if the stars are above the line of the same color, just pretraining a bigger model beats using more test-time compute.

The takeaway here is roughly that you can replace parameters with inference flops for easy math problems but not hard ones. There’s also a trend where huge amounts of test-time compute start beating the larger parameter count, but it’s a more complex story because they used a single, larger model here and really you’d want to scale up your param and pretraining token counts differently for different regimes.

This paper was great in that it changed my thinking about what works and where we’re headed in multiple ways:

1. Simple strategies like weighted voting just keep getting better as you scale up test-time compute. Reliable improvements from simple techniques are a giant neon sign for big tech companies saying “invest more in me!” and I suspect that’s what we’ll see. This is especially likely because, based on eyeballing the plots, the diminishing returns don’t look at flat as the Chinchilla power laws (note that the y-axes are not log scale).

2. There’s a regime where you’re limited by your ability to recognize a good solution, not your ability to generate one. This regime is the core of why synthetic data works and doesn’t violate the information processing inequality. And now we’re seeing this pattern emerge for inference in general, not just inference to produce synthetic data.

3. In case it wasn’t clear from the amount of synthetic data in the Llama 3.1 paper, the ratio of test-time to training-time compute demand is increasing.

4. Batch size 1 inference is becoming a less important workload. We’ve put a lot of effort into clever band-aids like speculative decoding, but the future might just be best-of-k parallel generations as standard practice for many problems.

5. Tree search with Process Reward Models is a legit strategy that people have independently replicated.

6. Many simple models may be replaced with compound systems consisting of proposer and verifier modules that can be trained and improved separately. I’m sure developers will still do the simplest possible thing in many cases, but I could see “train or prompt engineer a verifier” joining the standard bag of techniques like “add RAG,” “finetune on my data,” “stick some examples in the prompt,” etc.

Of course, the above extrapolates quite a bit from a set of experiments that focus just on math problems—especially since math is unusually conducive to having a verifier. But this is how I now see the trajectories we’re on.

Large Language Monkeys: Scaling Inference Compute with Repeated Sampling

Another Google paper about scaling up test-time compute. This one also examines code datasets, not just math.

Across a variety of coding tasks, they find that sampling more generations keeps increasing the probability of at least one generation solving the problem. This is not surprising for small numbers of generations, but it’s not obvious a priori that this would keep helping even at 10000+ generations.

They also find that whether you should use a smaller model with more generations or a bigger model depends on the task.

In some cases, using a smaller model + more attempts can work better than using a big, state-of-the-art model. A cool example of this is using DeepSeek’s code model to solve simple (but real!) GitHub issues in 5 attempts instead of Claude or GPT-4o in 1 attempt.

The holy grail here though is finding a clean scaling formula that lets you predict the benefits of any given amount of test-time compute. They sort of find this.

The quality of the fit depends on the model and task, but it looks like the probability of not generating any correct solutions decays exponentially with the number of generations.

This is almost really intuitive; with fixed probability p of generating the right answer, we have a negative binomial distribution and probability 1 - (1-p)^k of generating at least 1 right answer in k attempts. But what’s weird here is that it’s not just exp(constant * k). It uses a power of k instead. I’m not sure what to make of this; I’m pretty sure the generations are iid, so you can’t really get this to happen on a per-problem basis. I’m therefore reading this as a statement about the distribution of difficulty across problems in the dataset.

An interesting related finding is that different families of models have different pass@k scaling curve shapes. So you might not get a some perfect, universal formula, but you can at least estimate large-scale behavior from smaller-scale experiments.

The final result I’ll highlight here is that choosing the correct answer from all the generations is hard. Even in cases like math and code where you nominally have a way to check answers, you often get incorrect tests or evaluation code.

My takeaways from this paper are similar to those from the previous one, so I won’t repeat them. But it’s super nice to see two independent teams replicating the same high-level findings, and even better than they’re doing so on somewhat different tasks and datasets.1

Are More LLM Calls All You Need? Towards Scaling Laws of Compound Inference Systems

They find that aggregating more LLM outputs to make a final prediction sometimes makes the predictions worse after a certain number of outputs.

The reason this happens is that some problems are easy and some problems are hard. Aggregating more outputs is always better for easy problems, but often worse for hard problems. Once you’ve maxed out your chances of getting all the easy questions right for a given dataset, you get a regime change where the decreasing probability of getting the hard questions right takes over.

The way I’m thinking about this is that voting reduces the variance of the estimator, which helps when your probability of being right is >50% but hurts when it’s <50%. More votes means you never get unlucky missing something you should have gotten right, but also never get lucky getting something you should have gotten wrong. Worst case, you have bimodal questions with either a 99.99% to get it right or a 49% chance to get it right, so that reducing variance only costs you good luck and not bad luck.

So how do we reconcile these results with the above two papers finding monotonic scaling, or even closed-form scaling formulas?

The key seems to be the choice of aggregation methods. This paper uses simple (filtered) majority voting, which the other papers also found to not scale as well as alternatives.

That said, I buy their theoretical analysis since it leans on simple assumptions. So I would expect that it’s possible to get non-monotonic scaling under some circumstances even with fancier reward-model-based aggregation functions. And even if it’s not, this paper gave me clearer thinking about ensembling.

It’s also encouraging to see independent empirical and theoretical evidence that sample difficulty is a key factor to look at when reasoning about test-time computation.

When Scaling Meets LLM Finetuning: The Effect of Data, Model and Finetuning Method

To what extent to different factors, like model size or finetuning dataset size, affect the quality of your finetuned model?

To answer this question, they did a ton of experiments on 1-16B param LLMs for various multilingual tasks with various finetuning methods.

First, finetuned model quality is more sensitive to model size than finetuning dataset size. This holds across full model tuning (FMT), prompt tuning, and LoRA. I’m reading this as more evidence for the superficial alignment hypothesis.

Same goes for the finding that pretraining dataset size matters more than finetuning dataset size (note that the rightmost subplots are constant because they don’t do any finetuning):

But even finetuning dataset size matters way more than the number of parameters added by your PEFT method. This makes sense since LoRA and prompt tuning add so few params relative to the size of the original model.

Overall, they manage to fit power law curves to these findings with minimal error even on held-out scaling configurations.

Although the coefficients in the formula vary by finetuning and method and task. Ideally, you’d like some universal constants that let you avoid curve fitting on new problems—but I would be shocked if real-world data admitted such a clean shortcut.

There’s a ton of detail in this paper and it’s worth staring at thoroughly if you care a lot about finetuning. There are always question marks around how well the findings generalize to other problems, downstream metrics, etc, but sweeping this many {finetuning methods x model sizes x pretraining data sizes x finetuning data sizes} is already unusually comprehensive and a great contribution.

Scaling Law with Learning Rate Annealing

So you know how we have power law scaling curves to let you estimate the final loss for an LLM training run given a bunch of other, smaller runs?

What would be even better is if we had a formula that also predicted loss within a training run, taking into account the learning rate schedule. Besides letting us save a lot of compute by predicting the outcomes of training runs early, this would also shed light on what design choices actually matter in a learning rate schedule.

We now have such a formula:

Let’s break down what this says and discuss what they do and don’t demonstrate. The core claim is that the learning rate schedule can be characterized by two constants: the area under the schedule (S1) and the sum of the momentums of decreases in the learning rate (S2).

The S2 term is the weird one and makes more sense when you think of it as an integrated exponential moving average (EMA) of learning rate decreases. This says that bigger, earlier drops reduce the final loss more.

They arrived at this EMA formula based on empirical observations of the relationship between learning rate and loss. The loss follows the learning rate (including increases!), but with a small, setup-specific delay.

I saw this same relationship of learning rate driving quality in our cyclic learning rates paper, but I’m actually not sure if we saw the delay. The raw loss curves are lost to time/account migrations, but I don’t think we saw it with downstream metrics (part of our Figure 1 clipped below:)

But regardless, this observed lag between learning rate changes and loss changes is what motivates their use of learning rate momentum in their formula.

So how well does this formula work? It’s not perfect but it’s pretty good, even on funky learning rate schedules.

It also plays nicely with parameter and dataset size scaling if you just add the classic power law terms for those features.

In principle, this means that you could fit your parameter and dataset size scaling coefficients based on just one or two runs, using each run to collect many samples instead of just the final loss.

When they try this though, they get way a worse scaling exponent than the Chinchilla paper (~.21 instead of ~.34). But it’s unclear if this is a result of the training setup or the proposed sampling of multiple points per curve.

One impressive aspect of their results is that they fit their coefficients not on the whole training curve, or even a prefix of the training curve, but on entirely separate training curves. They fit their parameters on loss curves that used simple learning rate schedules and generalize from there to other schedules.

As an aside, they have an interesting result showing that gradient norms end up proportional to learning rate. You might expect the updates to be proportional to the LR since the LR is, you know, multiplied in at the end. But the gradients growing with LR is not obvious. I think (?) this suggests nearly-constant curvature in the update directions, so that the mapping of step size to change in gradient norm is fixed.

Back to the main results, their reduction of learning rate schedules to a couple statistics and a simple formula has a lot of explanatory power. In fact, they have 8 straight subsections describing how they explain observed phenomenon related to learning rate schedules.

This is really cool, and exactly the kind of work we need more of as a field. Getting something resembling Maxwell’s equations for deep learning is one of the most important problems we have, and this paper—assuming it replicates—looks like a big chunk of progress on this front.

And I mean a big chunk—we’ve gone decades without a good model for how learning rate schedules work and proposed dozens if not hundreds of schedules, and this paper just comes along and unifies seemingly all of it with two terms and four free parameters.

That said, I don’t think this is the last word on the problem. First, as they point out, their formula is game-able in a way that you wouldn’t expect if we’ve fully nailed the “physics” of learning rate schedules.

First, running steps with a learning rate of 0 still decreases their predicted loss (though note that they propose a variant to mitigate this). Second, if I’m thinking through it correctly, the best learning rate schedule according to their formula should be just the highest LR that doesn’t diverge for most of training followed by an instant drop to ~0. Maybe this actually is better, but the ubiquity of smooth decreases makes me skeptical of rewarding sudden plummets. Finally, I can’t tell how good a fit we should expect from just extracting some sensible features from a learning rate time series and fitting some parameters to them—if there are lots of ways to do this with just as good an R^2 value, then it’s less likely that this formula reflects some underlying physics of neural net training.

So I’d be surprised if we all end up using this exact formula to fit training curves—but now that this paper has shown that such fits are possible, I do expect some extremely similar formula to become standard practice.

The dev team is healthy and the training run is hitting its quality targets. But boy are there a lot of failures to debug late at night…

The Llama 3 Herd of Models

In case you missed it, Meta released a family of models that are about as good as the best proprietary models.

This is a huge development for the field and also 73-page paper so…today’s post is mostly going to be about this paper, with less coverage of the literature in general.

Architecture

• What’s interesting here is that it’s not interesting—just a vanilla transformer.

  • Of particular note is that they didn’t use sparse mixture of experts (MoE), which has a strong track record of improving cost-quality tradeoffs. This suggests they’re serious about the “becoming a standard” angle, since it’s much harder to do training, finetuning, and serving well with MoE.

  • It might also be a result of their interconnect bottlenecking their training on gradient communication, which would mean their training time is proportional to the parameter count. This would make MoE a bad deal during training.

    

Scaling

• They’re trying to predict (and control!) benchmark numbers—not just pretraining loss—through their scaling choices.

  

• To make these predictions, they often fit a sigmoid to downstream metrics like accuracy that saturate. Basically a generalized linear model.

  

• They ran their own scaling experiments and didn’t just use the Chinchilla constants. The power law fit they get is pretty good, but there does seem to be a little curvature (meaning you want a higher token to parameter ratio at higher compute budgets).

  

• They massively overtrained the small models but the 405B model is slightly undertrained based on their scaling results.

  

Infrastructure

• They trained on 8 pods of 3072 H100s, with full all-to-all connectivity within each pod but 7x lower bandwidth across the pods. They only used 16k gpus at a time though (not 24k), possibly because this is a production cluster rather than a research cluster.

  

• They only have two 8xH100 boxes per rack. This is pretty standard and probably a result of thermal limits. This low density is mostly okay because space is cheap compared to hardware, but having everything spread out can require you to buy more expensive cables to deal with signal attenuation.

• They have their own fork of NCCL and their own ECMP variant.

  

• They saw 1-2% speed variations based on time of day as a result of temperature changes. Just in case you thought your own profiling attempts didn’t have enough lurking variables.

  

• Here’s their overall breakdown of outages/interruptions. They nominally trained for 885k H100-days and had 268 GPU hardware failures, for a GPU-level MTBF of about 3300 H100-days. This seems…way too long. Or maybe Meta is unusually good at datacenter design? Either way, it’s still a pretty high failure rate since you can’t hot swap the cards in an SXM box—you need 0 of 8 to fail for the server to work, so each server will fail a little less than once a year at this rate.

  

• Despite the failures, they got >90% training uptime with only 3 instances of heavy manual intervention in 54 days. Getting your uptime this high is really hard to do.

  • They did a ton of work to reduce job startup and checkpointing times. E.g., they checkpoint to a giant SSD cluster with 2TB/s of sustained bandwidth.

  • Part of their ability to fix and prevent issues came from integration between their NCCL fork and PyTorch to get more detailed debugging information.

Sharding and Parallelism

• They’re doing 4D parallelism with {tensor, context, pipeline, data} parallelism. These are in order of more local → more global. The TP (and maybe the CP?) are intra-box while the DP is inter-pod.

  

• They didn’t need any activation checkpointing. But to achieve this, they did need to manually control some tensor lifetimes and mess with some NCCL details (like TORCH_NCCL_AVOID_RECORD_STREAMS).

• They do interleaved pipeline parallelism, where each group of machines handles more than one pipeline stage. This reduces pipeline bubbles.

  

• Their context parallelism just plain allgathers the keys and values right there on the critical path. But this is okay-ish because the KV tensor is small thanks to grouped query attention.

• They use FSDP in Zero-2 mode, not Zero-3 mode. I.e., they keep the weight tensors materialized after the forward pass instead of re-gathering them in backward.

• They use f32 gradient accumulation, including for their reduce-scatters.

• Overall, they get around 40% MFU. You’d like 60%+ for a dense model within a single pod, but considering that they have to allreduce across pods and pay other scale taxes, this is impressive.

  

Pretraining

• Over the first 252M tokens, they increased the batch size from 4M to 8M tokens and the sequence length from 4096 to 8192. They later increased the batch size to 16M tokens after 2.87T tokens.

• They dynamically adjusted the data mix during training based on how intermediate evals were looking on downstream tasks. They also added fresh web data late in training to increase the knowledge cutoff.

• They did 800B tokens of long context pretraining, gradually increasing the context length from 8k to 128k. They kept it at a given context length until short-context eval numbers fully recovered and the model solved “needle in a haystack” perfectly at the new length.

• They used 40M tokens of extra high quality content at the very end.

Post-training

• They did 6 rounds of iteratively post-training the model and using the new model to collect human feedback and generate possible responses. They use the best and worst1 responses (according to a reward model) as a contrastive pair for supervised finetuning.

• They chose DPO over PPO since it’s not only more compute efficient, but also “performed better, especially on instruction following benchmarks like IFEval.”

• They let human annotators edit responses in addition to rating them, and treat edited responses as a separate, higher tier of quality in their contrastive training.

• Since most of their post-training data is LLM-generated, they add a bunch of rule-based cleaning steps to account for the LLMs’ quirks. E.g., they downsample outputs that start by apologizing.

• Even for real data, they have elaborate pipelines that try to deduplicate documents, score them for quality, and estimate their difficulty. Within a given topic cluster, they greedily take documents that appear different from documents added-so-far in order of quality * difficulty.

• Besides the synthetic data generation loop, they have a bunch of static, high quality data they’ve collected. Some of it is offline synthetic data that’s used to help with particular capabilities like coding or function calling.

  

Downstream task capabilities

Coding

• They forked the model near the end of training and did continued pretraining on 1T tokens of code-heavy data to get a “code expert” model. They then post-trained this model much like the main model. This code expert lets them generate millions of synthetic code-debugging chats.

• They generate their code data by:

  • Seeding a prompt with some real example code as inspiration to get diverse but realistic coding problems.

  • Having a model try to solve the problem and generate tests for it.

  • Filtering out responses that fail linting or the tests.

• For code examples generated during the post-training rejection sampling part, you can’t always pull out the code cleanly. So here they just use an LLM as a judge. They also go out of their way to clean up the most challenging (and therefore valuable) examples so that they pass this filter.

• They augment their set of coding examples by translating problems between different programming languages.

Multilingual ability

• They mostly didn’t train on machine translations, but made an exception for math/reasoning data. This data didn’t hit as many quality issues when translated, maybe because there’s less cultural context or linguistic subtlety here.

• They turned multilingual NLP datasets for various tasks into chats.

• They used parallel texts from stuff like Wikipedia pages.

Math and Reasoning

• They took a ton of mathematical documents and turned them into questions and answers for supervised finetuning.

• They manually looked at what the model was bad at and got humans to generate training data to help with that.

• They learn stepwise reward models to filter out data with incorrect reasoning traces and Monte Carlo Tree Search to generate attempted reasoning traces.

• They prompt the model to generate code to embody its reasoning and filter out responses where the code is broken.

• They try to get correct examples by prompting the model to fix broken ones.

Long Context

• With long contexts, they mostly want the model to be good at question answering, summarization, and coding.

• Almost all the data is synthetic since it’s so slow to have humans annotate it.

• For summarization, they hierarchically summarize 8k chunks, then summarize the summaries, and so on.

• For question answering, they have the model generate question-answer pairs based on random 8k snippets from long documents, but then use the whole document as part of the sample. They also generate some QA pairs based on their above generated summaries.

• For long-context coding, they rip out a file from a repo and have model try to generate the file based on how it’s used in other files.

• They use a 999:1 ratio of short-context to long-context data to trade off short- and long-context output quality.

Factuality / Hallucinations

• They generate training data for “knowing what you don’t know” by generating factual questions from snippets of the pretraining data and using an LLM to check whether generated answers are correct or not.

  

Function Calling / Tool Use

• The models are trained to call a search engine API, call the Wolfram Alpha API, or invoke a Python interpreter. They’re also trained for zero-shot tool use given just an in-context function signature and docstring.

• All their tool use is executed through the Python interpreter. There are custom functions for each tool it can use, and new tools should be Python functions.

• They *don’t* use model generations plus rejection sampling here—just human annotations and SFT on tool use examples from previous Llama 3 models.

Inference

• You can’t fit the 405B model on 8 H100s in bf16, so they use pipeline parallelism. This is probably the least bad option, but gets wrecked by batch size 1 inference. You not only pay for 2 nodes, but only get to use the hardware of 1 of them at a time.

• With fp8, it fits on one node and you get a big cost + speed win.

  

• The first surprising part of their fp8 quantization is that they clamp the row-wise scale factors at 1200. This avoids an edge case where high-entropy tokens cause large activations, which in turn cause underflow in their row if the scaling factor isn’t clamped.

• The second surprise is that they don’t quantize the attention parameters.

I didn’t even get to the vision or speech sections here, but you’ve been reading this for a while and I’ve been reading this paper for hours so I’m gonna truncate the methodology summarization here.

Analysis

What stands out to me more than anything is the huge amount of synthetic data involved. Synthetic data is at the core of everything after pretraining. It’s also a result of highly customized pipelines. We used to incorporate domain knowledge by engineering features, but now we’re instead engineering sampling and filtering pipelines.

I’d expect these synthetic data pipelines to get better as models improve, which has three interesting consequences:

1. There is now a clear, tight feedback loop of AI progress directly causing AI progress in a virtuous cycle. And there’s no end in sight.

2. It’s likely that there will be way more inference compute than training compute in the future, since only a small fraction of the generated data passes the filters.

3. We could be in for winner-take-all dynamics, where the organizations with the best model can generate the best data.

   1. Although Llama 3 405B being open-weights means that no one actually has this structural advantage at the moment. The advantage instead lies in the proprietary code and internal know-how that lets you actually run the data generation pipeline.

4. There may be a more efficient way to turn money into model quality. Even with Chinchilla-optimal scaling, you get brutal diminishing returns where you have to do ~64x as much compute each time you want to halve the reducible error.2 So what you actually do is allocate some of your money to raw scale, and some to extremely expensive labor to shift the cost-quality curve. But it could be the case that the marginal returns on compute for synthetic data are higher than either of these other investments, at least to a point.

In terms of non-technical implications…having one player just give away for free what other firms are spending building of dollars to produce tends to go poorly for the other firms. I know people at most of the foundation model companies and I hope things work out well for them, but the structural economics aren’t looking great—at least as long as Meta sticks to its open source guns.

The good news for foundation model companies is that:

1. Query volume is a structural advantage. More people are hitting [Azure] OpenAI than other providers, and this gives them both higher batching efficiency and more data.

2. If differentiation is impossible, brand becomes a moat. E.g., you can’t beat Coca Cola because there is no way to improve soda enough to overcome their brand power. Maybe the AI APIs will all be about the same and people will just use OpenAI or Google because that’s what they know about.

3. Meta doesn’t have that compelling a case for open weights compared to past examples, so maybe they’ll stop at some point. E.g., Google could give away Gmail, Google Docs, etc, because they had a clear path from user data to profit. But Meta doesn’t directly make more money from people using Llama 3. They just want to make their stack a standard and commoditize their complements (i.e., crush the firms who bid up the prices of GPUs and labor).

4. The sheer amount of revenue going into foundation model companies could get so high that even Meta can’t justify keeping up. We’re nowhere near this, but suppose someone builds a $100B dollar cluster. Could Meta really get enough marginal ad revenue from matching that to make it worthwhile? The digital advertising TAM is not as large as the “all knowledge work” TAM and this could start to matter.

The Illusion of State in State-Space Models

From a theory perspective, most state-space models can’t solve simple problems that involve tracking basic state. Examples include evaluating a boolean formula, composing permutations, or representing the current state of a chess board after a history of moves.

This also applies to transformers, but, interestingly, not RNNs.

This kinda makes sense given that SSMs can be expressed as convolutions. If your function can be expressed as a pure function of the previous layer with no sequential dependencies, it stands to reason that you can’t capture certain sequential dependencies.

Better & Faster Large Language Models via Multi-token Prediction

They add more heads to the LLM to predict future tokens kind of like Medusa, but they don’t necessarily use these extra heads during inference.

Instead, they just use this as a different training objective and find that it lifts model quality for large enough models.

They do show, though, that you can use the extra heads during inference as candidate tokens for speculative decoding. This can yield a pretty significant speedup in the best case.

It’s a bit unclear why this multi-headed training helps quality, but they conjecture that it has to do with emphasizing important tokens in the loss.

A little intellectually unsatisfying, but I’m always a fan of practical time-to-accuracy wins.

DiPaCo: Distributed Path Composition

This paper is a bunch of DeepMind people asking “what would it look like to build the most scalable training setup possible?”

Their answer is roughly that each of your layers would be a mixture-of-experts, but with all the routing decisions determined at the input. You can think of a full sequence of routing decisions as a “path” through the model.

They also choose to route chunks of tokens and not individual tokens, optionally with different train- and test-time behavior. E.g., we might do more granular routing at test time or do top-k with k>1 during training to give experts more data.

The other big difference from typical MoE is that they materialize each possible path through the experts separately—maybe even on separate hardware in separate datacenters.

This is nice and parallel, but means that there are different copies of a given expert in different places. How do we synchronize them?

Their answer is using DiLoCo. This algorithm has each path optimize its params separately and then uses SGDM to sync across all the replicas, with the summed local parameter changes since the last sync treated as a gradient. This lets them sync weights rarely—maybe once every few hundred local steps.

To build this system, they largely lean on a smart task scheduler and Google’s distributed filesystem. This design is surprising to me because

1. usually we all live in SPMD-land with deterministic, synchronous semantics, and

2. usually we try to keep the filesystem or object storage out of the critical path.

For point (1), I’m not sure whether they’re getting sacrificing synchronous semantics and/or determinism. But for point (2), it looks like they’re syncing experts across paths so rarely that they can get away with using GFS and CPU parameter servers to read, combine, and broadcast all the local updates.

It’s a little unclear how to evaluate this approach since it’s so different from common practices. But it passes a bunch of sanity checks—e.g., you’re better off having bigger models and more data.

Similarly, they find that it does, in fact, help to tie parameters across the different paths, instead of just having each path be a separate network.

Side note: I like the below table because it backs up something I think I found but never looked at rigorously: if you scale up the expert count too high, eventually the reduced data per expert hurts more than the higher parameter count helps.

What I like about this paper is that it’s ambitious. There’s so much incentive to work on tweaks to existing practices, and it’s so rare to have people think from first principles about what the best case scenario is. And it’s rarer still for people to turn that thought into a real, working system.

How to Train Data-Efficient LLMs

They explore how to do pretraining data curation well and use their findings get much better pretraining data.

The main tradeoff they explore is that of quality vs diversity. I.e., do you want only the cleanest samples, or samples that cover everything, even if they’re uglier? To explore this, they introduce two methods at the extreme ends of these tradeoffs: Ask-LLM and Density.

Ask-LLM feeds each passage into a pretrained LLM and asks it whether it looks like good pretraining data. The probablity assigned to the “yes” response is the passage’s quality score.

Density uses kernel density estimation to score passages, and then does Inverse Propensity Sampling to upsample passages in less dense parts of the distribution.

So how do these (and other filtering heuristics) work when holding dataset size and training hparams constant? For ~1e20 to 1e21 FLOP training runs, Density and and random subsampling consistently offer some of the lowest perplexities. When evaluating perplexity on high quality data rather than everything (middle subplots), Ask-LLM often does even better. The “overscaling” metric on the right is a normalized score across 111 downstream evaluation tasks, and shows Ask-LLM, Density, and sometimes K-means-prototype-based deduplication doing the best.

When looking at perplexity on high a quality eval set as a function of training time, Ask-LLM consistently works better than perplexity-based sample filtering.

Digging into what’s going on under the hood, they find that Ask-LLM yields nearly uncorrelated sample scores when asking different LLMs, while perplexity-based scoring has high correlation across all LLMs. The latter result isn’t that surprising since perplexity estimates should all converge to the true conditional token distribution. But the lack of score correlation for Ask-LLM is pretty surprising. It also suggests that maybe one could ensemble the different predictions to do even better (or that all of these models are too small, since the two largest models start to become correlated).

Also, the appendix is full of tons of examples of good and bad samples, along with the percentile of quality that both Ask-LLM and perplexity filtering assign to each. If you’re serious about understanding LLM data quality, this is pure gold and I would literally go through all 31 examples.

Besides showing the raw text and scores in this appendix, they also have analysis of the different trends and failure modes. E.g., it often takes a big LLM to recognize when high-perplexity passages are high quality but just talking about obscure topics.

A final interesting point they highlight in the discussion is that Ask-LLM is really expensive as filtering heuristics go. This is because it involves inference in a large LLM for every passage. Depending on how much data you filter vs train on and your training vs inference utilization, this could make your data filtering cost more than your pretraining.

Along with the growing importance of synthetic data, this suggests that inference compute will come to dominate training compute even more in the coming years. Although they do note that data filtering can be a one-off operation that gets amortized over many pretraining runs.

Overall, this is one of those papers to study in detail and come back to repeatedly if you care about LLM data quality.

Scaling Laws For Dense Retrieval

They found power law scaling for retrieval metrics with respect to model size and other factors. Here it is for model size:

And here it is for dataset size:

They also find that, when labeling ground truth (query, passage) pairs, different annotation approaches have different results. The highest-quality automated approach has the steepest slope, but, interestingly, all other methods have similar slopes.

When combining their different results, they get reasonably accurate predictions for what loss you’ll get with larger numbers of annotations at a given model size.

Power law scaling strikes again. This is more evidence that this scaling isn’t just a quirk of decoder-only models or next-token prediction. But since these results are still using text data, we don’t get evidence about which of {data distribution, optimization, transformer architecture, etc}, are necessary factors.

Learning to Generate Instruction Tuning Datasets for Zero-Shot Task Adaptation

High-quality Instruction tuning data is expensive—think tens of $$ per sample. What if we could generate these samples programmatically instead?

The idea here is to:

1. Gather a bunch of datasets from P3 that require a context to answer

   

2. Standardize them into a bunch of meta-templates that include a task type, context, question/prompt about the context, and correct answer.

   

3. Train a model on these standardized samples to go from {task type, context} to {question/prompt, correct answer}.

This lets you feed in regular old passages of text and get instruction + ground truth response pairs from them.

How well does it work? Holding dataset size and hyperparameters constant, using their generated samples works better than no instruction tuning and continued pretraining on the task data (TAPT).

You can also use to their generated data to instruction tune on top of already instruction-tuned models and consistently get further lift. This isn’t true of continued pretraining on task data.

Besides the technical impact of making instruction tuning data cheaper, there’s an interesting non-technical angle here around licensing. The pipeline of:

1. training a sample-creation model on non-commercial data and then

2. using that model to generate your own data

may be a workaround for non-commercial data licensing terms.1 This is because you’re getting value from the non-commercial data but technically not training your final, commercial model on it. In particular, this might let you indirectly train on outputs from other models that forbid this.

What I’m not sure about is whether this is just a temporary loophole that will get plugged in the next round of licenses or an arms race where people will always find a way to work around non-commercial licenses. I’m guessing it’s the former but we’ll see how it plays out. If it’s the latter, that’s a little bit bad for model providers, but mostly bad for training data providers like Scale, Surge, etc, who gain the most from these restrictive terms.

Toward Inference-optimal Mixture-of-Expert Large Language Models

Most work on LLM scaling has only considered training cost and dense models. What if you look at MoE models and take into account inference cost?

To figure out the best course of action, you first need to measure how MoE models’ loss varies with expert count E, base model param count N, and token count D:

This lets you estimate the loss for a given set of architecture hyperparameters. But if we want to minimize inference cost subject to a quality constraint, we also need to know how these architecture hparams map to inference cost.

Under the assumption that you can use whatever inference batch size and device count you want, this is a reasonably straightforward profiling problem (at least with various inference configuration + runtime choices fixed).

Putting these curves together with the training curves, you can derive tradeoffs between training FLOPs, loss, and inference cost.

Of course, getting these exact curves for a real-world MoE deployment is much more painful. The main issue is that you have to characterize your workload and SLAs in a lot more detail. What’s your query arrival rate? What’s the biggest burst in queries you have to be able to handle? Do you need both latency and throughput guarantees? Are these guarantees per user, per organization, or aggregate? Etc. And you also have to think about training utilization in addition to FLOPs.

But that said, this sort of detailed characterization is exactly what people end up needing in practice, and it’s great to see papers doing systematic studies of important questions.

Blending Is All You Need: Cheaper, Better Alternative to Trillion-Parameters LLM

They find that a simple mechanism for combining the outputs from different pretrained models lifts user engagement + retention metrics significantly on their chat platform.

What’s crazy is that the method of combining outputs is to just choose a model from a fixed set uniformly at random for each response. So after each time you type something as a user, you have a 1/N chance of getting a response from each of the N pretrained models.

I’m pretty baffled as to why this works so well. I was expecting to read a positive result for ensembling the logits or something, which would already be interesting because there’s been limited work on ensembling LLMs. But this result is even stronger because it just chooses one model at random without even conditioning on the input or trying to intelligently select which one to use.

My best guess is that there’s some property of how people use chat interfaces that makes diversity of responses really important. Maybe one model refuses to answer, but then the user insists and the next model agrees to answer?

Whatever’s going on, this is a bizarre datapoint, and that’s often the most valuable kind in science.

Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models

They propose to alternately sample new data from their model and then finetune the model using the new data, with the finetuning continuing as long as the metrics keep going up on a held-out eval set.

To ensure they’re getting good samples, they:

1. filter out generations that yield incorrect solutions (doable because they focus on math and coding tasks),

2. use real problems from the dataset as prompts and only generate solutions,

3. take at most 10 passing solutions per problem so that easy problems don’t get overrepresented, and

4. treat the problem as the prompt and the solution as the target sequence, so that they teach the model to generate solutions given problems but not to generate problems.

Despite its simplicity, their method can lift math and coding accuracy significantly if run for 1-3 iterations. Each iteration here includes both sample generation and finetuning.

A Tale of Tails: Model Collapse as a Change of Scaling Laws

Is power law scaling going to go away as more of the world’s text gets generated by LLMs?

First, why would this happen? The problem is that, with a finite number of samples, you end up losing the tails of your distribution (illustrated below with original Wikipedia vs samples from a model trained on Wikipedia).

They study this problem both in theory and practice, and arrive at a mix of good and bad news. As they show in a simplified model:

1. If the tail of the training distribution gets narrowed, you end up with test error that saturates instead of going down (left).

2. If you feed a model’s generations back into it multiple times, you get additive degradation in the test error proportional to the number of cycles of self-consumption (middle).

3. If you include even a little bit of real data in the training distribution, so that the tails just get downscaled and not cut off completely, you end up with a long plateau followed by a scaling exponent that looks like the one you’d get on the original data (right).

The results on real LLMs aren’t quite as clean, but do show similar sorts of degradation.

Especially interesting is their demonstration that having just 2% of the data come from the non-AI-generated distribution can significantly (though not completely) mitigate the damage for a LLaMA 2 model.

Besides the above, they have a bunch of theorems characterizing all of this behavior much more precisely.

What I’m most curious about regarding AI-generated data contamination is the extent to which this is best understood as a technical problem vs a linguistic non-problem.

Like, if I grow up with 80% of what I read generated by LLMs, is it wrong for new LLMs to match that distribution? If this new distribution is now what I think of as natural, why should models seek to replicate some other “clean” distribution? Plus, given how poorly written a lot of content on the internet is compared to LLM outputs, is having more of the latter a bad thing?

On the other hand, having a few LLMs generate a huge share of the world’s text gives a small number of companies + data annotators disproportionate power over common language. This might grant them influence in other ways too.

So depending on how it all plays out, the increasing prevalence of LLM-generated text could simply be a case of linguistic evolution. Like the emergence of new slang with each generation, it might be neither good nor bad, but just different.

Language models scale reliably with over-training and on downstream tasks

The Chinchilla scaling curves suggest you should train your model with a token to parameter ratio around 20. What happens if you use a larger number?

They show that, whatever ratio you use, you still get a power law with the same exponent as optimal scaling; using a sub-optimal token to parameter ratio only worsens the leading constants. Or at least this holds when the token and parameter scaling exponents are the same, as happens in practice.

Empirically, they also observe that there’s an exponential relationship between accuracy on downstream tasks and the validation loss typically described by power laws.

You can plug these relationships together to get an exponentially decaying scaling curve for downstream error rate. They fit the decay constant, initial error, and final error parameters to observed error rates.

Fitting these curves lets them accurately predict the downstream error rates of larger models and training budgets.

Note that the error rates they’re predicting are averages across various downstream tasks. Individual tasks can be high-variance.

A really thorough and practically relevant study. I’d never worked through the math for the scaling exponents at different token to parameter ratios, so it’s interesting to see that it’s just naturally the same with no assumptions needed on top of the standard power law scaling formula. As far as the part about curve fitting eval suite averages and extrapolating…I’ll just say I’m unusually confident these findings will replicate.

No "Zero-Shot" Without Exponential Data: Pretraining Concept Frequency Determines Multimodal Model Performance

Let’s say you have a pretrained image + text model and you ask it to classify some new images with no finetuning or in-context examples—i.e., you do zero-shot classification. If it works well, is this because the model “generalized” really well to these novel classes, or because it actually just saw a ton of examples of these classes during pretraining?

To answer this question, we need some generalization of “classes” that works across different datasets and tasks. Our generalization will be “concepts,” where concepts are nouns that show up anywhere in the pretraining data and at least five times in the eval set.

It turns out that, across a ton of models, downstream datasets, and tasks, there’s a clear log-linear relationship between how many times the concept shows up in image-text pairs during pretraining and how well the model does on the downstream task. That is, you need a multiplicative increase in pretraining concept frequency to get an additive increase in downstream metrics.

This log-linear trend holds even when you filter out the samples in the pretraining data that are most similar to the downstream samples (left) and when you artificially mess with the pretraining concept distribution so that it differs from the downstream concept distribution (right).

Since they have rich concept annotations for tons of image-text pairs across a variety of datasets, they also do some exploratory data analysis. First, they find that a ton of image-text pairs are “misaligned” in common datasets, as defined by having no shared concepts in the image and text.

They also find that the concept distributions across different datasets are correlated. YFCC15M is especially similar to everything else.

Finally, they create a dataset (“Let it Wag!”) from the image-text pairs corresponding to the rarest concepts and evaluate how well various models do on it. Short answer: not so well.

While the results of this paper are scoped to image-text models, this makes me super curious whether this log-linear relationship holds for text and other models. If so:

1. It makes LLM “zero-shot” results seem much less impressive

2. It makes LLM zero-shot efficacy more predictable a priori

3. It bodes well for the de facto strategy of various model providers: avoid the need for out-of-distribution generalization by just training on everything ever, so that everything is in-distribution.

But even if this result doesn’t generalize to other model and data types, it’s still super interesting. We usually see power law relationships in deep learning, in which a multiplicative change in input gets you a multiplicative change in output. Seeing instead an additive change in output is surprising, and I’d love to know what the root cause of this is.

Within computer science, I see three main research archetypes: Scientists, Mathematicians, and Inventors. Understanding what archetypes you tend towards and what archetypes your project falls under can help you choose what to work on, what sorts of results you need to write a great paper, and what pitfalls to avoid.

All mental models are wrong, but I hope you find this one useful.

Scientists

“Science” research strives to create new knowledge through experiments. “Science" papers often ask a question in the title or abstract and answer it with thorough experiments and analysis. They may also put forth a conjecture, such as the Lottery Ticket Hypothesis.

Good science papers:

• Ask useful questions

• Improve the reader’s understanding of some phenomenon, rather than merely reporting numbers. In the best case, the findings become an integral part of the reader’s mental models going forward.

  • E.g., ever since reading this paper, I’ve always thought of neural net loss landscapes as being subquadratic and made sense of results like the edge of stability in this light.

• Make nuanced claims that don’t step outside what their experiments demonstrate

• Highlight what can't be concluded from their results

• Attempt to harmonize their results with related studies

  • Example 1: Chinchilla harmonizing with OpenAI scaling laws

  • Example 2: Scaling Laws for Fine-Grained Mixture of Experts harmonizing with Unified Scaling Laws for Routed Language Models

One “science” trick I learned from Jonathan Frankle and Mike Carbin is to choose experiments that are certain to be informative no matter the outcome, rather than betting on a particular outcome. E.g., when Zack was dynamically varying masking probabilities for BERT pretraining, the result would be interesting whether increasing or decreasing the masking rate was better (or if it didn’t matter at all, or if a fixed masking rate was the best, etc).

Few people write these papers and those that do tend to have above-average experience and experimental resources, making the average quality high. However, asking a useless question, failing to control experiments properly, or failing to convey the circumstances under which the results hold can all render these papers useless.

These papers can also be harder to justify in industry because they don’t move any metrics. They rely on being informative enough that they help you design interventions that do move the metrics.

If you want one-stop-shops of good science papers in machine learning, look at stuff from Behnam Neyshabur, Sara Hooker, and Jonathan Frankle.

Mathematicians

Many papers focus on proving theorems.

What’s great about these papers is that their results are the strongest and most trustworthy. They’re not just usually true, or true in some poorly-defined circumstance; unless there’s a flaw in the proofs, they are universally, unequivocally true. They’re also likely to be timeless, since their dependencies are standard mathematical machinery, rather than particular software.

The bad news is that these papers are the least likely to be useful in practice. While theoretical guarantees can be necessary in domains like differential privacy and causal inference, typically a theory paper’s utility is serving as scaffolding for practical algorithms. E.g., one should almost never embed vectors with random projections, but the Johnson–Lindenstrauss lemma gives us a floor for how well fancier embedding methods should work.

Math research tends to have the highest skill barrier and be heavily tied to some mix of intelligence and thousands of hours of math experience. At least from my non-theoretician perspective, it also helps your impact a lot to choose problems of practical relevance and explain your derivations well.

My personal favorite math papers are Restricted Strong Convexity Implies Weak Submodularity and Locality-Sensitive Hashing Using Stable Distributions, both of which were approachable enough to teach me new ideas while also proving strong scaffolding results for important problems.

Invention

In “invention” projects, we try to make numbers go up and/or to achieve falsifiable properties like fault tolerance. "Invention" papers feature tables and plots wherein one's proposed method has higher numbers than baselines.

Let’s start with the bad news about this category.

First, most of the claimed improvements are just products of multiple hypothesis testing, graduate student descent, or other confounding factors. Of those that reproduce, many merely appear better thanks to insufficient measurement—e.g., saving FLOPs but not wall time.1

These not-so-useful papers tend to cluster around a few topics with low knowledge and implementation barriers, creating a vicious cycle in which:

1. There’s so much overcrowding that no reviewer can expect comparisons to more than a tiny fraction of existing approaches, which

2. lowers the barrier to publication, which

3. enhances the overcrowding.

This can also stop progress on the problem, which keeps barriers to entry low. See, e.g., the sad state of neural network pruning and my big list of discouraging meta-analyses.

However, the best "invention" papers introduce practical methods that are widely adopted. These few successful papers are, to a large extent, what justify the entire research enterprise.

Some examples of these methods in deep learning are:

• Skip connections

• Adam

• Transformers

• The Chinchilla scaling rules

• Instruction tuning

If you want a lot of examples of how to write these papers, Eamonn Keogh’s papers up through ~2014 are the best source I know of. This one might be my favorite paper ever. In deep learning, Noam Shazeer is probably the GOAT.

Invention Subtypes: Edison vs Tesla

I’ve spent the most time thinking about how to be effective at the Inventor archetype. Different people have different research styles here, but I’ve found that much of the variation reduces to a single spectrum: Edison vs Tesla. Oversimplifying2 to highlight the contrast, we have:

Edison

Edison was a speed experimentalist. He seemed to be literally addicted to the variable-ratio reinforcement of trying out different ideas at his lab bench until one worked. He would go until exhaustion forced him to nap under his desk, and he wouldn't return home to his family for weeks at a time.

He didn't sit and derive how to build the light bulb on a whiteboard. He didn't wait for materials science to become better understood. He just tried every plausible substance as a filament until something worked.

That's not to say there was no intuition or understanding behind his process—just that his approach was to forge ahead at top speed with whatever knowledge was already available or discovered along the way.

In empirical deep learning research, this looks like: constantly seeking more compute, making sure your GPUs never sit idle, building tooling to run huge sweeps, disregarding theory papers, and thinking of velocity in terms of experiments run.

Tesla

Tesla was, in some ways, the opposite. He was known to design entire machines in his mind's eye, down to the last screw.

He had high conviction in what was possible and how to do it ahead of time. His predictions weren't always correct, but were sometimes prophetic—e.g., he identified that, with a big enough tower, he could transmit text, sound, and even stock prices all over the world in real time.

In empirical machine learning research, this often looks like systems research—you have some insight about an inefficiency, assumption, or empirical phenomenon that others haven't spotted, you design the core of the system/algorithm before writing any code, and you labor for months or years to make your system real. There's an experimental component, but you roughly know what the results will be based on your detailed analysis of the workload and hardware. It’s implementation, not experimentation, that’s the bottleneck.3

The most extreme example of this style I’ve experienced was with Multiplying Matrices Without Multiplying—I spent years thinking about how to do this, and the moment I had the core insight about how to do the tree-based encoding, it was clear it would be a win. There was ~0 uncertainty aside from the exact distortion numbers—it took months to build, but it immediately worked, with no experiments needed along the way.

For more examples of this style, I’d check out papers about optimized kernels or distributed deep learning, such as those from Daniel Lemire or the DeepSpeed Team.

Conclusion

Many great papers (and researchers) have elements of all three archetypes. And most impactful projects have a mix of rapid experimentation and fundamental insights. But hopefully this breakdown helps you think about what you most enjoy, where your strengths lie, and how to maximize the impact of your work.

P.S.: My analysis and examples are heavily overfit to deep learning, so I’d be curious to hear comments on the extent to which all this applies to other fields. Or just constructive criticism in general.

Before we jump in, one quick clarification from last week: I mentioned how it was an interesting marketing lesson from DBRX development how we spent a bunch of time adding to MegaBlocks but then people ended up associating it with Mistral because they released an MoE first. A couple people said this part was “super spicy” (maybe because of the phrasing of the tweet I screenshotted?) so just to be explicit: I see this purely as a case study in marketing + open source, and no one at Databricks has anything against Mistral at all. We have a ton of respect for them and their models, and I hope they get to succeed as a startup too.

Anyway…

GoEX: Perspectives and Designs Towards a Runtime for Autonomous LLM Applications

This is one of those rare papers that really changed my thinking about where we’re headed.

I’ve thought for a long time about how LLM-generated code would need to be sandboxed, and how you probably want to sanitize the inputs in various ways, and how you might want to postprocess or constrain the outputs to do something like guarantee that you’re getting valid JSON.

This paper points out that what we really want is not a bunch of separate pieces, but a cohesive LLM Runtime that enforces certain properties, even in the presence of an untrusted LLM agent. Some example properties include:

1. Ensuring your LLM agent’s actions are “reversible,” meaning that you can always undo them if you don’t like them.

2. If you can’t get full reversibility, enforcing “damage confinement,” so that the persistent side-effects are well-understood and the rest of the system can work around them. E.g., you can’t un-send an email, but you might have a whitelist of recipients and MIME types.

3. Ensuring that cryptographic secrets don’t get shared with third-party LLM APIs.

4. Sandboxing code execution (think VMs, Docker, WASM, intercepting syscalls, limiting dependencies, etc).

To achieve these properties, they lean on a lot of ideas from classic data systems, like ACID transactions and commits.

While the paper reads like a position paper on LLM agents initially, you realize halfway through that this is hardcore systems research. Like, look at all these different runtime components they built just to safely hit a REST API:

Similarly, they end up with some interesting design tradeoffs around damage control/reversibility. The easy case is if you can lean on a database to just handle undo operations for you. The harder case is if the actions are too complex to cleanly undo, and you have to snapshot different versions of your system state and restore them.

I’m excited about this for a few reasons.

First, it’s just much cleaner thinking about what we need in order to make agents safe to deploy. I’d previously only thought about a hodgepodge of hammers and nails, but formulating this as “We need to design an LLM runtime with a particular threat model that guarantees particular properties” is much clearer.

The second is that this clean formulation means the field will pick up on this idea and make real progress. We’ll see systems researchers for the next 5+ years gradually tightening the guarantees, reducing the overheads, expanding the functionality, etc. We’ll see open source LLM runtimes, with auditable security properties. We’ll see vendor solutions that integrate with containers and/or hypervisors. We’ll see public REST and GraphQL APIs start to add reversibility features to facilitate calls from agents. And more.

The last reason is that I see this as a promising attack on AI doom scenarios. For a long time, this literature has focused on extremely high-level failure modes (e.g., humans being disempowered as we grow dependent on AI) and properties of agents (e.g., being aligned with human values). I feel like this paper has “compiled” many abstract concerns down to the level of verifiable system properties—like, if we can “undo” an agent’s actions, it probably can’t kill us all.

On a more meta level, it also reinforces my conviction in the following maxim:

Never solve with machine learning that which can be solved with regular code.1

For example, people spent decades trying to come up with some algorithmic alternative to backpropagation, in part to avoid the (biologically implausible) storage of activations and synchronous updates. But it turns out we can just devise systems improvements like sharded data parallelism, pipeline parallelism, etc, that make these downsides fine in practice. It’s much easier to make working algorithms fast than to find fast algorithms that work.

In the case of this paper and AI doom, I feel like we’ve spent decades thinking about high-level problems like specifying human values and “target loading” and detecting “deceptive alignment,” and what’s actually going to happen is we’re just going to sandbox the crap out of our agents, ensure their actions can be rolled back, and build good observability and monitoring tools. And we’ll train them on a few million samples of human preference data for good measure, although that’s as much a UX play as a defense-in-depth one. In short, whatever we haven’t solved with RL, philosophy, activation probing, etc., will get solved with strong systems work.

Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws

How much knowledge can you pack into a model with a certain number of parameters?

To answer this question, they consider synthetic datasets that present facts generated from a fixed set of possibilities. In the first variant, they turn these facts into plausible-sounding biographies:

In the second variant, they just kind of list out (name, attribute, value) tuples as English.

What’s nice about these synthetic datasets is that they let you lower and upper bound the entropy of the data, as well as precisely quantify how many bits of “knowledge” the model has learned.

So what happens when you train a bunch of models of different sizes for different numbers of passes through these facts? It turns out that models consistently saturate at just over 2 bits of knowledge per parameter. Not all architectures achieve this bound, but none of them significantly exceed it. Each dot below is a (num layers, num heads) heads tuple for a GPT2-like model. For each number of facts, some architectures fall along the 2 bits/param line—just above it with 1k epochs and just under it with 100.

These same patterns hold for LLaMA/Mistral-like networks, not just GPT2. It also works with GPT2’s MLPs made 4x smaller or eliminated completely.

If you only show the models the facts 100 times instead of 1000, architecture changes start to matter though. Getting rid of the MLPs is especially bad.

In this low-training regime, you can close the gap between GPT2 (better) and LLaMA (worse) by making LLama use a standard MLP instead of a GLU, tying the input embeddings and output softmax, and using the GPT2 Tokenizer.

They also find that quantizing models to int8 doesn’t hurt the 2 bits per param result, while quantizing to 4 bits brings the peak down to ~.7bits/param. It’s not clear whether this is inevitable though or just a limitation of how the quantization happened.

They also investigated what happens when you have 32 experts (topk=1, expansion_ratio=1) instead of a regular MLP (expansion_ratio=4). Despite being super sparse, the MoE only gets 1.3-1.5x fewer peak bits per param. The ratio also gets closer to 1 as the amount of training increases. This suggests to me that the experts might just be undertrained, see they’ve each seen about 1/32 as much data.

Finally, they explored what happens when you add in a bunch of “junk” data. Here, “junk” means the same sort of synthetic data, but drawn from a vast set of names such that you almost never get repeats—just a bunch of one-off facts.

First, making only 1/8 of your dataset non-junk can hurt the model by way more than a factor of 8. This ratio does improve as you keep training though, and eventually the model focuses on the non-junk (repeated) knowledge.

In contrast, if the junk data uses a much smaller set of names in its facts (so that it’s much more repetitive than the regular data), the capacity ratio isn’t affected much. This suggests that deduplicating your dataset might be much less important than removing one-off claims in the data.

If you can’t filter out the junk, you can at least help the model learn to ignore it by prepending sequences with a special token to indicate whether they’re from the regular distribution or the junk distribution. Given this cue, the model learns to ignore the junk.

I see a few actionable takeaways here:

• It might help to prepend a “domain” token to your sequences telling the model what data source they’re from (e.g., wikipedia). Unclear what you do with this at inference time—might need to specify the “domain” as part of your prompt.

• It might be more important to filter out non-repeated data than often-repeated data.

• MoEs seem to need more training per param than dense models, though they’re surprisingly resilient to not having this.

• We probably won’t get below 2 bits per parameter when compressing models without a large drop in model quality.

I also just feel like this is one of those papers I’ll keep coming back to as data to be explained when new science-of-deep-learning papers come along.

Lastly, I’m curious if we can push past 2 bits/param with more training. I could imagine that, say 10T tokens

Scaling Laws for Fine-Grained Mixture of Experts

Should you have bigger experts or more experts in your MoE model? Or more generally, how does your loss scale with expert size and count?

They consider in particular the effects of a “granularity” parameter, where G > 1 corresponds to chopping up each expert into G smaller experts with the same total parameter count.

They also consider the ratio of params in all the experts to params in a typical FFN with one expert and G=1, which they term the expansion rate. In typical MoEs, this is just the number of experts. But with more, smaller experts, we have to divide by the granularity G.

So what do they find? First:

• There’s a power law relationship between granularity and loss (straight line on leftmost subplot).

• Using a granularity of 4 doesn’t mess with the normal power law scaling wrt param count (straight lines in middle subplot).

• Different granularity values don’t change the returns on training for more tokens (lines all similar in the right subplot).

Based on these interactions with the granularity G (or lack thereof), they find that a certain formula involving some power law relationships fits their observations:

This formula turns the granularity into a multiplier on the loss contribution from the parameter count, as opposed to contributing a separate loss term. Otherwise, it looks like the typical irreducible loss + additive param count loss + additive {token, step} count loss.

In practice though, you might not want the most granular routing possible because it starts to hurt your hardware utilization. This is because you’ll have:

1. More data transmitted into your all-to-alls to route tokens to and from experts

2. More activation memory for all your different copies of the tokens (at least with expert parallelism)

3. Smaller experts, which means smaller matmuls

4. More routing overhead

Taking into account the extra FLOPs from the routing, they derive optimal training settings for minimal loss subject to a FLOP constraint:

They also reconcile their results with those of Unified Scaling Laws for Routed Language Models by pointing out that you have to jack up the token to parameter ratio at larger model sizes. This previous paper held token count constant, which explains why they found diminishing benefits from MoE at larger scales.

We didn’t sweep enough hparams to get a formula, but we also found that more granular experts tended to do better—hence our use of 16 choose 4 in DBRX rather than 8 choose 2. The real limitation here though is not the FLOPs from routing, but the memory and allgather overhead; this means that I’d expect the granularities derived in the above table to be overestimates of what you should use in practice.

But either way, studies of how to scale are super important for maximizing and efficiency and model quality2, and this one is quite well-executed.

Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

Why apply every transformer block to every token when presumably some tokens need more computation than others?

A classic approach to this is early exit, where we insert extra output heads at various depths and stop processing a token once one of the output heads has high enough confidence in its prediction.

But this approach adds a bunch of auxiliary heads (of size equal to our vocab size for language modeling…), yields stochastic work for each input, and induces weird and variable microbatch sizes across depths as tokens get dropped.

A cleaner approach is to only feed, say, 50% of the tokens into each residual branch. This yields a fixed, convenient batch size for each block and no extra output heads. And that’s exactly what this paper does.

The question is just how to pick the tokens to operate on. They solve this by performing a top-k operation with a learned router.

Or rather, they would implement it with a top-k op if they weren’t applying this to causal decoder models. In the causal case we encounter at test time, we can’t know if this time step’s token will be in the top-k until we see all the future tokens, which isn’t allowed.

So what they actually do is make the router a binary classifier whose labels are whether the token was in the top-k. This classifier makes an immediate operate vs skip decision for each token at each time step.

So how well does this work? At small scale, they can improve in wall clock time vs loss with mixture of depths. The best setting seems to be only applying their method to every other block, but skipping those blocks for 7/8ths of the tokens.

The way to get better results out of their method isn’t necessarily to just apply it to a given architecture though. You want to scale up the parameter count to compensate for the reduced FLOPs/token.

You might notice that all of the above results are use training loss where the ground truth top-k routing applies. What happens during autoregressive sampling, where we’re forced to switch to our trained binary classifier? Well…looks like we’re fine?

Finally, they test whether their Mixture of Depths (MoD) method plays nicely with Mixture of Experts (MoE), which they together refer to as MoDE. There aren’t super precise scaling formulas here, but it looks like the benefits from each compose. It’s not clear whether you’re better off combining them in the most obvious way or fusing the MoD and MoE routers; however, the latter is apparently better than just reducing the capacity factors and relying on token dropping to jankily approximate MoD.

We’ll have to see whether other people can reproduce it, but this seems like a significant win. And I’m bullish on this working since we’ve seen evidence for even more extreme token dropping in vision transformers.

Another benefit of this method that they didn’t really talk about is the big inference-time speedup you’d get during batch size 1 autoregressive decoding. Since you’re generating one token at a time, any block that the token skips requires zero work, including loading in the weights from RAM. With 12.5% of tokens selected in every other layer and no routing overhead this could get you a 1/(.5 + .125*.5) = 1.78x speedup.

Interestingly, this benefit disappears with speculative decoding or large-batch-size decoding, since you only get to skip loading the weights if all the tokens at this generation step skip it.

Side note: this was fewer papers but more depth than usual. How do you all feel about this?

Now that DBRX is out, I’m no longer in deadline mode and can actually write again.1 Speaking of which:

Introducing DBRX: A New State-of-the-Art Open LLM

We trained a 132B param MoE model on 12T tokens that beats all other open models and the original version of GPT-3.5.

We evaluated the crap out of this thing and I could regurgitate the claims point by point here, but what I think is most valuable to share is some perspective from building it.

People sometimes have this impression that big LLMs are like this space race where we stand around at whiteboards having breakthroughs and the company with the most talent or insights gets the best model.

The reality is much less sexy than that. We’re basically all getting the same power law scaling behavior, and it’s just a matter of:

• Parameter count

• Training data quantity / epoch count

• Training data quality

• Whether it’s an MoE model or not

• Hyperparameter tuning

The real innovation is in dataset construction and the systems you build to scale efficiently.

It also takes immense skill to debug the training, with all sorts of subtleties arising from all over the stack. E.g., a lot of weird errors turn out to be symptoms of expert load imbalance and the resulting variations in memory consumption across devices and time steps. Another fun lesson was that, if your job restarts at exactly the same rate across too many nodes, you can DDOS your object store—and not only that, but do so in such a way that the vendor’s client library silently eats the error. I could list war stories like this for an hour (and I have).

Another non-obvious point about building huge LLMs is that there are a few open design/hparam choices we all pay attention to when another LLM shop releases details about their model. This is because we all assume they ablated the choice, so more labs making a given decision is evidence that it’s a good idea. E.g., seeing that MegaScale used parallel attention increased our suspicion that we could get away with it too (although so far we’ve always found it’s too large a quality hit).

In that spirit, here’s DBRX (pronounced2 “D-B-R-X” or “D-B-Rex”) in the low-dimensional space of LLM design choices. We ablated almost all of these at smaller scale.

• LayerNorm, not RMSNorm

• QK clipping to fix loss spikes, not QK LayerNorm (AFAIK no one else does this because Vitaliy invented it)

• LION optimizer, not Adam

• A lot of data curation. I can’t comment on what we trained on (or even what we didn’t train on!) because lawsuits have started flying. But suffice it to say, we could have made certain numbers go up a lot more if we were less principled. We also didn’t train on any test sets (at least intentionally), although we did do targeted dataset construction with certain benchmarks in mind.

• RoPE rather than Alibi. Alibi was never worse quality and we used it for MPT, but RoPE has better software support across hardware.

• 4-of-16 mixture of experts instead of 2-of-8 using MegaBlocks. This granularity is consistent with the view of MoE as a sparse matrix product. The funny/frustrating story here is that we’d been adding to + using MegaBlocks for months when Mixtral came out and got all the brand recognition from it.

  

• Serial Attention + FFNs (not parallel)

• No attempt to hide the system prompt. Here it is, with comments.

In terms of results, we have a ton of tables in the technical blog post, but what really stood out to me was how flaky different evals are and how hard it is to get your eval harness right. Like, for a while, we logically could not get 20% of the HumanEval questions right because of a bug in our eval code.

Besides raw quality, DBRX is also fast. Since it’s an MoE with only 36B active params for a given token, the number of FLOPs you have to do for inference is much lower than for a 70B model.

We also optimized the crap out of our training stack. We only talked about dense models in this post to not spoil the MoE part, but most of that content applies here. I can’t tell you absolute MFU, but we managed to get about a 1.4x speedup from in-house optimizations after squeezing everything we could out of off-the-shelf libraries.

Besides the technical lessons learned, I think this model has interesting implications for the ecosystem.

For one, it’s going to be increasingly hard for open models to one-up each other. Last year, you could get noticed by releasing a 7B model. At that scale, a startup can justify it as a pure marketing expense. But if you need to drop $10M+ on the training…that’s hard to fit into a marketing budget. Only the companies with serious ROI can justify that. This excludes almost everyone but:

• Popular API providers who can amortize the pretraining cost across many paid queries

• Companies helping enterprises build custom models who need to pipeclean their training stacks (us and Microsoft, plus a few other startups trying)

• Companies selling the promise of “AGI” who need SotA-ish results to raise their next round.

• Companies who profit from LLM progress in general and needed to train a big model anyway (Meta, Google, various Chinese firms)

It’s also an interesting datapoint around AI progress. DBRX would have been the world’s best LLM 15 months ago—and it could be much better if we had just chosen to throw more money at it. Since OpenAI had GPT-3 before MosaicML even existed as a company, this indicates that the gap between the leaders and others is narrowing (as measured in time, not necessarily model capabilities). My read is that everyone is hitting the same power law scaling, and there was just a big ramp-up time for people to build good training and serving infra initially. That’s not to say that no meaningful innovations will happen—just that they’re extremely rare (e.g., MoE) or incremental (e.g., the never-ending grind of data quality).

But anyway, hope you all enjoy it! Here’s the GitHub repo, 🤗 model, and interactive demo (it says CPU, but it’s running on a GPU in our inference stack).

Mechanistic Design and Scaling of Hybrid Architectures

There have been a bunch of proposed non-transformer architectures in the past few months, like Hyena, Mamba, RWKV, Griffen, and more. Instead of proposing yet another SSM/RNN-like model, they try to devise a principled way of making architecture choices.

The first part of their work is devising a set of simple synthetic tasks that test how well models can learn different primitives. Think selective copying, fuzzy matching, or memorizing key → value mappings.

They evaluate various small models with different architectures to check how these architectural choices affect accuracy across these tasks.

A few patterns stick out in these results. First, you’re better off using a mix of different types of token mixing—e.g., not just attention, but attention along with something more SSM-like at another layer.

You also want larger heads rather than more heads. I’m reading this as more mixing of feature dimensions, although the details vary by architecture.

As you would hope, mixture of experts is beneficial.

The real tests though is whether accuracy on these simple benchmarks is predictive of model quality when the architecture and training are scaled up. The correlations are imperfect for each individual task, but when you average across all of them, the toy tasks become a good predictor.

They also assess whether the finding that you should use other token mixing schemes to supplement attention holds at scale. Looks like it does, with the optimal fraction of attention being around 25% of the relevant layers.

My favorite experiment is this investigation of how perplexity varies with state size. By “state size”, we mean the KV history for transformers or the latent state variables for SSM/RNN variants. You get clean fairly clean power law scaling (exponent -.28) within a given class of architectures, but there’s no overall Pareto frontier that everything lies on. It looks like multi-headed SSM variants can get the same perplexity as other variants that need 10-100x larger states. However, the overall appeal of models depends on whether you’re looking at compute (left) or state size (right).

Also, the ordering between different models’ speed vs quality tradeoffs can vary across tasks. E.g., SSMs look a lot better than transformers when you operate on individual bytes without using a tokenizer.

I’m a fan of papers that try to get to the truth instead of sell you on yet another method, and this is no exception. I’m also excited that someone swept state size vs perplexity and just plain characterized the scaling curves for all these different approaches.

I’m also a fan of trying to get tighter feedback loops around design choices, since that’s a super important problem in practice.

And lastly, it’s nice to see corroboration that simple synthetic tasks can be a decent proxy for real data in certain cases.

Speculative Streaming: Fast LLM Inference without Auxiliary Models

In speculative decoding, we use a lightweight model to estimate what the next tokens are and our real model to “check” these tokens. With some clever rejection sampling, you can reduce the number of forward passes through the main model without changing the output distribution at all.

But…why have a separate model to generate the next token proposals? Why not just stick some extra classifier heads on top of our main model, one per future token position you want to generate?

Well, you can. That’s more or less Medusa.

The problem with this is that the classifier heads have a lot of parameters and, if you’re Apple, you’re trying to do this on an iPhone. So you’d rather find a way to reuse the same classifier head you already have.

To do this, you have to instead generate a bunch of extra token embeddings to classify, one per speculated token. They derive these embeddings from the observed tokens’ embeddings partway through the network.

They feed these derived embeddings in starting at one of the later layers and have them attend to the observed tokens, as well as earlier speculated tokens. To increase their probability of having some speculated tokens accepted, they speculate a tree of possible completions instead of just one long completion. Hopefully at least some prefix of the tree will be accepted. This makes the attention masks complicated, but these masks are at least fixed across all iterations.

I’m not sure why (perhaps because of the modified training objective?), but this approach doesn’t preserve the semantics of typical autoregressive decoding. Although it seems to work about as well across various tasks.

Getting people to change their model semantics is always an uphill battle, but it does seem unnecessary to have a separate model to generate the proposal distribution. And if you’re going to get a proposal distribution out of your target model without (significantly) increasing the parameter count, this seems like a sensible way to do it.

The Unreasonable Ineffectiveness of the Deeper Layers

Consistent with previous work, they find that the representations in later layers of the network tend to be more and more similar to each other. Does this mean we can just remove a lot of these layers?

Sort of. For downstream tasks, you can often remove a big suffix of the network (excluding the final classifier head) without hurting model quality for downstream tasks. This is especially true if you run some QLoRA to “heal” the network. But for validation loss on C4, it’s tradeoffs all the way down (see how linear the lower right is).

I think of successive residual blocks as performing steps of gradient descent in activation space.3 So I’m reading this as suggesting that, for many tasks, the samples are so separable that the fine activation adjustments from the later layers aren’t needed.

Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

Similar to Branch-Train-Merge and it’s follow-up paper, but instead of ensembling the different “branch” LLMs at test time or weighted averaging them into one dense LLM, they use the different models to initialize the experts for a single model.

They use a slightly atypical load balancing loss, which replaces each expert’s token count with its total softmax weight without the top-k operation. I’m not sure why that’s desirable—maybe it increases the gradient flow to routing weights for experts that weren’t selected?

A component of this loss that seems counter-intuitive at first glance is that they scale the sum over the N experts by N, rather than 1/N. But the reason you do this is that:

1. the average softmax output scales as 1/N and

2. b) we’re squaring these outputs for each expert. This means the per-expert loss scales as 1/N^2.

So summing over N experts still yields a total loss that scales as N * 1/N^2 = 1/N. Therefore multiplying by N gets us consistent loss magnitude across expert counts.

In terms of results, what they did was start with a LLaMA-2 7B and finetune three copies of it: one on math, one on code, and one on wikipedia. They trained each one on about 200B tokens, and then used these three FFNs, along with the original LLaMA FFNs, to construct 4 experts in each transformer block. They also initialized the attention weights to the average of the weights across all the finetunes.

So how well does it work? On average, it’s way better than just using the initial LLaMA model, as well as each of the individual finetunes. It’s also a bit better than sparse upcycling.

However, the sparse upcycling baseline apparently used far fewer tokens. They gave it more compute time, but it looks like their MoE utilization was so low that a lot of their gains came from just keeping the model dense for a higher fraction of training. For maybe the first time in my life, I wish someone had held FLOPs constant instead of compute time.

Still, their method and sparse upcycling seem to yield a much better training time vs quality tradeoff than any alternative.

They also have a bunch of good ablations. It seems like their sampling based routing (annealed to determinism) works better than more conventional routing (as proposed in the Switch Transformer paper).

The load balancing loss seems to make weirdly little difference overall, but a huge difference for GSM8K (math) and HumanEval (coding).

It’s not clear why these two datasets would be so affected, but they do visualize what’s going on really nicely.

FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design

You may have heard that NVIDIA’s next generation of GPUs will support fp6 (i.e., 6-bit floating point numbers). And this is the first paper I’ve seen getting fp6 to work in practice.

Concretely, they wrote a fused CUDA kernel that dequantizes fp6 values to fp16 before doing a regular fp16 matmul with fp32 accumulators. This is similar to existing 4-bit and 8-bit kernels for low-bit inference, but with a tricky new format.

There are a bunch of reasons this fp6 kernel is hard. First, 6-bit values aren’t byte aligned. CUDA threads can’t just load a certain 6 bits from memory, and if the values for different threads span the same byte, you can end up with redundant loads. The solution to this is to pack together runs of many 6 bit values so that you can always address byte boundaries.

But even this is hard because you want to do wide, power-of-2 reads to get the best memory bandwidth utilization. If you were to load, say, 3 bytes = 4 fp6 values per thread, you’d be doing a 32*3 = 96B read (since CUDA threads operate in groups of 32). But you’d really like to do at least a 128B read. So the lowest multiple of 6 elements is 3x 128B reads = 384B = 512 fp6 values minimum to get nice, wide reads.

But what you need to load in your matmul is a tile, of size 64x64 in their case. So you’re only going to get runs of 64 contiguous values, not 512. So…what do you do?

In short, we’ll shuffle the order of the weights to enable long, contiguous reads while still respecting the ordering required by tensor cores. Essentially, we’re just pre-packing the weight tensor while also taking into account bank conflicts.

Having the layout already permuted for tensor core loads enables an additional optimization, where we keep the data as fp6 in shared memory and only decompress it a slice at a time in registers as needed. This reduces our shared memory consumption and can improve occupancy.

Their dequantization also does some nice bit twiddling to convert fp6 values to fp16 four at a time with only a few instructions.

As you would expect, using their smaller weights lets them run inference faster in the memory-bandwidth-bound regime (low batch sizes). Below, the x-axis is batch size in all cases.

Since there are no model quality measurements here, it’s not clear how w6a16 inference compares to other configurations. But that’s fine—just getting the speed maxed out is already plenty of contribution for one paper.

Personally, I’m not sold on non-byte-aligned scalar quantization as the “right answer”—at that level of software complexity, you might as well be doing vector quantization. But since the new NVIDIA cards are going to support fp6, it certainly does make sense to see how well we can milk that.

They built a crazy fast hardware accelerator based on my approximate matrix multiplication paper. By “crazy fast,” I mean they get 15x higher area efficiency and 25x better power efficiency vs conventional matrix multiply acclerators (think GPUs) when holding transistor technology constant. At ~1.5x better power efficiency per hardware generation, this is about 8 generations of gains at once.

The catch is that they lost some model accuracy, at least for the ResNet-9 they tested on. And this is with layer-by-layer approximation of a pretrained model and subsequent finetuning, not just with post-training quantization.

That said, it’s probably a large net win; a rule of thumb I like is that you can trade 1% accuracy for 2x speed, so ~20x for 2% is much stronger than typical tradeoffs—though we of course need many more experiments and baselines to be sure.

So how do they achieve these gains?

Interlude: MADDNESS background

Before talking about their accelerator, we need to go over the approximate matrix product algorithm they’re accelerating. Basically, we’re going to approximate the product AB by lossy compressing A and B in a special format that’s smaller and that we can operate on super fast.

Common lossy compression techniques for this problem are pruning, quantization, and factorization. These help, but it turns out you can do way better with a different approach (spoiler: vector quantization).

The intuition for this approach is as follows: what if we knew that A only had a few unique rows that repeated many times?

In this case, we could just replace A with the set of unique rows and a bunch of indices, identifying which of the unique rows was in each position.

Given this representation of A, we could also compress B a ton. The observations here are:

1. a matrix product is just a bunch of dot products between rows of A and columns of B, and

2. the rows of A are repeated a bunch of times. So why recompute the same dot product when we could just memorize it?

So what we do is replace each column of B with a lookup table. Each entry in the table is its dot product with one of the unique rows. This is lossless compression for a given set of unique rows.

So now we’ve turned A into a bunch of indices and B into a bunch of tables. How do we multiply them? Easy—it’s just a bunch of table lookups. All we do is:

for i in range(A.shape[0]):
    for j in range(B.shape[1]):
        out[i, j] = B[A[i], j]

That’s it.

Now there’s one big problem with this: the matrix A probably doesn’t have a bunch of exact-duplicate rows. We’re going to solve this in two ways:

1. We’ll cluster the rows of A. This gets us prototypes/centroids instead of “unique rows.” The fact that a true row isn’t identical to its centroid is what makes the algorithm lossy.

2. We’ll split A into a bunch of disjoint subspaces and cluster within each subspace. Clustering won’t capture much of the variance in a 1024-dimensional space, but it can capture lot of it in, say, a 4-dimensional space. We’ll have many low-dimensional clusters instead of a few high-dimensional ones.

When the clustering algorithm is K-means, this algorithm is called Product Quantization (PQ).

This works super well and is a foundational algorithm in information retrieval.

The code to create the centroids looks like this, and runs in an offline training phase before we receive our matrices:

def learn_centroids(
    A_train, # training matrix; e.g., activations for training inputs
    num_subspaces: int,
    num_centroids: int = 16) -> array:

    N, D = A_train.shape  
    subspace_len = D // num_subspaces
    centroids = empty(num_subspaces, K, D / num_subspaces)
    for s in num_subspaces:
        A_subs = A[:, (s * subspace_len):(s+1)*subspace_len]
        centroids[s] = k_means(A_subs, num_centroids)
    return centroids

At runtime, here’s how it encodes A:

def encode_A(A, centroids) -> array:
    N, D = A.shape
    num_subspaces, num_centroids, subspace_len = centroids.shape
    out = empty(N, num_subspaces)
    for s in num_subspaces:
        A_subs = A[:, (s * subspace_len):(s+1)*subspace_len]
        l2_dists = A_subs @ centroids[s].T - l2_norm(centroids[s].T, dim=0)
        out[:, s] = argmin(l2_dists, dim=1)  # N x 1
    return out

And how it encodes B1:

def encode_B(B, centroids) -> array:
    D, M = B.shape
    num_subspaces, num_centroids, subspace_len = centroids.shape
    out = empty(num_subspaces, num_centroids, M)
    for s in num_subspaces:
        B_subs = B[(s * subspace_len):(s+1)*subspace_len] # subs_len x M
        out[s] = centroids[s] @ B_subs  # 1 x M
    return out

to estimate dot products, product quantization just generalizes our single

out[i, j] = B[A[i], j]

to a sum over the different subspaces:

out[i, j] = sum([ [B[s, A[i], j] for s in range(num_subspaces)])

And that’s product quantization.

MADDNESS, the algorithm we’re accelerating, adds three optimizations on top of this:

1. It averages across subspaces at the instruction level instead of summing for better numerics.

2. It lets the centroids span across subspaces and jointly optimizes them using ridge regression, which makes the compression less lossy.

3. It uses probably2 the fastest locality sensentive hash function there is to assign rows to clusters, instead of the brute-force comparison to all centroids.

The third one is the most interesting. I won’t go over it but it’s essentially just learning a binary decision tree, assigning rows to leaves, and calling each leaf a cluster.

These changes let it approximate matrix products ridiculously fast even on current CPUs.

And this is without good hardware support. With better hardware support, it would of course be even better.

But is this algorithm conducive to hardware acceleration? Yes, very much so. This is because it:

1. Lets us do table lookups instead of multiply-adds. Multipliers scale quadratically in the number of bits involved and are pretty complicated circuits. Table lookups are just multiplexers, which are way simpler.

2. Compresses the matrices a ton, reducing memory usage.

3. Preserves the layout and access patterns of a dense matmul, including the tiling/data reuse. A MADDNESS kernel looks exactly like a dense GEMM kernel except that the matrices have different element sizes.

And now you understand the state of the art in approximate matrix multiplication!3

Hardware

So…back to the paper at hand. They add hardware acceleration to the centroid assignment function by loading the relevant indices in parallel and store-forwarding the left- vs right-child decision at each level of the tree. This gets them a throughput of one centroid assignment per cycle with a latency of four cycles.

They similarly accelerate the table lookups by using a multiplexer and a low-bitwidth adder. They tried fp16 table entries with fp32 accumulators, but found that int8 + int24 worked just as well.

They also introduce some algorithmic ideas to make the accuracy not terrible, since just naively approximating each layer in isolation goes badly. They:

• gradually replace layers from beginning to end,

• keep the first and last layers in fp16, and

• relax the hash function assignment to be differentiable during finetuning.

I’m not too calibrated on energy dissipation numbers, but here’s how much they use for all you hardware people:

So what does this all mean?

First, the potential power and area savings are real. An H100 uses 400W to do ~2e15 int8 multiply-adds per second, for a total of 5 TOp/s/W. Meanwhile, this paper got 43 TOp/s/W. That’s over 8x better despite the former using 4nm process technology (~2021) while the latter uses 14nm (~2014).

But no one has managed to do this without accuracy loss, and probably no one will4. The reason is that vector quantization / table-lookup-based approaches almost always operate in the high compression regime. You need small tables to fit in registers or L1 cache, which means each table is, in practice, either 16 or 256 elements. This in turn makes the table indices either 4 or 8 bits. If we compress one scalar into 4 bits, that’s already aggressive scalar quantization. But if we’re compressing multiple scalars into 4 bits—that’s a lot of compression.

Putting it all together, the way I think about this is that:

1. Vector quantization nearly guarantees better space efficiency than scalar quantization. This is because vector quantization can learn to just scalar quantize as a special case, but exploits mutual information across dimensions if it’s there.

2. Consequently, MADDNESS and friends are strictly better than binarization, in that they’re more expressive but no more expensive.5

3. However, you can only compress a lot with vector quantization, not a little. So you probably need to train with more parameters/tokens to hold accuracy constant.

4. Vector quantization currently lacks support in tensor cores and other hardware, so it’s only useful today on CPUs and for bandwidth-bound operations like matrix-vector products.

I have like 95% confidence that vector quantization will win in the end, since the information per bit is so much higher. But it will probably take more papers like this before it gets on the roadmaps of major accelerator vendors.

Exponentially Faster Language Modelling

Here’s another ambitious paper that incorporates some similar ideas. It’s largely a follow-up to their recent Fast Feedforward Networks paper that applies this technique to a Crammed BERT and improves some details.

The basic idea here is to do extremely granular mixture of experts, with small slices of weight matrix columns (or even single columns) treated as experts.

To make this work, they:

1. Arrange the experts into a balanced binary tree, with each expert corresponding to a leaf node.

2. Have multiple disjoint trees + sets of experts, and sum the outputs from each.

3. Do soft routing during training. This routes every token to every expert, but with a small routing weight for nearly all experts.

4. Randomly flip the routing weights during training. E.g., if a token would give 90% weight to one child node and 10% to the other, this is swapped.

5. Add a “hardening” loss to anneal towards hard expert assignment by the end of training.

Here’s the training and test-time pseudocode:

When finetuning a pretrained BERT on GLUE tasks for 24 GPU-hours, they get almost the same accuracy as the dense model, even when only 12 out of the 4095 output features get activated for each token.

Because feature-level sparsity gets you both smaller inputs and smaller outputs, this can yield enormous speedups. You can’t really use the tensor cores once the experts get small enough, but with only 12 neurons used per layer, you can overcome this constant factor with sheer sparsity.

We’ll need to see clean {training time, inference time} vs accuracy curves and independent reproduction before we can be sure this is a real win, so here’s a spectrum of ways I think this could play out.

In the most optimistic case, there’s so much activation sparsity in common networks that this granular approach to conditional computation is usually just a good idea (at least for inference). It’s effectively the MoE-as-block-sparsity framing, with higher sparsity prioritized over higher utilization.

A moderately optimistic take is that they’ve found some combination of:

• A better way of doing MoE routing, especially at large expert counts

• More evidence that increased granularity in MoE experts is helpful.

• More evidence that summing outputs from disjoint sets of experts is helpful.

And the pessimistic take is of course that, like the vast majority of papers, this won’t reproduce and/or won’t even look good once you understand the experimental details.6 The biggest question mark I see is that apparently their dense GPU baseline is their own extremely slow kernel. This is totally legitimate as a way to make the baseline implementation closely match their method’s implementation, but it does mean that the baseline is probably 10x+ slower than what you’d get normally.

In any but the pessimistic case, this has an interesting relationship to the vector quantization stuff above. Their differentiable binary trees could be exactly what we need to get tree-based hashing differentiable. Further, their finding that hashing input rows using multiple learned binary trees exactly lines up with what worked well in MADDNESS. Maybe just wishful thinking, but two groups arriving at similar solutions independently seems like mild evidence that there’s something there.

QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models

Okay, I want more papers like this.

What I like is that they exploit the fact that we often have an explicit decompression step between loading weights from RAM and using them in our ops. When this is the case, you don’t need your compressed representation to just be quantized, sparse, or factorized—it can be the result of an arbitrary lossy compression algorithm. As I’ve said before, this buys you way more degrees of freedom and enables much better compression.

The overall setup here is that they do GPTQ on batches of experts, with ternary quantization in the most aggressive case. They then run this ternary encoding through the rest of their compression pipeline.

This pipeline uses dictionary coding with fixed codewords. The codewords are formed by computing the empirical probability of a zero in the ternary encoding and materializing the 2^16 most probable sequences of codes assuming the elements are iid and the +1 and -1 ternary values are equiprobable.

This lets them encode many weights with a single 16-bit code. One subtlety here though is that they also need to handle sequences of values that aren’t among the most probable 2^16. They do this by having variable-length dictionary entries, with each entry containing a header that stores how many symbols it contains. As long as all codes of length at most k are present in the dictionary for some k, this lets you encode any sequence of ternary values.

On top of this encoding, they sort these 2^16 codewords by estimated probability to maximize the number of cache hits they get. They also wrote a fused decode + matvec kernel to avoid an unnecessary roundtrip to RAM.

To make all this work, they had to solve some other hard problems along the way. E.g., they have to fetch the relevant weights lazily from disk to avoid CPU OOMs for trillion-parameter models.

Another enhancement they introduce is ignoring special tokens like padding when estimating the Hessians for quantization. They also ignore the tokens immediately preceding the special tokens.

Combining all these ideas, they manage to get wall-time speedups for matvecs compared to bf16.

In terms of space, they get a ton of compression—even when only compressing the MoE layers.

This compression yields minimal change in C4 validation perplexity for these same large MoE models.

The post-training quantization literature is too devoid of apples-to-apples comparisons to know what the state of the art is, but I do really like how they surfaced a bunch of systems challenges and recognized that they’re solving a compression problem rather than a quantization problem per se. I’m personally skeptical of the aggressive scalar quantization → lossless dictionary coding pipeline, since in my experience, you tend to do better just lossy compressing with vector quantization (and VQ lets you do matmuls/matvecs even faster)—but that’s conjecture, not criticism.

Zephyr: Direct Distillation of LM Alignment

Instead of doing RLHF or supervised finetuning on LLM-generated instruction-response pairs, they a) have GPT-4 rank various LLM-generated responses and b) finetune their model to generate the highest-rank one instead of the others.

This reduces to a simple objective function with no separate RLHF reward model to train.

Applying their method to Mistral-7B lets it beat LLaMA 2 70b pretty consistently.

Their approach also does better than various other models and methods on both chat and academic benchmarks.

Importantly, these gains seem to actually come from their method, as opposed to lurking variables.

Probably just my confirmation bias, but seems like even more evidence that we don’t need all the complexity of RLHF.

Controlled Decoding from Language Models

Instead of finetuning an LLM to act differently or crafting a different prompt, they propose to use an auxiliary scoring model to bias the logits. This auxiliary model can be trained on other data (i.e., off-policy). This approach lets you steer the outputs on a per-input or per-user basis without having to engineer a prompt or keep around a ton of different PEFT weights/adapters.

I don’t know yet whether this is the right approach for a given problem, but it’s a design point I hadn’t really thought about before—effectively a generaliztion of constrained decoding to biased decoding. I could see a bunch of community-created scoring models pop up for a given tokenizer that you could plug into any model.

Efficient Numerical Algorithm for Large-Scale Damped Natural Gradient Descent

Do you need to solve an equation that looks like this?

If so, you’re in luck because there’s now a better algorithm for it. No inverses other than for a lower triangle matrix, which doesn’t really count.

I’m not sure what to do with this (maybe stick it in a second-order optimizer somehow?), but I’m always on the lookout for better building blocks.

Test-Time Distribution Normalization for Contrastively Learned Vision-language Models

Before using the dot product to measure similarity in embedding space for CLIP models at test time, you should subtract off the mean {image, text} embeddings.

This more closely resembles the training objective. Not doing this subtraction gets you a zeroth-order approximation of this objective, while doing it gets you a first-order approximation.

Adding this correction at test time consistently improves the accuracy for various tasks.

Obviously needs to be independently replicated, but these sorts of observations about something being broken are probably the methodological contributions I trust the most.

Future Lens: Anticipating Subsequent Tokens from a Single Hidden State

How well can you predict upcoming tokens given only the hidden states for the latest token?

The answer is: kinda sorta well? The last few layers are decent for the next token while the middle layers are a bit better for tokens after that.

I was hoping the answer would be “extremely well,” because that would be promising for skipping or early exiting decoding steps. But if we actually do need the semantics of generating each token serially, that’s a lot harder to speed up.

Transformer Memory as a Differentiable Search Index

Similar to the Neural Corpus Indexer, they just train a model to output the IDs of which docs are relevant instead of having a conventional embedding + search setup.

They manage to get awesome retrieval results out of this, although it’s unclear how the speed vs accuracy curves look.

I’m taking this as more evidence that training a model to spit out document IDs instead of doing conventional retrieval might be promising.

Fast Inner-Product Algorithms and Architectures for Deep Neural Network Accelerators

This paper is about computing matrix products with more adds but fewer multiplies, in the same vein as Strassen’s algorithm and Winograd convolution. I don’t understand why the algorithm works, but here it is:

The lets them do an MxN = MxK @ KxN matmul using just (MNK + MK + NK)/2 multiplications instead of the typical MNK multiplications.

They also talk about how this method enables more efficient matmul acceleration in hardware, and how it works well with neural net inference. In the inference case, you can precompute the β_j terms and fuse the b_ij terms into the biases.

I don’t know enough to comment on the hardware aspects, but apparently their method saves a bunch of registers and makes it easy to raise your clock frequency compared to previous approaches.

If future hardware could use 2x fewer multiply units for the same matmul throughput, that would certainly be a big area and power win.

Simplifying Transformer Blocks

They found that you could rip out a bunch of components for BERT models up to 300M params while preserving time-to-accuracy on GLUE. In particular, they remove:

1. Skip connections

2. LayerNorms

3. The projection matrix for the values in the attention

4. The linear after the attention

They also put the FFN and attention blocks in parallel.

To make this work, they bias the attentions such that the attention module is an identity function during initialization.

Other than that, they evaluate it with essentially the same hparams as a baseline model and find that the time vs loss is a bit better.

I’m surprised they got rid of the values projection matrix instead of either the query or key projections, because the latter two are redundant7. But regardless, these sort of simplifications will be a nice win if others can reproduce them.

Jorge: Approximate Preconditioning for GPU-efficient Second-order Optimization

What we’d often like to do in gradient-based optimization is Newton’s method, which jumps straight to where the minimum should be under a local quadratic approximation of your loss landscape.

But the actual Hessian matrix is of size num_parameters x num_parameters, which is too big and expensive to compute. So we approximate it by pretending it’s diagonal (c.f. Adam and friends) or, sometimes, with a more algorithm.

This paper introduces an inverse Hessian approximation that’s cool because it doesn’t require computing any matrix inverses. Similar to Shampoo, it pretends the Hessian matrix decomposes into separate matrices for each layer and maintains an approximate Kronecker factorization of the Hessian. But instead of just doing this via a moving average of gradient matrix products, it does updates that look more like computing the elements of a Neumann series.

I (and presumably many others) have tried to avoid the inverse here using Neumann series, but it’s hard to get around the constraint that the preconditioning matrices have norms less than one. They solve this by dynamically adjusting the β parameters based on the current Frobenius norms. Computing these norms adds some overhead, but less than any of the matrix products.

With this constraint addressed, the derivation is pretty straightforward. They just observe that you can rewrite the Shampoo update using a recurrence relation with an auxiliary matrix,

and then get rid of the matrix exponent using the binomial theorem. And…that’s it. It’s as simple as you’re ever gonna get in an optimization paper.

They get better accuracy than SGD and Adam across various benchmarks by just porting the hparams from SGD via simple conversion formulas. They also only update the preconditioning matrices every 50 iterations, so that the overhead is small enough to stay within ~10% of SGD’s runtime.

If you tune the length of the training to exploit the faster convergence, their method can improve time-to-accuracy significantly.

I’m unusually optimistic about this optimizer. There will be a bunch of questions about how to get it to scale when parameter matrices are sharded, but I buy that a) second-order optimizers can be a win if you get the systems side right, and b) the simplest-possible second-order optimizer without matrix inverses should be especially promising.

SpecTr: Fast Speculative Decoding via Optimal Transport

They do faster speculative decoding for low-batch-size LLM inference without messing with the sampling distribution.

If you’re not familiar with speculative decoding, it basically just exploits the observations that:

1. LLM inference at low batch sizes gets horrible utilization, but

2. you get better utilization if you generate several tokens at once.

The problem is that you normally can’t generate multiple tokens at once because each token generation needs to condition on the previous output. So what speculative decoding does is:

1. generate a sequence of proposed tokens from a small helper model, and

2. have the main model generate the conditional distributions for all these tokens at once as if the earlier proposed tokens were part of the output.

You then have both the true model and helper model’s conditional probabilities for each token and you can rejection sample to get the same math as just using the main model. If the helper model’s distribution matches the main model well, you reject infrequently and you can keep a long prefix (or all) of the proposed tokens.

That’s regular speculative decoding. What this paper asks is, “what if we sampled multiple proposed sequences from the helper model at once?”

They do a ton of math to make this work and prove stuff about it, but essentially they get to accept longer prefixes on average while still preserving the sampling distribution.

I’d have to stare at this a lot longer to really understand it, but this looks like a win in the low inference batch size regime.

The Synergy of Speculative Decoding and Batching in Serving Large Language Models

They point out that the amount of speculative decoding you try to do should vary based on how many queries you have batched at a given moment. This is because running forward passes for more tokens has increasing marginal cost as you approach being compute-bound.

Exploiting this observation can reduce your average latency a bit for certain traffic patterns.

I’m actually surprised how little difference this makes, but that’s not a criticism. Thorough experiments with trustworthy conclusions are much more important than how much the number changes. Plus this paper improved by thinking about what the ideal inference system should do.

Position Interpolation Improves ALiBi Extrapolation

Normally AliBi linearly decays the attention between tokens in proportion to the absolute difference in their positions within the sequence. When trying to get a model to use a longer sequence length than it was trained on, they propose to instead decay attention values based on the relative difference in token positions, where “relative” means normalizing by sequence length.

This simple change can improve length generalization significantly.

Towards End-to-end 4-Bit Inference on Generative Large Language Models

They do 4-bit inference with 4-bit matmul kernels that quantize and dequantize before each matmul. This is more aggressive than typical 4-bit quantization that just keeps the weights in 4 bits at rest and upcasts to bf16 before doing any math.

The quantization is mostly just regular GPTQ / linear quantization, but they handle outliers intelligently:

• First, they use a calibration set of samples to identify which columns have outliers features and handle these separately.

• Secondly, they rearrange the weight matrices so that these columns are contiguous and get quantized last. The “last” aspect is important because GPTQ tends to accumulate more error in later columns; but if these columns aren’t getting quantized at all, we dodge our worst-case behavior entirely.

• Third and finally, they use 8-bit quantization for the down projection linears since these have the most variance.

This approach manages to get a 2.5x speedup vs fp16 inference on LLaMA 2 70B with 8-bit quantization and a 3x speedup with 4-bit inference. These are pretty close to the ideal speedups. They do need pretty large matrix products to hide the quantization and dequantization overheads though.

They don’t quite preserve the accuracy of an fp16 model, but they get closer than various baselines.

Microscaling Data Formats for Deep Learning

Regular old linear quantization with 8-bit power-of-two scale factors just kinda works during training.

By “just kinda works,” I mean you can quantize all the matmul inputs with ~0 accuracy loss.

Basically, if we had hardware support for block-specific scale factors, we could just do everything with 8-bit (or even narrower) formats. This isn’t news for inference and is weaker than the unit scaling claims for training, but it is new to see that simple quantization just works out of the box with no special math or initialization.

FP8-LM: Training FP8 Large Language Models

They do a bunch of careful numerical changes to get more of the training process to work in fp8. First, they track tensor-specific scale factors for the gradients to keep them from overflowing or underflowing. When they detect an overflow, they scale down by a factor of two, and otherwise gradually increase the scale slightly at each iteration.

Before allreducing gradients, they allgather the scale factors and have every device use the smallest scale factor. They then undo this scaling before feeding the gradients into the optimizer. This lets them allreduce (and allgather) in fp8, saving 2x communication vs bf16.

They also save some space in the optimizer. Instead of having 4 bytes for each of the {master weight, gradient, momentum, variance} tensors in Adam, they use fp8 gradients, fp8 momentums, 16-bit variance estimates, and 16-bit master weights. Importantly, the 16-bit values have associated scale factors, rather than just using the default fp16 range. You can do math to find the effective step sizes at which 16-bit master weights will break (i.e., a learning rate that yields steps around 2^-(num_mantissa_bits + 1) as large as the weights will go badly; even 24-bit master weights are often sketchy for finetuning), but this is a nice savings if your steps are large enough.

They also propose to shard entire tensors across devices instead of slices of tensors. This scares me from a complexity and interconnect utilization perspective, but does give a clean story for how to handle all the scale factors.

Even with all these space savings, they still get relatively little model quality loss compared to using bf16 for everything.

They also make training go ~30% faster, which likely yields a time-to-accuracy improvement:

It looks like mostly-fp8 training is going to become a standard practice within the next year or so, at least for the big industry labs.

Emergent Mixture-of-Experts: Can Dense Pre-trained Transformers Benefit from Emergent Modular Structures?

They use the observation that a set of sparsely-gated MoE FFNs can be seen as a big activation-sparse matmul, similar to MegaBlocks. But instead of using this to define an MoE module at the start, they train a regular network and then approximate it using an MoE router. To do this, they cluster the columns of the first linear, use the centroids of each cluster as routing weights, and then only apply the rows from the second matrix associated with the chosen cluster’s features. You can also look at this as doing an approximate matrix product using vector quantization.

If you apply this technique to a pretrained model before finetuning, you can get more accuracy lift than both the dense model and at least one conventional MoE baseline.

This is especially interesting to me because I’ve gotten related negative results. Namely, if you add the inner products with the routing weights to the outputs (so that you’re doing residual vector quantization instead of their regular vector quantization), it’s not quite as good as Switch-style MoE routing. But this was with the MoE routing there from the start instead of added at the end. It wouldn’t be surprising if imposing the block sparse structure at initialization isn’t as good as “pruning” to a block sparse structure, even if the pruning happens all at once.

In-Context Pretraining: Language Modeling Beyond Document Boundaries

Why concatenate random documents when you can concatenate related documents?

Turns out this improves pretraining, as measured by eval perplexity on randomly concatenated documents after 300B pretraining tokens.

It also improves downstream metrics, including making retrieval augmentation work better sometimes.

Closed-Form Diffusion Models

I’d need to stare at this a lot longer to grok it, but they devised a decent diffusion-based sampling procedure that needs no training, works pretty well, and runs faster on a CPU than a neural net-based diffusion model does on GPU.

They also derived a closed-form solution for the optimal diffusion procedure under Gaussian noise in terms of test error, and argue that most diffusion models generalize because they have a smoothness bias and don’t recover this optimal solution.

TLM: Token-Level Masking for Transformers

They propose to randomly drop whole tokens in each attention head.

This can work better than other attention masking schemes.

You probably don’t want to use more than one of these schemes though. As is common with regularization, you get diminishing returns from composing multiple methods.

Battle of the Backbones: A Large-Scale Comparison of Pretrained Models across Computer Vision Tasks

They did a big ol’ bake-off between different pretrained vision models. They summarize their main findings well so I won’t repeat them:

They also have recommendations for practitioners regarding which models you should probably use at different sizes and for different tasks.

Matryoshka Diffusion Models

Their U-Net takes in multiple resolutions at once and pools information across them to denoise them all jointly.

This approach, along with focusing much of the training budget on lower resolutions, gets them a much better training steps vs image quality tradeoff than various baselines (theirs is “MDM”).

CapsFusion: Rethinking Image-Text Data at Scale

They observe that raw image captions are often noisy, but synthetic captions are vague and don’t capture world knowledge that may have been in the raw captions.

To get the best of both worlds, they propose a pipeline to combine raw and synthetic captions using ChatGPT. To avoid querying ChatGPT for their entire dataset, they just query it a million times and then finetune a LLama model to combine captions using the ChatGPT outputs.

Their improved captions yield significantly better models, as measured across a variety of benchmarks.

FlashDecoding++: Faster Large Language Model Inference on GPUs

They do faster transformer inference vs various baselines by making a few changes to the decoding logic.

First, they use a single max exponent for the whole softmax instead of adaptively increasing it as successive tiles get processed. This exponent has to be set carefully, but you avoid some data dependencies this way. To ensure they preserve the math, they fall back to using a larger max if they detect an overflow.

To set this maximum value appropriately, they look at the empirical distributions of the softmax inputs.

They also get better utilization by writing kernels for matvecs (and skinny matmuls) that increase occupancy via smaller tile sizes + by software pipelining loads into a separate buffer.

They choose whether to use their kernels or other kernels at runtime based on some heuristics.

These optimizations together apparently yield large speed gains for both prefill and decoding.

ToolChain*: Efficient Action Space Navigation in Large Language Models with A* Search

Given some means of scoring intermediate and final outputs when iteratively querying an LLM, we can use classical AI algorithms like A* search to arrive at better answers.

This can work way better than other tree-based approaches, especially when you take into account runtime.

Mostly I think this is super interesting on a meta level, since it shows that:

1. Classical AI algorithms are getting relevant for mainstream AI again.

2. We’re having more success by carefully structuring our programs around LLMs, as opposed to, say, offering as little structure as possible and just having the LLM figure it out. This is good from a safety and interpretability perspective.

This was a lot of fun and I’m definitely down to do more podcasts.

Side note: I’m super paranoid about exaggerating my similarity search contributions in this conversation, so see this footnote if you care.1

xVal: A Continuous Number Encoding for Large Language Models

They encode numbers as a single token, but with the embedding magnitude scaled by the number’s magnitude. At the decoder, they read the embedding scale whenever the token is predicted to be a number.

This approach often works better than alternatives for various regression tasks, despite adding fewer parameters to the embedding tables.

They also have a good discussion of limitations. You shouldn’t use this encoding when you need a huge dynamic range, or when you need many digits of precision. Although appropriate pre- and post-processing of numbers could help mitigate these limitations.

Maybe this is how we should all be handling numbers in LLMs?

Sparse Backpropagation for MoE Training

They did math to come up with better MoE routing, where “better” means higher accuracy in fewer training steps.

The math/notation is pretty dense and their Github repo is private, so here’s what I think is happening; apologies if I mess it up.

First, you can think of the MoE layer as “sampling” a subset of the experts. The sampling might happen to be a (noisy) top-k in practice, but it’s still sampling. So we want to minimize the expected loss over this sampling distribution.

Applying the product rule to the rightmost expression, we see that the gradient of the loss can be written as two summations. You can think of the first summation as requiring us to differentiate with respect to the sampling probabilities, and the second as differentiating with respect to the routing weights conditioned on the sampling distribution.

They observe that computing the second summation with backprop works great, but backproping through a sampling step is sketchy. So what most MoE models currently do is just ignore the first (sampling) term of the true gradient.

Before getting to their solution, let’s think about what we’d like to do. What we’d like to know are the outputs of all the experts, and therefore how the expected MoE layer output changes as we increase the probability of sampling a given expert.

But computing the average output under the current distribution requires us to route every token through every expert, which kinda defeats the point of sparse MoE. So instead of using the true average output as the baseline, we could instead just use all zeros as our baseline. With a lot of experts and no nonlinearity at the end to add bias, this is a sane approximation of the true average anyway.

This lets us approximate the gradient using only the experts we actually select in our top k, instead of all of them. But it still requires us to run an all-zeros activation tensor through the rest of our network. And ain’t nobody got time for that. So we’re going to introduce a first-order approximation to this difference: we’ll compute the slope (gradient wrt the actual expert output, provided by backprop) and scale it by the distance (expert output minus zero = expert output).

You can do a little bit better by computing the gradient halfway between the expert output and zero, instead of all the way at the expert output. In practice, they just halve the expert output in forward and double the gradient in backward.

This halving introduces a train vs test discrepancy though, so they only use the midpoint method…during training, maybe? They might be varying whether they sample from the softmax distribution or use hard top-k during training.

They also mask out all but the top-k entries before computing the softmax and add a trainable, token-independent scale factor for each expert.

This masking is apparently really important. So are the learnable scales and only sometimes using the midpoint version, although those ablations are shown in the figure at the top.

In terms of wall time, their method adds no measurable overhead, which is awesome.

This might be an immediate and large win for anyone training MoE models.

GrowLength: Accelerating LLMs Pretraining by Progressively Growing Training Length

Progressive resizing helped us set the ResNet-50 time-to-accuracy record, and it turns out it works for LLMs too. At least for sub-billion parameter models.

It’s not as good as training at the full sequence length for a given token budget—but it is better for a fixed time budget, presumably because their workload is spending a lot of time in the attention ops.

Training and inference of large language models using 8-bit floating point

They show that simple scaling of the linear layers lets you convert them to fp8 with ~0 loss of model quality. The scaling is just coming up with a tensor-wide scale factor, and they explore both scaling adaptively for each tensor and via one overall exponent bias.

Apparently you can use the same exponent bias across all the weights and activations and preserve 99.5% of the accuracy across a pretty wide range of exponent values. The exact range varies by task, but just setting it to zero worked in every case; this held for both finetuning and inference.

Graphcore’s been crushing it with fp8 work lately. Also, they actually trained these models on their hardware, which is impressive.

Soft Convex Quantization: Revisiting Vector Quantization with Convex Optimization

They do better vector quantization in a VQ autoencoder. In the forward pass, they try to get better (soft) centroid assignments by solving a convex optimization problem that treats the centroids as fixed. In the backward pass, they use the implicit function theorem + the fact that the solver converged to get gradients wrt the centroids.

This seems to work better than simpler vector quantization for a few image tasks.

Pretty cool application of implicit layers. Also might help us improve vector quantization as a model compression technique in general.

A Long Way to Go: Investigating Length Correlations in RLHF

Turns out RLHF is largely just telling models to produce longer outputs.

You can even get a lot of the RLHF benefits by just plain rewarding length.

LightSeq: Sequence Level Parallelism for Distributed Training of Long Context Transformers

First, they observe that standard gradient checkpointing plays poorly with Flash Attention. The problem is that Flash Attention throws away its intermediate outputs even if you’re already in the middle of rematerializing the attention block (when you’d rather keep the intermediate outputs to save compute).

“flash attention will not materialize the intermediate values during the forward, and will recompute it during the backward, regardless of the re-computation strategy in the outer system level.”

To mitigate this, they checkpoint the Flash Attention outputs. This means they don’t have to run another forward through Flash Attention, and therefore only end up rematerializing its intermediate outputs once.

This can improve training speed significantly at long sequence lengths.

Second, they propose to use sequence parallelism for everything. This is nice in that it lets you do everything but the attention locally, and do tensor-parallel attention with more devices than attention heads. But it requires some complex compute + communication scheduling logic within the attention computation.

Their changes let them support much longer sequence lengths on a given GPU than Megatron,

and train faster in many cases.

This made me spend a bunch of time thinking about the optimal way to parallelize attention, which is great. The improved Flash Attention checkpointing also seems like an easy win for anyone doing large-scale transformer training.

The Cost of Down-Scaling Language Models: Fact Recall Deteriorates before In-Context Learning

If you reduce the parameter count in an LLM, it tends to lose recall of facts before it gets worse at learning from examples in the prompt. This holds for parameter count reductions via both pruning and using a smaller dense model.

Mistral 7B

They highlight interesting aspects of what’s likely the best 7 billion parameter model right now.

The main thing they talk about is adding a sliding window restriction to grouped query attention. This just means that each token can attend only to the W-1 previous tokens for some constant W.

This restriction enables some cool systems optimizations, including allocating a fixed-size circular buffer for the KV cache.

It also lets them skip a lot of the regular N^2 computation during prefill (i.e., processing the prompt in parallel before sequential decoding begins).

They also talk a bit about their system prompt and how they can query the model itself to determine whether it should refuse to answer a given prompt.

They don’t really say anything about the data or how the model was trained, but apparently it beats LLaMA 2 13B across a variety of benchmarks.

There’s not much detail here besides their system prompt and the sliding window attention, but at least that’s something—it’s certainly more than what many of the big labs are releasing about their models.

Take a Step Back: Evoking Reasoning via Abstraction in Large Language Models

They propose to have language models ask themselves (and answer) a more general question than the prompt, and use the results to refine their initial answer.

This apparently increases accuracy a lot across a variety of knowledge and reasoning tasks.