Transformers reloaded

Author: stefanradev93

bayesflow/networks/embeddings/__init__.py

@@ -1,2 +1,3 @@
 from .fourier_embedding import FourierEmbedding
 from .time2vec import Time2Vec
+from .recurrent_embedding import RecurrentEmbedding

bayesflow/networks/embeddings/recurrent_embedding.py

@@ -0,0 +1,45 @@
+import keras
+from keras.saving import register_keras_serializable as serializable
+
+from bayesflow.types import Tensor
+from bayesflow.utils import expand_tile
+
+
+@serializable(package="bayesflow.networks")
+class RecurrentEmbedding(keras.Layer):
+    """Implements a recurrent network for embedding time."""
+
+    def __init__(self, embed_dim: int = 8, embedding: str = "lstm"):
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.embedding = embedding
+        if embedding == "lstm":
+            self.embedder = keras.layers.LSTM(embed_dim, return_sequences=True)
+        elif embedding == "gru":
+            self.embedder = keras.layers.GRU(embed_dim, return_sequences=True)
+        else:
+            raise ValueError(f"Unknown embedding type {embedding}. Must be in ['lstm', 'gru']")
+
+    def call(self, x: Tensor, t: Tensor = None) -> Tensor:
+        """Creates time representations and concatenates them to x.
+
+        Parameters:
+        -----------
+        x   : Tensor of shape (batch_size, sequence_length, dim)
+            The input sequence.
+        t   : Tensor of shape (batch_size, sequence_length)
+            Vector of times
+
+        Returns:
+        --------
+        emb : Tensor
+            Embedding of shape (batch_size, sequence_length, embed_dim + 1)
+        """
+
+        if t is None:
+            t = keras.ops.linspace(0, keras.ops.shape(x)[1], keras.ops.shape(x)[1], dtype=x.dtype)
+            t = expand_tile(t, keras.ops.shape(x)[0], axis=0)
+
+        emb = self.embedder(t[..., None])
+        return keras.ops.concatenate([x, emb], axis=-1)

bayesflow/networks/embeddings/time2vec.py

@@ -58,8 +58,7 @@ def call(self, x: Tensor, t: Tensor = None) -> Tensor:
         Returns:
         --------
         emb : Tensor
-            Embedding of shape (batch_size, fourier_emb_dim) if `include_identity`
-            is False, else (batch_size, fourier_emb_dim+1)
+            Embedding of shape (batch_size, sequence_length, num_periodic_features + 1)
         """
 
         if t is None:

bayesflow/networks/transformers/fusion_transformer.py

@@ -5,7 +5,6 @@
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
 
-from ..embeddings import Time2Vec
 from ..summary_network import SummaryNetwork
 
 from .mab import MultiHeadAttentionBlock
@@ -28,15 +27,17 @@ def __init__(
         kernel_initializer: str = "he_normal",
         use_bias: bool = True,
         layer_norm: bool = True,
-        t2v_embed_dim: int = 8,
         template_type: str = "lstm",
         bidirectional: bool = True,
         template_dim: int = 128,
         **kwargs,
     ):
-        """Creates a fusion transformer used to flexibly compress time series. If the time intervals vary across
-        batches, it is highly recommended that your simulator also returns a "time" vector denoting absolute or
-        relative time.
+        """Creates a fusion transformer used to flexibly compress time series and learn additional time embeddings
+        using a recurrent neural network. If the time intervals vary across batches, it is highly recommended that
+        your simulator also returns a "time" vector appended to the simulator outputs.
+
+        Important: This network needs at least 2 transformer blocks and will generally be slower than the
+        corresponding TimeSeriesTransformer.
 
         Parameters
         ----------
@@ -73,6 +74,8 @@ def __init__(
         template_dim         : int, optional (default - 128)
             Only used if ``template_type`` in ['lstm', 'gru']. The number of hidden
             units (equiv. output dimensions) of the recurrent network.
+        time_axis     : int, optional (default - None)
+            The time axis (e.g., -1 for last axis) from which to grab the time vector that goes into t2v.
         **kwargs : dict
             Additional keyword arguments passed to the base layer.
         """
