Three Layers of Groundwork for an ML Researcher: Starting from Kimi's Attention Residuals

I've been thinking about what foundational abilities a competent ML researcher should have. Without a solid foundation, talk of innovative ideas, good taste and the like is a castle in the air.

Let's look at frontier work and see what parts a good piece of modern ML architecture research contains, what abilities we need in order to understand it, reproduce it and even build on it, and how to develop those abilities.

Three layers

A complete piece of modern ML architecture research usually has three layers:

1. The model-definition layer (what to compute): what change does it make to the architecture, and why does that make the model better?
2. The algorithm layer (how to compute it): keeping the result the same, how can the computation be reorganized to need less compute and fewer memory reads and writes?
3. The system layer (where to compute it): how do you map the computation onto concrete hardware and clusters, and solve the problems of memory, communication, parallelism and deployment?

The three layers are also judged by different measures: the model layer by quality (loss, benchmarks), the algorithm layer by theoretical cost (compute, memory traffic), and the system layer by measurement (throughput, latency, memory footprint, cost). The boundaries are not always clear-cut either: work like FlashAttention spans the algorithm and system layers at once.

Coming up with a good new idea is often only the first step. An idea that works on a small model is not thereby trainable on a large cluster, or able to serve users online at low latency; those are hard engineering problems for the algorithm and system layers to solve (which shows how much a CS foundation matters).

Below I use Kimi's Attention Residuals (AttnRes) as an example of how these three layers show up in a concrete piece of architecture research, and then turn round to how one might learn the abilities each layer needs.

Kimi's Attention Residuals as an example

The idea is actually clear and easy to grasp: Attention Residuals (AttnRes for short) replaces the residual connections between a Transformer's layers, changing them from "add up every layer's output with equal weight" to "do softmax attention along the depth direction", much as the Transformer improved on the RNN along the token direction.

Layer one: the model definition

The basic motivation: in the common PreNorm structure, the hidden state keeps growing with depth, so each layer's output makes up a smaller and smaller share of it, and the information from early layers gets drowned out.

Starting from this idea, the Kimi team built Full AttnRes; then, to save memory and cross-GPU communication (an engineering-driven idea), they built Block AttnRes; and they brought in similar baselines such as DenseFormer and mHC for comparison.

Once the idea is proposed, you still have to show that it really works: scaling laws to show the gain holds across scales, and ablations to show which component is actually doing the work.

Layer two: the algorithm

Now consider: how can the new architecture be made to compute faster?

For example, the queries of all the layers in a block first compute the inter-block attention in one batch, then each layer handles the intra-block part in turn, and the two are merged with online softmax. (Inspired by FlashAttention: the result is exactly the same as computing it directly; only the order of computation changes.) The batching is possible because each layer's query is a learned fixed vector (the paper calls it a pseudo-query) that doesn't depend on the input, so they are all known before the computation starts.

Then you measure what the improvement is actually worth: the paper first works out, in theory, the memory traffic per layer for each residual mechanism, and then measures the end-to-end latency overhead (see layer three).

Layer three: the system

On the system side, what can be improved?

1. Training

Large-scale training uses pipeline parallelism: the model's layers are cut into several segments that are assigned to several GPUs and computed in sequence; each segment is called a stage. Block AttnRes has to pass the accumulated block representations (the sum of the outputs of the layers within each block) on to later stages. The naive approach passes the whole history at every hand-off, so communication grows with the square of the number of stages; and these block representations have to be kept until the later layers are done with them, which also increases memory.

The fixes:

• Cross-stage caching: in an interleaved pipeline, data goes round the GPUs several times, and each GPU processes the same data several times over. The block representations received the first time stay local, and afterwards only the new part is sent; the amount sent each time then no longer grows with the number of stages, and can be hidden behind computation.
• Memory: activation checkpointing stores only each layer's input, and the intermediate results are recomputed during backpropagation, trading one extra computation for less storage.

Result: for the same model, how much slower does training get when the standard residual is replaced with AttnRes? Without pipeline parallelism, negligibly; with it, under 4% end to end. The end-to-end result shows this is not a fanciful improvement but one that can actually be trained in practice.

2. Inference

Long-context prefill (reading in the whole prompt at once) is the most memory-hungry part: every token has to store the 8 preceding block representations (8 being the number of blocks Block AttnRes splits the network into), which for 128K tokens comes to about 15 GB. If the merge in the second phase is computed separately, it also costs an extra round of memory reads and writes.

The fixes:

