word2vec学习 spark版

参考资料：

http://ir.dlut.edu.cn/NewsShow.aspx?ID=291

http://www.douban.com/note/298095260/

http://machinelearning.wustl.edu/mlpapers/paper_files/BengioDVJ03.pdf

https://code.google.com/p/word2vec/

https://spark.apache.org/docs/latest/mllib-feature-extraction.html#word2vec

word2vec是NLP领域的重要算法，它的功能是将word用K维的dense vector来表达，训练集是语料库，不含标点，以空格断句。因此可以看作是种特征处理方法。

主要优点：

• 加法操作。
• 高效。单机可处理1小时2千万词。

google的开源版本比较权威，地址（ http://word2vec.googlecode.com/svn/trunk/ ），不过我以spark版本学习的。

I.背景知识

Distributed representation，word的特征表达方式，通过训练将每个词映射成 K 维实数向量(K 一般为模型中的超参数)，通过词之间的距离(比如 cosine 相似度、欧氏距离等)来判断它们之间的语义相似度。

语言模型：n-gram等。

II.模型

0.word window构成context，对于一个单词i，以$u_i$表示，它作为别的单词的context时用$v_i$表示（也即它作为context的表示是不同的）。只有word window内的word才被认为是context，并且是顺序无关的。

1.概率模型为\[ P=\sum lot p(u_i) ，\]i表示位置（或单词），也即各单词出现概率的累积函数。

2.以skip gram为例（CBOW条件概率反过来），则位置i的单词出现概率为

\[ p(u_i)=\sum_{-c\leq j\leq c,j\neq 0} p(v_{i+j}|u_{i}) \]

表示位置i只和其context有关。

3.条件概率$p(v_{i+j}|u_i)$ 通过softmax实现K维向量到概率的转化表达。

III.优化

最开始使用神经网络，后来用层次softmax等来降低时间复杂度。还用了很多trick，比如ExpTable。

a) 删除隐藏层

b) 使用Hierarchical softmax 或negative sampling

c) 去除小于minCount的词

d)预先计算ExpTable

e) 根据一下公式算出每个词被选出的概率，如果选出来则不予更新。此方法可以节省时间而且可以提高非频繁词的准确度。

\[ prob(w)=1-\large(\sqrt{\frac{t}{f(w)}}+\frac{t}{f(w)}\large)  \]   其中$t$为设定好的阈值，$f(w)$ 为$w$出现的频率。

f) 选取邻近词的窗口大小不固定。有利于更加偏重于离自己近的词进行更新。

g)  多线程，无需考虑互斥。

