src/transformers/models/yoso/modeling_yoso.py/YosoCumulation/backward
class YosoCumulation:    def backward(ctx, grad):
        grad = to_contiguous(grad)

        query_mask, key_mask, expectation, query, key, value = ctx.saved_tensors
        config = ctx.config

        hash_code_len = config["hash_code_len"]

        weighted_exp = torch.matmul(grad, value.transpose(-1, -2)) * expectation
        grad_query = torch.matmul(weighted_exp, (hash_code_len / 2) * key)
        grad_key = torch.matmul(weighted_exp.transpose(-1, -2), (hash_code_len / 2) * query)
        grad_value = torch.matmul(expectation.transpose(-1, -2), grad)

        return None, None, grad_query, grad_key, grad_value, None