Skip to content

laplace.curvature #

CurvatureInterface #

CurvatureInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Interface to access curvature for a model and corresponding likelihood. A CurvatureInterface must inherit from this baseclass and implement the necessary functions jacobians, full, kron, and diag. The interface might be extended in the future to account for other curvature structures, for example, a block-diagonal one.

Parameters:

  • model (torch.nn.Module or `laplace.utils.feature_extractor.FeatureExtractor`) –

    torch model (neural network)

  • likelihood (('classification', 'regression'), default: 'classification' ) –
  • last_layer (bool, default: False ) –

    only consider curvature of last layer

  • subnetwork_indices (LongTensor, default: None ) –

    indices of the vectorized model parameters that define the subnetwork to apply the Laplace approximation over

  • dict_key_x (str, default: 'input_ids' ) –

    The dictionary key under which the input tensor x is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

  • dict_key_y (str, default: 'labels' ) –

    The dictionary key under which the target tensor y is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

Attributes:

  • lossfunc (MSELoss or CrossEntropyLoss) –
  • factor (float) –

    conversion factor between torch losses and base likelihoods For example, \(\frac{1}{2}\) to get to \(\mathcal{N}(f, 1)\) from MSELoss.

Source code in laplace/curvature/curvature.py 
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
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
):
    assert likelihood in [Likelihood.REGRESSION, Likelihood.CLASSIFICATION]
    self.likelihood: Likelihood | str = likelihood
    self.model: nn.Module = model
    self.last_layer: bool = last_layer
    self.subnetwork_indices: torch.LongTensor | None = subnetwork_indices
    self.dict_key_x = dict_key_x
    self.dict_key_y = dict_key_y

    if likelihood == "regression":
        self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
            MSELoss(reduction="sum")
        )
        self.factor: float = 0.5
    else:
        self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
            CrossEntropyLoss(reduction="sum")
        )
        self.factor: float = 1.0

    self.params: list[nn.Parameter] = [
        p for p in self._model.parameters() if p.requires_grad
    ]
    self.params_dict: dict[str, nn.Parameter] = {
        k: v for k, v in self._model.named_parameters() if v.requires_grad
    }
    self.buffers_dict: dict[str, torch.Tensor] = {
        k: v for k, v in self.model.named_buffers()
    }

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a dense curvature (approximation) in the form of a \(P \times P\) matrix \(H\) with respect to parameters \(\theta \in \mathbb{R}^P\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    Hessian approximation (parameters, parameters)

Source code in laplace/curvature/curvature.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a dense curvature (approximation) in the form of a \\(P \\times P\\) matrix
    \\(H\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        Hessian approximation `(parameters, parameters)`
    """
    raise NotImplementedError

kron #

kron(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    N: int,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Kron]

Compute a Kronecker factored curvature approximation (such as KFAC). The approximation to \(H\) takes the form of two Kronecker factors \(Q, H\), i.e., \(H \approx Q \otimes H\) for each Module in the neural network permitting such curvature. \(Q\) is quadratic in the input-dimension of a module \(p_{in} \times p_{in}\) and \(H\) in the output-dimension \(p_{out} \times p_{out}\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

  • N (int) –

    total number of data points

Returns:

  • loss ( Tensor ) –
  • H ( `laplace.utils.matrix.Kron` ) –

    Kronecker factored Hessian approximation.

Source code in laplace/curvature/curvature.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def kron(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    N: int,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, Kron]:
    """Compute a Kronecker factored curvature approximation (such as KFAC).
    The approximation to \\(H\\) takes the form of two Kronecker factors \\(Q, H\\),
    i.e., \\(H \\approx Q \\otimes H\\) for each Module in the neural network permitting
    such curvature.
    \\(Q\\) is quadratic in the input-dimension of a module \\(p_{in} \\times p_{in}\\)
    and \\(H\\) in the output-dimension \\(p_{out} \\times p_{out}\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`
    N : int
        total number of data points

    Returns
    -------
    loss : torch.Tensor
    H : `laplace.utils.matrix.Kron`
        Kronecker factored Hessian approximation.
    """
    raise NotImplementedError

diag #

diag(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a diagonal Hessian approximation to \(H\) and is represented as a vector of the dimensionality of parameters \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    vector representing the diagonal of H

Source code in laplace/curvature/curvature.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def diag(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a
    vector of the dimensionality of parameters \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        vector representing the diagonal of H
    """
    raise NotImplementedError

GGNInterface #

GGNInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
    num_samples: int = 1,
)

Bases: CurvatureInterface

Generalized Gauss-Newton or Fisher Curvature Interface. The GGN is equal to the Fisher information for the available likelihoods. In addition to CurvatureInterface, methods for Jacobians are required by subclasses.

Parameters:

  • model (torch.nn.Module or `laplace.utils.feature_extractor.FeatureExtractor`) –

    torch model (neural network)

  • likelihood (('classification', 'regression'), default: 'classification' ) –
  • last_layer (bool, default: False ) –

    only consider curvature of last layer

  • subnetwork_indices (Tensor, default: None ) –

    indices of the vectorized model parameters that define the subnetwork to apply the Laplace approximation over

  • dict_key_x (str, default: 'input_ids' ) –

    The dictionary key under which the input tensor x is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

  • dict_key_y (str, default: 'labels' ) –

    The dictionary key under which the target tensor y is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

  • stochastic (bool, default: False ) –

    Fisher if stochastic else GGN

  • num_samples (int, default: 1 ) –

    Number of samples used to approximate the stochastic Fisher

Source code in laplace/curvature/curvature.py 
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
    num_samples: int = 1,
) -> None:
    self.stochastic: bool = stochastic
    self.num_samples: int = num_samples

    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

kron #

kron(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    N: int,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Kron]

Compute a Kronecker factored curvature approximation (such as KFAC). The approximation to \(H\) takes the form of two Kronecker factors \(Q, H\), i.e., \(H \approx Q \otimes H\) for each Module in the neural network permitting such curvature. \(Q\) is quadratic in the input-dimension of a module \(p_{in} \times p_{in}\) and \(H\) in the output-dimension \(p_{out} \times p_{out}\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

  • N (int) –

    total number of data points

Returns:

  • loss ( Tensor ) –
  • H ( `laplace.utils.matrix.Kron` ) –

    Kronecker factored Hessian approximation.

Source code in laplace/curvature/curvature.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def kron(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    N: int,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, Kron]:
    """Compute a Kronecker factored curvature approximation (such as KFAC).
    The approximation to \\(H\\) takes the form of two Kronecker factors \\(Q, H\\),
    i.e., \\(H \\approx Q \\otimes H\\) for each Module in the neural network permitting
    such curvature.
    \\(Q\\) is quadratic in the input-dimension of a module \\(p_{in} \\times p_{in}\\)
    and \\(H\\) in the output-dimension \\(p_{out} \\times p_{out}\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`
    N : int
        total number of data points

    Returns
    -------
    loss : torch.Tensor
    H : `laplace.utils.matrix.Kron`
        Kronecker factored Hessian approximation.
    """
    raise NotImplementedError

_get_mc_functional_fisher #

_get_mc_functional_fisher(f: Tensor) -> Tensor

Approximate the Fisher's middle matrix (expected outer product of the functional gradient) using MC integral with self.num_samples many samples.

Source code in laplace/curvature/curvature.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def _get_mc_functional_fisher(self, f: torch.Tensor) -> torch.Tensor:
    """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
    using MC integral with `self.num_samples` many samples.
    """
    F = 0

    for _ in range(self.num_samples):
        if self.likelihood == "regression":
            # N(y | f, 1)
            y_sample = f + torch.randn(f.shape, device=f.device, dtype=f.dtype)
            grad_sample = f - y_sample  # functional MSE grad
        else:  # classification with softmax
            y_sample = torch.distributions.Multinomial(logits=f).sample()
            # First functional derivative of the loglik is p - y
            p = torch.softmax(f, dim=-1)
            grad_sample = p - y_sample

        F += (
            1
            / self.num_samples
            * torch.einsum("bc,bk->bck", grad_sample, grad_sample)
        )

    return F

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full GGN \(P \times P\) matrix as Hessian approximation \(H_{ggn}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    GGN (parameters, parameters)

Source code in laplace/curvature/curvature.py 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full GGN \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ggn}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        GGN `(parameters, parameters)`
    """
    Js, f = self.last_layer_jacobians(x) if self.last_layer else self.jacobians(x)
    H_lik = (
        self._get_mc_functional_fisher(f)
        if self.stochastic
        else self._get_functional_hessian(f)
    )

    if H_lik is not None:
        H = torch.einsum("bcp,bck,bkq->pq", Js, H_lik, Js)
    else:  # The case of exact GGN for regression
        H = torch.einsum("bcp,bcq->pq", Js, Js)
    loss = self.factor * self.lossfunc(f, y)

    return loss.detach(), H.detach()

EFInterface #

EFInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvatureInterface

Interface for Empirical Fisher as Hessian approximation. In addition to CurvatureInterface, methods for gradients are required by subclasses.

Parameters:

  • model (torch.nn.Module or `laplace.utils.feature_extractor.FeatureExtractor`) –

    torch model (neural network)

  • likelihood (('classification', 'regression'), default: 'classification' ) –
  • last_layer (bool, default: False ) –

    only consider curvature of last layer

  • subnetwork_indices (Tensor, default: None ) –

    indices of the vectorized model parameters that define the subnetwork to apply the Laplace approximation over

  • dict_key_x (str, default: 'input_ids' ) –

    The dictionary key under which the input tensor x is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

  • dict_key_y (str, default: 'labels' ) –

    The dictionary key under which the target tensor y is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

Attributes:

  • lossfunc (MSELoss or CrossEntropyLoss) –
  • factor (float) –

    conversion factor between torch losses and base likelihoods For example, \(\frac{1}{2}\) to get to \(\mathcal{N}(f, 1)\) from MSELoss.

Source code in laplace/curvature/curvature.py 
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
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
):
    assert likelihood in [Likelihood.REGRESSION, Likelihood.CLASSIFICATION]
    self.likelihood: Likelihood | str = likelihood
    self.model: nn.Module = model
    self.last_layer: bool = last_layer
    self.subnetwork_indices: torch.LongTensor | None = subnetwork_indices
    self.dict_key_x = dict_key_x
    self.dict_key_y = dict_key_y

    if likelihood == "regression":
        self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
            MSELoss(reduction="sum")
        )
        self.factor: float = 0.5
    else:
        self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
            CrossEntropyLoss(reduction="sum")
        )
        self.factor: float = 1.0

    self.params: list[nn.Parameter] = [
        p for p in self._model.parameters() if p.requires_grad
    ]
    self.params_dict: dict[str, nn.Parameter] = {
        k: v for k, v in self._model.named_parameters() if v.requires_grad
    }
    self.buffers_dict: dict[str, torch.Tensor] = {
        k: v for k, v in self.model.named_buffers()
    }

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

