跳转至

locking pretrained weights via deep low-rank residual distillation

会议: NeurIPS 2026
arXiv: 2605.10777
代码: https://github.com/apple-aiml-research/ml-dlr-lock
领域: 模型压缩
关键词: 模型锁定, 低秩残差蒸馏, 自动微分不对称性, 激活内存屏障, 权值保护

一句话总结

针对开源权重模型极易被恶意微调或越狱篡改的难题,本文提出 DLR-Lock,利用自动微分在推理与训练时的显存不对称性,将预训练模型中的宽 MLP 替换为等参数量的深度低秩残差网络,在保持前向推理能力与 FLOPs 几乎不变的前提下,使反向传播激活显存呈深度线性增长、优化曲率恶化且反向计算严重减速,从而实现抗自适应攻击的模型锁定。

研究背景与动机

开源权重语言模型在推动学术研究与工业部署方面发挥了巨大作用,使用者不仅能够自由适配异构软硬件平台,还能低成本地进行领域微调与下游定制。然而,开源权重的完全开放性也带来了不可逆的安全风险:一旦权重公开发布,发布方就彻底失去了对模型的控制权,任何恶意攻击者都可以通过极小规模的微调轻易消除模型的安全对齐防护,甚至将其改造为恶意攻击工具。现有的开源协议虽然具备法律层面的使用约束,但在技术底层缺乏有效的执行与防护手段。

这种安全困境的核心技术矛盾在于:如果攻击者拥有全部模型权重与架构细节(符合密码学中的柯克霍夫原则),任何基于对称性变换的结构防御(例如符号缩放、可逆矩阵插入)都能被攻击者通过 SVD 分解或范数重平衡轻松逆转;而如果尝试通过双层优化在损失曲面上人为构建“难以优化的屏障”或病态 Hessian,防御者自身在求解该解时也将陷入优化陷阱,且攻击者拥有完全相同的梯度信息,导致基于优化防御的方案本质上是自相矛盾且易被攻破的。

本文的切入角度避开了脆弱的权重空间变换与损失曲面篡改,转而挖掘深度学习系统更底层的计算本质——自动微分在推理与训练之间的显存不对称性。前向推理只需单次流式传递并丢弃中间激活,而反向传播为了计算梯度必须在显存中保留所有前向激活值。核心 idea:将预训练模型中的每个宽浅 SwiGLU MLP 替换为等参数量但深度极大的低秩残差网络(DLR-Net),利用两阶段蒸馏精确保持输入输出映射,从而在推理几乎无感的同时,对反向传播强加不可逾越的激活显存瓶颈、架构失配优化阻碍与不成比例的反向计算延迟。

方法详解

整体框架

DLR-Lock 的核心思想是将 Transformer 结构中计算密度最高的 MLP 前馈子层进行等价但计算不对称的重构。标准的 Transformer 模块中,前向计算在经过注意力子层后通过 RMSNorm 与 SwiGLU MLP 产生残差更新。DLR-Lock 在不改动注意力子层、Embedding 与输出层的前提下,将整个归一化及 MLP 残差子块用一个由 \(L\) 层极窄瓶颈构成的深度低秩残差网络(DLR-Net)完全替换。

在模型发布前,防御者执行两阶段蒸馏流水线:第一阶段在各层独立执行模块级回归(Module-wise Regression),使初始化的 DLR-Net 快速拟合原 MLP 子块的映射函数;第二阶段冻结注意力等主干参数,在全模型尺度执行端到端 Logit 蒸馏(Logit Distillation),以消除各层微小误差在前向传播累积带来的困惑度退化。替换后的模型在推理阶段经由定制融合算子编译后保持接近基线的延迟;但在微调训练阶段,微调者(攻击者)必须反向穿透这数十层甚至上百层低秩残差块,付出成倍激增的显存开销与反向时间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["预训练 Transformer 骨干输入<br/>(冻结 Attention 与 Embedding)"] --> B["阶段 1:模块级回归<br/>独立拟合各层 MLP 残差映射"]
    B --> C["阶段 2:端到端 Logit 蒸馏<br/>Top-k KL 消除多层累积误差"]
    C --> D["锁定模型交付<br/>(DLR-Net 替换全部 SwiGLU MLP)"]
    D --> E["防御机制 1:深度激活显存线性累积<br/>反向传播强制保留每一小层激活"]
    E --> F["防御机制 2:架构失配导致优化受阻<br/>超深残差打破标准微调学习率"]
    F --> G["防御机制 3:反向计算非对称减速<br/>高显存带宽依赖使微调反向不可加速"]

