CTSketch: Compositional Tensor Sketching for Scalable Neurosymbolic Learning

Originality: [...] Having said that, the techniques only apply to the symbol-inference part, and are disconnected from the neural part of the pipeline.

The reason why we use the program summaries and tensor approximations is to connect the neural component with the symbolic program in a way that is end-to-end differentiable. In that sense, CTSketch is connected to the neural part of the pipeline.

Quality: [...] Something that could be of interest and that is currently missing is the impact of the program decomposition on the final results.

We performed an ablation study on comparing the performance gains from program decomposition and tensor sketching with sum-16. We vary the level of decomposition and compare accuracy, per-digit accuracy, and per-epoch training time using full-rank vs rank-2 sketches. Here 16 refers to no decomposition, whereas (8, 2) refers to 2-layer decomposition into two sum-8 at the first layer and a sum-2 at the second layer.

By comparing the full-rank (FULL) and rank-2 sketches configurations under the same decomposition structure, we observe that sketching offers substantial space savings, albeit with some loss in accuracy. As we scale to higher-dimensional input spaces, memory becomes a critical bottleneck—program decomposition alone is often insufficient. This highlights the necessity of tensor sketching for scaling.

Within full-rank, the gains from greater degree of decomposition are clear, both in terms of training time and memory efficiency. However, this trend reverses with sketches, due to the use of the RBF kernel at the output of each program layer. For reference, we report results on using full-rank tensors with RBF kernels (FULL+RBF) which also show that accuracy drops with more program layers

With regards to related work / future work, the authors may find interesting work on tensor decompositions for probabilistic inference in general.

Thank you for all these references. We will look into them for automating program decomposition.

"Inference in CTSketch does not aggregate proofs", is an odd statement in a sense that the summations performed within CTSketch are aggregating in a similar fashion. I.e., you compute Eq 2.2, which is aggregating 'proofs'.

We agree that this statement is inaccurate. What we meant to say is that CTSketch does not aggregate proofs following the execution of a logic program, as is done in Scallop and DeepProbLog. CTSketch instead aggregates proofs through its tensor operations. We will clarify this point in the revised version of the paper.

The paper claims that they are the first to use program decompositions to scale neurosymbolic learning. However, this is not the case.

We will clarify in the revision that we are not the first to introduce general program decomposition. We do believe that CTSketch is the first neurosymbolic framework to use a decomposed program, represented as tensors, to improve the efficiency of inference. While a Scallop sum64 program would probably not benefit from decomposing it as 6 layers of sum programs (a timeout will still occur), CTSketch can drastically improve its efficiency when programs are decomposed.

De Smet and Zuidberg Dos Martires [1] already used program decomposition to scale to similar problem sizes. Albeit for a restricted set of neurosymbolic problems. This reference is missing in the paper.

We will refer to PLIA in the Related Works section, explaining PLIA adapts FFT for efficient addition of probability distributions over integers, showing high speed-ups for programs expressible with integer arithmetics.

The claim that logic programming languages do not support external API calls seems a bit strange. Logic programming is Turing-complete meaning they should probably be able to express an API call.

We made this claim in the context of differentiable logic programs not being compatible with calls to APIs. For example, one could not write a DeepProbLog program that calls GPT-4 as part of a neurosymbolic learning framework because it would not be end-to-end differentiable. However, some logic programming languages do offer foreign function interfaces like SWI-Prolog, which allows HTTP requests.

Tensors (used in this paper) and arithmetic/probabilistic circuits (used in many other works) exhibit many similarities. It would be nice to see a discussion on these similarities in the context of the paper.

Indeed, there are several similarities between the tensors used in CTSketch and probabilistic circuits (PCs). PCs decompose their probability distribution computation with sum and product nodes. Product nodes in PCs compute the product of their inputs, where each input corresponds to a probability over a disjoint set of random variables. This is similar to how CTSketch takes the product of probabilities corresponding to every possible input combination. Sum nodes in PCs compute a probability distribution over the same set of random variables. This is similar to how for each possible output, CTSketch adds the probabilities of the input combinations that resulted in this output. Unlike PCs, CTSketch has the ability to compute an approximate output distribution extremely efficiently with tensor sketching. However, their similarities are important to highlight in the related work section, and we plan to do so in the next revision.

How was the full rank tensor implemented? Was it a dense multilinear function? The reason I am asking is because recent advances have shown that one can efficiently implement arithmetic circuits (ie sparse multilinear functions, ie sparse tensors) efficiently on the GPU and scale neurosymbolic learning without any approximations [2].

