Add 'zipf' option to specify the exponent of the zipf distribution

Author: danieldk

finalfrontier-utils/src/bin/ff-train.rs

@@ -92,6 +92,7 @@ static MAXN: &str = "maxn";
 static MODEL: &str = "model";
 static NS: &str = "ns";
 static THREADS: &str = "threads";
+static ZIPF_EXPONENT: &str = "zipf";
 
 // Argument constants
 static CORPUS: &str = "CORPUS";
@@ -145,6 +146,13 @@ fn config_from_matches<'a>(matches: &ArgMatches<'a>) -> Config {
                 .or_exit("Cannot parse number of negative samples", 1)
         })
         .unwrap_or(5);
+    let zipf_exponent = matches
+        .value_of(ZIPF_EXPONENT)
+        .map(|v| {
+            v.parse()
+                .or_exit("Cannot parse exponent zipf distribution", 1)
+        })
+        .unwrap_or(0.5);
 
     Config {
         context_size,
@@ -159,6 +167,7 @@ fn config_from_matches<'a>(matches: &ArgMatches<'a>) -> Config {
         buckets_exp,
         negative_samples,
         lr,
+        zipf_exponent,
     }
 }
 
@@ -249,6 +258,13 @@ fn parse_args() -> ArgMatches<'static> {
                 .help("Number of threads (default: logical_cpus / 2)")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name(ZIPF_EXPONENT)
+                .long("zipf")
+                .value_name("EXP")
+                .help("Exponent Zipf distribution for negative sampling (default: 0.5)")
+                .takes_value(true),
+        )
         .arg(
             Arg::with_name(CORPUS)
                 .help("Tokenized corpus")

finalfrontier/src/config.rs

@@ -96,4 +96,9 @@ pub struct Config {
 
     /// The initial learning rate.
     pub lr: f32,
+
+    /// Exponent in zipfian distribution.
+    ///
+    /// This is s in *f(k) = 1 / (k^s H_{N, s})*.
+    pub zipf_exponent: f64,
 }

finalfrontier/src/model.rs

@@ -1,5 +1,6 @@
 //! Embedding prediction model.
 
+use std::f64;
 use std::fs::File;
 use std::io::{BufReader, Read, Seek, SeekFrom, Write};
 use std::iter::Enumerate;
@@ -223,6 +224,7 @@ where
         buckets_exp,
         negative_samples,
         lr,
+        zipf_exponent: f64::NAN,
     })
 }
 
@@ -350,6 +352,7 @@ mod tests {
         min_n: 3,
         model: ModelType::SkipGram,
         negative_samples: 5,
+        zipf_exponent: 0.5,
     };
 
     #[test]

finalfrontier/src/sampling.rs

@@ -10,7 +10,7 @@ pub trait RangeGenerator: Iterator<Item = usize> {
 /// Exponent to use for the Zipf's distribution.
 ///
 /// This is the exponent s in f(k) = 1 / (k^s H_{N, s})
-const ZIPF_RANGE_GENERATOR_EXPONENT: f64 = 0.5;
+const ZIPF_RANGE_GENERATOR_DEFAULT_EXPONENT: f64 = 0.5;
 
 /// An iterator that draws from *[0, n)* with integer weights.
 ///
@@ -86,6 +86,7 @@ where
 /// is proportional to its frequency.
 pub struct ZipfRangeGenerator<R> {
     upper_bound: usize,
+    exponent: f64,
     rng: R,
     dist: ZipfDistribution,
 }
@@ -97,8 +98,9 @@ where
     fn clone(&self) -> Self {
         ZipfRangeGenerator {
             upper_bound: self.upper_bound,
+            exponent: self.exponent,
             rng: self.rng.clone(),
-            dist: ZipfDistribution::new(self.upper_bound, ZIPF_RANGE_GENERATOR_EXPONENT).unwrap(),
+            dist: ZipfDistribution::new(self.upper_bound, self.exponent).unwrap(),
         }
     }
 }
@@ -107,11 +109,17 @@ impl<R> ZipfRangeGenerator<R>
 where
     R: Rng,
 {
+    #[allow(dead_code)]
     pub fn new(rng: R, upper: usize) -> Self {
+        Self::new_with_exponent(rng, upper, ZIPF_RANGE_GENERATOR_DEFAULT_EXPONENT)
+    }
+
+    pub fn new_with_exponent(rng: R, upper_bound: usize, exponent: f64) -> Self {
         ZipfRangeGenerator {
-            upper_bound: upper,
+            upper_bound,
+            exponent,
             rng,
-            dist: ZipfDistribution::new(upper, ZIPF_RANGE_GENERATOR_EXPONENT).unwrap(),
+            dist: ZipfDistribution::new(upper_bound, exponent).unwrap(),
         }
     }
 }

finalfrontier/src/sgd.rs

@@ -59,7 +59,11 @@ where
 
         let range_gen = BandedRangeGenerator::new(
             reseed_on_clone.clone(),
-            ZipfRangeGenerator::new(reseed_on_clone.clone(), model.vocab().len()),
+            ZipfRangeGenerator::new_with_exponent(
+                reseed_on_clone.clone(),
+                model.vocab().len(),
+                model.config().zipf_exponent,
+            ),
             band_size as usize,
         );
 

finalfrontier/src/train_model.rs

@@ -191,6 +191,7 @@ mod tests {
         min_n: 3,
         model: ModelType::SkipGram,
         negative_samples: 5,
+        zipf_exponent: 0.5,
     };
 
     #[test]

finalfrontier/src/vocab.rs

@@ -273,6 +273,7 @@ mod tests {
         min_n: 3,
         model: ModelType::SkipGram,
         negative_samples: 5,
+        zipf_exponent: 0.5,
     };
 
     #[test]

man/ff-train.1.md

@@ -87,7 +87,12 @@ OPTIONS
 
 :   The number of thread to use during training for parallelization. The
     default is to use half of the logical CPUs of the machine.
-	
+
+`--zipf` *EXP*
+
+:   Exponent *s* used in the Zipf distribution `p(k) = 1 / (k^s H_N)` for
+    negative sampling. Default: 0.5
+
 EXAMPLES
 ========