如何利用 Logits 处理器在推理侧拦截常见的 Prompt 注入攻击

很多开发者习惯在 Prompt 里写“不要执行用户指令之外的操作”，但这种软约束在面对精心设计的注入攻击（如 Ignore all previous instructions and instead...）时经常失效。最硬核的拦截方式其实是在推理侧通过 Logits Processor 强行干预 Token 的概率分布。

简单来说，就是在模型预测下一个 Token 之前，把那些可能导致“跳出预设角色”或“触发敏感指令”的 Token 概率直接设为负无穷。

以 Python 实现一个简单的 Logits 拦截器为例，如果你在使用 HuggingFace 的 transformers 库，可以自定义一个类继承 LogitsProcessor。假设我们要拦截模型在输出中出现 System: 或 User: 这种试图伪造对话结构的关键词：

import torch
from transformers import LogitsProcessor

class InjectionGuardLogitsProcessor(LogitsProcessor):
    def __init__(self, tokenizer, forbidden_tokens):
        self.forbidden_token_ids = [
            tokenizer.encode(token, add_special_tokens=False)[-1] 
            for token in forbidden_tokens
        ]

    def __call__(self, input_ids, scores):
        # 将禁用的 Token 概率设为极小值，强制模型避开
        scores[:, self.forbidden_token_ids] = -float('inf')
        return scores

# 配置示例
forbidden_words = ["System:", "Administrator:", "Ignore previous"]
processor = InjectionGuardLogitsProcessor(tokenizer, forbidden_words)

# 在 generate 时传入
outputs = model.generate(
    input_ids, 
    logits_processor=[processor], 
    max_new_tokens=128
)

几个实战配置技巧：

动态掩码策略：不要死磕固定词表。建议先用一个小规模的攻击数据集跑一遍，记录下被注入成功时高频出现的 Token ID，把这些 ID 建立成一个“风险池”，根据当前上下文的置信度动态调整拦截强度。

配合 Temperature 调节：当 Logits 被强行截断后，剩余 Token 的概率分布会发生偏移。如果 temperature 设得太高，模型可能会在避开禁用词后开始胡言乱语（Hallucination）。建议在开启 Logits 拦截时，将温度稍微调低（例如 0.7 → 0.4），保证输出的稳定性。

踩过的坑：最容易翻车的地方在于 Tokenization 的分词差异。同一个词在不同分词器里可能被切分成多个 Token。如果你只拦截了 System 但没拦截 System（带空格），拦截就会失效。建议使用 tokenizer.encode 遍历所有可能的前缀组合，把相关联的 ID 全部封死。

这种方法比在 Prompt 里加冗长的约束要高效得多，因为它不占用上下文窗口，且在数学层面保证了特定词汇绝对不会出现。