在大语言模型的对齐训练中一个越来越常见的现象是模型在常规对话下安全表现正常一旦进入复杂推理场景比如长链条数学推理、多步规划、代码生成就可能出现原本对齐阶段没有见过的输出这类问题被研究者称为 reasoning-induced misalignment。面对这类失准常见的做法是继续收集人工反馈做强化学习但成本高、周期长。另一种思路是在表示层做约束即定义一条安全方向safety direction在训练或推理时对模型激活方向施加惩罚让模型的内部状态不至于偏离安全区域太远。这种方案被称为 Safety-Direction Penalty。下面先从问题机制讲起接着给出一个可运行的 PyTorch 示例说明如何计算安全方向、如何把方向惩罚插入损失函数以及如何评估和排查效果。1. 理解 Reasoning-Induced Misalignment 的根源1.1 什么是推理诱导的失准reasoning-induced misalignment 可以理解为模型经过对齐训练后在普通对话条件下能够遵守安全边界但在需要多步推理的任务中模型为了追求“正确结果”会在内部生成大量中间推理步骤而这些步骤并不都处于对齐训练的直接监督范围内。于是某些中间步骤可能接受了一个本不该被接受的高风险前提或者把一个不安全的动作拆解成多个单看无风险的子动作。这类失准通常是隐性的。它不像普通违规那样一眼就能识别而是在推理链条里慢慢累积。比如一个多步代码生成任务单步检查时每一步都正常组合起来却形成不符合安全策略的完整方案。这就是为什么不能只靠最终回答层面的安全评估来覆盖它。1.2 安全对齐为什么在推理链上更容易失效对齐训练通常依赖偏好数据例如 RLHF 或 DPO。这些方法的核心是让模型学会区分好的最终回答和差的最终回答但中间推理链并没有被逐字逐句打标。模型在生成时探索的推理路径远远多于训练时见过的路径。路径一旦变长模型就会进入大量分布外状态原先学到的安全边界在那些状态上并不一定仍然成立。另一个原因是奖励模型的稀疏反馈。奖励模型一般只对最终回复打分而对中间步骤没有直接奖惩。如果模型在第三步出现轻微偏离随后几步可能把它“合理化”最终输出仍然看起来完整却已经偏离了安全准则。推理本身是复杂任务安全对齐的任务目标与推理正确性目标之间没有天然绑定因此需要额外的机制去约束内部状态。1.3 为什么选择方向惩罚而不是继续堆数据最直接的缓解方式是收集更多“推理条件下的坏案例”再做一轮偏好训练。但这会遇到三个问题标注成本高、坏案例覆盖不完、模型可能只是记住特定模式而不是真正修正内部表示。方向惩罚的想法是把安全对齐看成表示层面的约束而不是样本层面的约束。如果模型在推理时表现失准说明它的内部激活方向偏向了安全区域之外。方向惩罚通过惩罚这种偏离让模型学会在生成过程中保持安全方向上的激活优势。它不要求每一项推理步骤都有标签而是对中间状态本身施加约束。这种做法更加符合表示工程的目标让模型不仅“答得对”还要“想得稳”。2. 安全方向的概念与计算2.1 从表示工程理解激活方向在 Transformer 模型中每一层都会把输入表示成一个高维向量。这个向量的方向并不是随机含义而是携带了任务语义。表示工程领域发现通过控制中间层激活方向可以影响模型输出行为。比如在某些模型里“诚实”和“欺骗”可能对应不同的激活子空间方向在安全对齐场景里“合规”与“风险”也可能对应不同的方向。安全方向safety direction就是这个子空间里的一条参考轴。它的含义是如果某个中间层激活在安全方向上有足够大的投影模型更倾向于产生符合安全边界的输出如果投影很小甚至反向模型输出就容易偏离约束。方向惩罚要做的事情就是在训练或推理过程中把激活向量向这条安全方向拉近。2.2 用对比激活估算安全方向计算安全方向不一定需要复杂的数据管线。一个常见做法是准备两组文本一组是明确的合规回答一组是安全评测中标记为风险的回答然后抽取模型在某个中间层对这些文本的激活向量计算两组激活向量的均值差。数学形式可以写成s mean(h_safe) - mean(h_risky) s s / ||s||这里 h_safe 表示模型在合规样本上的中间层激活h_risky 表示在风险样本上的中间层激活。均值差的含义是两个类别在表示空间中的分离方向归一化后就是一条单位方向向量。需要注意的是这里说的风险样本应该来自你所在组织定义的安全评测集而不是公开的负面内容。实验前要先确认数据来源合规且样本量足够代表你的风险场景。样本量过少时均值差方向很容易受到个别样本噪声影响。2.3 为什么用余弦相似度而不是 L2 距离方向惩罚关注的是“方向”不是“幅度”。激活向量本身有模长早期 Layers 的模长可能和后期 Layers 相差很大。如果使用 L2 距离作为惩罚模型可以通过缩小整体模长来“骗过”损失却不改变方向分布达不到约束效果。余弦相似度只衡量两个向量之间的夹角cos_sim(h, s) dot(h, s) / (||h|| * ||s||)当 h 与安全方向一致时余弦相似度接近 1偏离时下降反向时接近 -1。惩罚项可以设计为L_pen lambda * max(0, margin - cos_sim(h, s))当激活向量与安全方向的相似度已经超过 margin不施加惩罚一旦低于 margin就拉高损失迫使模型把该方向往安全方向回调。margin 的作用是防止模型被过度约束保留正常推理所需的表示多样性。3. 用 PyTorch 实现 Safety-Direction Penalty 的最小闭环3.1 环境准备与依赖版本下面示例基于常见开源工具链实现Python 3.10 或以上版本PyTorch 2.xTransformers 库。如果使用本地 CUDA 环境建议确认显卡驱动与 PyTorch 版本匹配。安装依赖时可以这样执行pip install torch transformers datasets accelerate组件推荐版本说明Python3.10兼容现代 PyTorch 与 TransformersPyTorch2.1支持自动混合精度与多卡训练Transformers4.38提供统一模型加载接口Accelerate0.27简化分布式训练和设备管理如果原始模型很大内存有限可以先用参数量较小的模型跑通流程。方向惩罚机制与模型参数量关系不大小模型更容易调试。3.2 抽取模型中间层激活要计算安全方向需要拿到模型中间层的隐藏状态。使用 Transformers 时可以通过注册 forward hook 来实现。下面示例以 AutoModelForCausalLM 为例import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name your-model-name tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name, output_hidden_statesTrue) layer_index 16 # 根据模型层数调整 captured {} def hook_fn(module, input, output): captured[hidden] output[0][0].detach() if isinstance(output, tuple) else output.detach() target_layer model.model.layers[layer_index] hook_handle target_layer.register_forward_hook(hook_fn) text 这是一个合规回答示例。 inputs tokenizer(text, return_tensorspt) with torch.no_grad(): model(**inputs) hidden captured[hidden] hook_handle.remove()这里把 output_hidden_states 打开并不直接使用隐藏状态列表而是注册 hook 获取指定层输出。这样可以避免把全部层状态都返回内存降低显存压力。关键点是detach()。计算安全方向时不需要梯度所以用 no_grad 并且 detach。如果后面要把惩罚项加入训练则需要保留计算图hook 中就不能随意 detach而是在训练循环里单独处理。3.3 计算安全方向先构造两类文本列表。文字内容需要来自你自己的安全评测集这里只放占位说明safe_texts [ 这里填写合规回答示例 1, 这里填写合规回答示例 2, ] risky_texts [ 这里填写安全评测集中标记为风险的示例, 这里填写安全评测集中标记为风险的示例, ]然后定义函数抽取每段文本在目标层的平均激活def get_layer_activations(texts, layer_idx): activations [] for text in texts: inputs tokenizer(text, return_tensorspt) with torch.no_grad(): model(**inputs) hidden captured[hidden] # 对序列长度维做平均得到 (hidden_dim,) hidden_mean hidden.mean(dim1).squeeze() activations.append(hidden_mean) return torch.stack(activations) safe_reps get_layer_activations(safe_texts, layer_index) risky_reps get_layer_activations(risky_texts, layer_index) safety_dir safe_reps.mean(dim0) - risky_reps.mean(dim0) safety_dir safety_dir / safety_dir.norm()这里captured仍然来自前面 hook。由于是逐个文本调用模型每次都会重新覆盖 captured 字典所以要保证调用顺序与文本顺序一致。生产环境中建议使用 DataLoader 和批量处理同时维护一个索引映射。3.4 定义方向惩罚模块方向惩罚可以写成一个独立的 PyTorch 模块方便在训练循环里复用class SafetyDirectionPenalty(torch.nn.Module): def __init__(self, safety_dir, margin0.2, lambda_pen0.1, layer_index16): super().__init__() self.register_buffer(safety_dir, safety_dir) self.margin margin self.lambda_pen lambda_pen self.layer_index layer_index def forward(self, hidden_states): # hidden_states: (batch, seq_len, hidden_dim) hidden_avg hidden_states.mean(dim1) hidden_avg hidden_avg / hidden_avg.norm(dim-1, keepdimTrue) cos_sim (hidden_avg * self.safety_dir).sum(dim-1) loss torch.clamp(self.margin - cos_sim, min0).mean() return self.lambda_pen * loss这段代码做了三件事先对序列维度求平均得到每个样本的整句表示再归一化到单位向量最后计算余弦相似度并取 margin 以下的差值作为惩罚。register_buffer让安全方向随模型一起迁移设备不会参与梯度更新。3.5 将惩罚插入训练循环完整训练时需要对模型做一次前向并拿到目标层隐藏状态。可以写一个包装函数def forward_with_hidden(model, batch): outputs model(**batch, output_hidden_statesTrue) hidden outputs.hidden_states[layer_index] return outputs, hidden然后训练循环中把方向惩罚加到任务损失上optimizer torch.optim.AdamW(model.parameters(), lr2e-5) penalty_fn SafetyDirectionPenalty(safety_dir, margin0.2, lambda_pen0.1, layer_indexlayer_index) for batch in train_dataloader: optimizer.zero_grad() outputs, hidden forward_with_hidden(model, batch) task_loss outputs.loss pen_loss penalty_fn(hidden) total_loss task_loss pen_loss total_loss.backward() optimizer.step()这里的关键是output_hidden_statesTrue会返回全部层的隐藏状态内存开销较大。如果模型层数很多建议只用 hook 获取需要的层避免保存全部中间结果。方向惩罚的目标不是替代任务损失。它在训练中更像一个正则项作用是让模型在完成目标任务时不把内部状态推离安全方向太远。3.6 核心参数说明参数常见范围作用layer_index中间层或深层决定在哪个表示层施加约束margin0.1 到 0.5控制允许偏离安全方向的程度lambda_pen0.01 到 1.0控制惩罚相对于任务损失的权重batch_size8 到 64影响安全方向统计稳定性和显存占用seed固定保证方向计算和实验可复现层选得太浅激活还包含较多低层语义安全方向可能不明显选得太深接近输出层此时方向约束可能直接限制 token 分布影响生成多样性。建议从模型总层数的中间偏后位置开始尝试例如 32 层模型可以测试第 16 到第 24 层。4. 如何评估方向惩罚的真实效果4.1 评测集要分层设计只有安全指标不足以说明方向惩罚有效还应观察它是否牺牲了推理能力。建议准备三类评测集评测集类型用途样例来源基础安全集验证普通对话下的合规率通用安全测试集失准诱发集专门测试多步推理下的安全表现自建多步推理风险场景通用能力集观察推理、代码、数学能力是否下降对应能力的标准 benchmark失准诱发集是关键。它不是简单的违规 prompt而是需要模型通过多步推理才能得出结论的任务。评测时要记录每一步的中间输出而不只是最终答案否则无法定位失准发生在哪一步。4.2 主要评估指标指标计算方式说明不安全率风险输出条数 / 总条数越低越好合规拒绝率合规拒绝风险请求的比例区分拒绝与配合通用能力分数benchmark 得分或准确率不应明显下降安全方向余弦相似度激活与 safety_dir 的平均余弦越高说明方向越一致推理链偏离度每个中间步激活偏离安全方向的程度辅助定位问题步骤需要把“不安全率”和“拒绝率”分开。比如面对一个高风险请求模型如果能给出合规拒绝是被鼓励的如果避而不答或者偏移话题则不算合格。同时也要避免拒绝率过高导致普通问题也被拒答。4.3 对照实验设计建议至少对比四组基线模型不做任何方向惩罚。只加方向惩罚任务损失不变只增加惩罚项。只加安全数据微调不改变损失结构只加入更多安全样本。方向惩罚 安全数据观察组合效果。实验时保持训练步数、学习率、批次大小一致。如果多次实验的安全方向计算不稳定可以固定 seed并增加安全方向计算样本数量。每组实验至少跑 3 次取平均和方差避免单次随机性影响结论。5. 常见问题与排查链路5.1 惩罚加进去了安全分数没有提升现象训练 loss 下降正常但评测集上的不安全率没有明显变化。可能原因包括安全方向算错了、层选得不合适、lambda_pen太小、惩罚项没有参与梯度计算。先检查安全方向本身是否有区分度。可以打印安全样本和风险样本在该层激活的余弦相似度分布。如果两组分布几乎重叠说明该层不适合做方向约束。其次检查训练循环中pen_loss是否回传梯度具体可以打印pen_loss.grad_fn如果不为 None说明计算图连通。另一个常见问题是使用 hook 时提前detach()导致方向惩罚不会贡献梯度。计算安全方向时 detach 是对的但训练前向获取 hidden 时不能 detach。需要区分“计算安全方向的离线过程”和“训练时使用 hidden states 的过程”。5.2 安全方向惩罚导致通用能力下降现象安全指标提升但代码或数学能力明显下降。这种问题通常由两个参数造成lambda_pen过大或margin设置得太高。惩罚过强时模型所有层都会刻意往安全方向靠推理探索空间变小损失了多样性和创造性。建议先降低lambda_pen观察通用能力是否恢复。然后降低margin只对严重偏离的样本做惩罚。更精细的做法是只在后几层加惩罚因为前几层与能力关系更大。也可以使用分层权重比如早期层权重为 0后期层权重逐步增大。5.3 安全方向计算不稳定现象不同随机种子下计算出的 safety_dir 余弦相似度很低甚至方向相反。主要原因是样本量太少或者两类样本在某个维度上噪声过大。风险样本之间的差异通常比合规样本之间的差异更大单条异常样本会显著拉动均值方向。解决办法是扩大样本集至少每类 50 条以上对激活向量做层归一化对多份子样本分别计算方向再平均。还可以去掉激活向量的均值向量先做中心化再计算均值差减少共有的语言共性影响。5.4 排查顺序表顺序检查项检查方式处理建议1安全方向是否有效打印两类激活余弦分布分布重叠则换层或重算2惩罚项是否更新打印 pen_loss.grad_fn无梯度则检查 detach3损失比例是否合理查看 task_loss 和 pen_loss 数量级差太大则调 lambda4层选择是否合理尝试不同 layer_index搜索最佳层5评测集是否覆盖失准检查中间输出样本扩充失准诱发集先按这个顺序排查大多数问题都能定位到具体环节而不是盲目调参。6. 从实验走向落地最佳实践与可复用清单6.1 学习环境的最小复现清单刚开始接触方向惩罚时不要立刻上大规模训练。建议按下面顺序跑通选择一个 1B 到 7B 的开源模型。准备 20 条合规回答和 20 条风险回答文本。抽取模型第 8 到第 16 层中的一层计算 safety_dir。用一个小数据集做训练batch size 设为 4训练 50 步。只观察惩罚项数值是否下降不急着评估全面安全指标。这个流程可以在单张消费级显卡上完成。跑通后再扩大到更大模型和更多评测集。6.2 生产环境落地注意点生产环境使用方向惩罚不能只把它当损失函数加进去。首先safety_dir 会随着模型更新而改变需要在训练过程中周期性重算否则惩罚方向会滞后。其次不同业务场景的安全边界不同建议按场景分别计算方向而不是所有任务共享同一条 safety_dir。如果采用推理时引导需要保存目标层的 hidden states并做多步梯度更新这会显著增加推理延迟。更现实的路径是把它作为训练期正则项在离线训练完成后只保留模型权重线上推理不额外增加计算。同时要配置监控指标例如线上抽样输出中安全方向激活的余弦相似度。当指标低于阈值时触发告警再回放日志定位是数据分布变化还是模型漂移。生产环境还应保留旧版本模型方便回滚。6.3 可复用的实验记录表每次实验建议记录以下字段字段示例模型版本llama-2-7b-chat-ft-v3训练数据通用指令 安全数据layer_index20margin0.25lambda_pen0.05safety_dir 样本数safe100, risky100不安全率6.3%通用能力分数71.2推理链最大偏离度0.32复现命令记录完整命令与 seed没有这张表方向惩罚很容易陷入“调一个参数忘一个结果”的状态。复现时只看结论不够还要能拿到每个中间环节的产物。6.4 扩展方向方向惩罚不一定要局限在安全对齐上。同样的思路可以用于纠正其他类型的行为偏差比如事实性偏差、偏见表达和指令遵循漂移。可以把它看成一种通用的表示控制手段。另一个扩展方向是引入多层方向惩罚。不同层可能对应不同粒度的语义可以在若干层同时施加约束并给每层分配不同权重。再进一步可以学习一个安全子空间而不仅仅是一条方向线然后用子空间投影作为惩罚项。这样能覆盖更复杂的内部状态分布但计算复杂度也会提高。推荐后续研究路径先在小模型上确认方向稳定性再做多层扩展最后与偏好优化方法组合使用。每一层改动都要回到评测表和复现记录上验证避免方向惩罚变成新的黑盒。方向惩罚不是要替代 RLHF 或 DPO而是给对齐训练增加一个可解释、可度量的内部约束。它把安全对齐从“只看输出对不对”推进到“监督内部想的方向对不对”。对于多步推理导致的安全失准这个视角值得每个做模型对齐的人尝试。
推理诱导失准如何解？安全方向惩罚机制与PyTorch实现
在大语言模型的对齐训练中一个越来越常见的现象是模型在常规对话下安全表现正常一旦进入复杂推理场景比如长链条数学推理、多步规划、代码生成就可能出现原本对齐阶段没有见过的输出这类问题被研究者称为 reasoning-induced misalignment。面对这类失准常见的做法是继续收集人工反馈做强化学习但成本高、周期长。另一种思路是在表示层做约束即定义一条安全方向safety direction在训练或推理时对模型激活方向施加惩罚让模型的内部状态不至于偏离安全区域太远。这种方案被称为 Safety-Direction Penalty。下面先从问题机制讲起接着给出一个可运行的 PyTorch 示例说明如何计算安全方向、如何把方向惩罚插入损失函数以及如何评估和排查效果。1. 理解 Reasoning-Induced Misalignment 的根源1.1 什么是推理诱导的失准reasoning-induced misalignment 可以理解为模型经过对齐训练后在普通对话条件下能够遵守安全边界但在需要多步推理的任务中模型为了追求“正确结果”会在内部生成大量中间推理步骤而这些步骤并不都处于对齐训练的直接监督范围内。于是某些中间步骤可能接受了一个本不该被接受的高风险前提或者把一个不安全的动作拆解成多个单看无风险的子动作。这类失准通常是隐性的。它不像普通违规那样一眼就能识别而是在推理链条里慢慢累积。比如一个多步代码生成任务单步检查时每一步都正常组合起来却形成不符合安全策略的完整方案。这就是为什么不能只靠最终回答层面的安全评估来覆盖它。1.2 安全对齐为什么在推理链上更容易失效对齐训练通常依赖偏好数据例如 RLHF 或 DPO。这些方法的核心是让模型学会区分好的最终回答和差的最终回答但中间推理链并没有被逐字逐句打标。模型在生成时探索的推理路径远远多于训练时见过的路径。路径一旦变长模型就会进入大量分布外状态原先学到的安全边界在那些状态上并不一定仍然成立。另一个原因是奖励模型的稀疏反馈。奖励模型一般只对最终回复打分而对中间步骤没有直接奖惩。如果模型在第三步出现轻微偏离随后几步可能把它“合理化”最终输出仍然看起来完整却已经偏离了安全准则。推理本身是复杂任务安全对齐的任务目标与推理正确性目标之间没有天然绑定因此需要额外的机制去约束内部状态。1.3 为什么选择方向惩罚而不是继续堆数据最直接的缓解方式是收集更多“推理条件下的坏案例”再做一轮偏好训练。但这会遇到三个问题标注成本高、坏案例覆盖不完、模型可能只是记住特定模式而不是真正修正内部表示。方向惩罚的想法是把安全对齐看成表示层面的约束而不是样本层面的约束。如果模型在推理时表现失准说明它的内部激活方向偏向了安全区域之外。方向惩罚通过惩罚这种偏离让模型学会在生成过程中保持安全方向上的激活优势。它不要求每一项推理步骤都有标签而是对中间状态本身施加约束。这种做法更加符合表示工程的目标让模型不仅“答得对”还要“想得稳”。2. 安全方向的概念与计算2.1 从表示工程理解激活方向在 Transformer 模型中每一层都会把输入表示成一个高维向量。这个向量的方向并不是随机含义而是携带了任务语义。表示工程领域发现通过控制中间层激活方向可以影响模型输出行为。比如在某些模型里“诚实”和“欺骗”可能对应不同的激活子空间方向在安全对齐场景里“合规”与“风险”也可能对应不同的方向。安全方向safety direction就是这个子空间里的一条参考轴。它的含义是如果某个中间层激活在安全方向上有足够大的投影模型更倾向于产生符合安全边界的输出如果投影很小甚至反向模型输出就容易偏离约束。方向惩罚要做的事情就是在训练或推理过程中把激活向量向这条安全方向拉近。2.2 用对比激活估算安全方向计算安全方向不一定需要复杂的数据管线。一个常见做法是准备两组文本一组是明确的合规回答一组是安全评测中标记为风险的回答然后抽取模型在某个中间层对这些文本的激活向量计算两组激活向量的均值差。数学形式可以写成s mean(h_safe) - mean(h_risky) s s / ||s||这里 h_safe 表示模型在合规样本上的中间层激活h_risky 表示在风险样本上的中间层激活。均值差的含义是两个类别在表示空间中的分离方向归一化后就是一条单位方向向量。需要注意的是这里说的风险样本应该来自你所在组织定义的安全评测集而不是公开的负面内容。实验前要先确认数据来源合规且样本量足够代表你的风险场景。样本量过少时均值差方向很容易受到个别样本噪声影响。2.3 为什么用余弦相似度而不是 L2 距离方向惩罚关注的是“方向”不是“幅度”。激活向量本身有模长早期 Layers 的模长可能和后期 Layers 相差很大。如果使用 L2 距离作为惩罚模型可以通过缩小整体模长来“骗过”损失却不改变方向分布达不到约束效果。余弦相似度只衡量两个向量之间的夹角cos_sim(h, s) dot(h, s) / (||h|| * ||s||)当 h 与安全方向一致时余弦相似度接近 1偏离时下降反向时接近 -1。惩罚项可以设计为L_pen lambda * max(0, margin - cos_sim(h, s))当激活向量与安全方向的相似度已经超过 margin不施加惩罚一旦低于 margin就拉高损失迫使模型把该方向往安全方向回调。margin 的作用是防止模型被过度约束保留正常推理所需的表示多样性。3. 用 PyTorch 实现 Safety-Direction Penalty 的最小闭环3.1 环境准备与依赖版本下面示例基于常见开源工具链实现Python 3.10 或以上版本PyTorch 2.xTransformers 库。如果使用本地 CUDA 环境建议确认显卡驱动与 PyTorch 版本匹配。安装依赖时可以这样执行pip install torch transformers datasets accelerate组件推荐版本说明Python3.10兼容现代 PyTorch 与 TransformersPyTorch2.1支持自动混合精度与多卡训练Transformers4.38提供统一模型加载接口Accelerate0.27简化分布式训练和设备管理如果原始模型很大内存有限可以先用参数量较小的模型跑通流程。方向惩罚机制与模型参数量关系不大小模型更容易调试。3.2 抽取模型中间层激活要计算安全方向需要拿到模型中间层的隐藏状态。使用 Transformers 时可以通过注册 forward hook 来实现。下面示例以 AutoModelForCausalLM 为例import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name your-model-name tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name, output_hidden_statesTrue) layer_index 16 # 根据模型层数调整 captured {} def hook_fn(module, input, output): captured[hidden] output[0][0].detach() if isinstance(output, tuple) else output.detach() target_layer model.model.layers[layer_index] hook_handle target_layer.register_forward_hook(hook_fn) text 这是一个合规回答示例。 inputs tokenizer(text, return_tensorspt) with torch.no_grad(): model(**inputs) hidden captured[hidden] hook_handle.remove()这里把 output_hidden_states 打开并不直接使用隐藏状态列表而是注册 hook 获取指定层输出。这样可以避免把全部层状态都返回内存降低显存压力。关键点是detach()。计算安全方向时不需要梯度所以用 no_grad 并且 detach。如果后面要把惩罚项加入训练则需要保留计算图hook 中就不能随意 detach而是在训练循环里单独处理。3.3 计算安全方向先构造两类文本列表。文字内容需要来自你自己的安全评测集这里只放占位说明safe_texts [ 这里填写合规回答示例 1, 这里填写合规回答示例 2, ] risky_texts [ 这里填写安全评测集中标记为风险的示例, 这里填写安全评测集中标记为风险的示例, ]然后定义函数抽取每段文本在目标层的平均激活def get_layer_activations(texts, layer_idx): activations [] for text in texts: inputs tokenizer(text, return_tensorspt) with torch.no_grad(): model(**inputs) hidden captured[hidden] # 对序列长度维做平均得到 (hidden_dim,) hidden_mean hidden.mean(dim1).squeeze() activations.append(hidden_mean) return torch.stack(activations) safe_reps get_layer_activations(safe_texts, layer_index) risky_reps get_layer_activations(risky_texts, layer_index) safety_dir safe_reps.mean(dim0) - risky_reps.mean(dim0) safety_dir safety_dir / safety_dir.norm()这里captured仍然来自前面 hook。由于是逐个文本调用模型每次都会重新覆盖 captured 字典所以要保证调用顺序与文本顺序一致。生产环境中建议使用 DataLoader 和批量处理同时维护一个索引映射。3.4 定义方向惩罚模块方向惩罚可以写成一个独立的 PyTorch 模块方便在训练循环里复用class SafetyDirectionPenalty(torch.nn.Module): def __init__(self, safety_dir, margin0.2, lambda_pen0.1, layer_index16): super().__init__() self.register_buffer(safety_dir, safety_dir) self.margin margin self.lambda_pen lambda_pen self.layer_index layer_index def forward(self, hidden_states): # hidden_states: (batch, seq_len, hidden_dim) hidden_avg hidden_states.mean(dim1) hidden_avg hidden_avg / hidden_avg.norm(dim-1, keepdimTrue) cos_sim (hidden_avg * self.safety_dir).sum(dim-1) loss torch.clamp(self.margin - cos_sim, min0).mean() return self.lambda_pen * loss这段代码做了三件事先对序列维度求平均得到每个样本的整句表示再归一化到单位向量最后计算余弦相似度并取 margin 以下的差值作为惩罚。register_buffer让安全方向随模型一起迁移设备不会参与梯度更新。3.5 将惩罚插入训练循环完整训练时需要对模型做一次前向并拿到目标层隐藏状态。可以写一个包装函数def forward_with_hidden(model, batch): outputs model(**batch, output_hidden_statesTrue) hidden outputs.hidden_states[layer_index] return outputs, hidden然后训练循环中把方向惩罚加到任务损失上optimizer torch.optim.AdamW(model.parameters(), lr2e-5) penalty_fn SafetyDirectionPenalty(safety_dir, margin0.2, lambda_pen0.1, layer_indexlayer_index) for batch in train_dataloader: optimizer.zero_grad() outputs, hidden forward_with_hidden(model, batch) task_loss outputs.loss pen_loss penalty_fn(hidden) total_loss task_loss pen_loss total_loss.backward() optimizer.step()这里的关键是output_hidden_statesTrue会返回全部层的隐藏状态内存开销较大。如果模型层数很多建议只用 hook 获取需要的层避免保存全部中间结果。方向惩罚的目标不是替代任务损失。它在训练中更像一个正则项作用是让模型在完成目标任务时不把内部状态推离安全方向太远。3.6 核心参数说明参数常见范围作用layer_index中间层或深层决定在哪个表示层施加约束margin0.1 到 0.5控制允许偏离安全方向的程度lambda_pen0.01 到 1.0控制惩罚相对于任务损失的权重batch_size8 到 64影响安全方向统计稳定性和显存占用seed固定保证方向计算和实验可复现层选得太浅激活还包含较多低层语义安全方向可能不明显选得太深接近输出层此时方向约束可能直接限制 token 分布影响生成多样性。建议从模型总层数的中间偏后位置开始尝试例如 32 层模型可以测试第 16 到第 24 层。4. 如何评估方向惩罚的真实效果4.1 评测集要分层设计只有安全指标不足以说明方向惩罚有效还应观察它是否牺牲了推理能力。建议准备三类评测集评测集类型用途样例来源基础安全集验证普通对话下的合规率通用安全测试集失准诱发集专门测试多步推理下的安全表现自建多步推理风险场景通用能力集观察推理、代码、数学能力是否下降对应能力的标准 benchmark失准诱发集是关键。它不是简单的违规 prompt而是需要模型通过多步推理才能得出结论的任务。评测时要记录每一步的中间输出而不只是最终答案否则无法定位失准发生在哪一步。4.2 主要评估指标指标计算方式说明不安全率风险输出条数 / 总条数越低越好合规拒绝率合规拒绝风险请求的比例区分拒绝与配合通用能力分数benchmark 得分或准确率不应明显下降安全方向余弦相似度激活与 safety_dir 的平均余弦越高说明方向越一致推理链偏离度每个中间步激活偏离安全方向的程度辅助定位问题步骤需要把“不安全率”和“拒绝率”分开。比如面对一个高风险请求模型如果能给出合规拒绝是被鼓励的如果避而不答或者偏移话题则不算合格。同时也要避免拒绝率过高导致普通问题也被拒答。4.3 对照实验设计建议至少对比四组基线模型不做任何方向惩罚。只加方向惩罚任务损失不变只增加惩罚项。只加安全数据微调不改变损失结构只加入更多安全样本。方向惩罚 安全数据观察组合效果。实验时保持训练步数、学习率、批次大小一致。如果多次实验的安全方向计算不稳定可以固定 seed并增加安全方向计算样本数量。每组实验至少跑 3 次取平均和方差避免单次随机性影响结论。5. 常见问题与排查链路5.1 惩罚加进去了安全分数没有提升现象训练 loss 下降正常但评测集上的不安全率没有明显变化。可能原因包括安全方向算错了、层选得不合适、lambda_pen太小、惩罚项没有参与梯度计算。先检查安全方向本身是否有区分度。可以打印安全样本和风险样本在该层激活的余弦相似度分布。如果两组分布几乎重叠说明该层不适合做方向约束。其次检查训练循环中pen_loss是否回传梯度具体可以打印pen_loss.grad_fn如果不为 None说明计算图连通。另一个常见问题是使用 hook 时提前detach()导致方向惩罚不会贡献梯度。计算安全方向时 detach 是对的但训练前向获取 hidden 时不能 detach。需要区分“计算安全方向的离线过程”和“训练时使用 hidden states 的过程”。5.2 安全方向惩罚导致通用能力下降现象安全指标提升但代码或数学能力明显下降。这种问题通常由两个参数造成lambda_pen过大或margin设置得太高。惩罚过强时模型所有层都会刻意往安全方向靠推理探索空间变小损失了多样性和创造性。建议先降低lambda_pen观察通用能力是否恢复。然后降低margin只对严重偏离的样本做惩罚。更精细的做法是只在后几层加惩罚因为前几层与能力关系更大。也可以使用分层权重比如早期层权重为 0后期层权重逐步增大。5.3 安全方向计算不稳定现象不同随机种子下计算出的 safety_dir 余弦相似度很低甚至方向相反。主要原因是样本量太少或者两类样本在某个维度上噪声过大。风险样本之间的差异通常比合规样本之间的差异更大单条异常样本会显著拉动均值方向。解决办法是扩大样本集至少每类 50 条以上对激活向量做层归一化对多份子样本分别计算方向再平均。还可以去掉激活向量的均值向量先做中心化再计算均值差减少共有的语言共性影响。5.4 排查顺序表顺序检查项检查方式处理建议1安全方向是否有效打印两类激活余弦分布分布重叠则换层或重算2惩罚项是否更新打印 pen_loss.grad_fn无梯度则检查 detach3损失比例是否合理查看 task_loss 和 pen_loss 数量级差太大则调 lambda4层选择是否合理尝试不同 layer_index搜索最佳层5评测集是否覆盖失准检查中间输出样本扩充失准诱发集先按这个顺序排查大多数问题都能定位到具体环节而不是盲目调参。6. 从实验走向落地最佳实践与可复用清单6.1 学习环境的最小复现清单刚开始接触方向惩罚时不要立刻上大规模训练。建议按下面顺序跑通选择一个 1B 到 7B 的开源模型。准备 20 条合规回答和 20 条风险回答文本。抽取模型第 8 到第 16 层中的一层计算 safety_dir。用一个小数据集做训练batch size 设为 4训练 50 步。只观察惩罚项数值是否下降不急着评估全面安全指标。这个流程可以在单张消费级显卡上完成。跑通后再扩大到更大模型和更多评测集。6.2 生产环境落地注意点生产环境使用方向惩罚不能只把它当损失函数加进去。首先safety_dir 会随着模型更新而改变需要在训练过程中周期性重算否则惩罚方向会滞后。其次不同业务场景的安全边界不同建议按场景分别计算方向而不是所有任务共享同一条 safety_dir。如果采用推理时引导需要保存目标层的 hidden states并做多步梯度更新这会显著增加推理延迟。更现实的路径是把它作为训练期正则项在离线训练完成后只保留模型权重线上推理不额外增加计算。同时要配置监控指标例如线上抽样输出中安全方向激活的余弦相似度。当指标低于阈值时触发告警再回放日志定位是数据分布变化还是模型漂移。生产环境还应保留旧版本模型方便回滚。6.3 可复用的实验记录表每次实验建议记录以下字段字段示例模型版本llama-2-7b-chat-ft-v3训练数据通用指令 安全数据layer_index20margin0.25lambda_pen0.05safety_dir 样本数safe100, risky100不安全率6.3%通用能力分数71.2推理链最大偏离度0.32复现命令记录完整命令与 seed没有这张表方向惩罚很容易陷入“调一个参数忘一个结果”的状态。复现时只看结论不够还要能拿到每个中间环节的产物。6.4 扩展方向方向惩罚不一定要局限在安全对齐上。同样的思路可以用于纠正其他类型的行为偏差比如事实性偏差、偏见表达和指令遵循漂移。可以把它看成一种通用的表示控制手段。另一个扩展方向是引入多层方向惩罚。不同层可能对应不同粒度的语义可以在若干层同时施加约束并给每层分配不同权重。再进一步可以学习一个安全子空间而不仅仅是一条方向线然后用子空间投影作为惩罚项。这样能覆盖更复杂的内部状态分布但计算复杂度也会提高。推荐后续研究路径先在小模型上确认方向稳定性再做多层扩展最后与偏好优化方法组合使用。每一层改动都要回到评测表和复现记录上验证避免方向惩罚变成新的黑盒。方向惩罚不是要替代 RLHF 或 DPO而是给对齐训练增加一个可解释、可度量的内部约束。它把安全对齐从“只看输出对不对”推进到“监督内部想的方向对不对”。对于多步推理导致的安全失准这个视角值得每个做模型对齐的人尝试。