diff --git a/README.rst b/README.rst
index f2df90e..9712611 100644
--- a/README.rst
+++ b/README.rst
@@ -12,7 +12,7 @@
 
 
 
-This is a Python package for CStree models :footcite:p:`duarte2021representation`, a family of graphical causal models that encode context-specific dependence for multivariate multinomial distributions.
+This is a Python package for CStree models :footcite:p:`duarte2021representation` :footcite:p:`rios2024scalable` a family of graphical causal models that encode context-specific dependence for multivariate multinomial distributions.
 
 .. As not all staged tree models admit this property, CStrees are a subclass that provides a transparent, intuitive and compact representation of context-specific causal information.
 
diff --git a/docs/_static/bibtex.bib b/docs/_static/bibtex.bib
index 1fda854..a687d8b 100644
--- a/docs/_static/bibtex.bib
+++ b/docs/_static/bibtex.bib
@@ -1,7 +1,17 @@
-@article{duarte2021representation,
+@misc{duarte2021representation,
   title={Representation of context-specific causal models with observational and interventional data},
   author={Duarte, Eliana and Solus, Liam},
-  journal={arXiv preprint arXiv:2101.09271},
+  archivePrefix={arXiv},
   year={2021},
-  url={https://arxiv.org/abs/2101.09271}
+  primaryClass={stat.ML},
+  eprint={2101.09271}
 }
+
+@misc{rios2024scalable,
+      title={Scalable Structure Learning for Sparse Context-Specific Causal Systems}, 
+      author={Felix L. Rios and Alex Markham and Liam Solus},
+      year={2024},
+      eprint={2402.07762},
+      archivePrefix={arXiv},
+      primaryClass={stat.ML}
+}
\ No newline at end of file
diff --git a/docs/source/learn_demo.ipynb b/docs/source/learn_demo.ipynb
index 5d6eb58..ddf22c5 100644
--- a/docs/source/learn_demo.ipynb
+++ b/docs/source/learn_demo.ipynb
@@ -12,7 +12,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This notebook shows how to learn a CStree from observational data using an exhaustive search. procedure. "
+    "This notebook shows how to learn a CStree from observational data using an exhaustive search procedure. "
    ]
   },
   {
diff --git a/src/cstrees/__meta__.py b/src/cstrees/__meta__.py
index 6bd7bb5..c16fa39 100644
--- a/src/cstrees/__meta__.py
+++ b/src/cstrees/__meta__.py
@@ -1,3 +1,3 @@
 # Automatically created. Please do not edit.
-__version__ = '1.1.1'
+__version__ = '1.1.2'
 __author__ = ''
diff --git a/src/cstrees/learning.py b/src/cstrees/learning.py
index d8ff61f..5761a08 100644
--- a/src/cstrees/learning.py
+++ b/src/cstrees/learning.py
@@ -162,10 +162,10 @@ def _optimal_staging_at_level(
                 staging_score = context_scores["scores"][var]["None"]
                 continue
 
-            # here we (=I) anyway extract just the context, so the stage format
-            # is a bit redundant.
-            stage_context = sc._stage_to_context_key(stage, order)
-            score = context_scores["scores"][var][stage_context]
+            # here we (=I) anyway extract just the context, so the stage format is a bit redundant.
+            stage_context = sc._stage_to_context_key(stage, order) 
+            if stage_context in context_scores["scores"][var]:
+                score = context_scores["scores"][var][stage_context]
             staging_score += score
 
         # Update the max score and the max staging
diff --git a/src/cstrees/scoring.py b/src/cstrees/scoring.py
index 9cf1be7..3ad390c 100644
--- a/src/cstrees/scoring.py
+++ b/src/cstrees/scoring.py
@@ -518,7 +518,8 @@ def order_score_tables(
                 for stage in staging:
                     # OK! even when restricting to some possible cvars
                     stage_context = _stage_to_context_key(stage, subset)
-                    staging_unnorm_post += context_scores["scores"][var][stage_context]
+                    if stage_context in context_scores["scores"][var]:
+                        staging_unnorm_post += context_scores["scores"][var][stage_context]
 
                 if i == 0:
                     order_scores["scores"][var][subset_str] = staging_unnorm_post