• Splitting along the sequence: attention along depth only happens among a single token's own layers, and different tokens don't interact. So the tokens can be split up across the 8 GPUs of tensor parallelism (cutting each layer's matrices up so several GPUs compute them together), with each GPU storing only its own share, which brings it down to about 1.9 GB per GPU; then the prompt is processed in 16K chunks (chunked prefill), bringing it below 0.3 GB.
• Folding the merge into communication: tensor parallelism already does a reduce-scatter and an all-gather in every layer, and between those two steps each GPU holds exactly a share of the tokens. Use that gap to do the merge and RMSNorm locally, with no extra communication.

Result: for the same model, how much slower does inference get with AttnRes? Under typical inference workloads, end-to-end latency rises by less than 2%.

How the three layers connect

Constraints from the lower layers often turn round and decide the design of the upper ones.

We can start from two perspectives. One is the perspective of model capability (Model), whose objective is the quality of the model, measured by loss, benchmarks, behaviour and so on; the other is the perspective of systems engineering (System), which is about how to serve users faster at inference, how to iterate on experiments faster and so on, measured by latency, throughput, memory cost and so on.

With only the model view, you end up with weights that perform very well but can't meet the latency, memory or cost requirements; with only the system view, you end up with a very efficient model that can't learn enough.

In AttnRes there are three places where the system decides the model:

• Block AttnRes itself: under pipeline parallelism Full AttnRes has to pass O(Ld) layer outputs, which today's interconnect bandwidth can't carry, and that is why blocks exist.
• Pseudo-queries don't depend on the input: this is what makes the batched computation in layer two possible. In the ablations, input-dependent queries actually did better, but they weren't adopted.
• The number of blocks is fixed at about 8: Full AttnRes works best, and finer blocking comes closer to it; the authors say plainly that fixing N≈8 is for infra efficiency, and that they'll move to finer granularity once the hardware loosens up.

The same structure shows up in academic work too. Depth Memory (LDM), proposed after AttnRes by another, academic team, can use exact recursive least squares (RLS) along depth because layers are computed sequentially anyway, so it doesn't need the parallelizable approximations used along the sequence; and RLS happens to be invertible, so during backpropagation the state can be reconstructed backwards from the last layer, saving per-layer storage.

But the system layers of the two papers answer different questions. Kimi is an industrial team validating on a 48B MoE with the aim of using it in production models, so it cares about pipeline communication, long-context prefill and deployment latency; Depth Memory comes from an academic group, which cares about single-GPU kernel speed and activation memory. What question the system layer answers often depends on where the researchers stand.

So when reading an architecture paper, it's worth asking of every design choice: is this trade-off for quality, or for the system? Systems engineering ability is absolutely indispensable; many designs are driven by engineering.

(It brings to mind a popular saying: "teaching a researcher how to do good engineering is far harder than teaching an engineer how to do good research.")

How to learn the three layers

As a researcher you don't need to go deep on all three layers; for each one, it's enough to be able to read it (to know why things are designed the way they are) and to use the corresponding simple frameworks. You can also choose one area to study in depth, and collaborate with friends who have different strengths.

1. The model-definition layer

(a) The abilities needed

• Know the components of modern models, and what problem each design solves and at what cost: attention and its variants (MHA, GQA, MLA, sparse and linear attention, state-space models), where normalization goes, residuals and the flow of information along depth, MoE, positional encoding.
• Understand the regularities of training: scaling laws; training stability (initialization, learning rate and warmup, loss spikes); how hyperparameters transfer across scale.
• Read papers closely: reconstruct the authors' argument, and distinguish "what the authors claim" from "what the experimental data actually support"; and put yourself in their place: if you were the author, how would you argue and experiment?
• Design rigorous experiments: controlled variables, ablations, multiple seeds, trends across several sizes, tuning each method separately, so that small-scale experiments give conclusions you can trust.
• Find good problems: look where the data and the received wisdom don't match (where an author's narrative doesn't match your own experiments, where two papers' data contradict each other, and so on), or find an entry point in an analogy from another field.

(b) Resources

• Courses: Stanford CS336 (Language Modeling from Scratch); the architecture and pre-training parts of Berkeley's Scalable AI.
• Classic code: compact implementations like nanoGPT; models like OLMo whose code, data and training logs are all public.
• Classic papers: the original Transformer paper; the Chinchilla scaling law; technical reports of recent open models, such as DeepSeek-V3 and Qwen3.

2. The algorithm layer

(a) The abilities needed

• Estimate cost: for a given operator, estimate the compute and memory traffic, and judge where the bottleneck is.
• Reorder computation: use associativity, blocking, and conversion between recurrent and parallel forms to find algorithms with the same result at lower cost. Classic examples are online softmax and the chunked parallel algorithms for linear attention.
• Numerical linear algebra and numerical stability: matrix factorizations, least squares, low-rank structure; knowing where things overflow or lose precision at low precision.
• Translate a model change into an algorithmic problem: can a new structure be parallelized? Can it be blocked? Can the intermediate results go unstored?

(b) Resources

