tinygemma4

tinygemma4 is a deliberately tiny, text-only Gemma 4 model trained from scratch on TinyStories. It is intended for architecture compatibility checks, inference-engine testing, and small-scale language-model experiments.

This is not a useful assistant model. It was trained on simple synthetic stories and should be expected to produce short, child-story-like completions with limited coherence.

Model Details

• Architecture: Gemma4ForCausalLM
• Parameters: 4,964,752
• Vocabulary: 8192-token byte-level BPE
• Context length in config: 2048
• Training block size: 256
• Hidden size: 128
• Per-layer input hidden size: 16
• Layers: 12
• Attention heads: 4
• KV heads: 1
• Head dimension: 32
• MLP intermediate size: 384
• Sliding window: 128
• Full attention layers: 4, 8, 12
• Embeddings: tied
• MoE: disabled
• Multimodal components: none
• Weight dtype: float32
• Tensor format: safetensors

The checkpoint is saved in ordinary Hugging Face Transformers format. Any runtime with a correct Gemma 4 text implementation and support for these small dimensions should be able to load it.

Training

• Dataset: roneneldan/TinyStories
• Training file: TinyStoriesV2-GPT4-train.txt
• Validation file: TinyStoriesV2-GPT4-valid.txt
• Final training step: 300000
• Optimizer: Muon + AdamW (Muon for non-embedding matrix parameters; AdamW for embeddings, norms, and scalar parameters)
• Learning-rate schedule: cosine decay from 3.5e-4 to 3.5e-5, with 1,000 warmup steps
• Batch size: 32
• Random seed: 1337
• Hardware: AMD Radeon RX 9070 XT, ROCm PyTorch for Windows
• Training dtype: bf16 autocast where available

Evaluation

Validation was measured during training on held-out TinyStories text with the local training script:

• Validation loss: 1.3825
• Validation perplexity: 3.98

These numbers are only for this training setup. Evaluation used 50 shuffled validation batches of 32 sequences at a 256-token block size. They are not general language-understanding benchmarks.

The final checkpoint's best tested sampling settings were temperature 0.45, top-p 0.90, top-k 50, and repetition penalty 1.03. On a small fixed sampling suite, the model remained fluent but still showed repetition, contradictions, pronoun errors, and weak long-range causality.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "ApexDevelopment/tinygemma4"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")

inputs = tokenizer("Once upon a time,", return_tensors="pt")
outputs = model.generate(
    **inputs,
    max_new_tokens=80,
    do_sample=True,
    temperature=0.45,
    top_p=0.90,
    top_k=50,
    repetition_penalty=1.03,
    pad_token_id=tokenizer.pad_token_id,
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Limitations

• The model is tiny and heavily capacity-limited.
• It is trained only on synthetic TinyStories text.
• It is not instruction tuned.
• It is not safety tuned.
• It can repeat, contradict itself, confuse characters or objects, and produce malformed story fragments.
• Its coherent range is usually much shorter than the configured 2,048-token context length.
• It should be used for experimentation and testing, not production.

Data and License Notes

The training dataset card lists TinyStories under cdla-sharing-1.0. This model was trained from scratch; it does not contain Gemma weights from Google or weights from TinyLLama-v0.

Weights are released under the license declared in the metadata above. Users are responsible for checking whether their intended use is compatible with the dataset license and applicable law.

Inspiration

This project was inspired by Maykeye/TinyLLama-v0, but uses a Gemma 4 text configuration instead of Llama.