如何用PyTorch计算线性回归模型的Hessian矩阵以获取参数标准误

假设出于某些原因，我希望使用PyTorch拟合线性回归，现有示例代码如下，请问如何计算该模型的Hessian矩阵，最终得到参数估计的标准误？

PyTorch线性回归实现代码 #

import torch 
import torch.nn as nn
# 设置随机种子
torch.manual_seed(42)
# 定义模型
class OLS_pytorch(nn.Module):
    def __init__(self, X, Y):
        super(OLS_pytorch, self).__init__()
        self.X = X
        self.Y = Y
        self.beta = nn.Parameter(torch.ones(X.shape[1], 1, requires_grad=True))
        self.intercept = nn.Parameter(torch.ones(1, requires_grad=True))
        self.loss = nn.MSELoss()
        
    def forward(self):
        return self.X @ self.beta + self.intercept
    
    def fit(self, lr=0.01, epochs=1000):
        optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        for epoch in range(epochs):
            optimizer.zero_grad()
            loss = self.loss(self.forward(), self.Y)
            loss.backward()
            optimizer.step()
            if epoch % 10 == 0:
                print(f"Epoch {epoch} loss: {loss.item()}")
        return self

生成数据并拟合模型 #

# 生成数据    
X = torch.randn(100, 1)
Y = 2 * X + 3 + torch.randn(100, 1)
# 拟合模型
model = OLS_pytorch(X, Y)
model.fit()
# 提取参数 
model.beta, model.intercept

# 输出结果：
# Epoch 980 loss: 0.7803605794906616
# Epoch 990 loss: 0.7803605794906616
# (Parameter containing:
# tensor([[2.0118]], requires_grad=True),
# Parameter containing:
# tensor([3.0357], requires_grad=True))

R语言对比实现 #

使用相同数据，R语言的lm()函数不仅能得到一致的参数估计，还可以直接获取Hessian矩阵并计算标准误：

ols <- lm(Y ~ X, data = xy)
ols$coefficients
# (Intercept)           X 
#    3.035674    2.011811 
vcov(ols)
#               (Intercept)             X
# (Intercept)  0.0079923921 -0.0004940884
# X           -0.0004940884  0.0082671053

summary(ols)
# Coefficients:
#             Estimate Std. Error t value Pr(>|t|)    
# (Intercept)  3.03567    0.08940   33.96   <2e-16 ***
# X            2.01181    0.09092   22.13   <2e-16 ***
# ---
# Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1

适配R输出的标准误计算实现 #

以下是能匹配R中lm()输出标准误的PyTorch实现：

# 生成预测值   
y_pred = model.X @ model.beta + model.intercept
N = len(model.Y)
sigma_hat = torch.sum((y_pred - model.Y)**2)/ (N-2) # 2是估计的参数数量

from torch.autograd.functional import hessian
def loss(beta, intercept):
    y_pred = model.X @ beta + intercept
    return model.loss(y_pred, model.Y)

# 计算Hessian矩阵
H = torch.Tensor(hessian(loss, (model.beta, model.intercept)))
# 计算参数的标准误
vcov = torch.sqrt(torch.diag(sigma_hat*torch.inverse(H/2)) )
print(vcov)
# 输出：tensor([0.9092, 0.8940], grad_fn=<SqrtBackward0>)

内容的提问来源于Stack Exchange，提问作者Álvaro A. Gutiérrez-Vargas