DiarizationLM-Gemma-4

This is not an officially supported Google product.

Overview

DiarizationLM is a Large Language Model framework designed to post-process, correct, and optimize automatic speech recognition (ASR) and speaker diarization outputs.

google/DiarizationLM-Gemma-4-E4B-v1 is built on Google's Gemma 4 E4B (4B dense parameters) foundation model and fine-tuned with Locality-Preserving Oracle Supervision across all 4 canonical speaker diarization benchmark corpora:

1. Fisher English (2-speaker conversational telephone speech)
2. Callhome American English (2–5 speaker informal telephone conversations)
3. ICSI Meeting Corpus (3–9 speaker academic research meetings)
4. AMI Meeting Corpus (4-speaker tabletop meetings)

• Foundation model: google/gemma-4-E4B
• Open-source library & scripts: https://github.com/google/speaker-id/tree/master/DiarizationLM

Unlike earlier models trained exclusively on 2-speaker telephone data (such as google/DiarizationLM-8b-Fisher-v2) or naive multi-domain SFT (which suffers from long-monologue speaker identity drift in 4+ speaker meetings), google/DiarizationLM-Gemma-4-E4B-v1 is trained on locality-preserving oracle targets that teach the model to correct lexical turn boundaries and backchannels (1 ≤ L ≤ 5 words) while preserving acoustic speaker anchors across long monologues (L ≥ 6 words). As a result, it achieves statistically significant (p < 0.0001) WDER and cpWER improvements simultaneously across all four benchmarks under standard out-of-the-box Transcript-Preserving Speaker Transfer (transfer_llm_completion), despite having half the parameter count (4B vs. 8B).

Training Configuration

• Base Architecture: Gemma 4 E4B (42 layers, hidden size 2560, hybrid 5:1 sliding-window and global attention, 262,144 vocabulary size)
• LoRA Adapter: Rank r = 256 applied to all attention (q_proj, k_proj, v_proj, o_proj), MLP (gate_proj, up_proj, down_proj), and per-layer input (per_layer_input_gate, per_layer_projection, per_layer_model_projection) linear projections, merged into 16-bit (bfloat16) base weights and serialized to 4-bit GGUF (Q4_K_M and Q4_0)
• Training Objective: Completion-only cross-entropy loss ( --> [eod])
• Training Data: 51,063 Fisher + 20,762 Multi-Corpus (Callhome, ICSI, AMI) locality-preserving prompt-completion pairs
• Optimization: 10,000 steps, global batch size 8, AdamW (beta1 = 0.9, beta2 = 0.99), peak learning rate 1.5e-4 with 500-step linear warmup and cosine decay on 8 Google Cloud TPU v5p chips
• Prompt Segmentation Length: 4,000 characters (maximal sequence length 2,560 tokens)

Model Files Included

• model.safetensors: Merged 16-bit (bfloat16) Hugging Face transformers weights (~16.0 GB)
• DiarizationLM-Gemma-4-E4B-v1-q4_k_m.gguf: (Recommended GGUF) Serialized 4-bit K-quant Medium (Q4_K_M) GGUF model (~5.30 GB) for llama.cpp / Ollama / llama-cpp-python (uses 256-element super-blocks in Q4_K with sensitive attn_v / ffn_down and embedding matrices retained in 6-bit Q6_K)
• DiarizationLM-Gemma-4-E4B-v1-q4_0.gguf: Serialized legacy 4-bit (Q4_0) GGUF model (~5.15 GB) for llama.cpp / Ollama / llama-cpp-python
• config.json, generation_config.json, tokenizer.json, tokenizer_config.json, processor_config.json, chat_template.jinja, special_tokens_map.json: Tokenizer and model configuration files

Benchmark Performance (with 95% Bootstrap Confidence Intervals)

All metrics below are micro-averaged across the full evaluation sets using the USM + turn-to-diarize baseline and scored via Hungarian-matching dynamic programming (diarizationlm.compute_metrics_on_json_dict). Ranges in brackets indicate 95% non-parametric bootstrap confidence intervals (B = 10,000 conversation-level resamples):

Usage

1. Python (transformers + diarizationlm)

First, install the required packages:

pip install transformers diarizationlm

Run inference on a GPU with bfloat16:

from diarizationlm import utils
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "google/DiarizationLM-Gemma-4-E4B-v1"

HYPOTHESIS = (
    " Hello, how are you doing  today? I am doing well."
    " What about  you? I'm do