Add sparsemax and SparsemaxSemiring

yzhangcs · yzhangcs · commit cabb0f02979b · 2021-12-25T17:28:32.000+08:00
diff --git a/supar/structs/fn.py b/supar/structs/fn.py
@@ -230,4 +230,28 @@ def backward(ctx, grad_output):
         return None, None
 
 
+class SparsemaxFunction(Function):
+
+    @staticmethod
+    def forward(ctx, x, dim=-1):
+        ctx.dim = dim
+        sorted_x, _ = x.sort(dim, True)
+        z = sorted_x.cumsum(dim) - 1
+        k = x.new_tensor(range(1, sorted_x.size(dim) + 1)).view(-1, *[1] * (x.dim() - 1)).transpose(0, dim)
+        k = (k * sorted_x).gt(z).sum(dim, True)
+        tau = z.gather(dim, k - 1) / k
+        p = torch.clamp(x - tau, 0)
+        ctx.save_for_backward(k, p)
+        return p
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        k, p, dim = *ctx.saved_tensors, ctx.dim
+        grad = grad_output.masked_fill(p.eq(0), 0)
+        grad = torch.where(p.ne(0), grad - grad.sum(dim, True) / k, grad)
+        return grad, None
+
+
 sampled_logsumexp = SampledLogsumexp.apply
+
+sparsemax = SparsemaxFunction.apply
diff --git a/supar/structs/semiring.py b/supar/structs/semiring.py
@@ -4,7 +4,7 @@
 
 import torch
 from supar.utils.common import MIN
-from supar.structs.fn import sampled_logsumexp
+from supar.structs.fn import sampled_logsumexp, sparsemax
 
 
 class Semiring(object):
@@ -260,3 +260,15 @@ class SampledSemiring(LogSemiring):
     @classmethod
     def sum(cls, x, dim=-1):
         return sampled_logsumexp(x, dim)
+
+
+class SparsemaxSemiring(LogSemiring):
+    r"""
+    Sparsemax semiring :math:`<\mathrm{sparsemax}, +, -\infty, 0>`
+    :cite:`martins-etal-2016-sparsemax,mensch-etal-2018-dp,correia-etal-2020-efficient`.
+    """
+
+    @staticmethod
+    def sum(x, dim=-1):
+        p = sparsemax(x, dim)
+        return x.mul(p).sum(dim) - p.norm(p=2, dim=dim)
