DeepSeek-OCR 2: Arquitectura Visual Causal Flow y codificador basado en LLM

DeepSeek-AI ha liberado esta semana DeepSeek-OCR 2, una evolución significativa de su modelo de OCR enfocado en documentos complejos. La actualización introduce una nueva arquitectura de codificador llamada DeepEncoder V2, la cual sustituye el enfoque tradicional basado en Vision Transformers (ViT) o CLIP por un modelo de lenguaje (LLM) compacto de 500 millones de parámetros adaptado para visión. Este cambio estructural permite implementar un "flujo visual causal", mejorando la capacidad del modelo para interpretar la lógica de lectura en diseños no lineales.

Rendimiento en Benchmarks

DeepSeek-OCR 2 reporta mejoras sustanciales en OmniDocBench v1.5, un conjunto de evaluación compuesto por 1,355 documentos diversos (revistas, papers académicos, informes).

Implementación del Visual Causal Flow

El documento técnico "DeepSeek-OCR 2: Visual Causal Flow" detalla la motivación detrás del cambio arquitectónico, citando las limitaciones del escaneo rasterizado (de arriba a abajo, de izquierda a derecha) en documentos con maquetación compleja.

"Presentamos DeepSeek-OCR 2 para investigar la viabilidad de un nuevo codificador —DeepEncoder V2— capaz de reordenar dinámicamente los tokens visuales basándose en la semántica de la imagen. Los modelos de lenguaje-visión (VLMs) convencionales procesan invariablemente los tokens visuales en un orden rígido de escaneo rasterizado [...] con codificación posicional fija cuando se alimentan a los LLMs. Sin embargo, esto contradice la percepción visual humana, que sigue patrones de escaneo flexibles pero semánticamente coherentes impulsados por estructuras lógicas inherentes."

Sobre la implementación técnica específica del encoder y su mecanismo de atención:

"Sustituimos el componente CLIP en DeepEncoder por una arquitectura estilo LLM [...]. Al personalizar la máscara de atención, los tokens visuales utilizan atención bidireccional mientras que las queries aprendibles adoptan atención causal. Cada token de query puede, por tanto, atender a todos los tokens visuales y a las queries precedentes, permitiendo un reordenamiento causal progresivo sobre la información visual."

El diagrama siguiente contrasta el escaneo rasterizado tradicional con el Visual Causal Flow de DeepEncoder V2:

La evolución de DeepSeek-OCR se enmarca en la necesidad de generar datos de entrenamiento de alta calidad para el pre-entrenamiento de LLMs. Mientras que la versión anterior dependía de un encoder visual estándar (CLIP) que aplanaba los parches de imagen en una secuencia 1D estática, la versión 2 intenta mitigar el sesgo inductivo que esto introduce.

Al utilizar un modelo de lenguaje pequeño como codificador visual, DeepSeek valida una arquitectura de "doble causalidad": un primer paso de razonamiento causal en el encoder para ordenar la visión, seguido del razonamiento causal estándar del decoder para generar texto. Esto es relevante para tareas donde la estructura espacial no se correlaciona linealmente con el flujo de lectura, como en tablas, fórmulas matemáticas incrustadas o diseños de revista con múltiples columnas.

Arquitectura e Implementación

El pipeline de DeepSeek-OCR 2 se estructura en tres componentes principales, optimizados para mantener la eficiencia de compresión de su predecesor.

1. Vision Tokenizer

El sistema utiliza un compresor de imágenes híbrido de 80M de parámetros basado en SAM-base y capas convolucionales. Este módulo reduce la dimensionalidad espacial, logrando una compresión de tokens de 16x. Esto es crítico para reducir la carga en el KV Cache durante la inferencia.

• Input: Imágenes de $1024 \times 1024$ (vista global) y $768 \times 768$ (crops locales).
• Output: Embeddings de dimensión reducida (896 dim).

2. DeepEncoder V2 (LLM as Vision Encoder)

Se instancia utilizando la arquitectura de Qwen2-0.5B. Se añaden learnable queries al final de la secuencia visual.

• Estrategia de Tokens: Se mantiene una cardinalidad igual entre los tokens visuales y las consultas causales.
• Multi-crop: El número total de tokens visuales alimentados al LLM oscila entre 256 y 1120, dependiendo del número de crops locales (0 a 6).

3. Decoder (DeepSeek-MoE)

El decoder no sufre cambios respecto a la versión anterior. Se utiliza una arquitectura Mixture of Experts (MoE) con 3B de parámetros totales y aproximadamente 570M activos. Para profundizar en el funcionamiento de esta arquitectura, puede consultarse el análisis sobre Mixture of Experts: Sparse Activation y Estrategias de Routing en LLMs.

El forward pass se describe como:

$$O = \mathcal{D}(\pi_{Q}(\mathcal{T}^{L}(\mathcal{E}(I) \oplus Q_{0}; M)))$$

Donde $\pi_{Q}$ es el operador de proyección que extrae únicamente los últimos $n$ tokens (las consultas causales) para ser procesados por el decoder de lenguaje $\mathcal{D}$.

Código de Inferencia

El siguiente fragmento ilustra la interfaz de uso utilizando la librería transformers, destacando la simplicidad de la integración a pesar de la complejidad interna del encoder.

from transformers import AutoModel, AutoTokenizer
import torch

model_name = 'deepseek-ai/DeepSeek-OCR-2'

# Carga del modelo con soporte para Flash Attention 2
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
    model_name, 
    _attn_implementation='flash_attention_2', 
    trust_remote_code=True, 
    use_safetensors=True
)
model = model.eval().cuda().to(torch.bfloat16)

# Prompt con tag especial de grounding para estructura
prompt = "<image>\n<|grounding|>Convert the document to markdown. "
image_file = 'document_layout.jpg'

# Inferencia con estrategia de crop dinámico
res = model.infer(
    tokenizer, 
    prompt=prompt, 
    image_file=image_file, 
    base_size=1024, 
    image_size=768, 
    crop_mode=True  # Habilita multi-view
)

Detalles de Integración

• Inferencia: El modelo es compatible con transformers y soporta aceleración mediante Flash Attention 2.
• vLLM: Se menciona soporte y guía para aceleración de inferencia y procesamiento de PDFs mediante vLLM en el repositorio oficial.

Fuentes

• Repositorio en GitHub (DeepSeek-OCR-2)
• Paper: DeepSeek-OCR 2: Visual Causal Flow
• Modelo en Hugging Face