mmd median heuristic

mackelab · Feb 7, 2024 · 47b93c3 · 47b93c3
1 parent c8a2fbd
commit 47b93c3
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/labproject/metrics/MMD_torch.py b/labproject/metrics/MMD_torch.py
@@ -27,6 +27,12 @@ def compute_rbf_mmd(x, y, bandwidth=1.0):
     mmd = torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel)
     return mmd
 
+@register_metric("mmd_rbf_median_heuristic")
+def compute_rbf_mmd_median_heuristic(x, y):
+    # https://arxiv.org/pdf/1707.07269.pdf
+    median = torch.median(torch.cdist(x, y))
+    return compute_rbf_mmd(x, y, median)
+
 
 @register_metric("mmd_rbf_auto")
 def compute_rbf_mmd_auto(x, y, bandwidth=1.0):