Learning notes [Gumbel softmax]

gumbel softmax

• Used for processing argmax Non differentiable case

• Solutions ： introduce gumbel Distribution . Use... In forward propagation argmax, Used in backward gradient return gumbel_softmax Calculation

• Code

def gumbel_softmax(logits: Tensor, tau: float = 1, hard: bool = False, eps: float = 1e-10, dim: int = -1) -> Tensor:
    ...
    gumbels = (
        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log()
    )  # ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)

    if hard:
        # Straight through.
        index = y_soft.max(dim, keepdim=True)[1]
        y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
        ret = y_hard - y_soft.detach() + y_soft
    else:
        # Reparametrization trick.
        ret = y_soft
    return ret

gumbel_softmax Temperature is introduced in t, t The smaller it is ,softmax The closer you get to One-hot. To train stability , commonly t Will take a larger number , Then gradually shrink .
The content is reproduced from gumbel softmax