CRIA-LM-75M

Cria (noun): a baby llama, alpaca, vicuña, or guanaco. Pronounced ˈkrē-ə.

~ Merriam-Webster

cRia-LM-75M is a 75.7M-parameter base language model built as a Relaxed Recursive Transformer (RRT). It uses a shared 11-layer recurrent block evaluated twice, with pass-specific LoRA parameters on the second traversal.

Training was carried out in three stages. Stage 1 established the 2K base model over 10B tokens. Stage 2 continued training with a 2B-token budget and a capability-focused data curriculum; the released Stage 2 checkpoint is step 10,000, corresponding to about 1.31B continuation tokens. Stage 3 extended the context window from 2,048 to 4,096 tokens with a 50M-token run on codelion/sutra-1B. The released checkpoint continues from that long-context stage with 35 learned residual-gain scalars and one epoch over the HellaSwag, ARC-Easy, ARC-Challenge, PIQA, SciQ and SWAG training splits.

This is a base model. It has not been instruction tuned and does not use a chat template. The Transformer layer design follows HuggingFaceTB/SmolLM2-135M, with QK-Norm added for training stability.

Model Details

Architecture

cRia-LM-75M uses 13 unique Transformer layers arranged as:

1 prelude + (11 shared layers x 2 recurrent passes) + 1 coda

This gives an effective depth of 24 Transformer layers while storing the main parameters for only 13 unique layers.

The 11-layer recurrent block is shared across both passes. On the second pass, the recurrent linear projections receive rank-172 LoRA updates. The two traversals therefore share the base weights while retaining pass-specific capacity.

Each Transformer layer uses:

• hidden size 576
• 9 query heads and 3 KV heads with grouped-query attention
• head dimension 64
• SwiGLU feed-forward network with intermediate size 1,536
• RMSNorm
• rotary position embeddings
• QK-Norm

Factorized Tied Embedding

The input embedding and language-model readout are tied through a rank-210 factorization:

49,152 x 210
210 x 576

The same factors are used for token lookup and output projection. This keeps the full 49,152-token vocabulary while reducing the parameter cost of the embedding and output head.

Training

Stage 1 — Base pretraining

Stage 1 trained the model for 10B tokens at a sequence length of 2,048. The objective combined next-token cross-entropy with logit-level knowledge distillation from HuggingFaceTB/SmolLM2-360M, with dynamic CE/KD scale balancing.

The Stage 1 sampling mix was:

Stage 1 used Muon for matrix parameters and AdamW for the remaining parameter groups. The learning-rate schedule was warmup-stable-decay: 1% warmup, decay beginning at 80% of the run, and the final 20% used for decay.

Stage 2 — Capability mid-training

Stage 2 used a 2B-token training budget at 2,048-token context. The released Stage 2 checkpoint is step 10,000, or 1,310,720,000 tokens at 131,072 tokens per optimizer step.

The phase used two-teacher distillation: same-tokenizer KL from HuggingFaceTB/SmolLM2-360M and cross-tokenizer supervision from Qwen/Qwen3.5-0.8B-Base using X-Token. The run used the H-KL cross-tokenizer objective after the tokenizer-coverage audit, with dynamic CE/KD balancing retained around the combined distillation loss.

Phase 2A mix

Phase 2A is the capability-injection mix. It remains active through 65% of the planned Stage 2 run. The original configuration assigned 8% to Nemotron QA, but that source was unavailable when the run started. The loader dropped it and renormalized the remaining weights; the table below shows the effective mix that was actually sampled.

Phase 2B mix

Phase 2B is the cooldown mix. It removes the noisier sources and shifts more weight toward higher-quality web, educational, encyclopedic, science, and math data.

Stage 2 WSD schedule

Stage 1 had already completed a full decay, so Stage 2 re-warmed to a lower peak rather than returning to the original Stage 1 learning rate.

After the stable region, the LR multiplier follows the run's square-root WSD decay:

lr_scale = max(0.02, 1 - sqrt(decay_progress))

Because the data-mix switch occurs at 65%, the final portion of Phase 2A already runs under LR decay. The selected step-10,000 checkpoint lands just after the switch to Phase 2B.

Stage 3 — 2K to 4K context extension

Stage 3 increases the supported context length from 2,048 to 4,096 tokens. The run was configured for 50M tokens on codelion/sutra-1B, with the long-context checkpoint used for this release selected at step 100, roughly 15M tokens into the run. The released model then continued with 35 learned residual-gain scalars for one epoch over the HellaSwag, ARC-Easy, ARC-Challenge, PIQA, SWAG, and SciQ training splits.

The context-length curriculum used the following token fractions:

The training path uses Apple's Cut Cross Entropy directly against cRia's factorized tied output head. The cce_exact implementation was used to avoid gradient filtering while reducing the memory cost of materializing full-vocabulary logits for CE computation.

