ENAS-pytorch项目中共享RNN模型深度解析

概述

本文将深入分析ENAS-pytorch项目中的共享RNN模型实现，该模型是Efficient Neural Architecture Search(ENAS)算法的核心组件之一。共享RNN模型通过权重共享机制，使得不同子模型可以复用同一组参数，大幅提升了神经架构搜索的效率。

模型架构

1. 基础组件

1.1 EmbeddingDropout

EmbeddingDropout类是对标准PyTorch Embedding层的扩展，实现了词嵌入级别的dropout。与传统dropout不同，它会在嵌入矩阵中随机置零整行，相当于在输入序列中"丢弃"某些词。

特点：

• 在训练时按概率dropout随机屏蔽部分词嵌入
• 可通过scale参数对保留的词嵌入进行额外缩放
• 测试时不进行dropout操作

1.2 LockedDropout

LockedDropout实现了时间步锁定的dropout，即在序列处理过程中，对同一时间步的所有元素使用相同的dropout掩码。这种技术在处理序列数据时比标准dropout表现更好。

2. 核心RNN模型

2.1 初始化结构

模型初始化时创建了以下关键组件：

• 编码器(encoder): 将输入token转换为嵌入向量
• 解码器(decoder): 将隐藏状态映射回输出空间
• 输入到隐藏的线性变换(w_xc, w_xh)
• 隐藏到隐藏的权重矩阵(w_hc_raw, w_hh_raw)

2.2 权重共享机制

模型通过w_h和w_c两个字典实现了节点间的连接权重共享：

• w_h: 存储节点间的隐藏状态变换权重
• w_c: 存储节点间的门控变换权重
• 这些权重在搜索过程中被所有可能的架构共享

2.3 DropConnect实现

模型使用_get_dropped_weights函数实现了DropConnect技术，即在训练时随机丢弃权重矩阵中的连接而非激活值。

前向传播过程

1. 输入处理

1. 输入序列通过EmbeddingDropout层转换为嵌入向量
2. 应用LockedDropout进行输入dropout(如果启用)

2. 循环处理

模型对每个时间步执行以下操作：

1. 通过cell方法计算当前时间步的输出和隐藏状态
2. 对隐藏状态进行梯度裁剪，防止爆炸
3. 收集所有时间步的输出和隐藏状态

3. 输出处理

1. 对输出应用LockedDropout
2. 通过解码器将隐藏状态映射到输出空间

核心单元(Cell)实现

cell方法实现了ENAS论文中描述的RNN单元计算过程：

1. 初始化第一个节点的计算：

   f[0] = self.get_f(dag[-1][0].name)
   c[0] = F.sigmoid(self.w_xc(x) + F.linear(h_prev, self.w_hc, None))
   h[0] = (c[0]*f[0](self.w_xh(x) + F.linear(h_prev, self.w_hh, None)) +
          (1 - c[0])*h_prev)
   

2. 广度优先遍历计算图中的所有节点：

   • 对每个节点应用对应的激活函数(ReLU、tanh等)
   • 计算门控值和隐藏状态
   • 收集所有叶节点的输出

3. 对叶节点输出取平均作为单元最终输出

关键技术点

1. 权重共享：所有可能的架构共享同一组权重参数，极大提高了搜索效率
2. DropConnect：在权重矩阵上而非激活值上应用dropout，提供更强的正则化
3. 梯度裁剪：对隐藏状态的范数进行裁剪，防止梯度爆炸
4. 批归一化：在训练模式下使用批归一化稳定训练过程
5. 参数初始化：根据模式(训练/评估)使用不同的初始化范围

模型特点

1. 实现了ENAS论文中描述的RNN架构搜索空间
2. 支持多种激活函数(ReLU、tanh、sigmoid等)
3. 提供了完整的dropout和dropconnect实现
4. 包含隐藏状态稳定技术(裁剪、批归一化)
5. 实现了权重绑定(tie_weights)选项以减少参数量

通过这种共享RNN模型的实现，ENAS算法能够高效地探索RNN架构空间，找到在特定任务上表现优异的网络结构。