#

可配置的 Transformer 组件

9import copy
10
11import torch.nn as nn
12
13from labml.configs import BaseConfigs, option, calculate, aggregate
14from labml_helpers.module import Module
15from .feed_forward import FeedForward
16from .mha import MultiHeadAttention
17from .models import EmbeddingsWithPositionalEncoding, EmbeddingsWithLearnedPositionalEncoding, TransformerLayer, \
18    Encoder, Decoder, Generator, EncoderDecoder

#

FFN 配置

在feed_forward.py 中定义了一个位置前馈网络。

21class FeedForwardConfigs(BaseConfigs):

#

位置前馈层

31    ffn: FeedForward

#

嵌入的特征数量

33    d_model: int

#

隐藏层中的特征数量

35    d_ff: int = 2048

#

Dropout 率

37    dropout: float = 0.1

#

位置前馈层中的激活函数

39    activation: nn.Module = 'ReLU'

#

是否应对 FFN 层进行门控

41    is_gated: bool = False

#

第一个全连接层是否具有可学习的偏置

43    bias1: bool = True

#

第二个全连接层是否具有可学习的偏置

45    bias2: bool = True

#

门控的全连接层是否具有可学习的偏置

47    bias_gate: bool = False

#

预定义的 GLU 变体

49    glu_variant: str = 'none'

#

ReLU 激活函数

$max (0, x)$

52@option(FeedForwardConfigs.activation, 'ReLU')
53def _ffn_activation_relu():

#

59    return nn.ReLU()

#

GELU 激活函数

$x Φ (x)$ 其中， $Φ (x) = P (X \leq x), X \sim N (0, 1)$

这是在论文《 Gaussian Error Linear Units 》中介绍的。

62@option(FeedForwardConfigs.activation, 'GELU')
63def _ffn_activation_gelu():

#

71    return nn.GELU()

#

初始化前馈网络

74@option(FeedForwardConfigs.ffn, 'default')
75def _feed_forward(c: FeedForwardConfigs):

#

79    return FeedForward(c.d_model, c.d_ff,
80                       dropout=c.dropout,
81                       activation=c.activation,
82                       is_gated=c.is_gated,
83                       bias1=c.bias1,
84                       bias2=c.bias2,
85                       bias_gate=c.bias_gate)

#

GLU 变体

这些是在论文《 GLU Variants Improve Transformer 》中包含的各种带门控隐藏层的 FFN 变体。我们已按照论文规定省略了偏置项。

#

带门控线性单元的 FFN

$FF N_{G LU} (x) (x, W_{1}, V, W_{2}) = (σ (x W_{1}) \otimes x V) W_{2}$

95aggregate(FeedForwardConfigs.glu_variant, 'GLU',
96          (FeedForwardConfigs.is_gated, True),
97          (FeedForwardConfigs.bias1, False),
98          (FeedForwardConfigs.bias2, False),
99          (FeedForwardConfigs.bias_gate, False),
100          (FeedForwardConfigs.activation, nn.Sigmoid()))

#

带双线性隐藏层的 FFN

$FF N_{B i l in e a r} (x) (x, W_{1}, V, W_{2}) = (x W_{1} \otimes x V) W_{2}$

105aggregate(FeedForwardConfigs.glu_variant, 'Bilinear',
106          (FeedForwardConfigs.is_gated, True),
107          (FeedForwardConfigs.bias1, False),
108          (FeedForwardConfigs.bias2, False),
109          (FeedForwardConfigs.bias_gate, False),
110          (FeedForwardConfigs.activation, nn.Identity()))

#

带 ReLU 门的 FFN

$FF N_{R e G LU} (x) (x, W_{1}, V, W_{2}) = (max (0, x W_{1}) \otimes x V) W_{2}$

115aggregate(FeedForwardConfigs.glu_variant, 'ReGLU',
116          (FeedForwardConfigs.is_gated, True),
117          (FeedForwardConfigs.bias1, False),
118          (FeedForwardConfigs.bias2, False),
119          (FeedForwardConfigs.bias_gate, False),
120          (FeedForwardConfigs.activation, nn.ReLU()))

#

带 GELU 门的 FFN

$FF N_{GEG LU} (x) (x, W_{1}, V, W_{2}) = (GELU (x W_{1}) \otimes x V) W_{2}$

125aggregate(FeedForwardConfigs.glu_variant, 'GEGLU',
126          (FeedForwardConfigs.is_gated, True),
127          (FeedForwardConfigs.bias1, False),
128          (FeedForwardConfigs.bias2, False),
129          (FeedForwardConfigs.bias_gate, False),
130          (FeedForwardConfigs.activation, nn.GELU()))

#

带 Swish 门的 FFN

$FF N_{Sw i G LU} (x) (x, W_{1}, V, W_{2}) = (Swish_{1} (x W_{1}) \otimes x V) W_{2}$ 其中， $Swish_{β} (x) = x σ (β x)$

136aggregate(FeedForwardConfigs.glu_variant, 'SwiGLU',
137          (FeedForwardConfigs.is_gated, True),
138          (FeedForwardConfigs.bias1, False),
139          (FeedForwardConfigs.bias2, False),
140          (FeedForwardConfigs.bias_gate, False),
141          (FeedForwardConfigs.activation, nn.SiLU()))

#

Transformer 配置

这定义了 Transformer 的配置。这些配置是通过可选择的函数进行计算的。它们是惰性加载的，因此只有必要的模块才会被计算。

144class TransformerConfigs(BaseConfigs):

#

注意力头数量

156    n_heads: int = 8

#

Transformer 嵌入大小

