ruvector-dag 0.1.0

Directed Acyclic Graph (DAG) structures for query plan optimization with neural 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

//! EWC++: Elastic Weight Consolidation to prevent forgetting

use ndarray::Array1;

#[derive(Debug, Clone)]
pub struct EwcConfig {
    pub lambda: f32,  // Importance weight (2000-15000)
    pub decay: f32,   // Fisher decay rate
    pub online: bool, // Use online EWC
}

impl Default for EwcConfig {
    fn default() -> Self {
        Self {
            lambda: 5000.0,
            decay: 0.99,
            online: true,
        }
    }
}

pub struct EwcPlusPlus {
    config: EwcConfig,
    fisher_diag: Option<Array1<f32>>,
    optimal_params: Option<Array1<f32>>,
    task_count: usize,
}

impl EwcPlusPlus {
    pub fn new(config: EwcConfig) -> Self {
        Self {
            config,
            fisher_diag: None,
            optimal_params: None,
            task_count: 0,
        }
    }

    /// Consolidate current parameters after training
    pub fn consolidate(&mut self, params: &Array1<f32>, fisher: &Array1<f32>) {
        if self.config.online && self.fisher_diag.is_some() {
            // Online EWC: accumulate Fisher information
            let current_fisher = self.fisher_diag.as_ref().unwrap();
            self.fisher_diag =
                Some(current_fisher * self.config.decay + fisher * (1.0 - self.config.decay));
        } else {
            self.fisher_diag = Some(fisher.clone());
        }

        self.optimal_params = Some(params.clone());
        self.task_count += 1;
    }

    /// Compute EWC penalty for given parameters
    pub fn penalty(&self, params: &Array1<f32>) -> f32 {
        match (&self.fisher_diag, &self.optimal_params) {
            (Some(fisher), Some(optimal)) => {
                let diff = params - optimal;
                let weighted = &diff * &diff * fisher;
                0.5 * self.config.lambda * weighted.sum()
            }
            _ => 0.0,
        }
    }

    /// Compute gradient of EWC penalty
    pub fn penalty_gradient(&self, params: &Array1<f32>) -> Option<Array1<f32>> {
        match (&self.fisher_diag, &self.optimal_params) {
            (Some(fisher), Some(optimal)) => {
                let diff = params - optimal;
                Some(self.config.lambda * fisher * &diff)
            }
            _ => None,
        }
    }

    /// Compute Fisher information from gradients
    pub fn compute_fisher(gradients: &[Array1<f32>]) -> Array1<f32> {
        if gradients.is_empty() {
            return Array1::zeros(0);
        }

        let dim = gradients[0].len();
        let mut fisher = Array1::zeros(dim);

        for grad in gradients {
            fisher = fisher + grad.mapv(|x| x * x);
        }

        fisher / gradients.len() as f32
    }

    pub fn has_prior(&self) -> bool {
        self.fisher_diag.is_some()
    }

    pub fn task_count(&self) -> usize {
        self.task_count
    }
}