uLM-28M

uLM-28M is a 28.8M-parameter decoder-only causal language model released with the paper Micro Language Models Enable Instant Responses.

It is designed for instant local response initiation: the model generates a short, contextually grounded opening of approximately 4–8 words on a resource-constrained device, and a larger cloud model continues the response.

For representative behavior, use the tokenizer's chat template, sampling, and short-prefix generation as shown below.

Model Details

The model was trained from scratch on chat-style data. Further training and evaluation details are provided in the accompanying paper.

Intended Generation Pattern

uLM-28M is intended to generate only the first few words of a response:

User request
    ↓
uLM-28M generates a 4–8 word local prefix
    ↓
A larger model continues from that committed prefix

For example:

Question: What is the capital of France?
Local prefix: The capital of France is Paris.

The exact output may vary because the recommended decoding configuration uses sampling.

Compatibility

This release is validated with Transformers 4.x.

Transformers 5.x is not currently supported because of changes to custom-model weight tying and cache interfaces.

Install a supported version with:

pip install "transformers>=4.45,<5" torch safetensors sentencepiece

Usage

This repository contains a custom Transformers architecture, so loading it requires trust_remote_code=True.

from __future__ import annotations

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    StoppingCriteria,
    StoppingCriteriaList,
)

REPO_ID = "Sensente/uLM-28M"
MAX_PREFIX_WORDS = 8


class PrefixWordStoppingCriteria(StoppingCriteria):
    """Stop once the generated continuation contains enough words."""

    def __init__(
        self,
        tokenizer,
        prompt_length: int,
        max_words: int,
    ):
        self.tokenizer = tokenizer
        self.prompt_length = prompt_length
        self.max_words = max_words

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
        **kwargs,
    ) -> bool:
        generated_ids = input_ids[0, self.prompt_length:]
        generated_text = self.tokenizer.decode(
            generated_ids,
            skip_special_tokens=True,
        )
        return len(generated_text.strip().split()) >= self.max_words


def choose_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


device = choose_device()

tokenizer = AutoTokenizer.from_pretrained(
    REPO_ID,
    trust_remote_code=True,
    use_fast=False,
)

model = AutoModelForCausalLM.from_pretrained(
    REPO_ID,
    trust_remote_code=True,
    dtype="auto",
).to(device).eval()

conversation = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

rendered_prompt = tokenizer.apply_chat_template(
    conversation,
    tokenize=False,
    add_generation_prompt=True,
)

inputs = tokenizer(
    rendered_prompt,
    return_tensors="pt",
    add_special_tokens=False,
    return_token_type_ids=False,
)
inputs.pop("token_type_ids", None)

inputs = {
    name: tensor.to(device)
    for name, tensor in inputs.items()
}

prompt_length = inputs["input_ids"].shape[1]

stopping_criteria = StoppingCriteriaList(
    [
        PrefixWordStoppingCriteria(
            tokenizer=tokenizer,
            prompt_length=prompt_length,
            max_words=MAX_PREFIX_WORDS,
        )
    ]
)

torch.manual_seed(0)

with torch.no_grad():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=32,
        do_sample=True,
        temperature=0.5,
        top_p=0.9,
        use_cache=True,
        stopping_criteria=stopping_criteria,
        pad_token_id=tokenizer.eos_token_id,
    )

generated_ids = output_ids[0, prompt_length:]

prefix = tokenizer.decode(
    generated_ids,
    skip_special_tokens=True,
).strip()

prefix = " ".join(prefix.split()[:MAX_PREFIX_WORDS])

print(prefix)

Prompting and Decoding Notes

For representative behavior:

1. Use tokenizer.apply_chat_template(...).
2. Set add_generation_prompt=True.
3. Tokenize the rendered prompt with add_special_tokens=False.
4. Do not pass token_type_ids.
5. Use sampling with temperature=0.5 and top_p=0.9.
6. Stop after approximately 4–8 generated words.

The following is useful for structural testing but is not representative of the intended application:

model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=32,
)

Long greedy continuations may become repetitive because uLM-28M was optimized for short response initiation rather than standalone long-form generation.

Custom Transformers Code

The repository includes:

configuration_swen.py
modeling_swen.py

The model is registered through the auto_map field in config.json:

{
  "AutoConfig": "configuration_swen.SwenConfig",
  "AutoModelForCausalLM": "modeling_swen.SwenForCausalLM"
}

Because trust_remote_code=True executes Python code from the model repository, users should review the implementation and pin a specific revision in production:

REVISION = "<commit-sha>"

model = AutoModelForCausalLM.from_pretrained(
    REPO_ID,
    trust_remote_code=True,
    revision=REVISION,
    dtype="auto",
)

Conversion and Release Validation

The Hugging Face release was converted from the original uLM checkpoint and validated both locally and after upload to the Hugging Face Hub.

The validation included:

• comparison of 75 source and converted state tensors;
• maximum absolute weight difference of 0.0;
• comparison of incremental KV-cache logits with full recomputation;
• identical generation with use_cache=True and use_cache=False;
• token-identical demo-style generation between the original implementation and the Hugging Face implementation;
• clean remote loading through AutoConfig, AutoTokenizer, and AutoModelForCausalLM.

For the test question:

What is the capital of France?

the original implementation and the Hugging Face release both generated:

The capital of France is Paris.

with identical token IDs:

[1168, 5118, 297, 6850, 327, 4462, 16, 2]

These checks establish conversion and packaging fidelity. They are not a substitute for task-level quality evaluation.

Intended Uses

uLM-28M is intended for research and experimentation involving:

• low-latency local response initiation;
• hybrid local-cloud language generation;
• resource-constrained edge deployment;
• committed-prefix generation;
• semantic handoff from a micro model to a larger model;
• studies of perceived response latency.

Out-of-Scope Uses

uLM-28M is not intended to:

• independently produce long-form answers;
• replace a general-purpose instruction-tuned assistant;
• provide reliable factual answers without verification;
• make medical, legal, financial, or safety-critical decisions;
• generate high-stakes content without human review.

Limitations

uLM-28M is an extremely small language model. Its capabilities are substantially more limited than those of larger language models.

Known limitations include:

• limited factual knowledge;
• limited reasoning ability;
• weak long-form coherence;
• possible repetition during long or greedy generation;
• sensitivity to prompt formatting and decoding parameters;
• possible incorrect, misleading, biased, or inappropriate outputs;
• limited safety alignment compared with larger production assistants.

The model's short prefix is immediately committed in the intended local-cloud framework. Downstream systems should therefore implement continuation and recovery strategies for incorrect or poorly aligned prefixes.

Training and Evaluation

For details on the model family, training setup, datasets, evaluation tasks, local-cloud continuation framework, and error-recovery methods, see:

• Paper: Micro Language Models Enable Instant Responses
• Code and demo: Sensente/micro_language_model_project

License

The model weights and custom model implementation in this repository are released under the MIT License.

See the LICENSE file for the full license text.

Third-party dependencies remain subject to their respective licenses.

Citation

@misc{cheng2026micro,
  title         = {Micro Language Models Enable Instant Responses},
  author        = {Wen Cheng and Tuochao Chen and Karim Helwani and
                   Sriram Srinivasan and Luke Zettlemoyer and
                   Shyamnath Gollakota},
  year          = {2026},
  eprint        = {2604.19642},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2604.19642}
}