A New Trick Slashes the Cost of Shrinking Giant AI Models
Researchers have found a way to train smaller, cheaper AI models from giant ones using a fraction of the memory previously required, opening the door to experiments that once needed a warehouse of hardware.
Key points
- Knowledge distillation, copying a large AI model's behaviour into a smaller one, typically demands hundreds of specialist chips and careful coordination.
- A new paper caches the big model's outputs once upfront, so it never has to run again during training, cutting memory use significantly.
- A second technique, called a fused chunked KL loss, processes training data in small slices instead of all at once, keeping peak memory almost flat even at very long inputs.
- At 32,000 tokens of context, peak memory falls from 85.2 gigabytes to 5.45 gigabytes, a 15.6-times reduction.
- The full code for the memory-efficient loss function is publicly available on GitHub.
The biggest AI models in the world are almost impossibly large. Kimi-K3, released earlier this year, has 2.8 trillion parameters (think of parameters as the millions of tiny dials inside a model that are tuned during training) and needs roughly three terabytes of specialised memory just to load. That is more than twenty high-end consumer graphics cards' worth of space.
Because running models that size is so expensive, researchers routinely compress them into smaller versions that cost less to operate. The standard method is called knowledge distillation: you run the big "teacher" model on a pile of text, record how it would respond, then train a smaller "student" model to copy those responses as closely as possible. The student ends up much cheaper to run while keeping most of the teacher's quality.
Why is this so hard to do cheaply?
The expense comes from two places. First, the standard approach keeps both the teacher and the student loaded in memory at the same time. Second, comparing their outputs requires building an enormous grid: one row for every word in the model's vocabulary (often more than 100,000 words), one column for every position in the input. Multiply those together across a long piece of text and you quickly fill every chip you have.
As a concrete example: one popular open model, gpt-oss-120b, has a vocabulary of 201,088 words. At a sequence length of 32,000 tokens (roughly a short novel's worth of text) and a modest batch size of four examples at once, the teacher's output tensor alone takes up about 50 gigabytes. Add the student model, gradients and everything else the training process needs, and a single training step can spike to roughly 250 gigabytes, more than an H200, one of the most powerful AI chips available today, can hold.
What did the researchers change?
Two things, working together. Both are described in a paper published by the team at CompactifAI, first shared on Hugging Face.
The first change is offline distillation. Instead of re-running the teacher at every training step, the researchers run it once and save only its top 100 most likely word predictions per position. That cache is written to disk. The teacher is then switched off completely, and the student trains against the saved cache alone. Because the cache can be reused across many different experiments, the expensive teacher run becomes a one-time cost.
The second change is the fused chunked KL loss. KL divergence is a mathematical measure of how different two probability distributions are, here, how different the student's word predictions are from the teacher's. The standard way to compute it builds that full vocabulary-times-sequence grid in one go and holds all of it in memory. The new method slices the sequence into small chunks, calculates and discards each chunk before moving to the next, and recomputes what it needs during the backward pass rather than storing it.
| Method | Peak memory (8K context) | Iteration time |
|---|---|---|
| Online distillation | 102.8 GB | 25.9 s |
| Offline, dense KL | 78.3 GB | 18.5 s |
| Offline, forward-chunked KL | 61.8 GB | 18.4 s |
| Offline, fused chunked KL | 58.3 GB | 20.2 s |
All four methods produce nearly identical training results, meaning the quality loss from using the cached top-100 predictions instead of the full teacher distribution is negligible.
Does this matter for longer documents?
Yes, and the gap grows quickly. The memory advantage of the fused chunked approach is modest at 8,000 tokens but dramatic at longer lengths. At 32,000 tokens, peak memory drops from 85.2 gigabytes with the standard approach to 5.45 gigabytes with the new one. The dense loss fails completely at 64,000 tokens. At 256,000 tokens, the fused chunked method uses 11.6 gigabytes against 134.2 gigabytes for the next-best alternative.
In plain terms: experiments that previously required a rack of hundreds of chips can now run on a single high-end card. That makes it realistic for smaller research teams and companies to build compressed models from the largest open-source AI systems available today.
Common questions
Does the smaller model end up worse than the original?
In the benchmarks reported, the student model's training loss was nearly identical across all four methods, including the new cheap approach. Quality held up despite using only the top 100 word predictions per position instead of the full teacher output.
Can anyone use this technique?
The chunked loss code is open source and posted publicly on GitHub by CompactifAI. Anyone with a single modern AI chip and a copy of a large open-source model can, in principle, run their own distillation experiments with it.
Who benefits first?
Small AI research teams and companies that want to compress large open-source models for cheaper deployment will feel the difference most quickly. End users may benefit later as more affordable compressed models reach products and services.
