Baseball-LLM — Kevin Jones

Overview

The goal was to create and train a large language model from scratch using the Transformer architecture, then tailor the model to answer baseball questions.

The central experimentCompare a general language model with a baseball-specialized version and observe what specialization improves — and what it costs.

Model design

The model follows the Transformer architecture from Attention Is All You Need. Instead of using PyTorch's built-in Transformer block, the implementation builds multi-headed self-attention with linear and embedding layers to develop a more fundamental understanding of the architecture.

Hyperparameters were chosen to follow OpenAI's GPT-2-small design, including a 124-million-parameter architecture and the GPT-2 tokenizer. A custom byte-pair tokenizer was also explored but proved too slow without the optimizations used by tiktoken.

Pretraining

The model was initialized with random weights and pretrained on the FineWeb-Edu dataset from Hugging Face. Training ran for several days on an NVIDIA 4070 GPU.

GPU memory limited each step to 2¹² tokens, so gradient accumulation was used to simulate a substantially larger batch size. The project also used the learning-rate decay approach described in the GPT-2 paper.

Validation-loss plots and training artifacts are available in the GitHub repository.

Fine-tuning

After pretraining, the model was fine-tuned using the Nectar dataset from Berkeley NEST. The data was preprocessed with an end-of-text token, and training continued for several days with gradient accumulation.

Over the course of fine-tuning, the model began to learn how to respond as an assistant by answering user questions.

Baseball specialization

The next step was continued pretraining on the official 2021 baseball rules, followed by fine-tuning on roughly 1,000 baseball-specific question-and-answer pairs.

The intent was to improve baseball performance without completely losing the capabilities learned during the general training stage.

Results

The baseball-tuned model answered baseball questions substantially better than the general model. In the project's numerical comparison, the baseball model achieved an average probability of 0.43 for the correct next token on held-out baseball QA pairs, versus 0.29 for the general model — about 50% higher.

The trade-off was equally clear: baseball specialization reduced performance on general questions, with the model often trying to relate unrelated questions back to baseball.

TakeawaySpecialization worked, but it came with a loss of generality — a concrete example of the trade-offs that appear when fine-tuning small language models toward a narrow domain.