Laplace

pytest lint format

</div>

laplace 包可以为整个神经网络、神经网络的子网络或仅最后一层应用拉普拉斯近似。该包支持后验近似、边际似然估计和各种后验预测计算。库文档可在 https://aleximmer.github.io/Laplace 获取。

此外还有一篇相关论文《Laplace Redux — 轻松实现贝叶斯深度学习》（https://arxiv.org/abs/2106.14806），介绍了该库，提供了拉普拉斯近似的入门知识，回顾了它在深度学习中的应用，并通过实验展示了其多功能性和竞争力。使用我们的库时，请考虑引用这篇论文：

@inproceedings{laplace2021,
  title={Laplace Redux--Effortless {B}ayesian Deep Learning},
  author={Erik Daxberger and Agustinus Kristiadi and Alexander Immer
          and Runa Eschenhagen and Matthias Bauer and Philipp Hennig},
  booktitle={{N}eur{IPS}},
  year={2021}
}

论文中实验的复现代码也已公开；它提供了如何使用我们的库进行预测不确定性量化、模型选择和持续学习的示例。

[!IMPORTANT] 作为用户，不应期望 Laplace 能自动工作。应该尝试不同的 Laplace 选项（hessian_factorization、prior precision tuning method、predictive method、backend 等）。可以查看使用 Laplace 的各种论文，了解如何根据具体应用/问题设置这些选项。

安装

[!IMPORTANT] 我们假设 Python 版本 >= 3.9，因为较低版本（即将）被弃用。为了完全兼容，还需要 PyTorch 2.0 及以上版本。

要使用 pip 安装 laplace，请运行以下命令：

pip install laplace-torch

此外，如果您想使用 asdfghjkl 后端，请通过以下方式安装：

pip install git+https://git@github.com/wiseodd/asdl@asdfghjkl

设置开发环境

出于开发目的，例如如果您想做出贡献，请按照以下步骤操作：

安装 uv
然后克隆此仓库并安装开发依赖项：

git clone git@github.com:aleximmer/Laplace.git
uv sync --all-extras

laplace-torch 现在可以在可编辑模式下使用，例如，您可以运行：

uv run python example/regression_example.py

# 或等效地：
source .venv/bin/activate
python example/regression_example.py

[!NOTE] 请参阅贡献指南。我们期待您的贡献！

使用示例

简单用法

在以下示例中，加载了一个预训练模型，然后将拉普拉斯近似拟合到训练数据（使用对所有参数的对角 Hessian 近似），并使用交叉验证 "gridsearch" 优化先验精度。之后，使用 "probit" 预测方法的结果 LA 进行分类预测。

[!IMPORTANT] Laplace 期望所有数据加载器（如下面的 train_loader 和 val_loader）是 PyTorch DataLoader 的实例。每个批次，next(iter(data_loader)) 必须是标准的 (X, y) 张量或包含至少在 Laplace 构造函数中指定的 dict_key_x 和 dict_key_y 键的字典类对象。

[!IMPORTANT] 所有数据加载器中的总数据点数必须可通过 len(train_loader.dataset) 访问。

[!IMPORTANT] 在 optimize_prior_precision 中，确保参数与预测时在 la(x, ...) 中传递的参数匹配。

from laplace import Laplace

# 预训练模型
model = load_map_model()

# 用户指定的 LA 风格
la = Laplace(model, "classification",
             subset_of_weights="all",
             hessian_structure="diag")
la.fit(train_loader)
la.optimize_prior_precision(
    method="gridsearch",
    pred_type="glm",
    link_approx="probit",
    val_loader=val_loader
)

# 用户指定的预测近似
pred = la(x, pred_type="glm", link_approx="probit")

边际似然

边际似然可用于模型选择 [10]，并且对先验精度或观测噪声等连续超参数可微分。在这里，我们拟合库默认的 KFAC 最后一层 LA，并对对数边际似然进行微分。

from laplace import Laplace

# 未训练或预训练模型
model = load_model()

# 默认为推荐的最后一层 KFAC LA：
la = Laplace(model, likelihood="regression")
la.fit(train_loader)

# 关于先验精度和观测噪声的 ML
ml = la.log_marginal_likelihood(prior_prec, obs_noise)
ml.backward()

在 LLM 上使用 Laplace

[!TIP] 该库还支持 Huggingface 模型和参数高效微调。完整说明请参见 examples/huggingface_examples.py 和 examples/huggingface_examples.md。

