Add use_causal_mask to MultiHeadAttention and approximate GELU

include/fdeep/import_model.hpp

@@ -1329,11 +1329,11 @@ namespace internal {
         const get_param_f&, const nlohmann::json& data,
         const std::string& name)
     {
+        bool approximate = false;
         if (json_obj_has_member(data, "config") && json_obj_has_member(data["config"], "approximate") && !data["config"]["approximate"].is_null()) {
-            const bool approximate = data["config"]["approximate"];
-            assertion(approximate == false, "Gelu with approximate = True is not supported.");
+            approximate = data["config"]["approximate"];
         }
-        return std::make_shared<gelu_layer>(name);
+        return std::make_shared<gelu_layer>(name, approximate);
     }
 
     inline activation_layer_ptr create_softsign_layer(
@@ -1473,6 +1473,9 @@ namespace internal {
         const std::size_t key_dim = data["config"]["key_dim"];
         const std::size_t value_dim = data["config"]["value_dim"];
         const bool use_bias = data["config"]["use_bias"];
+        const bool use_causal_mask = json_obj_has_member(data["config"], "use_causal_mask")
+            ? data["config"]["use_causal_mask"].get<bool>()
+            : false;
         const auto weight_shapes = create_vector<std::vector<std::size_t>>(fplus::bind_1st_of_2(
                                                                                create_vector<std::size_t, decltype(create_size_t)>, create_size_t),
             get_param(name, "weight_shapes"));
@@ -1485,7 +1488,8 @@ namespace internal {
             },
             weight_shapes, weight_values);
         return std::make_shared<multi_head_attention_layer>(name,
-            num_heads, key_dim, value_dim, use_bias, weights_and_biases);
+            num_heads, key_dim, value_dim, use_bias, use_causal_mask,
+            weights_and_biases);
     }
 
     inline layer_ptr create_lstm_layer(const get_param_f& get_param,

include/fdeep/layers/gelu_layer.hpp

@@ -17,16 +17,20 @@ namespace internal {
 
     class gelu_layer : public activation_layer {
     public:
-        explicit gelu_layer(const std::string& name)
+        explicit gelu_layer(const std::string& name, bool approximate = false)
             : activation_layer(name)
+            , approximate_(approximate)
         {
         }
 
     protected:
         tensor transform_input(const tensor& in_vol) const override
         {
-            return transform_tensor(gelu_activation, in_vol);
+            return approximate_
+                ? transform_tensor(gelu_approximate_activation, in_vol)
+                : transform_tensor(gelu_activation, in_vol);
         }
+        bool approximate_;
     };
 
 }

include/fdeep/layers/multi_head_attention_layer.hpp

@@ -19,11 +19,13 @@ namespace internal {
     public:
         explicit multi_head_attention_layer(const std::string& name,
             std::size_t num_heads, std::size_t key_dim, std::size_t value_dim,
-            bool use_bias, const std::vector<tensor>& weights_and_biases)
+            bool use_bias, bool use_causal_mask,
+            const std::vector<tensor>& weights_and_biases)
             : layer(name)
             , num_heads_(num_heads)
             , key_dim_(key_dim)
             , value_dim_(value_dim)
+            , use_causal_mask_(use_causal_mask)
             , query_dense_(create_dense_layers(weights_and_biases, use_bias, num_heads, 0, key_dim, name + "_query_dense"))
             , value_dense_(create_dense_layers(weights_and_biases, use_bias, num_heads, 2, value_dim, name + "_value_dense"))
             , key_dense_(create_dense_layers(weights_and_biases, use_bias, num_heads, 1, key_dim, name + "_key_dense"))
@@ -88,12 +90,35 @@ namespace internal {
             // https://dmol.pub/dl/attention.html#multi-head-attention-block
             // https://github.com/keras-team/keras/blob/v2.14.0/keras/layers/attention/multi_head_attention.py
             // https://gist.github.com/sevagh/b71d253a347a9b59c026580625452fc5
-            const tensor scores = dot_product_tensors(query, transpose(key), std::vector<int>({ 2, 1 }), false);
+            tensor scores = dot_product_tensors(query, transpose(key), std::vector<int>({ 2, 1 }), false);
             const std::size_t query_size = query.shape().depth_;
-            const tensor distribution = softmax(transform_tensor(fplus::multiply_with(1 / std::sqrt(query_size)), scores));
+            scores = transform_tensor(fplus::multiply_with(1 / std::sqrt(query_size)), scores);
+            if (use_causal_mask_) {
+                apply_causal_mask(scores);
+            }
+            const tensor distribution = softmax(scores);
             return dot_product_tensors(distribution, value, std::vector<int>({ 2, 1 }), false);
         }
 
+        static void apply_causal_mask(tensor& scores)
+        {
+            // Scores have shape (..., T, S) with width=T (query positions)
+            // and depth=S (key positions). Mask out s > t.
+            const auto& s = scores.shape();
+            const float_type neg_inf = -std::numeric_limits<float_type>::infinity();
+            for (std::size_t d5 = 0; d5 < s.size_dim_5_; ++d5) {
+                for (std::size_t d4 = 0; d4 < s.size_dim_4_; ++d4) {
+                    for (std::size_t y = 0; y < s.height_; ++y) {
+                        for (std::size_t t = 0; t < s.width_; ++t) {
+                            for (std::size_t k = t + 1; k < s.depth_; ++k) {
+                                scores.set_ignore_rank(tensor_pos(d5, d4, y, t, k), neg_inf);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
     protected:
         tensors apply_impl(const tensors& input) const override
         {
@@ -111,6 +136,7 @@ namespace internal {
         std::size_t num_heads_;
         std::size_t key_dim_;
         std::size_t value_dim_;
+        bool use_causal_mask_;
         std::vector<dense_layer> query_dense_;
         std::vector<dense_layer> value_dense_;
         std::vector<dense_layer> key_dense_;

include/fdeep/recurrent_ops.hpp

@@ -75,6 +75,15 @@ namespace internal {
         return static_cast<float_type>(0.5) * x * (static_cast<float_type>(1) + static_cast<float_type>(std::erf(x / std::sqrt(static_cast<float_type>(2)))));
     }
 
+    inline float_type gelu_approximate_activation(float_type x)
+    {
+        // 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+        const float_type sqrt_2_over_pi = static_cast<float_type>(0.7978845608028654);
+        const float_type c = static_cast<float_type>(0.044715);
+        const float_type inner = sqrt_2_over_pi * (x + c * x * x * x);
+        return static_cast<float_type>(0.5) * x * (static_cast<float_type>(1) + std::tanh(inner));
+    }
+
     inline float_type softsign_activation(float_type x)
     {
         return x / (std::abs(x) + static_cast<float_type>(1));

keras_export/convert_model.py

@@ -11,11 +11,21 @@
 
 import numpy as np
 import numpy.typing  # pylint: disable=unused-import
+import keras
 from keras import backend as K, Layer
 from keras.layers import Input, Embedding, CategoryEncoding
 from keras.models import Model, load_model
 from keras.src import Functional
 
+
+@keras.saving.register_keras_serializable(package="fdeep")
+def gelu_approximate(x):  # type: ignore[no-untyped-def]
+    """Tanh-approximation form of GELU, registered so saved models that use
+    it can be loaded by this converter. The conversion step rewrites it to a
+    plain ``gelu`` activation with an extra ``approximate=True`` flag in the
+    layer config, which the C++ runtime picks up."""
+    return keras.activations.gelu(x, approximate=True)
+
 __author__ = "Tobias Hermann"
 __copyright__ = "Copyright 2017, Tobias Hermann"
 __license__ = "MIT"
@@ -968,6 +978,62 @@ def calculate_hash(model: Model) -> str:
     return hash_m.hexdigest()
 
 
+_ACTIVATION_FUNCTION_REWRITES = {
+    'fdeep>gelu_approximate': ('gelu', {'approximate': True}),
+}
+
+
+def rewrite_custom_activations(arch: Any) -> None:
+    """Walk the serialized architecture and replace registered custom
+    activation functions with a plain string activation plus extra config
+    flags, so the C++ runtime sees a known activation type."""
+    if isinstance(arch, dict):
+        cfg = arch.get('config')
+        if isinstance(cfg, dict):
+            act = cfg.get('activation')
+            if isinstance(act, dict) and act.get('class_name') == 'function':
+                rewrite = _ACTIVATION_FUNCTION_REWRITES.get(act.get('config'))
+                if rewrite is not None:
+                    new_name, extra = rewrite
+                    cfg['activation'] = new_name
+                    cfg.update(extra)
+        for v in arch.values():
+            rewrite_custom_activations(v)
+    elif isinstance(arch, list):
+        for v in arch:
+            rewrite_custom_activations(v)
+
+
+def inject_mha_call_kwargs(arch: Any, model: Model) -> None:
+    """``use_causal_mask`` is a per-call kwarg of ``MultiHeadAttention``, not a
+    layer config field, so ``model.to_json()`` doesn't include it. Look the
+    flag up from the layer's first inbound node and bake it into the JSON
+    config so the C++ runtime can read it."""
+    if isinstance(arch, dict):
+        if arch.get('class_name') == 'MultiHeadAttention':
+            cfg = arch.get('config') or {}
+            layer_name = cfg.get('name')
+            if layer_name:
+                try:
+                    layer = model.get_layer(layer_name)
+                except (ValueError, KeyError):
+                    layer = None
+                if layer is not None and layer._inbound_nodes:
+                    kwargs = layer._inbound_nodes[0].arguments.kwargs
+                    if kwargs.get('use_causal_mask'):
+                        cfg['use_causal_mask'] = True
+                    if kwargs.get('attention_mask') is not None:
+                        raise NotImplementedError(
+                            f"MultiHeadAttention layer {layer_name!r} was "
+                            "called with an explicit attention_mask, which "
+                            "the frugally-deep runtime does not yet support.")
+        for v in arch.values():
+            inject_mha_call_kwargs(v, model)
+    elif isinstance(arch, list):
+        for v in arch:
+            inject_mha_call_kwargs(v, model)
+
+
 def model_to_fdeep_json(model: Model, no_tests: bool = False) -> Mapping[str, Any]:
     """Convert any Keras model to the frugally-deep model format."""
 
@@ -983,6 +1049,8 @@ def model_to_fdeep_json(model: Model, no_tests: bool = False) -> Mapping[str, An
     json_output = {}
     print('Converting model architecture.')
     json_output['architecture'] = json.loads(model.to_json())
+    rewrite_custom_activations(json_output['architecture'])
+    inject_mha_call_kwargs(json_output['architecture'], model)
     json_output['image_data_format'] = K.image_data_format()
     json_output['input_shapes'] = list(map(get_layer_input_shape_tensor_shape, get_model_input_layers(model)))
     json_output['output_shapes'] = list(map(keras_shape_to_fdeep_tensor_shape, as_list(model.output_shape)))

keras_export/generate_test_models.py

@@ -5,6 +5,11 @@
 import sys
 from typing import Tuple, List, Union
 
+# convert_model lives next to this script; importing it pulls in the
+# ``gelu_approximate`` serializable activation that the converter then
+# rewrites to ``gelu`` with ``approximate=True``.
+from convert_model import gelu_approximate
+
 import numpy as np
 from keras import activations
 from keras.layers import ActivityRegularization
@@ -413,6 +418,8 @@ def get_test_model_exhaustive() -> Model:
     outputs.append(Dense(3, use_bias=True)(inputs[14]))
     outputs.append(Dense(4, use_bias=False)(inputs[16]))
     outputs.append(Dense(4, use_bias=False, activation='tanh')(inputs[18]))
+    outputs.append(Dense(4, use_bias=True, activation='gelu')(inputs[18]))
+    outputs.append(Dense(4, use_bias=True, activation=gelu_approximate)(inputs[18]))
     outputs.append(Dense(4, use_bias=False)(inputs[20]))
 
     outputs.append(Reshape(((2 * 3 * 4 * 5 * 6),))(inputs[0]))
@@ -541,6 +548,20 @@ def get_test_model_exhaustive() -> Model:
         num_heads=2, key_dim=3, value_dim=5,
         use_bias=True, output_shape=None, attention_axes=None)(inputs[49], inputs[50], inputs[51]))
 
+    # use_causal_mask=True: triangular attention so position i only attends to <=i.
+    outputs.append(MultiHeadAttention(
+        num_heads=1, key_dim=2, value_dim=None,
+        use_bias=True, output_shape=None, attention_axes=None)(
+            inputs[50], inputs[50], use_causal_mask=True))
+    outputs.append(MultiHeadAttention(
+        num_heads=2, key_dim=3, value_dim=4,
+        use_bias=True, output_shape=None, attention_axes=None)(
+            inputs[50], inputs[50], inputs[50], use_causal_mask=True))
+    outputs.append(MultiHeadAttention(
+        num_heads=3, key_dim=2, value_dim=2,
+        use_bias=False, output_shape=None, attention_axes=None)(
+            inputs[50], inputs[50], use_causal_mask=True))
+
     # GroupedQueryAttention: shared Q/K/V seq, separate K/V seq, with/without gate.
     outputs.append(GroupQueryAttention(head_dim=4, num_query_heads=6,
         num_key_value_heads=2)(inputs[49], inputs[49], inputs[49]))