RoPE theta remains 100,000 and no RoPE scaling is applied. The model has no learned absolute position table, so moving from 2K to 4K does not add parameters.

Evaluation

The complete standard suite was evaluated zero-shot with FP32 model tensors and FP32 log-softmax. CUDA TF32 matmuls were enabled.

The average is the unweighted arithmetic mean of the 14 standard accuracy scores, ArithMark-3 normalized accuracy, and BananaMind raw accuracy.

Comparison with established small language models

We don't make any claims about inferiority or superiority of any model. While we took care of avoiding leakage from validation and test sets, we DID use the train splits of several of these benchmarks in our pre-training corpus. Compared to other models who had a longer training regime on non-benchmark data, it is possible real performance may lag behind what these benchmarks indicate (LAMBADA perplexity and MMLU already shows this trend somewhat).

Inference Performance

Measurements used FP32 weights, batch size 1, synchronized CUDA medians, PyTorch 2.12.0+cu130, WSL2, and an RTX 4070 Ti.

For repeated unpadded greedy requests of one fixed shape, the reusable CUDA graph runner reached 318.77 tokens/s with a 128-token prompt and 64 generated tokens. The graph is inference-only, is fixed to its configured batch/prompt/output shape, and must be rebuilt after training changes model parameters.

Usage

The model ships with a custom Transformers implementation, so trust_remote_code=True is required for the model. The tokenizer is loaded directly from the cRia repository.

model.generate() uses the recursion-wise KV cache by default. use_cache=True is shown explicitly below so the inference path is clear.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "sz14/cRia-LM-75M"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=torch.float32,
).to(device).eval()

prompt = "In mathematics, a vector is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=False,
        use_cache=True,
    )

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

Explicit KV-cache usage

For custom decode loops, a cache can also be created and passed through past_key_values directly:

prompt = "In mathematics, a vector is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
cache = model.new_cache(batch=inputs.input_ids.shape[0])

with torch.no_grad():
    # Prefill the prompt and populate the cache.
    out = model(
        **inputs,
        past_key_values=cache,
        use_cache=True,
    )

    # Decode one token using only the new token plus cached K/V states.
    next_token = out.logits[:, -1].argmax(dim=-1, keepdim=True)
    out = model(
        input_ids=next_token,
        past_key_values=cache,
        use_cache=True,
    )

print("cached positions:", cache.pos)

Pass use_cache=False to model.generate() to fall back to full-prefix recomputation.

KV-cache implementation

The recurrent block shares weights across passes, but the two passes do not share keys and values. The second traversal uses its own LoRA-adapted projections, so each recurrent pass receives a separate cache slot. With the current architecture this produces 24 cache slots, matching the model's 24 effective layers.

The current implementation supports:

• cached prompt prefill and token-by-token decode
• batched generation with left padding
• attention masks over the full cached key length
• RoPE positions derived from the cumulative attention mask for padded batches
• beam-search cache reordering
• explicit cache cropping/resetting
• a hard 4,096-token cache/context bound with clear errors beyond the configured window

A standalone model.generate_cached(...) helper is also available for direct incremental decoding outside the standard Transformers generation path.

Intended Use

cRia-LM-75M is intended primarily for:

• research on recursive parameter sharing
• experiments with compact language models
• further pretraining and domain adaptation
• supervised fine-tuning
• small language-model backbones
• architecture and knowledge-distillation research

Because this is a base model, prompts are treated as ordinary text continuation rather than instructions.

Limitations

cRia-LM-75M is a small base language model and should not be expected to match substantially larger pretrained models.

Known limitations include:

• weak multi-step reasoning
• limited factual knowledge
• potential factual errors and hallucinations
• no instruction-following training
• no chat template
• English-focused training
• a 4,096-token maximum context window

The model should not be treated as a reliable source of factual information or used without additional validation in high-stakes applications.

Architecture and Training Lineage

The Transformer layer design follows HuggingFaceTB/SmolLM2-135M, released by Hugging Face under the Apache 2.0 license.

HuggingFaceTB/SmolLM2-360M was used for same-tokenizer logit distillation in the initial training and Stage 2. Stage 2 additionally used Qwen/Qwen3.5-0.8B-Base for cross-tokenizer X-Token distillation.

The recursive parameter-sharing approach is based on:

Bae et al., "Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA," arXiv:2410.20672.

Stage 2 cross-tokenizer distillation uses:

Sreenivas et al., "X-Token: Projection-Guided Cross-Tokenizer Knowledge Distillation," arXiv:2605.21699.

Stage 3 uses Apple's Cut Cross Entropy implementation from Wijmans et al., "Cut Your Losses in Large-Vocabulary Language Models," arXiv:2411.09009v2.

License

cRia-LM-75M is released under the Apache License 2.0.