158    d_model: int = 512

#

层数

160    n_layers: int = 6

#

Dropout 率

162    dropout: float = 0.1

#

源词汇表中的 token 数量（用于 token 嵌入）

164    n_src_vocab: int

#

目标词汇表中的 token 数量（用于生成预测的 logits ）

166    n_tgt_vocab: int

#

编码器自注意力

169    encoder_attn: MultiHeadAttention = 'mha'

#

解码器自注意力

171    decoder_attn: MultiHeadAttention = 'mha'

#

解码器记忆与注意力

173    decoder_mem_attn: MultiHeadAttention = 'mha'

#

可配置的前馈层

176    ffn: FeedForwardConfigs

#

编码器层

179    encoder_layer: TransformerLayer = 'default'

#

解码器层

181    decoder_layer: TransformerLayer = 'default'

#

由多个编码器层组成的编码器

184    encoder: Encoder = 'default'

#

由多个解码器层组成的编码器

186    decoder: Decoder = 'default'

#

源数据的嵌入层

189    src_embed: Module = 'fixed_pos'

#

目标数据的嵌入层（用于解码器）

191    tgt_embed: Module = 'fixed_pos'

#

用于预测的 Logit 生成器

194    generator: Generator = 'default'

#

编码器-解码器

197    encoder_decoder: EncoderDecoder

#

多头注意力

201def _mha(c: TransformerConfigs):
202    return MultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
203
204
205calculate(TransformerConfigs.encoder_attn, 'mha', _mha)
206calculate(TransformerConfigs.decoder_attn, 'mha', _mha)
207calculate(TransformerConfigs.decoder_mem_attn, 'mha', _mha)

#

相对多头注意力

211def _relative_mha(c: TransformerConfigs):
212    from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
213    return RelativeMultiHeadAttention(c.n_heads, c.d_model)
214
215
216calculate(TransformerConfigs.encoder_attn, 'relative', _relative_mha)
217calculate(TransformerConfigs.decoder_attn, 'relative', _relative_mha)
218calculate(TransformerConfigs.decoder_mem_attn, 'relative', _relative_mha)

#

创建前馈层配置

221@option(TransformerConfigs.ffn, 'default')
222def _feed_forward(c: TransformerConfigs):

#

226    conf = FeedForwardConfigs()
227    conf.set_default(FeedForwardConfigs.d_model, func=lambda: c.d_model)
228    conf.set_default(FeedForwardConfigs.dropout, func=lambda: c.dropout)
229    return conf

#

编码器层

232@option(TransformerConfigs.encoder_layer, 'default')
233def _encoder_layer(c: TransformerConfigs):

#

237    return TransformerLayer(d_model=c.d_model, self_attn=c.encoder_attn,
238                            src_attn=None, feed_forward=copy.deepcopy(c.ffn.ffn),
239                            dropout_prob=c.dropout)

#

解码器层

242@option(TransformerConfigs.decoder_layer, 'default')
243def _decoder_layer(c: TransformerConfigs):

#

247    return TransformerLayer(d_model=c.d_model, self_attn=c.decoder_attn,
248                            src_attn=c.decoder_mem_attn, feed_forward=copy.deepcopy(c.ffn.ffn),
249                            dropout_prob=c.dropout)

#

编码器

252@option(TransformerConfigs.encoder, 'default')
253def _encoder(c: TransformerConfigs):

#

257    return Encoder(c.encoder_layer, c.n_layers)

#

解码器

260@option(TransformerConfigs.decoder, 'default')
261def _decoder(c: TransformerConfigs):

#

265    return Decoder(c.decoder_layer, c.n_layers)

#

Logit 生成器

268@option(TransformerConfigs.generator, 'default')
269def _generator(c: TransformerConfigs):

#

273    return Generator(c.n_tgt_vocab, c.d_model)

#

固定位置嵌入

使用固定位置编码进行源嵌入

277@option(TransformerConfigs.src_embed, 'fixed_pos')
278def _src_embed_with_positional(c: TransformerConfigs):

#

282    return EmbeddingsWithPositionalEncoding(c.d_model, c.n_src_vocab)

#

使用固定位置编码进行目标嵌入

285@option(TransformerConfigs.tgt_embed, 'fixed_pos')
286def _tgt_embed_with_positional(c: TransformerConfigs):

#

290    return EmbeddingsWithPositionalEncoding(c.d_model, c.n_tgt_vocab)

#

可学习的位置嵌入

使用可学习的位置编码进行嵌入

294@option(TransformerConfigs.src_embed, 'learned_pos')
295def _src_embed_with_learned_positional(c: TransformerConfigs):

#

299    return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_src_vocab)

#

使用可学习的位置编码进行目标嵌入

302@option(TransformerConfigs.tgt_embed, 'learned_pos')
303def _tgt_embed_with_learned_positional(c: TransformerConfigs):

#

307    return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_tgt_vocab)

#

无位置嵌入

没有位置编码的源嵌入

311@option(TransformerConfigs.src_embed, 'no_pos')
312def _src_embed_without_positional(c: TransformerConfigs):

#

316    return nn.Embedding(c.n_src_vocab, c.d_model)

#

319@option(TransformerConfigs.tgt_embed, 'no_pos')
320def _tgt_embed_without_positional(c: TransformerConfigs):
321    return nn.Embedding(c.n_tgt_vocab, c.d_model)
322
323
324@option(TransformerConfigs.encoder_decoder, 'default')
325def _encoder_decoder(c: TransformerConfigs):
326    return EncoderDecoder(c.encoder, c.decoder, c.src_embed, c.tgt_embed, c.generator)