JustANormalTinkerer/hayai-ocr-v2-alpha

This is the EXTREMELY ALPHA version of hayai-ocr-v2. This is like after 4 epochs and the language decoder was initialized from scratch. It has repetition issues, it flat out dies in korean and chinese sometimes. It is pretty good at Japanese though.

It is NOT suitable for production yet. I will continue training this and this will eventually replace hayai-ocr once it matures.

Here is the model architecture:

1. Vision Encoder

• Backbone: google/siglip2-base-patch16-naflex (SigLIP2).
• Naflex Batching: It natively supports images of varying aspect ratios and sizes without requiring uniform resizing. The encoder outputs visual features alongside spatial shape and attention mask tensors.
• Feature Dimension: The vision encoder produces features with a dimension of 768.

2. Vision-Language Bridge (Projector)

• Architecture: A Multi-Layer Perceptron (MLPProjector) consisting of two linear layers with a GELU activation function in between.
• Function: Maps the 768-dimensional visual representation into the 512-dimensional hidden space utilized by the decoder.

3. Causal Decoder

The text decoder (VisualCausalOCRDecoder) is a 12-layer custom Transformer decoder designed for autoregressive token-by-token generation:

• Layer Count: 12 decoder layers.
• Hidden Dimension: 512.
• Feed-Forward Network Dimension: 2048.
• Normalization: RMSNorm is used instead of LayerNorm, applied at the input of the attention block, feed-forward block, and before the final output head.
• Grouped-Query Attention (GQA):
  • Configured with 8 Query heads and 2 Key/Value heads. This reduces the memory footprint of the KV cache during generation.
  • Uses a head dimension of 64.
  • Applies RMSNorm directly to the Projected Query and Key states (q_norm and k_norm) before computing attention.
• SwiGLU with Clamping:
  • The feed-forward network uses a SwiGLU activation variant.
  • To prevent instability, the intermediate activations are clamped: the gate tensor is capped at a maximum of 10.0, and the up-projection tensor is clamped between -10.0 and 10.0.
• Positional Embeddings: Rotary Position Embeddings (RoPE) are applied to the Query and Key states in the attention blocks.
• Attention Masking:
  • Uses a block-causal mask.
  • The visual features act as a prefix context where all visual tokens can attend to one another, but cannot attend to future text tokens.
  • The text tokens are causally masked, meaning they can attend to all visual features and prior text tokens, but not future text tokens.

4. Tokenizer & Output Vocab

• Tokenizer: A custom UTF-8 byte-level tokenizer (ByteTokenizer).
• Vocabulary Size: 260.
  • Bytes 0–255 represent standard UTF-8 raw bytes.
  • 4 special tokens: BOS (256), EOS (257), PAD (258), and UNK (259).
• Weight Tying: The input token embedding matrix weight is shared with the final linear output projection head (output_head).