Custom Autograd Function: SiLU

Write the backward pass of an activation yourself and plug it into autograd.

SiLU (also called swish) is the activation in many modern networks:

silu(x) = x * sigmoid(x)

Autograd can already differentiate it. Here you do the work by hand, with a torch.autograd.Function. A Function has two static methods:

• forward(ctx, x) computes the output. Use ctx.save_for_backward(...) to keep the tensors that the backward pass needs.
• backward(ctx, grad_output) gets dL/d(output) and must return dL/dx. Read the saved tensors from ctx.saved_tensors.

The derivative, with s = sigmoid(x), is:

d silu / dx = s * (1 + x * (1 - s))

Then write silu(x), which applies the Function with SiLU.apply(x).

Function signature

class SiLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x): ...

    @staticmethod
    def backward(ctx, grad_output): ...


def silu(x):

Example

Input

x = torch.tensor([0.0, 2.0], requires_grad=True)
silu(x).sum().backward()

Output

silu(x)  ->  tensor([0.0000, 1.7616])
x.grad   ->  tensor([0.5000, 1.0908])

Reasoning

At x = 0, s = 0.5, so the output is 0 and the slope is 0.5 * (1 + 0) = 0.5. At x = 2, s = 0.8808, the output is 2 * 0.8808 = 1.7616 and the slope is 0.8808 * (1 + 2 * 0.1192) = 1.0908.

Constraints and notes

• x can have any shape and any float dtype. Keep both.
• backward must multiply by grad_output. The tests use a loss that is not a plain sum, so the incoming gradient is not all ones.
• Use torch.sigmoid. It is stable for large |x|. Do not divide two exponentials, because exp(100) / exp(100) is inf / inf, which is nan.
• The tests run torch.autograd.gradcheck on your Function in float64.