new(transformer): implement TransformerEncoderLayer and add test

test_transformer.jl

@@ -4,22 +4,28 @@ true || include("transformers.jl")
 
 using Random
 
-const EMBED_DIM = 10
+const INPUT_DIM = 10
-const NUM_HEADS = 2
+const NUM_HEADS = 1
-const KV_DIM = 12
 
-rng = TaskLocalRNG()
+rng = Random.default_rng()
 
-l = MultiheadAttention(EMBED_DIM, NUM_HEADS, kvdim = KV_DIM)
+model = TransformerEncoderLayer(10, 2)
-@info "layer" l
+@info "model" model
 
-ps = LuxCore.initialparameters(rng, l)
+ps, st = LuxCore.setup(rng, model)
-st = LuxCore.initialstates(rng, l)
+@info "parameters" ps
-@info "parameters and states" ps st
+@info "status" st
 
-q = rand(rng, Float32, (EMBED_DIM,))
+# Unbatched
-k = rand(rng, Float32, (KV_DIM,))
+x = randn(rng, Float32, (10,))
-v = rand(rng, Float32, (KV_DIM,))
+@info "input" size(x)
-@info "q k v" summary.((q, k, v))
 
-l((; q, k, v), ps, st)
+y, st = model(x, ps, st)
+@info "output" y st
+
+# Batched
+x = randn(rng, Float32, (10, 10))
+@info "input" size(x)
+
+y, st = model(x, ps, st)
+@info "output" y st

transformers.jl