kron #

kron(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    N: int,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Kron]

Compute a Kronecker factored curvature approximation (such as KFAC). The approximation to \(H\) takes the form of two Kronecker factors \(Q, H\), i.e., \(H \approx Q \otimes H\) for each Module in the neural network permitting such curvature. \(Q\) is quadratic in the input-dimension of a module \(p_{in} \times p_{in}\) and \(H\) in the output-dimension \(p_{out} \times p_{out}\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

  • N (int) –

    total number of data points

Returns:

  • loss ( Tensor ) –
  • H ( `laplace.utils.matrix.Kron` ) –

    Kronecker factored Hessian approximation.

Source code in laplace/curvature/curvature.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def kron(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    N: int,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, Kron]:
    """Compute a Kronecker factored curvature approximation (such as KFAC).
    The approximation to \\(H\\) takes the form of two Kronecker factors \\(Q, H\\),
    i.e., \\(H \\approx Q \\otimes H\\) for each Module in the neural network permitting
    such curvature.
    \\(Q\\) is quadratic in the input-dimension of a module \\(p_{in} \\times p_{in}\\)
    and \\(H\\) in the output-dimension \\(p_{out} \\times p_{out}\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`
    N : int
        total number of data points

    Returns
    -------
    loss : torch.Tensor
    H : `laplace.utils.matrix.Kron`
        Kronecker factored Hessian approximation.
    """
    raise NotImplementedError

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full EF \(P \times P\) matrix as Hessian approximation \(H_{ef}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H_ef ( Tensor ) –

    EF (parameters, parameters)

Source code in laplace/curvature/curvature.py 
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full EF \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ef}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H_ef : torch.Tensor
        EF `(parameters, parameters)`
    """
    Gs, loss = self.gradients(x, y)
    Gs, loss = Gs.detach(), loss.detach()
    H_ef = torch.einsum("bp,bq->pq", Gs, Gs)
    return self.factor * loss.detach(), self.factor * H_ef

AsdlInterface #

AsdlInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvatureInterface

Interface for asdfghjkl backend.

Source code in laplace/curvature/asdl.py 
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
):
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a dense curvature (approximation) in the form of a \(P \times P\) matrix \(H\) with respect to parameters \(\theta \in \mathbb{R}^P\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    Hessian approximation (parameters, parameters)

Source code in laplace/curvature/curvature.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a dense curvature (approximation) in the form of a \\(P \\times P\\) matrix
    \\(H\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        Hessian approximation `(parameters, parameters)`
    """
    raise NotImplementedError

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_\theta f(x;\theta)\) at current parameter \(\theta\) using asdfghjkl's gradient per output dimension.

