Augmented AI: Distilling a 7B Teacher into a 70M On-Device Student

Introduction

Modern language models have become extraordinarily capable, and almost entirely out of reach of the hardware most people actually own. A 7 billion parameter model needs gigabytes of RAM and serious compute, which means the power of conversational AI sits behind cloud APIs. This project takes the opposite path. It builds a student model with only 70 million parameters that learns to imitate a 7 billion parameter teacher, then deploys that student where the teacher can never go: on device, in the browser, on the cheapest hardware you can find.

This is not a toy. A 100x reduction in parameters, trained with careful distillation, produces a model small enough to load in under a second that still carries the conversational behavior of its teacher. The architecture is a classic GPT style transformer, the training pipeline is transparent, and the result ships as a working web interface anyone can open and talk to.

Why Distillation

The Size Problem

A 7B parameter model in full precision occupies roughly 28 GB. Even quantized to 4 bits it needs around 3.5 GB, which immediately disqualifies most phones, cheap laptops, and browser based products. The developer ecosystem has largely accepted this as a law of nature: big models live in the cloud, small devices talk to them over the network.

Distillation challenges that assumption. The student is not trained to memorize the teacher’s outputs. It is trained to reproduce the teacher’s behavior, its distribution over next tokens, which compresses the teacher’s knowledge into a fraction of the parameters. The teacher becomes a living textbook; the student becomes a working expert.

The Privacy Argument

Every conversation sent to a cloud model is data in someone else’s hands. For a medical question, a business plan, or a child’s homework, that is a meaningful trade. A distilled on device model keeps every query local. The inference runs where you are. The data never leaves the device. Privacy stops being a policy promise and becomes an architectural property.

The Economics

Inference in the cloud costs money per token. At scale, that cost defines what products can exist. A 70M parameter model runs on CPUs, batteries, and browsers, with zero marginal cost. Students get an assistant that works offline. Clinics get a screening tool that never drops its connection. Developers get a model they can embed anywhere, forever, for free.

System Architecture

The project is organized as a closed loop. The teacher produces synthetic conversations. Those conversations train the student. The student then serves users through a web interface. Each stage depends on the previous one, and the whole loop is designed to be rerun whenever better teachers or better datasets appear.

     ┌────────────────────────┐      ┌────────────────────────┐      ┌────────────────────────┐
     │     Teacher Model      │      │     Synthetic Data     │      │     Student Model      │
     │     (7B LLM, GGUF)     │►────◄│       Generation       │►────◄│    (70M GPT Style)     │
     │                        │      │                        │      │                        │
     └────────────────────────┘      └────────────────────────┘      └────────────────────────┘
                  │                                                               │
                  ────────────────────────────────┴────────────────────────────────
                                                  │
                                   Distillation Training Pipeline

The Teacher Model

The teacher is a 7B parameter LLM served locally through llama.cpp and a FastAPI endpoint. Running the teacher on the same machine as the pipeline has two practical virtues. First, synthetic data generation never hits rate limits or API costs. Second, the teacher’s logits, the raw probability distributions over the vocabulary, are directly accessible, which is exactly what distillation needs.

The Student Model

The student is a 70M parameter GPT style transformer, written from scratch with a straightforward model definition. Its internals follow the canonical recipe, scaled down:

               ┌────────────────────────────────────────────────────────────┐
               │                       Student Model                        │
               │                                                            │
               │────────────────────────────────────────────────────────────│
               │  Embedding Layer                                           │
               │  [Vocab: 50258]  Token Emb + Pos Emb + LN                  │
               │────────────────────────────────────────────────────────────│
               │  Transformer Decoder Blocks                                │
               │  ┌──────────────────────────────────────────────┐          │
               │  │      Multi-Head Self-Attention (Causal)      │          │
               │  │         Feed Forward Network (2048)          │          │
               │  │             Residual + LayerNorm             │          │
               │  └──────────────────────────────────────────────┘          │
               │  × 6 Layers                                                │
               │────────────────────────────────────────────────────────────│
               │  Language Modeling Head                                    │
               │  Linear(512, 50258) + Softmax                              │
               └────────────────────────────────────────────────────────────┘

Every component is declared explicitly in the model file: the vocabulary of 50,258 tokens (matching the tokenizer), an embedding layer, six decoder blocks each containing causal self attention and a feed forward network, and a language modeling head that projects back into token space. Nothing is hidden in a framework abstraction. A reader can trace the entire forward pass by hand.

The Training Pipeline

Training joins the teacher’s soft knowledge with the student’s own grounding:

  1. Generate Synthetic Data
     └── Teacher generates 1000+ Q&A pairs
  2. Save Teacher Logits
     └── Cache soft labels for distillation
  3. Distill Train Student
     └── 20 epochs, KL + CE loss
  4. Evaluate Learning
     └── Test on held-out conversational tasks
  5. Serve via API/UI
     └── Web demo for interaction

The pipeline is explicit about its steps: generate a thousand question and answer pairs from the teacher, cache the teacher’s logits as soft labels, distill the student across twenty epochs using a combination of KL divergence loss and cross entropy loss, evaluate on held out conversational tasks, and finally serve the result through the API and web interface.

Getting Started

Prerequisites

• Python 3.10 or newer.
• Git.
• A machine capable of running a 7B model for the training phase. A GPU is strongly recommended for the distillation run, though the pipeline functions without one.
• Roughly 8 GB of disk space for models, datasets, and checkpoints.

Clone and Install

git clone https://github.com/AbduljabbarBXR/Augmented-Ai.git
cd Augmented-Ai
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Download the Models

The teacher model is distributed as a GGUF file for llama.cpp compatibility. Place it under the models directory and note the path; the configuration file will need it.