IV.spark源码分析

 /**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

 package org.apache.spark.mllib.feature

 import java.lang.{Iterable => JavaIterable}

 import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.apache.spark.Logging
 import org.apache.spark.SparkContext._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.util.Utils
 import org.apache.spark.util.random.XORShiftRandom
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer

 /**
  *  Entry in vocabulary
  */
 private case class VocabWord(
                               var word: String,
                               var cn: Int,
                               var point: Array[Int],
                               var code: Array[Int],
                               var codeLen:Int
                               )

 /**
  * :: Experimental ::
  * Word2Vec creates vector representation of words in a text corpus.
  * The algorithm first constructs a vocabulary from the corpus
  * and then learns vector representation of words in the vocabulary.
  * The vector representation can be used as features in
  * natural language processing and machine learning algorithms.
  *
  * We used skip-gram model in our implementation and hierarchical softmax
  * method to train the model. The variable names in the implementation
  * matches the original C implementation.
  *
  * For original C implementation, see https://code.google.com/p/word2vec/
  * For research papers, see
  * Efficient Estimation of Word Representations in Vector Space
  * and
  * Distributed Representations of Words and Phrases and their Compositionality.
  */
 @Experimental
 class Word2VectorEX extends Serializable with Logging {

   private var vectorSize =
   private var startingAlpha = 0.025
   private var numPartitions =
   private var numIterations =
   private var seed = Utils.random.nextLong()

   /**
    * Sets vector size (default: 100).
    */
   def setVectorSize(vectorSize: Int): this.type = {
     this.vectorSize = vectorSize
     this
   }

   /**
    * Sets initial learning rate (default: 0.025).
    */
   def setLearningRate(learningRate: Double): this.type = {
     this.startingAlpha = learningRate
     this
   }

   /**
    * Sets number of partitions (default: 1). Use a small number for accuracy.
    */
   def setNumPartitions(numPartitions: Int): this.type = {
     require(numPartitions > , s"numPartitions must be greater than 0 but got $numPartitions")
     this.numPartitions = numPartitions
     this
   }

   /**
    * Sets number of iterations (default: 1), which should be smaller than or equal to number of
    * partitions.
    */
   def setNumIterations(numIterations: Int): this.type = {
     this.numIterations = numIterations
     this
   }

   /**
    * Sets random seed (default: a random long integer).
    */
   def setSeed(seed: Long): this.type = {
     this.seed = seed
     this
   }

   private val EXP_TABLE_SIZE =
   private val MAX_EXP =
   private val MAX_CODE_LENGTH =
   private val MAX_SENTENCE_LENGTH = 

   /** context words from [-window, window] */
   private val window =             //context 范围限定

   /** minimum frequency to consider a vocabulary word */
   private val minCount =            //过滤单词阈值

   private var trainWordsCount =           //语料库总共词量（计重复出现）
   private var vocabSize =                  //词表内单词总数
   private var vocab: Array[VocabWord] = null      //词表
   private var vocabHash = mutable.HashMap.empty[String, Int]      //词表反查索引

   private def learnVocab(words: RDD[String]): Unit = {           //构造词表，统计更新上面四个量
     vocab = words.map(w => (w, ))
       .reduceByKey(_ + _)
       .map(x => VocabWord(
       x._1,
       x._2,
       new Array[Int](MAX_CODE_LENGTH),
       new Array[Int](MAX_CODE_LENGTH),
       ))
       .filter(_.cn >= minCount)
       .collect()
       .sortWith((a, b) => a.cn > b.cn)

     vocabSize = vocab.length
     var a =
     while (a < vocabSize) {
       vocabHash += vocab(a).word -> a
       trainWordsCount += vocab(a).cn
       a +=
     }
     logInfo("trainWordsCount = " + trainWordsCount)
   }

   private def createExpTable(): Array[Float] = {        //指数运算查表
     val expTable = new Array[Float](EXP_TABLE_SIZE)
     var i =
     while (i < EXP_TABLE_SIZE) {
       val tmp = math.exp((2.0 * i / EXP_TABLE_SIZE - 1.0) * MAX_EXP)
       expTable(i) = (tmp / (tmp + 1.0)).toFloat
       i +=
     }
     expTable
   }

   private def createBinaryTree(): Unit = {
     val count = new Array[Long](vocabSize *  + )
     val binary = new Array[Int](vocabSize *  + )
     val parentNode = new Array[Int](vocabSize *  + )
     val code = new Array[Int](MAX_CODE_LENGTH)
     val point = new Array[Int](MAX_CODE_LENGTH)
     var a =
     while (a < vocabSize) {
       count(a) = vocab(a).cn
       a +=
     }
     while (a <  * vocabSize) {
       count(a) = 1e9.toInt
       a +=
     }
     var pos1 = vocabSize -
     var pos2 = vocabSize

     var min1i =
     var min2i = 

     a =
     while (a < vocabSize - ) {
       if (pos1 >= ) {
         if (count(pos1) < count(pos2)) {
           min1i = pos1
           pos1 -=
         } else {
           min1i = pos2
           pos2 +=
         }
       } else {
         min1i = pos2
         pos2 +=
       }
       if (pos1 >= ) {
         if (count(pos1) < count(pos2)) {
           min2i = pos1
           pos1 -=
         } else {
           min2i = pos2
           pos2 +=
         }
       } else {
         min2i = pos2
         pos2 +=
       }
       count(vocabSize + a) = count(min1i) + count(min2i)
       parentNode(min1i) = vocabSize + a
       parentNode(min2i) = vocabSize + a
       binary(min2i) =
       a +=
     }
     // Now assign binary code to each vocabulary word
     var i =
     a =
     while (a < vocabSize) {
       var b = a
       i =
       while (b != vocabSize *  - ) {
         code(i) = binary(b)
         point(i) = b
         i +=
         b = parentNode(b)
       }
       vocab(a).codeLen = i
       vocab(a).point() = vocabSize -
       b =
       while (b < i) {
         vocab(a).code(i - b - ) = code(b)
         vocab(a).point(i - b) = point(b) - vocabSize
         b +=
       }
       a +=
     }
   }

   /**
    * Computes the vector representation of each word in vocabulary.
    * @param dataset an RDD of words
    * @return a Word2VecModel
    */
   def fit[S <: Iterable[String]](dataset: RDD[S]): Word2VectorModel = {

     val words = dataset.flatMap(x => x)       //拉成词序列，句话断点通过Iterable来表征

     learnVocab(words)        //学习词库

     createBinaryTree()

     val sc = dataset.context

     val expTable = sc.broadcast(createExpTable())
     val bcVocab = sc.broadcast(vocab)
     val bcVocabHash = sc.broadcast(vocabHash)

     val sentences: RDD[Array[Int]] = words.mapPartitions { iter =>            //按句子划分，单词以Int表征
       new Iterator[Array[Int]] {
         def hasNext: Boolean = iter.hasNext

         def next(): Array[Int] = {
           var sentence = new ArrayBuffer[Int]
           var sentenceLength =
           while (iter.hasNext && sentenceLength < MAX_SENTENCE_LENGTH) {
             val word = bcVocabHash.value.get(iter.next())
             word match {
               case Some(w) =>
                 sentence += w
                 sentenceLength +=
               case None =>
             }
           }
           sentence.toArray
         }
       }
     }

     //Hierarchical Softmax
     val newSentences = sentences.repartition(numPartitions).cache()
     val initRandom = new XORShiftRandom(seed)
     val syn0Global =
       Array.fill[Float](vocabSize * vectorSize)((initRandom.nextFloat() - 0.5f) / vectorSize)
     val syn1Global = new Array[Float](vocabSize * vectorSize)
     var alpha = startingAlpha
     for (k <-  to numIterations) {
       val partial = newSentences.mapPartitionsWithIndex { case (idx, iter) =>
         val random = new XORShiftRandom(seed ^ ((idx + ) << ) ^ ((-k - ) << ))    //随机梯度下降
         val syn0Modify = new Array[Int](vocabSize)
         val syn1Modify = new Array[Int](vocabSize)
         val model = iter.foldLeft((syn0Global, syn1Global, , )) {
           case ((syn0, syn1, lastWordCount, wordCount), sentence) =>
             var lwc = lastWordCount
             var wc = wordCount
             if (wordCount - lastWordCount > ) {
               lwc = wordCount
               // TODO: discount by iteration?
               alpha =
                 startingAlpha * ( - numPartitions * wordCount.toDouble / (trainWordsCount + ))
               if (alpha < startingAlpha * 0.0001) alpha = startingAlpha * 0.0001
               logInfo("wordCount = " + wordCount + ", alpha = " + alpha)
             }
             wc += sentence.size
             var pos =
             while (pos < sentence.size) {
               val word = sentence(pos)
               val b = random.nextInt(window)
               // Train Skip-gram
               var a = b
               while (a < window *  +  - b) {
                 if (a != window) {
                   val c = pos - window + a
                   if (c >=  && c < sentence.size) {
                     val lastWord = sentence(c)
                     val l1 = lastWord * vectorSize
                     val neu1e = new Array[Float](vectorSize)
                     // Hierarchical softmax
                     var d =
                     while (d < bcVocab.value(word).codeLen) {
                       val inner = bcVocab.value(word).point(d)
                       val l2 = inner * vectorSize
                       // Propagate hidden -> output
                       var f = blas.sdot(vectorSize, syn0, l1, , syn1, l2, )
                       if (f > -MAX_EXP && f < MAX_EXP) {
                         val ind = ((f + MAX_EXP) * (EXP_TABLE_SIZE / MAX_EXP / 2.0)).toInt
                         f = expTable.value(ind)
                         val g = (( - bcVocab.value(word).code(d) - f) * alpha).toFloat
                         blas.saxpy(vectorSize, g, syn1, l2, , neu1e, , )
                         blas.saxpy(vectorSize, g, syn0, l1, , syn1, l2, )
                         syn1Modify(inner) +=
                       }
                       d +=
                     }
                     blas.saxpy(vectorSize, 1.0f, neu1e, , , syn0, l1, )
                     syn0Modify(lastWord) +=
                   }
                 }
                 a +=
               }
               pos +=
             }
             (syn0, syn1, lwc, wc)
         }
         val syn0Local = model._1
         val syn1Local = model._2
         // Only output modified vectors.
         Iterator.tabulate(vocabSize) { index =>
           if (syn0Modify(index) > ) {
             Some((index, syn0Local.slice(index * vectorSize, (index + ) * vectorSize)))
           } else {
             None
           }
         }.flatten ++ Iterator.tabulate(vocabSize) { index =>
           if (syn1Modify(index) > ) {
             Some((index + vocabSize, syn1Local.slice(index * vectorSize, (index + ) * vectorSize)))
           } else {
             None
           }
         }.flatten
       }
       val synAgg = partial.reduceByKey { case (v1, v2) =>
         blas.saxpy(vectorSize, 1.0f, v2, , v1, )
         v1
       }.collect()
       var i =
       while (i < synAgg.length) {
         val index = synAgg(i)._1
         if (index < vocabSize) {
           Array.copy(synAgg(i)._2, , syn0Global, index * vectorSize, vectorSize)
         } else {
           Array.copy(synAgg(i)._2, , syn1Global, (index - vocabSize) * vectorSize, vectorSize)
         }
         i +=
       }
     }
     newSentences.unpersist()

     val word2VecMap = mutable.HashMap.empty[String, Array[Float]]
     var i =
     while (i < vocabSize) {
       val word = bcVocab.value(i).word
       val vector = new Array[Float](vectorSize)
       Array.copy(syn0Global, i * vectorSize, vector, , vectorSize)
       word2VecMap += word -> vector
       i +=
     }

     new Word2VectorModel(word2VecMap.toMap)
   }

   /**
    * Computes the vector representation of each word in vocabulary (Java version).
    * @param dataset a JavaRDD of words
    * @return a Word2VecModel
    */
   def fit[S <: JavaIterable[String]](dataset: JavaRDD[S]): Word2VectorModel = {
     fit(dataset.rdd.map(_.asScala))
   }

 }

 /**
  * :: Experimental ::
  * Word2Vec model
  */
 @Experimental
 class Word2VectorModel private[mllib] (
                                      private  val model: Map[String, Array[Float]]) extends Serializable {

   private def cosineSimilarity(v1: Array[Float], v2: Array[Float]): Double = {
     require(v1.length == v2.length, "Vectors should have the same length")
     val n = v1.length
     val norm1 = blas.snrm2(n, v1, )
     val norm2 = blas.snrm2(n, v2, )
     if (norm1 ==  || norm2 == ) return 0.0
     blas.sdot(n, v1, , v2,) / norm1 / norm2
   }

   /**
    * Transforms a word to its vector representation
    * @param word a word
    * @return vector representation of word
    */
   def transform(word: String): Vector = {
     model.get(word) match {
       case Some(vec) =>
         Vectors.dense(vec.map(_.toDouble))
       case None =>
         throw new IllegalStateException(s"$word not in vocabulary")
     }
   }

   /**
    * Find synonyms of a word
    * @param word a word
    * @param num number of synonyms to find
    * @return array of (word, similarity)
    */
   def findSynonyms(word: String, num: Int): Array[(String, Double)] = {
     val vector = transform(word)
     findSynonyms(vector,num)
   }

   /**
    * Find synonyms of the vector representation of a word
    * @param vector vector representation of a word
    * @param num number of synonyms to find
    * @return array of (word, cosineSimilarity)
    */
   def findSynonyms(vector: Vector, num: Int): Array[(String, Double)] = {
     require(num > , "Number of similar words should > 0")
     // TODO: optimize top-k
     val fVector = vector.toArray.map(_.toFloat)
     model.mapValues(vec => cosineSimilarity(fVector, vec))
       .toSeq
       .sortBy(- _._2)
       .take(num + )
       .tail
       .toArray
   }

   def getModel(): Map[String, Array[Float]] = {
      model
   }

 }