PyTorch中如何创建具有特定模式的可学习张量？

这类带约束模式的可学习参数不需要把完整张量的所有元素都设为独立参数，只需要注册真正独立的可学习变量，前向传播时动态生成符合模式的完整张量即可，全程自动支持反向传播，也能严格保证张量模式符合要求。

核心实现思路 #

• 你给出的A属于对称Toeplitz矩阵的变种，独立可学习参数只有3个：a0、a1、a2，没有冗余参数，也不会出现不符合模式的参数值。
• 用PyTorch原生的nn.Parameter注册独立参数，再通过张量操作生成完整A矩阵，所有构造操作都在计算图内，梯度会自动回传到独立参数上，不需要手动写反向传播逻辑。

完整代码示例 #

import torch
import torch.nn as nn

class PatternedLinearTransform(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # Q用标准nn.Linear实现，对应Y = Q(X+A)的线性变换部分
        self.linear_q = nn.Linear(input_dim, output_dim)
        # 注册A的3个独立可学习参数，可自定义初始化策略
        self.a0 = nn.Parameter(torch.randn(1))
        self.a1 = nn.Parameter(torch.randn(1))
        self.a2 = nn.Parameter(torch.randn(1))
        self.input_dim = input_dim

    def forward(self, x):
        # 第一步：生成符合模式的A矩阵，和输入x在同一个设备上
        A = torch.full((self.input_dim, self.input_dim), fill_value=self.a2, device=x.device)
        # 填充主对角线为a0
        A.diagonal().fill_(self.a0)
        # 填充上下次对角线为a1
        A.diagonal(offset=1).fill_(self.a1)
        A.diagonal(offset=-1).fill_(self.a1)
        
        # 第二步：执行线性变换Y=Q(X+A)，适配batch维度
        # 假设输入x形状为 [batch_size, input_dim]，给A扩展batch维度后相加
        x_add_a = x + A.unsqueeze(0)
        y = self.linear_q(x_add_a)
        return y

扩展优化说明 #

• 如果输入维度很大，可以不显式生成完整A矩阵，直接通过广播+索引修改计算X+A的结果，节省显存：先把x整体加a2，再对每个样本的主对角元素加a0 - a2、次对角元素加a1 - a2，计算结果完全一致。
• 如果A的模式更复杂，比如有更多对角带的参数，只需要调整独立参数的注册逻辑和张量填充规则即可，训练流程不需要改动。

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