@@ -3,6 +3,10 @@ using Random: AbstractRNG
 using Lux
 
 """
+    MultiheadAttention{F} <: LuxCore.AbstractLuxLayer
+
+Multi head attention layer used in Transformer model.
+
 # Fields
 - `embed_dim`: Queue vector length. `q_len`
 - `kvdim::Int`: Key vector and value vector length.
@@ -63,6 +67,7 @@ Constructor.
 MultiheadAttention has no states.
 
 # Inputs
+NamedTuple of these three variables.
 - `q`: ``Q``
 - `k`: ``K``
 - `v`: ``V``
@@ -97,6 +102,18 @@ function LuxCore.initialstates(::AbstractRNG, ::MultiheadAttention)
     NamedTuple()
 end
 
+function LuxCore.parameterlength(l::MultiheadAttention)
+    dim_weight_q = l.embed_dim * l.num_heads * l.qdim
+    dim_weight_k = l.embed_dim * l.num_heads * l.kvdim
+    dim_weight_v = l.v_embed_dim * l.num_heads * l.kvdim
+    dim_weight_o = l.embed_dim * l.v_embed_dim * l.num_heads
+    dim_weight_q + dim_weight_k + dim_weight_v + dim_weight_o
+end
+
+function LuxCore.statelength(l::MultiheadAttention)
+    0
+end
+
 function (l::MultiheadAttention)(x::NamedTuple, ps, _st::NamedTuple)
     if size(x.q, 1) != l.embed_dim
         ArgumentError(
@@ -125,9 +142,132 @@ function (l::MultiheadAttention)(x::NamedTuple, ps, _st::NamedTuple)
     # [y] = (v_dim, q_len, batch_size...)
     # [α] = (kv_len, q_len, nheads, batch_size...)
     y, α = dot_product_attention(q, k, v; nheads = l.num_heads)
-    ps.weight_o * y
+    ps.weight_o * y, _st
 end
 
-# struct TransformerEncoder <: LuxCore.AbstractLuxContainerLayer{}
+"""
-# end
+    TransformerEncoderLayer <: LuxCore.AbstractLuxLayer
-#
+
+One layer of encoder block of Transformer model.
+
+# Structure
+
+It consists of two sublayers, `MultiheadAttention` and feed forward network(two Dense layers).
+Both of them are wrapped with residual connection and followed by [`Lux.Dropout`](@ref) and [`Lux.LayerNorm`](@ref).
+
+They are combined using [`Chain`](@ref) and [`SkipConnection`](@ref).
+
+See the constructor of this layer or displayed layer to see the structure.
+
+## MultiheadAttention sublayer
+
+```
+   ========= SkipConnection ==========
+  -> MultiheadAttention -> Dropout -->
+--                                     +(add) -> LayerNorm
+  -> -------------------------------->
+```
+
+## FeedForwardNetworkSubLayer
+
+The core is two chained `Dense` layer, which has internal dimension `feedforward_dim`.
+Then this is wrapped with `SkipConnection` and followed by `LayerNorm`.
+
+# Parameters & States
+
+Since this layer is implemented as [`LuxCore.AbstractLuxContainerLayer`](@ref),
+both parameters and states have structured as the container layer.
+So see the "Structure" section for more detail.
+
+# Inputs
+
+AbstractArray with `size(x, 1) == model_dim`. (same as Dense)
+
+"""
+struct TransformerEncoderLayer{LSA, LFFN} <: LuxCore.AbstractLuxContainerLayer{(
+    :sublayer_self_attention,
+    :sublayer_feed_forward_network,
+)}
+    sublayer_self_attention::LSA
+    sublayer_feed_forward_network::LFFN
+end
+
+"""
+    TransformerEncoderLayer(
+        model_dim::Int,
+        num_heads::Int;
+        feedforward_dim::Int = 2048,
+        dropout::T1 = 0.1,
+        activation::F = relu,
+        layer_norm_eps::T2 = 1e-5,
+    ) where {F, T1 <: AbstractFloat, T2 <: AbstractFloat}
+
+Constructor of [`TransformerEncoderLayer`](@ref).
+
+# Arguments
+- `model_dim::Int`: model input size
+- `num_heads::Int`: number of heads of MultiheadAttention
+- `feedforward_dim::Int = 2048`: dimension of feed forward network
+- `dropout::T1 = 0.1`: dropout rate for all Dropout layers
+- `activation::F = relu`: activation used in feed forward network
+- `layer_norm_epsilon::T2 = 1e-5`: eps of LayerNorm
+"""
+function TransformerEncoderLayer(
+    model_dim::Int,
+    num_heads::Int;
+    feedforward_dim::Int = 2048,
+    dropout::T1 = 0.1,
+    activation::F = relu,
+    layer_norm_epsilon::T2 = 1.0f-5,
+) where {F, T1 <: AbstractFloat, T2 <: AbstractFloat}
+    sublayer_self_attention = let
+        layer_split = Lux.WrappedFunction(x -> (q = x, k = x, v = x))
+        layer_self_attention = MultiheadAttention(model_dim, num_heads)
+        layer_dropout = Lux.Dropout(dropout)
+        layer_residual_connection = Lux.SkipConnection(
+            Lux.Chain(; layer_split, layer_self_attention, layer_dropout),
+            +,
+        )
+        Lux.Chain(;
+            layer_residual_connection,
+            layer_layer_norm = Lux.LayerNorm(model_dim; epsilon = layer_norm_epsilon),
+            name = "SelfAttentionSubLayer",
+        )
+    end
+    sublayer_feed_forward_network = let
+        layer_feed_forward_network = Lux.Chain(
+            Lux.Dense(model_dim => feedforward_dim, activation),
+            Lux.Dropout(dropout),
+            Lux.Dense(feedforward_dim => model_dim, activation),
+            Lux.Dropout(dropout),
+        )
+        layer_residual_connection = Lux.SkipConnection(layer_feed_forward_network, +)
+        Lux.Chain(;
+            layer_residual_connection,
+            layer_layer_norm = Lux.LayerNorm(model_dim; epsilon = layer_norm_epsilon),
+            name = "FeedForwardNetworkSubLayer",
+        )
+    end
+    TransformerEncoderLayer(sublayer_self_attention, sublayer_feed_forward_network)
+end
+
+function (encoder::TransformerEncoderLayer)(x, ps, st)
+    x, st_sublayer_self_attention = Lux.apply(
+        encoder.sublayer_self_attention,
+        x,
+        ps.sublayer_self_attention,
+        st.sublayer_self_attention,
+    )
+    x, st_sublayer_feed_forward_network = Lux.apply(
+        encoder.sublayer_feed_forward_network,
+        x,
+        ps.sublayer_feed_forward_network,
+        st.sublayer_feed_forward_network,
+    )
+
+    return x,
+    (
+        sublayer_self_attention = st_sublayer_self_attention,
+        sublayer_feed_forward_network = st_sublayer_feed_forward_network,
+    )
+end