🚀 BlockDiffuse: Fully Parallel Latent Space Reasoning Generation

TL;DR: BlockDiffuse is a non-autoregressive / block-autoregressive generative framework that generates 100 tokens simultaneously in continuous latent space using Rectified Flow Matching and a Diffusion Transformer (DiT) conditioned on intermediate layers of modern LLMs (Qwen/Qwen2.5-0.5B-Instruct).

---

⚡ Key Highlights & Benchmark Results

All benchmarks measured on a single consumer NVIDIA GeForce RTX 4070 Laptop GPU (8GB VRAM):

Generation Mode	Target Size	ODE Steps / Block	Numerical Solver	Latency (ms)	Throughput (tokens/sec)	VRAM Footprint
Single-Block Parallel	100 tokens	8 ODE steps	DPM-Solver + TFE	1,730.60 ms	57.78 tok/s	3,674 MB
Multi-Block Autoregressive	200 tokens	8 ODE steps / block	DPM-Solver + TFE	1,279.20 ms	156.35 tok/s	3,789 MB

---

🏗️ Architecture Overview

Prompt Prefix ──► Frozen Qwen2.5 (Layers 1..12) ──► Continuous Context c [L_p x 896]
                                                          │
Initial Gaussian Noise z_0 [100 x 896] ~ N(0, I) ─────────┤
                                                          ▼
                                             BlockDiffuse DiT (8 Layers, 14 Heads)
                                             - AdaLN-Zero Timestep Conditioning
                                             - Continuous RoPE Positional Encoding
                                             - Rectified Flow (v-prediction)
                                                          │
                                                          ▼
                                            Predicted Latents z_1 [100 x 896]
                                                          │
                                                          ▼
                                             Deep Proj Head (3-Layer SwiGLU MLP)
                                                          │
                                                          ▼
                                            Pre-Head RMSNorm + Frozen LM Head
                                                          │
                                                          ▼
                                            Discrete Next 100 Tokens in Parallel

1. Base LLM Backbone

• Model: Qwen/Qwen2.5-0.5B-Instruct
• Representation Layer: Layer 12 (mid-layer context extraction, $d_{\text{model}} = 896$).
• Head: Frozen LM head with vocab size $151{,}936$.

2. Diffusion Transformer (DiT)

• Depth: 8 Transformer Blocks.
• Attention: 14 heads (head dimension 64, matches $d_{\text{model}} = 896$).
• Initialization: Direct parameter transfer from layers 6–11 of Qwen2.5-0.5B.
• Modulation: AdaLN-Zero modulates scale and shift parameters based on timestep $t \in [0, 1]$.

3. Flow Matching & Multi-Objective Training

Rectified Flow straight-line trajectory: $$z_t=(1-t)z_0+t{z}_1,v_t=\frac{d{z}_t}{dt}=z_1-z_{0}z\_t\; =\; (1\; -\; t)\; z\_0\; +\; t\; z\_1,\; \backslash quad\; v\_t\; =\; \backslash frac\{dz\_t\}\{dt\}\; =\; z\_1\; -\; z\_0$$

Trained under composite multi-loss: $${\mathcal{L}}_{\text{total}}={\lambda }_{\text{FM}}{\mathcal{L}}_{\text{FM}}+{\lambda }_{\text{disp}}{\mathcal{L}}_{\text{disp}}+{\lambda }_{\text{KL}}{\mathcal{L}}_{\text{KL}}+{\lambda }_{\text{CE}}{\mathcal{L}}_{\text{CE}}+{\lambda }_{\text{NN}}{\mathcal{L}}_{\text{NN}}\backslash mathcal\{L\}\_\{\backslash text\{total\}\}\; =\; \backslash lambda\_\{\backslash text\{FM\}\}\; \backslash mathcal\{L\}\_\{\backslash text\{FM\}\}\; +\; \backslash lambda\_\{\backslash text\{disp\}\}\; \backslash mathcal\{L\}\_\{\backslash text\{disp\}\}\; +\; \backslash lambda\_\{\backslash text\{KL\}\}\; \backslash mathcal\{L\}\_\{\backslash text\{KL\}\}\; +\; \backslash lambda\_\{\backslash text\{CE\}\}\; \backslash mathcal\{L\}\_\{\backslash text\{CE\}\}\; +\; \backslash lambda\_\{\backslash text\{NN\}\}\; \backslash mathcal\{L\}\_\{\backslash text\{NN\}\}$$

---

💻 Quickstart: Inference

1. Clone & Setup

git clone https://github.com/Hooshaai/BlockDiffuse.git
cd BlockDiffuse
pip install -r requirements.txt

2. Download Checkpoint from Hugging Face

from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download(
    repo_id="tahamajs/BlockDiffuse",
    filename="blockdiffuse_final.pt"
)
print("Checkpoint downloaded to:", ckpt_path)

3. Run Parallel Multi-Block Generation

python inference.py \
    --model Qwen/Qwen2.5-0.5B-Instruct \
    --checkpoint ./checkpoints_improved/blockdiffuse_final.pt \
    --prompt "<|im_start|>system\nYou are a helpful assistant that solves problems step by step.<|im_end|>\n<|im_start|>user\nA bookstore has 140 books on Monday. On Tuesday, they sell 45 books. On Wednesday, they receive 80 books. How many remain?<|im_end|>\n<|im_start|>assistant\n" \
    --max_blocks 2 \
    --steps 8 \
    --solver dpm_solver \
    --use_tfe \
    --tfe_seeds 3

---

📜 Citation

@article{blockdiffuse2026,
  title={BlockDiffuse: Fully Parallel Latent Space Reasoning Generation with Diffusion Transformers},
  author={Hooshaai Research},
  journal={GitHub / HuggingFace Technical Report},
  year={2026},
  url={https://github.com/Hooshaai/BlockDiffuse}
}