sokuto_server/guardrails/
attention.rs

1//! # 構造的分離アテンションマスクモジュール (Segment-Isolated Attention)
2//!
3//! 因果マスキングを持たない双方向エンコーダ (ModernBERT 等) において、
4//! ユーザー入力末尾に含まれる悪意ある再帰的指示 (「指示を無視せよ」等) が
5//! 文頭のシステム指示 (Instructions / Criteria) の注意重みを後方から汚染することを
6//! 物理的に防止する半双方向アテンションマスクを提供する。
7//!
8//! ## アテンション結合仕様
9//! - $S \to S$ (システム指示同士): 許可 ($0.0$)
10//! - $S \to U$ (システム指示がユーザー入力を精査): 許可 ($0.0$)
11//! - $U \to U$ (ユーザー入力同士): 許可 ($0.0$)
12//! - $U \to S$ (ユーザー入力がシステム指示を変調): **物理遮断 ($-10000.0$)**
13
14/// マスク値の定数定義。
15pub const ATTN_MASK_ALLOW: f32 = 0.0;
16/// 物理遮断を表すアテンションバイアス値 ($-\infty$ 相当)。
17pub const ATTN_MASK_BLOCK: f32 = -10000.0;
18
19/// セグメント種別。
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum SegmentType {
22    /// システム指示 (Instructions, Criteria, System Marker)。
23    System,
24    /// 外部ユーザー入力 (User State, Untrusted Prompt)。
25    User,
26    /// パディング領域。
27    Padding,
28}
29
30/// 構造的分離アテンションマスク設定。
31#[derive(Debug, Clone, PartialEq)]
32pub struct AttentionIsolationConfig {
33    /// 構造的分離アテンションマスクを有効にするか。
34    pub enabled: bool,
35    /// 遮断時に加算するバイアス値 (デフォルト: -10000.0)。
36    pub block_bias: f32,
37}
38
39impl Default for AttentionIsolationConfig {
40    fn default() -> Self {
41        Self {
42            enabled: true,
43            block_bias: ATTN_MASK_BLOCK,
44        }
45    }
46}
47
48/// 系列長 $L$ のセグメント割り当てに基づいて 2 次元のアテンションマスク $(L, L)$ を生成する。
49///
50/// 行 $i$ (Query) から列 $j$ (Key) へのアテンション重みを決定する。
51///
52/// # 引数
53/// - `segments`: 各トークン位置のセグメント種別スライス。
54/// - `config`: アテンション分離設定。
55///
56/// # 戻り値
57/// サイズ $(L, L)$ のフラット化されたベクトルまたは 2 次元ベクトル。
58pub fn build_segment_isolated_attention_mask(
59    segments: &[SegmentType],
60    config: &AttentionIsolationConfig,
61) -> Vec<Vec<f32>> {
62    let len = segments.len();
63    let mut mask = vec![vec![ATTN_MASK_ALLOW; len]; len];
64
65    if !config.enabled {
66        // パディングのみ遮断
67        for row in mask.iter_mut() {
68            for (j, seg) in segments.iter().enumerate() {
69                if *seg == SegmentType::Padding {
70                    row[j] = config.block_bias;
71                }
72            }
73        }
74        return mask;
75    }
76
77    for (i, &query_seg) in segments.iter().enumerate() {
78        for (j, &key_seg) in segments.iter().enumerate() {
79            if key_seg == SegmentType::Padding {
80                // パディングへの注意は常に遮断
81                mask[i][j] = config.block_bias;
82            } else if query_seg == SegmentType::User && key_seg == SegmentType::System {
83                // ユーザー入力からシステム指示へのアテンションを物理遮断 (U -> S)
84                mask[i][j] = config.block_bias;
85            } else {
86                // それ以外はすべて相互作用を許可 (S -> S, S -> U, U -> U)
87                mask[i][j] = ATTN_MASK_ALLOW;
88            }
89        }
90    }
91
92    mask
93}
94
95/// システム指示長とユーザー入力長からフラット化された 1 次元アテンションマスク配列 $(L \times L)$ を生成する。
96///
97/// テンソル入力 (ONNX Runtime 等) への直接変換に最適化されている。
98///
99/// # 引数
100/// - `sys_len`: システム指示トークン数。
101/// - `user_len`: ユーザー入力トークン数。
102/// - `pad_len`: パディングトークン数。
103/// - `config`: アテンション分離設定。
104pub fn build_flat_attention_mask(
105    sys_len: usize,
106    user_len: usize,
107    pad_len: usize,
108    config: &AttentionIsolationConfig,
109) -> Vec<f32> {
110    let total_len = sys_len + user_len + pad_len;
111    let mut flat = vec![ATTN_MASK_ALLOW; total_len * total_len];
112
113    let block_val = if config.enabled {
114        config.block_bias
115    } else {
116        ATTN_MASK_ALLOW
117    };
118
119    for i in 0..total_len {
120        let is_query_user = i >= sys_len && i < sys_len + user_len;
121
122        for j in 0..total_len {
123            let idx = i * total_len + j;
124            let is_key_system = j < sys_len;
125            let is_key_pad = j >= sys_len + user_len;
126
127            if is_key_pad {
128                flat[idx] = config.block_bias;
129            } else if config.enabled && is_query_user && is_key_system {
130                // User -> System 遮断
131                flat[idx] = block_val;
132            } else {
133                flat[idx] = ATTN_MASK_ALLOW;
134            }
135        }
136    }
137
138    flat
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn test_segment_isolated_attention_mask_behavior() {
147        let config = AttentionIsolationConfig::default();
148        let segments = vec![
149            SegmentType::System,  // 0: Instruction
150            SegmentType::System,  // 1: Criteria
151            SegmentType::User,    // 2: User Input
152            SegmentType::User,    // 3: User Input
153            SegmentType::Padding, // 4: Pad
154        ];
155
156        let mask = build_segment_isolated_attention_mask(&segments, &config);
157
158        // S -> S (0 -> 1): 許可
159        assert_eq!(mask[0][1], ATTN_MASK_ALLOW);
160        // S -> U (0 -> 2): 許可 (システムがユーザー入力を精査)
161        assert_eq!(mask[0][2], ATTN_MASK_ALLOW);
162        // U -> U (2 -> 3): 許可
163        assert_eq!(mask[2][3], ATTN_MASK_ALLOW);
164
165        // U -> S (2 -> 0, 3 -> 1): 物理遮断 (ユーザーがシステム指示を変調不可)
166        assert_eq!(mask[2][0], ATTN_MASK_BLOCK);
167        assert_eq!(mask[2][1], ATTN_MASK_BLOCK);
168        assert_eq!(mask[3][0], ATTN_MASK_BLOCK);
169        assert_eq!(mask[3][1], ATTN_MASK_BLOCK);
170
171        // 任意 -> Padding: 物理遮断
172        assert_eq!(mask[0][4], ATTN_MASK_BLOCK);
173        assert_eq!(mask[2][4], ATTN_MASK_BLOCK);
174    }
175
176    #[test]
177    fn test_flat_attention_mask_parity() {
178        let config = AttentionIsolationConfig::default();
179        let sys_len = 2;
180        let user_len = 2;
181        let pad_len = 1;
182        let total = sys_len + user_len + pad_len;
183
184        let flat = build_flat_attention_mask(sys_len, user_len, pad_len, &config);
185        assert_eq!(flat.len(), total * total);
186
187        // User (i=2) -> System (j=0) は遮断
188        assert_eq!(flat[2 * total], ATTN_MASK_BLOCK);
189        // System (i=0) -> User (j=2) は許可
190        assert_eq!(flat[2], ATTN_MASK_ALLOW);
191        // Key is Padding (j=4) は遮断
192        assert_eq!(flat[4], ATTN_MASK_BLOCK);
193    }
194}
sokuto_server/guardrails/attention.rs

sokuto_server/guardrails/
attention.rs
