文档分类的分层转换器：模型实现错误，提取注意力权重

Question

我正在尝试在 Keras/tensorflow 中实现用于文档分类的分层转换器，其中：

(1) 词级转换器生成每个句子的表示，以及每个词的注意力权重，并且，

(2) 句子级转换器使用 (1) 的输出来生成每个文档的表示，以及每个句子的注意力权重，最后，

(3) (2) 生成的文档表示用于对文档进行分类（在以下示例中，属于或不属于给定类）。

我正在尝试按照 Yang 等人 (https://www.cs.cmu.edu/~./hovy/papers/16HLT-hierarchical-attention-networks.pdf) 的方法对分类器进行建模，但将 GRU 和注意力层替换为转换器。

我正在使用来自 https://keras.io/examples/nlp/text_classification_with_transformer/ 的 Apoorv Nandan 的变压器实现。

我有两个问题要感谢社区的帮助：

(1) 我在上层（句子）级别模型中遇到一个我无法解决的错误（详情和代码如下）

(2) 我不知道如何提取单词和句子级别的注意力权重，并重视如何最好地做到这一点的建议。

我是 Keras 和这个论坛的新手，因此对于明显的错误深表歉意，并提前感谢您的帮助。

这是一个可重现的例子，指出我遇到错误的地方：

首先，在 Nandan 之后，建立多头注意力、变换器和令牌/位置嵌入层。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import pandas as pd
import numpy as np

class MultiHeadSelfAttention(layers.Layer):
    def __init__(self, embed_dim, num_heads=8):
        super(MultiHeadSelfAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}"
            )
        self.projection_dim = embed_dim // num_heads
        self.query_dense = layers.Dense(embed_dim)
        self.key_dense = layers.Dense(embed_dim)
        self.value_dense = layers.Dense(embed_dim)
        self.combine_heads = layers.Dense(embed_dim)

    def attention(self, query, key, value):
        score = tf.matmul(query, key, transpose_b=True)
        dim_key = tf.cast(tf.shape(key)[-1], tf.float32)
        scaled_score = score / tf.math.sqrt(dim_key)
        weights = tf.nn.softmax(scaled_score, axis=-1)
        output = tf.matmul(weights, value)
        return output, weights

    def separate_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.projection_dim))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs):
        # x.shape = [batch_size, seq_len, embedding_dim]
        batch_size = tf.shape(inputs)[0]
        query = self.query_dense(inputs)  # (batch_size, seq_len, embed_dim)
        key = self.key_dense(inputs)  # (batch_size, seq_len, embed_dim)
        value = self.value_dense(inputs)  # (batch_size, seq_len, embed_dim)
        query = self.separate_heads(
            query, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)
        key = self.separate_heads(
            key, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)
        value = self.separate_heads(
            value, batch_size
        )  # (batch_size, num_heads, seq_len, projection_dim)
        attention, weights = self.attention(query, key, value)
        attention = tf.transpose(
            attention, perm=[0, 2, 1, 3]
        )  # (batch_size, seq_len, num_heads, projection_dim)
        concat_attention = tf.reshape(
            attention, (batch_size, -1, self.embed_dim)
        )  # (batch_size, seq_len, embed_dim)
        output = self.combine_heads(
            concat_attention
        )  # (batch_size, seq_len, embed_dim)
        return output

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout_rate, name=None):
        super(TransformerBlock, self).__init__(name=name)
        self.att = MultiHeadSelfAttention(embed_dim, num_heads)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(dropout_rate)
        self.dropout2 = layers.Dropout(dropout_rate)

    def call(self, inputs, training):
        attn_output = self.att(inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim, name=None):
        super(TokenAndPositionEmbedding, self).__init__(name=name)
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def call(self, x):
        maxlen = tf.shape(x)[-1]
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

对于本示例，数据为 10,000 个文档，每个文档被截断为 15 个句子，每个句子最多包含 60 个单词，已转换为 1-1000 的整数标记。

X 是包含这些标记的 3-D 张量 (10000, 15, 60)。 y 是包含文档类别（1 或 0）的一维张量。就本例而言，X 和 y 之间没有关系。

以下生成示例数据：

max_docs = 10000
max_sentences = 15
max_words = 60

X = tf.random.uniform(shape=(max_docs, max_sentences, max_words), minval=1, maxval=1000, dtype=tf.dtypes.int32, seed=1)

