cohere-asr-ja-lite

LiteASR style compression applied to efwkjn/cohere-asr-ja. The encoder linear layers are compressed using low rank decomposition, reducing size/latency by 44%/~8% with minimal impact on CER.

This model uses a custom inference engine based on PyTorch/torch.compile. And is probably the fastest implementation for offline Cohere, transcribing 273 videos (68.5 hours) in 82.9s (full model 89.5s) on a single RTX 6000 Blackwell with enc8/dec512 batch size. Comparing to hf-audio longform leaderboard results on the official model has WER within margin of error. Leaderboard RTFx is from H200 (a stronger GPU) runs mentioned here. WER change on cv17 en test (n=16393) is statistically insignificant - one-tail trfms < mine p=0.41.

Benchmarks

Latency (s) on JSUT-book (6 files, 4257s, 135 segments) best@5 16GB GPU

See my most recent whisper finetune for more evals

Usage

from huggingface_hub import snapshot_download
model_dir = 'lite'
snapshot_download('efwkjn/cohere-asr-ja-lite', cache_dir=model_dir, local_dir=model_dir)

import importlib
import torch
from tokenizers import Tokenizer

cohere_asr = importlib.import_module(model_dir + '.cohere_asr')
file = 'audio.wav'
device = 'cuda'
fe = cohere_asr.CohereAsrFeatureExtractor(model_dir, device='cpu')
model = cohere_asr.CohereAsr.from_pretrained(model_dir, device)
tokenizer: Tokenizer = Tokenizer.from_file(model_dir + '/tokenizer.json')
prompt = tokenizer.encode(
    '<|startofcontext|><|startoftranscript|><|emo:undefined|>'
    '<|ja|><|ja|><|pnc|><|noitn|><|notimestamp|><|nodiarize|>',
    add_special_tokens=False,
).ids
prompts = torch.tensor(prompt, dtype=torch.int64, device=device)[None, :].expand(1, len(prompt))
for lst in model.transcribe([file], fe, prompts, compile=False):
    print(''.join(tokenizer.decode_batch([x.tolist() for x in lst])))

Performance Notes

• Above benchmarks are with defaults except compile=True and batch size
• Feature extractor can be offloaded if cpu is slow
• If using cpu fe, torch.set_num_threads(1) can be useful, increase workers instead for multithreading
• Workers (default 2 loaders, 4 extractors) may need to be increased based on audio decoding/feature extractor performance relative to gpu
• Audio preprocessing is throttled to maintain a buffer
• Use max-autotune for extra speed
• Compile padding can be tuned: fe.pad_to, model.max_audio_clip, pad='decoder'
• block_size (self, cross) can be tuned to reduce kv cache padding reads (I like 64, 128)
• Pointwise convs were changed to linears because gemm kernels profiled slightly faster
• Encoder batch size split from decoder due to inefficiencies in large batch conv kernels and better pipelining
• Possible optimisations: cache encoder attention p, custom kernels for low rank, cache prefill

Analysis

[!Note] The following is a tech report style analysis of the method. This isn't useful to run the model. Understanding of the original paper is assumed.

Methodology

The original paper performs PCA with SVD but this requires storing activations needing \( 1280(1280*5+5120),\textrm{per layer}*32,\textrm{layers}*1500,\textrm{tokens}*2,\textrm{fp16}=1.4\textrm{GB} \) per sample. We could do it online, running all samples in a single batch and decomposing after each linear but this still requires a non-insignificant amount of memory as samples increases.

Instead I reimplement the compression script to do PCA with eigendecomposition of the covariance matrix with a simple 2 pass method. First pass is quick, determining the mean using \( (1280*5+5120)*32,\textrm{layers}4,\textrm{fp32}=1.5\textrm{MB} \) per minibatch and reduced after. Second pass accumulates, using a constant \( (1280^25+5120^2)*32,\textrm{layers}*4,\textrm{fp32}=4.4GB \) for any amount of calibration samples (and can be offloaded to CPU).

Evaluation uses a custom inference implementation with pure greedy search no timestamps, missing minor steps like suppress_tokens. CER is calculated with whisper.normalizers.BasicTextNormalizer; removing whitespace; and using jiwer.process_characters().cer doing a weighted mean instead of arithmetic mean. All other details as implemented in the original.

