PyTorch中实现带二进制{0,1}权重的自定义线性层（基于STE）的问题咨询

我看了你尝试在PyTorch里实现0/1二进制权重的掩码层，还想用Straight-Through Estimator（STE）解决反向传播的问题，先帮你梳理下当前代码里的核心问题，再给出修正后的实现思路~

首先，你的binarize函数的STE实现完全错了 #

你当前写的return weight_bin.detach() + weight + weight.detach()逻辑混乱，完全不符合STE的核心思想。STE的本质是：

前向传播：用二值化后的离散权重（0/1）计算输出
反向传播：跳过二值化的截断操作，让梯度直接流回原始的连续权重参数

正确的STE实现应该是这样的：

def binarize(self, weight):
    # 前向传播：生成0/1的二值权重
    weight_bin = (weight >= 0).float()
    # STE核心：前向输出二值权重，反向梯度直接流向原始weight
    return weight_bin.detach() + (weight - weight.detach())

解释下这个公式：

• 前向传播时，weight - weight.detach()的数值为0，所以最终输出是二值化后的weight_bin
• 反向传播时，weight_bin.detach()的梯度为0，梯度会直接作用在原始的weight上，相当于跳过二值化操作，用连续权重的梯度来更新参数

其次，训练时观察的权重不对 #

你在训练循环里打印的是model.custom.weight.detach().cpu().numpy()，这是原始的连续权重参数，不是实际生效的二值化掩码。如果要观察模型实际用的0/1权重，应该打印二值化后的结果：

# 替换原来的打印代码
binary_weights = model.custom.binarize(model.custom.weight).detach().cpu().numpy()
status = "epoch: {}\n custom layer binary weights: {}".format(epoch, binary_weights)
print(status)

修正后的完整CustomLinear类 #

class CustomLinear(nn.Module):
    def __init__(self, input_dim):
        super(CustomLinear, self).__init__()
        # 也可以换成torch.rand(input_dim)（0-1均匀分布），初始值更直观
        self.weight = nn.Parameter(torch.randn(input_dim))  # 正态分布初始化也可行

    def binarize(self, weight):
        # 生成0/1二值权重
        weight_bin = (weight >= 0).float()
        # 正确的STE实现
        return weight_bin.detach() + (weight - weight.detach())

    def forward(self, input):
        binary_weight = self.binarize(self.weight)
        # 逐元素相乘实现特征掩码（符合你的需求）
        output = input * binary_weight
        return output

补充几个实用细节 #

1. 模型初始化规范：你原来的CustomLinear(input_dim, input_dim)在__init__里没用到传入的参数，修正后的类直接接收input_dim，更符合PyTorch的编码习惯
2. 权重观察优化：可以同时打印原始连续权重和二值化权重，对比观察更新过程：

   # 训练循环里的打印逻辑优化
   raw_weights = model.custom.weight.detach().cpu().numpy()
   binary_weights = model.custom.binarize(model.custom.weight).detach().cpu().numpy()
   status = "epoch: {}\n raw weights: {}\n binary weights: {}".format(epoch, raw_weights, binary_weights)
   print(status)
   print("="*50)
   

3. 训练趋势说明：训练一段时间后，原始连续权重会逐渐向正/负无穷收敛，只要权重符号不变，二值化后的掩码结果就会稳定在0或1，这是STE训练离散参数的正常现象

备注：内容来源于stack exchange，提问作者Qba Liu