apply recency weighting to evaluation

Author: L-M-Sherlock

src/batch_shuffle.rs

@@ -110,7 +110,7 @@ mod tests {
     use super::*;
     use crate::{
         convertor_tests::anki21_sample_file_converted_to_fsrs,
-        dataset::{prepare_training_data, simple_weighted_fsrs_items},
+        dataset::{constant_weighted_fsrs_items, prepare_training_data},
     };
 
     #[test]
@@ -120,7 +120,7 @@ mod tests {
             .sorted_by_cached_key(|item| item.reviews.len())
             .collect();
         let (_pre_train_set, train_set) = prepare_training_data(train_set);
-        let dataset = FSRSDataset::from(simple_weighted_fsrs_items(train_set));
+        let dataset = FSRSDataset::from(constant_weighted_fsrs_items(train_set));
         let batch_size = 512;
         let seed = 114514;
         let device = NdArrayDevice::Cpu;

src/convertor_tests.rs

@@ -1,5 +1,5 @@
 use crate::convertor_tests::RevlogReviewKind::*;
-use crate::dataset::{simple_weighted_fsrs_items, FSRSBatcher};
+use crate::dataset::{constant_weighted_fsrs_items, FSRSBatcher};
 use crate::dataset::{FSRSItem, FSRSReview};
 use crate::optimal_retention::{RevlogEntry, RevlogReviewKind};
 use crate::test_helpers::NdArrayAutodiff;
@@ -388,7 +388,7 @@ fn conversion_works() {
         ]
     );
 
-    let mut weighted_fsrs_items = simple_weighted_fsrs_items(fsrs_items);
+    let mut weighted_fsrs_items = constant_weighted_fsrs_items(fsrs_items);
 
     let device = NdArrayDevice::Cpu;
     let batcher = FSRSBatcher::<NdArrayAutodiff>::new(device);

src/dataset.rs

@@ -271,7 +271,8 @@ pub(crate) fn sort_items_by_review_length(items: Vec<WeightedFSRSItem>) -> Vec<W
     items
 }
 
-pub(crate) fn simple_weighted_fsrs_items(items: Vec<FSRSItem>) -> Vec<WeightedFSRSItem> {
+#[cfg(test)]
+pub(crate) fn constant_weighted_fsrs_items(items: Vec<FSRSItem>) -> Vec<WeightedFSRSItem> {
     items
         .into_iter()
         .map(|item| WeightedFSRSItem { weight: 1.0, item })
@@ -300,7 +301,7 @@ mod tests {
     fn from_anki() {
         use burn::data::dataloader::Dataset;
 
-        let dataset = FSRSDataset::from(sort_items_by_review_length(simple_weighted_fsrs_items(
+        let dataset = FSRSDataset::from(sort_items_by_review_length(constant_weighted_fsrs_items(
             anki21_sample_file_converted_to_fsrs(),
         )));
         assert_eq!(

src/inference.rs

@@ -6,8 +6,7 @@ use burn::nn::loss::Reduction;
 use burn::tensor::{Data, Shape, Tensor};
 use burn::{data::dataloader::batcher::Batcher, tensor::backend::Backend};
 
-use crate::dataset::FSRSBatcher;
-use crate::dataset::{simple_weighted_fsrs_items, FSRSBatch};
+use crate::dataset::{recency_weighted_fsrs_items, FSRSBatch, FSRSBatcher};
 use crate::error::Result;
 use crate::model::Model;
 use crate::training::BCELoss;
@@ -210,7 +209,7 @@ impl<B: Backend> FSRS<B> {
         if items.is_empty() {
             return Err(FSRSError::NotEnoughData);
         }
-        let items = simple_weighted_fsrs_items(items);
+        let items = recency_weighted_fsrs_items(items);
         let batcher = FSRSBatcher::new(self.device());
         let mut all_retention = vec![];
         let mut all_labels = vec![];
@@ -232,10 +231,10 @@ impl<B: Backend> FSRS<B> {
             all_weights.push(batch.weights);
             izip!(chunk, pred, true_val).for_each(|(item, p, y)| {
                 let bin = item.item.r_matrix_index();
-                let (pred, real, count) = r_matrix.entry(bin).or_insert((0.0, 0.0, 0.0));
+                let (pred, real, weight) = r_matrix.entry(bin).or_insert((0.0, 0.0, 0.0));
                 *pred += p;
                 *real += y;
-                *count += 1.0;
+                *weight += item.weight;
             });
             progress_info.current += chunk.len();
             if !progress(progress_info) {
@@ -244,13 +243,13 @@ impl<B: Backend> FSRS<B> {
         }
         let rmse = (r_matrix
             .values()
-            .map(|(pred, real, count)| {
-                let pred = pred / count;
-                let real = real / count;
-                (pred - real).powi(2) * count
+            .map(|(pred, real, weight)| {
+                let pred = pred / weight;
+                let real = real / weight;
+                (pred - real).powi(2) * weight
             })
             .sum::<f32>()
-            / r_matrix.values().map(|(_, _, count)| count).sum::<f32>())
+            / r_matrix.values().map(|(_, _, weight)| weight).sum::<f32>())
         .sqrt();
         let all_retention = Tensor::cat(all_retention, 0);
         let all_labels = Tensor::cat(all_labels, 0).float();
@@ -282,7 +281,7 @@ impl<B: Backend> FSRS<B> {
         if items.is_empty() {
             return Err(FSRSError::NotEnoughData);
         }
-        let items = simple_weighted_fsrs_items(items);
+        let items = recency_weighted_fsrs_items(items);
         let batcher = FSRSBatcher::new(self.device());
         let mut all_predictions_self = vec![];
         let mut all_predictions_other = vec![];
@@ -499,17 +498,17 @@ mod tests {
         ]))?;
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
-        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.211007, 0.037216]);
+        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.212817, 0.034676]);
 
         let fsrs = FSRS::new(Some(&[]))?;
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
-        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.216286, 0.038692]);
+        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.217251, 0.036590]);
 
         let fsrs = FSRS::new(Some(PARAMETERS))?;
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
-        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.203049, 0.027558]);
+        assert_approx_eq([metrics.log_loss, metrics.rmse_bins], [0.203552, 0.025646]);
 
         let (self_by_other, other_by_self) = fsrs
             .universal_metrics(items.clone(), &DEFAULT_PARAMETERS, |_| true)