Granite-Docling-258M-BitNet-1.58b

This repository hosts a BitNet b1.58 ternary (-1, 0, +1) quantized edition of IBM's ibm-granite/granite-docling-258M (Idefics3 architecture).

Trained via End-to-End Multimodal Knowledge Distillation with Straight-Through Estimators (STE) on the Edge-Intelligence framework. This model quantizes matrix multiplication operations in the language decoder into multiplication-free ternary addition and subtraction. This reduces the parameter memory footprint by ~3.2× - 4.4× and achieves a 89.12% logit cosine similarity to the original teacher model over the full 100,352 token vocabulary.

📊 Benchmark Highlights

🏗️ Model Architecture

• Base Model: ibm-granite/granite-docling-258M
• Vision Backbone: SigLIP2 Vision Encoder (frozen BF16 ViT backbone, preserving high-resolution document OCR and layout token embeddings)
• Multimodal Connector: Pixel Shuffle Projector (linear dimension connector, 4× spatial reduction)
• Language Decoder: Granite 165M Decoder (GQA 8:2, 12 layers, hidden size 768) converted completely to BitLinear b1.58 layers
• Ternary Packing: 5 ternary weights $\in {-1, 0, +1}$ packed per byte ($3^5 = 243 \le 255$)
• Total Parameters: ~258M (Decoder: ~165M, Vision + Connector: ~93M)

🔬 Training & Distillation Methodology

1. Straight-Through Estimator (STE) Quantization

Weights are ternary-quantized dynamically with per-tensor mean absolute scale: $$\mathbf{W}_{\text{ternary}} = \text{Round}\left(\text{Clip}\left(\frac{\mathbf{W}}{\gamma}, -1, 1\right)\right), \quad \gamma = \frac{1}{nm} |\mathbf{W}|1$$ Activations are scaled to 8-bit integers per token using absolute maximum scaling: $$\mathbf{X}{\text{int8}} = \text{Clip}\left(\text{Round}\left(\mathbf{X} \times \frac{127}{\max(|\mathbf{X}|)}\right), -128, 127\right)$$

2. Multi-Objective Distillation Loss

• Supervised Cross-Entropy ($\mathcal{L}_{\text{CE}}$): Next-token prediction on valid assistant DocTags (converged to $\sim 0.39 - 0.49$).
• KL-Divergence Distillation ($\mathcal{L}_{\text{KD}}$): Full-vocabulary logit matching at temperature $T = 2.0$ (converged to $\sim 0.84 - 0.96$).
• Feature Representation Alignment ($\mathcal{L}_{\text{feat}}$): Cosine similarity minimization over decoder hidden states.

📚 Training Datasets

Trained directly on official IBM document extraction and layout datasets:

1. HuggingFaceM4/DoclingMatix (50%): Real multi-page document images, PDFs, invoices, forms, structured document QA, and DocTags layout annotations.
2. ds4sd/SynthCodeNet (20%): Synthesized code blocks with pixel bounding-box DocTags (<code><loc_...>).
3. ds4sd/SynthFormulaNet (15%): Dense LaTeX mathematical formulas with layout coordinates (<formula><loc_...>).
4. ds4sd/SynthChartNet (15%): Charts, infographics, and table structure extractions with cell coordinates.

📈 Training History & Token Statistics

The model underwent four cumulative training stages with an effective batch size of 64 document images per step.

• Average multimodal sequence length: ~953 tokens/sample (vision patches + prompt + DocTags).
• Directly optimized DocTags output: ~93 tokens/sample.
• Total Training Tokens Processed: ~506.3 Million multimodal tokens across all cumulative phases.
• Convergence State: The model reached diminishing returns with sustained logit cosine similarity of 0.8912, saturating the parameter capacity of 1.58-bit ternary representations.

🖥️ Hardware & Training Infrastructure

The training pipeline was executed on the following infrastructure:

• Hardware: NVIDIA GB10 (12 GB VRAM)
• Active VRAM Usage: ~3.0 GB active / ~11.0 GB reserved (via selective token masking & gradient checkpointing)
• Optimizer: 8-bit PagedAdamW (bitsandbytes) with weight decay $1 \times 10^{-4}$
• Learning Rate Schedule: Warm restart cosine annealing ($\text{LR}{\max} = 1.5 \times 10^{-4}$, $\text{LR}{\min} = 1 \times 10^{-5}$)
• Throughput: ~0.03 steps/s (~33 seconds per accumulated batch of 64 pages)

🛠️ Usage via Hugging Face (transformers)

import torch
from transformers import AutoProcessor, Idefics3ForConditionalGeneration
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "tevfikk/granite-docling-258M-bitnet"

processor = AutoProcessor.from_pretrained(model_id)
model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map=device
)

image = Image.open("sample_document.png").convert("RGB")
prompt = (
    "<|start_of_role|>user<|end_of_role|><image>Convert this page to docling.<|end_of_text|>\n"
    "<|start_of_role|>assistant<|end_of_role|>"
)

inputs = processor(text=prompt, images=[image], return_tensors="pt").to(device)
with torch.no_grad():
    generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False)

output = processor.batch_decode(
    generated_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=False
)[0]
print(output)
# Output: <doctag><section_header_level_1>INVOICE #9812</section_header_level_1>...

⚡ Standalone C99 Edge Binary (eif-runtime)

For microcontrollers, Raspberry Pi, and edge accelerators without Python or CUDA:

# 1. Clone & build the zero-dependency C99 runtime
git clone https://github.com/tevfik/eif-runtime.git
cd eif-runtime && mkdir build && cd build && cmake .. && make -j

# 2. Run at 156+ tokens/sec using AVX2/NEON multiplication-free kernel
./bin/eif_infer --model granite_docling_bitnet_dense.eifm --image sample_document.png

📜 Citation & Credits

• Base Model: IBM Granite Docling 258M
• BitNet Architecture: The Era of 1-bit LLMs (Microsoft Research)
• Distillation Framework & Runtime: Edge-Intelligence Project & eif-runtime