Yes, the full rank tensor is implemented as a dense multilinear function. For a program $P: R_1 \times \dots \times R_n \rightarrow Y$, the full-rank tensor would be $\phi: R_1 \times \dots \times R_n$. This tensor can be represented as a arithmetic circuit consisting of three layers with $\sum_i |R_i|$ leaves, $\prod_i |R_i|$ nodes in layer 2, and $|Y|$ root nodes, similar to Figure 9 in [2]. KLay allows a time efficient neurosymbolic learning with this circuit, whereas CTSketch sketches $\phi$ for scalable inference..

To clearly state how inference is done with the sketches, we further explain the computation of the equation between lines 135 and 136.

$

v &= \sum_a^{|R_1|} \sum_b^{|R_2|} \sum_x^2 {p_1}[a] {p_2}[b] {t_1}[a, x] {t_2}[x, b] \\ v &= \sum_x^2 \left( \sum_a^{|R_1|} {p_1}[a] {t_1}[a, x] \right) \left( \sum_b^{|R_2|} {p_2}[b] {t_2}[x, b] \right) \\ v &= ( p_1^\top t_1 ) \cdot ( t_2 p _2 )

$

We have vectors $p_1:|R_1|$ and $p_2:|R_2|$ that are probability distributions, and 2D matrices $t_1:|R_1|\times 2$ and $t_2:|R_2|\times 2$ (where 2 is the sketching rank). Instead of taking the outer product of $t_1 \times t_2$ and reconstructing the tensor, we can rearrange the order and do vector-matrix multiplication ($p_1$ and $t_1$, $p_2$ and $t_2$) and take the inner product of the two.

This might imply that scaling is mainly due to program decomposition and not low-rank tensor approxiamtions?

We performed an ablation study on comparing the performance gains from program decomposition and tensor sketching with sum-16. We vary the level of decomposition and compare accuracy, per-digit accuracy, and per-epoch training time using full-rank vs rank-2 sketches. Here 16 refers to no decomposition, whereas (8, 2) refers to 2-layer decomposition into two sum-8 at the first layer and a sum-2 at the second layer.

By comparing the full-rank (FULL) and rank-2 sketches configurations under the same decomposition structure, we observe that sketching offers substantial space savings, albeit with some loss in accuracy. As we scale to higher-dimensional input spaces, memory becomes a critical bottleneck—program decomposition alone is often insufficient. This highlights the necessity of tensor sketching for scaling.

Within full-rank, the gains from greater degree of decomposition are clear, both in terms of training time and memory efficiency. However, this trend reverses with sketches, due to the use of the RBF kernel at the output of each program layer. For reference, we report results on using full-rank tensors with RBF kernels (FULL+RBF) which also show that accuracy drops with more program layers.

Is it not a bit of cheating to give direct supervision to the low-rank tensor decompositions of the sub-programs compared to what other methods are getting as training signal?

We do not provide any direct supervision using intermediate labels. We make the same assumptions as other neurosymbolic frameworks, which assume the domain of symbols (inputs to the program) is known. We use this assumption to randomly sample the symbols, prior to training and irrelevant to the actual training set labels, to initialize the program tensor.

You mentioned that program synthesis could potentially come to the rescue in that case. Could you argue more about it?

As a starting point for synthesizing decomposed programs, we asked Gemini 2.5 Flash how it would decompose Sum16: “How would you decompose the program that adds 16 digits, into a tree-based structure, using sub-programs that take as few inputs as possible…” Its response was the same as what we came up with:

Depth: 4 layers (excluding the initial input layer).

Components:

• Layer 0 (Inputs): 16 single-digit inputs.
• Layer 1: 8 ADD sub-programs. Each takes 2 single-digit inputs.
• Layer 2: 4 ADD sub-programs. Each takes 2 multi-digit inputs from Layer 1.
• Layer 3: 2 ADD sub-programs. Each takes 2 multi-digit inputs from Layer 2.
• Layer 4 (Root): 1 ADD sub-program. Takes 2 multi-digit inputs from Layer 3.

After two steps of prompting, we were also able to get a pairwise comparison decomposed program for Visudo. We think that LLMs could be a first step towards automating the decomposition process.

Firstly, the manual decomposition of tasks necessary for scaling (also mentioned by the authors themselves, which I appreciate) seems indeed a significant limitation for the applicability of the method outside of the simple domains investigated in this work.

