SlowGuess/NeurIPS_2023_OCR huggingface dataset:
which is evaluated at a given point  $(x,y)$  in the domain. The complete space  $\mathbb{R}^p\times \mathcal{U}^{(n)}$  is often called the  $n$ -th order jet space [11].
A system of differential equations is a set of  $l$  differential equations  $\Delta : \mathbb{R}^p \times \mathcal{U}^{(n)} \to \mathbb{R}^l$  of the independent and dependent variables with dependence on the derivatives up to a maximum order of  $n$ :
$$
\Delta_ {\nu} (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}) = 0, \quad \nu = 1, \dots , l. \tag {13}
$$
A smooth solution is thus a function  $f$  such that for all points in the domain of  $x$ :
$$
\Delta_ {\nu} (\boldsymbol {x}, \operatorname {p r} ^ {(n)} f (\boldsymbol {x})) = 0, \quad \nu = 1, \dots , l. \tag {14}
$$
In geometric terms, the system of differential equations states where the given map  $\Delta$  vanishes on the jet space, and forms a subvariety
$$
Z _ {\Delta} = \left\{\left(\boldsymbol {x}, \boldsymbol {u} ^ {(n)}\right): \Delta \left(\boldsymbol {x}, \boldsymbol {u} ^ {(n)}\right) = 0 \right\} \subset \mathbb {R} ^ {p} \times \mathcal {U} ^ {(n)}. \tag {15}
$$
Therefore to check if a solution is valid, one can check if the prolongation of the solution falls within the subvariety  $Z_{\Delta}$ . As an example, consider the one-dimensional heat equation
$$
\Delta = u _ {t} - c u _ {x x} = 0. \tag {16}
$$
We can check that  $f(x,t) = \sin (x)e^{-ct}$  is a solution by forming its prolongation and checking if it falls within the subvariety given by the above equation:
$$
\operatorname {p r} ^ {(2)} f (x, t) = \left(\sin (x) e ^ {- c t}; \cos (x) e ^ {- c t}, - c \sin (x) e ^ {- c t}; - \sin (x) e ^ {- c t}, - c \cos (x) e ^ {- c t}, c ^ {2} \sin (x) e ^ {- c t}\right),
$$
$$
\Delta (x, t, \boldsymbol {u} ^ {(n)}) = - c \sin (x) e ^ {- c t} + c \sin (x) e ^ {- c t} = 0. \tag {17}
$$
A.1 Symmetry Groups and Infinitesimal Invariance
A symmetry group  $G$  for a system of differential equations is a set of local transformations to the function which transform one solution of the system of differential equations to another. The group takes the form of a Lie group, where group operations can be expressed as a composition of one-parameter transforms. More rigorously, given the graph of a solution  $\Gamma_f$  as defined in Eq. (10), a group operation  $g \in G$  maps this graph to a new graph
$$
g \cdot \Gamma_ {f} = \left\{\left(\tilde {\boldsymbol {x}}, \tilde {\boldsymbol {u}}\right) = g \cdot \left(\boldsymbol {x}, \boldsymbol {u}\right): \left(\boldsymbol {x}, \boldsymbol {u}\right) \in \Gamma_ {f} \right\}, \tag {18}
$$
where  $(\tilde{\pmb{x}},\tilde{\pmb{u}})$  label the new coordinates of the solution in the set  $g\cdot \Gamma_f$ . For example, if  $\pmb {x} = (x,t)$ ,  $u = u(x,t)$ , and  $g$  acts on  $(\pmb {x},u)$  via
$$
(x, t, u) \mapsto (x + \epsilon t, t, u + \epsilon),
$$
then  $\tilde{u} (\tilde{x},\tilde{t}) = u(x,t) + \epsilon = u(\tilde{x} -\epsilon \tilde{t},\tilde{t}) + \epsilon$  , where  $(\tilde{x},\tilde{t}) = (x + \epsilon t,t)$
Note, that the set  $g \cdot \Gamma_f$  may not necessarily be a graph of a new  $x$ -valued function; however, since all transformations are local and smooth, one can ensure transformations are valid in some region near the identity of the group.
As an example, consider the following transformations which are members of the symmetry group of the differential equation  $u_{xx} = 0$ .  $g_1(t)$  translates a single spatial coordinate  $x$  by an amount  $t$  and  $g_2$  scales the output coordinate  $u$  by an amount  $e^r$ :
$$
g _ {1} (t) \cdot (x, u) = (x + t, u), \tag {19}
$$
$$
g _ {2} (r) \cdot (x, u) = (x, e ^ {r} \cdot u).
$$
It is easy to verify that both of these operations are local and smooth around a region of the identity, as sending  $r, t \to 0$  recovers the identity operation. Lie theory allows one to equivalently describe
the potentially nonlinear group operations above with corresponding infinitesimal generators of the group action, corresponding to the Lie algebra of the group. Infinitesimal generators form a vector field over the total space  $\Omega \times \mathcal{U}$ , and the group operations correspond to integral flows over that vector field. To map from a single parameter Lie group operation to its corresponding infinitesimal generator, we take the derivative of the single parameter operation at the identity:
$$
\left. \boldsymbol {v} _ {g} \right| _ {(x, u)} = \frac {d}{d t} g (t) \cdot (x, u) \Big | _ {t = 0}, \tag {20}
$$
where  $g(0)\cdot (x,u) = (x,u)$
To map from the infinitesimal generator back to the corresponding group operation, one can apply the exponential map
$$
\exp (t \boldsymbol {v}) \cdot (\boldsymbol {x}, \boldsymbol {u}) = g (t) \cdot (\boldsymbol {x}, \boldsymbol {u}), \tag {21}
$$
where  $\exp : \mathfrak{g} \to G$ . Here,  $\exp(\cdot)$  maps from the Lie algebra,  $\mathfrak{g}$ , to the corresponding Lie group,  $G$ . This exponential map can be evaluated using various methods, as detailed in Appendix B and Appendix E.
Returning to the example earlier from Equation (19), the corresponding Lie algebra elements are
$$
\boldsymbol {v} _ {g _ {1}} = \partial_ {x} \leftrightarrow g _ {1} (t) \cdot (x, u) = (x + t, u), \tag {22}
$$
$$
\boldsymbol {v} _ {g _ {2}} = u \partial_ {u} \leftrightarrow g _ {2} (r) \cdot (x, u) = (x, e ^ {r} \cdot u).
$$
Informally, Lie algebras help simplify notions of invariance as it allows one to check whether functions or differential equations are invariant to a group by needing only to check it at the level of the derivative of that group. In other words, for any vector field corresponding to a Lie algebra element, a given function is invariant to that vector field if the action of the vector field on the given function evaluates to zero everywhere. Thus, given a symmetry group, one can determine a set of invariants using the vector fields corresponding to the infinitesimal generators of the group. To determine whether a differential equation is in such a set of invariants, we extend the definition of a prolongation to act on vector fields as
$$
\left. \Pr^ {(n)} \boldsymbol {v} \right| _ {\left(\boldsymbol {x}, \boldsymbol {u} ^ {(n)}\right)} = \left. \frac {d}{d \epsilon} \right| _ {\epsilon = 0} \Pr^ {(n)} [ \exp (\epsilon \boldsymbol {v}) ] (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}). \tag {23}
$$
A given vector field  $\pmb{v}$  is therefore an infinitesimal generator of a symmetry group  $G$  of a system of differential equations  $\Delta_{\nu}$  indexed by  $\nu \in \{1,\dots ,l\}$  if the prolonged vector field of any given solution is still a solution:
$$
\Pr^ {(n)} \boldsymbol {v} \left[ \Delta_ {\nu} \left(\boldsymbol {x}, \boldsymbol {u} ^ {(n)}\right) \right] = 0, \quad \nu = 1, \dots , l, \quad \text {w h e n e v e r} \Delta (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}) = 0. \tag {24}
$$
For sake of convenience and brevity, we leave out many of the formal definitions behind these concepts and refer the reader to [11] for complete details.
A.2 Deriving Generators of the Symmetry Group of a PDE
Since symmetries of differential equations correspond to smooth maps, it is typically easier to derive the particular symmetries of a differential equation via their infinitesimal generators. To derive such generators, we first show how to perform the prolongation of a vector field. As before, assume we have  $p$  independent variables  $x^1, \ldots, x^p$  and  $l$  dependent variables  $u^1, \ldots, u^l$ , which are a function of the dependent variables. Note that we use superscripts to denote a particular variable. Derivatives with respect to a given variable are denoted via subscripts corresponding to the indices. For example, the variable  $u_{112}^1$  denotes the third order derivative of  $u^1$  taken twice with respect to the variable  $x^1$  and once with respect to  $x^2$ . As stated earlier, the prolongation of a vector field is defined as the operation
$$
\left. \Pr^ {(n)} \boldsymbol {v} \right| _ {\left(\boldsymbol {x}, \boldsymbol {u} ^ {(n)}\right)} = \frac {d}{d \epsilon} \Bigg | _ {\epsilon = 0} \Pr^ {(n)} [ \exp (\epsilon \boldsymbol {v}) ] (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}). \tag {25}
$$
To calculate the above, we can evaluate the formula on a vector field written in a generalized form. I.e., any vector field corresponding to the infinitesimal generator of a symmetry takes the general form
$$
\boldsymbol {v} = \sum_ {i = 1} ^ {p} \xi^ {i} (\boldsymbol {x}, \boldsymbol {u}) \frac {\partial}{\partial x ^ {i}} + \sum_ {\alpha = 1} ^ {q} \phi_ {\alpha} (\boldsymbol {x}, \boldsymbol {u}) \frac {\partial}{\partial u ^ {\alpha}}. \tag {26}
$$
Throughout, we will use Greek letter indices for dependent variables and standard letter indices for independent variables. Then, we have that
$$
\Pr^ {(n)} \boldsymbol {v} = \boldsymbol {v} + \sum_ {\alpha = 1} ^ {q} \sum_ {\boldsymbol {J}} \phi_ {\alpha} ^ {\boldsymbol {J}} (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}) \frac {\partial}{\partial u _ {\boldsymbol {J}} ^ {\alpha}}, \tag {27}
$$
where  $\pmb{J}$  is a tuple of dependent variables indicating which variables are in the derivative of  $\frac{\partial}{\partial u_{J}^{\alpha}}$ . Each  $\phi_{\alpha}^{J}(\pmb{x},\pmb{u}^{(n)})$  is calculated as
$$
\phi_ {\alpha} ^ {\boldsymbol {J}} (\boldsymbol {x}, \boldsymbol {u} ^ {(n)}) = \prod_ {i \in \boldsymbol {J}} D _ {i} \left(\phi_ {a} - \sum_ {i = 1} ^ {p} \xi^ {i} u _ {i} ^ {\alpha}\right) + \sum_ {i = 1} ^ {p} \xi^ {i} u _ {\boldsymbol {J}, i} ^ {\alpha}, \tag {28}
$$
where  $u_{\mathbf{J},i}^{\alpha} = \partial u_{\mathbf{J}}^{\alpha} / \partial x^{i}$  and  $D_{i}$  is the total derivative operator with respect to variable  $i$  defined as
$$
\boldsymbol {D} _ {i} P (x, u ^ {(n)}) = \frac {\partial P}{\partial x ^ {i}} + \sum_ {i = 1} ^ {q} \sum_ {\boldsymbol {J}} u _ {\boldsymbol {J}, i} ^ {\alpha} \frac {\partial P}{\partial u _ {\boldsymbol {J}} ^ {\alpha}}. \tag {29}
$$
After evaluating the coefficients,  $\phi_{\alpha}^{J}(x,u^{(n)})$  , we can substitute these values into the definition of the vector field's prolongation in Equation (27). This fully describes the infinitesimal generator of the given PDE, which can be used to evaluate the necessary symmetries of the system of differential equations. An example for Burgers' equation, a canonical PDE, is presented in the following.
A.3 Example: Burgers' Equation
Burgers' equation is a PDE used to describe convection-diffusion phenomena commonly observed in fluid mechanics, traffic flow, and acoustics [43]. The PDE can be written in either its "potential" form or its "viscous" form. The potential form is
$$
u _ {t} = u _ {x x} + u _ {x} ^ {2}. \tag {30}
$$
Cautionary note: We derive here the symmetries of Burgers' equation in its potential form since this form is more convenient and simpler to study for the sake of an example. The equation we consider in our experiments is the more commonly studied Burgers' equation in its standard form which does not have the same Lie symmetry group (see Table 4). Similar derivations for Burgers' equation in its standard form can be found in example 6.1 of [44].
Following the notation from the previous section,  $p = 2$  and  $q = 1$ . Consequently, the symmetry group of Burgers' equation will be generated by vector fields of the following form
$$
\boldsymbol {v} = \xi (x, t, u) \frac {\partial}{\partial x} + \tau (x, t, u) \frac {\partial}{\partial t} + \phi (x, t, u) \frac {\partial}{\partial u}, \tag {31}
$$
where we wish to determine all possible coefficient functions,  $\xi (t,x,u)$ ,  $\tau (x,t,u)$ , and  $\phi (x,t,u)$  such that the resulting one-parameter sub-group  $\exp (\varepsilon \mathbf{v})$  is a symmetry group of Burgers' equation.
To evaluate these coefficients, we need to prolong the vector field up to  $2^{\mathrm{nd}}$  order, given that the highest-degree derivative present in the governing PDE is of order 2. The  $2^{\mathrm{nd}}$  prolongation of the vector field can be expressed as
$$
\operatorname {p r} ^ {(2)} \boldsymbol {v} = \boldsymbol {v} + \phi^ {x} \frac {\partial}{\partial u _ {x}} + \phi^ {t} \frac {\partial}{\partial u _ {t}} + \phi^ {x x} \frac {\partial}{\partial u _ {x x}} + \phi^ {x t} \frac {\partial}{\partial u _ {x t}} + \phi^ {t t} \frac {\partial}{\partial u _ {t t}}. \tag {32}
$$
Applying this prolonged vector field to the differential equation in Equation (30), we get the infinitesimal symmetry criteria that
$$
\Pr^ {(2)} \boldsymbol {v} [ \Delta (x, t, \boldsymbol {u} ^ {(2)}) ] = \phi^ {t} - \phi^ {x x} + 2 u _ {x} \phi^ {x} = 0. \tag {33}
$$
To evaluate the individual coefficients, we apply Equation (28). Next, we substitute every instance of  $u_{t}$  with  $u_{x}^{2} + u_{xx}$ , and equate the coefficients of each monomial in the first and second-order
Table 3: Monomial coefficients in vector field prolongation for Burgers' equation.
Monomial	Coefficient
1	φt = φxx
ux	2φx + 2(φxu -ξxx) = -ξt
ux2	2(φu -ξx) - τxx + (φuu - 2ξxu) = φu - τt
ux3	-2τx - 2ξu - 2τxu -ξuu = -ξu
ux4	-2τu - τuu = -τu
uxx	-τxx + (φu - 2ξx) = φu - τt
uxuxx	-2τx - 2τxu - 3ξu = -ξu
ux2uxx	-2τu - τuu - τu = -2τu
ux2	-τu = -τu
uxt	-2τx = 0
uxxt	-2τu = 0
derivatives of  $u$  to find the pertinent symmetry groups. Table 3 below lists the relevant monomials as well as their respective coefficients.
Using these relations, we can solve for the coefficient functions. For the case of Burgers' equation, the most general infinitesimal symmetries have coefficient functions of the following form:
$$
\xi (t, x) = k _ {1} + k _ {4} x + 2 k _ {5} t + 4 k _ {6} x t \tag {34}
$$
$$
\tau (t) = k _ {2} + 2 k _ {4} t + 4 k _ {6} t ^ {2} \tag {35}
$$
$$
\phi (t, x, u) = \left(k _ {3} - k _ {5} x - 2 k _ {6} t - k _ {6} x ^ {2}\right) u + \gamma (x, t) \tag {36}
$$
where  $k_{1},\ldots ,k_{6}\in \mathbb{R}$  and  $\gamma (x,t)$  is an arbitrary solution to Burgers' equation. These coefficient functions can be used to generate the infinitesimal symmetries. These symmetries are spanned by the six vector fields below:
$$
\boldsymbol {v} _ {1} = \partial_ {x} \tag {37}
$$
$$
\boldsymbol {v} _ {2} = \partial_ {t} \tag {38}
$$
$$
\boldsymbol {v} _ {3} = \partial_ {u} \tag {39}
$$
$$
\boldsymbol {v} _ {4} = x \partial_ {x} + 2 t \partial_ {t} \tag {40}
$$
$$
\boldsymbol {v} _ {5} = 2 t \partial_ {x} - x \partial_ {u} \tag {41}
$$
$$
\boldsymbol {v} _ {6} = 4 x t \partial_ {x} + 4 t ^ {2} \partial_ {t} - (x ^ {2} + 2 t) \partial_ {u} \tag {42}
$$
as well as the infinite-dimensional subalgebra:  $\pmb{v}_{\gamma} = \gamma (x,t)e^{-u}\partial_{u}$ . Here,  $\gamma (x,t)$  is any arbitrary solution to the heat equation. The relationship between the Heat equation and Burgers' equation can be seen, whereby if  $u$  is replaced by  $w = e^{u}$ , the Cole-Hopf transformation is recovered.
B Exponential map and its approximations
As observed in the previous section, symmetry groups are generally derived in the Lie algebra of the group. The exponential map can then be applied, taking elements of this Lie algebra to the corresponding group operations. Working within the Lie algebra of a group provides several benefits. First, a Lie algebra is a vector space, so elements of the Lie algebra can be added and subtracted to yield new elements of the Lie algebra (and the group, via the exponential map). Second, when generators of the Lie algebra are closed under the Lie bracket of the Lie algebra (i.e., the generators form a basis for the structure constants of the Lie algebra), any arbitrary Lie point symmetry can be obtained via an element of the Lie algebra (i.e. the exponential map is surjective onto the connected component of the identity) [11]. In contrast, composing group operations in an arbitrary, fixed sequence is not guaranteed to be able to generate any element of the group. Lastly, although not extensively detailed here, the "strength," or magnitude, of Lie algebra elements can be measured using an appropriately selected norm. For instance, the operator norm of a matrix could be used for matrix Lie algebras.
In certain cases, especially when the element  $\pmb{v}$  in the Lie algebra consists of a single basis element, the exponential map  $\exp (\pmb {v})$  applied to that element of the Lie algebra can be calculated explicitly. Here, applying the group operation to a tuple of independent and dependent variables results in the so-called Lie point transformation, since it is applied at a given point  $\exp (\epsilon \pmb {v})\cdot (x,f(x))\mapsto (x',f(x)')$ . Consider the concrete example below from Burger's equation.
Example B.1 (Exponential map on symmetry generator of Burger's equation). The Burger's equation contains the Lie point symmetry  $\pmb{v}_{\gamma} = \gamma (x,t)e^{-u}\partial_{u}$  with corresponding group transformation  $\exp (\epsilon \pmb{v}_{\gamma})\cdot (x,t,u) = (x,t,\log (e^{u} + \epsilon \gamma))$
Proof. This transformation only changes the  $u$  component. Here, we have
$$
\begin{array}{l} \exp \left(\epsilon \gamma e ^ {- u} \partial_ {u}\right) u = u + \sum_ {k = 1} ^ {n} \left(\epsilon \gamma e ^ {- u} \partial_ {u}\right) ^ {k} \cdot u \tag {43} \\ = u + \epsilon \gamma e ^ {- u} - \frac {1}{2} \epsilon^ {2} \gamma^ {2} e ^ {- 2 u} + \frac {1}{3} \epsilon^ {3} \gamma^ {3} e ^ {- 3 u} + \dots \\ \end{array}
$$
Applying the series expansion  $\log (1 + x) = x - \frac{x^2}{2} +\frac{x^3}{3} -\dots$  , we get
$$
\begin{array}{l} \exp \left(\epsilon \gamma e ^ {- u} \partial_ {u}\right) u = u + \log \left(1 + \epsilon \gamma e ^ {- u}\right) \\ = \log (e ^ {u}) + \log (1 + \epsilon \gamma e ^ {- u}) \tag {44} \\ = \log \left(e ^ {u} + \epsilon \gamma\right). \\ \end{array}
$$
In general, the output of the exponential map cannot be easily calculated as we did above, especially if the vector field  $\pmb{v}$  is a weighted sum of various generators. In these cases, we can still apply the exponential map to a desired accuracy using efficient approximation methods, which we discuss next.
B.1 Approximations to the exponential map
For arbitrary Lie groups, computing the exact exponential map is often not feasible due to the complex nature of the group and its associated Lie algebra. Hence, it is necessary to approximate the exponential map to obtain useful results. Two common methods for approximating the exponential map are the truncation of Taylor series and Lie-Trotter approximations.
Taylor series approximation Given a vector field  $\pmb{v}$  in the Lie algebra of the group, the exponential map can be approximated by truncating the Taylor series expansion of  $\exp (\pmb {v})$ . The Taylor series expansion of the exponential map is given by:
$$
\exp (\boldsymbol {v}) = \operatorname {I d} + \boldsymbol {v} + \frac {1}{2} \boldsymbol {v} \cdot \boldsymbol {v} + \dots = \sum_ {n = 0} ^ {\infty} \frac {\boldsymbol {v} ^ {n}}{n !}. \tag {45}
$$
To approximate the exponential map, we retain a finite number of terms in the series:
$$
\exp (\boldsymbol {v}) = \sum_ {n = 0} ^ {k} \frac {\boldsymbol {v} ^ {n}}{n !} + o \left(\| \boldsymbol {v} \| ^ {k}\right), \tag {46}
$$
where  $k$  is the order of the truncation. The accuracy of the approximation depends on the number of terms retained in the truncated series and the operator norm  $\| \pmb{v} \|$ . For matrix Lie groups, where  $\pmb{v}$  is also a matrix, this operator norm is equivalent to the largest magnitude of the eigenvalues of the matrix [45]. The error associated with truncating the Taylor series after  $k$  terms thus decays exponentially with the order of the approximation.
Two drawbacks exist when using the Taylor approximation. First, for a given vector field  $v$ , applying  $v \cdot f$  to a given function  $f$  requires algebraic computation of derivatives. Alternatively, derivatives can also be approximated through finite difference schemes, but this would add an additional source of error. Second, when using the Taylor series to apply a symmetry transformation of a PDE to a starting solution of that PDE, the Taylor series truncation will result in a new function, which is not necessarily a solution of the PDE anymore (although it can be made arbitrarily close to a solution by increasing the truncation order). Lie-Trotter approximations, which we study next, approximate the exponential map by a composition of symmetry operations, thus avoiding these two drawbacks.
Lie-Trotter series approximations The Lie-Trotter approximation is an alternative method for approximating the exponential map, particularly useful when one has access to group elements directly, i.e. the closed-form output of the exponential map on each Lie algebra generator), but they are non-commutative. To provide motivation for this method, consider two elements  $X$  and  $\mathbf{Y}$  in the Lie algebra. The Lie-Trotter formula (or Lie product formula) approximates the exponential of their sum [22, 46].
$$
\exp (\boldsymbol {X} + \boldsymbol {Y}) = \lim  _ {n \rightarrow \infty} \left[ \exp \left(\frac {\boldsymbol {X}}{n}\right) \exp \left(\frac {\boldsymbol {Y}}{n}\right)\right] ^ {n} \approx \left[ \exp \left(\frac {\boldsymbol {X}}{k}\right) \exp \left(\frac {\boldsymbol {Y}}{k}\right)\right] ^ {k}, \tag {47}
$$
where  $k$  is a positive integer controlling the level of approximation.
The first-order approximation above can be extended to higher orders, referred to as the Lie-Trotter-Suzuki approximations. Though various different such approximations exist, we particularly use the following recursive approximation scheme [47, 23] for a given Lie algebra component  $\pmb{v} = \sum_{i=1}^{p} \pmb{v}_i$ .
$$
\mathcal {T} _ {2} (\boldsymbol {v}) = \exp \left(\frac {\boldsymbol {v} _ {1}}{2}\right) \cdot \exp \left(\frac {\boldsymbol {v} _ {2}}{2}\right) \dots \exp \left(\frac {\boldsymbol {v} _ {p}}{2}\right) \exp \left(\frac {\boldsymbol {v} _ {p}}{2}\right) \cdot \exp \left(\frac {\boldsymbol {v} _ {p - 1}}{2}\right) \dots \exp \left(\frac {\boldsymbol {v} _ {1}}{2}\right),
$$
$$
\mathcal {T} _ {2 k} (\boldsymbol {v}) = \mathcal {T} _ {2 k - 2} (u _ {k} \boldsymbol {v}) ^ {2} \cdot \mathcal {T} _ {2 k - 2} ((1 - 4 u _ {k}) \boldsymbol {v}) \cdot \mathcal {T} _ {2 k - 2} (u _ {k} \boldsymbol {v}) ^ {2}, \tag {48}
$$
$$
u _ {k} = \frac {1}{4 - 4 ^ {1 / (2 k - 1)}}.
$$
To apply the above formula, we tune the order parameter  $p$  and split the time evolution into  $r$  segments to apply the approximation  $\exp(\pmb{v}) \approx \prod_{i=1}^{r} \mathcal{T}_p(\pmb{v}/r)$ . For the  $p$ -th order, the number of stages in the Suzuki formula above is equal to  $2 \cdot 5^{p/2-1}$ , so the total number of stages applied is equal to  $2r \cdot 5^{p/2-1}$ .
These methods are especially useful in the context of PDEs, as they allow for the approximation of the exponential map while preserving the structure of the Lie algebra and group. Similar techniques are used in the design of splitting methods for numerically solving PDEs [48, 49]. Crucially, these approximations will always provide valid solutions to the PDEs, since each individual group operation in the composition above is itself a symmetry of the PDE. This is in contrast with approximations via Taylor series truncation, which only provide approximate solutions.
As with the Taylor series approximation, the  $p$ -th order approximation above is accurate to  $o(\| \boldsymbol{v} \|^p)$  with suitably selected values of  $r$  and  $p$  [23]. As a cautionary note, the approximations here may fail to converge when applied to unbounded operators [50, 51]. In practice, we tested a range of bounds to the augmentations and tuned augmentations accordingly (see Appendix E).
C VICReg Loss
In our implementations, we use the VICReg loss as our choice of SSL loss [9]. This loss contains three different terms: a variance term that ensures representations do not collapse to a single point, a covariance term that ensures different dimensions of the representation encode different data, and an invariance term to enforce similarity of the representations for pairs of inputs related by an augmentation. We go through each term in more detail below. Given a distribution  $\mathcal{T}$  from which to draw augmentations and a set of inputs  $\boldsymbol{x}_i$ , the precise algorithm to calculate the VICReg loss for a batch of data is also given in Algorithm 1.