如何用矩阵/向量乘法加速PyTorch几何衰减(Adstock)张量操作？

我想通过矩阵/向量乘法加速下面的张量操作，有没有高效方案？要求兼容形状为0（torch.Size([])）的张量，但我没法初始化这类张量。

当前实现代码 #

def adstock_geometric(x: torch.Tensor, theta: float):
    x_decayed = torch.zeros_like(x)
    x_decayed[0] = x[0]

    for xi in range(1, len(x_decayed)):
        x_decayed[xi] = x[xi] + theta * x_decayed[xi - 1]

    return x_decayed

def adstock_multiple_samples(x: torch.Tensor, theta: torch.Tensor):

    listtheta = theta.tolist()
    if isinstance(listtheta, float):
        return adstock_geometric(x=x,
                                 theta=theta)
    x_decayed = torch.zeros((100, 112, 1))
    for idx, theta_ in enumerate(listtheta):
        x_decayed_one_entry = adstock_geometric(x=x,
                                                theta=theta_)
        x_decayed[idx] = x_decayed_one_entry
    return x_decayed

if __name__ == '__main__':
    ones = torch.tensor([1])
    hundreds = torch.tensor([idx for idx in range(100)])
    x = torch.tensor([[idx] for idx in range(112)])
    ones = adstock_multiple_samples(x=x,
                                    theta=ones)
    hundreds = adstock_multiple_samples(x=x,
                                        theta=hundreds)
    print(ones)
    print(hundreds)

高效实现方案 #

核心思路 #

几何adstock的计算本质是因果加权累积和：每个位置t的结果等于x[t] + θ*x[t-1] + θ²*x[t-2] + ... + θ^t*x[0]。我们可以通过构造批量权重矩阵+矩阵乘法实现完全向量化，彻底消除Python循环，同时原生兼容0维张量。

向量化实现代码 #

import torch

def adstock_geometric_vectorized(x: torch.Tensor, theta: torch.Tensor):
    # 统一x的维度，确保为二维（适配矩阵乘法）
    x = x.unsqueeze(0) if x.dim() == 1 else x
    T = x.size(1)
    
    # 构造权重序列：θ^0, θ^1, ..., θ^{T-1}，适配广播形状
    weights = theta ** torch.arange(T, device=x.device, dtype=x.dtype)
    weights = weights.view(1, T, 1)
    
    # 生成下三角掩码，确保只累积当前及之前的元素
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype)).unsqueeze(0)
    
    # 构造因果权重矩阵：每个位置(i,j)对应θ^(i-j)（i>=j时有效）
    weight_matrix = weights * mask.flip(dims=[1])
    
    # 矩阵乘法计算结果，还原输入形状
    x_decayed = torch.matmul(weight_matrix, x)
    return x_decayed.squeeze(0) if x_decayed.size(0) == 1 else x_decayed

def adstock_multiple_samples_vectorized(x: torch.Tensor, theta: torch.Tensor):
    # 统一theta形状：0维转1维，保持批量维度
    theta = theta.view(-1)
    N = theta.size(0)
    T = x.size(0)
    
    # 扩展x形状以支持批量广播
    x = x.unsqueeze(0).unsqueeze(-1)
    
    # 构造批量权重序列：(N, T, 1)
    arange = torch.arange(T, device=x.device, dtype=x.dtype)
    weights = theta.view(N, 1, 1) ** arange.view(1, T, 1)
    
    # 生成下三角掩码：(1, T, T)
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=x.dtype)).unsqueeze(0)
    
    # 构造批量权重矩阵：(N, T, T)
    weight_matrix = weights * mask.flip(dims=[1])
    
    # 批量矩阵乘法计算，整理输出形状
    x_decayed = torch.matmul(weight_matrix, x).squeeze(-1)
    return x_decayed.squeeze(0) if N == 1 else x_decayed

测试验证（含0维张量） #

if __name__ == '__main__':
    x = torch.tensor([[idx] for idx in range(112)])
    
    # 测试0维theta
    theta_0d = torch.tensor(0.5)
    result_0d = adstock_multiple_samples_vectorized(x, theta_0d)
    print("0维theta结果形状:", result_0d.shape)
    
    # 测试单元素1维theta
    theta_1d_single = torch.tensor([1])
    result_1d_single = adstock_multiple_samples_vectorized(x, theta_1d_single)
    print("单元素1维theta结果形状:", result_1d_single.shape)
    
    # 测试多元素1维theta
    theta_1d_multi = torch.tensor([idx/100 for idx in range(100)])
    result_1d_multi = adstock_multiple_samples_vectorized(x, theta_1d_multi)
    print("多元素1维theta结果形状:", result_1d_multi.shape)

方案优势 #

• 速度提升：消除Python循环，利用PyTorch的GPU加速（若可用），批量处理100个theta的速度比原实现快一个数量级以上。
• 兼容性强：自动适配0维、1维theta张量，无需额外类型判断逻辑。
• 形状友好：单theta时返回与输入x一致的形状，多theta时返回批量形状，符合直觉。

内容的提问来源于stack exchange，提问作者richardhansson