We are also currently investigating CTSketch’s applicability to real-world domains such as predicting long-horizon sepsis from electronic health record (EHR) data. We are working on training a transformer to predict future EHR data from previous entries, and using a SOFA score program (which measures the severity of organ failure in a patient) with CTSketch to add an additional learning signal for predicting sepsis.

You mention in the paper that the construction of this vector could also be automated (using a subset of the training), but this would require labels for the intermediate representations, which would be against the initial proposition of working in an end-to-end setting.

We did not mean to imply that the construction of the task tensor would use a subset of the training data. What we meant was that the task tensor can use a subset of all possible structured input combinations, which doesn’t require any labeled data. For example, to initialize a sum2 tensor, we would need to know that the possible values (symbols) for both inputs are integers between 0 and 9, and we could sample some subset of the possible symbol combinations (and their corresponding outputs) to populate the summary tensor.

On another note, for my understanding, using low-rank approximation allows the avoidance of the instantiation of the full (sub-) task tensor in memory. However, this still needs to be computed fully (see Equation between Lines 135-136). So, if from a computation perspective the LoRa approximation only worsens the efficiency, where do the improvements shown in Section 4.5 come?

Due to properties of tensor products, we can compute the value of the Equation between lines 135-136 without reconstructing the tensor.

$

v &= \sum_a^{|R_1|} \sum_b^{|R_2|} \sum_x^2 {p_1}[a] {p_2}[b] {t_1}[a, x] {t_2}[x, b] \\ v &= \sum_x^2 \left( \sum_a^{|R_1|} {p_1}[a] {t_1}[a, x] \right) \left( \sum_b^{|R_2|} {p_2}[b] {t_2}[x, b] \right) \\ v &= ( p_1^\top t_1 ) \cdot ( t_2 p _2 )

$

$p_1:|R_1|$ and $p_2:|R_2|$ are vectors representing probability distributions, and $t_1:|R_1|\times 2$ and $t_2:|R_2|\times 2$ are rank-2 sketches. By rearranging the equation, we see that, instead of taking the output product of $t_1 \times t_2$, we can obtain the result with two vector-matrix multiplications and taking the inner product of the resulting vectors.

I suggest that the authors elaborate more on this regard in the manuscript to increase its soundness, possibly including additional ablations regarding the effect of the single components (LoRa approximation, sub-task decomposition) on both accuracy and computational efficiency for at least some of the investigated tasks.

We performed an ablation study on comparing the performance gains from program decomposition and tensor sketching with sum-16. We vary the level of decomposition and compare accuracy, per-digit accuracy, and per-epoch training time using full-rank vs rank-2 sketches. Here 16 refers to no decomposition, whereas (8, 2) refers to 2-layer decomposition into two sum-8 at the first layer and a sum-2 at the second layer.

By comparing the full-rank (FULL) and rank-2 sketches configurations under the same decomposition structure, we observe that sketching offers substantial space savings, albeit with some loss in accuracy. As we scale to higher-dimensional input spaces, memory becomes a critical bottleneck—program decomposition alone is often insufficient. This highlights the necessity of tensor sketching for scaling.

Within full-rank, the gains from greater degree of decomposition are clear, both in terms of training time and memory efficiency. However, this trend reverses with sketches, due to the use of the RBF kernel at the output of each program layer. For reference, we report results on using full-rank tensors with RBF kernels (FULL+RBF) which also show that accuracy drops with more program layers.

More generally, I also question the generalizability of encoding the task semantics with a single input-output tensor, which is, by construction, only going to work with known input examples. For instance, a network trained on sum4 would not generalize to sum5 out of the box, requiring the manual extension of its entries to cover those examples which are not in sum4 \cap sum5

We want to make it clear which components do and do not generalize to different tasks. It is true that a program summary tensor for sum4, representing the program’s discrete inputs and outputs, would not generalize to sum5. However, the same could be said for existing neurosymbolic approaches, i.e., a DeepProbLog program that sums 2 digits cannot be used to train sum4. However, the neural network used to train sum2 in CTSketch (with sketching) does generalize to sum4 at test time. Over 10 random seeds, we obtained an average of 92.5% accuracy and 98.1% digit accuracy on sum4, using a neural network trained on sum2.

Regarding the writing, I think that the quality of the paper is acceptable, but some parts are not extremely clear (especially Sections 1 and 2) due to notation clutter or imprecise/verbose explanations.