Whisper

Testing official checkpoints calibrated on English released by efficient-speech with Japanese (using their hf impl), results in some romaji instead of a proper transcriptions. See link for CER results.

I start with efwkjn/whisper-ja-760M calibrating on 512 samples generated with my training code. Figures 1/2 uses the FLEURS ja_jp dev set filtered for duplicate text (n=135).

<img src="figure_1.png" width="60%" alt="Figure 1">

Figure 1: FLEURS dev CER vs encoder size, coloured by MLP θ.

<img src="figure_2.png" width="60%" alt="Figure 2">

Figure 2: FLEURS dev CER vs θ, coloured by encoder size.

CER mainly jumps from IR exploding due to looping. Inspecting some problematic generations indicates all clips are extremely low volume and likely out of distribution from the calibration data.

<img src="figure_3.png" width="60%" alt="Figure 3">

Figure 3: JSUT basic5000 CER vs θ, coloured by encoder size.

In Figure 3 I run another grid search evaluating the first 135 samples of JSUT basic5000 for a comparison on clean audio. There is still a cliff from looping but at higher compression levels. In the previous 2 searches, there is a shift at attn θ=0.9884, reason is unclear (maybe denoising effect).

<img src="figure_4.png" width="60%" alt="Figure 4">

Figure 4: Calibration samples powers of 2 (8-512) vs shared θ attn/MLP (0.992-1.0), coloured by encoder size.

Next I ablate the number of calibration samples sweeping θ and show results in figure 4. As samples grow, encoder size required to hit a certain θ increases but CER to size with 32+ samples is roughly equal. With the 128 sample data I do a small grid search around the attn θ=0.9884 area and find the same behaviour seen previously.

Conformer

In limitations of the original paper they mention the method is not applied to Conformer convolution layers. But the method can be used on pointwise convs which are mathematically identical to linear layers (with a transposed representation); and depthwise convs have negligible params/flops. In Figure 5 I apply this to efwkjn/cohere-asr-ja.

<img src="figure_5.png" width="60%" alt="Figure 5">

Figure 5: FLEURS dev CER vs encoder size, coloured by convolution θ.

Convolution layers are significantly more sensitive to compression where the Pareto front is dominated by no to low compression. Attention layers can take more compression than MLP, but θ needs to be lower for similar compression/CER effects compared to Whisper. Unlike Whisper, CER increases gradually with occasional failures throughout where the decoder outputs no text.

So far we grouped θ by layer type, next I separate each tensor and search with NSGA-II optimising for latency/KLD. I calculate KLD on 256 calibration samples and microbench each replacement (with default PyTorch) for a latency cost. k<64 is excluded for no speedup (memory bound without custom kernels), and some values just above multiples of 256 for being slower than full rank, e.g. k>512 for 1280×1280 linears.

<img src="figure_6.png" width="60%" alt="Figure 6">

Figure 6: FLEURS train CER vs encoder size.

To improve convergence of the 480D space over random init, I seed the population with minimal compression and slowly limit latency. At 0.025 KLD (got bored) I generate 512 new samples to reduce overfitting and search around that cost for a final model. In Figure 6, I select points from both the GS and BO, and benchmark on FLEURS train set (n=2292). Despite GS points optimised on FLEURS dev; and no FLEURS in calibration, BO substantially outperforms GS. I choose the highlighted point to release as cohere-asr-ja-lite.

<img src="figure_7.png" width="60%" alt="Figure 7">

Figure 7: Compression ratios for cohere-asr-ja-lite.

Compression ratios and derived θ of cohere-asr-ja-lite are shown in Figure 7. Derived θ indicates per type θ is suboptimal with large variation within types especially in early layers which get underestimated. Attention value/out also take significantly less compression than query/key.

To evaluate quality loss, I compare clips from the 9 shortform datasets above (n=131590) with clip-wise bootstrap method. Mean difference is 0.17±0.06 CER (CI assumes clip independence, but effect of dataset bias should be small with similarity of models).

Conclusion

LiteASR is cheap and works reasonably well for a couple minutes of compute given sensible θs. But within the same structure can be substantially improved with a few days of Bayesian optimisation. For a future Whisper finetune I will use structured pruning instead to maximise performance, preliminary healing results show better loss on a smaller model compared to LiteASR.