DOCSEARCHbyAlgolia
Submit Docs

entmax Documentation

repository·master·Indexed 19 days ago

https://github.com/deep-spin/entmax

A PyTorch implementation of entmax and entmax loss families. It provides sparse probability mappings that generalize softmax and cross-entropy, including sparsemax, entmax15, entmax_bisect, normmax_bisect, and budget_bisect, allowing for exact zero probabilities in attention mechanisms and classification tasks.

Tokens
674
Snippets
3
Records
3
Agent score
17%

What's inside entmax

  2. Compute gradients with respect to alpha for adaptive sparsity

    master

    The entmax_bisect function supports automatic differentiation with respect to the alpha parameter. This allows you to learn the sparsity level ($\alpha$) during training.

    To do this, ensure the alpha tensor has requires_grad=True before passing it to entmax_bisect.

    from torch.autograd import grad
from entmax import entmax_bisect
import torch

x = torch.tensor([[-1, 0, 0.5], [1, 2, 3.5]])
alpha = torch.tensor(1.33, requires_grad=True)

p = entmax_bisect(x, alpha)

# Compute gradient of a specific element of the output with respect to alpha
g = grad(p[0, 0], alpha)
# Returns: (tensor(-0.2562),)

  3. Use entmax sparse probability mappings

    master

    The entmax package provides several sparse probability mapping functions that generalize softmax. These functions can produce exact zeros, unlike standard softmax.

    Available functions include:

    • sparsemax: 2-entmax implementation.
    • entmax15: 1.5-entmax implementation.
    • entmax_bisect: Generic $\alpha$-entmax using a bisection-based algorithm.
    • normmax_bisect: $\alpha$-normmax using a bisection-based algorithm.
    • budget_bisect: A transformation that handles $k$-subsets budget via bisection.

    All functions accept a dim argument to specify the dimension along which the mapping is computed.

    import torch
from entmax import sparsemax, entmax15, entmax_bisect, normmax_bisect, budget_bisect

x = torch.tensor([-2, 0, 0.5])

# Standard softmax (non-sparse)
# Out: tensor([0.0486, 0.3592, 0.5922])

# Sparsemax (2-entmax)
# Out: tensor([0.0000, 0.2500, 0.7500])

# 1.5-entmax
# Out: tensor([0.0000, 0.3260, 0.6740])

# Generic alpha-normmax
# Out: tensor([0.0000, 0.3110, 0.6890])

# Budget-based transformation
# Out: tensor([0.0000, 1.0000, 1.0000])