Parameters:

  • x (Tensor or MutableMapping(dict, UserDict)) –

    input data (batch, input_shape) on compatible device with model if torch.Tensor. If MutableMapping, then at least contains self.dict_key_x. The latter is specific for reward modeling.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/asdl.py 
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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using asdfghjkl's gradient per output dimension.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping (e.g. dict, UserDict)
        input data `(batch, input_shape)` on compatible device with model if torch.Tensor.
        If MutableMapping, then at least contains `self.dict_key_x`.
        The latter is specific for reward modeling.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    Js = list()
    for i in range(self.model.output_size):

        def closure():
            self.model.zero_grad()
            f = self.model(x)
            loss = f[:, i].sum()
            loss.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            return f

        Ji, f = batch_gradient(
            self.model,
            closure,
            return_outputs=True,
            batch_size=self._get_batch_size(x),
        )
        if self.subnetwork_indices is not None:
            Ji = Ji[:, self.subnetwork_indices]
        Js.append(Ji)
    Js = torch.stack(Js, dim=1)
    return Js, f

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using asdfghjkl's backend.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • loss ( Tensor ) –
  • Gs ( Tensor ) –

    gradients (batch, parameters)

Source code in laplace/curvature/asdl.py 
 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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using asdfghjkl's backend.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    loss : torch.Tensor
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    """

    def closure():
        self.model.zero_grad()
        loss = self.lossfunc(self.model(x), y)
        loss.backward()
        return loss

    Gs, loss = batch_gradient(
        self.model, closure, return_outputs=True, batch_size=self._get_batch_size(x)
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

_get_batch_size #

_get_batch_size(
    x: Tensor | MutableMapping[str, Tensor | Any]
) -> int | None

ASDL assumes that all leading dimensions are the batch size by default (batch_size = None). Here, we want to specify that only the first dimension is the actual batch size. This is the case for LLMs.

Source code in laplace/curvature/asdl.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
def _get_batch_size(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
) -> int | None:
    """
    ASDL assumes that all leading dimensions are the batch size by default (batch_size = None).
    Here, we want to specify that only the first dimension is the actual batch size.
    This is the case for LLMs.
    """
    if isinstance(x, MutableMapping):
        return x[self.dict_key_x].shape[0]
    else:
        return None  # Use ASDL default behavior

AsdlGGN #

AsdlGGN(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
)

Bases: AsdlInterface, GGNInterface

Implementation of the GGNInterface using asdfghjkl.

Source code in laplace/curvature/asdl.py 
435
436
437
438
439
440
441
442
443
444
445
446
447
448
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
):
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )
    self.stochastic = stochastic

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_\theta f(x;\theta)\) at current parameter \(\theta\) using asdfghjkl's gradient per output dimension.

Parameters:

  • x (Tensor or MutableMapping(dict, UserDict)) –

    input data (batch, input_shape) on compatible device with model if torch.Tensor. If MutableMapping, then at least contains self.dict_key_x. The latter is specific for reward modeling.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/asdl.py 
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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using asdfghjkl's gradient per output dimension.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping (e.g. dict, UserDict)
        input data `(batch, input_shape)` on compatible device with model if torch.Tensor.
        If MutableMapping, then at least contains `self.dict_key_x`.
        The latter is specific for reward modeling.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    Js = list()
    for i in range(self.model.output_size):

        def closure():
            self.model.zero_grad()
            f = self.model(x)
            loss = f[:, i].sum()
            loss.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            return f

        Ji, f = batch_gradient(
            self.model,
            closure,
            return_outputs=True,
            batch_size=self._get_batch_size(x),
        )
        if self.subnetwork_indices is not None:
            Ji = Ji[:, self.subnetwork_indices]
        Js.append(Ji)
    Js = torch.stack(Js, dim=1)
    return Js, f

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using asdfghjkl's backend.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • loss ( Tensor ) –
  • Gs ( Tensor ) –

    gradients (batch, parameters)

Source code in laplace/curvature/asdl.py 
 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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using asdfghjkl's backend.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    loss : torch.Tensor
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    """

    def closure():
        self.model.zero_grad()
        loss = self.lossfunc(self.model(x), y)
        loss.backward()
        return loss

    Gs, loss = batch_gradient(
        self.model, closure, return_outputs=True, batch_size=self._get_batch_size(x)
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full GGN \(P \times P\) matrix as Hessian approximation \(H_{ggn}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    GGN (parameters, parameters)

Source code in laplace/curvature/curvature.py 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full GGN \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ggn}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        GGN `(parameters, parameters)`
    """
    Js, f = self.last_layer_jacobians(x) if self.last_layer else self.jacobians(x)
    H_lik = (
        self._get_mc_functional_fisher(f)
        if self.stochastic
        else self._get_functional_hessian(f)
    )

    if H_lik is not None:
        H = torch.einsum("bcp,bck,bkq->pq", Js, H_lik, Js)
    else:  # The case of exact GGN for regression
        H = torch.einsum("bcp,bcq->pq", Js, Js)
    loss = self.factor * self.lossfunc(f, y)

    return loss.detach(), H.detach()

_get_mc_functional_fisher #

_get_mc_functional_fisher(f: Tensor) -> Tensor

Approximate the Fisher's middle matrix (expected outer product of the functional gradient) using MC integral with self.num_samples many samples.

Source code in laplace/curvature/curvature.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def _get_mc_functional_fisher(self, f: torch.Tensor) -> torch.Tensor:
    """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
    using MC integral with `self.num_samples` many samples.
    """
    F = 0

    for _ in range(self.num_samples):
        if self.likelihood == "regression":
            # N(y | f, 1)
            y_sample = f + torch.randn(f.shape, device=f.device, dtype=f.dtype)
            grad_sample = f - y_sample  # functional MSE grad
        else:  # classification with softmax
            y_sample = torch.distributions.Multinomial(logits=f).sample()
            # First functional derivative of the loglik is p - y
            p = torch.softmax(f, dim=-1)
            grad_sample = p - y_sample

        F += (
            1
            / self.num_samples
            * torch.einsum("bc,bk->bck", grad_sample, grad_sample)
        )

    return F

_get_batch_size #

_get_batch_size(
    x: Tensor | MutableMapping[str, Tensor | Any]
) -> int | None

ASDL assumes that all leading dimensions are the batch size by default (batch_size = None). Here, we want to specify that only the first dimension is the actual batch size. This is the case for LLMs.

Source code in laplace/curvature/asdl.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
def _get_batch_size(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
) -> int | None:
    """
    ASDL assumes that all leading dimensions are the batch size by default (batch_size = None).
    Here, we want to specify that only the first dimension is the actual batch size.
    This is the case for LLMs.
    """
    if isinstance(x, MutableMapping):
        return x[self.dict_key_x].shape[0]
    else:
        return None  # Use ASDL default behavior

AsdlEF #

AsdlEF(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: AsdlInterface, EFInterface

Implementation of the EFInterface using asdfghjkl.

Source code in laplace/curvature/asdl.py 
458
459
460
461
462
463
464
465
466
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
):
    super().__init__(model, likelihood, last_layer, None, dict_key_x, dict_key_y)

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_\theta f(x;\theta)\) at current parameter \(\theta\) using asdfghjkl's gradient per output dimension.

Parameters:

  • x (Tensor or MutableMapping(dict, UserDict)) –

    input data (batch, input_shape) on compatible device with model if torch.Tensor. If MutableMapping, then at least contains self.dict_key_x. The latter is specific for reward modeling.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/asdl.py 
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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using asdfghjkl's gradient per output dimension.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping (e.g. dict, UserDict)
        input data `(batch, input_shape)` on compatible device with model if torch.Tensor.
        If MutableMapping, then at least contains `self.dict_key_x`.
        The latter is specific for reward modeling.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    Js = list()
    for i in range(self.model.output_size):

        def closure():
            self.model.zero_grad()
            f = self.model(x)
            loss = f[:, i].sum()
            loss.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            return f

        Ji, f = batch_gradient(
            self.model,
            closure,
            return_outputs=True,
            batch_size=self._get_batch_size(x),
        )
        if self.subnetwork_indices is not None:
            Ji = Ji[:, self.subnetwork_indices]
        Js.append(Ji)
    Js = torch.stack(Js, dim=1)
    return Js, f

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using asdfghjkl's backend.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • loss ( Tensor ) –
  • Gs ( Tensor ) –

    gradients (batch, parameters)

Source code in laplace/curvature/asdl.py 
 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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using asdfghjkl's backend.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    loss : torch.Tensor
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    """

    def closure():
        self.model.zero_grad()
        loss = self.lossfunc(self.model(x), y)
        loss.backward()
        return loss

    Gs, loss = batch_gradient(
        self.model, closure, return_outputs=True, batch_size=self._get_batch_size(x)
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full EF \(P \times P\) matrix as Hessian approximation \(H_{ef}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H_ef ( Tensor ) –

    EF (parameters, parameters)

Source code in laplace/curvature/curvature.py 
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full EF \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ef}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H_ef : torch.Tensor
        EF `(parameters, parameters)`
    """
    Gs, loss = self.gradients(x, y)
    Gs, loss = Gs.detach(), loss.detach()
    H_ef = torch.einsum("bp,bq->pq", Gs, Gs)
    return self.factor * loss.detach(), self.factor * H_ef

_get_batch_size #

_get_batch_size(
    x: Tensor | MutableMapping[str, Tensor | Any]
) -> int | None

ASDL assumes that all leading dimensions are the batch size by default (batch_size = None). Here, we want to specify that only the first dimension is the actual batch size. This is the case for LLMs.

Source code in laplace/curvature/asdl.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
def _get_batch_size(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
) -> int | None:
    """
    ASDL assumes that all leading dimensions are the batch size by default (batch_size = None).
    Here, we want to specify that only the first dimension is the actual batch size.
    This is the case for LLMs.
    """
    if isinstance(x, MutableMapping):
        return x[self.dict_key_x].shape[0]
    else:
        return None  # Use ASDL default behavior

AsdlHessian #

AsdlHessian(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: AsdlInterface

Source code in laplace/curvature/asdl.py 
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model,
        likelihood,
        last_layer,
        subnetwork_indices=None,
        dict_key_x=dict_key_x,
        dict_key_y=dict_key_y,
    )

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_\theta f(x;\theta)\) at current parameter \(\theta\) using asdfghjkl's gradient per output dimension.

Parameters:

  • x (Tensor or MutableMapping(dict, UserDict)) –

    input data (batch, input_shape) on compatible device with model if torch.Tensor. If MutableMapping, then at least contains self.dict_key_x. The latter is specific for reward modeling.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/asdl.py 
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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using asdfghjkl's gradient per output dimension.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping (e.g. dict, UserDict)
        input data `(batch, input_shape)` on compatible device with model if torch.Tensor.
        If MutableMapping, then at least contains `self.dict_key_x`.
        The latter is specific for reward modeling.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    Js = list()
    for i in range(self.model.output_size):

        def closure():
            self.model.zero_grad()
            f = self.model(x)
            loss = f[:, i].sum()
            loss.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            return f

        Ji, f = batch_gradient(
            self.model,
            closure,
            return_outputs=True,
            batch_size=self._get_batch_size(x),
        )
        if self.subnetwork_indices is not None:
            Ji = Ji[:, self.subnetwork_indices]
        Js.append(Ji)
    Js = torch.stack(Js, dim=1)
    return Js, f

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using asdfghjkl's backend.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • loss ( Tensor ) –
  • Gs ( Tensor ) –

    gradients (batch, parameters)

Source code in laplace/curvature/asdl.py 
 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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using asdfghjkl's backend.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    loss : torch.Tensor
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    """

    def closure():
        self.model.zero_grad()
        loss = self.lossfunc(self.model(x), y)
        loss.backward()
        return loss

    Gs, loss = batch_gradient(
        self.model, closure, return_outputs=True, batch_size=self._get_batch_size(x)
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

_get_batch_size #

_get_batch_size(
    x: Tensor | MutableMapping[str, Tensor | Any]
) -> int | None

ASDL assumes that all leading dimensions are the batch size by default (batch_size = None). Here, we want to specify that only the first dimension is the actual batch size. This is the case for LLMs.

Source code in laplace/curvature/asdl.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
def _get_batch_size(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
) -> int | None:
    """
    ASDL assumes that all leading dimensions are the batch size by default (batch_size = None).
    Here, we want to specify that only the first dimension is the actual batch size.
    This is the case for LLMs.
    """
    if isinstance(x, MutableMapping):
        return x[self.dict_key_x].shape[0]
    else:
        return None  # Use ASDL default behavior

BackPackInterface #

BackPackInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvatureInterface

Interface for Backpack backend.

Source code in laplace/curvature/backpack.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

    extend(self._model)
    extend(self.lossfunc)

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a dense curvature (approximation) in the form of a \(P \times P\) matrix \(H\) with respect to parameters \(\theta \in \mathbb{R}^P\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    Hessian approximation (parameters, parameters)

Source code in laplace/curvature/curvature.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a dense curvature (approximation) in the form of a \\(P \\times P\\) matrix
    \\(H\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        Hessian approximation `(parameters, parameters)`
    """
    raise NotImplementedError

kron #

kron(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    N: int,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Kron]

Compute a Kronecker factored curvature approximation (such as KFAC). The approximation to \(H\) takes the form of two Kronecker factors \(Q, H\), i.e., \(H \approx Q \otimes H\) for each Module in the neural network permitting such curvature. \(Q\) is quadratic in the input-dimension of a module \(p_{in} \times p_{in}\) and \(H\) in the output-dimension \(p_{out} \times p_{out}\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

  • N (int) –

    total number of data points

Returns:

  • loss ( Tensor ) –
  • H ( `laplace.utils.matrix.Kron` ) –

    Kronecker factored Hessian approximation.

Source code in laplace/curvature/curvature.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def kron(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    N: int,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, Kron]:
    """Compute a Kronecker factored curvature approximation (such as KFAC).
    The approximation to \\(H\\) takes the form of two Kronecker factors \\(Q, H\\),
    i.e., \\(H \\approx Q \\otimes H\\) for each Module in the neural network permitting
    such curvature.
    \\(Q\\) is quadratic in the input-dimension of a module \\(p_{in} \\times p_{in}\\)
    and \\(H\\) in the output-dimension \\(p_{out} \\times p_{out}\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`
    N : int
        total number of data points

    Returns
    -------
    loss : torch.Tensor
    H : `laplace.utils.matrix.Kron`
        Kronecker factored Hessian approximation.
    """
    raise NotImplementedError

diag #

diag(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a diagonal Hessian approximation to \(H\) and is represented as a vector of the dimensionality of parameters \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    vector representing the diagonal of H

Source code in laplace/curvature/curvature.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def diag(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a
    vector of the dimensionality of parameters \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        vector representing the diagonal of H
    """
    raise NotImplementedError

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\) using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/backpack.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    to_stack = []
    for i in range(model.output_size):
        model.zero_grad()
        out = model(x)
        with backpack(BatchGrad()):
            if model.output_size > 1:
                out[:, i].sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            else:
                out.sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            to_cat = []
            for param in model.parameters():
                to_cat.append(param.grad_batch.reshape(x.shape[0], -1))
                delattr(param, "grad_batch")
            Jk = torch.cat(to_cat, dim=1)
            if self.subnetwork_indices is not None:
                Jk = Jk[:, self.subnetwork_indices]
        to_stack.append(Jk)
        if i == 0:
            f = out

    model.zero_grad()
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)
    if model.output_size > 1:
        J = torch.stack(to_stack, dim=2).transpose(1, 2)
    else:
        J = Jk.unsqueeze(-1).transpose(1, 2)

    return (J, f) if enable_backprop else (J.detach(), f.detach())

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using Backpack's BatchGrad. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/backpack.py 
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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using Backpack's BatchGrad. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """
    f = self.model(x)
    loss = self.lossfunc(f, y)
    with backpack(BatchGrad()):
        loss.backward()
    Gs = torch.cat(
        [p.grad_batch.data.flatten(start_dim=1) for p in self._model.parameters()],
        dim=1,
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping,
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute selected output Jacobians using BackPACK's BatchGrad.

Parameters:

  • x (Tensor) –

    Input data of shape (batch, input_shape) on the model's device.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/backpack.py 
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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping,
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute selected output Jacobians using BackPACK's BatchGrad.

    Parameters
    ----------
    x : torch.Tensor
        Input data of shape `(batch, input_shape)` on the model's device.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != x.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    columns = []
    first_output = None

    for column in range(outputs.shape[1]):
        model.zero_grad(set_to_none=True)
        output = model(x)
        if first_output is None:
            first_output = output

        selected = output.gather(
            1, outputs[:, column : column + 1].to(output.device)
        ).sum()

        with backpack(BatchGrad()):
            selected.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            columns.append(
                torch.cat(
                    [
                        parameter.grad_batch.reshape(x.shape[0], -1)
                        for parameter in model.parameters()
                        if parameter.requires_grad
                    ],
                    dim=1,
                )
            )
            for parameter in model.parameters():
                if hasattr(parameter, "grad_batch"):
                    delattr(parameter, "grad_batch")

    model.zero_grad(set_to_none=True)
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)

    result = torch.stack(columns, dim=1)
    assert first_output is not None

    return (
        (result, first_output)
        if enable_backprop
        else (
            result.detach(),
            first_output.detach(),
        )
    )

BackPackGGN #

BackPackGGN(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
)

Bases: BackPackInterface, GGNInterface

Implementation of the GGNInterface using Backpack.

Source code in laplace/curvature/backpack.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
):
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )
    self.stochastic = stochastic

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\) using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/backpack.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    to_stack = []
    for i in range(model.output_size):
        model.zero_grad()
        out = model(x)
        with backpack(BatchGrad()):
            if model.output_size > 1:
                out[:, i].sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            else:
                out.sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            to_cat = []
            for param in model.parameters():
                to_cat.append(param.grad_batch.reshape(x.shape[0], -1))
                delattr(param, "grad_batch")
            Jk = torch.cat(to_cat, dim=1)
            if self.subnetwork_indices is not None:
                Jk = Jk[:, self.subnetwork_indices]
        to_stack.append(Jk)
        if i == 0:
            f = out

    model.zero_grad()
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)
    if model.output_size > 1:
        J = torch.stack(to_stack, dim=2).transpose(1, 2)
    else:
        J = Jk.unsqueeze(-1).transpose(1, 2)

    return (J, f) if enable_backprop else (J.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping,
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute selected output Jacobians using BackPACK's BatchGrad.

Parameters:

  • x (Tensor) –

    Input data of shape (batch, input_shape) on the model's device.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/backpack.py 
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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping,
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute selected output Jacobians using BackPACK's BatchGrad.

    Parameters
    ----------
    x : torch.Tensor
        Input data of shape `(batch, input_shape)` on the model's device.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != x.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    columns = []
    first_output = None

    for column in range(outputs.shape[1]):
        model.zero_grad(set_to_none=True)
        output = model(x)
        if first_output is None:
            first_output = output

        selected = output.gather(
            1, outputs[:, column : column + 1].to(output.device)
        ).sum()

        with backpack(BatchGrad()):
            selected.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            columns.append(
                torch.cat(
                    [
                        parameter.grad_batch.reshape(x.shape[0], -1)
                        for parameter in model.parameters()
                        if parameter.requires_grad
                    ],
                    dim=1,
                )
            )
            for parameter in model.parameters():
                if hasattr(parameter, "grad_batch"):
                    delattr(parameter, "grad_batch")

    model.zero_grad(set_to_none=True)
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)

    result = torch.stack(columns, dim=1)
    assert first_output is not None

    return (
        (result, first_output)
        if enable_backprop
        else (
            result.detach(),
            first_output.detach(),
        )
    )

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using Backpack's BatchGrad. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/backpack.py 
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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using Backpack's BatchGrad. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """
    f = self.model(x)
    loss = self.lossfunc(f, y)
    with backpack(BatchGrad()):
        loss.backward()
    Gs = torch.cat(
        [p.grad_batch.data.flatten(start_dim=1) for p in self._model.parameters()],
        dim=1,
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full GGN \(P \times P\) matrix as Hessian approximation \(H_{ggn}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    GGN (parameters, parameters)

Source code in laplace/curvature/curvature.py 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full GGN \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ggn}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        GGN `(parameters, parameters)`
    """
    Js, f = self.last_layer_jacobians(x) if self.last_layer else self.jacobians(x)
    H_lik = (
        self._get_mc_functional_fisher(f)
        if self.stochastic
        else self._get_functional_hessian(f)
    )

    if H_lik is not None:
        H = torch.einsum("bcp,bck,bkq->pq", Js, H_lik, Js)
    else:  # The case of exact GGN for regression
        H = torch.einsum("bcp,bcq->pq", Js, Js)
    loss = self.factor * self.lossfunc(f, y)

    return loss.detach(), H.detach()

_get_mc_functional_fisher #

_get_mc_functional_fisher(f: Tensor) -> Tensor

Approximate the Fisher's middle matrix (expected outer product of the functional gradient) using MC integral with self.num_samples many samples.

Source code in laplace/curvature/curvature.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def _get_mc_functional_fisher(self, f: torch.Tensor) -> torch.Tensor:
    """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
    using MC integral with `self.num_samples` many samples.
    """
    F = 0

    for _ in range(self.num_samples):
        if self.likelihood == "regression":
            # N(y | f, 1)
            y_sample = f + torch.randn(f.shape, device=f.device, dtype=f.dtype)
            grad_sample = f - y_sample  # functional MSE grad
        else:  # classification with softmax
            y_sample = torch.distributions.Multinomial(logits=f).sample()
            # First functional derivative of the loglik is p - y
            p = torch.softmax(f, dim=-1)
            grad_sample = p - y_sample

        F += (
            1
            / self.num_samples
            * torch.einsum("bc,bk->bck", grad_sample, grad_sample)
        )

    return F

BackPackEF #

BackPackEF(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: BackPackInterface, EFInterface

Implementation of EFInterface using Backpack.

Source code in laplace/curvature/backpack.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

    extend(self._model)
    extend(self.lossfunc)

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\) using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/backpack.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\)
    using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    to_stack = []
    for i in range(model.output_size):
        model.zero_grad()
        out = model(x)
        with backpack(BatchGrad()):
            if model.output_size > 1:
                out[:, i].sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            else:
                out.sum().backward(
                    create_graph=enable_backprop, retain_graph=enable_backprop
                )
            to_cat = []
            for param in model.parameters():
                to_cat.append(param.grad_batch.reshape(x.shape[0], -1))
                delattr(param, "grad_batch")
            Jk = torch.cat(to_cat, dim=1)
            if self.subnetwork_indices is not None:
                Jk = Jk[:, self.subnetwork_indices]
        to_stack.append(Jk)
        if i == 0:
            f = out

    model.zero_grad()
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)
    if model.output_size > 1:
        J = torch.stack(to_stack, dim=2).transpose(1, 2)
    else:
        J = Jk.unsqueeze(-1).transpose(1, 2)

    return (J, f) if enable_backprop else (J.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping,
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute selected output Jacobians using BackPACK's BatchGrad.

Parameters:

  • x (Tensor) –

    Input data of shape (batch, input_shape) on the model's device.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/backpack.py 
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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping,
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute selected output Jacobians using BackPACK's BatchGrad.

    Parameters
    ----------
    x : torch.Tensor
        Input data of shape `(batch, input_shape)` on the model's device.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    if isinstance(x, MutableMapping):
        raise ValueError("BackPACK backend does not support dict-like inputs!")

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != x.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    input_grad = (
        x.grad.detach().clone() if x.is_leaf and x.grad is not None else None
    )
    model = extend(self.model)
    columns = []
    first_output = None

    for column in range(outputs.shape[1]):
        model.zero_grad(set_to_none=True)
        output = model(x)
        if first_output is None:
            first_output = output

        selected = output.gather(
            1, outputs[:, column : column + 1].to(output.device)
        ).sum()

        with backpack(BatchGrad()):
            selected.backward(
                create_graph=enable_backprop, retain_graph=enable_backprop
            )
            columns.append(
                torch.cat(
                    [
                        parameter.grad_batch.reshape(x.shape[0], -1)
                        for parameter in model.parameters()
                        if parameter.requires_grad
                    ],
                    dim=1,
                )
            )
            for parameter in model.parameters():
                if hasattr(parameter, "grad_batch"):
                    delattr(parameter, "grad_batch")

    model.zero_grad(set_to_none=True)
    if x.is_leaf and x.requires_grad:
        x.grad = input_grad
    CTX.remove_hooks()
    _cleanup(model)

    result = torch.stack(columns, dim=1)
    assert first_output is not None

    return (
        (result, first_output)
        if enable_backprop
        else (
            result.detach(),
            first_output.detach(),
        )
    )

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\) using Backpack's BatchGrad. Note that BackPACK doesn't play well with torch.func, so this method has to be overridden.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/backpack.py 
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
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
    \\(\\theta\\) using Backpack's BatchGrad. Note that BackPACK doesn't play well
    with torch.func, so this method has to be overridden.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """
    f = self.model(x)
    loss = self.lossfunc(f, y)
    with backpack(BatchGrad()):
        loss.backward()
    Gs = torch.cat(
        [p.grad_batch.data.flatten(start_dim=1) for p in self._model.parameters()],
        dim=1,
    )
    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]
    return Gs, loss

full #

full(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
) -> tuple[Tensor, Tensor]

Compute the full EF \(P \times P\) matrix as Hessian approximation \(H_{ef}\) with respect to parameters \(\theta \in \mathbb{R}^P\). For last-layer, reduced to \(\theta_{last}\)

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H_ef ( Tensor ) –

    EF (parameters, parameters)

Source code in laplace/curvature/curvature.py 
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def full(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the full EF \\(P \\times P\\) matrix as Hessian approximation
    \\(H_{ef}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
    For last-layer, reduced to \\(\\theta_{last}\\)

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H_ef : torch.Tensor
        EF `(parameters, parameters)`
    """
    Gs, loss = self.gradients(x, y)
    Gs, loss = Gs.detach(), loss.detach()
    H_ef = torch.einsum("bp,bq->pq", Gs, Gs)
    return self.factor * loss.detach(), self.factor * H_ef

CurvlinopsInterface #

CurvlinopsInterface(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvatureInterface

Interface for Curvlinops backend. https://github.com/f-dangel/curvlinops

Source code in laplace/curvature/curvlinops.py 
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

diag #

diag(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a diagonal Hessian approximation to \(H\) and is represented as a vector of the dimensionality of parameters \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    vector representing the diagonal of H

Source code in laplace/curvature/curvature.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def diag(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a
    vector of the dimensionality of parameters \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        vector representing the diagonal of H
    """
    raise NotImplementedError

CurvlinopsGGN #

CurvlinopsGGN(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
)

Bases: CurvlinopsInterface, GGNInterface

Implementation of the GGNInterface using Curvlinops.

Source code in laplace/curvature/curvlinops.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    stochastic: bool = False,
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )
    self.stochastic = stochastic

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

_get_mc_functional_fisher #

_get_mc_functional_fisher(f: Tensor) -> Tensor

Approximate the Fisher's middle matrix (expected outer product of the functional gradient) using MC integral with self.num_samples many samples.

Source code in laplace/curvature/curvature.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def _get_mc_functional_fisher(self, f: torch.Tensor) -> torch.Tensor:
    """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
    using MC integral with `self.num_samples` many samples.
    """
    F = 0

    for _ in range(self.num_samples):
        if self.likelihood == "regression":
            # N(y | f, 1)
            y_sample = f + torch.randn(f.shape, device=f.device, dtype=f.dtype)
            grad_sample = f - y_sample  # functional MSE grad
        else:  # classification with softmax
            y_sample = torch.distributions.Multinomial(logits=f).sample()
            # First functional derivative of the loglik is p - y
            p = torch.softmax(f, dim=-1)
            grad_sample = p - y_sample

        F += (
            1
            / self.num_samples
            * torch.einsum("bc,bk->bck", grad_sample, grad_sample)
        )

    return F

CurvlinopsEF #

CurvlinopsEF(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvlinopsInterface, EFInterface

Implementation of EFInterface using Curvlinops.

Source code in laplace/curvature/curvlinops.py 
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

CurvlinopsHessian #

CurvlinopsHessian(
    model: Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
)

Bases: CurvlinopsInterface

Implementation of the full Hessian using Curvlinops.

Source code in laplace/curvature/curvlinops.py 
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    last_layer: bool = False,
    subnetwork_indices: torch.LongTensor | None = None,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
) -> None:
    super().__init__(
        model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
    )

jacobians #

jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta} f(x;\theta)\) at current parameter \(\theta\), via torch.func.

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • enable_backprop (bool, default: = False ) –

    whether to enable backprop through the Js and f w.r.t. x

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, parameters, outputs)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
 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
def jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
    via torch.func.

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    enable_backprop : bool, default = False
        whether to enable backprop through the Js and f w.r.t. x

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, parameters, outputs)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """

    def model_fn_params_only(params_dict, buffers_dict):
        out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
        return out, out

    Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
        self.params_dict, self.buffers_dict
    )

    # Concatenate over flattened parameters
    Js = [
        j.flatten(start_dim=-p.dim())
        for j, p in zip(Js.values(), self.params_dict.values())
    ]
    Js = torch.cat(Js, dim=-1)

    if self.subnetwork_indices is not None:
        Js = Js[:, :, self.subnetwork_indices]

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

last_layer_jacobians #

last_layer_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute Jacobians \(\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})\) only at current last-layer parameter \(\theta_{\textrm{last}}\).

Parameters:

  • x (Tensor) –
  • enable_backprop (bool, default: False ) –

Returns:

  • Js ( Tensor ) –

    Jacobians (batch, outputs, last-layer-parameters)

  • f ( Tensor ) –

    output function (batch, outputs)

Source code in laplace/curvature/curvature.py 
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
def last_layer_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
    only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).

    Parameters
    ----------
    x : torch.Tensor
    enable_backprop : bool, default=False

    Returns
    -------
    Js : torch.Tensor
        Jacobians `(batch, outputs, last-layer-parameters)`
    f : torch.Tensor
        output function `(batch, outputs)`
    """
    f, phi = self.model.forward_with_features(x)
    bsize = phi.shape[0]
    output_size = int(f.numel() / bsize)

    # calculate Jacobians using the feature vector 'phi'
    p = next(self.model.parameters())
    identity = (
        torch.eye(output_size, device=p.device, dtype=p.dtype)
        .unsqueeze(0)
        .tile(bsize, 1, 1)
    )
    # Jacobians are batch x output x params
    Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
    if self.model.last_layer.bias is not None:
        Js = torch.cat([Js, identity], dim=2)

    return (Js, f) if enable_backprop else (Js.detach(), f.detach())

selected_output_jacobians #

selected_output_jacobians(
    x: Tensor | MutableMapping[str, Tensor | Any],
    outputs: Tensor,
    enable_backprop: bool = False,
) -> tuple[Tensor, Tensor]

Compute parameter Jacobians for selected output coordinates per input.

The default computes all output Jacobians before selecting coordinates. Backends may override it to avoid materializing the full Jacobian.

Parameters:

  • x (Tensor or MutableMapping) –

    Input data on the model's device, with a leading batch dimension.

  • outputs (Tensor) –

    Output indices for each input, shaped (batch,) or (batch, selected_outputs).

  • enable_backprop (bool, default: False ) –

    Whether to retain gradients through the Jacobians and model outputs.

Returns:

  • Js ( Tensor ) –

    Selected Jacobians of shape (batch, selected_outputs, parameters).

  • f ( Tensor ) –

    Model outputs of shape (batch, outputs).

Source code in laplace/curvature/curvature.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def selected_output_jacobians(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    outputs: torch.Tensor,
    enable_backprop: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute parameter Jacobians for selected output coordinates per input.

    The default computes all output Jacobians before selecting coordinates.
    Backends may override it to avoid materializing the full Jacobian.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Input data on the model's device, with a leading batch dimension.
    outputs : torch.Tensor
        Output indices for each input, shaped `(batch,)` or
        `(batch, selected_outputs)`.
    enable_backprop : bool, default=False
        Whether to retain gradients through the Jacobians and model outputs.

    Returns
    -------
    Js : torch.Tensor
        Selected Jacobians of shape `(batch, selected_outputs, parameters)`.
    f : torch.Tensor
        Model outputs of shape `(batch, outputs)`.
    """
    # Use the output shape observed at this input. Reward models return two
    # training logits but one output for an individual inference input.
    jacobians, values = CurvatureInterface.jacobians(
        self, x, enable_backprop=enable_backprop
    )

    if outputs.ndim == 1:
        outputs = outputs[:, None]

    if outputs.shape[0] != jacobians.shape[0]:
        raise ValueError("Output selectors must match the input batch size.")

    selected = torch.gather(
        jacobians,
        dim=1,
        index=outputs.to(jacobians.device, dtype=torch.long)
        .unsqueeze(-1)
        .expand(-1, -1, jacobians.shape[-1]),
    )

    return selected, values

gradients #

gradients(
    x: Tensor | MutableMapping[str, Tensor | Any], y: Tensor
) -> tuple[Tensor, Tensor]

Compute batch gradients \(\nabla_\theta \ell(f(x;\theta, y)\) at current parameter \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape) on compatible device with model.

  • y (Tensor) –

Returns:

  • Gs ( Tensor ) –

    gradients (batch, parameters)

  • loss ( Tensor ) –
Source code in laplace/curvature/curvature.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def gradients(
    self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
    current parameter \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)` on compatible device with model.
    y : torch.Tensor

    Returns
    -------
    Gs : torch.Tensor
        gradients `(batch, parameters)`
    loss : torch.Tensor
    """

    def loss_single(x, y, params_dict, buffers_dict):
        """Compute the gradient for a single sample."""
        x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
        output = torch.func.functional_call(
            self.model, (params_dict, buffers_dict), x
        )
        loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
        return loss, loss

    grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
    batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))

    batch_grad, batch_loss = batch_grad_fn(
        x, y, self.params_dict, self.buffers_dict
    )
    Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)

    if self.subnetwork_indices is not None:
        Gs = Gs[:, self.subnetwork_indices]

    loss = batch_loss.sum(0)

    return Gs, loss

diag #

diag(
    x: Tensor | MutableMapping[str, Tensor | Any],
    y: Tensor,
    **kwargs: dict[str, Any]
)

Compute a diagonal Hessian approximation to \(H\) and is represented as a vector of the dimensionality of parameters \(\theta\).

Parameters:

  • x (Tensor) –

    input data (batch, input_shape)

  • y (Tensor) –

    labels (batch, label_shape)

Returns:

  • loss ( Tensor ) –
  • H ( Tensor ) –

    vector representing the diagonal of H

Source code in laplace/curvature/curvature.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def diag(
    self,
    x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
    y: torch.Tensor,
    **kwargs: dict[str, Any],
):
    """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a
    vector of the dimensionality of parameters \\(\\theta\\).

    Parameters
    ----------
    x : torch.Tensor
        input data `(batch, input_shape)`
    y : torch.Tensor
        labels `(batch, label_shape)`

    Returns
    -------
    loss : torch.Tensor
    H : torch.Tensor
        vector representing the diagonal of H
    """
    raise NotImplementedError