@@ -82,9 +85,6 @@ def __init__(
         # Ensure all tuple-settings have the same length
         check_lengths_same(embed_dims, num_heads, mlp_depths, mlp_widths)
 
-        # Initialize Time2Vec embedding layer
-        self.time2vec = Time2Vec(t2v_embed_dim)
-
         # Construct a series of set-attention blocks
         self.attention_blocks = []
         for i in range(len(embed_dims)):
@@ -121,17 +121,13 @@ def __init__(
 
         self.output_projector = keras.layers.Dense(summary_dim)
 
-    def call(self, input_sequence: Tensor, time: Tensor = None, training: bool = False, **kwargs) -> Tensor:
+    def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
 
         Parameters
         ----------
         input_sequence  : Tensor
             Input of shape (batch_size, sequence_length, input_dim)
-        time            : Tensor
-            Time vector of shape (batch_size, sequence_length), optional (default - None)
-            Note: time values for Time2Vec embeddings will be inferred on a linearly spaced
-            interval between [0, sequence length], if no time vector is specified.
         training        : boolean, optional (default - False)
             Passed to the optional internal dropout and spectral normalization
             layers to distinguish between train and test time behavior.
@@ -145,12 +141,12 @@ def call(self, input_sequence: Tensor, time: Tensor = None, training: bool = Fal
             Output of shape (batch_size, set_size, output_dim)
         """
 
-        inp = self.time2vec(input_sequence, t=time)
-        template = self.template_net(inp, training=training)
+        template = self.template_net(input_sequence, training=training)
 
+        rep = input_sequence
         for layer in self.attention_blocks[:-1]:
-            inp = layer(inp, inp, training=training, **kwargs)
+            rep = layer(rep, rep, training=training, **kwargs)
 
-        summary = self.attention_blocks[-1](keras.ops.expand_dims(template, axis=1), inp, training=training, **kwargs)
+        summary = self.attention_blocks[-1](keras.ops.expand_dims(template, axis=1), rep, training=training, **kwargs)
         summary = self.output_projector(keras.ops.squeeze(summary, axis=1))
         return summary

bayesflow/networks/transformers/mab.py

@@ -40,10 +40,7 @@ def __init__(
 
         self.input_projector = layers.Dense(embed_dim)
         self.attention = layers.MultiHeadAttention(
-            key_dim=embed_dim,
-            num_heads=num_heads,
-            dropout=dropout,
-            use_bias=use_bias,
+            key_dim=embed_dim, num_heads=num_heads, dropout=dropout, use_bias=use_bias, output_shape=embed_dim
         )
         self.ln_pre = layers.LayerNormalization() if layer_norm else None
         self.mlp = MLP(

bayesflow/networks/transformers/time_series_transformer.py

@@ -4,7 +4,7 @@
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
 
-from ..embeddings import Time2Vec
+from ..embeddings import Time2Vec, RecurrentEmbedding
 from ..summary_network import SummaryNetwork
 
 from .mab import MultiHeadAttentionBlock
@@ -24,12 +24,14 @@ def __init__(
         kernel_initializer: str = "he_normal",
         use_bias: bool = True,
         layer_norm: bool = True,
-        t2v_embed_dim: int = 8,
+        time_embedding: str = "time2vec",
+        time_embed_dim: int = 8,
+        time_axis: int = None,
         **kwargs,
     ):
         """Creates a regular transformer coupled with Time2Vec embeddings of time used to flexibly compress time series.
         If the time intervals vary across batches, it is highly recommended that your simulator also returns a "time"
-        vector denoting absolute or relative time.
+        vector appended to the simulator outputs and specified via the "time_axis" argument.
 
         Parameters
         ----------
@@ -53,8 +55,14 @@ def __init__(
             Whether to include a bias term in the dense layers.
         layer_norm : bool, optional (default - True)
             Whether to apply layer normalization after the attention and MLP layers.
-        t2v_embed_dim : int, optional (default - 8)
-            The dimensionality of the Time2Vec embedding.
+        time_embedding  : str, optional (default - "time2vec")
+            The type of embedding to use. Must be in ["time2vec", "lstm", "gru"]
+        time_embed_dim  : int, optional (default - 8)
+            The dimensionality of the Time2Vec or recurrent embedding.
+        time_axis     : int, optional (default - None)
+            The time axis (e.g., -1 for last axis) from which to grab the time vector that goes into the embedding.
+            If an embedding is provided and time_axis is None, a uniform time interval between [0, sequence_len]
+            will be assumed.
         **kwargs : dict
             Additional keyword arguments passed to the base layer.
         """
@@ -65,7 +73,14 @@ def __init__(
         check_lengths_same(embed_dims, num_heads, mlp_depths, mlp_widths)
 
         # Initialize Time2Vec embedding layer
-        self.time2vec = Time2Vec(t2v_embed_dim)
+        if time_embedding is None:
+            self.time_embedding = None
+        elif time_embedding == "time2vec":
+            self.time_embedding = Time2Vec(num_periodic_features=time_embed_dim - 1)
+        elif time_embedding in ["lstm", "gru"]:
+            self.time_embedding = RecurrentEmbedding(time_embed_dim, time_embedding)
+        else:
+            raise ValueError("Embedding not found!")
 
         # Construct a series of set-attention blocks
         self.attention_blocks = []
@@ -89,17 +104,15 @@ def __init__(
         self.pooling = keras.layers.GlobalAvgPool1D()
         self.output_projector = keras.layers.Dense(summary_dim)
 
-    def call(self, input_sequence: Tensor, time: Tensor = None, training: bool = False, **kwargs) -> Tensor:
+        self.time_axis = time_axis
+
+    def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
 
         Parameters
         ----------
         input_sequence  : Tensor
             Input of shape (batch_size, sequence_length, input_dim)
-        time            : Tensor
-            Time vector of shape (batch_size, sequence_length), optional (default - None)
-            Note: time values for Time2Vec embeddings will be inferred on a linearly spaced
-            interval between [0, sequence length], if no time vector is specified.
         training        : boolean, optional (default - False)
             Passed to the optional internal dropout and spectral normalization
             layers to distinguish between train and test time behavior.
@@ -113,8 +126,17 @@ def call(self, input_sequence: Tensor, time: Tensor = None, training: bool = Fal
             Output of shape (batch_size, set_size, output_dim)
         """
 
-        # Concatenate learnable time embedding to input sequence
-        inp = self.time2vec(input_sequence, t=time)
+        if self.time_axis is not None:
+            t = input_sequence[..., self.time_axis]
+            indices = list(range(keras.ops.shape(input_sequence)[-1]))
+            indices.pop(self.time_axis)
+            inp = keras.ops.take(input_sequence, indices, axis=-1)
+        else:
+            t = None
+            inp = input_sequence
+
+        if self.time_embedding:
+            inp = self.time_embedding(inp, t=t)
 
         # Apply self-attention blocks
         for layer in self.attention_blocks: