论文How Large a Vocabulary Does Text Classification Need? A Variational Approach to Vocabulary Selection的代码分析

论文How Large a Vocabulary Does Text Classification Need? A Variational Approach to Vocabulary Selection的代码实现，我有一篇关于这个论文的介绍，这里介绍它的具体实现。

引入及简介

论文How Large a Vocabulary Does Text Classification Need? A Variational Approach to Vocabulary Selection的代码实现，我有一篇关于这个论文的介绍，这里介绍它的具体实现。

这项目的代码实现给我看麻了，肯定是那种大组的祖传代码，2019年的论文，用的是python2+tensorflow0.X（甚至不是1.X），一堆上古写法，看一点就得谷歌一会儿。

当然，好过那种不给代码的。由于实际上我没弄懂变分估计，我只是按照它的代码把整个计算逻辑搬过去了而已。

（本工作进行于2020年5月）

通过变分法估计词的重要性分布需要在具体的任务中进行，在不同的任务下，单词的重要性分布可能会发生变化。因此本项目算是个多任务学习模型，损失由主任务损失和变分词嵌入模型损失组成。

核心模型

这篇论文的核心是带变分估计的词嵌入，故其核心代码是项目中的VarDropoutEmbedding类，但是可能是因为祖传代码，这个项目的解耦做的很屑，除了VarDropoutEmbedding类，还有许多机制放在了主model中，甚至还有许多机制放在了训练的外围代码中。

VarDropoutEmbedding的初始化代码如下：

def __init__(self, input_size, layer_size, batch_size, name="embedding"):
    self.name = name
    self.input_size = input_size
    self.layer_size = layer_size
    self.batch_size = batch_size
    self.logdropout_init = tf.random_uniform_initializer(0, 3)
    # 全词汇的嵌入矩阵
    self.embedding_mean = tf.get_variable(name, [self.input_size, self.layer_size])
    # 词的dropout值
    self.embedding_logdropout_ratio = tf.get_variable(name + "_ratio", [self.input_size, 1], initializer=self.logdropout_init)
    self.eps = tf.random_normal([self.batch_size, 1, self.layer_size], 0.0, 1.0)

由代码可知，模型为所有的词创建了词嵌入矩阵，即嵌入模型训练时用到了整个语料库的所有的词，因为你并不能预先假设哪些词没有用（可以进行一定的预先切割）。

并按词的数量创建了一个dropout向量，每个词的dropout值在0~3间随机生成

VarDropoutEmbedding的调用化代码如下：

def __call__(self, input_data, sample=False, mask=None):
    if sample:
        output_mean = tf.nn.embedding_lookup(self.embedding_mean, input_data)
        output_logdropout = tf.nn.embedding_lookup(self.clip(self.embedding_logdropout_ratio), input_data)
        output_std = tf.exp(0.5 * output_logdropout) * output_mean
        output = output_mean + output_std * self.eps
    elif mask is None:
        output = tf.nn.embedding_lookup(self.clip(self.embedding_mean), input_data)
    else:
        output = tf.nn.embedding_lookup(mask * self.clip(self.embedding_mean), input_data)
    return output

当调用嵌入层时，如果是使用采样模式（由sample标记控制），就使用完整的词向量矩阵，查询每个词的向量和每个词的dropout值，并进行一定的转换（应该是变分的机制）。

如果不是采样模式，且没有通过外部传入mask的话，则直接查询嵌入；如果外部传入了mask，则用外部传入的mask（这条主要在验证和测试时使用）。

其中self.clip()是将dropout值限制到-10到10之间的一个内部方法。

代码纠缠处其之一

上面说到，本模型的损失由主任务模型和变分模型损失组成，变分估计是推断未知数据的概率分布，而衡量两个概率分布之间的差一半用KL散度，这个项目把KL散度的计算的一部分放在了VarDropoutEmbedding层中：

def regularizer(self):
    k1, k2, k3 = 0.63576, 1.8732, 1.48695
    log_alpha = self.clip(self.embedding_logdropout_ratio)  # [all_word, 1]
    KLD = -tf.reduce_sum(k1 * tf.sigmoid(k2 + k3 * log_alpha) - 0.5 * tf.nn.softplus(-log_alpha) - k1)
    return KLD

注意，只是一部分，因为这里算出来的这个KL散度并不能直接用，在主模型中还需要处理这个KL散度。

