yehezkielhaganta/NimoLM-244M-base

NimoLM-244M

NimoLM-244M is a compact decoder-only language model developed from scratch, designed to explore hybrid attention architectures that combine Gated Linear Attention (GLA) with conventional causal self-attention.

The current model is approximately 244M parameters and uses a custom Byte-Level BPE tokenizer with a vocabulary size of 65,536 (2¹⁶).

Status: Pretraining in progress

The reported training metrics in this README represent the current training state and are not the final results.

---

Model Overview

Component	Configuration
Architecture	Decoder-only Transformer
Parameters	~244M
Vocabulary	65,536
Hidden Size	1,024
Attention Heads	16
Head Dimension (SA)	64
FNN Dimension	3,072
Activation	SwiGLU
Normalization	RMSNorm
Positional Encoding	RoPE
Attention Architecture	Hybrid GLA + Causal Self-Attention
Hybrid Blocks	3
GLA per Hybrid Block	3
Self-Attention per Hybrid Block	1
Total GLA Blocks	9
Total Self-Attention Blocks	4
GLA Mode	Chunk
GLA Key Dimension	512
GLA Value Dimension	1,024
GLA Heads	16
Bias	Disabled for main Linear layers
Weight Tying	Input embedding ↔ LM Head
Linear Attention Library	flash-linear-attention[cuda]

The architecture follows a hybrid design where each hybrid block contains:

GLA
 ↓
GLA
 ↓
GLA
 ↓
Self-Attention

This design intentionally combines the efficiency of linear attention with the global token interaction capability of conventional self-attention.

Model Architecture

NimoLM uses a 3:1 GLA-to-self-attention ratio inside each hybrid block.

The GLA layers are responsible for the majority of sequence processing, while the self-attention layer provides periodic full token-to-token interaction.

Conceptually, the architecture can be represented as:

NimoForCausalLM
│
├── Token Embedding
│   └── 65,536 → 1,024
│
├── Initial Self-Attention Block
│
├── 3 × Hybrid Block
│   ├── 3 × GLA Block
│   └── 1 × Self-Attention Block
│
├── Final RMSNorm
│
└── LM Head
    └── 1,024 → 65,536

Tokenizer

NimoLM uses a custom Byte-Level BPE tokenizer trained specifically for this project.

Training Corpus

The tokenizer was trained using a subset of the FineWeb-Edu 10BT dataset.

Approximately 970,000 documents were selected from the FineWeb-Edu sample/10BT dataset to construct the tokenizer training corpus.

The tokenizer vocabulary was intentionally set to: 65,536 = 2¹⁶ tokens

The vocabulary size was selected primarily to provide a convenient power-of-two vocabulary size for the model architecture.

Tokenizer Configuration

Algorithm: Byte-Level BPE
Vocabulary size: 65,536
Normalization: NFKC
Pre-tokenizer: ByteLevel
Decoder: ByteLevel
Special tokens:
<|pad|>
<|eos|>
<|system|>
<|user|>
<|assistant|>

The tokenizer configuration confirms ByteLevel pre-tokenization with add_prefix_space=false, while the decoder uses ByteLevel decoding.

The tokenizer vocabulary begins with the project's special tokens followed by the Byte-Level BPE vocabulary.

Tokenizer Evaluation

A preliminary comparison against the GPT-2 tokenizer was performed on the FineWeb-Edu corpus.

The custom tokenizer showed approximately 4% fewer tokens than the GPT-2 tokenizer on the evaluated data.

This suggests that the custom vocabulary provides slightly better token compression for the corpus used in this project.

This tokenizer evaluation is preliminary. A more comprehensive evaluation across multiple datasets and domains is planned for future experiments.

Pretraining

NimoLM is pretrained with a next-token prediction objective.

Dataset

The main pretraining dataset is:

FineWeb-Edu 10BT

The training data is stored as pre-tokenized NumPy shards and separated into:

train/
validation/

The training dataloader uses a sequence length of 2,048 tokens.

Training Objective

The model is trained using standard causal language modeling:

x = [t₁, t₂, ..., tₙ]
y = [t₂, t₃, ..., tₙ₊₁]

The loss function is cross-entropy over the next-token prediction task.

Training Configuration

Parameter	Value
Optimizer	AdamW
Learning Rate	6e-4
Betas	(0.9, 0.95)
Epsilon	1e-8
Weight Decay	0.1
Peak LR	6e-4
Minimum LR	6e-5
Warmup	500 steps
LR Schedule	Linear Warmup + Cosine Decay
Gradient Clipping	1.0
Mixed Precision	FP16
Effective Batch Size	524,288 tokens
Sequence Length	2,048
Target Training Steps	18,088

The optimizer is AdamW with separate weight-decay handling for matrix parameters and non-matrix parameters.

The learning-rate schedule uses linear warmup followed by cosine decay toward 10% of the maximum learning rate.

The effective batch size is configured as approximately 524K tokens per optimizer step.

Current Pretraining Results

Training is currently ongoing.

The current training log contains 10,750 steps out of the planned 18,088 steps.

At the latest recorded step:

Metric	Current
Training Step	10,750
Train Loss	3.0898
Validation Loss	3.0361
Learning Rate	2.606e-4
Tokens / Optimizer Step	524,288
Approx. Tokens Processed	5.64B

Validation

The current best recorded validation loss is:

Validation Loss: 3.0361
Perplexity: ~20.82

The validation loss is computed periodically using a separate validation split. The current training configuration evaluates every 250 steps using 50 validation batches.

Important: These are intermediate results. Training has not reached the planned 18,088 steps, so the final validation loss and perplexity may change.

Learning curve

Hugging Face Integration

NimoLM is packaged for use with the Hugging Face Transformers ecosystem.

Example:

pip install flash-linear-attention[cuda]

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "yehezkielhaganta/NimoLM-244M-base"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

model = model.cuda()

inputs = tokenizer(
    "rose is red violets are blue",
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=150,
        temperature=1,
        top_k=50,
        do_sample=True,
    )

print(tokenizer.decode(output[0]))

The model uses custom modeling code to expose the NimoLM architecture through the Transformers API.

Research Motivation

The main objective of NimoLM is to investigate whether a relatively small language model can benefit from combining linear and full attention mechanisms.

The central research question is:

Can a hybrid GLA + self-attention architecture provide a useful efficiency/performance trade-off at the ~244M parameter scale?

Rather than relying entirely on full self-attention, NimoLM allocates most of its attention layers to GLA while retaining periodic self-attention layers.

This creates a practical experimental setting for studying:

• Hybrid attention ratios
• Training efficiency
• Memory usage
• Language modeling performance
• Long-context behavior
• Scaling behavior at small model sizes

Acknowledgements

NimoLM is inspired by recent research into efficient sequence modeling, including Gated Linear Attention, hybrid attention architectures, and modern decoder-only language models.

The current GLA implementation makes use of the Flash Linear Attention project for efficient GPU kernels and training-ready linear-attention components.