Skip to content

torchebm.samplers.langevin_dynamics

Langevin Dynamics Sampler Module.

LangevinDynamics

Bases: BaseSampler

Langevin Dynamics sampler.

Update rule:

\[ x_{t+1} = x_t - \eta \nabla_x U(x_t) + \sqrt{2\eta} \epsilon_t \]

Parameters:

Name Type Description Default
model BaseModel

Energy-based model to sample from.

 required
step_size Union[float, BaseScheduler]

Step size for gradient descent.

 0.001
noise_scale Union[float, BaseScheduler]

Scale of Gaussian noise injection.

 1.0
decay float

Damping coefficient (not supported).

 0.0
dtype dtype

Data type for computations.

 float32
device Optional[Union[str, device]]

Device for computations.

 None
Example 
1
2
3
4
5
6
7
from torchebm.samplers import LangevinDynamics
from torchebm.core import DoubleWellEnergy
import torch

energy = DoubleWellEnergy()
sampler = LangevinDynamics(energy, step_size=0.01, noise_scale=1.0)
samples = sampler.sample(n_samples=100, dim=2, n_steps=500)
Source code in torchebm/samplers/langevin_dynamics.py 
 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class LangevinDynamics(BaseSampler):
    r"""
    Langevin Dynamics sampler.

    Update rule:

    \[
    x_{t+1} = x_t - \eta \nabla_x U(x_t) + \sqrt{2\eta} \epsilon_t
    \]

    Args:
        model: Energy-based model to sample from.
        step_size: Step size for gradient descent.
        noise_scale: Scale of Gaussian noise injection.
        decay: Damping coefficient (not supported).
        dtype: Data type for computations.
        device: Device for computations.

    Example:
        ```python
        from torchebm.samplers import LangevinDynamics
        from torchebm.core import DoubleWellEnergy
        import torch

        energy = DoubleWellEnergy()
        sampler = LangevinDynamics(energy, step_size=0.01, noise_scale=1.0)
        samples = sampler.sample(n_samples=100, dim=2, n_steps=500)
        ```
    """

    def __init__(
        self,
        model: BaseModel,
        step_size: Union[float, BaseScheduler] = 1e-3,
        noise_scale: Union[float, BaseScheduler] = 1.0,
        decay: float = 0.0,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = None,
        *args,
        **kwargs,
    ):
        super().__init__(model=model, dtype=dtype, device=device)

        if isinstance(step_size, BaseScheduler):
            self.register_scheduler("step_size", step_size)
        else:
            if step_size <= 0:
                raise ValueError("step_size must be positive")
            self.register_scheduler("step_size", ConstantScheduler(step_size))

        if isinstance(noise_scale, BaseScheduler):
            self.register_scheduler("noise_scale", noise_scale)
        else:
            if noise_scale <= 0:
                raise ValueError("noise_scale must be positive")
            self.register_scheduler("noise_scale", ConstantScheduler(noise_scale))

        self.decay = decay
        self.integrator = EulerMaruyamaIntegrator(device=self.device, dtype=self.dtype)

    @torch.no_grad()
    def sample(
        self,
        x: Optional[torch.Tensor] = None,
        dim: int = 10,
        n_steps: int = 100,
        n_samples: int = 1,
        thin: int = 1,
        return_trajectory: bool = False,
        return_diagnostics: bool = False,
        *args,
        **kwargs,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
        """
        Generates samples using Langevin dynamics.

        Args:
            x (Optional[torch.Tensor]): The initial state to start sampling from. If `None`,
                a random state is created.
            dim (int): The dimension of the state space (if `x` is not provided).
            n_steps (int): The number of MCMC steps to perform.
            n_samples (int): The number of parallel chains/samples to generate.
            thin (int): The thinning factor (not currently supported).
            return_trajectory (bool): Whether to return the full sample trajectory.
            return_diagnostics (bool): Whether to return sampling diagnostics.

        Returns:
            Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
                - The final samples.
                - If `return_trajectory` is `True`, the full trajectory.
                - If `return_diagnostics` is `True`, a tuple of samples and diagnostics.
        """

        self.reset_schedulers()

        if x is None:
            x = torch.randn(n_samples, dim, dtype=self.dtype, device=self.device)
        else:
            x = x.to(device=self.device, dtype=self.dtype)
            dim = x.shape[-1]
            n_samples = x.shape[0]

        if return_trajectory:
            trajectory = torch.empty(
                (n_samples, n_steps, dim), dtype=self.dtype, device=self.device
            )

        if return_diagnostics:
            diagnostics = self._setup_diagnostics(dim, n_steps, n_samples=n_samples)

        drift = lambda x_, t_: -self.model.gradient(x_)
        with self.autocast_context():
            for i in range(n_steps):
                self.step_schedulers()
                state = {"x": x}
                x = self.integrator.step(
                    state=state,
                    step_size=self.get_scheduled_value("step_size"),
                    noise_scale=self.get_scheduled_value("noise_scale"),
                    drift=drift,
                )["x"]

                if return_trajectory:
                    trajectory[:, i, :] = x

                if return_diagnostics:
                    if n_samples > 1:
                        mean_x = x.mean(dim=0, keepdim=True)
                        var_x = torch.clamp(
                            x.var(dim=0, unbiased=False, keepdim=True),
                            min=1e-10,
                            max=1e10,
                        )
                    else:
                        mean_x = x
                        var_x = torch.zeros_like(x)
                    energy = self.model(x)
                    diagnostics[i, 0, :, :] = (
                        mean_x if n_samples > 1 else mean_x.unsqueeze(0)
                    )
                    diagnostics[i, 1, :, :] = (
                        var_x if n_samples > 1 else var_x.unsqueeze(0)
                    )
                    diagnostics[i, 2, :, :] = energy.view(-1, 1).expand(n_samples, dim)

        if return_trajectory:
            if return_diagnostics:
                return trajectory, diagnostics
            return trajectory
        if return_diagnostics:
            return x, diagnostics
        return x

    def _setup_diagnostics(
        self, dim: int, n_steps: int, n_samples: int = None
    ) -> torch.Tensor:
        if n_samples is not None:
            return torch.empty(
                (n_steps, 3, n_samples, dim), device=self.device, dtype=self.dtype
            )
        return torch.empty((n_steps, 3, dim), device=self.device, dtype=self.dtype)

sample(x=None, dim=10, n_steps=100, n_samples=1, thin=1, return_trajectory=False, return_diagnostics=False, *args, **kwargs)

Generates samples using Langevin dynamics.

Parameters:

Name Type Description Default
x Optional[Tensor]

The initial state to start sampling from. If None, a random state is created.

 None
dim int

The dimension of the state space (if x is not provided).

 10
n_steps int

The number of MCMC steps to perform.

 100
n_samples int

The number of parallel chains/samples to generate.

 1
thin int

The thinning factor (not currently supported).

 1
return_trajectory bool

Whether to return the full sample trajectory.

 False
return_diagnostics bool

Whether to return sampling diagnostics.

 False

Returns:

Type Description
Union[Tensor, Tuple[Tensor, List[dict]]]

Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]: - The final samples. - If return_trajectory is True, the full trajectory. - If return_diagnostics is True, a tuple of samples and diagnostics.
Source code in torchebm/samplers/langevin_dynamics.py 
 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@torch.no_grad()
