如何在PyTorch中高效实现可微分的前向填充(ffill)？

问题背景 #

需要对形状为N×L×C（批量、序列长度、通道）的张量实现类似pandas中ffill的前向填充逻辑，核心要求：

• 各通道序列相互独立，可等效为处理(N*C)×L的张量
• 保留PyTorch变量的可微分性
• 规避原方法中O(L²)级别的内存与时间开销

高效实现方案 #

以下是**O(L)时间复杂度、O(L)空间复杂度**的实现，完全适配GPU计算，且全程保留张量的可微分性：

def tensor_ffill(t: torch.Tensor) -> torch.Tensor:
    """
    对形状为[Batch, Channel, Time]的张量执行前向填充（ffill）
    逻辑等效于pandas的ffill，高效适配GPU，保留可微分性
    
    参数:
    - t: 输入张量，形状为[Batch, Channel, Time]
    
    返回:
    - 填充后的张量，形状与输入完全一致
    """
    # 重塑为[Batch*Channel, Time]，独立处理每个序列
    batch_channel_total = t.shape[0] * t.shape[1]
    t_reshaped = t.reshape(batch_channel_total, -1)
    
    # 标记非零元素位置
    non_zero_mask = t_reshaped != 0
    
    # 构造索引张量：非零位置保留原索引，零位置设为-1
    seq_indices = torch.arange(t_reshaped.shape[1], device=t.device).repeat(batch_channel_total, 1)
    seq_indices = torch.where(non_zero_mask, seq_indices, -1)
    
    # 累积取最大索引：每个零位置会被最近的左侧非零索引覆盖
    cum_max_indices = seq_indices.cummax(dim=1)[0]
    
    # 处理序列开头为零的情况：将-1索引替换为0，保留初始零值
    cum_max_indices = torch.where(cum_max_indices == -1, 0, cum_max_indices)
    
    # 根据索引取值完成填充
    filled_tensor = t_reshaped.gather(dim=1, index=cum_max_indices)
    
    # 恢复原张量形状
    return filled_tensor.reshape(t.shape)

验证示例 #

示例1：单序列测试 #

t1 = torch.tensor([[[0,1,2,0,0,3,0,4,0,0,0,5,6,0]]])
filled_t1 = tensor_ffill(t1)
print(filled_t1.squeeze().tolist())
# 输出: [0,1,2,2,2,3,3,4,4,4,4,5,6,6]

示例2：多序列测试 #

t2 = torch.tensor([[[0, 1, 0, 3]], [[1, 2, 0, 3]]])
filled_t2 = tensor_ffill(t2)
print(filled_t2.squeeze().tolist())
# 输出: [[0, 1, 1, 3], [1, 2, 2, 3]]

方案对比 #

• 原方法通过构造T×T矩阵实现，时间与空间复杂度均为O(L²)，当序列长度L较大时，GPU内存占用会急剧膨胀，计算效率极低
• 本方法仅使用O(L)的额外空间，所有操作均为向量化的GPU友好运算，计算速度提升显著，同时完全保留张量的可微分特性

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