DiarizationLM-Gemma-4-E4B-v1

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 (<prompt> --> <completion> [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 = (
    "<speaker:1> Hello, how are you doing <speaker:2> today? I am doing well."
    " What about <speaker:1> you? I'm doing well, too. Thank you."
)

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, device_map="cuda")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda"
)

print("Tokenizing input...")
inputs = tokenizer([HYPOTHESIS + " --> "], return_tensors="pt").to("cuda")

print("Generating completion...")
outputs = model.generate(
    **inputs,
    max_new_tokens=int(inputs.input_ids.shape[1] * 1.2),
    do_sample=False,
    use_cache=True,
)

print("Decoding completion...")
completion = tokenizer.batch_decode(
    outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True
)[0]
completion = utils.truncate_suffix_and_tailing_text(completion, " [eod]")

print("Transferring completion to hypothesis text...")
transferred_completion = utils.transfer_llm_completion(completion, HYPOTHESIS)

print("========================================")
print("Hypothesis:", HYPOTHESIS)
print("========================================")
print("Completion:", completion)
print("========================================")
print("Transferred completion:", transferred_completion)
print("========================================")

2. GGUF (llama.cpp)

You can also run the quantized DiarizationLM-Gemma-4-E4B-v1-q4_k_m.gguf model directly with llama.cpp:

llama-cli \
  -m DiarizationLM-Gemma-4-E4B-v1-q4_k_m.gguf \
  -p "<speaker:1> Hello, how are you doing <speaker:2> today? I am doing well. What about <speaker:1> you? I'm doing well, too. Thank you. --> " \
  --temp 0.0 \
  -n 128

Citation

@inproceedings{wang24h_interspeech,
  title     = {{DiarizationLM: Speaker Diarization Post-Processing with Large Language Models}},
  author    = {Quan Wang and Yiling Huang and Guanlong Zhao and Evan Clark and Wei Xia and Hank Liao},
  year      = {2024},
  booktitle = {Interspeech 2024},
  pages     = {3754--3758},
  doi       = {10.21437/Interspeech.2024-209},
}