def sample(
    self,
    x: Optional[torch.Tensor] = None,
    dim: int = 10,
    n_steps: int = 100,
    n_samples: int = 1,
    thin: int = 1,
    return_trajectory: bool = False,
    return_diagnostics: bool = False,
    *args,
    **kwargs,
) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
    """
    Generates samples using Langevin dynamics.

    Args:
        x (Optional[torch.Tensor]): The initial state to start sampling from. If `None`,
            a random state is created.
        dim (int): The dimension of the state space (if `x` is not provided).
        n_steps (int): The number of MCMC steps to perform.
        n_samples (int): The number of parallel chains/samples to generate.
        thin (int): The thinning factor (not currently supported).
        return_trajectory (bool): Whether to return the full sample trajectory.
        return_diagnostics (bool): Whether to return sampling diagnostics.

    Returns:
        Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
            - The final samples.
            - If `return_trajectory` is `True`, the full trajectory.
            - If `return_diagnostics` is `True`, a tuple of samples and diagnostics.
    """

    self.reset_schedulers()

    if x is None:
        x = torch.randn(n_samples, dim, dtype=self.dtype, device=self.device)
    else:
        x = x.to(device=self.device, dtype=self.dtype)
        dim = x.shape[-1]
        n_samples = x.shape[0]

    if return_trajectory:
        trajectory = torch.empty(
            (n_samples, n_steps, dim), dtype=self.dtype, device=self.device
        )

    if return_diagnostics:
        diagnostics = self._setup_diagnostics(dim, n_steps, n_samples=n_samples)

    drift = lambda x_, t_: -self.model.gradient(x_)
    with self.autocast_context():
        for i in range(n_steps):
            self.step_schedulers()
            state = {"x": x}
            x = self.integrator.step(
                state=state,
                step_size=self.get_scheduled_value("step_size"),
                noise_scale=self.get_scheduled_value("noise_scale"),
                drift=drift,
            )["x"]

            if return_trajectory:
                trajectory[:, i, :] = x

            if return_diagnostics:
                if n_samples > 1:
                    mean_x = x.mean(dim=0, keepdim=True)
                    var_x = torch.clamp(
                        x.var(dim=0, unbiased=False, keepdim=True),
                        min=1e-10,
                        max=1e10,
                    )
                else:
                    mean_x = x
                    var_x = torch.zeros_like(x)
                energy = self.model(x)
                diagnostics[i, 0, :, :] = (
                    mean_x if n_samples > 1 else mean_x.unsqueeze(0)
                )
                diagnostics[i, 1, :, :] = (
                    var_x if n_samples > 1 else var_x.unsqueeze(0)
                )
                diagnostics[i, 2, :, :] = energy.view(-1, 1).expand(n_samples, dim)

    if return_trajectory:
        if return_diagnostics:
            return trajectory, diagnostics
        return trajectory
    if return_diagnostics:
        return x, diagnostics
    return x