y = tf.random.uniform(shape=(max_docs,), minval=0, maxval=2, dtype=tf.dtypes.int32, seed=1)

这里我尝试构造词级编码器，在https://keras.io/examples/nlp/text_classification_with_transformer/之后：

# Lower level (produce a representation of each sentence):

embed_dim = 100 # Embedding size for each token
num_heads = 2  # Number of attention heads
ff_dim = 64  # Hidden layer size in feed forward network inside transformer
L1_dense_units = 100 # Size of the sentence-level representations output by the word-level model
dropout_rate = 0.1
vocab_size=1000

word_input = layers.Input(shape=(max_words,), name='word_input') 
word_embedding = TokenAndPositionEmbedding(maxlen=max_words, vocab_size=vocab_size, 
                                           embed_dim=embed_dim, name='word_embedding')(word_input) 
word_transformer = TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=ff_dim, 
                                    dropout_rate=dropout_rate, name='word_transformer')(word_embedding)
word_pool = layers.GlobalAveragePooling1D(name='word_pooling')(word_transformer) 
word_drop = layers.Dropout(dropout_rate,name='word_drop')(word_pool)
word_dense = layers.Dense(L1_dense_units, activation="relu",name='word_dense')(word_drop)
word_encoder = keras.Model(word_input, word_dense) 

word_encoder.summary()

看起来这个单词编码器的工作原理是为了生成每个句子的表示。在这里，在第一个文档上运行，它会生成一个形状为 (15, 100) 的张量，其中包含代表 15 个句子中的每一个的向量：

word_encoder(X[0]).shape

我的问题是将其连接到更高（句子）级别的模型，以生成文档表示。

尝试将单词编码器应用于文档中的每个句子时，我收到错误“NotImplementedError”。对于解决此问题的任何帮助，我将不胜感激，因为错误消息并未提供有关具体问题的信息。

在将单词编码器应用于每个句子之后，目标是应用另一个转换器来为每个句子生成注意力权重，以及用于执行分类的文档级表示。由于上面的错误，我无法确定模型的这部分是否可以工作。

最后，我想为每个文档提取单词和句子级别的注意力权重，并希望得到有关如何执行此操作的建议。

提前感谢您的任何见解。

# Upper level (produce a representation of each document):

L2_dense_units = 100

sentence_input = layers.Input(shape=(max_sentences, max_words), name='sentence_input') 

# This is the line producing "NotImplementedError":
sentence_encoder = tf.keras.layers.TimeDistributed(word_encoder, name='sentence_encoder')(sentence_input) 

sentence_transformer = TransformerBlock(embed_dim=L1_dense_units, num_heads=num_heads, ff_dim=ff_dim, 
                               dropout_rate=dropout_rate, name='sentence_transformer')(sentence_encoder)
sentence_dense = layers.TimeDistributed(Dense(int(L2_dense_units)),name='sentence_dense')(sentence_transformer)
sentence_out = layers.Dropout(dropout_rate)(sentence_dense)
preds = layers.Dense(1, activation='sigmoid', name='sentence_output')(sentence_out)

model = keras.Model(sentence_input, preds) 
model.summary()

Answer 1

在尝试与您做同样的事情时，我也遇到了 NotImplementedError。问题是 Keras 的 TimeDistributed 层需要知道其内部自定义层的输出形状。所以你应该在你的自定义层中添加 compute_output_shape 方法。

在您的情况下，MultiHeadSelfAttention、TransformerBlock 和 TokenAndPositionEmbedding 层应包括：

class MultiHeadSelfAttention(layers.Layer):
    ...
    def compute_output_shape(self, input_shape):
        # it does not change the shape of its input
        return input_shape

class TransformerBlock(layers.Layer):
    ...
    def compute_output_shape(self, input_shape):
        # it does not change the shape of its input
        return input_shape

class TokenAndPositionEmbedding(layers.Layer):
    ...
    def compute_output_shape(self, input_shape):
        # it changes the shape from (batch_size, maxlen) to (batch_size, maxlen, embed_dim)
        return input_shape + (self.pos_emb.output_dim,)

添加这些方法后，您应该可以运行您的代码了。

关于你的第二个问题，我不确定，但也许你可以在 MultiHeadSelfAttention 和 TransformerBlock 的调用方法中返回从 MultiHeadSelfAttention 的注意力方法返回的“权重”变量。这样您就可以在构建模型的位置访问它。