tiny-llm-java

An educational autoregressive language model implemented in pure Java.

This project is not a production-scale LLM. It is a compact, readable codebase for learning how tokenization, training, text generation, persistence, and serving can fit together in a small Java application.

• A small mlp language model
• A small decoder-style transformer language model
• A Spring Boot REST API
• A CLI for training and text generation
• Model save/load support with JSON snapshots
• BPE-style subword tokenization

The project predicts the next token from the previous contextSize tokens.

It currently supports two model types:

• mlp: embeddings + fully connected hidden layer + softmax
• transformer: token embeddings + position embeddings + multi-layer multi-head self-attention + residual connections + LayerNorm + feed-forward network + softmax

Recent training features include:

• mini-batch training
• Adam optimization
• AdamW-style decoupled weight decay
• label smoothing
• gradient clipping
• validation split
• early stopping
• ReduceLROnPlateau-style scheduling
• linear warmup
• cosine decay
• Transformer dropout

Trained models can be saved and restored together with tokenizer vocabulary and BPE merge rules. Legacy character-only snapshots and older single-layer Transformer snapshots remain backward compatible.

• Java 21+
• Maven 3.9+

Run tests:

mvn test

Train with the bundled sample corpus:

mvn exec:java

Run the Transformer explicitly:

mvn exec:java -Dexec.args="--model=transformer --epochs=120 --heads=4 --layers=2 --validation-split=0.1 --patience=8 --lr-patience=3 --lr-factor=0.5 --min-lr=0.0005 --warmup-epochs=5 --cosine-decay=true --clip=1.0 --dropout=0.1 --weight-decay=0.01 --label-smoothing=0.1 --prompt=Java --generate=160 --lr=0.005 --temperature=0.9"

Run the MLP model:

mvn exec:java -Dexec.args="--model=mlp --epochs=120 --prompt=Java --generate=160"

Train from an external corpus:

mvn exec:java -Dexec.args="--model=transformer --corpus=/absolute/path/to/corpus.txt --epochs=150 --prompt=AI --generate=200"

Save a trained model:

mvn exec:java -Dexec.args="--model=transformer --heads=4 --layers=2 --epochs=120 --save=./saved/my-transformer.json --prompt=Java"

Load a saved model and generate text:

mvn exec:java -Dexec.args="--load=./saved/my-transformer.json --prompt=Java --generate=120 --temperature=0.8"

• --model=TYPE: mlp or transformer
• --corpus=PATH: training text file; uses the bundled sample if omitted
• --epochs=N: number of training epochs
• --context=N: context window size
• --embed=N: embedding dimension
• --hidden=N: hidden dimension
• --heads=N: number of Transformer attention heads
• --layers=N: number of Transformer blocks
• --batch=N: mini-batch size
• --validation-split=FLOAT: fraction of data reserved for validation; 0 disables it
• --patience=N: early stopping patience; 0 disables it
• --lr-patience=N: patience before reducing the learning rate on a validation plateau; 0 disables it
• --lr-factor=FLOAT: learning-rate reduction factor where 0 < factor < 1
• --min-lr=FLOAT: minimum learning rate
• --warmup-epochs=N: number of warmup epochs
• --cosine-decay=BOOL: enables cosine decay after warmup
• --clip=FLOAT: global norm gradient clipping threshold; 0 disables it
• --dropout=FLOAT: Transformer dropout rate during training; 0 disables it
• --weight-decay=FLOAT: AdamW-style decoupled weight decay; 0 disables it
• --label-smoothing=FLOAT: label smoothing rate; 0 disables it
• --lr=FLOAT: learning rate
• --prompt=TEXT: generation prompt
• --generate=N: approximate number of generated characters or subword units
• --temperature=FLOAT: sampling temperature
• --seed=N: random seed
• --save=PATH: path to write a trained model snapshot
• --load=PATH: path to load a trained model snapshot

Start the Spring Boot API:

mvn spring-boot:run

Health check:

curl http://localhost:8080/api/health

Train a Transformer model:

curl -X POST http://localhost:8080/api/models/train \
  -H "Content-Type: application/json" \
  -d '{
    "modelType": "transformer",
    "epochs": 80,
    "contextSize": 6,
    "embeddingDim": 24,
    "hiddenDim": 64,
    "numHeads": 4,
    "numLayers": 2,
    "batchSize": 8,
    "validationSplit": 0.1,
    "earlyStoppingPatience": 8,
    "learningRateSchedulerPatience": 3,
    "learningRateSchedulerFactor": 0.5,
    "minLearningRate": 0.0005,
    "warmupEpochs": 5,
    "cosineDecayEnabled": true,
    "gradientClipNorm": 1.0,
    "dropoutRate": 0.1,
    "weightDecay": 0.01,
    "labelSmoothing": 0.1,
    "learningRate": 0.005,
    "seed": 42,
    "corpus": "Train a tiny language model in Java. Try generating text in Java."
  }'

Generate text from a trained model:

curl -X POST http://localhost:8080/api/models/generate \
  -H "Content-Type: application/json" \
  -d '{
    "modelId": "YOUR_MODEL_ID",
    "prompt": "Java",
    "generateLength": 120,
    "temperature": 0.8,
    "seed": 44
  }'

List saved models:

curl http://localhost:8080/api/models

Get one model summary:

curl http://localhost:8080/api/models/YOUR_MODEL_ID

Explicitly save an in-memory model:

curl -X POST http://localhost:8080/api/models/YOUR_MODEL_ID/save

Models trained through the Spring Boot API are automatically saved to data/models by default and reloaded on the next startup.

Snapshots include:

• model weights
• tokenizer vocabulary
• BPE merge rules
• training configuration
• training loss history

This makes it possible to restart the application and continue generating text with the same tokenization behavior.

Example style of generated output:

Java builds a small language model.
The model learns local text patterns and predicts the next token step by step.

Because the bundled corpus is small, generated text is short and unstable. Better output usually comes from more data, more training time, and careful tuning of width, depth, and optimization settings.