Spam filter performance and accuracy improvements (part 8)

crates/nlp/src/classifier/adam.rs

@@ -0,0 +1,111 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{Optimizer, model::FhClassifier};
+
+pub struct Adam {
+    parameters: Vec<f32>,
+    bias: f32,
+    learning_rate: f32,
+    beta1: f32,
+    beta2: f32,
+    epsilon: f32,
+    t: f32,
+    m0: Vec<f32>,
+    v0: Vec<f32>,
+    m_bias: f32,
+    v_bias: f32,
+
+    // Step info
+    bias2_sqrt: f32,
+    alpha_t: f32,
+}
+
+impl Adam {
+    pub fn new(n_parameters: usize, learning_rate: f32) -> Self {
+        Adam {
+            parameters: vec![0.0; n_parameters],
+            learning_rate,
+            beta1: 0.9,
+            beta2: 0.999,
+            epsilon: 1e-8,
+            t: 0.0,
+            m0: vec![0.0; n_parameters],
+            v0: vec![0.0; n_parameters],
+            m_bias: 0.0,
+            v_bias: 0.0,
+            bias: 0.0,
+            bias2_sqrt: 0.0,
+            alpha_t: 0.0,
+        }
+    }
+
+    pub fn with_hyperparams(mut self, beta1: f32, beta2: f32, epsilon: f32) -> Self {
+        self.beta1 = beta1;
+        self.beta2 = beta2;
+        self.epsilon = epsilon;
+        self
+    }
+
+    pub fn with_initial_weights(self, value: f32) -> Self {
+        Adam {
+            parameters: vec![value; self.parameters.len()],
+            ..self
+        }
+    }
+}
+
+impl Optimizer for Adam {
+    #[inline(always)]
+    fn step(&mut self) {
+        self.t += 1.0;
+        let bias1 = 1.0 - self.beta1.powf(self.t);
+        self.bias2_sqrt = (1.0 - self.beta2.powf(self.t)).sqrt();
+        self.alpha_t = self.learning_rate / bias1;
+    }
+
+    #[inline(always)]
+    fn update_param(&mut self, i: usize, g: f32) {
+        self.m0[i] = self.beta1 * self.m0[i] + (1.0 - self.beta1) * g;
+        self.v0[i] = self.beta2 * self.v0[i] + (1.0 - self.beta2) * g * g;
+        self.parameters[i] -=
+            self.alpha_t * self.m0[i] / (self.v0[i].sqrt() / self.bias2_sqrt + self.epsilon);
+    }
+
+    #[inline(always)]
+    fn update_bias(&mut self, g: f32) {
+        self.m_bias = self.beta1 * self.m_bias + (1.0 - self.beta1) * g;
+        self.v_bias = self.beta2 * self.v_bias + (1.0 - self.beta2) * g * g;
+        self.bias -=
+            self.alpha_t * self.m_bias / (self.v_bias.sqrt() / self.bias2_sqrt + self.epsilon);
+    }
+
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        self.parameters[idx]
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        self.bias
+    }
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.parameters[idx]
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self.parameters.clone(),
+            bias: self.bias,
+        }
+    }
+
+    fn num_parameters(&self) -> usize {
+        self.parameters.len()
+    }
+}

crates/nlp/src/classifier/feature.rs

@@ -4,82 +4,160 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
-use nohash::BuildNoHashHasher;
 use std::collections::HashMap;
 use xxhash_rust::xxh3::xxh3_64_with_seed;
 
