Merge pull request #551 from beshkenadze/fix/lid-ecapa-parity-upstream

fix(lid): align ECAPA inference with SpeechBrain

Author: lucasnewman

mlx_audio/lid/models/ecapa_tdnn/ecapa_tdnn.py

@@ -36,7 +36,7 @@ def __init__(self, in_dim: int, out_dim: int):
         self.norm = nn.BatchNorm(out_dim)
 
     def __call__(self, x: mx.array) -> mx.array:
-        return nn.relu(self.norm(self.linear(x)))
+        return self.norm(nn.leaky_relu(self.linear(x), negative_slope=0.01))
 
 
 class DNN(nn.Module):
@@ -66,6 +66,7 @@ def __init__(self, config: ModelConfig):
 
     def __call__(self, x: mx.array) -> mx.array:
         out = mx.squeeze(x, axis=1)
+        out = nn.leaky_relu(out, negative_slope=0.01)
         out = self.norm(out)
         out = self.DNN(out)
         out = self.out(out)
@@ -121,10 +122,16 @@ def __call__(self, mel_features: mx.array) -> mx.array:
         Returns:
             Log-probabilities ``[batch, num_classes]``.
         """
-        embeddings = self.embedding_model(mel_features)
+        normalized_mel_features = self.sentence_mean_normalize(mel_features)
+        embeddings = self.embedding_model(normalized_mel_features)
         embeddings = mx.expand_dims(embeddings, axis=1)
         return self.classifier(embeddings)
 
+    @staticmethod
+    def sentence_mean_normalize(mel_features: mx.array) -> mx.array:
+        """Mirror SpeechBrain's sentence-level mean-only InputNormalization."""
+        return mel_features - mx.mean(mel_features, axis=1, keepdims=True)
+
     def predict(
         self,
         audio: mx.array,

mlx_audio/lid/tests/test_lid.py

@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 import mlx.core as mx
+import mlx.nn as nn
 import numpy as np
 
 
@@ -329,6 +330,35 @@ def test_forward_log_probs_sum(self):
         total = float(mx.sum(probs[0]).item())
         self.assertAlmostEqual(total, 1.0, places=3)
 
+    def test_sentence_mean_normalize_centers_each_mel_bin(self):
+        mel = mx.array([[[1.0, 3.0], [3.0, 5.0], [5.0, 7.0]]])
+        normalized = self.Model.sentence_mean_normalize(mel)
+        mean_per_bin = mx.mean(normalized, axis=1)
+        mx.eval(mean_per_bin)
+
+        self.assertAlmostEqual(float(mean_per_bin[0, 0].item()), 0.0, places=5)
+        self.assertAlmostEqual(float(mean_per_bin[0, 1].item()), 0.0, places=5)
+
+    def test_classifier_matches_speechbrain_order(self):
+        model = self.Model(self.config)
+        model.eval()
+        classifier = model.classifier
+        x = mx.random.normal((1, 1, self.config.embedding_dim))
+
+        expected = mx.squeeze(x, axis=1)
+        expected = nn.leaky_relu(expected, negative_slope=0.01)
+        expected = classifier.norm(expected)
+        expected = classifier.DNN.block_0.linear(expected)
+        expected = nn.leaky_relu(expected, negative_slope=0.01)
+        expected = classifier.DNN.block_0.norm(expected)
+        expected = classifier.out(expected)
+        expected = mx.log(mx.softmax(expected, axis=-1) + 1e-10)
+
+        actual = classifier(x)
+        mx.eval(expected, actual)
+
+        self.assertTrue(mx.allclose(actual, expected, atol=1e-5, rtol=1e-5).item())
+
     def test_predict_returns_sorted(self):
         model = self.Model(self.config)
         labels = {str(i): f"lang_{i}" for i in range(10)}