modilify/Modilify-Mk2

Modilify Mk2

Modilify Mk2 is a text diffusion model with learned trajectory memory. Built on the DiffusionGemma 26B-A4B text backbone, it develops a response through repeated revisions of a rolling text canvas. Its memory follows both the current revision process and the content already committed to the response.

This repository provides the complete model weights trained on 78M tokens, tokenizer, and a PyTorch / Transformers runtime.

Architecture

Modilify Mk2 works on a 256-token canvas. During each denoising pass, it revises candidate tokens together and uses their confidence and uncertainty to select a prefix to commit. The canvas then advances to continue the response.

Two timescales of GDN2 memory connect these passes. Working memory tracks changes within the canvas, with states for individual positions and the canvas as a whole. Persistent memory incorporates experience from committed tokens and carries it into later parts of the response. The model learns these memory components alongside dense and expert LoRA adapters.

Training supervises response tokens across the valid canvas, with additional objectives for confidence calibration, prefix readiness, and response termination.

Property	Value
Backbone	DiffusionGemma 26B-A4B, text only
Adaptation	Dense and expert LoRA
Memory	Dual-timescale GDN2 trajectory memory
Canvas size	256 tokens
Training step	4400
Adaptation tokens	78,332,410 (prompt + response, excluding padding)
Weight precision	BF16, with selected memory parameters in FP32
Weight files	16 Safetensors shards, approximately 48.23 GiB
Framework	PyTorch / Hugging Face Transformers

Quickstart

From the downloaded model directory, use Python 3.12 and install the dependencies:

python -m pip install -r requirements.txt
python inference.py --prompt "Why is the sky blue?" --stream

The CLI selects CUDA, Apple Silicon MPS, or CPU according to availability. Choose a device explicitly with --device cuda, --device mps, or --device cpu. In addition to the weights, allow memory for the prompt cache, trajectory states, and generation workspace. Local inference has been exercised on CPU and MPS.

Use with Transformers

The model loads through AutoModelForCausalLM. Set model_path to the downloaded directory or Modilify/Modilify-Mk2 on the Hub.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "./"
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True,
    dtype="auto",
    device_map=device,
    attn_implementation="sdpa",
).eval()

messages = [{"role": "user", "content": "Explain how rainbows form."}]
inputs = tokenizer.apply_chat_template(
    messages,
    enable_thinking=True,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
).to(model.device)

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=256,
        max_denoising_steps=256,
        seed=42,
    )

response_ids = output.sequences[0, inputs.input_ids.shape[1]:]
print(tokenizer.decode(response_ids, skip_special_tokens=True))

Thinking and conversations

The chat template supports system, user, and assistant turns. Pass the conversation as messages and set enable_thinking=True to enable reasoning before the answer. The CLI enables thinking by default; --no-thinking disables it.

Generation controls

max_new_tokens sets the response length budget, while max_denoising_steps sets the number of canvas revisions. Both default to 256. Increase the budgets for longer responses; reaching either limit can leave a response unfinished. output.stop_reason reports how generation ended.

Use seed for a single request or seeds=[...] for batch requests. For batches, set tokenizer.padding_side = "left" and pass the attention mask returned by the tokenizer. The returned output.sequences includes the input prompt, and output.generated_lengths gives each response's token count.

Generation runs through the included diffusion implementation. The weights include the backbone, adapters, and memory components; from_pretrained loads them together. dtype="auto" preserves the model's mixed precision settings.

Use and limitations

Modilify Mk2 is released for non-commercial research and experimentation with text diffusion, iterative response refinement, and learned generation memory. It accepts text input and produces text, including an optional reasoning channel.

Broad capability and safety benchmark results are not yet available for this checkpoint. Generated responses can contain factual errors, repetition, or incomplete reasoning. Review outputs before using them for consequential decisions. The adaptation dataset is not included in this release.

License and credits

The model builds on Google DeepMind's DiffusionGemma and integrates with Hugging Face Transformers. Attribution and upstream licenses are listed in NOTICE.md.

Modilify's contributions are available under the Modilify Non-Commercial Model License 1.0. Commercial use requires a separate agreement.