TypeScript Decoder-Only with Autograd Implementation

This project provides a complete, from-scratch implementation of a Decoder-Only architecture (the underlying structure of modern Large Language Models like the GPT series) using native TypeScript. It does not depend on any deep learning frameworks (such as TensorFlow.js or ONNX.js). Instead, it relies purely on native Float32Array computations to implement both Forward Pass and Autograd (Automatic Differentiation) Backward Pass, making it an ideal resource for demonstrating and understanding the inner workings of Transformer models.

Directory Structure

• tensor.ts - Core tensor computation module. Includes basic tensor operations (e.g., matmul, add, softmax) and a lightweight Autograd Engine based on topological sorting.
• layers.ts - Foundational neural network layers: Fully Connected (Linear), Layer Normalization (LayerNorm), and Position-Wise Feed-Forward Network (FFN).
• positional_encoding.ts - Sine/Cosine-based positional encoding module.
• attention.ts - Scaled Dot-Product Attention and Multi-Head Attention (MultiHeadAttention), supporting Causal Masking.
• decoder_layer.ts - A single Decoder Layer structure (comprising Masked Self-Attention and FFN).
• decoder.ts - The full Transformer Decoder model. Integrates the Language Model Head (LM Head) to output logits over the vocabulary.
• loss.ts - Implementation of Cross-Entropy Loss with forward and backward passes.
• optim.ts - Stochastic Gradient Descent (SGD) optimizer with gradient zeroing and parameter update functionality.
• test.ts - Unit tests and a training example with a 5-epoch loop.

Forward Pass Details

During model inference or training, data moves through 5 core stages:

1. Input & Embedding Stage

• Token Embedding: Maps discrete Token IDs to dense floating-point feature vectors [batch_size, seq_len, d_model] using an embedding weight table. The result is scaled by $\sqrt{d_{model}}$.
• Positional Encoding: Generates absolute positional encoding matrices using sine and cosine functions and adds them to the feature vectors. This allows the otherwise order-agnostic attention mechanism to perceive the sequence order.

2. Masked Multi-Head Self-Attention

• Q, K, V Generation & Splitting: Features are mapped via three linear layers into Query, Key, and Value, then split into num_heads independent heads to focus on different information dimensions.
• Attention Score Calculation: Computes $Q \times K^T / \sqrt{d_k}$.
• Causal Masking: Adds a mask matrix (upper triangle set to $-\infty$) to ensure a token only attends to itself and preceding tokens, preventing it from "peeking" at future words (essential for autoregressive generation).
• Weighted Fusion: Applies Softmax to scores to get weights, then multiplies by $V$. Heads are concatenated and mapped back to the original dimension.

3. Residual Connection & Layer Normalization (Add & Norm)

• Residual Connection: Adds the input x to the output of the sublayer x + sublayer(x), creating a "highway" for gradients to prevent vanishing gradients in deep networks.
• Layer Normalization: Standardizes features (subtracting mean, dividing by standard deviation) and applies learnable scale ($\gamma$) and shift ($\beta$) parameters to accelerate convergence.

4. Position-wise Feed-Forward Network (FFN)

• After attention-based context mixing, features at each position independently enter the FFN.
• The FFN consists of two linear layers with a ReLU activation in between, projecting features to a higher dimension (typically 4x) for non-linear transformation before compressing back.
• Each FFN is also followed by a Residual Connection and LayerNorm.

5. Output Stage (LM Head)

• Final features [batch_size, seq_len, d_model] are passed to the LM Head.
• This linear layer maps features from d_model to vocab_size, outputting Logits which can be converted to probabilities for the next-token prediction.

Training & Backpropagation Workflow

1. Loss Computation

• Softmax: Converts Logits into a probability distribution summing to 1.
• Negative Log Likelihood (NLL Loss): Extracts the predicted probability of the Target Token and takes its negative log. Lower loss indicates better prediction.
• The average loss across the batch is computed as a Scalar Loss Tensor.

2. Backpropagation (Autograd)

• Computation Graph: Every operation (matmul, add, etc.) implicitly records its parent nodes (.prev) and registers a gradient calculation closure (_backward).
• Execution: Calling loss.backward() triggers the Autograd engine to build a topological sort and initialize the loss gradient to 1.0.
• Chain Rule: Executes all _backward functions in reverse order. Errors propagate via the chain rule and accumulate in the .grad attribute of all learnable parameters (e.g., w_q, w_k, weight, gamma).

3. Optimizer Step

• Update (optim.step()): The SGD optimizer uses the accumulated gradients to update parameters: $W = W - \text{lr} \times \text{grad}$.
• Zero Grad (optim.zero_grad()): Resets all .grad attributes to 0 before the next batch to prevent gradient accumulation across iterations.

Usage

Run the test script to verify the training loop and loss reduction:

# Ensure tsx is installed or use npx
npx tsx tests/test.ts

Example output showing decreasing loss:

Epoch 1/5 | Loss: 4.5663 | Time: 12 ms
Epoch 2/5 | Loss: 4.3809 | Time: 2 ms
Epoch 3/5 | Loss: 4.1999 | Time: 1 ms
...

Inference & Generation

Once trained, the model can perform autoregressive text generation:

1. Prompt: Provide an initial sequence (e.g., [10]).
2. Forward Pass: The model outputs Logits of shape [1, seq_len, vocab_size].
3. Extract Last Token: We only care about the dictionary scores for the last position Logits[0, -1, :].
4. Sampling: Pick the Token ID with the highest score (Greedy Search).
5. Autoregressive Loop: Append the new token to the prompt and repeat until the desired length is reached.

License

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