首先，我们需要包装预训练模型，使 forward 方法接受字典类输入。请注意，当您遍历 Huggingface 数据加载器时，默认情况下就会得到这种输入。使用字典类输入很方便，因为不同的模型有不同数量的输入（例如，GPT 类 LLM 只接受 input_ids，而 BERT 类模型同时接受 input_ids 和 attention_mask 等）。在这个 forward 方法中，您可以进行常规的预处理，如将张量输入移动到正确的设备。

class MyGPT2(nn.Module):
    def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
        super().__init__()
        config = GPT2Config.from_pretrained("gpt2")
        config.pad_token_id = tokenizer.pad_token_id
        config.num_labels = 2
        self.hf_model = GPT2ForSequenceClassification.from_pretrained(
            "gpt2", config=config
        )

    def forward(self, data: MutableMapping) -> torch.Tensor:
        device = next(self.parameters()).device
        input_ids = data["input_ids"].to(device)
        attn_mask = data["attention_mask"].to(device)
        output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
        return output_dict.logits

然后你可以"选择"要对哪些LLM参数应用拉普拉斯近似,方法是关闭"不需要"的参数的梯度。例如,我们可以复制一个最后一层的拉普拉斯近似:(在实际使用中,请使用Laplace(..., subset_of_weights='last_layer', ...)来代替,尽管如此!)

model = MyGPT2(tokenizer)
model.eval()

# 只为最后一层启用梯度
for p in model.hf_model.parameters():
    p.requires_grad = False
for p in model.hf_model.score.parameters():
    p.requires_grad = True

la = Laplace(
    model,
    likelihood="classification",
    # 只会影响最后一层,因为只有它是梯度启用的
    subset_of_weights="all",
    hessian_structure="diag",
)
la.fit(dataloader)
la.optimize_prior_precision()

test_data = next(iter(dataloader))
pred = la(test_data)

这很有用,因为我们可以只对参数高效微调的权重应用LA。例如,我们可以固定LLM本身,只对LoRA权重应用拉普拉斯近似。Huggingface会自动关闭非LoRA权重的梯度。

def get_lora_model():
    model = MyGPT2(tokenizer)  # 注意我们没有禁用梯度
    config = LoraConfig(
        r=4,
        lora_alpha=16,
        target_modules=["c_attn"],  # 对注意力权重应用LoRA
        lora_dropout=0.1,
        bias="none",
    )
    lora_model = get_peft_model(model, config)
    return lora_model

lora_model = get_lora_model()

# 在这里像往常一样训练...

lora_model.eval()

lora_la = Laplace(
    lora_model,
    likelihood="classification",
    subset_of_weights="all",
    hessian_structure="diag",
    backend=AsdlGGN,
)

test_data = next(iter(dataloader))
lora_pred = lora_la(test_data)

子网络拉普拉斯

这个例子展示了如何仅对神经网络中的子网络拟合拉普拉斯近似(同时将所有其他参数固定在它们的MAP估计值),这是在[11]中提出的。它还举例说明了指定要执行推理的子网络的不同方法。

首先,我们使用SubnetLaplace,在这里我们通过生成活跃模型参数的索引列表来指定子网络。

from laplace import Laplace

# 预训练模型
model = load_model()

# 指定子网络的不同方法示例
# 通过向量化模型参数的索引
#
# 示例1: 选择幅度最大的128个参数
from laplace.utils import LargestMagnitudeSubnetMask
subnetwork_mask = LargestMagnitudeSubnetMask(model, n_params_subnet=128)
subnetwork_indices = subnetwork_mask.select()

# 示例2: 指定定义子网络的层
from laplace.utils import ModuleNameSubnetMask
subnetwork_mask = ModuleNameSubnetMask(model, module_names=["layer.1", "layer.3"])
subnetwork_mask.select()
subnetwork_indices = subnetwork_mask.indices

# 示例3: 通过自定义子网络索引手动定义子网络
import torch
subnetwork_indices = torch.tensor([0, 4, 11, 42, 123, 2021])

# 使用指定的子网络索引定义和拟合子网络LA
la = Laplace(model, "classification",
             subset_of_weights="subnetwork",
             hessian_structure="full",
             subnetwork_indices=subnetwork_indices)
la.fit(train_loader)

除了SubnetLaplace,如前所述,你也可以使用Laplace(..., subset_of_weights='last_layer')只处理最后一层,这使用了LLLaplace。作为第三种方法,你可以通过禁用固定模型参数的梯度来定义子网络。不同的方法针对不同的使用场景。每种方法都有优缺点,详情请参见这个讨论。总结如下:

禁用梯度: 在特定类型的层/参数上执行拉普拉斯的通用方法,例如在带有LoRA的LLM中。也可以用来模拟LLLaplace。对于这种方法,始终使用subset_of_weights='all'。
- 通过禁用梯度选择子网比SubnetLaplace更有效,因为它避免了首先计算完整的雅可比矩阵
- 禁用梯度只能在Parameter级别执行,而不能针对单个权重,所以这并不涵盖SubnetLaplace提供的所有情况,如Largest*SubnetMask或RandomSubnetMask
LLLaplace: 针对最后一层的特定代码,性能得到改进(#145)
SubnetLaplace: 更细粒度的分区,如LargestMagnitudeSubnetMask

序列化

与普通的torch一样,我们支持两种序列化数据的方式。

一种是熟悉的state_dict方法。在这里你需要保存和重新创建model和Laplace。使用这种方法来长期存储模型和共享拟合好的Laplace实例。

# 保存模型和Laplace实例
torch.save(model.state_dict(), "model_state_dict.bin")
torch.save(la.state_dict(), "la_state_dict.bin")

# 加载序列化数据
model2 = MyModel(...)
model2.load_state_dict(torch.load("model_state_dict.bin"))
la2 = Laplace(model2, "classification",
              subset_of_weights="all",
              hessian_structure="diag")
la2.load_state_dict(torch.load("la_state_dict.bin"))

第二种方法是保存整个Laplace对象,包括self.model。这种方法不那么繁琐,更方便,因为你可以将训练好的模型和拟合好的Laplace数据存储在一起,但也有一些缺点。在实验过程中用于快速保存-加载周期。

# 保存Laplace,包括la.model
torch.save(la, "la.pt")

# 加载两者
torch.load("la.pt")

一些Laplace变体,如LLLaplace,在使用默认的pickle模块(即torch.save()和torch.load()使用的模块)序列化时可能会遇到问题(AttributeError: Can't pickle local object ...)。在这种情况下,dill包会很有用。

import dill

torch.save(la, "la.pt", pickle_module=dill)

对于这两种方法,你可以自由切换设备,例如当你在GPU上训练但想在CPU上运行预测时。在这种情况下,使用

torch.load(..., map_location="cpu")

[!警告] 目前，该库始终假设模型的输出张量形状为(batch_size, ..., n_classes)，因此在图像输出的情况下，您需要将其从NCHW重新排列为NHWC。

结构

laplace包由两个主要组件组成：

laplace.BaseLaplace的子类，实现了不同的稀疏结构：不同的权重子集（'all'、'subnetwork'和'last_layer'）以及Hessian近似的不同结构（'full'、'kron'、'lowrank'、'diag'和'gp'）。这导致了目前有_十种_可用选项：laplace.FullLaplace、laplace.KronLaplace、laplace.DiagLaplace、laplace.FunctionalLaplace，以及相应的最后一层变体laplace.FullLLLaplace、laplace.KronLLLaplace、laplace.DiagLLLaplace和laplace.FunctionalLLLaplace（它们都是laplace.LLLaplace的子类），laplace.SubnetLaplace（仅支持'full'和'diag'的Hessian近似）以及laplace.LowRankLaplace（仅支持对'all'权重进行推断）。所有这些都可以通过laplace.Laplace函数方便地访问。
laplace.curvature中的后端，提供对应稀疏结构的Hessian近似的访问，例如对角线GGN。

此外，该包还提供了用于将神经网络分解为特征提取器和最后一层的工具，用于LLLaplace子类（laplace.utils.feature_extractor）以及有效处理Kronecker因子的工具（laplace.utils.matrix）。

最后，该包为SubnetLaplace实现了几种选择/指定子网络的选项（作为laplace.utils.subnetmask.SubnetMask的子类）。自动子网络选择策略包括：均匀随机选择（laplace.utils.subnetmask.RandomSubnetMask）、按最大参数幅度选择（LargestMagnitudeSubnetMask）以及按最大边缘参数方差选择（LargestVarianceDiagLaplaceSubnetMask和LargestVarianceSWAGSubnetMask）。除此之外，还可以手动指定子网络，通过列出要进行Laplace推断的模型参数名称（ParamNameSubnetMask）或模块名称（ModuleNameSubnetMask）。

可扩展性

要扩展laplace包，可以设计新的BaseLaplace子类，例如具有块对角Hessian结构的Laplace。还可以实现自定义的子网络选择策略作为SubnetMask的新子类。

另外，扩展或集成后端（curvature.curvature的子类）可以为Laplace近似提供不同的Hessian近似。例如，目前基于Curvlinops和原生torch.func（以前称为functorch）的curvature.CurvlinopsInterface、基于BackPACK的curvature.BackPackInterface以及基于ASDL的curvature.AsdlInterface都是可用的。

何时使用哪个后端

[!提示] 每个后端都有自己的特点/行为。根据您的模型和应用，可以使用以下指南来选择合适的后端。

小型、简单的MLP，或最后一层Laplace： 任何后端都应该效果不错。如果hessian_factorization = 'kron'，推荐使用CurvlinopsGGN或CurvlinopsEF，但对于其他因子分解方法效率较低。
使用PEFT（如LoRA）的LLM： 推荐使用AsdlGGN和AsdlEF。
连续贝叶斯优化： 推荐使用CurvlinopsGGN/EF和BackpackGGN/EF，因为它们是唯一支持对Jacobian进行反向传播的后端。

[!注意] 对于全部和对角线因子分解，curvlinops后端效率较低。此外，由于它们依赖于torch.func.jacrev和torch.func.vmap的组合，计算大型模型的Jacobian也效率较低！最后，curvlinops仅为nn.Linear和nn.Conv2d模块（包括更大模块内的这些模块，如Attention）计算K-FAC（hessian_factorization = 'kron'）。

[!注意] BackPack后端仅限于表示为nn.Sequential的模型。此外，它们与归一化层不兼容。

文档

文档可在此处获取，或者可以在本地生成和/或查看：

# 假设已克隆仓库
uv sync --all-extras
# 创建文档并写入html
uv run bash update_docs.sh
# ..或直接提供文档服务
uv run pdoc --http 0.0.0.0:8080 laplace --template-dir template

贡献

非常欢迎提交拉取请求。请遵循以下指南：

通过uv sync --all-extras安装Laplace，这将安装ruff和运行测试及构建文档所需的所有依赖项。
使用ruff作为自动格式化工具。请参考以下makefile并通过make ruff运行。请注意，ruff check --fix和ruff format的顺序很重要！
同样使用ruff作为代码检查工具。在打开拉取请求之前，请手动修复所有的代码检查错误/警告。
以Python文档字符串、类型提示的形式全面记录您的更改，并在适用的情况下在./examples子目录中提供代码/markdown示例。
提供尽可能多的测试用例。确保所有测试用例都通过。

欢迎提出问题、错误报告和想法！

有用的链接

使用uv发布包：https://docs.astral.sh/uv/guides/publish/

参考文献

这个程序包依赖于对神经网络的拉普拉斯近似的多项改进，这最初是由MacKay提出的[1]。如果您通过我们的laplace库使用了他们提出的任何方法，请考虑引用相应的论文。

[1] MacKay, DJC. [反向传播网络的实用贝叶斯框架]. 神经计算 1992.
[2] Gibbs, M. N. [用于回归和分类的贝叶斯高斯过程]. 博士论文 1997.
[3] Snoek, J., Rippel, O., Swersky, K., Kiros, R., Satish, N., Sundaram, N., Patwary, M., Prabhat, M., Adams, R. [使用深度神经网络的可扩展贝叶斯优化]. ICML 2015.
[4] Ritter, H., Botev, A., Barber, D. [神经网络的可扩展拉普拉斯近似]. ICLR 2018.
[5] Foong, A. Y., Li, Y., Hernández-Lobato, J. M., Turner, R. E. [贝叶斯神经网络中的"中间"不确定性]. ICML UDL研讨会 2019.
[6] Khan, M. E., Immer, A., Abedi, E., Korzepa, M. [近似推断将深度网络转化为高斯过程]. NeurIPS 2019.
[7] Kristiadi, A., Hein, M., Hennig, P. [即使只是一点点贝叶斯，也能修复ReLU网络中的过度自信]. ICML 2020.
[8] Immer, A., Korzepa, M., Bauer, M. [通过局部线性化改进贝叶斯神经网络的预测]. AISTATS 2021.
[9] Sharma, A., Azizan, N., Pavone, M. [为深度神经网络勾勒曲率以实现高效的分布外检测]. UAI 2021.
[10] Immer, A., Bauer, M., Fortuin, V., Rätsch, G., Khan, EM. [用于深度学习模型选择的可扩展边际似然估计]. ICML 2021.
[11] Daxberger, E., Nalisnick, E., Allingham, JU., Antorán, J., Hernández-Lobato, JM. [通过子网络推断实现贝叶斯深度学习]. ICML 2021.