NeurIPS 2024 | 水印與高效推理如何兩全其美？最新理論：這做不到
本文第一作者為畢業于馬里蘭大學計算機系的博士胡正冕，其導師為 Heng Huang。他的主要研究方向是采樣與機器學習理論，曾在 ICML、NeurIPS、ICLR、EMNLP 等頂會上發表多篇論文。郵箱: huzhengmian@gmail.com
近日，DeepMind 團隊將水印技術和投機采樣（speculative sampling）結合，在為大語言模型加入水印的同時，提升其推理效率，降低推理成本，因此適合用于大規模生產環境。這項研究發表在了 Nature 雜志上，給出了兩種具體的結合方法，分別在水印檢測效果和生成速度上達到了現有最優水平，但無法同時兼顧兩者。
無獨有偶，另一組來自馬里蘭大學的研究人員針對同一個問題進行了理論角度的深入分析。他們發現了一個「不可行」定理，證明了不存在一個算法可以同時達到最優的推理效率和最高的水印強度。因此，任何水印系統都必須在這兩個目標之間進行權衡。這項名為「Inevitable Trade-off between Watermark Strength and Speculative Sampling Efficiency for Language Models」的研究已被 NeurIPS 2024 會議接收。
- 論文地址: https://arxiv.org/pdf/2410.20418
- 代碼倉庫: https://github.com/xiaoniu-578fa6bff964d005/AcceleratedUnbiasedWatermark
無偏水印與投機采樣方法
無偏水印方法 [1] 是一種將水印嵌入到生成文本的技術，它在理論上保證不影響生成文本的質量和多樣性，并讓水印無法被人察覺。這種水印方法可用于版權保護和來源追蹤。其核心思想是在生成過程中對候選 token 的概率進行重新加權，并要求調整后的概率分布在對水印的隨機私鑰進行平均后，在數學期望上與原始分布相同，從而避免引入偏差。
投機采樣方法 [2,3] 是一種加速大語言模型推理的技術。此方法利用較小的草稿模型快速生成草稿序列，再通過目標模型對這些草稿進行驗證和修正。草稿模型規模較小，生成速度更快但質量較低；目標模型并行進行驗證，接受符合條件的部分。這種方法可以在保持生成質量的同時，顯著提升推理效率，降低計算成本，已成為大規模應用中不可或缺的工具。
「不可行」定理
研究者提出了一個兩次重加權框架，分別對草稿模型和目標模型的輸出分布進行重加權。這個框架具有很高的自由度，兩次重加權函數可以不同，驗證算法也可以任意選擇，只要滿足最終算法保持輸出概率分布，從而維持輸出質量即可。
在此框架下，研究者嚴格描述并證明了一個「不可行」定理：當詞匯表大小大于 2 時，任何試圖同時保持水印強度和加速效果的方法都必然使用兩個平凡的重加權函數。
這個定理揭示了水印強度和投機采樣效率之間存在一個不可避免的權衡，必須在兩者之間做出選擇。任何潛在的方法都無法同時達到在不考慮加速情況下最好的水印技術的檢測效率，和在不考慮水印情況下最好的投機技術的加速效果。此結果也為未來的研究指明了方向：在設計新算法時研究者需要考慮這個理論限制，根據應用需求在兩個目標之間進行平衡。
兩種結合方法
即便「不可行」定理揭示了水印強度和采樣效率之間存在不可避免的權衡，此研究給出了兩種具體的結合水印和投機采樣的方法，使得其分別滿足以下兩點性質：
1. 保持水印強度的方法優先確保水印的可檢測性，即使這可能會犧牲一定的采樣效率。這種方法首先對目標模型和草稿模型的輸出概率分布進行重加權并從中采樣，然后基于兩個重加權分布進行投機采樣。
2. 保持采樣效率的方法優先保證生成速度，即使水印的強度可能有所降低。這種方法仍然對目標模型和草稿模型的輸出概率分布進行重加權并從中采樣，但接下來基于兩個未經過重加權的分布進行投機采樣。
這兩種方法體現了不同的取舍，使用者可以根據具體應用場景和需求進行選擇。
實驗結果
研究人員在多種文本生成任務以及多種語言模型上進行了實驗。他們引入了兩個指標：平均每步接受的 token 數（AATPS）衡量投機采樣效率，每個 token 的平均負對數 P 值（ANLPPT）衡量水印強度。
實驗結果驗證了理論分析的正確性，證實了水印強度和采樣效率之間確實存在權衡。保持水印強度的方法在水印強度上與傳統無偏水印方法相當，但采樣效率相比原始投機采樣方法有所下降；保持采樣效率的方法在采樣效率上與原始投機采樣方法相當，但水印強度有明顯降低。這表明根據實際需求，可以選擇優先保證水印強度或采樣效率的策略。
總結
此研究從理論角度證明了水印可檢測性和投機采樣效率之間存在根本沖突，這種沖突并不僅限于現在已有的方法，而是一個普遍規律。鑒于投機采樣已經成為降低推理成本不可或缺的工具，高可檢測水印系統帶來的推理成本增加在未來將持續存在。
將水印技術與投機采樣結合這個行為雖然可以讓水印更接近實用，但同時也可能帶來一些倫理問題，如未經披露的跟蹤行為。在實際應用中，人們應當謹慎、合乎倫理地應用無偏水印方法，并向用戶明確說明其存在，工作原理和意義。
