Neerudi Sai Vikas
Resume ↗
← All projects
Research · 2024 – now

A language model, built from scratch

Wrote and trained a 21M-parameter GPT-style model in PyTorch, tokenizer included. Now using it to test ways to make small models cheaper to run.

Role
AI Systems Engineer & Applied ML Researcher
Built with
Python, PyTorch, Custom BPE Tokenizer, RoPE, KV-Cache, Shannon Entropy Patching, CSA2 Compression, FastAPI, Docker

Results

  • 21MShipped BaselineFrom-scratch PyTorch decoder + RoPE + KV-Cache
  • O(1)BPE Pair CountingTuple-hashed merge table tokenizer from scratch
  • H(Xₜ|X₍<ₜ₎)Axis 1: Input PatchingDynamic byte patching via conditional entropy
  • CSA2Axis 2: KV CompressionFull / Reindex / Reuse cross-layer KV sharing

By replacing uniform subword allocation with entropy-triggered byte patches and sharing latent KV representations across Full/Reindex/Reuse layer blocks (CSA2), target 100M–150M agentic SLMs can fit multi-turn tool trajectories in edge/single-GPU HBM with >60% lower KV memory footprint.

The problem

Standard autoregressive Transformers waste FLOPs on predictable subword sequences due to static frequency-based BPE merges while suffering from linear HBM KV-cache bottlenecks on long multi-turn agent trajectories.

In more detail

Autoregressive GPT-style decoder trained from scratch in PyTorch (21M BPE checkpoint shipped), serving as a controlled testbed for two-axis compression toward a 100M–150M agentic SLM: Shannon conditional-entropy input patching and cross-layer KV-cache reuse (CSA2).

How it works

┌──────────────────────────────────────────────────────────────────────────────┐
│ AXIS 1: INPUT REPRESENTATION (Static BPE vs. Shannon Entropy Patching)       │
│                                                                              │
│  • Shipped Baseline: Custom BPE Tokenizer (O(1) Tuple-Hashed Merge Table)    │
│  • Active R&D: Local Entropy Model P(x_t | x_{<t})                           │
│    Patch Boundary Triggered when H(X_t | X_{<t}) > τ_entropy (Surprise Gate) │
└──────────────────────────────────────┬───────────────────────────────────────┘
                                       │ Embeddings + Rotary Pos (RoPE θ_i)
                                       ▼
┌──────────────────────────────────────────────────────────────────────────────┐
│ AXIS 2: DECODER STACK & CSA2 CROSS-LAYER KV-CACHE COMPRESSION                │
│                                                                              │
│  ┌────────────────────────────────────────────────────────────────────────┐  │
│  │ Layer l_0 [FULL MODE]:    Compute Latent KV + Top-k Sparse Indices I_0 │  │
│  └───────────────────────────────────┬────────────────────────────────────┘  │
│                                      ├──► Share Compressed KV + Reindex      │
│  ┌───────────────────────────────────▼────────────────────────────────────┐  │
│  │ Layer l_1 [REINDEX MODE]: Reuse Parent KV Cache • Recompute Sparse I_1 │  │
│  └───────────────────────────────────┬────────────────────────────────────┘  │
│                                      ├──► Zero New KV Allocation + Reuse I_1 │
│  ┌───────────────────────────────────▼────────────────────────────────────┐  │
│  │ Layer l_2 [REUSE MODE]:   Reuse Both Parent KV Cache & Sparse Indices  │  │
│  └────────────────────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────┬───────────────────────────────────────┘
                                       ▼
┌──────────────────────────────────────────────────────────────────────────────┐
│ FastAPI AUTOREGRESSIVE INFERENCE SERVER (KV-Cached Streaming Generation)     │
└──────────────────────────────────────────────────────────────────────────────┘

Key decisions

  1. Shipped Baseline (21M Parameter PyTorch Checkpoint)

    Implemented a complete decoder-only Transformer from first principles in PyTorch—including custom Byte-Pair Encoding (BPE) merge tables with O(1) tuple-hashed pair counting, Multi-Head Attention (MHA), Rotary Position Embeddings (RoPE), and autoregressive KV-cached inference containerized in FastAPI.

  2. Active Axis 1 — Input Representation (Shannon Conditional Entropy vs. Static BPE)

    Investigating dynamic byte/character patching guided by local Shannon conditional entropy H(X_t | X_{<t}) (inspired by Byte Latent Transformer) so model compute scales dynamically with information surprise rather than fixed frequency merges.

  3. Active Axis 2 — Inference Memory (Cross-Layer KV & Sparse Index Reuse via CSA2)

    Evaluating downscaled variants of DeepSeek-V4.1-Flash's Compressed Sparse Attention 2 (CSA2) across Full, Reindex, and Reuse layer modes plus teacher-student knowledge distillation to minimize per-token HBM KV footprint on long agent trajectories.

One request, step by step

Further reading