主模型

主模型中涉及变分嵌入模型的代码如下：

...
# 创建变分嵌入模型
self.embedding = VarDropoutEmbedding(vocabulary_size, emb_size, batch_size)

# 获取变分嵌入层中的mask
if variational:
    self.mask = tf.cast(tf.less(self.embedding.embedding_logdropout_ratio, self.threshold), tf.float32)
    self.sparsity = tf.nn.zero_fraction(self.mask)
...
self.weight_decay = tf.placeholder(tf.float32, shape=(), name="weight_decay")

# 查找input的词嵌入
if variational:
    if is_training:
        self.x_emb = tf.expand_dims(self.embedding(self.x, sample=True, mask=None), -1)
    else:
        self.x_emb = tf.expand_dims(self.embedding(self.x, sample=False, mask=self.mask), -1)
        self.embedding_matrix = self.embedding.zeroed_embedding(self.mask)
...
# 计算变分损失（权重递减）
if variational:
    self.reg_loss = self.weight_decay * self.embedding.regularizer()
else:
    self.reg_loss = tf.constant(0., dtype=tf.float32)
...
# 主任务代码
...
# 计算总损失
if is_training:
    with tf.name_scope("loss"):
        self.cross_entropy = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits, labels=self.y))
        self.loss = self.cross_entropy + self.reg_loss
        self.optimizer = tf.compat.v1.train.AdamOptimizer(self.learning_rate).minimize(self.loss, global_step=self.global_step, var_list=trainables)

主模型会将外部传入的threshold值送入到变分嵌入层，并获得mask矩阵。KL散度将会在主模型内通过weight_decay进行权重递减。获得主任务的输出后将会计算主任务的损失，并将两个损失相加。

按我的习惯，这个损失要么在变分嵌入层内完成，不要暴露到外面来，事实上这个reg_loss的计算与主模型无关。

要么就把所有损失都放在最外层的外围代码中，不要在模型内部写损失计算逻辑，不利于解耦。

这是我觉的第二个纠结的地方。

外围代码

前面的代码块里有一个weight_decay值，这个值是从外部给来的，仅用于执行对于KD散度的权重递减（本代码没有收录adam优化器）。

weight_decay通过get_decay_rate函数计算，并feed_dict传入（feed_dict是坏文明），代码如下：

for epochs, x_batch, y_batch in train_batches:        
    cur_decay, learning_rate = get_decay_rate(epochs)
    train_feed_dict = {
        model.x: x_batch,
        model.y: y_batch,
        model.weight_decay: cur_decay,
        model.learning_rate: learning_rate,
        model.threshold: 3.0,
        model.l1_threshold: 1e-4
    }

由以上代码可知，get_decay_rate函数的入参是epochs，即训练轮次。

按照训练轮次进行权重递减是常见操作，但是它这里手撸权重递减，让我后面在keras里使用这个方法时，改得意识模糊。因为keras中epoch值仅有框架本身能掌握，没有提供向模型内传入当前训练的epoch的方法。

并且由于tensorflow1.X默认不是eager execution模式，而itag是用keras实现的，后端是1.x版本的tensorflow，直接改成eager execution有要修一堆bug，也没法自己简单的实现入参。

最后我在itag中新增了一个input，然后在data generator中生成epoch作为假数据输入实现。

以下是get_decay_rate函数

def get_decay_rate(epoch):
    if "char" in args.model:
        return 0, 1e-3
    else:
        if "rnn" in args.model:
            learning_rate = 5e-3
        else:
            learning_rate = 5e-3
        if args.l1:
            return 1e-6, learning_rate
        elif args.variational:
            if "rnn" in args.model:
                small_decay = 1e-5
            else:
                small_decay = 1e-5
            large_deacy = 1e-5
            start_decay = 40
            interval = (large_deacy - small_decay) / (NUM_EPOCHS - start_decay)
            if epoch < start_decay:
                cur_decay = small_decay
            else:
                cur_decay = interval * (epoch - small_decay) + small_decay
            return cur_decay, learning_rate
        else:
            return 0, learning_rate

根据使用的主任务模型的不同，KL散度的decay和学习率也有所不同。

总结

上古代码真难看。

feed_dict是坏文明。

KL散度实际计算逻辑分了三块真反人类。