Thank you for this feedback. To clarify, in Equation 2.1, we are considering a program with $n$ inputs, so we have probability distributions over each input $p_1$ through $p_n$. $r_i$ refers to the $i$th discrete input in the input combination $r_1, \dots, r_n$ that results in output $\hat{y}$. So $p_i[r_i]$ refers to the probability of this input value. Since this notation is confusing, we will revise it to $\text{WMC}(\hat{y} \mid p_1, \dots, p_n) = \sum_{\mathbf{r} \in \mathcal{R},~ c(\mathbf{r}) = \hat{y}} \prod_{i=1}^n p_i(r_i)$

where $r=(r_1, \dots,r_n​) \in R_1 \times \dots \times R_n$. We hope that this will remove some ambiguity around indexing. For Figure 2, we thought that populating these matrices with values would make the figure look cluttered, but we can see why only explaining how they are populated could be unclear. In the revised version, we will add these numbers to the figure either in the main body or in the appendix. We can fix Figure 1 to make it clear that the 4th digit is also an input to the neural network, and we will change the section title to “Bound on Approximation Error”

[..] tensor decomposition is only applicable when the full program summary tensor is computationally affordable.

In the paper, we construct the full program summary tensor before applying sketching, and mention streaming tensor sketching methods—where sketches are computed without explicitly forming the full tensor—as a direction for future work. However, in response to your comment, we conducted an additional experiment on the sum-16 task without program decomposition (results reported below). In this setting, the full summary tensor would have size $10^{16}$, making direct construction computationally infeasible. Instead, we employed a cross approximation technique to learn rank-2 sketches directly from a subset of input-output samples, without forming the full tensor. We used the implementation provided by the tntorch Python package. A more thorough evaluation of how cross approximation can be integrated into CTSketch is left for future work.

Program decomposition itself is well-known and commonly used in the neuro-symbolic field such as in Terpret.

We should clarify that CTSketch isn’t the first work to propose program decomposition. TerpreT does have a compositional rule framework which guides their program synthesis. But TerpreT is a language designed for inductive program synthesis, and its programs can be used in the neurosymbolic learning setting. However, we believe that CTSketch is the first neurosymbolic learning framework to use a decomposed program, represented as tensors, to improve the efficiency of inference. While a Scallop sum64 program would probably not benefit from decomposing it as 6 layers of sum programs (a timeout will still occur), CTSketch can drastically improve its efficiency when programs are decomposed. We will clarify in the revision that we are introducing program decomposition with tensor summaries, not just general program decomposition.

The authors did not state the ranks to approximate each program in experiments as well.

We report the ranks for each task in Table 2 (Appendix C) along with other hyperparameters.

It is hard to tell if the scalability comes from program decomposition or tensor operation approximation. [...] Are there more results regarding scalabilities e.g., for more programs and with more metrics?

We performed an ablation study on comparing the performance gains from program decomposition and tensor sketching with sum-16. We vary the level of decomposition and compare accuracy, per-digit accuracy, and per-epoch training time using full-rank vs rank-2 sketches. Here 16 refers to no decomposition, whereas (8, 2) refers to 2-layer decomposition into two sum-8 at the first layer and a sum-2 at the second layer.

By comparing the full-rank (FULL) and rank-2 sketches configurations under the same decomposition structure, we observe that sketching offers substantial space savings, albeit with some loss in accuracy. As we scale to higher-dimensional input spaces, memory becomes a critical bottleneck—program decomposition alone is often insufficient. This highlights the necessity of tensor sketching for scaling.

Within full-rank, the gains from greater degree of decomposition are clear, both in terms of training time and memory efficiency. However, this trend reverses with sketches, due to the use of the RBF kernel at the output of each program layer. For reference, we report results on using full-rank tensors with RBF kernels (FULL+RBF) which also show that accuracy drops with more program layers. Sketching ranks are reported in Table 2 (Appendix C) along with other hyperparameters.

Can tensor operation approximation speed up general programs such as solving sudoku?

Encoding programs as tensors provides a faster way to do inference and gradient propagation on general programs compared to other neurosymbolic frameworks. Tensor sketching provides additional benefit in terms of memory, in trade-off of performance. For high-dimensional inputs, smaller sized tensors allow for faster inference.

Can tensor operation approximation be extended to more realistic settings where the sub-program tensor summary is infeasible?

The additional results reported above include solving sum-16 without program decomposition, for which using a full program tensor is computationally infeasible. Furthermore, the sum-8 sub-program tensor when decomposing into sum-8 and sum-2 also does not fit in memory. In both cases, sketching those tensors provides a solution.