-pub struct Sample {
-    pub(super) features: Features,
-    pub(super) class: f32,
+#[derive(Debug)]
+pub struct Sample<T> {
+    pub features: Vec<T>,
+    pub class: f32,
 }
 
-pub struct Features(pub(super) HashMap<u32, f32, BuildNoHashHasher<u32>>);
-
-pub struct FeatureBuilder {
-    pub(super) features_mask: u32,
+pub struct FhFeatureBuilder {
+    pub(super) weight_mask: u64,
 }
 
-pub trait Feature {
+#[derive(Debug)]
+pub struct FhFeature {
+    pub idx: usize,
+    pub weight: f32,
+}
+
+#[derive(Debug)]
+pub struct CcfhFeature {
+    pub idx_w1: usize,
+    pub idx_w2: usize,
+    pub idx_i: usize,
+    pub weight: f32,
+}
+
+pub struct CcfhFeatureBuilder {
+    pub(super) weight_mask: u64,
+    pub(super) indicator_mask: u64,
+}
+
+pub trait FeatureWeight {
+    fn idx(&self) -> usize;
+    fn weight(&self) -> f32;
+    fn weight_mut(&mut self) -> &mut f32;
+}
+
+pub trait UnprocessedFeature {
     fn prefix(&self) -> u16;
     fn value(&self) -> &[u8];
 }
 
-impl FeatureBuilder {
-    pub fn scale<I: Feature>(&self, features: &mut HashMap<I, f32>) {
+impl FeatureWeight for FhFeature {
+    fn weight(&self) -> f32 {
+        self.weight
+    }
+
+    fn weight_mut(&mut self) -> &mut f32 {
+        &mut self.weight
+    }
+
+    fn idx(&self) -> usize {
+        self.idx
+    }
+}
+
+impl FeatureWeight for CcfhFeature {
+    fn weight(&self) -> f32 {
+        self.weight
+    }
+
+    fn weight_mut(&mut self) -> &mut f32 {
+        &mut self.weight
+    }
+
+    fn idx(&self) -> usize {
+        self.idx_w1
+    }
+}
+
+impl FeatureBuilder for FhFeatureBuilder {
+    type Feature = FhFeature;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> FhFeature {
+        let hash1 = xxh3_64_with_seed(bytes, 0);
+        let sign = if hash1 & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+        FhFeature {
+            idx: (hash1 & self.weight_mask) as usize,
+            weight: sign * weight,
+        }
+    }
+}
+
+impl FeatureBuilder for CcfhFeatureBuilder {
+    type Feature = CcfhFeature;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> CcfhFeature {
+        let hash1 = xxh3_64_with_seed(bytes, 0);
+        let hash2 = xxh3_64_with_seed(bytes, 0x9E3779B97F4A7C15);
+        let hash3 = xxh3_64_with_seed(bytes, 0x517CC1B727220A95);
+        let sign = if hash3 & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+        CcfhFeature {
+            idx_w1: (hash1 & self.weight_mask) as usize,
+            idx_w2: (hash2 & self.weight_mask) as usize,
+            idx_i: (hash3 & self.indicator_mask) as usize,
+            weight: sign * weight,
+        }
+    }
+}
+
+pub trait FeatureBuilder {
+    // Feature type associated type
+    type Feature: FeatureWeight;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> Self::Feature;
+
+    fn scale<I: UnprocessedFeature>(&self, features: &mut HashMap<I, f32>) {
         // Log frequency scaling
         for x in features.values_mut() {
             *x = x.ln_1p();
         }
     }
 
-    pub fn build<I: Feature>(
+    fn build<I: UnprocessedFeature>(
         &self,
-        features: &HashMap<I, f32>,
+        features_in: &HashMap<I, f32>,
         account_id: Option<u32>,
-    ) -> Features {
-        // Do the "hash trick"
-        let mut features_map =
-            HashMap::with_capacity_and_hasher(features.len(), BuildNoHashHasher::default());
+    ) -> Vec<Self::Feature> {
+        let mut features_out = Vec::with_capacity(features_in.len());
         let mut buf = Vec::with_capacity(2 + 4 + 63);
-        for (feature, count) in features {
+        for (feature, count) in features_in {
             buf.extend_from_slice(&feature.prefix().to_be_bytes());
             buf.extend_from_slice(feature.value());
-
-            let big_hash = xxh3_64_with_seed(&buf, 0);
-            let hash = big_hash as u32 & self.features_mask;
-            let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
-
-            *features_map.entry(hash).or_default() += sign * count;
+            features_out.push(self.build_feature(&buf, *count));
 
             if let Some(account_id) = account_id {
                 buf.extend_from_slice(&account_id.to_be_bytes());
-                let big_hash = xxh3_64_with_seed(&buf, 0);
-                let hash = big_hash as u32 & self.features_mask;
-                let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
-
-                *features_map.entry(hash).or_default() += sign * count;
+                features_out.push(self.build_feature(&buf, *count));
             }
+
             buf.clear();
         }
 
         // L2 normalization
-        let sum_of_squares = features_map
-            .values()
-            .map(|&x| x as f64 * x as f64)
+        let sum_of_squares = features_out
+            .iter()
+            .map(|f| f.weight() as f64 * f.weight() as f64)
             .sum::<f64>();
         if sum_of_squares > 0.0 {
             let norm = sum_of_squares.sqrt() as f32;
-            for x in features_map.values_mut() {
-                *x /= norm;
+            for feature in &mut features_out {
+                *feature.weight_mut() /= norm;
             }
         }
 
-        Features(features_map)
+        features_out
     }
 }
 
-impl Sample {
-    pub fn new(features: Features, class: bool) -> Self {
+impl<T> Sample<T> {
+    pub fn new(features: Vec<T>, class: bool) -> Self {
         Self {
             features,
             class: if class { 1.0 } else { 0.0 },
@@ -87,14 +165,8 @@ impl Sample {
     }
 }
 
-impl AsRef<Sample> for Sample {
-    fn as_ref(&self) -> &Sample {
-        self
-    }
-}
-
-impl AsRef<Features> for Features {
-    fn as_ref(&self) -> &Features {
+impl<T> AsRef<Sample<T>> for Sample<T> {
+    fn as_ref(&self) -> &Sample<T> {
         self
     }
 }

crates/nlp/src/classifier/ftrl.rs

@@ -0,0 +1,132 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{Optimizer, model::FhClassifier};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct Ftrl {
+    alpha: f64,
+    beta: f64,
+    l1_ratio: f64,
+    l2_ratio: f64,
+    zn: Vec<Zn>,
+    zn_bias: Zn,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Clone, Copy, Debug, Default)]
+pub struct Zn {
+    z: f32,
+    n: f64,
+}
+
+impl Ftrl {
+    pub fn new(n_features: usize) -> Self {
+        Ftrl {
+            alpha: 2.0,
+            beta: 1.0,
+            l1_ratio: 0.001,
+            l2_ratio: 0.0001,
+            zn: vec![Zn::default(); n_features],
+            zn_bias: Zn::default(),
+        }
+    }
+
+    pub fn with_hyperparams(mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
+        self.alpha = alpha;
+        self.beta = beta;
+        self.l1_ratio = l1_ratio;
+        self.l2_ratio = l2_ratio;
+        self
+    }
+
+    pub fn set_hyperparams(&mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) {
+        self.alpha = alpha;
+        self.beta = beta;
+        self.l1_ratio = l1_ratio;
+        self.l2_ratio = l2_ratio;
+    }
+
+    pub fn with_initial_weights(self, value: f32) -> Self {
+        Ftrl {
+            zn: vec![Zn { z: value, n: 0.0 }; self.zn.len()],
+            ..self
+        }
+    }
+}
+
+impl Optimizer for Ftrl {
+    #[inline(always)]
+    fn update_param(&mut self, idx: usize, grad: f32) {
+        let zn = &mut self.zn[idx];
+        let current_w = if zn.z.abs() as f64 <= self.l1_ratio {
+            0.0
+        } else {
+            -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+        };
+        let grad = grad as f64;
+        let grad_sq = grad * grad;
+        let sigma = ((zn.n + grad_sq).sqrt() - zn.n.sqrt()) / self.alpha;
+        zn.z += (grad - sigma * current_w as f64) as f32;
+        zn.n += grad_sq;
+    }
+
+    #[inline(always)]
+    fn update_bias(&mut self, grad: f32) {
+        let current_bias = -self.zn_bias.z
+            / ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32;
+        let grad = grad as f64;
+        let grad_sq = grad * grad;
+        let sigma = ((self.zn_bias.n + grad_sq).sqrt() - self.zn_bias.n.sqrt()) / self.alpha;
+        self.zn_bias.z += (grad - sigma * current_bias as f64) as f32;
+        self.zn_bias.n += grad_sq;
+    }
+
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        let zn = self.zn[idx];
+        if zn.z.abs() as f64 <= self.l1_ratio {
+            0.0
+        } else {
+            -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+        }
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        -self.zn_bias.z / ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32
+    }
+
+    fn step(&mut self) {}
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.zn[idx].z
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self
+                .zn
+                .iter()
+                .map(|zn| {
+                    if zn.z.abs() as f64 <= self.l1_ratio {
+                        0.0
+                    } else {
+                        -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                            / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+                    }
+                })
+                .collect(),
+            bias: self.get_bias(),
+        }
+    }
+
+    fn num_parameters(&self) -> usize {
+        self.zn.len()
+    }
+}

crates/nlp/src/classifier/mod.rs

@@ -4,5 +4,46 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
+use crate::classifier::model::FhClassifier;
+
+pub mod adam;
 pub mod feature;
+pub mod ftrl;
+pub mod model;
+pub mod reservoir;
 pub mod sgd;
+pub mod train;
+
+const MAX_DLOSS: f32 = 1e4;
+
+pub trait Optimizer {
+    fn step(&mut self);
+    fn update_param(&mut self, i: usize, g: f32);
+    fn update_bias(&mut self, g: f32);
+    fn get_param(&self, idx: usize) -> f32;
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32;
+    fn get_bias(&self) -> f32;
+    fn build_classifier(&self) -> FhClassifier;
+    fn num_parameters(&self) -> usize;
+}
+
+#[inline(always)]
+fn sigmoid(z: f32) -> f32 {
+    let z = z.clamp(-35.0, 35.0);
+    if z >= 0.0 {
+        1.0 / (1.0 + (-z).exp())
+    } else {
+        let exp_z = z.exp();
+        exp_z / (1.0 + exp_z)
+    }
+}
+
+#[inline(always)]
+fn gradient(y: f32, p: f32) -> f32 {
+    if p > -16.0 {
+        let exp_tmp = (-p).exp();
+        ((1.0 - y) - y * exp_tmp) / (1.0 + exp_tmp)
+    } else {
+        p.exp() - y
+    }
+}

crates/nlp/src/classifier/model.rs

@@ -0,0 +1,109 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{
+    feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder},
+    sigmoid,
+};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct FhClassifier {
+    pub(crate) parameters: Vec<f32>,
+    pub(crate) bias: f32,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct CcfhClassifier {
+    pub(crate) parameters: Vec<f32>,
+    pub(crate) indicators: Vec<f32>,
+    pub(crate) bias: f32,
+}
+
+impl FhClassifier {
+    pub fn predict_proba_sample(&self, features: &[FhFeature]) -> f32 {
+        let mut z: f32 = 0.0;
+
+        for f in features {
+            z += self.parameters[f.idx] * f.weight;
+        }
+
+        sigmoid(z + self.bias)
+    }
+
+    pub fn predict(&self, features: &[FhFeature]) -> f32 {
+        if self.predict_proba_sample(features) > 0.7 {
+            1.0
+        } else {
+            0.0
+        }
+    }
+
+    pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
+    where
+        I: IntoIterator,
+        I::Item: AsRef<Vec<FhFeature>>,
+    {
+        test.into_iter()
+            .map(|features| self.predict(features.as_ref()))
+            .collect()
+    }
+
+    pub fn feature_builder(&self) -> FhFeatureBuilder {
+        FhFeatureBuilder {
+            weight_mask: (self.parameters.len() - 1) as u64,
+        }
+    }
+
+    pub fn parameters(&self) -> &[f32] {
+        &self.parameters
+    }
+
+    pub fn bias(&self) -> f32 {
+        self.bias
+    }
+}
+
+impl CcfhClassifier {
+    pub fn predict_proba_sample(&self, features: &[CcfhFeature]) -> f32 {
+        let mut z: f32 = 0.0;
+        for f in features {
+            let q = self.indicators[f.idx_i];
+            let v1 = self.parameters[f.idx_w1];
+            let v2 = self.parameters[f.idx_w2];
+            z += (q * v1 + (1.0 - q) * v2) * f.weight;
+        }
+        sigmoid(z + self.bias)
+    }
+
+    pub fn predict(&self, features: &[CcfhFeature]) -> f32 {
+        if self.predict_proba_sample(features) >= 0.5 {
+            1.0
+        } else {
+            0.0
+        }
+    }
+
+    pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
+    where
+        I: IntoIterator,
+        I::Item: AsRef<Vec<CcfhFeature>>,
+    {
+        test.into_iter()
+            .map(|features| self.predict(features.as_ref()))
+            .collect()
+    }
+
+    pub fn feature_builder(&self) -> CcfhFeatureBuilder {
+        CcfhFeatureBuilder {
+            weight_mask: (self.parameters.len() - 1) as u64,
+            indicator_mask: (self.indicators.len() - 1) as u64,
+        }
+    }
+
+    pub fn is_active(&self) -> bool {
+        !self.parameters.is_empty()
+    }
+}

crates/nlp/src/classifier/reservoir.rs

@@ -0,0 +1,81 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use rand::{Rng, seq::IndexedRandom};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct SampleReservoir<T> {
+    pub spam: SampleReservoirClass<T>,
+    pub ham: SampleReservoirClass<T>,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct SampleReservoirClass<T> {
+    pub buffer: Vec<T>,
+    pub total_seen: u64,
+}
+
+impl<T: Clone + Eq> SampleReservoir<T> {
+    pub fn update_reservoir(&mut self, item: &T, is_spam: bool, capacity: usize) {
+        let class = if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        };
+
+        class.total_seen += 1;
+
+        if class.buffer.len() < capacity {
+            class.buffer.push(item.clone());
+        } else if let Some(buf) = class
+            .buffer
+            .get_mut(rand::rng().random_range(0..class.total_seen as usize))
+        {
+            *buf = item.clone();
+        }
+    }
+
+    pub fn replay_samples(
+        &mut self,
+        count_needed: usize,
+        is_spam: bool,
+    ) -> impl Iterator<Item = &T> {
+        (if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        })
+        .buffer
+        .choose_multiple(&mut rand::rng(), count_needed)
+    }
+
+    pub fn remove_sample(&mut self, item: &T, is_spam: bool) {
+        let class = if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        };
+
+        if let Some(pos) = class.buffer.iter().position(|x| x == item) {
+            class.buffer.swap_remove(pos);
+        }
+    }
+}
+
+impl<T> Default for SampleReservoir<T> {
+    fn default() -> Self {
+        SampleReservoir {
+            spam: SampleReservoirClass {
+                buffer: Vec::new(),
+                total_seen: 0,
+            },
+            ham: SampleReservoirClass {
+                buffer: Vec::new(),
+                total_seen: 0,
+            },
+        }
+    }
+}

crates/nlp/src/classifier/sgd.rs

@@ -4,163 +4,124 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
-use crate::classifier::feature::{FeatureBuilder, Features, Sample};
-use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
+use crate::classifier::{Optimizer, gradient, model::FhClassifier};
 
 #[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
-pub struct TextClassifier {
-    weights: Vec<f32>,
-    intercept: f32,
+pub struct Sgd {
+    parameters: Vec<f32>,
+    bias: f32,
+    alpha: f64,
+    l1_ratio: f64,
+    l2_ratio: f64,
+    t: f64,
+    w_scale: f32,
+    optimal_init: f64,
+    eta: f32,
+    u: f32,
+    q: Vec<f32>,
 }
 
-const MAX_DLOSS: f32 = 1e4;
-
-impl TextClassifier {
-    pub fn new(n_features: usize) -> Self {
-        TextClassifier {
-            weights: vec![0.0; n_features],
-            intercept: 0.0,
-        }
-    }
-
-    pub fn fit(&mut self, samples: &mut [impl AsRef<Sample>], n_epochs: usize, alpha: f32) {
-        let mut rng = StdRng::seed_from_u64(42);
-        let mut t = 1;
-        let mut w_scale = 1.0;
-
-        // Heuristic to initialize 'optimal' learning rate
+impl Sgd {
+    pub fn new(n_features: usize, alpha: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
         let typw = (1.0 / alpha.sqrt()).sqrt();
-        let initial_eta0 = typw / 1.0_f32.max(gradient(1.0, -typw));
+        let initial_eta0 = typw / 1.0_f64.max(gradient(1.0, -typw as f32) as f64);
         let optimal_init = 1.0 / (initial_eta0 * alpha);
 
-        for _ in 0..n_epochs {
-            samples.shuffle(&mut rng);
-
-            for sample in samples.iter() {
-                // Prediction
-                let sample = sample.as_ref();
-                let mut dot: f32 = 0.0;
-                for (idx, feature) in &sample.features.0 {
-                    dot += self.weights[*idx as usize] * *feature;
-                }
-                let p = (dot * w_scale) + self.intercept;
-                let eta = 1.0 / (alpha * (optimal_init + (t as f32) - 1.0));
-
-                // Compute Loss & Gradient
-                let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
-
-                // Lazy weight decay
-                w_scale *= 1.0 - (eta * alpha);
-
-                // Update weights
-                let update = -eta * dloss;
-                if update != 0.0 {
-                    let scaled_update = update / w_scale;
-
-                    for (idx, feature) in &sample.features.0 {
-                        self.weights[*idx as usize] += scaled_update * *feature;
-                    }
-
-                    self.intercept += update;
-                }
-
-                // Rescale weights if w_scale is too small or too large
-                if !(1e-6..=1e6).contains(&w_scale) {
-                    for w in &mut self.weights {
-                        *w *= w_scale;
-                    }
-                    w_scale = 1.0;
-                }
-
-                t += 1;
-            }
+        Sgd {
+            parameters: vec![0.0; n_features],
+            bias: 0.0,
+            alpha,
+            l1_ratio,
+            l2_ratio,
+            t: 0.0,
+            w_scale: 1.0,
+            optimal_init,
+            eta: initial_eta0 as f32,
+            u: 0.0,
+            q: vec![0.0; n_features],
         }
+    }
 
-        if w_scale != 1.0 {
-            for w in &mut self.weights {
-                *w *= w_scale;
+    pub fn with_initial_parameters(self, value: f32) -> Self {
+        Sgd {
+            parameters: vec![value; self.parameters.len()],
+            ..self
+        }
+    }
+
+    fn maybe_rescale(&mut self) {
+        if !(1e-6..=1e6).contains(&self.w_scale) {
+            for w in &mut self.parameters {
+                *w *= self.w_scale;
+            }
+            self.w_scale = 1.0;
+        }
+    }
+
+    #[inline(always)]
+    fn apply_l1_penalty(&mut self) {
+        if self.l1_ratio > 0.0 {
+            for (z, q) in self.parameters.iter_mut().zip(self.q.iter_mut()) {
+                let z_orig = *z;
+                let scaled_z = *z * self.w_scale;
+                if scaled_z > 0.0 {
+                    *z = (*z - (self.u + *q) / self.w_scale).max(0.0);
+                } else if scaled_z < 0.0 {
+                    *z = (*z + (self.u - *q) / self.w_scale).min(0.0);
+                }
+                *q += self.w_scale * (z_orig - *z);
             }
         }
     }
+}
 
-    pub fn predict_proba_sample(&self, features: &Features) -> f32 {
-        let mut z: f32 = 0.0;
-        for (idx, feature) in &features.0 {
-            z += self.weights[*idx as usize] * *feature;
-        }
-        z += self.intercept;
-
-        sigmoid(z)
+impl Optimizer for Sgd {
+    fn step(&mut self) {
+        self.t += 1.0;
+        self.eta = (1.0 / ((self.alpha) * (self.optimal_init + self.t - 1.0))) as f32;
+        self.w_scale *= 1.0 - ((1.0 - self.l1_ratio) as f32 * self.eta * self.l2_ratio as f32);
+        self.u += self.eta * self.l1_ratio as f32 * self.alpha as f32;
     }
 
-    pub fn predict(&self, features: &Features) -> f32 {
-        let proba = self.predict_proba_sample(features);
-        if proba >= 0.5 { 1.0 } else { 0.0 }
+    fn update_param(&mut self, i: usize, g: f32) {
+        self.parameters[i] += (-self.eta * g) / self.w_scale;
     }
 
-    pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
-    where
-        I: IntoIterator,
-        I::Item: AsRef<Features>,
-    {
-        test.into_iter()
-            .map(|features| self.predict(features.as_ref()))
-            .collect()
+    fn update_bias(&mut self, g: f32) {
+        self.bias += -self.eta * g;
+        self.maybe_rescale();
+        self.apply_l1_penalty();
     }
 
-    pub fn feature_builder(&self) -> FeatureBuilder {
-        FeatureBuilder {
-            features_mask: (self.weights.len() - 1) as u32,
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        self.parameters[idx] * self.w_scale
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        self.bias
+    }
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.parameters[idx]
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self.parameters.iter().map(|w| w * self.w_scale).collect(),
+            bias: self.bias,
         }
     }
 
-    pub fn is_active(&self) -> bool {
-        !self.weights.is_empty()
+    fn num_parameters(&self) -> usize {
+        self.parameters.len()
     }
 }
 
-#[inline(always)]
-fn gradient(y: f32, p: f32) -> f32 {
-    if p > -16.0 {
-        let exp_tmp = (-p).exp();
-        ((1.0 - y) - y * exp_tmp) / (1.0 + exp_tmp)
-    } else {
-        p.exp() - y
-    }
-}
-
-#[inline(always)]
-fn sigmoid(z: f32) -> f32 {
-    if z >= 0.0 {
-        1.0 / (1.0 + (-z).exp())
-    } else {
-        let exp_z = z.exp();
-        exp_z / (1.0 + exp_z)
-    }
-}
-
-/*#[inline(always)]
-fn loss(y: f32, p: f32) -> f32 {
-    log1pexp(p) - y * p
-}
-
-#[inline(always)]
-fn log1pexp(x: f32) -> f32 {
-    if x <= -16.0 {
-        x.exp()
-    } else if x <= 16.0 {
-        (1.0 + x.exp()).ln()
-    } else {
-        x
-    }
-}*/
-
 #[cfg(test)]
 pub mod tests {
-    use crate::classifier::{
-        feature::{Feature, Sample},
-        sgd::TextClassifier,
-    };
     use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
     use std::{
         collections::HashMap,
@@ -169,6 +130,178 @@ pub mod tests {
         time::Instant,
     };
 
+    use crate::classifier::{
+        Optimizer,
+        adam::Adam,
+        feature::{
+            CcfhFeature, CcfhFeatureBuilder, FeatureBuilder, FhFeature, FhFeatureBuilder, Sample,
+            UnprocessedFeature,
+        },
+        ftrl::Ftrl,
+        train::{CcfhTrainer, FhTrainer},
+    };
+
+    #[test]
+    fn text_classifier() {
+        let reader = BufReader::new(
+            File::open("/Users/me/code/playground/phishing_email.csv")
+                .expect("Could not open file"),
+        );
+        let mut samples = Vec::with_capacity(1024);
+
+        let time = Instant::now();
+
+        for line in reader.lines().skip(1) {
+            let line = line.unwrap();
+            let (text, class) = line.trim().rsplit_once(',').unwrap();
+            //let (class, text) = line.trim().split_once(',').unwrap();
+            let text = text.trim_start_matches('"').trim_end_matches('"');
+
+            samples.push((text.to_string(), class == "1"));
+        }
+
+        println!("Loaded {} samples in {:?}", samples.len(), time.elapsed());
+
+        samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+        let (train_samples, test_samples) = train_test_split(&samples, 0.2);
+
+        println!(
+            "Training samples: {}, Testing samples: {}",
+            train_samples.len(),
+            test_samples.len()
+        );
+
+        const FH_SIZE: usize = 16;
+        const CCFH_SIZE: usize = FH_SIZE - 2;
+        let mut rng = StdRng::seed_from_u64(42);
+
+        let fh_builder = FhFeatureBuilder {
+            weight_mask: (1 << FH_SIZE) - 1,
+        };
+        let mut fh_train_samples = build_fh_samples(train_samples.as_slice(), &fh_builder);
+        fh_train_samples.shuffle(&mut rng);
+        let fh_test_samples = build_fh_samples(test_samples.as_slice(), &fh_builder);
+        let ccfh_builder = CcfhFeatureBuilder {
+            weight_mask: (1 << FH_SIZE) - 1,
+            indicator_mask: (1 << CCFH_SIZE) - 1,
+        };
+        let mut ccfh_train_samples = build_ccfh_samples(train_samples.as_slice(), &ccfh_builder);
+        ccfh_train_samples.shuffle(&mut rng);
+        let ccfh_test_samples = build_ccfh_samples(test_samples.as_slice(), &ccfh_builder);
+
+        fh_model_stats(
+            "FTRL",
+            FhTrainer::new(Ftrl::new(1 << FH_SIZE)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "FTRL + FTRL",
+            CcfhTrainer::new(
+                Ftrl::new(1 << FH_SIZE),
+                Ftrl::new(1 << CCFH_SIZE).with_initial_weights(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );
+
+        fh_model_stats(
+            "Adam",
+            FhTrainer::new(Adam::new(1 << FH_SIZE, 0.01)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "Adam + Adam",
+            CcfhTrainer::new(
+                Adam::new(1 << FH_SIZE, 0.01),
+                Adam::new(1 << CCFH_SIZE, 0.01).with_initial_weights(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );
+
+        /*fh_model_stats(
+            "SGD",
+            FhTrainer::new(Sgd::new(1 << FH_SIZE, 0.0001, 0.0, 0.0001)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "FTRL + SGD",
+            CcfhTrainer::new(
+                Ftrl::new(1 << FH_SIZE),
+                Sgd::new(1 << CCFH_SIZE, 0.0001, 0.0, 0.0001).with_initial_parameters(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );*/
+    }
+
+    fn fh_model_stats(
+        name: &str,
+        mut model: FhTrainer<impl Optimizer>,
+        train_samples: &[Sample<FhFeature>],
+        test_samples: &[Sample<FhFeature>],
+    ) {
+        print!("⏳ Training {}... ", name);
+        let time = Instant::now();
+        let mut batch = Vec::new();
+        for sample in train_samples {
+            batch.push(sample);
+            if batch.len() == 128 {
+                model.fit(&mut batch, 5);
+                batch.clear();
+            }
+        }
+        if !batch.is_empty() {
+            model.fit(&mut batch, 5);
+        }
+        println!(" trained in {:?}", time.elapsed());
+        let y_pred = model
+            .build_classifier()
+            .predict_batch(test_samples.iter().map(|s| &s.features));
+        let y_train: Vec<f32> = test_samples.iter().map(|s| s.class).collect();
+        println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred));
+        println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0));
+        println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0));
+        println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0));
+    }
+
+    fn ccfh_model_stats(
+        name: &str,
+        mut model: CcfhTrainer<impl Optimizer, impl Optimizer>,
+        train_samples: &[Sample<CcfhFeature>],
+        test_samples: &[Sample<CcfhFeature>],
+    ) {
+        print!("⏳ Training {}... ", name);
+        let time = Instant::now();
+        let mut batch = Vec::new();
+        for sample in train_samples {
+            batch.push(sample);
+            if batch.len() == 128 {
+                model.fit(&mut batch, 5);
+                batch.clear();
+            }
+        }
+        if !batch.is_empty() {
+            model.fit(&mut batch, 5);
+        }
+        println!(" trained in {:?}", time.elapsed());
+        let y_pred = model
+            .build_classifier()
+            .predict_batch(test_samples.iter().map(|s| &s.features));
+        let y_train: Vec<f32> = test_samples.iter().map(|s| s.class).collect();
+        println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred));
+        println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0));
+        println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0));
+        println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0));
+    }
+
     fn accuracy_score(y_true: &[f32], y_pred: &[f32]) -> f32 {
         y_true
             .iter()
@@ -231,15 +364,19 @@ pub mod tests {
         }
     }
 
-    fn train_test_split(data: &[Sample], test_size: f32) -> (Vec<&Sample>, Vec<&Sample>) {
-        let mut class_0: Vec<&Sample> = Vec::new();
-        let mut class_1: Vec<&Sample> = Vec::new();
+    #[allow(clippy::type_complexity)]
+    pub fn train_test_split(
+        data: &[(String, bool)],
+        test_size: f32,
+    ) -> (Vec<(&String, bool)>, Vec<(&String, bool)>) {
+        let mut class_0: Vec<(&String, bool)> = Vec::new();
+        let mut class_1: Vec<(&String, bool)> = Vec::new();
 
-        for sample in data {
-            if sample.class == 0.0 {
-                class_0.push(sample);
+        for (sample, class) in data {
+            if !*class {
+                class_0.push((sample, *class));
             } else {
-                class_1.push(sample);
+                class_1.push((sample, *class));
             }
         }
 
@@ -260,7 +397,49 @@ pub mod tests {
         (train, test)
     }
 
-    impl Feature for String {
+    pub fn build_fh_samples(
+        data: &[(&String, bool)],
+        builder: &FhFeatureBuilder,
+    ) -> Vec<Sample<FhFeature>> {
+        let mut samples = Vec::with_capacity(data.len());
+
+        for (text, class) in data {
+            let mut sample: HashMap<String, f32> = HashMap::new();
+            for word in text.split_whitespace() {
+                *sample.entry(word.to_string()).or_default() += 1.0;
+            }
+            builder.scale(&mut sample);
+            samples.push(Sample {
+                features: builder.build(&sample, 12345.into()),
+                class: if *class { 1.0 } else { 0.0 },
+            });
+        }
+
+        samples
+    }
+
+    pub fn build_ccfh_samples(
+        data: &[(&String, bool)],
+        builder: &CcfhFeatureBuilder,
+    ) -> Vec<Sample<CcfhFeature>> {
+        let mut samples = Vec::with_capacity(data.len());
+
+        for (text, class) in data {
+            let mut sample: HashMap<String, f32> = HashMap::new();
+            for word in text.split_whitespace() {
+                *sample.entry(word.to_string()).or_default() += 1.0;
+            }
+            builder.scale(&mut sample);
+            samples.push(Sample {
+                features: builder.build(&sample, 12345.into()),
+                class: if *class { 1.0 } else { 0.0 },
+            });
+        }
+
+        samples
+    }
+
+    impl UnprocessedFeature for String {
         fn prefix(&self) -> u16 {
             0
         }
@@ -269,61 +448,4 @@ pub mod tests {
             self.as_bytes()
         }
     }
-
-    #[test]
-    fn sgd_classifier() {
-        let reader = BufReader::new(
-            File::open("/Users/me/code/playground/phishing_email.csv")
-                .expect("Could not open file"),
-        );
-        let mut samples = Vec::with_capacity(1024);
-
-        let mut model = TextClassifier::new(1 << 20);
-        let builder = model.feature_builder();
-
-        let time = Instant::now();
-
-        for line in reader.lines().skip(1) {
-            let line = line.unwrap();
-            let (text, class) = line.trim().rsplit_once(',').unwrap();
-            let text = text.trim_start_matches('"').trim_end_matches('"');
-
-            let mut sample: HashMap<String, f32> = HashMap::new();
-            for word in text.split_whitespace() {
-                *sample.entry(word.to_string()).or_default() += 1.0;
-            }
-            builder.scale(&mut sample);
-            samples.push(Sample {
-                features: builder.build(&sample, None),
-                class: class
-                    .parse()
-                    .unwrap_or_else(|_| panic!("Invalid class value: {line}")),
-            });
-        }
-
-        println!("Loaded {} samples in {:?}", samples.len(), time.elapsed());
-
-        samples.shuffle(&mut StdRng::seed_from_u64(42));
-
-        let (mut train_samples, test_samples) = train_test_split(&samples, 0.2);
-
-        println!(
-            "Training samples: {}, Testing samples: {}",
-            train_samples.len(),
-            test_samples.len()
-        );
-
-        println!("Training SGD Classifier...");
-        let time = Instant::now();
-        model.fit(&mut train_samples, 1000, 0.0001);
-        println!("SGD Classifier trained in {:?}", time.elapsed());
-
-        let y_pred = model.predict_batch(test_samples.iter().map(|s| &s.features));
-        let y_train: Vec<f32> = test_samples.iter().map(|s| s.class).collect();
-
-        println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred));
-        println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0));
-        println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0));
-        println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0));
-    }
 }

crates/nlp/src/classifier/train.rs

@@ -0,0 +1,157 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{
+    MAX_DLOSS, Optimizer,
+    feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder, Sample},
+    gradient,
+    model::{CcfhClassifier, FhClassifier},
+};
+use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct FhTrainer<T: Optimizer> {
+    pub optimizer: T,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct CcfhTrainer<W: Optimizer, I: Optimizer> {
+    pub w_optimizer: W,
+    pub i_optimizer: I,
+}
+
+impl<T: Optimizer> FhTrainer<T> {
+    pub fn new(optimizer: T) -> Self {
+        FhTrainer { optimizer }
+    }
+
+    pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<FhFeature>>], num_epochs: usize) {
+        for _ in 0..num_epochs {
+            samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+            for sample in samples.iter() {
+                let sample = sample.as_ref();
+                let mut dot: f32 = 0.0;
+                for f in &sample.features {
+                    dot += self.optimizer.get_param(f.idx) * f.weight;
+                }
+                let p = dot + self.optimizer.get_bias();
+                let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
+
+                self.optimizer.step();
+
+                for f in &sample.features {
+                    self.optimizer.update_param(f.idx, dloss * f.weight);
+                }
+
+                self.optimizer.update_bias(dloss);
+            }
+        }
+    }
+
+    pub fn feature_builder(&self) -> FhFeatureBuilder {
+        FhFeatureBuilder {
+            weight_mask: (self.optimizer.num_parameters() - 1) as u64,
+        }
+    }
+
+    pub fn build_classifier(&self) -> FhClassifier {
+        self.optimizer.build_classifier()
+    }
+
+    pub fn optimizer(&self) -> &T {
+        &self.optimizer
+    }
+
+    pub fn optimizer_mut(&mut self) -> &mut T {
+        &mut self.optimizer
+    }
+}
+
+impl<W: Optimizer, I: Optimizer> CcfhTrainer<W, I> {
+    pub fn new(w_optimizer: W, i_optimizer: I) -> Self {
+        CcfhTrainer {
+            w_optimizer,
+            i_optimizer,
+        }
+    }
+
+    pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<CcfhFeature>>], num_epochs: usize) {
+        for _ in 0..num_epochs {
+            samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+            for sample in samples.iter() {
+                let sample = sample.as_ref();
+                let mut dot: f32 = 0.0;
+                for f in &sample.features {
+                    let q = self.i_optimizer.get_param(f.idx_i);
+                    let v1 = self.w_optimizer.get_param(f.idx_w1);
+                    let v2 = self.w_optimizer.get_param(f.idx_w2);
+                    dot += (q * v1 + (1.0 - q) * v2) * f.weight;
+                }
+                let p = dot + self.w_optimizer.get_bias();
+                let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
+
+                self.w_optimizer.step();
+                self.i_optimizer.step();
+
+                for f in &sample.features {
+                    let q = self.i_optimizer.get_param(f.idx_i);
+                    let v1 = self.w_optimizer.get_param(f.idx_w1);
+                    let v2 = self.w_optimizer.get_param(f.idx_w2);
+
+                    // Update weights
+                    let d_v1 = f.weight * q;
+                    let d_v2 = f.weight * (1.0 - q);
+                    self.w_optimizer.update_param(f.idx_w1, dloss * d_v1);
+                    self.w_optimizer.update_param(f.idx_w2, dloss * d_v2);
+
+                    // Update indicator
+                    let d_q = (v1 - v2) * f.weight;
+                    self.i_optimizer.update_param(f.idx_i, dloss * d_q);
+                    let fi = self.i_optimizer.get_param_mut(f.idx_i);
+                    *fi = fi.clamp(0.0, 1.0);
+                }
+
+                self.w_optimizer.update_bias(dloss);
+            }
+        }
+    }
+
+    pub fn feature_builder(&self) -> CcfhFeatureBuilder {
+        CcfhFeatureBuilder {
+            weight_mask: (self.w_optimizer.num_parameters() - 1) as u64,
+            indicator_mask: (self.i_optimizer.num_parameters() - 1) as u64,
+        }
+    }
+
+    pub fn build_classifier(&self) -> CcfhClassifier {
+        let w_classifier = self.w_optimizer.build_classifier();
+        let i_classifier = self.i_optimizer.build_classifier();
+
+        CcfhClassifier {
+            parameters: w_classifier.parameters,
+            indicators: i_classifier.parameters,
+            bias: w_classifier.bias,
+        }
+    }
+
+    pub fn w_optimizer(&self) -> &W {
+        &self.w_optimizer
+    }
+
+    pub fn w_optimizer_mut(&mut self) -> &mut W {
+        &mut self.w_optimizer
+    }
+
+    pub fn i_optimizer(&self) -> &I {
+        &self.i_optimizer
+    }
+
+    pub fn i_optimizer_mut(&mut self) -> &mut I {
+        &mut self.i_optimizer
+    }
+}