diff --git a/physgrad-comparison.ipynb b/physgrad-comparison.ipynb
index 2f9898c..d49d84a 100644
--- a/physgrad-comparison.ipynb
+++ b/physgrad-comparison.ipynb
@@ -212,14 +212,13 @@
     "For gradient descent, the simple gradient based update from equation {eq}`GD-update`\n",
     "in our setting gives the following update step in $\\mathbf{x}$:\n",
     "\n",
-    "$\\begin{aligned}\n",
-    "\\quad\n",
+    "$$\\begin{aligned}\n",
     "\\Delta \\mathbf{x} \n",
     "&= \n",
     "- \\eta ( J_{L} J_{\\mathbf{z}} )^T  \\\\\n",
     "&=\n",
     "- \\eta ( \\frac{\\partial L }{ \\partial \\mathbf{z} } \\frac{\\partial \\mathbf{z} }{ \\partial \\mathbf{x} }  )^T\n",
-    "\\end{aligned}$\n",
+    "\\end{aligned}$$\n",
     "\n",
     "where $\\eta$ denotes the step size parameter .\n",
     "\n",
@@ -324,15 +323,16 @@
     "\n",
     "For Newton's method, the update step is given by\n",
     "\n",
-    "$\\begin{aligned}\n",
-    "\\quad\n",
+    "$$\n",
+    "\\begin{aligned}\n",
     "\\Delta \\mathbf{x} &= \n",
     "- \\eta \\left( \\frac{\\partial^2 L }{ \\partial \\mathbf{x}^2 }  \\right)^{-1}\n",
     "  \\frac{\\partial L }{ \\partial \\mathbf{x} }\n",
     "\\\\\n",
     "&=\n",
     "- \\eta \\ H_L^{-1} \\ ( J_{L} J_{\\mathbf{z}} )^T\n",
-    "\\end{aligned}$\n",
+    "\\end{aligned}\n",
+    "$$\n",
     "\n",
     "Hence, in addition to the same gradient as for GD, we now need to evaluate and invert the Hessian of $\\frac{\\partial^2 L }{ \\partial \\mathbf{x}^2 }$.\n",
     "\n",
@@ -444,16 +444,15 @@
     "## Physical Gradients\n",
     "\n",
     "Now we also use inverse physics, i.e. the inverse of z:\n",
-    "$\\mathbf{z}^{-1}(\\mathbf{x}) = [x_0 \\ x_1^{1/2}]^T$, to compute the _physical gradient_. As a slight look-ahead to the next section, we'll use a Newton's step for $L$, and combine it with the inverse physics function to get an overall update.\n",
+    "$\\mathbf{z}^{-1}(\\mathbf{x}) = [x_0 \\ x_1^{1/2}]^T$, to compute the _physical gradient_. As a slight look-ahead to the next section, we'll use a Newton's step for $L$, and combine it with the inverse physics function to get an overall update. This gives an update step:\n",
     "\n",
-    "This gives an update step:\n",
-    "$\\begin{aligned}\n",
+    "$$\\begin{aligned}\n",
     "\\Delta \\mathbf{x} &= \n",
     "\\mathbf{z}^{-1} \\left( \\mathbf{z}(\\mathbf{x}) - \\eta\n",
     "  \\left( \\frac{\\partial^2 L }{ \\partial \\mathbf{z}^2 }  \\right)^{-1}\n",
     "  \\frac{\\partial L }{ \\partial \\mathbf{z} }\n",
     "\\right) - \\mathbf{x}\n",
-    "\\end{aligned}$\n",
+    "\\end{aligned}$$\n",
     "\n",
     "Below, we define our inverse function `fun_z_inv_analytic` (we'll come to a variant below), and then evaluate an optimization with the physical gradient for ten steps:\n"
    ]