\subsection{Additional Analysis and Impact of \texttt{torch.compile}}

Here, we compare the theoretically estimated performance improvements based on counting the number of forward-propagated vectors with the empirically measured performance.

\paragraph{The number of propagated vectors is a good empirical performance estimate.}
To estimate the performance ratio between standard and collapsed Taylor mode, we can use the number of additional vectors both modes propagate forward as we increase either the batch size or the number of Monte-Carlo samples.
This is a relatively simplistic proxy; \eg, it assumes that each vector adds the same computational load, which is inaccurate as vectors corresponding to higher coefficients require more work and memory (as the Faà di Bruno formula contains more terms in general).
Conversely, while incrementing the MC samples does add additional vectors that are propagated, it does not introduce additional cost to compute or store the derivatives, as they are already computed with just a single sample.
\Cref{tab:benchmark-ratios} summarizes the theoretical and empirical ratios.
We find them to align quite well, despite the overly simplistic assumptions.

\paragraph{Concrete example.}
Consider the exact Laplacian.
Adding one datum introduces {\color{tab-green}$2 + D$} versus {\color{tab-orange}$1 + 2D$} new vectors.
For $D=50$, their ratio is $\nicefrac{\color{tab-green}(2 + D)}{\color{tab-orange}(1 + 2D)} \approx 0.51$.
Empirically, we measure that adding one datum adds {\color{tab-orange}\inputMetricOnly{jet/exp/exp01_benchmark_laplacian/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_laplacian_vary_batch_size/jet_naive_best.txt}\,ms} to standard, and {\color{tab-green}\inputMetricOnly{jet/exp/exp01_benchmark_laplacian/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_laplacian_vary_batch_size/jet_simplified_best.txt}\,ms} to collapsed, Taylor mode (\cref{tab:benchmark}); the ratio of $\approx \!\!\!\inputMetricRatio{jet/exp/exp01_benchmark_laplacian/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_laplacian_vary_batch_size/jet_simplified_best.txt}{jet/exp/exp01_benchmark_laplacian/performance/architecture_tanh_mlp_768_768_512_512_1_device_cuda_dim_50_name_laplacian_vary_batch_size/jet_naive_best.txt}$ is close.

\input{tables/torch_ratios}

\paragraph{Compilation reduces memory, but not runtime.}
In \cref{tab:benchmark-compiled}, we repeat the benchmark from \cref{tab:benchmark} using \texttt{torch.compile}.
We observe that compiling can further reduce the memory footprint of all approaches for computing the Laplacian and weighted Laplacian, while the runtime remains roughly identical.
For the biharmonic operator, we observe that compilation leaves runtime and memory footprint unchanged.

\input{tables/torch_benchmark_compiled}

\input{figures/torch_benchmark_compiled}

\subsection{Rank-deficient Weighted Laplacian}\label{sec:rank-deficient-weighted-laplacian}

In the main text we use a weighted Laplacian with a full-rank weight matrix (\ie, $R \coloneqq \rank(\mD) = D$).
Since the weight matrix has full rank, the weighted Laplacian is as expensive as the unweighted Laplacian, and this is confirmed by our experiments.
To show that the weight matrix's rank indeed affects the cost, we experiment with a rank-deficient weight matrix in this section and also consider the ranks $R \in \{ \nicefrac{D}{2}, \nicefrac{D}{10} \}$.
\Cref{tab:benchmark-rank} contains the results of this analysis.
We observe that going from full to half-full rank roughly halves both the runtime and memory consumption for all implementations.
For small ranks, this linear relationship weakens because the fraction of computations that do not scale with $R$ grows.

\input{tables/torch_rank_benchmark}

%%% Local Variables:
%%% mode: LaTeX
%%% TeX-master: "../main"
%%% End:
