multihead_attention: update attention in accordance with linear2d

Author: OneAdder

src/nf/nf_cross_attention_layer.f90

@@ -42,14 +42,14 @@ module function cross_attention_layer_cons(sequence_length, model_dimension, n_h
     end if
     res % head_size = model_dimension / n_heads
 
-    res % query_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % key_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % value_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % output_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    call res % query_layer % init([0])
-    call res % key_layer % init([0])
-    call res % value_layer % init([0])
-    call res % output_layer % init([0])
+    res % query_layer = linear2d_layer(model_dimension)
+    res % key_layer = linear2d_layer(model_dimension)
+    res % value_layer = linear2d_layer(model_dimension)
+    res % output_layer = linear2d_layer(model_dimension)
+    call res % query_layer % init([sequence_length, model_dimension])
+    call res % key_layer % init([sequence_length, model_dimension])
+    call res % value_layer % init([sequence_length, model_dimension])
+    call res % output_layer % init([sequence_length, model_dimension])
 
     res % softmax_func = softmax()
   end function cross_attention_layer_cons

src/nf/nf_multihead_attention_submodule.f90

@@ -20,14 +20,14 @@ module function multihead_attention_layer_cons(sequence_length, model_dimension,
     end if
     res % head_size = model_dimension / n_heads
 
-    res % query_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % key_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % value_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % output_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    call res % query_layer % init([0])
-    call res % key_layer % init([0])
-    call res % value_layer % init([0])
-    call res % output_layer % init([0])
+    res % query_layer = linear2d_layer(model_dimension)
+    res % key_layer = linear2d_layer(model_dimension)
+    res % value_layer = linear2d_layer(model_dimension)
+    res % output_layer = linear2d_layer(model_dimension)
+    call res % query_layer % init([sequence_length, model_dimension])
+    call res % key_layer % init([sequence_length, model_dimension])
+    call res % value_layer % init([sequence_length, model_dimension])
+    call res % output_layer % init([sequence_length, model_dimension])
 
     res % softmax_func = softmax()
   end function multihead_attention_layer_cons

src/nf/nf_self_attention_layer.f90

@@ -42,14 +42,14 @@ module function self_attention_layer_cons(sequence_length, model_dimension, n_he
     end if
     res % head_size = model_dimension / n_heads
 
-    res % query_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % key_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % value_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    res % output_layer = linear2d_layer(sequence_length, model_dimension, model_dimension)
-    call res % query_layer % init([0])
-    call res % key_layer % init([0])
-    call res % value_layer % init([0])
-    call res % output_layer % init([0])
+    res % query_layer = linear2d_layer(model_dimension)
+    res % key_layer = linear2d_layer(model_dimension)
+    res % value_layer = linear2d_layer(model_dimension)
+    res % output_layer = linear2d_layer(model_dimension)
+    call res % query_layer % init([sequence_length, model_dimension])
+    call res % key_layer % init([sequence_length, model_dimension])
+    call res % value_layer % init([sequence_length, model_dimension])
+    call res % output_layer % init([sequence_length, model_dimension])
 
     res % softmax_func = softmax()
   end function self_attention_layer_cons

test/test_multihead_attention_layer.f90

@@ -16,6 +16,7 @@ program test_multihead_attention_layer
 
   attention = multihead_attention_layer(sequence_length=3, model_dimension=4, n_heads=2)
   call attention % init_base([0])
+  call set_weights(attention)
 
   call test_multihead_attention_split_heads(attention, sample_input, ok, split_heads_output)
   call test_multihead_attention_create_attention_matrix(attention, split_heads_output, ok)
@@ -30,6 +31,18 @@ program test_multihead_attention_layer
   call test_cross_attention(ok)
 
 contains
+  subroutine set_weights(attention)
+    type(multihead_attention_layer), intent(in out) :: attention
+    attention % query_layer % weights = 0.1
+    attention % key_layer % weights = 0.1
+    attention % value_layer % weights = 0.1
+    attention % output_layer % weights = 0.1
+    attention % query_layer % biases = 0.11
+    attention % key_layer % biases = 0.11
+    attention % value_layer % biases = 0.11
+    attention % output_layer % biases = 0.11
+  end subroutine set_weights
+
   subroutine test_multihead_attention_split_heads(attention, input, ok, output)
     type(multihead_attention_layer), intent(in) :: attention
     real, intent(in) :: input(:, :)
@@ -199,6 +212,7 @@ subroutine test_multihead_attention_forward_reallife_shape(ok)
 
     attention = multihead_attention_layer(sequence_length=148, model_dimension=512, n_heads=8)
     call attention % init_base([0])
+    call set_weights(attention)
 
     call attention % common_forward(input, input, input)
 
@@ -305,6 +319,14 @@ subroutine test_self_attention(ok)
 
     attention = self_attention_layer(sequence_length=2, model_dimension=3, n_heads=1)
     call attention % init([0])
+    attention % query_layer % weights = 0.1
+    attention % key_layer % weights = 0.1
+    attention % value_layer % weights = 0.1
+    attention % output_layer % weights = 0.1
+    attention % query_layer % biases = 0.11
+    attention % key_layer % biases = 0.11
+    attention % value_layer % biases = 0.11
+    attention % output_layer % biases = 0.11
 
     call attention % forward(input)
     output_flat = reshape(attention % output, shape(output_flat))
@@ -346,6 +368,14 @@ subroutine test_cross_attention(ok)
 
     attention = cross_attention_layer(sequence_length=2, model_dimension=3, n_heads=1)
     call attention % init([0])
+    attention % query_layer % weights = 0.1
+    attention % key_layer % weights = 0.1
+    attention % value_layer % weights = 0.1
+    attention % output_layer % weights = 0.1
+    attention % query_layer % biases = 0.11
+    attention % key_layer % biases = 0.11
+    attention % value_layer % biases = 0.11
+    attention % output_layer % biases = 0.11
 
     call attention % forward(input)
     output_flat = reshape(attention % output, shape(output_flat))