AI & Technology

Orchestral Manoeuvres in the Dark? Shrinking LLMs with self-compression

By Ed Plowman, CTO, Imagination Technologies

Edge devices are the ideal playground for large language models. They offer privacy, minimal latency and good reliability.  But delivering these experiences at sufficient speed on a battery constrained device like a smartphone is still a major challenge.  

One potential approach is, of course, hardware related: faster processors, more memory, bigger batteries. Another approach is software innovation: new techniques that help high quality models fit efficiently into the edge.  

A direct and surprisingly effective software technique is to prune complete channels (feature dimensions within a network) whose contribution to the task is negligible.  The question is how to prune precisely – and whether it can be done without manual input. 

What is self-compression? 

Self-compression is a training-time procedure that lets back-propagation decide the bit-width of every channel, so unhelpful ones fade away. This simultaneously reduces parameters, numerical precision to improve efficiency, and even the number of tuneable hyper-parameters, yet leaves predictive quality intact. 

In tests where the same idea is extended to a transformer, there are some intriguing results: when a channel’s learned precision slips all the way to zero bits, the resulting model becomes even more compact than one restricted to a fixed ternary code.   

Think of it like managing an orchestra about to go on a massive tour, but you only have one coach with a restricted number of seats and luggage space. 

You decide to let the musicians rehearse and figure out who should attend and how much luggage they can bring. During rehearsals you tell them: “The more important your part is to the performance, the more you can pack. If your part is barely audible, you’re going to have to pack small – or maybe not come at all.” 

As they practice, the musicians discuss and refine. The solo violinist clearly gets a seat and a large suitcase. The triangle player realises they aren’t actually contributing at all. Their allowance shrinks all the way to zero. They pack nothing and stay home.  

The orchestra gets smaller and lighter. The performance remains amazing. And the manager has done very little work to get to this point. 

The idea can go further still. We have explored the use of a block-based sparsity pattern. We have woven self-compression into the baseline model, examining the weight patterns it produces, and the implications for resource-constrained deployment. 

Self-compressing an LLM – an example 

The reference model we have used is nanoGPT [2], a compact GPT variant trained on the shakespeare_char corpus. With about 11 million trainable parameters, the model is small enough to run quickly, yet large enough to expose the full transformer computation pattern. 

The model contains: 

  • An embedding of each token to a multidimensional vector. 
  • Six identical transformer blocks with causal multi-head attention (including input and output linear layers), layer normalisation, and a feed-forward module containing a further two linear layers. 
  • An output stack comprising a final layer normalisation, a linear layer, and a softmax outputting probability over the possible output tokens. 

In transformer networks, over 90% of the weights – and therefore the bulk of memory bandwidth, DRAM consumption and power usage – reside in the four large linear layers within the transformer blocks. FSelf-compression was only targeted at those layers, leaving the remainder of the baseline model unchanged. 

Understanding which layers go sparse first provides useful insight into which layers of an LLM may be less important. This may allow targeting of future optimisations where redundancy naturally accumulates. 

How self-compression works 

Self-compression lets the network learn its own width and precision as part of conventional neural network training. Every output channel is quantised by a differentiable function. 

where bit depth b≥0  and scale exponent e are learnable parameters in the same sense as neural network weights.  A straight-through estimator treats the rounding derivative as 1, so that b and e receive normal gradients. This is easy to implement in a DL framework like PyTorch. 

Training minimises the original task loss L0 , but we include an extra size penalty Q . 

 

where Q is the average bits stored in the model, and  is a penalty term chosen by the user. 

For appropriate choice of y, the method retains baseline accuracy while cutting total model bits, all within an otherwise standard training setup. 

What happens after self-compression? 

As training progresses, we see the number of bits remaining in the network decreasing along with the validation loss. After starting training by assigning 4 bits per weight, it reduces by half within a few hundred epochs, going on to settle over time at around 0.55 bits per weight. 

Things get more interesting still when we look at the sparsity rate, and how it changes with depth. Sparsity increases as we go deeper into the model. This suggests that later layers are less information dense.  

We observed that the linear layers in the attention module become the sparsest in deeper blocks, with more than 95% of weights in two specific blocks. The feed-forward layers also became very sparse, with about 85% of weights removed. By contrast, the first block (0) keeps over half of its weights, likely because early layers are important for capturing basic patterns in the data. 

If this generalises to other language models and tasks, then it’s possible that a performance saving can be obtained even without self-compression, simply by reducing the number of features towards the end of the network. 

Conclusion 

Self-compression reduces both the bit-width and the number of active weights, while creating patterns like channel sparsity that are easy to understand and efficient to use on hardware. Early layers stay mostly dense to keep important information, while deeper layers become highly sparse. The few remaining weights stay small and close to zero. These results show that Self-Compression helps build models that are smaller, faster, and ready for use in resource-limited settings like edge devices. 

Self-Compression can also successfully shrink a transformer model, in this case a nanoGPT trained on character-level Shakespeare dataset, without harming its predictive quality. By letting the model decide which channels and weights to keep, the method avoids unnecessary manual tuning and produces a clean block-sparse structure that is easy to deploy on CPUs, GPUs, NPUs, and other hardware. This means that the same compact model can serve across the entire edge stack without further modifications and no additional manual effort for the programmer. 

To return to the orchestra analogy: block-based self-compression means that organisations can cross entire departments off the master list. It is no longer a case of saving on occasional moments of “luggage weight”; once the pruning has taken place once, the tour can be repeatedly and efficiently rolled out to any location. 

Future work could explore applying this method to larger language models, multi-modal transformers, or models fine-tuned for specific tasks, as well as combining self-compression with other techniques like knowledge distillation to push efficiency even further. 

References 

Szabolcs Cséfalvay and James Imber. Self-compressing neural networks
arXiv preprint arXiv:2301.13142, 2023. URL https://arxiv.org/abs/2301.13142

Andrej Karpathy. Nanogpt
https://github.com/karpathy/nanoGPT, 2023.

Sahitya Arya. Mastering decoder-only transformer: A comprehensive guide
https://www.analyticsvidhya.com/blog/2024/04/mastering-decoder-only-transformer-a-comprehensive-guide/, 2024

Szabolcs Csefalvay. Self-compressing neural networks
https://blog.imaginationtech.com/self-compressing-neural-networks , 2022.

Self-Compression of Language Models for Edge Inference, Edge AI Milan 2025
https://site.pheedloop.com/event/milan2025/sessions/SESYFUFEHW3Q3B0RE  

Author

Related Articles

Back to top button