vivekkopthsd/vajra-1

TxCat-GPT

A 703K-parameter GPT trained from scratch on ML/statistics/coding research text. Raw PyTorch throughout: custom BPE tokenizer, hand-written causal attention, no training frameworks, no pretrained weights.

Trained on a Kaggle P100 (torch 2.4.0+cu121). 3,000 steps take about a minute on that GPU.

Model

Param	Value
Parameters	702,948
Layers / heads / d_model	3 / 4 / 128
Context length	128 tokens
Vocabulary	356 (BPE, 100 merges, trained on 1.5M chars of the corpus)
Dropout	0.1

Training

• Corpus: 14.0M chars, 8.08M tokens. arXiv abstracts filtered to cs.LG, stat.ML, cs.CL, cs.AI, stat.ME, math.ST, plus the author's own transformer code.
• 3,000 steps, batch 64, AdamW lr 3e-4, weight decay 0.01, cosine schedule, grad clip 1.0.
• Val loss 1.995, val perplexity 7.35.
• Loss trajectory (perplexity): 321 at step 1, 30.7 at 300, 7.9 at 1500, 7.35 at 3000.

This is a small model on a small corpus. It learned the surface structure of ML abstracts: paper format, domain vocabulary, sentence scaffolding. It does not produce coherent text, and the generations should not be read as content. At this scale that is the expected outcome; the model's purpose is to demonstrate the pretraining pipeline end to end, not to generate usable prose.

Usage

Clone the repo and run the included CLI:

python txcat_gpt_infer.py "TITLE: Bias-variance tradeoff
ABSTRACT:"

Generates 160 tokens at temperature 0.8 on CPU. No GPU needed.

from txcat_gpt_infer import load_model, generate

text, obj = generate("def gradient_descent(X, y,", n_tokens=120, temperature=0.8)

txcat_gpt_infer.py contains the model definition, the BPE tokenizer, and the generation loop. It loads pytorch_model.bin, config.json, and merges.json from this repo.

Files

File	Contents
pytorch_model.bin	State dict (3.0 MB)
config.json	Architecture and training config
merges.json	BPE merge table
txcat_gpt_infer.py	Model, tokenizer, generation CLI

Reproducibility

• Training kernel: vivekkopthsd/txcat-gpt on Kaggle (private). Four revisions; the log documents each failure.
• Corpus: vivekkopthsd/txcat-ml-corpus on Kaggle (private).
• Hardware: P100, torch 2.4.0+cu121. The Kaggle default image ships cu128 builds that have no Pascal kernels, so the pin is required.

History

Two bugs were found and fixed during development; both are documented in the training kernel log.

1. Bidirectional attention. The first version used nn.TransformerEncoderLayer, which is bidirectional. It reported val perplexity 1.03, which was not a real language-modeling result: the model could attend to future tokens during training and collapsed at generation time. Replaced with hand-written causal attention (triangular mask). All numbers above are from the causal model.
2. BPE decoder recursion. Merged tokens that contain other merged tokens crashed the decoder (a str was passed to bytearray.extend). Fixed by decoding to bytes at every recursion level and decoding UTF-8 once at the top.

License

Weights trained from scratch on public arXiv abstracts (see arXiv terms of use) and the author's own code. Educational use.