AIMS-1-128M-it
AIMS-1-128M-it is a supervised fine-tuned 128M-parameter language model developed as an experiment in training a language model on a $200 budget.
The primary objective of the project is to investigate how far a relatively small, independently trained language model can go in learning the syntax and structure of English using a limited compute budget.
The model uses a custom architecture that borrows design ideas from modern language-model architectures such as Llama, Qwen, and GPT, while implementing its own model components and configuration.
Following pretraining, AIMS-1-128M was supervised fine-tuned (SFT) to follow a user/assistant conversational format and to produce responses to everyday conversational prompts.
Model specifications
- Parameters: 128,551,424
- Layers: 28
- Hidden width: 512
- FFN width: 1365
- Query heads: 16
- KV heads: 4
- Context length: 1024 tokens
- Vocabulary size: 50,260
- Residual dropout: 0.1
- RMSNorm epsilon: 1e-05
- Tokenizer: Custom GPT-2 byte-level BPE tokenizer exported from the training encoding
The architecture is custom-built and incorporates ideas found in several influential Transformer language-model families, including GPT, Llama, and Qwen. It should not be considered a direct implementation or derivative architecture of any one of these models.
Intended purpose
AIMS-1-128M-it was created primarily as a research and educational experiment around low-budget language-model training.
The central question behind the project is:
Can a useful understanding of English language syntax be learned by training a small language model with approximately $200 of compute?
The model is therefore intended to explore:
- Language-model training under strict compute constraints
- Learning of English syntax and linguistic patterns
- The capabilities and limitations of small language models
- Supervised fine-tuning of a custom language model
- Low-cost approaches to building and evaluating LLMs from scratch
The model is not intended to compete with large general-purpose commercial language models. Its primary value is as an experiment in understanding what can be achieved with a small model and limited resources.
Conversational capabilities
The base model was subsequently fine-tuned on conversational data to teach it the user/assistant interaction format.
The resulting SFT dataset contains examples across the following conversation categories:
| Category | Examples | Distribution |
|---|---|---|
| unknown_factual | 9,705 | 64.7% |
| casual_everyday | 1,103 | 7.4% |
| conversation_followup | 745 | 5.0% |
| social | 641 | 4.3% |
| hobbies | 637 | 4.2% |
| preferences | 514 | 3.4% |
| emotional_support | 441 | 2.9% |
| greetings | 410 | 2.7% |
| playful | 376 | 2.5% |
| unknown_personal | 133 | 0.9% |
| predictions | 131 | 0.9% |
| impossible_questions | 164 | 1.1% |
| Total | 15,000 | 100.0% |
The conversational fine-tuning is intended to improve the model's ability to respond in a user/assistant format. It does not imply that the model has reliable factual knowledge or robust reasoning capabilities.
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
repo_id = "AIMadeSimpleResearch/AIMS-1-128M-it"
tokenizer = AutoTokenizer.from_pretrained(
repo_id,
trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
repo_id,
trust_remote_code=True,
)
model.eval()
messages = [
{"role": "user", "content": "Hello! How are you?"}
]
inputs = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
)
output = model.generate(
**inputs,
max_new_tokens=40,
use_cache=False,
)
print(
tokenizer.decode(
output[0, inputs.input_ids.shape[1]:],
skip_special_tokens=True,
)
)
Consumers can install torch, transformers, and safetensors.
Custom-code loading is required.
This implementation currently has no KV cache. Generation therefore recomputes the prefix at each generation step. Keep the prompt plus generated tokens within the model's 1024-token context limit.
For batched generation, use left padding.
The padding adapter processes padded examples individually using the original AIMSModel and then restores the batch layout. This is slower than fully batched masked attention but avoids modifying the original model component implementation.
Temperature and top-p sampling
The model inherits generate() from Hugging Face's GenerationMixin.
Sampling can be enabled using the standard Transformers generation parameters:
model.eval()
messages = [
{"role": "user", "content": "Hello! How are you?"}
]
inputs = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
output = model.generate(
**inputs,
max_new_tokens=100,
do_sample=True,
temperature=0.7,
top_p=0.9,
top_k=0,
use_cache=False,
)
print(
tokenizer.decode(
output[0, inputs.input_ids.shape[1]:],
skip_special_tokens=True,
)
)
do_sample=Trueenables sampling instead of greedy decoding.temperature=0.7controls the sharpness of the token distribution.top_p=0.9limits sampling to the most probable tokens covering approximately 90% of the probability mass.top_k=0disables the additional top-k filter.use_cache=Falseis required by this implementation.
For greedy decoding, use do_sample=False and omit temperature and top-p.
Keep the padded prompt length plus max_new_tokens within the 1024-token context limit.
Chat format
The saved chat template follows the format used during supervised fine-tuning:
- User and assistant message markers
- Two newlines between messages
- EOS token after assistant replies
Only user and assistant roles with string content are supported.
The model was not trained using system or tool roles.
Generation stops at EOS token 50256.
No hosted-inference-provider compatibility is claimed.
Causal language-model loss
For causal language-model loss, pass unshifted labels, usually a copy of input_ids, with padding labels set to -100.
The model shifts labels internally.
The training notebook's already-shifted target tensors must not be supplied directly.
Training
AIMS-1-128M was developed with a primary constraint of approximately $200 in compute expenditure.
The training objective was to explore whether a small Transformer language model could acquire useful representations of English syntax and linguistic structure without the compute resources normally associated with large-scale language-model training.
The project consists of two major stages:
- Pretraining — learning language patterns and English linguistic structure from the pretraining corpus.
- Supervised fine-tuning — teaching the pretrained model to follow the user/assistant conversational format.
The selected SFT checkpoint is:
- Checkpoint:
model_epoch_2.pth - Reference notebook:
SFT-HF+Simple-Everyday-Conversations.ipynb
The SFT configuration used AdamW with:
- Learning rate:
5e-5 - Weight decay:
0.1 - Training epochs: 5
The selected intermediate checkpoint may precede completion of the full SFT schedule.
Training data
The SFT training pipeline uses conversational data from:
HuggingFaceTB/everyday-conversations-llama3.1-2kAIMadeSimpleResearch/simple-everyday-conversations
The conversational data covers everyday interactions, factual questions, follow-up conversations, preferences, hobbies, social conversations, emotional-support-style conversations, greetings, playful interactions, predictions, personal questions, and impossible questions.
The distribution of the 15,000 conversational examples is shown above.
The model card should not be interpreted as claiming that the model was trained to memorize the factual information contained in these datasets. The primary goal of the SFT stage was conversational behavior and user/assistant formatting.
Evaluation
Evaluation should be performed against the specific AIMS-1-128M checkpoint published in this repository.
Numerical verification of the exported model checks loading and computational fidelity; it does not measure language-model quality.
Evaluation results should therefore be interpreted in the context of the model's small parameter count, limited training budget, and training objectives.
In particular, conversational performance should not be inferred solely from successful model loading or numerical equivalence between the original and exported implementations.
Limitations
AIMS-1-128M is a small experimental language model and has substantial limitations.
The model may:
- Produce factually incorrect information
- Generate incoherent or repetitive text
- Fail to maintain conversational context
- Struggle with reasoning and multi-step tasks
- Produce grammatically incorrect sentences
- Hallucinate information
- Exhibit biases present in its training data
- Perform inconsistently across prompts
- Have limited knowledge compared with substantially larger language models
The model's ability to produce conversational responses should not be interpreted as evidence that it possesses reliable factual knowledge or general reasoning capabilities.
The model's relatively small context window and lack of KV caching also limit generation efficiency.
Architecture and implementation
AIMS-1-128M uses a custom Transformer-based architecture.
Its design incorporates ideas from multiple existing language-model architectures, including:
- GPT — autoregressive Transformer language modeling
- Llama — modern decoder architecture and normalization/design principles
- Qwen — contemporary Transformer architectural design choices
The implementation combines these ideas into a custom architecture rather than reproducing any single reference model architecture.
The original model computation is preserved in notebook_components.py. The Hugging Face integration is implemented separately in modeling_baby_aims.py.
The runtime implementation adapts the buffer lifecycle required for model loading while preserving the original computational behavior.
License
AIMS-1-128M is released under the Apache License 2.0.
The Apache-2.0 license applies to the model code and model artifacts released in this repository. Users are responsible for complying with the licenses and terms of the datasets, pretrained components, libraries, and other third-party materials used in the training or development process.
Attribution
AIMS-1-128M is developed by AIMadeSimple Research as part of an effort to explore accessible, low-cost, open research into language-model training.
The project focuses on making the process of building and understanding language models accessible without requiring the infrastructure or budgets normally associated with large-scale AI development.