An RNN that segments streams without fixing the cluster count

What it does

UIS-RNN is a PyTorch implementation of the Unbounded Interleaved-State Recurrent Neural Network, an algorithm that segments and clusters sequential data by learning from labeled examples. Though originally proposed for fully supervised speaker diarization—mapping audio embeddings to speaker labels—the core library is framed as a general sequence-clustering tool. You train a UISRNN model on observation sequences and their ground-truth cluster IDs, then ask it to infer labels on new, unseen sequences.

The interesting bit

The model sidesteps the classic clustering headache of guessing how many speakers (or clusters) exist ahead of time. Its “unbounded interleaved-state” architecture is designed for exactly that open-ended scenario. The README also notes that the training pipeline concatenates and block-wise shuffles sequences as a built-in data augmentation strategy when examples are scarce.

Key highlights

• Learns segmentation and clustering jointly from labeled sequences rather than relying on unsupervised distance metrics.
• Named for its “unbounded interleaved-state” design, it targets scenarios where the number of clusters is not known ahead of time.
• Accepts either a list of variable-length sequences or a single concatenated array for training, with specific flags for handling local versus global cluster IDs.
• Ships with toy data and integration tests, though the authors warn the open-source release differs slightly from the internal version used in the ICASSP paper.
• The authors explicitly state they cannot release the d-vector speaker encoder or proprietary training data used in the original paper, so you must bring your own embeddings.

Caveats

• The open-source implementation is slightly different from Google’s internal version, so paper results may not reproduce exactly.
• No pre-trained speaker recognition model or data is included; the library is essentially the clustering/segmentation engine, not an end-to-end diarization pipeline.
• Training on large datasets requires manually chunking calls to fit() and carefully balancing sequence sources across batches.

Verdict

Worth a look if you are building a supervised diarization or general sequence-clustering pipeline and already have your own feature extractor. Skip it if you need a turnkey speaker-diarization system out of the box.