FLUX.1架构的理解2——Transformer的注意力机制
文章目录
一、前言
仅供参考,未经实验验证。参考资料:
Demystifying Flux Architecture
论文地址:https://arxiv.org/pdf/2507.09595
二、FLUX.1架构
1.3 Transformer模型
Transformer 是一种最初为自然语言处理 (NLP) 引入的神经网络架构 [16]。Transformer 的核心创新是自注意力机制,它使模型能够衡量输入序列中所有元素之间的关系。此架构现已成为跨多种模态(包括文本、图像、视频和多模态系统)的许多最先进模型的基础。
在每个注意力块中,模型通过学习到的线性投影从输入嵌入中计算三个矩阵:查询(Q)、键(K)和值(V)。注意力分数计算如下:
Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dkQKT)V
此处, d k {d_k} dk是键向量的维度,softmax 操作确保注意力权重之和为 1。该机制允许模型在计算序列中每个新表示时,为不同 token 分配不同程度的重要性。
【公式含义】:
该公式描述了 Transformer
模型中自注意力(Self-Attention)机制的核心计算过程。它旨在量化输入序列中不同元素之间的关联强度,并据此加权聚合信息,为每个元素生成新的表示。【符号解释】:
- Attention ( Q , K , V ) \text{Attention}(Q, K, V) Attention(Q,K,V): 表示注意力机制的输出,是一个与 V V V 维度相同的矩阵,包含了加权后的信息。
- Q Q Q: Queries (查询) 矩阵。它由输入嵌入通过一个学习到的线性投影得到,代表了当前元素在寻找其他相关元素时的“提问”或“关注点”。
- K K K: Keys (键) 矩阵。它由输入嵌入通过另一个学习到的线性投影得到,代表了序列中所有元素各自的“身份”或“特征”,用于与查询进行匹配。
- V V V: Values (值) 矩阵。它由输入嵌入通过第三个学习到的线性投影得到,代表了序列中所有元素实际携带的“信息内容”。
- Q K T Q K^T QKT: 查询矩阵 Q Q Q 与键矩阵 K K K 的转置 K T K^T KT 的乘积。这计算了所有查询与所有键之间的点积相似度,得到一个注意力分数矩阵,表示每个查询对每个键的关注程度。
- d k d_k dk: 键向量的维度(dimension of key vectors)。
- d k \sqrt{d_k} dk: 键向量维度 d k d_k dk 的平方根。除以 d k \sqrt{d_k} dk 是为了缩放点积结果。当 d k d_k dk 很大时,点积结果的量级也会很大,可能导致 softmax函数的梯度变得非常小(进入饱和区),从而影响训练。通过缩放可以缓解这个问题,保持梯度的稳定性。
- softmax ( ⋅ ) \text{softmax}(\cdot) softmax(⋅): 软最大函数。它将注意力分数矩阵的每一行(对应一个查询对所有键的关注)转换为一个概率分布,确保所有权重之和为1。这样,每个元素对其他元素的关注度就被归一化为 0 到 1 之间的权重。
- ⋅ V \cdot V ⋅V: 将 softmax \text{softmax} softmax 得到的注意力权重与值矩阵 V V V 相乘。这意味着每个值向量都会根据其对应的注意力权重进行加权。权重高的值向量贡献更大,权重低的值向量贡献较小。
【公式解释】:
该公式不是一个推导公式,而是点积注意力(Dot-Product Attention)机制的定义和计算流程。它描述了
Transformer 模型如何通过计算输入序列中不同元素之间的相关性,来生成每个元素新的、上下文感知的表示。背景知识:
在传统的神经网络(如循环神经网络RNN)中,处理序列数据时,每个时间步的输出通常只依赖于当前输入和前一个时间步的隐藏状态。这使得模型难以捕捉长距离依赖关系。Transformer架构的核心创新在于引入了自注意力机制,它允许模型在处理序列的某个元素时,能够同时考虑到序列中的所有其他元素,并根据它们之间的相关性来动态地加权聚合信息。
公式的应用场景和详细解释:
该公式描述了自注意力计算的三个主要步骤:
计算注意力分数( Q K T Q K^T QKT):
- 首先,将输入序列中的每个元素(例如,一个词嵌入或图像块嵌入)通过三个不同的线性变换,分别映射成查询( Q Q Q)、键( K K K)和值( V V V)向量。
- 然后,通过计算查询矩阵 Q Q Q 和键矩阵 K K K 的转置 K T K^T KT 的点积,得到一个注意力分数矩阵。这个矩阵的每个元素 ( i , j ) (i, j) (i,j) 表示第 i i i 个查询(来自输入序列的第 i i i 个元素)对第 j j j 个键(来自输入序列的第 j j j
个元素)的相似度或相关性。点积越大,表示两者越相关。缩放( / d k / \sqrt{d_k} /dk):
- 将注意力分数矩阵的每个元素除以键向量维度 d k d_k dk 的平方根。这一步的目的是防止点积结果过大,从而避免 softmax \text{softmax} softmax 函数在输入值很大时进入梯度饱和区,导致梯度过小,影响模型的学习效率和稳定性。这是 Transformer论文《Attention Is All You Need》中提出的关键技巧。
归一化和加权求和( softmax ( ⋅ ) ⋅ V \text{softmax}(\cdot) \cdot V softmax(⋅)⋅V):
- 对缩放后的注意力分数应用 softmax \text{softmax} softmax 函数。这会将每一行(对应一个查询对所有键的关注)的分数转换成一个介于 0 到 1 之间的概率分布,所有概率之和为1。这些概率可以被视为权重,表示当前查询对序列中每个元素信息的关注程度。
- 最后,将这些归一化后的注意力权重与值矩阵 V V V 相乘。这意味着对于输入序列中的每个元素,其新的表示是所有值向量的加权和,权重由注意力机制动态计算得出。高度相关的元素(高注意力权重)其值向量对最终表示的贡献更大,而不相关的元素贡献较小。
具体应用示例:
- 文本处理: 假设输入是一个句子“The animal didn’t cross the street because it was too tired.”。当模型处理“it”这个词时,通过自注意力机制,它会计算“it”与其他所有词(“The”, “animal”, “didn’t”, “cross”, “the”, “street”, “because”, “was”, “too”,“tired”)之间的相关性。如果“it”与“animal”之间的注意力分数最高,那么在生成“it”的表示时,模型会更多地聚合“animal”的信息,从而理解“it”指的是“animal”。
- 图像处理(Vision Transformers, ViTs): 对于图像任务,图像首先被分割成固定大小的非重叠图像块(patches)。每个图像块被视为一个“token”。当模型处理某个图像块时,它会通过自注意力机制计算该图像块与图像中所有其他图像块之间的相关性。例如,在识别一张狗的图片时,模型可能会发现狗的头部图像块与身体图像块之间有很高的注意力权重,从而更好地理解图像的整体结构和对象。
【总结】:
该公式是 Transformer 架构中自注意力机制的核心,通过计算查询与键的点积相似度,并经过缩放和 softmax \text{softmax} softmax 归一化后,得到注意力权重,再将这些权重与值向量加权求和,从而为输入序列中的每个元素生成一个融合了上下文信息的新的表示。这一机制使得模型能够捕捉序列内部的长距离依赖关系,是Transformer 模型在多种模态任务中取得成功的关键。
在自注意力机制中,为什么选择点积作为计算查询和键之间相似度的方法,以及是否存在其他替代的相似度计算方式?
为什么选择点积作为查询和键之间相似度的方法?
在 Transformer 的自注意力机制中,选择点积 (Dot Product) 作为查询(Query, Q Q Q)和键(Key, K K K)之间相似度的计算方法,主要有以下几个核心原因:
计算效率高:
- 点积运算在现代硬件(如 GPU)上可以高度并行化。矩阵乘法是优化的操作,这使得 Transformer 能够高效地处理长序列。
- 相比其他更复杂的相似度计算,点积在计算上非常简洁和快速,尤其是在高维向量空间中。
数学直观性与几何解释:
- 在向量空间中,两个向量的点积与它们之间的夹角余弦值成正比(当向量归一化时,点积即为余弦相似度)。
- 点积可以直观地理解为两个向量在方向上的匹配程度以及它们的长度(或范数)的乘积。如果两个向量指向相似的方向(即它们“关注”或“代表”相似的信息),它们的点积会较大。如果它们方向正交,点积为零。这很自然地符合“相似度”的直观概念。
- 在 Transformer 的上下文中, Q Q Q 和 K K K 向量通过学习得到,它们被训练成能够编码语义信息,使得语义上相似的查询和键向量在特征空间中具有相近的方向。
可微性,利于梯度传播:
- 点积是一个完全可微的函数,这意味着模型可以通过反向传播算法有效地学习 Q Q Q, K K K, V V V 的线性投影矩阵。
- 这对于训练深度学习模型至关重要,因为梯度可以顺畅地流过注意力机制,优化模型参数。
与深度学习范式兼容:
- 点积作为一种线性操作,与神经网络中常用的线性层和激活函数能够很好地集成。
- 它允许模型通过训练来学习如何将输入映射到 Q Q Q 和 K K K 空间,从而自动发现哪些特征对于计算相关性是重要的。
缩放点积的必要性:
- 虽然点积本身是高效的,但原始 Transformer 论文中还引入了除以 d k \sqrt{d_k} dk 的缩放因子。这是因为在高维向量空间中,点积的值可能会变得非常大,导致 softmax \text{softmax} softmax 函数的输入值过大。当 softmax \text{softmax} softmax 函数的输入值过大时,其输出会变得非常接近 0 或 1,导致梯度消失,使得模型难以学习。
- 通过缩放,可以将点积结果的方差控制在一个更合理的范围内,从而避免 softmax \text{softmax} softmax 函数进入饱和区,维持梯度的稳定性,促进模型更好地训练。
替代的相似度计算方式:
除了点积,理论上存在多种方法可以计算查询和键之间的相似度。这些方法各有优缺点,有些也在其他注意力机制或神经网络架构中得到了应用。
加性注意力(Additive Attention / Concat Attention):
- 方法: 不像点积那样直接相乘,而是将 Q Q Q 和 K K K 向量拼接起来,然后通过一个前馈神经网络(通常包含一个 tanh \tanh tanh 激活函数)来计算相似度。
- 公式形式: e i j = v T tanh ( W q q i + W k k j ) e_{ij} = v^T \tanh(W_q q_i + W_k k_j) eij=vTtanh(Wqqi+Wkkj)
- 优点:
- 在低维空间中,加性注意力通常表现更好,因为它能够捕捉到点积可能遗漏的更复杂的非线性关系。
- 在某些早期的序列模型(如机器翻译中的 Bahdanau Attention)中表现良好。
- 缺点:
- 计算开销更大: 需要额外的参数(前馈网络的权重 W q , W k , v W_q, W_k, v Wq,Wk,v)和更复杂的非线性计算,导致计算效率低于点积注意力,尤其是在处理长序列时。
- 难以并行化:前馈网络的计算通常不如矩阵乘法那样容易并行化。
余弦相似度(Cosine Similarity):
- 方法: 计算两个向量夹角的余弦值。
- 公式形式: sim ( Q , K ) = Q ⋅ K ∣ ∣ Q ∣ ∣ ⋅ ∣ ∣ K ∣ ∣ \text{sim}(Q, K) = \frac{Q \cdot K}{||Q|| \cdot ||K||} sim(Q,K)=∣∣Q∣∣⋅∣∣K∣∣Q⋅K
- 优点:
- 只关注向量的方向,不受向量长度的影响。这在某些场景下可能更鲁棒,例如,如果向量的长度编码了不重要的信息。
- 与点积相似度高度相关,当向量被 L2 归一化后,点积相似度就等同于余弦相似度。
- 缺点:
- 需要额外的归一化步骤(计算范数),这会增加计算开销。
- 在 Transformer 的缩放点积注意力中,通过除以 d k \sqrt{d_k} dk 已经隐含地处理了长度对点积的影响,且通过训练可以学习到向量长度的意义。因此,额外强制归一化不一定是必需的,甚至可能限制模型学习某些有用信息(如重要的 token向量长度可能更大)。
广义点积注意力(General Dot Product Attention):
- 方法: 在点积计算前引入一个可学习的权重矩阵 W a W_a Wa。
- 公式形式: e i j = q i T W a k j e_{ij} = q_i^T W_a k_j eij=qiTWakj
- 优点: 增加了模型的表达能力,可以学习更复杂的相似度度量。
- 缺点: 引入了更多的参数,增加了模型复杂度和计算量。Transformer 采用的简化版本(无 W a W_a Wa)在实践中已经足够强大,且效率更高。
距离函数(Distance Functions):
- 方法: 使用欧氏距离、曼哈顿距离等来衡量向量之间的“不相似度”,然后将其转换为相似度(例如,通过取负数或指数函数)。
- 例如: sim ( Q , K ) = − distance ( Q , K ) \text{sim}(Q, K) = -\text{distance}(Q, K) sim(Q,K)=−distance(Q,K) 或 sim ( Q , K ) = exp ( − distance ( Q , K ) ) \text{sim}(Q, K) = \exp(-\text{distance}(Q, K)) sim(Q,K)=exp(−distance(Q,K))
- 优点: 直观易懂,在某些特定任务中可能适用。
- 缺点:
- 欧氏距离等距离函数在未经缩放处理时,其值域可能不适合直接送入 softmax \text{softmax} softmax。
- 计算通常比点积复杂,特别是当需要计算平方根时。
- 在深度学习中,点积在梯度计算和优化方面通常表现更好。
总结:
点积之所以成为 Transformer自注意力机制中的主流选择,是因为它在计算效率、数学直观性、可微性和与深度学习架构的兼容性之间取得了卓越的平衡。虽然存在其他理论上可行的相似度计算方法,但它们通常会带来额外的计算成本或模型复杂性,而缩放点积注意力在实践中已被证明在性能和效率上都非常出色。Transformer的成功很大程度上归功于这种简洁而强大的机制。
在FLUX.1模型中,自注意力机制是如何与扩散模型相结合的?它在处理图像的哪些特定方面(如细节、全局结构)发挥了关键作用?
FLUX.1 模型中自注意力机制与扩散模型的结合
FLUX.1 是一种基于 Rectified Flow (RF) 的扩散 Transformer模型,它在潜空间(latent space)中操作。与传统的 U-Net 架构扩散模型不同,FLUX.1 采用了完全基于Transformer 的架构来进行去噪过程(参见论文第 3 页 1.2 节和第 9 页 2.3 节)。
这种结合体现在以下几个关键点:
Transformer 作为去噪骨干 (Denoising Backbone):
- 在传统的扩散模型中,通常使用 U-Net 架构来预测噪声或速度向量。FLUX.1 用一个大型 Transformer 网络取代了整个 U-Net,这个 Transformer 的核心就是自注意力机制。
- 这个 Transformer 的任务是估计从噪声图像到清晰图像的速度向量 v θ v_\theta vθ(基于 Rectified Flow 框架),而不是像传统扩散模型那样预测噪声 ϵ \epsilon ϵ(参见论文第 4 页 1.4 节和第 9 页
2.3 节)。- 在推理阶段,模型从纯高斯噪声开始,通过一系列迭代步骤,利用 Transformer 预测的速度向量逐步去噪,最终生成清晰的潜在表示,再由 VAE 解码器转换回像素空间图像(参见论文第 8 页图 5)。
潜空间操作 (Latent Space Operation):
- FLUX.1 继承了 Stable Diffusion 等模型的思想,在压缩的潜空间中执行扩散和去噪过程(参见论文第 5 页 2.2 节)。
- 图像首先被预训练的自编码器编码成低维的潜在表示(latent image embeddings),文本提示也被文本编码器编码成文本嵌入(text embeddings)。
- Transformer 在这个潜空间中对这些潜在的图像和文本嵌入进行操作,大大提高了计算效率,同时保持了视觉保真度。
多模态注意力 (Multi-Modal Attention):
- FLUX.1 的 Transformer 结构支持多模态输入。它将文本嵌入和潜在图像嵌入拼接 (concatenate) 起来,然后通过一系列注意力块进行处理。
- 这种设计使得模型能够同时处理图像信息和文本条件信息,实现文本到图像的对齐(text-image alignment)。
- 论文特别指出,FLUX.1 使用了双流 (Double-Stream) 和单流 (Single-Stream) 两种 Transformer块,它们都采用多模态注意力机制,实现文本和图像模态之间的双向交互(参见论文第 9 页 2.3 节)。
条件机制 (Conditioning Mechanism):
- FLUX.1 通过 Adaptive Layer Normalization (AdaLN) 层将条件信息(如
timestep和pooled_projection,即 CLIP 文本编码器生成的整体提示嵌入)注入到 Transformer的每一层中(参见论文第 10 页 2.3.2 节和图 9)。- AdaLN 动态地生成归一化和调制参数,允许模型根据输入提示或指导信号调整其行为,从而实现对生成过程的精细控制。
自注意力机制在处理图像特定方面的关键作用
自注意力机制在 FLUX.1模型中,通过其独特的全局上下文建模能力,在处理图像的细节、全局结构以及文本与图像的对齐方面发挥了至关重要的作用。
捕捉全局结构 (Global Structure):
- 长距离依赖建模: 传统的卷积网络在捕捉图像的全局结构时需要多层堆叠,且感受野有限。而自注意力机制(特别是 Vision Transformer中的自注意力)能够直接计算图像中任意两个“token”(在这里是图像块的潜在表示)之间的关系,无论它们在空间上相距多远。
- 上下文理解: 通过计算每个图像 token 与所有其他图像 token 的相关性,自注意力机制能够建立图像中不同区域之间的长期依赖关系。例如,在生成一个汽车场景时,它能确保车轮位于车身下方,车窗位于车顶等,从而保证了生成图像的整体结构合理性和一致性。
- 处理复杂场景: 论文第 19 页图 5 展示了 FLUX.1 处理包含多个描述元素(如“十字路口”、“阴天”、“坐在轮椅上的老人”、“学步儿童”)的复杂提示的能力。自注意力机制能够将这些独立元素关联起来,构建出一个逻辑上连贯的复杂场景。
精炼局部细节 (Local Details):
- 虽然自注意力以其全局建模能力而闻名,但它也能够通过迭代细化过程来处理局部细节。在去噪的每个步骤中,Transformer 都会接收到当前噪声潜在表示,并预测一个更接近清晰图像的速度向量。这个过程是迭代的,意味着模型在每一步都能逐步完善图像的细节。
- 多头注意力 (Multi-Head Attention): 论文第 5 页提到,多头注意力机制用于捕捉不同类型的关系。这意味着不同的注意力头可能专注于不同尺度的信息,一些可能关注宏观结构,另一些则可能关注细微纹理或边缘等局部特征。
- 文本条件下的细节生成: 结合文本提示,自注意力机制能够确保生成的局部细节与文本描述精确对齐。例如,如果提示描述“一辆红色的跑车”,注意力机制能够确保生成的车身是红色且具有跑车的流线型细节。论文第18 页图 2 展示了 FLUX.1 能够区分“轻微下雪”和“大雪”的场景,这需要对细微的图像细节有精确的控制。
强大的文本-图像对齐 (Text-Image Alignment):
- 这是自注意力机制与多模态学习结合的直接体现。通过将文本嵌入和图像潜在嵌入拼接,并让它们在同一个注意力块中进行交互(跨模态注意力),模型可以:
- 语义理解: 确保图像内容与文本提示的语义信息高度一致。查询 (Q) 来自图像 token,而键 (K) 和值 (V) 来自文本嵌入,这使得图像生成过程能根据文本的语义信息进行调制(参见论文第 5 页 1.5 节)。
- 忠实度 (Faithfulness): 论文摘要和第 5 页 2.1 节都强调 FLUX.1 旨在实现忠实的文本-图像对齐。自注意力机制在此处发挥核心作用,它能确保图像的每个部分都与提示的相关文本描述保持一致。
- 处理稀有概念和模糊提示: 论文第 19 页图 3 和图 4 展示了 FLUX.1 处理模糊或稀有概念(如“不利天气”、“翻倒的汽车”、“船拖车”)的能力。这得益于 Transformer
的强大语言理解和生成能力,以及它将这些概念准确映射到视觉特征的能力。姿态位置编码 (Rotary Positional Embeddings, RoPE):
- FLUX.1 采用了 RoPE [25] 来注入位置信息(参见论文第 11 页)。RoPE 通过旋转查询和键向量来编码相对位置关系,这对于理解图像中不同部分的空间布局至关重要。
- 这种方法使得模型能够更好地泛化到未见过的序列长度,并支持将文本位置重新用于图像块,在视觉-语言模型中尤为关键。它帮助自注意力机制在考虑全局上下文的同时,精确地定位和组织图像中的元素。
综上所述,FLUX.1 通过将扩散模型与强大的 Transformer架构相结合,并利用自注意力机制的全局上下文建模、多模态交互以及精细的条件注入能力,成功地在生成高质量图像的同时,实现了卓越的文本-图像对齐和对图像细节及全局结构的精确控制。
在实践中,通常使用多头注意力机制,其中并行计算多组 Q/K/V 投影,以捕捉不同类型的关系。
对于视觉任务,Vision Transformers (ViTs) [17]将图像分块为固定大小的非重叠图像块(例如,16×16像素),将每个图像块展平,并将其投影到嵌入空间。然后,这些图像块嵌入通过标准的Transformer编码器进行处理,通常会添加位置编码以保留空间结构。
[17] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly et al., “An image is worth 16x16 words: Transformers for image recognition
at scale,” arXiv preprint arXiv:2010.11929, 2020.
这使得ViTs能够对图像中的长距离依赖进行建模,而无需卷积归纳偏置。
虽然最初的 Q/K/V表示都来自同一输入嵌入的投影,这一过程称为自注意力(Self-Attention),但在文本到图像生成中,Transformer 通常使用交叉注意力(Cross-Attention)来根据文本提示条件化图像合成。在此设置中,查询(Q)来自图像标记(image tokens)的投影,而键(K)和值(V)则来自文本嵌入的投影。这使得模型能够根据提示中的语义信息来调节图像生成。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐


所有评论(0)