Code/binary_kernels.py at main · AidanLiotard/Code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import Tensor

from ebm_binary import BinaryEnergyFunction, BinaryTransitionKernel, MarkovChainState


def _energy_and_grad(
    energy_function: BinaryEnergyFunction,
    visible: Tensor,
) -> tuple[Tensor, Tensor]:
    """Return ``E(x)`` and ``dE / dx`` for a batch of binary visibles."""

    with torch.enable_grad():
        visible_for_grad = visible.detach().clone().requires_grad_(True)
        energy = energy_function.energy(visible_for_grad)
        gradient = torch.autograd.grad(energy.sum(), visible_for_grad)[0]
    return energy.detach(), gradient.detach()


@torch.no_grad()
def _dmala_logits(
    visible: Tensor,
    energy_gradient: Tensor,
    *,
    beta: float,
    alpha: float,
) -> Tensor:
    """Return Bernoulli logits for the factorized DMALA proposal on ``{0, 1}^d``."""

    return -0.5 * beta * energy_gradient + 0.5 * (2.0 * visible - 1.0) / alpha


@torch.no_grad()
def _bernoulli_log_prob(sample: Tensor, logits: Tensor) -> Tensor:
    """Return ``log q(sample)`` for a factorized Bernoulli distribution."""

    return (
        sample * F.logsigmoid(logits)
        + (1.0 - sample) * F.logsigmoid(-logits)
    ).sum(dim=1)


class DMALAKernel(BinaryTransitionKernel):
    """Discrete MALA for binary EBMs using autograd-based visible gradients.
    """

    def __init__(
        self,
        alpha: float = 1.0,
        beta: float = 1.0,
    ) -> None:
        if alpha <= 0.0:
            raise ValueError("`alpha` must be positive.")
        if beta <= 0.0:
            raise ValueError("`beta` must be positive.")
        self.alpha = alpha
        self.beta = beta

    @torch.no_grad()
    def step(
        self,
        energy_function: BinaryEnergyFunction,
        state: MarkovChainState,
    ) -> tuple[MarkovChainState, float | None]:
        """Run one DMALA transition step."""

        visible = state.visible

        current_energy, current_gradient = _energy_and_grad(
            energy_function=energy_function,
            visible=visible,
        )
        forward_logits = _dmala_logits(
            visible=visible,
            energy_gradient=current_gradient,
            beta=self.beta,
            alpha=self.alpha,
        )
        proposal_prob = torch.sigmoid(forward_logits)
        proposal = torch.bernoulli(proposal_prob)

        proposal_energy, proposal_gradient = _energy_and_grad(
            energy_function=energy_function,
            visible=proposal,
        )
        reverse_logits = _dmala_logits(
            visible=proposal,
            energy_gradient=proposal_gradient,
            beta=self.beta,
            alpha=self.alpha,
        )

        forward_log_prob = _bernoulli_log_prob(proposal, forward_logits)
        reverse_log_prob = _bernoulli_log_prob(visible, reverse_logits)
        log_accept_prob = self.beta * (current_energy - proposal_energy) + reverse_log_prob - forward_log_prob

        accept = torch.log(torch.rand(visible.shape[0], device=visible.device)) < torch.clamp(
            log_accept_prob,
            max=0.0,
        )
        new_visible = torch.where(accept.unsqueeze(1), proposal, visible)
        return MarkovChainState(visible=new_visible), float(accept.float().mean().item())