Upload ARCHITECTURE.md with huggingface_hub
Browse files- ARCHITECTURE.md +70 -0
ARCHITECTURE.md
ADDED
|
@@ -0,0 +1,70 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# MiniLM: The 1.58-bit Architecture Deep Dive
|
| 2 |
+
|
| 3 |
+
MiniLM is not just a quantized model—it is a completely custom neural network architecture built from the ground up to natively operate in **1.58-bit (Ternary)** precision.
|
| 4 |
+
|
| 5 |
+
By heavily compressing the internal mathematics of the Transformer, we achieved a deep 12-layer model that fits entirely into **6.00 MB** of RAM, making it small enough to run on microcontrollers, smartwatches, and embedded IoT devices.
|
| 6 |
+
|
| 7 |
+
This document serves as a masterclass on exactly how MiniLM was engineered.
|
| 8 |
+
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
## 1. The Core Innovation: 1.58-bit Ternary Weights
|
| 12 |
+
|
| 13 |
+
In standard Large Language Models (like Llama 3 or GPT-4), the neural network's memory (its "weights") are stored as 16-bit floating-point numbers (`FP16`). A single layer can easily exceed gigabytes of RAM.
|
| 14 |
+
|
| 15 |
+
MiniLM uses the **BitNet 1.58b** architecture paradigm. We discard floating-point precision entirely. Every single internal weight in MiniLM's Linear layers is constrained to exactly three possible values:
|
| 16 |
+
* `-1`
|
| 17 |
+
* `0`
|
| 18 |
+
* `1`
|
| 19 |
+
|
| 20 |
+
Because $\log_2(3) \approx 1.58$, we call this a 1.58-bit model.
|
| 21 |
+
|
| 22 |
+
### Why is this revolutionary?
|
| 23 |
+
When you multiply a number by `-1`, `0`, or `1`, you aren't actually doing complex matrix multiplication. You are simply doing **Addition and Subtraction**.
|
| 24 |
+
If a weight is `1`, you add the input. If it is `-1`, you subtract the input. If it is `0`, you ignore it.
|
| 25 |
+
|
| 26 |
+
This means MiniLM replaces the most computationally expensive operation in AI (Floating Point Matrix Multiplication) with ultra-fast, hardware-efficient Integer Addition.
|
| 27 |
+
|
| 28 |
+
---
|
| 29 |
+
|
| 30 |
+
## 2. How We Trained It: The Straight-Through Estimator (STE)
|
| 31 |
+
|
| 32 |
+
You cannot train a ternary neural network using standard backpropagation, because the rounding function (clamping a value to -1, 0, or 1) has a derivative of zero almost everywhere. The gradient would instantly "die" and the model would never learn.
|
| 33 |
+
|
| 34 |
+
To solve this, we implemented a custom **Straight-Through Estimator (STE)**:
|
| 35 |
+
1. **Forward Pass:** We take the high-precision latent weights, calculate their mean, divide by a scaling factor (`beta`), and aggressively round them to `[-1, 0, 1]`. The forward calculations are performed using these ternary weights.
|
| 36 |
+
2. **Backward Pass:** When the loss calculates the error gradient, we *pretend* the rounding step never happened. We pass the gradient straight through to the high-precision latent weights.
|
| 37 |
+
|
| 38 |
+
This allows the high-precision weights to slowly adjust over time, until their rounded ternary counterparts snap into the optimal configuration.
|
| 39 |
+
|
| 40 |
+
---
|
| 41 |
+
|
| 42 |
+
## 3. Breaking the Depth Barrier: Weight Tying
|
| 43 |
+
|
| 44 |
+
Our initial 4-layer model fit into 3.93 MB and showed promising results, but 4 layers is incredibly shallow for an LLM to form coherent, long-form thoughts.
|
| 45 |
+
|
| 46 |
+
To solve this, we implemented **Weight Tying**.
|
| 47 |
+
In a standard LLM, the `Embedding Layer` (which turns words into vectors) and the `Output Head` (which turns vectors back into words) are two separate, massive matrices.
|
| 48 |
+
|
| 49 |
+
Because we used a 32,000 token vocabulary, these two matrices were consuming **over 85%** of our total parameter budget!
|
| 50 |
+
|
| 51 |
+
By mathematically tying the weights together (`model.head.weight = model.embedding.weight`), we instantly freed up 8 Million parameters. We re-invested this exact parameter budget to triple the depth of the neural network from 4 layers to **12 layers**, drastically improving output coherence without increasing the file size by a single byte.
|
| 52 |
+
|
| 53 |
+
---
|
| 54 |
+
|
| 55 |
+
## 4. Knowledge Distillation
|
| 56 |
+
|
| 57 |
+
Training a 1.58-bit model from absolute scratch using Next-Token Prediction is notoriously difficult and requires massive amounts of data and compute (100k+ steps).
|
| 58 |
+
|
| 59 |
+
Instead, we used **Knowledge Distillation**.
|
| 60 |
+
1. We loaded `HuggingFaceTB/SmolLM-135M-Instruct` as a "Teacher" model.
|
| 61 |
+
2. We forced MiniLM to use the exact same tokenizer as SmolLM.
|
| 62 |
+
3. For every prompt, the Teacher model output a rich probability distribution (logits) of what the next word should be.
|
| 63 |
+
4. We used `KLDivLoss` (KL Divergence) to force MiniLM to perfectly mimic the Teacher's probability distribution.
|
| 64 |
+
|
| 65 |
+
By learning from the Teacher's rich understanding of language rather than just a sparse one-hot encoded dataset, MiniLM converged in just **3,000 steps** on the TinyStories dataset!
|
| 66 |
+
|
| 67 |
+
---
|
| 68 |
+
|
| 69 |
+
## Conclusion
|
| 70 |
+
MiniLM is a testament to the future of Edge AI. By combining Ternary Quantization, Weight Tying, and Knowledge Distillation, we have packed the structural depth of a 12-layer Transformer into a file size smaller than an MP3 song.
|