trustformers_debug/interpretability/
attention.rs

1//! Attention analysis for transformer models
2//!
3//! This module provides comprehensive attention analysis capabilities including
4//! head specialization analysis, attention flow analysis, and pattern detection.
5
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Attention analysis result for transformer models
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct AttentionAnalysisResult {
13    /// Analysis timestamp
14    pub timestamp: DateTime<Utc>,
15    /// Attention weights by layer and head
16    pub attention_weights: HashMap<String, AttentionLayerResult>,
17    /// Token-to-token attention patterns
18    pub attention_patterns: AttentionPatterns,
19    /// Head specialization analysis
20    pub head_specialization: HeadSpecializationAnalysis,
21    /// Attention flow analysis
22    pub attention_flow: AttentionFlowAnalysis,
23    /// Key attention statistics
24    pub attention_stats: AttentionStatistics,
25}
26
27/// Attention analysis for a single layer
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct AttentionLayerResult {
30    /// Layer index
31    pub layer_index: usize,
32    /// Attention heads results
33    pub heads: HashMap<usize, AttentionHeadResult>,
34    /// Layer-level attention patterns
35    pub layer_patterns: LayerAttentionPatterns,
36    /// Layer attention statistics
37    pub layer_stats: LayerAttentionStats,
38}
39
40/// Attention analysis for a single head
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct AttentionHeadResult {
43    /// Head index
44    pub head_index: usize,
45    /// Attention matrix (token x token)
46    pub attention_matrix: Vec<Vec<f64>>,
47    /// Token attention scores
48    pub token_scores: Vec<TokenAttentionScore>,
49    /// Head specialization type
50    pub specialization_type: HeadSpecializationType,
51    /// Attention entropy
52    pub entropy: f64,
53    /// Maximum attention value
54    pub max_attention: f64,
55    /// Attention sparsity
56    pub sparsity: f64,
57}
58
59/// Token attention information
60#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct TokenAttentionScore {
62    /// Token text
63    pub token: String,
64    /// Token position
65    pub position: usize,
66    /// Attention received (sum of incoming attention)
67    pub attention_received: f64,
68    /// Attention given (sum of outgoing attention)
69    pub attention_given: f64,
70    /// Self-attention score
71    pub self_attention: f64,
72    /// Most attended tokens
73    pub most_attended: Vec<(String, f64)>,
74}
75
76/// Types of attention head specialization
77#[derive(Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq)]
78pub enum HeadSpecializationType {
79    /// Focuses on local dependencies
80    Local,
81    /// Focuses on long-range dependencies
82    Global,
83    /// Focuses on syntactic relationships
84    Syntactic,
85    /// Focuses on semantic relationships
86    Semantic,
87    /// Focuses on positional patterns
88    Positional,
89    /// Mixed or unclear specialization
90    Mixed,
91    /// Attention on special tokens
92    SpecialToken,
93}
94
95/// Overall attention patterns
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct AttentionPatterns {
98    /// Diagonal attention patterns (local dependencies)
99    pub diagonal_patterns: Vec<DiagonalPattern>,
100    /// Vertical attention patterns (token dependencies)
101    pub vertical_patterns: Vec<VerticalPattern>,
102    /// Block attention patterns
103    pub block_patterns: Vec<BlockPattern>,
104    /// Repetitive attention patterns
105    pub repetitive_patterns: Vec<RepetitivePattern>,
106}
107
108/// Diagonal attention pattern (local dependencies)
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct DiagonalPattern {
111    /// Layer and head identifier
112    pub layer_head: (usize, usize),
113    /// Diagonal offset
114    pub offset: i32,
115    /// Pattern strength
116    pub strength: f64,
117    /// Pattern coverage (how much of the diagonal)
118    pub coverage: f64,
119}
120
121/// Vertical attention pattern (specific token focus)
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct VerticalPattern {
124    /// Layer and head identifier
125    pub layer_head: (usize, usize),
126    /// Target token position
127    pub target_position: usize,
128    /// Pattern strength
129    pub strength: f64,
130    /// Number of tokens attending to target
131    pub attending_tokens: usize,
132}
133
134/// Block attention pattern (sequence segments)
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct BlockPattern {
137    /// Layer and head identifier
138    pub layer_head: (usize, usize),
139    /// Block start position
140    pub start_position: usize,
141    /// Block end position
142    pub end_position: usize,
143    /// Internal attention strength
144    pub internal_strength: f64,
145    /// External attention (leakage)
146    pub external_attention: f64,
147}
148
149/// Repetitive attention pattern
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct RepetitivePattern {
152    /// Layer and head identifier
153    pub layer_head: (usize, usize),
154    /// Pattern period
155    pub period: usize,
156    /// Pattern strength
157    pub strength: f64,
158    /// Number of repetitions
159    pub repetitions: usize,
160}
161
162/// Layer-level attention patterns
163#[derive(Debug, Clone, Serialize, Deserialize)]
164pub struct LayerAttentionPatterns {
165    /// Average attention distance
166    pub avg_attention_distance: f64,
167    /// Attention concentration (how focused)
168    pub concentration: f64,
169    /// Inter-head similarity
170    pub inter_head_similarity: f64,
171    /// Layer attention diversity
172    pub diversity: f64,
173}
174
175/// Layer attention statistics
176#[derive(Debug, Clone, Serialize, Deserialize)]
177pub struct LayerAttentionStats {
178    /// Mean attention value
179    pub mean_attention: f64,
180    /// Attention variance
181    pub attention_variance: f64,
182    /// Attention entropy
183    pub entropy: f64,
184    /// Number of significant attention connections
185    pub significant_connections: usize,
186    /// Attention sparsity ratio
187    pub sparsity_ratio: f64,
188}
189
190/// Head specialization analysis
191#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct HeadSpecializationAnalysis {
193    /// Specialization by layer
194    pub layer_specialization: HashMap<usize, Vec<HeadSpecializationType>>,
195    /// Head clustering results
196    pub head_clusters: Vec<HeadCluster>,
197    /// Specialization evolution across layers
198    pub specialization_evolution: SpecializationEvolution,
199    /// Head redundancy analysis
200    pub redundancy_analysis: HeadRedundancyAnalysis,
201}
202
203/// Cluster of similar attention heads
204#[derive(Debug, Clone, Serialize, Deserialize)]
205pub struct HeadCluster {
206    /// Cluster ID
207    pub cluster_id: usize,
208    /// Heads in cluster (layer, head)
209    pub heads: Vec<(usize, usize)>,
210    /// Cluster centroid pattern
211    pub centroid_pattern: Vec<f64>,
212    /// Cluster cohesion
213    pub cohesion: f64,
214    /// Cluster specialization type
215    pub specialization: HeadSpecializationType,
216}
217
218/// Evolution of specialization across layers
219#[derive(Debug, Clone, Serialize, Deserialize)]
220pub struct SpecializationEvolution {
221    /// Specialization distribution by layer
222    pub layer_distribution: HashMap<usize, HashMap<HeadSpecializationType, usize>>,
223    /// Specialization transitions
224    pub transitions: Vec<SpecializationTransition>,
225    /// Overall specialization trend
226    pub trend: SpecializationTrend,
227}
228
229/// Transition between specialization types
230#[derive(Debug, Clone, Serialize, Deserialize)]
231pub struct SpecializationTransition {
232    /// From layer
233    pub from_layer: usize,
234    /// To layer
235    pub to_layer: usize,
236    /// From specialization
237    pub from_specialization: HeadSpecializationType,
238    /// To specialization
239    pub to_specialization: HeadSpecializationType,
240    /// Transition strength
241    pub strength: f64,
242}
243
244/// Overall specialization trend
245#[derive(Debug, Clone, Serialize, Deserialize)]
246pub enum SpecializationTrend {
247    /// Increasing specialization with depth
248    Increasing,
249    /// Decreasing specialization with depth
250    Decreasing,
251    /// Stable specialization
252    Stable,
253    /// Mixed pattern
254    Mixed,
255}
256
257/// Head redundancy analysis
258#[derive(Debug, Clone, Serialize, Deserialize)]
259pub struct HeadRedundancyAnalysis {
260    /// Redundant head pairs
261    pub redundant_pairs: Vec<RedundantHeadPair>,
262    /// Overall redundancy score
263    pub redundancy_score: f64,
264    /// Pruning recommendations
265    pub pruning_recommendations: Vec<PruningRecommendation>,
266    /// Essential heads (cannot be pruned)
267    pub essential_heads: Vec<(usize, usize)>,
268}
269
270/// Pair of redundant attention heads
271#[derive(Debug, Clone, Serialize, Deserialize)]
272pub struct RedundantHeadPair {
273    /// First head (layer, head)
274    pub head1: (usize, usize),
275    /// Second head (layer, head)
276    pub head2: (usize, usize),
277    /// Similarity score
278    pub similarity: f64,
279    /// Redundancy type
280    pub redundancy_type: RedundancyType,
281}
282
283/// Type of redundancy between heads
284#[derive(Debug, Clone, Serialize, Deserialize)]
285pub enum RedundancyType {
286    /// Identical attention patterns
287    Identical,
288    /// Highly correlated patterns
289    Correlated,
290    /// Functionally equivalent
291    Functional,
292    /// Partially overlapping
293    Partial,
294}
295
296/// Recommendation for head pruning
297#[derive(Debug, Clone, Serialize, Deserialize)]
298pub struct PruningRecommendation {
299    /// Head to prune (layer, head)
300    pub head_to_prune: (usize, usize),
301    /// Confidence in pruning recommendation
302    pub confidence: f64,
303    /// Expected impact of pruning
304    pub expected_impact: PruningImpact,
305    /// Reason for recommendation
306    pub reason: String,
307}
308
309/// Expected impact of pruning a head
310#[derive(Debug, Clone, Serialize, Deserialize)]
311pub struct PruningImpact {
312    /// Expected performance drop
313    pub performance_drop: f64,
314    /// Memory savings
315    pub memory_savings: f64,
316    /// Computational savings
317    pub computational_savings: f64,
318    /// Risk level
319    pub risk_level: RiskLevel,
320}
321
322/// Risk level for pruning
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub enum RiskLevel {
325    /// Low risk
326    Low,
327    /// Medium risk
328    Medium,
329    /// High risk
330    High,
331    /// Critical risk
332    Critical,
333}
334
335/// Attention flow analysis
336#[derive(Debug, Clone, Serialize, Deserialize)]
337pub struct AttentionFlowAnalysis {
338    /// Flow paths through layers
339    pub flow_paths: Vec<AttentionFlowPath>,
340    /// Flow bottlenecks
341    pub bottlenecks: Vec<AttentionBottleneck>,
342    /// Flow efficiency metrics
343    pub efficiency_metrics: FlowEfficiencyMetrics,
344    /// Layer flow statistics
345    pub layer_flow_stats: HashMap<usize, LayerFlowStats>,
346}
347
348/// Attention flow path
349#[derive(Debug, Clone, Serialize, Deserialize)]
350pub struct AttentionFlowPath {
351    /// Path ID
352    pub path_id: String,
353    /// Starting token position
354    pub start_position: usize,
355    /// Ending token position
356    pub end_position: usize,
357    /// Flow steps through layers
358    pub flow_steps: Vec<LayerFlowStep>,
359    /// Total flow strength
360    pub total_strength: f64,
361    /// Path length
362    pub path_length: usize,
363}
364
365/// Flow step through a layer
366#[derive(Debug, Clone, Serialize, Deserialize)]
367pub struct LayerFlowStep {
368    /// Layer index
369    pub layer_index: usize,
370    /// Head index
371    pub head_index: usize,
372    /// Flow strength at this step
373    pub strength: f64,
374    /// Flow transformation type
375    pub transformation: FlowTransformation,
376    /// Tokens involved in this step
377    pub involved_tokens: Vec<usize>,
378}
379
380/// Type of flow transformation
381#[derive(Debug, Clone, Serialize, Deserialize)]
382pub enum FlowTransformation {
383    /// Direct attention flow
384    Direct,
385    /// Attention diffusion
386    Diffusion,
387    /// Attention concentration
388    Concentration,
389    /// Attention routing
390    Routing,
391}
392
393/// Attention bottleneck
394#[derive(Debug, Clone, Serialize, Deserialize)]
395pub struct AttentionBottleneck {
396    /// Bottleneck location (layer, head)
397    pub location: (usize, usize),
398    /// Bottleneck type
399    pub bottleneck_type: BottleneckType,
400    /// Severity of bottleneck
401    pub severity: f64,
402    /// Affected flow paths
403    pub affected_paths: Vec<String>,
404}
405
406/// Type of attention bottleneck
407#[derive(Debug, Clone, Serialize, Deserialize)]
408pub enum BottleneckType {
409    /// Information bottleneck
410    Information,
411    /// Attention bottleneck
412    Attention,
413    /// Capacity bottleneck
414    Capacity,
415    /// Flow bottleneck
416    Flow,
417}
418
419/// Flow efficiency metrics
420#[derive(Debug, Clone, Serialize, Deserialize)]
421pub struct FlowEfficiencyMetrics {
422    /// Overall flow efficiency
423    pub overall_efficiency: f64,
424    /// Information preservation ratio
425    pub information_preservation: f64,
426    /// Flow redundancy
427    pub flow_redundancy: f64,
428    /// Bottleneck impact
429    pub bottleneck_impact: f64,
430}
431
432/// Layer flow statistics
433#[derive(Debug, Clone, Serialize, Deserialize)]
434pub struct LayerFlowStats {
435    /// Incoming flow strength
436    pub incoming_flow: f64,
437    /// Outgoing flow strength
438    pub outgoing_flow: f64,
439    /// Flow retention ratio
440    pub retention_ratio: f64,
441    /// Flow transformation ratio
442    pub transformation_ratio: f64,
443}
444
445/// Attention statistics
446#[derive(Debug, Clone, Serialize, Deserialize)]
447pub struct AttentionStatistics {
448    /// Average attention entropy across all heads
449    pub avg_entropy: f64,
450    /// Attention concentration distribution
451    pub concentration_distribution: HashMap<String, f64>,
452    /// Sparsity distribution
453    pub sparsity_distribution: SparsityDistribution,
454    /// Key insights from attention analysis
455    pub insights: Vec<AttentionInsight>,
456}
457
458/// Sparsity distribution across layers and heads
459#[derive(Debug, Clone, Serialize, Deserialize)]
460pub struct SparsityDistribution {
461    /// Sparsity by layer
462    pub by_layer: HashMap<usize, f64>,
463    /// Sparsity by head
464    pub by_head: HashMap<(usize, usize), f64>,
465    /// Overall sparsity
466    pub overall_sparsity: f64,
467    /// Sparsity variance
468    pub sparsity_variance: f64,
469}
470
471/// Attention insight
472#[derive(Debug, Clone, Serialize, Deserialize)]
473pub struct AttentionInsight {
474    /// Insight type
475    pub insight_type: InsightType,
476    /// Insight description
477    pub description: String,
478    /// Confidence level
479    pub confidence: f64,
480    /// Supporting evidence
481    pub evidence: Vec<String>,
482}
483
484/// Type of attention insight
485#[derive(Debug, Clone, Serialize, Deserialize)]
486pub enum InsightType {
487    /// Head specialization pattern
488    HeadSpecialization,
489    /// Attention pattern discovery
490    PatternDiscovery,
491    /// Flow analysis finding
492    FlowAnalysis,
493    /// Redundancy detection
494    RedundancyDetection,
495    /// Performance optimization opportunity
496    OptimizationOpportunity,
497    /// Model behavior explanation
498    BehaviorExplanation,
499}
trustformers_debug/interpretability/attention.rs

trustformers_debug/interpretability/
attention.rs
