---
language: en
license: apache-2.0
library_name: pytorch
pipeline_tag: text-generation
tags:
  - diffusion
  - masked-diffusion
  - mdlm
  - psychology
  - mental-health
  - small-model
  - on-device
---

# psyche-diffuse-66m (production) & psyche-diffuse-120m

66.8M-parameter **masked diffusion language model** (MDLM, DiT backbone) for English psychology / mental-health dialogue, trained on one RTX 3090. The 120M sibling is published to make the paper's scaling null independently checkable. Paper: *Small-Scale Text Diffusion for a Psychology Assistant* (Denis Kim, 2026), PDF in the GitHub repo.

## ⚠️ Intended use & safety

Research artifact for studying small-scale text diffusion. **Not a therapist, not a medical device, not for crisis use.** Outputs are frequently vague, repetitive, and can be factually wrong (paper §9: blind homonym-drift floor 25–32%). For crisis support call 988 (US) or local equivalents. No real conversations or PII in training data.

## Files

| file | what |
|---|---|
| `model_66m_clean_sft.pt` | production checkpoint (DiT 512×12×8, MLP 1408; bf16) |
| `model_120m_clean_sft.pt` | scale checkpoint (640×15×10, 1664; bf16) |
| `spm_en.model` | 16k SentencePiece BPE (+`[MASK]`=16000) |
| `mdlm_chat.onnx` | fp32 mobile export of 66m |

## Use

```bash
git clone <github-repo>
python scripts/sample_mdlm_v2.py --checkpoint model_66m_clean_sft.pt \
  --prompt "user: I feel anxious about work. assistant:" --resp_len 110 --steps 28 \
  --temperature 0.9 --top_p 0.9   # + rep_penalty 1.3, neighbor-ban 2, no-repeat-3gram (defaults)
# 120m: add MDLM_HIDDEN=640 MDLM_LAYERS=15 MDLM_HEADS=10 MDLM_MLP_INTER=1664
```

Use confidence-ordered unmasking (v2 sampler). Random-order unmasking re-introduces repetition collapse (paper §5).

## Training

MDLM objective (mask prob t∼U(0.01,0.99), CE on masked); DiffuSeq-style SFT (response-only noising). Pretrain 20k steps / SFT 8k, batch 256×512, LR 4e-4 cosine, bf16 + AdamW8bit. Data: strict-filtered psychology corpus 0.774B tokens; SFT 261k pairs (assistant + mental-health + translated psychology dialogue). Corpora not redistributed; build scripts in repo.

## Eval (paper, audited)

725 paired prompts, MiniLM cosine-to-gold: 66m-clean 0.500; 120m 0.490 (d=−0.0094, win 0.479 → null at fixed data). Blind traps (3 annotators, κ .77–.87): 66m 9/28 vs 120m 7/28, p=0.77. distinct-2 0.993; rep-1 0.137; collapse 0.

## Citation

```bibtex
@misc{kim2026psyche, title={Small-Scale Text Diffusion for a Psychology Assistant}, author={Kim, Denis}, year={2026}, note={arXiv preprint}}
```