CoolFace
Modelpublic

tahamajs/BlockDiffuse

sourceHugging Faceapache-2.0updated 18d agoView on Hugging Face
0likes
Model Card

🚀 BlockDiffuse: Fully Parallel Latent Space Reasoning Generation

![License: Apache 2.0](https://opensource.org/licenses/Apache-2.0) ![GitHub Repository](https://github.com/Hooshaai/BlockDiffuse) ![HuggingFace Space](https://huggingface.co/spaces/tahamajs/BlockDiffuse-Blog) ![Dataset](https://huggingface.co/datasets/tahamajs/BlockDiffuse-Data)

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 ModeTarget SizeODE Steps / BlockNumerical SolverLatency (ms)Throughput (tokens/sec)VRAM Footprint
Single-Block Parallel100 tokens8 ODE stepsDPM-Solver + TFE1,730.60 ms57.78 tok/s3,674 MB
Multi-Block Autoregressive200 tokens8 ODE steps / blockDPM-Solver + TFE1,279.20 ms156.35 tok/s3,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: $$zt = (1 - t) z0 + t z1, \quad vt = \frac{dzt}{dt} = z1 - 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}}$$

💻 Quickstart: Inference

1. Clone & Setup

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

2. Download Checkpoint from Hugging Face

python
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

bash
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

bibtex
@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}
}