optirs-learned 0.3.1

OptiRS learned optimizers and meta-learning
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

// Configuration types for the Transformer-based Neural Optimizer
//
// This module defines configuration structures and enums that control
// the behavior of the Transformer optimizer including architecture settings,
// attention mechanisms, and optimization strategies.

use super::super::{LearnedOptimizerConfig, MetaOptimizationStrategy};

/// Configuration specific to Transformer optimizer
#[derive(Debug, Clone)]
pub struct TransformerOptimizerConfig {
    /// Base learned optimizer config
    pub base_config: LearnedOptimizerConfig,

    /// Model dimension (d_model)
    pub modeldim: usize,

    /// Number of attention heads
    pub numheads: usize,

    /// Feed-forward network dimension
    pub ff_dim: usize,

    /// Number of transformer layers
    pub num_layers: usize,

    /// Maximum sequence length
    pub max_sequence_length: usize,

    /// Attention dropout rate
    pub attention_dropout: f64,

    /// Feed-forward dropout rate
    pub ff_dropout: f64,

    /// Layer normalization epsilon
    pub layer_norm_eps: f64,

    /// Use pre-layer normalization
    pub pre_layer_norm: bool,

    /// Positional encoding type
    pub pos_encoding_type: PositionalEncodingType,

    /// Enable relative position bias
    pub relative_position_bias: bool,

    /// Use rotary position embedding
    pub use_rope: bool,

    /// Enable gradient checkpointing
    pub gradient_checkpointing: bool,

    /// Attention pattern optimization
    pub attention_optimization: AttentionOptimization,

    /// Multi-scale attention
    pub multi_scale_attention: bool,

    /// Cross-attention for multi-task learning
    pub cross_attention: bool,

    /// Memory efficiency mode
    pub memory_efficient: bool,
}

/// Types of positional encoding
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PositionalEncodingType {
    /// Sinusoidal position encoding
    Sinusoidal,

    /// Learned position embedding
    Learned,

    /// Rotary position embedding (RoPE)
    Rotary,

    /// Relative position encoding
    Relative,

    /// ALiBi (Attention with Linear Biases)
    ALiBi,
}

/// Attention optimization strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AttentionOptimization {
    /// Standard full attention
    Full,

    /// Sparse attention patterns
    Sparse,

    /// Linear attention approximation
    Linear,

    /// Local attention windows
    Local,

    /// Hierarchical attention
    Hierarchical,

    /// Adaptive attention sparsity
    Adaptive,
}

impl Default for TransformerOptimizerConfig {
    fn default() -> Self {
        Self {
            base_config: LearnedOptimizerConfig::default(),
            modeldim: 512,
            numheads: 8,
            ff_dim: 2048,
            num_layers: 6,
            max_sequence_length: 1024,
            attention_dropout: 0.1,
            ff_dropout: 0.1,
            layer_norm_eps: 1e-5,
            pre_layer_norm: true,
            pos_encoding_type: PositionalEncodingType::Sinusoidal,
            relative_position_bias: false,
            use_rope: false,
            gradient_checkpointing: false,
            attention_optimization: AttentionOptimization::Full,
            multi_scale_attention: false,
            cross_attention: false,
            memory_efficient: false,
        }
    }
}

impl Default for PositionalEncodingType {
    fn default() -> Self {
        Self::Sinusoidal
    }
}

impl Default for AttentionOptimization {
    fn default() -> Self {
        Self::Full
    }
}