Upload README.md with huggingface_hub

README.md

@@ -3,7 +3,6 @@ license: apache-2.0
 tags:
 - diffusion
 - flow-matching
-- rectified-flow
 - text-generation
 - reasoning
 - qwen2.5
@@ -13,47 +12,71 @@ language:
 - en
 ---
 
-# BlockDiffuse: Parallel Multi-Block Reasoning Generation in Latent Space
+# 🚀 BlockDiffuse: Fully Parallel Latent Space Reasoning Generation
 
-**BlockDiffuse** is a non-autoregressive / block-autoregressive diffusion framework that generates entire blocks of 100 tokens simultaneously in continuous latent space using Rectified Flow Matching.
+<div align="center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/text-generation.png" alt="Text Generation" width="400"/>
+</div>
+<br/>
 
-By coupling a Diffusion Transformer (DiT) with a frozen decoder-only Base LLM (`Qwen/Qwen2.5-0.5B-Instruct`), BlockDiffuse bypasses token-by-token sequential decoding, achieving parallel multi-token throughput.
+> **TL;DR:** BlockDiffuse is a novel generative framework that bypasses the autoregressive token-by-token bottleneck of LLMs. By combining a **Diffusion Transformer (DiT)** with **Rectified Flow Matching** and a frozen Qwen2.5-0.5B-Instruct, it generates **100-token blocks in parallel** in continuous latent space, achieving over **150 tokens/sec** on a laptop GPU!
 
-## Model Summary
+---
+
+## 📖 The Problem: The Autoregressive Bottleneck
+
+Traditional Large Language Models (LLMs) generate text autoregressively—one token at a time. While highly effective, this creates a severe latency bottleneck, especially for complex reasoning tasks (like math or coding) that require long "chain-of-thought" sequences before arriving at an answer. Because the model must perform a full forward pass for *every single token*, latency scales linearly with the sequence length.
+
+## 💡 The Solution: BlockDiffuse
+
+What if we could generate an entire paragraph of thoughts at once? 
+
+**BlockDiffuse** shifts text generation from discrete token classification to **continuous continuous latent space trajectory modeling**. 
+We map sequences of 100 discrete tokens into a continuous embedding space using the deep layers of a pre-trained LLM. Then, a DiT learns to denoise pure Gaussian noise into these 100-token latent representations simultaneously.
 
-- **Base LLM**: `Qwen/Qwen2.5-0.5B-Instruct`
-- **DiT Architecture**: 8 Transformer Blocks with 14 Attention Heads ($d_\text{model}=896$, head dim 64)
-- **Latent Space**: Mid-layer representations (layer 12) for prompt conditioning; target 100-token blocks supervised via Flow Matching.
-- **Initialization**: Transfer learning from layers 6–11 of Qwen2.5-0.5B.
-- **Projection Head**: Deep 3-layer MLP residual adapter before RMSNorm and discrete token decoding.
-- **Training Objectives**: Flow Matching Velocity MSE ($\mathcal{L}_\text{FM}$) + Dispersive Loss + Teacher KL Distillation + Discrete Token Cross-Entropy ($\mathcal{L}_\text{CE}$) + Contrastive InfoNCE Nearest-Neighbor Loss ($\mathcal{L}_\text{NN}$).
+### 🏗️ Architecture Deep Dive
+1. **Base LLM:** `Qwen/Qwen2.5-0.5B-Instruct` (Frozen). We use the first 12 layers for processing the prompt and the final language modeling head for decoding latents back to text.
+2. **Diffusion Transformer (DiT):** An 8-layer, 14-head ($d_\text{model}=896$) transformer initialized via transfer learning from Qwen's intermediate layers.
+3. **Deep Projection Head:** A 3-layer MLP residual adapter that perfectly bridges the gap between the diffusion output space and the discrete token space.
+4. **Flow Matching:** Instead of standard DDPM, we use **Rectified Flow Matching** with an ODE solver (DPM-Solver) allowing for ultra-fast generation in just 8 steps.
+5. **Loss Objectives:** A sophisticated multi-objective setup combining Velocity MSE ($\mathcal{L}_\text{FM}$), Teacher KL Distillation, Discrete Token Cross-Entropy ($\mathcal{L}_\text{CE}$), and Contrastive InfoNCE ($\mathcal{L}_\text{NN}$).
 
-## Benchmark Results on RTX 4070 Laptop GPU
+---
+
+## ⚡ Benchmark Results (RTX 4070 Laptop 8GB)
+
+BlockDiffuse achieves remarkable throughput by solving the ODE block-by-block instead of token-by-token.
 
-| Mode | Tokens | Diffusion Steps | Latency | Throughput |
+| Mode | Generation Size | Diffusion Steps | Latency | Throughput |
 | :--- | :--- | :--- | :--- | :--- |
-| **Single-Block Parallel** | 100 tokens | 8 ODE steps (DPM-Solver + TFE) | **1,730.60 ms** | **57.78 tokens/sec** |
-| **Multi-Block Autoregressive** | 200 tokens (2 blocks) | 8 ODE steps per block | **1,279.20 ms** | **156.35 tokens/sec** |
+| **Single-Block Parallel** | 100 tokens | 8 ODE steps | **1.73 seconds** | **57.78 tokens/sec** |
+| **Multi-Block Reasoning** | 200 tokens (2 blocks) | 8 ODE steps / block | **1.27 seconds** | **156.35 tokens/sec** |
+
+*(Using DPM-Solver with Training-Free Ensemble (TFE) k=3 seeds)*
+
+---
 
-## Quickstart & Inference
+## 💻 Quickstart & Inference
 
-To run inference using the official repository [Hooshaai/BlockDiffuse](https://github.com/Hooshaai/BlockDiffuse):
+Want to run this yourself? Clone the official repository and use our inference script!
 
 ```bash
+# 1. Clone the repository
 git clone https://github.com/Hooshaai/BlockDiffuse.git
 cd BlockDiffuse
 
+# 2. Run parallel multi-block reasoning
 python inference.py \
     --model Qwen/Qwen2.5-0.5B-Instruct \
-    --checkpoint ./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" \
+    --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\nJanet has 3 bags of 10 apples. She gives 5 apples to her friend and eats 2. How many apples does she have left?<|im_end|>\n<|im_start|>assistant\n" \
+    --max_blocks 2 \
     --steps 8 \
     --solver dpm_solver \
-    --use_tfe \
-    --tfe_seeds 3
+    --use_tfe
 ```
 
-## Citation
+## 📝 Citation
 
 ```bibtex
 @software{blockdiffuse2026,