Improved simformer advanced tutorial notebook with later modifications

Author: nMaax

docs/advanced_tutorials/22_simformer.ipynb

@@ -8,13 +8,13 @@
     "\n",
     "`sbi` incorporates also the Simformer, a novel framework to provide an \"all-in-one\" solution for Simulation-Based Inference.\n",
     "\n",
-    "Similarly to the NPSE and the FMPE algorithms, the Simformer leverage score-based generative models to sample from the posteriors. More specifically, the underlying architecture is a Transformer.\n",
+    "Similarly to the NPSE and the FMPE algorithms, the Simformer leverage score-based and flow-matching generative models to sample from the conditional distribution, leveraging a Transformer-based architecture.\n",
     "\n",
-    "While NPSE and FMPE replace Normalizing Flows as the core conditional density estimator for the posterior $p(\\theta|x)$, **Simformer** takes a more fundamental step. It aims to learn a single, unified score model that can implicitly represent **any joint distribution** of $\\theta$ and $x$, thereby unlocking a versatile range of inference capabilities from a single trained model.\n",
+    "While NPSE and FMPE limit themselves to predict the posterior $p(\\theta|x)$ distributuion, the **Simformer** takes a more fundamental step. It aims to learn a single, unified model that can implicitly represent **any joint distribution** of $\\theta$ and $x$, thereby unlocking a   range of inference capabilities from a single trained model.\n",
     "\n",
-    "At its heart, Simformer is a **Masked Conditional Score Model**. That is, instead of treating parameters ($\\theta$) and data ($x$) separately, Simformer rather expects them as a single input vector, such a concatenation $\\hat{\\mathbf{x}} = [\\theta, x]$. Which is then accompanied with two other tensors that represents dependencies between inputs:\n",
+    "At its heart, Simformer is a **Masked Conditional Model**. That is, instead of treating parameters ($\\theta$) and data ($x$) separately, Simformer rather expects them as a single input vector, such a concatenation $\\hat{\\mathbf{x}} = [\\theta, x]$. Which is then accompanied with two other tensors that represents dependencies between inputs:\n",
     "\n",
-    "- The **Condition Mask ($M_C$)**, which explicitly designates which variables in $\\hat{\\mathbf{x}}$ are **observed** and which are **latent (to be inferred)**.\n",
+    "- The **Condition Mask ($M_C$)**, which explicitly designates which variables in $\\hat{\\mathbf{x}}$ are **observed (to be conditioned on)** and which are **latent (to be inferred)**.\n",
     "- The **Edge Mask ($M_E$)**, which inject prior knowledge about the **causal or statistical dependencies** between variables in $\\hat{\\mathbf{x}}$. By specifying which variables can \"attend\" to others in the Transformer aligment mechanism.\n",
     "\n",
     "Most importantly, the Simformer can be trained using different condition and edge masks at different training steps, generating a final neural network which did not limit its learning process to a specific inference setting, but can be flexibly re-used with arbitrary condition and edge masks.\n",
@@ -25,6 +25,13 @@
     "- Gloeckler M. et al. \"All-in-one simulation-based inference.\" ICML 2024."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We first start importing dependencies and preparing a simulator."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -61,7 +68,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Differently from other sbi methods, the Simformer works by means unflattened variable features, i.e., we must train the Simformer on simulations that provide features of each variable in separate, trailing dimension."
+    "Differently from other sbi methods, the Simformer works by means unflattened variable features, i.e., we must train the Simformer on simulations that provide features of each variable in separate, trailing dimension. For our walkthrough we will implement the score-based Simformer, the flow-matching equivalent works by the same exact API."
    ]
   },
   {
@@ -85,7 +92,7 @@
    "source": [
     "As for NPSE, the Simformer approximates the posterior distribution by learning its score function, i.e., gradient of the log-density, using the denoising score matching loss. Refer to [19_flowmatching_and_scorematching](#) and [20_score_based_methods_new_features.ipynb](#) for understanding denoising score matching.\n",
     "\n",
-    "Note also that only the single-round version of Simformer is implemented currently."
+    ">Note also that currently only the single-round version of Simformer is implemented currently."
    ]
   },
   {
@@ -95,7 +102,7 @@
    "outputs": [],
    "source": [
     "# Instantiate Simformer and append simulations\n",
-    "inference = Simformer(device='gpu')\n",
+    "inference = Simformer()\n",
     "inference.append_simulations(inputs)"
    ]
   },
@@ -105,9 +112,8 @@
    "source": [
     "Notice how we only appended the `inputs` tensor, without providing any knowledge of dependencies within variables, or which are observed or latent (i.e., to be infered) in our problem. In other words, we did not pass any condition or edge mask.\n",
     "\n",
-    "By now, the Simformer will automatically generate condition and edge masks, more specifically:\n",
-    "\n",
-    "- Condition masks will be generated according to a $\\text{Bernoulli}(p=0.5)$, for degenerate cases in which the entire mask of a given sample collapse to a full latent set, or full observed—i.e. full ones or full zeroes, a mechanism that randomly flips one variable take place\n",
+    "Condition and edge masks will be rather generated at trainig time. Such generation can be ruled by the user passing either a `Callable`, a collection of masks (`list`, `set`, `Tensor`), or a fixed `Tensor`. If we do not specify anything sbi will automatically generate them:\n",
+    "- Condition masks will be generated according to a $\\text{Bernoulli}(p=0.5)$, for degenerate cases in which the entire mask of a given sample collapse to full observed—i.e. full ones, a mechanism that randomly flips one variable at random take place\n",
     "- Edge masks will be `None` (equivalent to a full tensor of ones, i.e., full attention)"
    ]
   },
@@ -118,7 +124,7 @@
    "outputs": [],
    "source": [
     "# Train the score estimator\n",
-    "score_estimator = inference.train()"
+    "score_estimator = inference.train(max_num_epochs=100)"
    ]
   },
   {
@@ -200,7 +206,7 @@
     "\n",
     "While the example provided is trivial, the Simformer becomes quite more interesting for more complex models, where there could exist more that two variables involved. And where the conditional distribution that one needs to infer is not assumed to be already choosen.\n",
     "\n",
-    "Nonetheless, the Simformer still allows to call the two common `build_posterior()` and `build_likelihood()` methods. As long as we specifiy which variables are intended as latent or observed from the posterior perspective."
+    "Nonetheless, the Simformer still allows to call the two common `build_posterior()` and `build_likelihood()` methods. As long as we first inform the Simformer which variables are intended as latent or observed **from the posterior perspective**."
    ]
   },
   {
@@ -216,7 +222,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This setting of latent and posterior indexes *do not affect* the behaviour of the `build_conditional()` method, as it only relies on the masks passed when it is called. It simply informs the Simformer on how to interpret variables when `build_posterior()` and `build_likelihood()` are called. Crucially, note that the indexes passed must be interpret as latent and observed from the posterior perspective. These can also be set at init time through `posterior_latent_idx` and `posterior_observed_idx` parameters. "
+    "This setting of latent and posterior indexes *do not affect* the behaviour of the `build_conditional()` method, as it only relies on the masks passed when it is called. It simply informs the Simformer on how to interpret variables when `build_posterior()` and `build_likelihood()` are called. Again, note that the indexes passed must be interpret as latent and observed **from the posterior perspective**. These can also be set at Simformer init time through `posterior_latent_idx` and `posterior_observed_idx` parameters. "
    ]
   },
   {
@@ -243,6 +249,13 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Indeed, the posterior just found is similar to the conditional infered in previous steps. We can similarly build the likelihood again."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,