让 PyTorch 的 sign、round、floor 自动可微：SLL-Core 一行代码搞定让 PyTorch 的

让 PyTorch 的 sign、round、floor 自动可微：SLL-Core 一行代码搞定

训练量化网络时，torch.round() 断了梯度？做硬阈值激活时，torch.sign() 无法 backward？别再手写 STE 了。

先来看一个让人崩溃的日常

做深度学习的你，一定遇到过这种情况：

import torch

x = torch.tensor([0.5, -0.3, 1.2], requires_grad=True)
y = torch.sign(x)          # 符号函数，很常用
loss = y.sum()
loss.backward()
print(x.grad)              # tensor([0., 0., 0.])

梯度全是 0。 参数根本更新不了。

不只是 sign，round、floor、ceil、argmax 这些离散算子都有同样的问题：它们几乎处处梯度为 0，标准反向传播直接失效。

而这类操作在深度学习里其实无处不在：

量化感知训练 (QAT)：round(x) 把权重映射到离散档位
硬阈值激活：(x > 0).float() 做二值化特征
符号网络：sign(x) 做梯度友好的二值近似
离散决策：argmax 选最优动作或类别

传统方案：要么改模型，要么梯度不准

面对这个问题，已经有几种 workaround，但都有明显代价：

方案	是否需要改代码	部署残留	梯度质量	调参难度
硬函数直接训练	✅ 无需改动	✅ 无	❌ 零梯度，训不动	❌ 无解
Sigmoid / Softmax 松弛	❌ 重写模型	❌ 有近似误差	⚠️ 梯度消失/爆炸	⚠️ 要调温度
Straight-Through Estimator (STE)	❌ 手写自定义梯度	✅ 无	⚠️ 方向常不准	⚠️ 容易震荡
Gumbel-Softmax	❌ 改模型结构	❌ 温度参数残留	⚠️ 高方差	⚠️ 慢

也就是：要么你改模型结构去适配一个"伪可微"的近似，要么你接受粗糙的梯度估计。

SLL：只在边界附近动手脚

SLL（Static Local Linearization，静态局部线性化）的思路非常直接：

不需要在整个定义域上做近似，只在决策边界附近的 ε-小区间内做局部线性化，其余区域严格保持原始硬逻辑。

以 sign 为例：

当 |x| > ε 时：SLL 就是硬 sign，输出 -1 或 1，梯度为 0
当 |x| ≤ ε 时：在边界附近做一条斜线过渡，梯度为常数 1/(2ε)

这样带来的好处是：

前向传播：远离边界时输出和硬函数完全一致，没有近似误差
反向传播：边界附近梯度是常数，不会出现 Sigmoid 那种梯度消失
部署阶段：训练完直接去掉 SLL，原始硬逻辑零开销恢复

代码：真正的一行解决

SLL-Core 把这个思路做成了一个零侵入的上下文管理器：

pip install sll-core

import torch
import sll

x = torch.tensor([-1.0, 0.0, 1.0], requires_grad=True)

with sll.linearize(eps=1e-2):     # ← 只加这一行
    y = torch.sign(x)              # 自动可微！
    z = torch.round(y * 10)
    loss = z.sum()
    loss.backward()

print(x.grad)                      # 梯度正常回传 ✅

离开 with 块后，torch.sign 自动恢复原始硬逻辑。你的模型代码完全不用改。

三种使用方式

1. 上下文管理器（推荐，零侵入）

with sll.linearize(eps=1e-2):
    y = model(x)       # 模型里的 sign/round/floor 自动可微
    loss = criterion(y, target)
    loss.backward()

2. 装饰器（函数级）

@sll.enable(eps=1e-2)
def quantized_model(x):
    return torch.sign(torch.round(x))

3. 显式调用（不 patch 全局）

y = sll.heaviside(x, eps=1e-2)
z = sll.argmax(x, dim=1, eps=1e-2)

让我们看效果：

数学原理：ε-局部线性化

以 Heaviside 阶跃函数为例，SLL 的数学形式非常简洁：

y(x) = \begin{cases} 0.5 + \frac{x}{2\epsilon} & \text{当 } |x| \leq \epsilon \\ H(x) & \text{其他} \end{cases}

其中 $H(x)$ 是原始 Heaviside 函数。当 $\epsilon \to 0$ 时， $y(x) \to H(x)$ ，最优解收敛到原始离散问题的最优解。

关键点：线性化只发生在边界附近的 ε-区间内，其余区域严格等于硬函数。这和 STE 那种"全局偷换梯度"的思路有本质区别。

实际应用场景

场景 1：量化感知训练 (QAT)

def quantize(x, levels=256):
    scale = (levels - 1) / (x.max() - x.min() + 1e-10)
    return torch.round((x - x.min()) * scale) / scale + x.min()

x = torch.randn(10, requires_grad=True)

with sll.linearize(eps=1e-3):
    y = quantize(x)                 # round 可微了
    loss = y.sum()
    loss.backward()

print("量化梯度:", x.grad)          # ✅ 正常回传

场景 2：带硬阈值激活的网络

class DiscreteModel(nn.Module):
    def forward(self, x):
        x = self.linear(x)
        return (x > 0).float()       # 硬阈值，原本不可微

# 训练时套一层 SLL，模型代码完全不用改
with sll.linearize(eps=1e-2):
    y = model(x)
    loss = criterion(y, target)
    loss.backward()

参数怎么选？

eps（默认 1e-3）：线性化区间半宽
- 越小：越接近硬函数，梯度区域越窄，适合对精度要求高的场景
- 越大：过渡越平滑，优化更稳定，适合训练初期
- 建议：从 1e-2 开始，根据收敛情况微调

注意事项

优先用 torch.sign(x) 而非 x.sign()：后者 SLL 会尽力拦截，但前者更可靠
比较运算符无法拦截：x > 0 不能被 patch，建议改用 sll.threshold(x, threshold=0.5)
部署零开销：训练完直接部署原始模型，不需要 SLL

总结

如果你正在做量化网络、符号网络、硬阈值激活，或者任何包含离散决策的 PyTorch 项目，SLL-Core 提供了一种不需要重写模型的可微分化方案。

核心就一句话：with sll.linearize():，然后你的离散模型就能 backward 了。

GitHub: github.com/jacksong-so…
PyPI: pip install sll-core
License: MIT

如果你在某个场景下用 SLL-Core 解决了梯度问题，欢迎去 GitHub 提 Issue 分享！