Hugging Face's logo

mkd-dev
/
keural-vision-encoder-poc

Image Feature Extraction
Transformers
Safetensors
English
keural_vision
vision
vision-encoder
image-text
contrastive-learning
custom-architecture
vlm
Model card Files
xet
 Community
keural-vision-encoder-poc / README.md
mkd-hika's picture
mkd-hika
Duplicate from mkd-hika/keural-vision-encoder-poc
96faa83
preview code
|
Raw History Blame Contribute Delete
7.07 kB
---
language: en
license: mit
library_name: transformers
tags:
- vision
- vision-encoder
- image-feature-extraction
- image-text
- contrastive-learning
- custom-architecture
- vlm
pipeline_tag: image-feature-extraction
---
<div align="center">
# 👁️ Keural Vision Encoder <sub>(PoC · V0.1)</sub>
**A 24.7M-parameter vision encoder trained _from scratch_ — no pretrained backbone, no CLIP weights. Structured tokens that plug into any LLM.**
<p>
<img src="https://img.shields.io/badge/MKD_Co.,_Ltd.-Keural-6d28d9?style=for-the-badge" alt="MKD">
<img src="https://img.shields.io/badge/status-V0.1_PoC-16a34a?style=for-the-badge" alt="Status">
</p>
<p>
<img src="https://img.shields.io/badge/params-24.7M-8b5cf6" alt="Params">
<img src="https://img.shields.io/badge/from-scratch-ec4899" alt="From scratch">
<img src="https://img.shields.io/badge/embed_dim-384-f59e0b" alt="Dim">
<img src="https://img.shields.io/badge/token_budget-adaptive-22d3ee" alt="ATB">
<img src="https://img.shields.io/badge/PyTorch-2.x-ee4c2c?logo=pytorch&logoColor=white" alt="PyTorch">
<img src="https://img.shields.io/badge/trained_on-1×_RTX_5090-38bdf8" alt="Hardware">
</p>
<br>
<img src="assets/keural_vision_encoder.gif" alt="Where the Keural encoder sits in the VLM pipeline" width="98%">
<br>
*The encoder (CNN Stem → ATB → Spatial Transformer) feeds the LevelAware Projector → Mistral-7B in the full [Keural VLM](https://huggingface.co/mkd-hika/keural-vlm-poc).*
</div>
---
## ✨ Key Innovations
**🎯 Adaptive Token Budget (ATB) Tokenization** — token count is a runtime parameter. Tokens are allocated to information-dense regions: a blank wall gets fewer, a dense document gets more.
```python
out = encoder(image, token_budget=64) # cheap
out = encoder(image, token_budget=256) # default
out = encoder(image, token_budget=1024) # full fidelity
```
**🪜 Hierarchical Concept Tokenization (HCT)** — every token carries a semantic level tag: global (whole-image), region (object-scale), or detail (fine-grained).
```python
out = encoder(image)
print(out.level_ids) # tensor of {0=global, 1=region, 2=detail}
```
---
## ⚙️ Model Specifications
| Property | Value |
|---|---|
| Parameters | **24.7M** (trained from scratch) |
| Architecture | CNN Stem + ATB Tokenizer + 12-layer Spatial Transformer |
| Embedding dim | 384 |
| Token budget | 256 (default) · runtime-adjustable |
| Token split | 5% global / 25% region / 70% detail |
| Precision | bfloat16 |
| Training data | CC3M + CC12M (~6.9M image-text pairs) |
| Hardware | 1× RTX 5090 (32 GB VRAM) |
---
## 🎨 Saliency & Token Placement
The ATB tokenizer concentrates tokens on salient regions. Left → right: original · saliency heatmap · token placement (global / region / detail).
<div align="center">
<img src="saliency_viz/step_36k/cat_saliency.png" alt="Saliency and token placement" width="88%">
</div>
---
## 🧪 Training
**Phase 1 — Vision Encoder Pretraining ✅ COMPLETE.** Trained from scratch for
~75,000 steps on CC3M + CC12M (~6.9M image–text pairs), SigLIP-style contrastive
objective, 1× RTX 5090. The frozen encoder is then integrated into the full VLM via
`LevelAwareProjector` (384 → 2048 → 4096) → **Mistral-7B-Instruct-v0.3** (4-bit NF4
QLoRA), fine-tuned with SFT (LLaVA-Instruct-150K, 30K steps).
<div align="center">
<img src="assets/training_loss_curves.png" alt="Training loss curves" width="88%">
</div>
---
## 📊 Benchmark Results — SFT-30K
Downstream VLM benchmarks using this encoder (frozen) + projector + Mistral-7B at the
**SFT-30K checkpoint** (supervised fine-tuning, 30,000 steps; **before** DPO).
Evaluated on 1,000 samples each where applicable. The Keural encoder is **12.4×
smaller** than LLaVA's CLIP encoder (307M).
| Benchmark | Keural SFT-30K | LLaVA 1.5 <sub>(307M enc)</sub> | LLaVA 1.6 <sub>(307M enc)</sub> |
|---|:---:|:---:|:---:|
| VQAv2 Accuracy | **12.9%** | 78.5% | 81.8% |
| POPE F1 | **66.9%** | 85.9% | 86.5% |
| MME Total Score | **704.3** | 1510.7 | 1519.3 |
| TextVQA Accuracy | **0.8%** | 58.2% | 64.9% |
| ScienceQA Accuracy <sub>(img)</sub> | **39.7%** | 66.8% | 70.6% |
> These are the **SFT-30K** numbers (measured from `outputs/eval/sft_30k/`). Applying
> DPO alignment (RLHF-V, 3K steps) improves every metric — e.g. VQAv2 → 43.6%,
> ScienceQA → 53.7%, MME → 838.8. Full **SFT+DPO** results, the comparison chart, and
> the complete VLM live on **[mkd-hika/keural-vlm-poc](https://huggingface.co/mkd-hika/keural-vlm-poc)**.
---
## ⚡ Adaptive Token Budget — latency vs budget
Encode latency vs token budget — trade visual detail for speed at runtime, no retraining.
<div align="center">
<img src="assets/token_budget_curve.png" alt="ATB token budget vs latency" width="80%">
</div>
---
## 📐 Output Contract
Every forward pass returns a `KeuralEncoderOutput`:
| Field | Shape | Description |
|---|---|---|
| `tokens` | `(B, N, 384)` | Token vectors |
| `attention_mask` | `(B, N)` | 1 = real token, 0 = padding |
| `level_ids` | `(B, N)` | 0=global, 1=region, 2=detail |
| `spatial_metadata` | `(B, N, 4)` | Per-token (cx, cy, w, h) in [0,1] |
| `saliency_scores` | `(B, N)` | Allocator confidence per token |
| `pooled` | `(B, 384)` | Single-vector image summary |
| `token_budget_used` | `(B,)` | Actual tokens used per image |
---
## 🚀 Usage
```python
import torch
from architecture.encoder import KeuralVisionEncoder
from keural_config import KeuralConfig
cfg = KeuralConfig.from_yaml("configs/keural_tiny_poc.yaml")
model = KeuralVisionEncoder.from_pretrained("mkd-hika/keural-vision-encoder-poc")
model.eval()
from PIL import Image
from torchvision import transforms
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
image = Image.open("image.jpg").convert("RGB")
x = transform(image).unsqueeze(0) # (1, 3, 256, 256)
with torch.no_grad():
out = model(x)
print(out.tokens.shape) # (1, 256, 384)
print(out.level_ids[0, :8]) # [0, 0, 1, 1, 1, 2, 2, 2]
print(out.pooled.shape) # (1, 384)
```
---
## 🗺️ Roadmap
| Phase | Params | Hardware | Status |
|---|:---:|:---:|:---:|
| **PoC** (this model) | 24.7M | 1× RTX 5090 | ✅ Phase 1 complete |
| Mid-level | ~230M | 8× H100 80 GB | 🔜 Planned |
| Commercial | ~1.1B | 64× H100 80 GB | 🔮 Future |
The mid-level model will use knowledge distillation from **SigLIP-400M** (natural images) and **InternViT-300M** (documents/OCR). Teachers are discarded after training — not part of the final model.
---
## 📝 Citation
```bibtex
@misc{keural_vision_encoder_2026,
title = {Keural Vision Encoder: Content-Adaptive Vision Encoding via Saliency-Guided Token Budgets},
author = {Barki, Hika and MKD Co., Ltd.},
year = {2026},
note = {V0.1 — Phase 1 complete},
}
```
---
## 📄 License
Code: MIT. Model weights trained on CC3M + CC12M — data licenses apply.
<div align="center">
**MKD Co., Ltd. — 2026**
</div>