• Papers: FlashAttention 1 and 2, the model of hardware-aware algorithm design; Mamba-2's SSD paper (Transformers are SSMs), a unified view of linear attention and state-space models.
• Blogs: Songlin Yang's DeltaNet series; Su Jianlin's "Scientific Spaces" (科学空间).
• Textbook: Trefethen and Bau, Numerical Linear Algebra.

3. The system layer

(a) The abilities needed

• Understand the GPU performance model: the memory hierarchy (HBM, SRAM, registers), tensor cores, the roofline, memory-bound versus compute-bound, kernel fusion.
• Do the memory accounting for training: how much the parameters, gradients, optimizer states and activations each take; how recomputation trades time for space.
• Understand distributed training: data parallelism and FSDP, tensor parallelism, pipeline parallelism, expert parallelism; what each has to send and what the communication costs.
• Understand inference systems: prefill and decode, the KV cache, batching and scheduling.
• Understand numerical precision: the trade-offs of BF16, FP8 and FP4.
• Use the tools: find bottlenecks with a profiler; write simple kernels in Triton when needed.

(b) Resources

• Intuition to start with: Horace He's Making Deep Learning Go Brrrr From First Principles.
• Large-model systems: How to Scale Your Model; Hugging Face's Ultra-Scale Playbook; the docs and source of vLLM and SGLang.
• Hands-on: the official Triton tutorials; the GPU MODE community lectures; Programming Massively Parallel Processors.
• Courses: the performance, parallelism and inference parts of UC Berkeley's Scalable AI; UC Berkeley CS 61C for the basics of computer architecture.

As for order, I'd suggest first building a basic understanding of ML, starting with the reading of papers and running of experiments in layer one, and when you find gaps in your knowledge, going selectively to good resources. Through papers and good courses, build up the performance intuition of layer three bit by bit, then go back to the algorithmic derivations of layer two: these two layers always alternate.

A practice project that runs through all three layers

The best way to learn is to tie the three layers together with one project, each step landing in one of the layers. You needn't design it from scratch; existing course assignments are very good:

• Stanford CS336's A2: implement FlashAttention in Triton, then do distributed training, which covers exactly the algorithm and system layers.
• UC Berkeley Scalable AI's A1 (performance engineering): the assignment package is open on the course homepage, and covers the basics of performance analysis.

Once you've done them, pick an architecture paper you're interested in and take one step in each layer: first implement it on a small model in PyTorch (model layer), then find the bottleneck with a profiler (system layer), and finally try reordering the computation or writing a fused kernel (algorithm and system layers).

Musings

The model layer decides what to do; the algorithm and system layers decide whether it can be scaled up and put into production. Good architecture research is the result of going back and forth between these three layers. The first layer takes the ability to model and experiment; the latter two have to do with MLSys.

But, but: going back to where it all began, what matters most is still the mindset. Do you truly love this field? Are you willing to slow down, quiet down, and move forward one step at a time?

ML research isn't easy, and the foundations have to be laid slowly, a little at a time. I recently read the "Slow Research" joint statement launched by professors, and I agree with it deeply. Yet nearly everyone is anxious now, and the rapid progress of AI works like a magnifying glass, making that anxiety even bigger. For me, the core of the anxiety is age: I always wish I could achieve something a little sooner.

Doing ML research may be, for me, what I was destined for: it is one of the callings I'm willing to pour my whole life into. After all the twists and turns, people do find their way back to where they belong. In my second year of high school, I remember discussing maths after class with a very close friend; we were both learning linear algebra then. Once we found out about machine learning, we said that after the gaokao we would definitely do research in it. It's a strange and wonderful thing: from the first time I came across the field, it has drawn me in, right up to today. Three years of undergrad went by, and I tried full-stack development, blockchain, quantitative investment research, even starting my own projects and building a brand: a long way round. Only in my fourth year did I come back here again.

I'm willing to take it slowly, because from the point of view of a life's work, waiting an extra year and getting there a year early make little difference. And what I love is doing research itself: being able to learn new things at the frontier every day and to think about them makes me truly, wonderfully happy.

But slowing down still asks you to show results; the way the real world judges you is deeply entrenched. On top of that, everyone has to face worldly trade-offs as they get older, the youth of your twenties is gone in the blink of an eye, and the pressure of the social clock forces us to keep our feet on reality, concrete and rough as it is.

I think that, strategically, you should raise your iteration rate as far as you can: while looking after your body and your life, make your learning and experiments more efficient, and focus your time on the right direction. When it comes to the research itself, though, you need to take it slowly and do things well, solidly :)

References

• Attention Residuals (Kimi Team, arXiv:2603.15031)
• Depth Memory (arXiv:2610.18809)
• Scalable AI, Fall 2026 (UC Berkeley)

(2026.10.16, Berkeley)