narrow-distill-1B

Narrow-Distill-1B

A 166M-parameter narrow-and-deep transformer (d=256, 200 layers), trained via knowledge distillation from Qwen2.5-Coder-14B, as part of a research experiment testing whether depth beats width at equal parameter budget.

Full experiment writeup, training code, and evaluation script: github.com/yozzaofficial/Narrow-Model

Architecture

Result

Compared against a Wide baseline (d=1024, 12 layers, 185M params) trained on the same distilled corpus, Narrow achieved better perplexity — supporting the hypothesis that depth aids generalization more than width in this regime.

Limitations

This is not a general-purpose language model. It uses a custom word-level tokenizer (not BPE, not HuggingFace AutoTokenizer compatible), has no chat/instruct format, and was trained on a relatively small corpus (~270K words). It is a research artifact for the Narrow-vs-Wide comparison, not a production-ready model.

Loading

import torch
import torch.nn as nn
import torch.nn.functional as F

VOCAB_SIZE, SEQ_LEN = 16384, 512

class Block(nn.Module):
    def __init__(self, d, n_heads, d_ff):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.Wq = nn.Linear(d, d, bias=False)
        self.Wk = nn.Linear(d, d, bias=False)
        self.Wv = nn.Linear(d, d, bias=False)
        self.Wo = nn.Linear(d, d, bias=False)
        self.ff1 = nn.Linear(d, d_ff, bias=False)
        self.ff2 = nn.Linear(d_ff, d, bias=False)
        self.n_heads, self.d_head = n_heads, d // n_heads

    def forward(self, x):
        B, T, D = x.shape
        h = self.norm1(x)
        q = self.Wq(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
        k = self.Wk(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
        v = self.Wv(h).reshape(B, T, self.n_heads, self.d_head).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.Wo(out.transpose(1, 2).reshape(B, T, D))
        return x + self.ff2(F.gelu(self.ff1(self.norm2(x))))

class NarrowLM(nn.Module):
    def __init__(self, vocab=VOCAB_SIZE, d=256, n_heads=8, n_layers=200, d_ff=1024):
        super().__init__()
        self.embed = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(SEQ_LEN, d)
        self.blocks = nn.ModuleList([Block(d, n_heads, d_ff) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d)
        self.lm_head = nn.Linear(d, vocab, bias=False)

    def forward(self, idx):
        B, T = idx.shape
        x = self.embed(idx) + self.pos(torch.arange(T, device=idx.device))
        for blk in self.blocks:
            x = blk(x)
        return self.lm_head(self.norm(x))

model = NarrowLM()
model.load_state_dict(torch.load("Narrow_distill_1B.pt", map_location="cpu", weights_only=False))
model.eval()

Context

Part of ChunkLLM, a research project on running large language models (14B+) on consumer hardware with limited RAM (16GB, Apple Silicon). This result informs the router architecture used in ChunkLLM.