PhaGen

PhaGen is a hierarchical masked diffusion model for bacteriophage DNA. It generates DNA on a masked canvas instead of left-to-right, so the same model produces sequences, representations and masked conditional scores.


Parameters	152,421,601
Stages	3, dimensions 512 / 256 / 196, depths 8 / 8 / 8
Token canvas	131,072 (`max_seq_len` 128 x 64 x 16)
Vocabulary	9 tokens: `[BOS]` A T C G `[EOS]` `[MASK]` `[PAD]` `[UNK]`
Precision	BF16 safetensors
Model type / classes	`phagen` / `PhaGenModel`, `PhaGenConfig`

Load

Requires PyTorch 2.5, Transformers 4.54.1, einops, accelerate and safetensors. The model code is custom, so trust_remote_code=True is required; review the supplied Python files first. The public classes are PhaGenModel and PhaGenConfig, implemented in modeling_phagen.py and configuration_phagen.py.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "YusenWithR/PhaGen"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True,
    attn_implementation="sdpa", torch_dtype=torch.float32,
).eval()

inputs = tokenizer("ACGTACGTAC", return_tensors="pt", add_special_tokens=False)
with torch.inference_mode():
    out = model(**inputs, is_causal=False)       # the tokenizer supplies attention_mask
print(out.logits.shape)                          # [1, 10, 9]
print([tuple(h.shape) for h in out.hidden_states])
# [(1, 2, 512), (1, 1, 65, 256), (1, 1, 64, 17, 196)] for this 10-base input.
# Inputs are padded to a multiple of 1024; the first stage grows for longer inputs.

AutoModelForCausalLM is the registered auto class for loading convenience. PhaGen is a masked diffusion model, not an autoregressive LM, and it does not implement generate() in the autoregressive sense.

Representations

Use model.embed for the paper's 964-dimensional sequence representation:

batch = tokenizer(
    ["ACGTACGTAC", "GATTACA"], padding=True,
    return_tensors="pt", add_special_tokens=False,
).to(model.device)
embeddings = model.embed(
    input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
)
print(embeddings.shape)  # torch.Size([2, 964]), float32, input order preserved

Generation

Diffusion sampling runs through diffusion_generate with the MDMGenerationConfig supplied in generation_utils.py:

from transformers.dynamic_module_utils import get_class_from_dynamic_module

MDMGenerationConfig = get_class_from_dynamic_module(
    "generation_utils.MDMGenerationConfig", repo)

torch.manual_seed(42)
config = MDMGenerationConfig(
    max_length=256, steps=32, alg="entropy", alg_temp=0.1,
    temperature=1.1, top_p=0.90, top_k=4,
    mask_token_id=tokenizer.mask_token_id, pad_token_id=tokenizer.pad_token_id,
    return_dict_in_generate=True,
)
prefix = torch.tensor([[tokenizer.bos_token_id]])
with torch.inference_mode():
    result = model.diffusion_generate(
        input_ids=prefix, attention_mask=torch.ones_like(prefix), generation_config=config)
print(tokenizer.decode(result.sequences[0], skip_special_tokens=True).replace(" ", "")[:60])

That runs in a few seconds on CPU. The paper production setting keeps the same sampling values and raises the budget to max_length=131072 with steps=256 and seed 42, which needs a GPU. alg accepts origin, entropy, topk_margin, p2 and maskgit_plus.

The canvas is a token budget, not an output DNA length. Extract complete records from the token IDs before removing special tokens:

# One token canvas at a time; returns ACGT strings in their original order.
contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=1)
for index, sequence in enumerate(contigs, start=1):
    print(f">contig_{index}\n{sequence}")

# The paper cohort retains complete records with at least 2000 DNA bases.
paper_contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=2000)

min_length counts DNA bases, excluding BOS/EOS. Only complete [BOS] ACGT... [EOS] spans are returned. A new BOS restarts an unfinished span; an internal PAD, MASK, UNK or other non-ACGT token invalidates the span. Empty, short and trailing unclosed spans are dropped. Short canvases may return an empty list. Decoding the entire canvas with skip_special_tokens=True is only a preview: it can join separate records and discard evidence of invalid content.

For several generated canvases, call the helper for each row and retain the row index alongside its records. The same extraction function is available without loading the weights via sequence_utils.extract_contigs(token_ids, tokenizer, min_length).

Training data

98,665 phage genomes (4,800,631,119 bases) assembled from MGV (52,836), GPD (29,768) and NCBI (16,061).

Intended use and limitations

Intended for research on phage genome modelling: unconditional generation, sequence representations for downstream probes, and masked conditional scoring.

The model is trained on phage DNA; behaviour on other genomes is untested.
Short runs can extract zero complete contigs; that is a property of the extraction protocol, not a failure.

Acknowledgement

PhaGen builds on the ideas of several outstanding open-source projects.

License details and upstream notices are in NOTICE.md.

Downloads last month: 4

Safetensors

Model size

0.2B params

Tensor type

BF16