在 Devin 中试用 DeepWiki 处理私有仓库
DeepWiki

tensorflow/models

菜单

Transformer 架构

相关源文件

本文档介绍了 TensorFlow Model Garden 中的 Transformer 架构，它是许多最先进 NLP 模型的基础。本文档涵盖了关键组件、实现细节以及代码库中提供的各种 Transformer 变体。

有关使用 Transformer 的具体 NLP 模型的信息，请参阅 NLP 模型和 BERT 及其变体。

Transformer 架构概览

Transformer 是一种深度学习架构，由 Vaswav 等人在论文 “Attention Is All You Need” 中提出。与循环神经网络不同，Transformer 使用自注意力机制并行处理所有输入 token，使其在建模长距离依赖关系方面更加高效。

Model Garden 主要在 official/nlp/modeling/layers 模块中实现 Transformer，其中包含注意力机制、编码器和解码器块以及位置嵌入等核心组件。

来源：official/nlp/modeling/layers/transformer.py official/nlp/modeling/layers/transformer_scaffold.py

核心组件

多头注意力

多头注意力机制是 Transformer 中的核心创新。它允许模型在不同位置联合关注来自不同表示子空间的信息。

代码库中的实现提供了几种注意力变体

  • 标准多头注意力 (Standard MultiHeadAttention)：用于大多数 Transformer 实现
  • 缓存注意力 (CachedAttention)：通过缓存键/值张量来优化解码
  • 多通道注意力 (MultiChannelAttention)：支持多个交叉注意力目标序列
  • TalkingHeadsAttention：在注意力头之间添加线性投影

来源：official/nlp/modeling/layers/attention.py26-107 official/nlp/modeling/layers/multi_channel_attention.py

位置嵌入

由于 Transformer 能够并行处理所有 token，因此它们缺乏固有的序列排序信息。位置嵌入会添加这种至关重要的位置上下文。

类型描述实现
绝对位置嵌入 (Absolute Position Embedding)添加基于绝对位置的学习嵌入PositionEmbedding
相对位置嵌入 (Relative Position Embedding)对 token 之间的相对距离进行编码RelativePositionEmbedding
相对位置偏置 (Relative Position Bias)根据相对位置添加偏置项RelativePositionBias

来源：official/nlp/modeling/layers/position_embedding.py

Transformer 编码器块

编码器块在进料向前网络之后应用自注意力，并带有残差连接和归一化。

来源：official/nlp/modeling/layers/transformer.py official/nlp/modeling/layers/transformer_scaffold.py

Transformer 解码器块

解码器块包含一个额外的交叉注意力层，该层关注编码器输出。

来源：official/nlp/modeling/layers/transformer.py157-457

Transformer 实现细节

关键类

Model Garden 提供了各种具有不同特征的 Transformer 实现

类目的特性
TransformerEncoderBlock标准 Transformer 编码器自注意力、残差连接、可配置的归一化
TransformerDecoderBlock标准 Transformer 解码器自注意力、交叉注意力、支持缓存解码
TransformerScaffold可定制的 Transformer可插入的注意力和前馈组件
ReZeroTransformer使用 ReZero 技术的变体使用学习到的缩放代替层归一化
MobileBertTransformer适用于 MobileBERT 的高效 Transformer用于提高效率的瓶颈架构

来源：official/nlp/modeling/layers/transformer.py official/nlp/modeling/layers/transformer_scaffold.py official/nlp/modeling/layers/rezero_transformer.py

代码组织

来源：official/nlp/modeling/layers/ official/nlp/modeling/networks/

配置系统

Transformer 的配置通过 official/nlp/configs/encoders.py 中的数据类进行管理。这些类为不同的 Transformer 变体提供了结构化的参数定义。

build_encoder 函数会使用这些配置来实例化相应的编码器实现。

来源：official/nlp/configs/encoders.py33-775

Seq2Seq Transformer

对于序列到序列任务，Model Garden 提供了编码器-解码器 Transformer 架构的完整实现。

Seq2Seq Transformer 的关键组件

  • 嵌入层 (Embedding Layer)：将输入 token 转换为向量并添加位置编码
  • TransformerEncoder：通过多个编码器块处理源序列
  • TransformerDecoder：利用编码器上下文生成输出序列
  • 束搜索 (Beam Search)：在推理过程中用于查找最佳输出序列

来源：official/nlp/modeling/models/seq2seq_transformer.py31-359

MobileBERT 和高效 Transformer

Model Garden 包含用于高效部署的优化 Transformer 变体。

MobileBERT 架构

MobileBERT 对标准 Transformer 架构进行了修改以提高效率。

MobileBERT 中的关键优化

  • 使用瓶颈结构来减少计算量
  • 采用线性变换代替层归一化
  • 堆叠多个前馈网络
  • 可选共享键-查询变换

来源：official/nlp/modeling/layers/mobile_bert_layers.py official/nlp/modeling/networks/mobile_bert_encoder.py

高级特性

缓存以实现高效解码

对于自回归生成任务，Transformer 实现支持键和值张量的缓存，以避免重复计算。

来源：official/nlp/modeling/layers/attention.py official/nlp/modeling/models/seq2seq_transformer.py298-359

Transformer Scaffold

TransformerScaffold 类提供了一种灵活的方式来定制 Transformer 层。

Scaffold 方法允许

  • 插入自定义注意力机制
  • 使用专门的前馈网络
  • 保持标准的 Transformer 结构
  • 在重用核心组件的同时创建研究型变体

来源：official/nlp/modeling/layers/transformer_scaffold.py27-359

使用示例

下面是一个创建 Transformer 编码器的简化示例

对于仅编码器模型（如 BERT），仅使用 Transformer 编码器组件。

来源：official/nlp/modeling/models/seq2seq_transformer.py360-651

另请参阅