yfyeung / PrunedCTC

# Pruned CTC
Official code for
*Pruned CTC for Memory-Efficient Large-Vocabulary ASR Training*.
 Vocabulary head and CTC loss memory (top) and runtime (bottom) for forward and backward passes. 
 Crosses mark extrapolated out-of-memory (OOM) points. 
Pruned CTC uses PyTorch and k2 to compute the CTC
loss directly from encoder states and a linear projection, avoiding a full
`(batch, time, vocabulary)` logit tensor. It combines:
- **Exact vocabulary reduction:** the CTC graph uses only the target tokens and
 blank. Probabilities are still normalized over the **full vocabulary**, and
 gradients include every vocabulary class.
- **Chunked projection and backward:** vocabulary chunks are recomputed during
 backward to reduce the memory used by the projection and loss activations.
- **Alignment pruning:** k2 prunes the alignment lattice using a configurable
 log-score beam. A finite beam makes the alignment sum approximate;
 vocabulary reduction itself remains exact.
The projection weights, their gradients, and optimizer states remain dense.
Memory savings depend on the vocabulary, batch, sequence lengths, chunk size,
and alignment beam.
## Installation
Requires **Python 3.10+**, **PyTorch 2.4+**, and a compatible **k2** build.
First install PyTorch and k2 for your Python version, device, and CUDA version.
Follow the
k2 installation guide;
k2 wheels are tied to specific PyTorch builds. A generic `pip install k2` can
select an older PyTorch dependency and replace an existing installation.
Install Pruned CTC from PyPI:
 [code] 
Or install from this repository:
 [code] 
## Quick start
 [code] 
## API
`pruned_ctc_loss` accepts the arguments below. Optional arguments are keyword-only.
| Argument | Default | Description |
| --- | --- | --- |
| `encoder_out` | Required | Encoder states of shape `(N, T, D)`. |
| `weight` | Required | Linear projection weights of shape `(V, D)`. |
| `bias` | Required | Projection bias of shape `(V,)`, or `None`. |
| `y` | Required | Two-axis `k2.RaggedTensor`: one target sequence per utterance, excluding `blank_id`. |
| `encoder_out_lens` | Required | Int32/int64 frame counts of shape `(N,)`, each in `[1, T]`. |
| `chunk_size` | `4096` | Positive number of vocabulary columns per chunk. Smaller chunks reduce temporary memory but may slow training. |
| `output_beam` | `100.0` | Positive, finite log-score beam. Larger beams retain more alignments and use more memory. |
| `max_states` | `100_000_000` | Positive lattice state limit. k2 may narrow the effective beam to stay within this limit. |
| `blank_id` | `0` | Blank token ID in the original vocabulary. |
`N`, `T`, `D`, and `V` denote batch size, padded sequence length, encoder
dimension, and vocabulary size, respectively.
The return value is a scalar `float32` **sum** over utterances. Nonfinite
per-utterance losses are excluded with a warning, including losses from
impossible target alignments. With finite inputs, impossible alignments
contribute zero loss and zero gradients. App