关键设计

1. 深度低秩残差网络设计:以极端深度换取自动微分显存膨胀

为了在参数量不显著膨胀的前提下最大化反向激活显存,DLR-Lock 将原本宽度为 \(d_{\mathrm{ff}}\) 的单层 SwiGLU 结构重塑为深度为 \(L\)、瓶颈秩为 \(r \ll d\) 的残差链。每个残差层输入首先经过 RMSNorm,再通过下投影矩阵 \(\bm{V}_i \in \mathbb{R}^{r \times d}\) 降维至低秩空间,经 SiLU 激活后由上投影矩阵 \(\bm{U}_i \in \mathbb{R}^{d \times r}\) 升维回模型维度,并乘以可学习的 ReZero 缩放因子 \(\alpha_i\)(初始化为 0)叠加至主干: $\(\bm{h}_{0} = \bm{z}, \quad \bm{h}_{i+1} = \bm{h}_{i} + \alpha_{i} \bm{U}_{i} \sigma\bigl(\bm{V}_{i} \mathrm{RMSNorm}(\bm{h}_{i})\bigr), \quad \bm{x}' = \bm{h}_{L}\)$ 在参数预算 \(P \approx 3 d \cdot d_{\mathrm{ff}}\) 固定的约束下,网络深度 \(L = \lfloor P / (2dr + c_{\mathrm{norm}}) \rfloor\)。由于 RMSNorm 在 \(d\) 维残差流上引入了非线性操作,即使攻击者尝试冻结 DLR-Net 参数仅微调注意力层,反向传播经过 RMSNorm 时依然必须在显存中保留每个残差层的输入 \(\bm{h}_i\),使得全模型单 token 激活显存从 \(\Theta(N d_{\mathrm{ff}})\) 剧增至 \(\Theta(N d L)\),产生高达 \(\Theta(d / r)\) 倍的显存放大因子。

2. 模块级回归与全局 Logit 蒸馏协同训练:解耦初始拟合与端到端精度修复

若直接端到端从随机初始化训练超深残差 MLP,极易陷入梯度弥散与严重的优化停滞。DLR-Lock 采用两阶段优化策略解耦这一难度。第一阶段在冻结主模型的条件下,收集各层注意力输出作为输入特征 \(\bm{z}_l\),令每层 DLR-Net \(g^{(l)}\) 独立拟合原始残差块输出 \(f_l(\bm{z})\),优化相对均方误差损失: $\(\mathcal{L}_{\mathrm{MSE}} = \frac{1}{n} \sum_{i=1}^{n} \frac{\|g^{(l)}(\bm{z}_i) - f_l(\bm{z}_i)\|^2}{\|f_l(\bm{z}_i)\|^2}\)$ 模块级回归不仅具备完全并行的训练效率,而且由于深度网络在表征能力上天然宽容浅层网络,能够快速逼近各层输入输出分布。第二阶段将所有 DLR-Net 组装回主干,冻结注意力与 Embedding 参数,仅微调 DLR-Net 权重,以预训练教师模型的顶层分布为基准,在 top-\(k\)(如 \(k=1000\))词表切片上最小化 KL 散度 \(\mathcal{L}_{\mathrm{KD}}\)。该阶段虽然需要承担深层反向传播,但防御者作为一次性代价平摊到所有下游用户中,有效消除了跨层误差复合效应。

3. 推理前向算子融合与反向计算不对称:构建单向推理加速通道

DLR-Net 引入的大量超薄矩阵乘法在原生 PyTorch 执行时会产生大量小算子发射开销与高带宽显存(HBM)频繁换入换出。为了让防御对合法推理用户完全透明,本文设计了单向算子融合策略:在 Hopper 架构上利用 TMA(张量显存加速器)与 WGMMA 指令,将残差流 \(\bm{h}_i\) 持续驻留在片上 SRAM 中,避免每层写回显存,并结合 CUDA Graph 消除主机端发射延迟。这一优化能够将推理前向延迟降低数倍,但该加速路径对微调攻击者完全失效——因为反向传播必须为外积梯度与链式法则回传保存全量的激活张量,其庞大的显存读写需求直接打碎了片上缓存流水线,使得微调训练的反向耗时无法获得与前向同等幅度的优化,将训练相对推理的时间膨胀比拉大到基线模型的数倍以上。

损失函数 / 训练策略

  • 阶段 1 超参数:AdamW 优化器,学习率 \(10^{-3}\),线性预热 5% 配合余弦衰减,权重衰减 \(10^{-5}\),梯度裁剪范数 1.0;每层独立迭代 500k 步,batch size 为 2,序列长度 2048。
  • 阶段 2 超参数:冻结除 DLR-Net 外所有参数,使用原模型作为 Teacher 进行 top-1000 KL 散度蒸馏;迭代 100k 步,学习率 \(10^{-3}\),batch size 为 4,序列长度 2048。
  • 架构设定:Qwen3-0.6B 采用 \(r=32, L=141\);1.7B 采用 \(r=64, L=142\);8B 采用 \(r=128, L=143\);32B 采用 \(r=256, L=149\)。深度保持在 140–150 层区间,将参数容量由秩 \(r\) 自适应吸收。

实验关键数据

主实验

实验基于 Qwen3-0.6B 模型,在 Nemotron-CC 验证集与 WikiText-103 上评测语言建模困惑度(PPL),并在 7 项常用评测基准(MMLU、ARC-Easy、ARC-Challenge、HellaSwag、WinoGrande、BoolQ、PIQA)上测试通用能力保持情况。

模型配置 Nemo PPL (↓) WT-103 PPL (↓) MMLU (%) (↑) ARC-E (%) (↑) ARC-C (%) (↑) HSw (%) (↑) WiGr (%) (↑) BoolQ (%) (↑) PIQA (%) (↑)
Pretrained (基准) 14.9 21.5 40.3 55.9 33.7 47.3 56.4 63.8 67.3
阶段1:仅模块回归 21.8 182.0 31.1 50.6 26.5 40.5 54.9 62.6 63.4
阶段2:完整 DLR-Lock 14.9 23.4 36.8 50.7 29.0 43.5 55.2 63.8 65.3
DLR-Lock (扩展蒸馏 300k) 14.8 21.5 39.8 54.2 32.5 46.5 56.1 63.8 66.8

消融实验

在 WikiText-103 上进行微调攻击评估,对比基线模型与锁定模型在开启梯度检查点(Gradient Checkpointing)下的训练收敛、FLOPs 开销与墙上时间。同时评估推理端通过专用融合 Kernel 带来的加速效果。

配置方案 微调最低 Loss (↓) 达标所需计算量 (FLOPs/Token) 达标相对代价 (×) 推理单步加速比 (\(t_f'/t_f\)) 训练单步加速比 (\(t_b'/t_b\))
基线模型 (SwiGLU) 2.337 6 FLOPs 1.0× (基准) 1.00 1.00
DLR-Lock (学习率 \(10^{-5}\)) 2.510 (收敛变慢) 8 FLOPs (需重计算) 6.4× (理论下界) 0.02 (0.6B, \(n=4\)) 0.09 (0.6B, \(n=4\))
DLR-Lock (学习率 \(3\times 10^{-5}\)) 2.499 (陷入平台) 8 FLOPs (需重计算) 6.4× (理论下界) 0.07 (0.6B, \(n=8192\)) 0.39 (0.6B, \(n=8192\))
DLR-Lock (学习率 \(10^{-4}\)) 梯度发散 (NaN/Inf) 不可收敛 \(\infty\) - -
攻击者 LoRA (\(\rho=16\) Attention) 2.638 (表现极差) 8 FLOPs (强制遍历超深) 显存爆显存/代价极高 - -

关键发现

  • 微调优化空间急剧收窄与发散:基线模型在 \(10^{-4}\) 学习率下能够正常快速收敛,而 DLR-Lock 模型在该学习率下直接剧烈震荡并完全发散;在最佳学习率下,锁定模型的收敛下限依然显著落后于基线模型,且达到相同性能所需的最少计算量下限为基线的 6.4 倍(在原生 Wall-clock 下更是达到了 109 倍的 GPU-hour 开销差距)。
  • 常见绕过攻击全部失效:攻击者尝试使用 LoRA 仅微调 Attention 时,即便完全冻结 DLR-Net,反向传播依然必须经过深度残差的 RMSNorm,单卡 80GB H100 在未开启梯度检查点时直接发生 OOM;开启检查点后由于超深重计算,微调速度依旧极其迟缓且性能无法达到未锁定水平。
  • 逆向蒸馏(Reverse Distillation)防线牢固:攻击者若试图先通过逆向蒸馏将 DLR-Net 还原为 SwiGLU MLP,其所需的最优算力预算高达 \(7.5 \times 10^{18}\) FLOPs,甚至超过了防御者初始锁定模型所需的 \(6.3 \times 10^{18}\) FLOPs(相当于单次微调任务开销的 50 倍以上),且由于缺乏防御者的专有预训练语料,逆向蒸馏的质量严重受限。

亮点与洞察

  • 防御视角降维打击:跳出了传统在权重参数空间做置换加密或在损失函数设计病态曲面的脆弱博弈,首次从系统级自动微分在“前向推理抛弃激活”与“反向传播必须驻留激活”的本质物理不对称性切入,建立了一套不可通过优化手段轻易逆转的结构性防御。
  • 柯克霍夫安全原则与理论可证明性:在攻击者完全知晓 DLR-Lock 架构与锁定机制的白盒设定下,从理论上推导出了微调显存膨胀下界 \(\Omega(d/r)\),使得无论攻击者如何调整优化器类型或超参数,都无法回避硬件底层显存与带宽的硬性开销。
  • 发布者单次沉没成本兑现全生命周期防护:利用两阶段蒸馏使得模型发布者在离线阶段承担一次性蒸馏成本,为下游合法用户提供几乎无感的推理体验,而将持续的计算与显存负担转嫁给潜在的恶意微调者。

局限与展望

  • 推理端单请求与短序列场景的绝对延迟仍然存在微小开销:在极短序列与 batch size 为 1 的场景下,由于 Attention 稀释效应较弱,即便使用了算子融合,由超深层带来的顺序执行延迟相比宽 MLP 仍存在一定残余开销(约 1.2–2×),需要更激进的片上跨层融合指令支持。
  • 锁定模型无法抵御完全从零预训练或从 Logit 完整蒸馏全新模型的高算力对手:当攻击者具备接近甚至超过预训练算力预算时,直接通过生成海量数据训练全新模型可以绕过一切权重锁定,但这属于所有开放模型防御在理论上的绝对上界。
  • 未来方向:探索跨 Transformer 层的联合低秩残差路由,以及将 DLR-Lock 与安全表征加噪(Representation Noising)相结合,形成架构硬件屏障与潜在表征屏障的双重深度防御体系。

相关工作与启发

  • vs SpecDef (Rosati et al., 2025): SpecDef 通过在预训练权重上最大化奇异值构建病态优化,但防御者优化病态 Hessian 的过程极其脆弱且会损害多任务通用泛化性;DLR-Lock 采用结构性替换与两阶段函数映射蒸馏,理论上保证了全任务能力的无损迁移。
  • vs Non-Fine-Tunable Foundation Models (Wang et al., 2026b): 后者依赖预训练期间嵌入二值抽签掩码(Lottery Ticket Mask),极大压缩了模型前向有效容量且容易被范数统计识别逆向;DLR-Lock 在后处理阶段执行且不损失模型参数容量。
  • vs 传统轻量化模块蒸馏 (Cheng et al., 2026; Liang et al., 2023b): 传统模块替换的目标是压窄网络降低推理 FLOPs;DLR-Lock 反其道而行之,在保持 FLOPs 相当的前提下,故意设计极深、极窄的低秩网络,将计算转化为主存带宽与显存占用的瓶颈,为安全防御提供了全新架构视角。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从自动微分的软硬件显存不对称性设计模型锁定,思路极具启发性与突破性。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 0.6B 到 32B 多个尺度,全面测试了通用基准、反向微调阻力、LoRA/逆向蒸馏等自适应攻击以及 H100 专用内核测试。
  • 写作质量: ⭐⭐⭐⭐⭐ 形式化定义严谨,柯克霍夫原则贯彻清晰,理论推导与工程实现剖析透彻。
  • 价值: ⭐⭐⭐⭐⭐ 为开源大模型在开放权重时代的合规发布与防止恶意滥用提供了高可行性的技术抓手。