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 peak memory sharply.
- A second technique, called a fused chunked KL loss, processes training data in small slices rather than all at once, keeping peak memory nearly flat even at very long inputs.
- At 32K 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 tiny dials inside a model that get tuned during training) and needs roughly three terabytes of specialised memory just to load. That's 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. The standard method is knowledge distillation: 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. Cheaper to run, nearly as capable.
Why is this so hard to do cheaply?
The expense has two sources. Keeping both the teacher and the student loaded in memory at once is the first. The second is that comparing their outputs requires building an enormous grid: one row for every word in the model's vocabulary, one column for every position in the input. Multiply those across a long piece of text and you fill every chip you have.
A concrete example: gpt-oss-120b has a vocabulary of 201,088 words. At a sequence length of 32K tokens and a batch size of four examples at once, the teacher's output tensor alone takes up about 50 gigabytes. Add the student model, gradients, optimizer states 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, can hold.
What did the researchers change?
Two things, working together. Both appear in a paper from CompactifAI, first shared on Hugging Face.
The first 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 goes to disk; the teacher is then switched off entirely. Because the cache is reusable across many experiments, the expensive teacher run becomes a one-time cost.
The second is the fused chunked KL loss. KL divergence is a mathematical measure of how different two probability distributions are, here meaning how far the student's word predictions stray from the teacher's. The standard calculation builds that full vocabulary-by-sequence grid in one go and holds all of it. The new approach works one slice of the sequence at a time, folds each slice into a running total, then discards it before moving on. The backward pass recomputes what it needs 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. The quality loss from using cached top-100 predictions instead of the full teacher distribution is negligible.
Does this matter for longer documents?
Yes, and the gap grows fast. The memory advantage of the fused chunked approach is modest at 8K tokens but dramatic beyond that. At 32K tokens, peak memory drops from 85.2 gigabytes with the standard approach to 5.45 gigabytes with the new one. Experiments that previously needed a rack of specialist chips can now run on a single high-end card. That's a realistic opening for smaller research teams and companies wanting to build compressed models from the largest open-source AI systems available today. We covered how Microsoft's Orchard framework pursues a similar goal, fitting capable agents into smaller models, in our 5 August story.
The honest caveat: the benchmark here uses Llama 3.1 8B as student. How well the technique holds at larger student sizes is a question the paper doesn't fully settle yet.
Common questions
Does the smaller model end up worse than the original?
In the benchmarks reported, training loss was nearly identical across all four methods, including the new cheap approach. Quality held despite using only the top 100 word predictions per position rather than the full teacher output.
Can anyone use this technique?
The chunked loss code is open source, 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.
Who benefits first?
Small research teams and companies wanting to compress large open-source models for cheaper deployment will feel the difference soonest. End users may benefit later as more affordable compressed models reach products and services.