The Dataset

Synthetic Generation

Synthetic data is the fuel of the pipeline. Rather than scraping the web or licensing datasets, the teacher is prompted to generate its own question and answer pairs across the domains the final product should handle. Each session produces structured conversations:

• User questions that resemble real usage, phrased naturally.
• Assistant answers in the teacher’s voice and format.
• Domain coverage that can be controlled by the prompt template, so the dataset is not a lucky accident but a designed artifact.

One thousand pairs is deliberately modest. It keeps the pipeline fast and the cost near zero, while demonstrating that distillation quality does not require million row datasets when the teacher is strong.

Logit Caching

The pipeline computes the teacher’s full logits over the training set once and stores them on disk. Distillation then runs without the teacher in memory, which shrinks training’s resource footprint dramatically. The student is trained against fixed soft labels rather than live queries.

Training the Student

Loss Function

The training objective combines two signals:

• The KL divergence between the student’s predicted distribution and the teacher’s cached logits, which transfers the teacher’s confidence patterns and inter token relationships.
• Standard cross entropy against the original tokens, which anchors the student to the ground truth text.

loss_total = alpha * loss_ce + (1 - alpha) * loss_kl

The alpha coefficient balances the two. High alpha leans on the hard text targets and behaves like ordinary language modeling. Low alpha trusts the teacher’s soft signal. The default configuration sits between the two, and the training script accepts the value as an argument so the balance can be swept experimentally.

Training Run

python train.py --epochs 20 --alpha 0.5 --batch-size 8

Twenty epochs over one thousand pairs is enough to see the student converge to coherent conversational behavior. Progress, loss curves, and sample generations are printed at each epoch so training is legible, not a black box.

Evaluation

The student is tested on held out conversational tasks it never saw during training. Evaluation looks for the behaviors that matter in a product: coherent replies, appropriate tone, and the ability to follow a simple instruction. Sample outputs are logged, making it easy to see what the model learned, and what it has not.

Serving the Model

FastAPI Endpoint

A FastAPI server exposes the student with a clean interface:

@app.post("/generate")
def generate(request: GenerationRequest):
    output = model.generate(
        request.prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature
    )
    return {"response": output}

Web Interface

The repository ships a web demo served by the same process, giving any browser in the room an instant chat interface to the student. Because the model is 70M parameters, responses arrive fast enough to feel live even on modest hardware.

python api.py

Open the local URL and start typing.

Deploying On Device

The payoff of the entire project is the deployment profile:

• The final model fits in tens of megabytes, quantized or not.
• Loading takes under a second on typical hardware.
• Inference runs on CPU; no GPU is required.
• Everything operates offline, with no API keys and no network calls.

That profile is what turns the student into a product surface: an embeddable assistant in a settings screen, a help widget on a website, a companion on a low powered phone. The same weight class that is a footnote in cloud scale discussions becomes the main event at the edge.

Results and Observations

The 100x parameter reduction is dramatic, and it is important to be honest about what it buys and what it costs. The student reproduces the teacher’s conversational style and short form correctness surprisingly well, which is the promise of distillation kept. It does not match the teacher’s long form reasoning or factual depth; those live in the parameter budget the student does not have.

What matters is that the trade is deliberate and tunable. More data, stronger teachers, and architecture tweaks all move the quality needle, and the pipeline is built to be rerun cheaply. This is a starting point for a research direction, not a finished destination.

Project Structure

Augmented-Ai/
├── README.md # Project overview and quickstart
├── requirements.txt # Python dependencies
├── models/ # Model checkpoints and GGUFs
│ └── teacher.gguf # 7B teacher model
├── data/
│ └── synthetic_dataset.json # Generated Q&A pairs
├── src/
│ ├── model.py # Student model definition
│ ├── tokenizer.py # Tokenizer utilities
│ ├── data_utils.py # Dataset loading and caching
│ ├── train.py # Distillation training script
│ └── api.py # FastAPI serving + web UI
└── notebooks/
 └── exploration.ipynb # Data and training exploration

Customizing the Pipeline

Changing the Teacher

Swap the GGUF file and regenerate the dataset. The prompt template in the generation script controls the domains the teacher covers, so editing that template repositions the entire product.

Changing the Student

The model definition in src/model.py is a plain PyTorch module. Width, depth, and vocabulary are parameters at the top of the file. A larger student closes some of the quality gap with the teacher at the cost of deployment size; a smaller one pushes the other direction.

Changing the Data Budget

More synthetic pairs improve the student, but the marginal return flattens. Sweep the pair count on a small slice first, then scale only what moves the metric.

Limitations

• The student inherits the teacher’s biases and blind spots, since it learns from the teacher’s outputs.
• Long form factual reasoning is the first casualty of the parameter reduction.
• Synthetic data quality depends on prompt design; a lazy prompt template produces a lazy dataset.
• The current demo targets conversational Q&A, not code generation or structured reasoning.

Future Work

• Training on larger and more diverse synthetic datasets.
• Comparing against alternative compression paths, such as quantization alone and pruning.
• A mobile runtime for the student, extending the web demo to native apps.
• Streaming generation in the web interface for a more live feel.
• Multi turn memory so conversations carry context across exchanges.

Contributing

Issues, pull requests, and ideas are welcome. The most useful contributions come from people who run the pipeline on real hardware and report what they find.

1. Fork the repository.
2. Create a feature branch.
3. Run the existing pipeline to verify the baseline.
4. Submit a pull request with a clear description of the change.

License

This project is licensed under the MIT License. See the LICENSE file for details.

Acknowledgments

Thanks to the llama.cpp project for making local large model inference practical, to the PyTorch team for the training stack, and to the open source community that keeps distillation research moving.