Learning Structure-Aware Representations of Dependent Types

Hi, and thank you for taking the time to review our paper.

... visual results on representation?

We provide a visualization of the empirical distributions of positive and negative lemmas in Figure 3, which indicates that positive lemmas are consistently ranked significantly higher than negative ones in the development set. We further supply additional visualizations in the author rebuttal pdf. We apologize if this format is inconvenient; we have produced interactive plots that are much more informative, but we have no way to make them available (author guidelines restrict us from uploading url links, even if anonymized; we have asked the area chair for permission, but we have received no response yet).

... open-sourcing the code in the future?

We have made our code available for review as supplementary material. We are committed to keeping our code open source and freely available after the reviewing process is concluded.

... search-based ATP ... more general setting

Indeed these are future directions we also wish to pursue, but definitely lie outside the scope of this paper. Although it should have been clear from the introduction that premise selection is just the most natural first step, we will also extend the conclusion of Section 6 to explicitly mention these points for future work.

• It is not clear how to connect BERT with the current dataset/experimental setup, can you elaborate? We remind you that the experimental setup does not involve textual representations of any kind.
• We have already provided this ablation; please inspect Table 1 (row: - structured attention), and paragraph Ablations in Section 5.2. As you say, a large part of the model gains are indeed due to the structured attention from the tree-structured positional encodings.

In the future, is it possible to scale up models designed with your architecture?

We missed this one, sorry. The most honest answer is "it depends". There are projects ongoing, like Dedukti (see Conclusion), which attempt a common formalization interface/meta-language between different type-theory-based proof assistants. If/when such a project takes hold, and upon larger libraries becoming accessible, it would make sense to upscale the model. Until then, and for the time being, it is very likely that there's little benefits from model scaling given the current data availability. Otherwise, from an architectural perspective, there's nothing really stopping one from training larger (or smaller) variations of the proposed model.

Thank you for your questions, we really appreciate your engagement.

• multiple positive samples

We don't have to do anything special! Think of some make-do proof that looks like A (B C), where A, B and C are lemmas. There are several possible holes: ? (B C), A (? C), A (B ?) and A ?. The last one contains two positive lemmas: B and C, so we want both of them to be picked by the premise selector.

• ... negative samples selected ...

Consider the previous example, and a context providing availability to lemmas A, B, C, D and E. For this present hole, A, D and E are negative as they do not appear inside of it.

• ... vanilla Transformer ...

We were also wondering the same thing; we already used a vanilla Transformer to catastrophic results. See the last row of Table 1. (ps: see also line 354)

• Supplement the baseline: I see. Can you provide results from training a pre-trained model such as BERT in conjunction with Equation 3? Although the comparison may not be entirely fair, scaling the model from the this paper is also challenging, so the best approach is to combine it with an existing representation model. In the future, is it possible to scale up models designed with your architecture?
• Ablation study: I suspect the model's performance comes from the position encoding. Could you provide an ablation study using a vanilla Transformer combined with binary-tree positional encoding[1] and Equation 3?

Reference:
[1]: Kogkalidis, K., Bernardy, J. P., & Garg, V. (2023). Algebraic Positional Encodings. arXiv preprint arXiv:2312.16045.

That's a fair question -- the datasets used in the cited paper are all without exception text-based. We kindly ask you to inspect pages 17 - 20 of the linked paper, which present example datapoints. The data points presented are pretty-printed proofs (i.e., strings), making them suitable for autoregressive language modeling. The data points we work with are not "strings that look like proofs" -- they are the "structural footprints of proofs". We explicitly refer to this this paper (in line 106), and compare their dataset extraction and modeling methodologies to ours (line 117+, and Section 4.2).

Scaling the model: Given the slow growth of the Agda library and limitations of Dedukti , why not consider choosing Mathlib instead? It has a vast amount of data and mature methods for acquiring large datasets, such as [1], which make scaling easier.

Reference: [1] Han, J. M., Rute, J., Wu, Y., Ayers, E. W., & Polu, S. (2021). Proof artifact co-training for theorem proving with language models. arXiv preprint arXiv:2102.06203.

Supplement the baseline:: Considering the remarkable effect of this positional encoding[1], I would like to see whether it is possible to inject this structural bias by taking a pre-trained model like BERT and continuing the training using contrastive learning with this positional encoding. If possible, could you provide the experimental results of vanilla transformers combined with other positional encodings, such as Yarn[2] or RoPE[3]?

Reference:
[1]: Kogkalidis, K., Bernardy, J. P., & Garg, V. (2023). Algebraic Positional Encodings. arXiv preprint arXiv:2312.16045.
[2]: YaRN: Efficient Context Window Extension of Large Language Models.
[3]: Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., & Liu, Y. (2024). Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568, 127063.

Hi again. We return with a RoPE-based baseline, as requested.

As you suggested, performance is much better than the vanilla sinusoidal positional encoding scheme. At the same time, it is still much worse than the tree-structured positional encoding we use (or any ablation, for that matter). Neither of the two observations is too surprising:

• we know RoPE empirically outperforms sinusoidal encodings from the literature
• RoPE is still a sequential encoding scheme; types are not sequences, and sequential models can't do them justice (section 4.2)

We address the limitations of hierarchical tree methods conceptually in lines 205-211. Practically, the papers you suggest are not suitable baselines due to their reliance on quadratic $O(n^2)$ attention, which is infeasible for our setup as discussed in lines 221-228. With lemma lengths ranging from $10^0$ to $10^4$ tokens and files counting $10^0$ to $10^3$ lemmas, quadratic attention results in $O(mn^2)$ complexity (for a batch size of 1), making it untenable. This is exactly why we chose a linear attention kernel (lines 233 - 248).

While hierarchical tree and linear attention could be combined theoretically, no existing implementation that we know of exists. We sincerely apologize, because your suggested experiments are valid and interesting, but we cannot implement, train and evaluate such a mechanism within the 12 hours left before the discussion period closes. Please also keep in mind that training alone takes approximately 6-8 hours (assuming the compute resources are immediately available).

Hi, and thanks for the review.

We’re very glad you appreciated our work. We acknowledge that the presentation can sometimes be dense. We aimed to keep the tone and language inclusive and to provide brief descriptions of any jargon where possible. However, the paper's topic sits at a narrow intersection of type theory, automated theorem proving, and ML, which unfortunately requires some familiarity with all three.

Although we hoped the concrete examples and figures in the Appendix would cover such concerns, it turns out that wasn't the case. To this end, we commit to extending section A.2 with some more explanations revolving around the topics you mention in your comment.

Hello, and thank you for your critical review. We will try to address some of your concerns below.

...why the dataset is necessary or valuable ...

Thank you very much for this question. The dataset is radically different to existing datasets, and while do point the fact out (e.g., in lines 8, 58 - 59, 152 - 157, 161 - 167, 376, 381, 384 - 385 etc.), it is true that we should have been more explicit about its uniqueness. It is unfortunate that we couldn't convey the points through our paper; let us try to do so here instead.

The dataset is necessary because:

• It is the first dataset that we know of that explicates the shape of program-proofs at the subtype level in a way that is structurally preserving. Whereas most datasets contain textual representations of formal proofs, our dataset exposes also the real underlying structure of the proof objects under scrutiny. As we extensively argue in Section 4.2 (lines 185 - 205), it is exactly this structure that distinguishes a proof from a string that happens to look like a proof. Not having access to this structure and a dataset that explicates it means the community is confined to working with the wrong kind of models and evaluating on the text-only benchmarks.

The dataset is valuable because:

• it represents proof objects according to Agda's term/type syntax, which is mind boggingly close to the real syntax of the underlying type theory. The type theory is a universal theoretical foundation; Agda is an implementation. Having representations close to the theoretical foundation means that the dataset has more chances of finding use in other related endeavors than a dataset that capitalizes on the surface syntax of e.g. Lean proofs.
• Along with Coq, Agda is one of the two most important programming languages for foundational mathematics on the planet today. Homotopy type theory and the univalent foundations program, i.e. the contemporary approaches to foundational mathematics, rely predominantly on these two languages to formalize and verify their findings. Expediting the progress of these two languages and their interface to ML is expediting the progress of mathematics.
• There is epistemic value in scientific plurality. By providing a dataset that highlights the structural elements of proofs in Agda, we enable the exploration of new and diverse approaches to proof verification and automated reasoning.

... one just could adopt the non-Agda datasets ...

We would just like to point out that this is largely non-trivial. There is no methodology that we know of (or can imagine) that would allow the conversion of text-formatted proofs of Coq into structure-faithful representations of Agda programs. The contrary is somewhat more likely.

The experimental setting is unclear for me.

We are very sorry to hear. Would that perhaps warrant a reduction in your very high confidence score? In either case, we will try to answer your questions below.

Is it fine-tuned on the Agda dataset?

The transformer is trained on the Agda dataset; we start from a blank slate (no pretraining) (line 298).

... is trained against the data in JSON format or text format?

Neither. The model takes a collection of ASTs as input (Section 4.2, line 301, etc.). The input to the model is the literal structure of the underlying Agda program. Each AST is a collection of position-specified type operators (zeroary, unary, binary) or references (intra- or inter-AST).

If the JSON format is adopted, I suspect it might be unfavorable for Transformer because it needs to parse JSON data by itself.

Very good point! This is exactly our point as well, except we are arguing against all text-based representations (Section 4.2). What we feed the model is neither the JSON format, nor the pretty printed strings of proof objects: it is their actual structure.

I'm unsure Transformer is appropriate as a baseline because ...

Also a very good point, and one we ourselves make in Section 4.2 (lines 185 - 188). The vanilla Transformer, and all LLM-related endeavors suffer from exactly these biases, regardless of whether they are generative or encoder-only. This is why we think our work is an important alternative take: one that emphasizes structural discipline and tools fit for the task rather than tasks fit for the tool.

What is the outputs of the proposed model? An AST of a proof term? Or, just a text?

Neither; our model produces a scalar value for each hole-lemma pair (line 301). Without the premise selection head attached to its top, the model produces a neural representation for each AST (hole or lemma).

What is the design principles of the proposed model?

The design principles are extensively motivated in Section 4.2. Basically: sequential architectures are not fit for the task; people simply use them because this is what seems to work now. Tree architectures are nice, but they are expensive and hard to parallelize and perform badly. Graph architectures are not a good conceptual match for the task either; the input is not just any arbitrary graph, but a very refined and structurally consistent one. Therefore we need an architecture that does the problem's structure justice. Our architecture does exactly that, and, from what it looks like, it works.

Figure 3

This is a paired violin plot. It shows the empirical distribution of the scores of positive ground truth items (above) vs the scores of negative ground truth items (below). It shows that positive items have a consistently higher score than negative items. We explain this in lines 318 - 324.

What is "the added cognitive load" mentioned in Section 5.2 (page 9)?

This is exactly referring to your earlier point; the model now has to parse sequentialized type representations, because it misses the structural positional encodings that make this parse structure explicit.

We hope this is helpful.

Thanks for the response. It resolves some of my concerns, especially the ones for the contributions on the dataset, so now I'm leaning towards acceptance.

As a critical reviewer, I recommend the revision clarifies the contributions described in the response for the dataset, in the introduction. I couldn't find these contributions in the introduction of the submission, which prevents me from understanding the benefits of the new dataset compared with the existing ones (especially reading the related work section, which mentions the internal type structure as a difference, but what it means was unclear). The response says that, with the quote of lines 58 - 59, the contributions are mentioned in the introduction. Perhaps the crucial word in the lines is "program-proofs," but it is difficult, at least for me, to find that this phrase means that the dataset made by the paper is more structured (and in what sense) than the existing datasets.

Furthermore, even after admitting the differences from the existing dataset, it is still not clear to me what the "subtype level" means. Does it mean what is called "inter-AST" in Section 4.2? In any case, it would be helpful to clarify what it means and clearly connect the contributions described in the introduction with the details in the remaining sections.

I'm unsure Transformer is appropriate as a baseline because ...

Also a very good point, and one we ourselves make in Section 4.2 (lines 185 - 188). The vanilla Transformer, and all LLM-related endeavors suffer from exactly these biases, regardless of whether they are generative or encoder-only. This is why we think our work is an important alternative take: one that emphasizes structural discipline and tools fit for the task rather than tasks fit for the tool.

I don't think this response addresses my concern of why only the vanilla LLM is used as a baseline and the other LLM-related models (claimed to better perform than the vanilla LLM) is not used.

Thank you for your engagement, and for revising your rejection.

... "subtype level" ..

Yes, that's correct! This is indeed one of the benefits of the extra structure. For example, consider the dependent type $\Pi_{x:A}(B~x)$, or in Agda syntax, (x : A) -> B x, which denotes a family of types B indexed by objects of type A. The dataset represents this as an AST that looks like:

DependentProduct
|--- Argument 
|    |--- Type: A  (this would actually be an inter-AST pointer to the AST defining A)
|--- Body
     |--- FunctionApplication
          |--- FunctionBody
               |--- Type: B  (this would be an inter-AST pointer to the AST defining B)
          |--- FunctionArgument
               |--- @0  (this would be an intra-AST pointer to the Argument above)

This illustrates how structural relations are fully specified at the subtype level, all the way down to typing primitives. This is what allows the neural system to operate on a structural rather than textual level. Practically, our embedding layer only has the tiniest number of symbols (i.e., only primitive type operators). Neural processing happens over complex structures made of these symbols.

..LLM is used as a baseline and the other LLM-related models

We do not use a LLM. We use a (non-pretrained) tiny, customized Transformer made of 6 million parameters. This is 1000 times less parameters than the LLMs that you cite and are currently in use (lines 563 - 566). The trained model is 26MB big. To make a fair architectural comparison, we compare against a Transformer of a similar size, because this is the architecture these LLMs rely on. Our comparison demonstrates that our model dramatically outperforms similarly sized Transformers. Comparing with overparameterized, closed-source, pretrained LLMs is beyond the scope of our work and offers no epistemic insights that we can think of.

Hello, and many thanks for your honest review.

We're glad you appreciated the value of the architecture and the potential in the results.

We acknowledge the weaknesses you spot. We tried our best to accommodate a reader that is not necessarily familiar with Agda, but does have prior exposition to ATP concepts and type theory. Given the focused nature of our contribution, it proved really hard to find a more inclusive tone/vocabulary without making the presentation entirely superficial.

... a positive to be an hole that is filled by a lemma ... a hole/lemma pair could have been a positive, it just wasn´t considered?

This is precisely on point. A positive is the specific lemma that fills that hole. Another lemma could possibly have filled that hole, but in practice the vast majority of "other" lemmas in context would not have been appropriate.

... shouldn´t it be about a specific hole in a specific sequence of steps?

That's a good question, and we're very happy you asked. As we explain in lines 73 - 74, the cool thing about Agda is that you don't have to write your proofs/programs sequentially -- you can write the entire proof as normal, but defer some parts to the future: these are holes. That means that a hole is not like an element of a sequence with left-only context, but an element of a syntax tree with a much wider context, spanning (i) the rest of the tree, but also (ii) all other trees in scope. This is precisely the factor motivating the design of our dataset and the architecture, and what makes both stand out from the established literature:

1. On the dataset front, we consider all proofs of the Agda libraries we train on, and we turn all possible subproofs within each proof into a hole.
2. On the architecture front, we take care to properly encode the entire context in a structure-aware fashion.

... fixed number of holes ... more specific

Our training configuration considers a batch of 32 holes-per-file, 1-file-per-batch. This is not the result of hyper-parameter optimization; it's just a value that maximized resource utilization without causing cuda OOM errors.

... how many positive and negative examples do you actually have

Across the training dataset, we have 53,668 positives and 15,026,433 negatives, i.e., positive lemmas make for 0.35% of the total lemmas.

Any specific reason to ignore GNNs?

We do in fact have a (very) short discussion on GNNs in lines 211 - 217. The idea is that GNNs capture the structure implicitly and in small, localized neighborhoods by virtue of the graph convolution kernel. If the graph is large (which is the case when dealing with complex dependent types), the model doesn't have a reliable way to account for the whole type's structure, which leads to coarse and structurally unrefined representations.

Fig 1, please explain

We are defining the type of naturals, N, and equip it with a binary additive operation. There are two ways to make a natural; either the $zero$, for which $zero + n = n$, and, inductively, via the $suc$ (succedent) function, for which it holds that $(m + 1) + n = (m + n) + 1$. We then proceed to prove that given $m$ and $n$ naturals, $m + n$ is equivalent to $n + m$. The proof follows a case-by-case analysis:

1. If $m$ and $n$ are both $zero$, we can prove our goal using reflexivity (both sides of the equation are syntactically equal).
2. If $m$ is $zero$ and $n$ is actually the succedent of some other $n'$, then we can invoke the proof of the commutativity of $m + n'$.
3. Similar to (2): If $n$ is $zero$ and $m$ is the successor of some other $m'$, then we can invoke the proof of the commutativity of $m' + n$.
4. If $m$ is the successor of some other $m'$ and $n$ is actually the successor of some other $n'$, then we appeal to a helper lemma about how the successor function $suc$ distributes across addition, before finally invoking the proof of the commutativity of $m + n'$.

It would also be interesting if we could have a baseline. say. are there comparable resuls for, say, Coq?

This would be very difficult to accomplish, given the vast difference between the datasets and the inherent information available in the actual content within those datasets.

Having said that, our methodology is, at least conceptually, applicable to any language based on Type Theory, hence it is possible to port our constructions to Coq and then start considering comparisons against previous works.

1.1 universal: ie language that uses type theory as its foundational core. It would be nice to mention a few such languages?

Duly noted, we will refer to the most commonplace functional programming languages and theorem provers based on Type Theory, e.g. Haskell/OCaml/Coq/Agda/Lean.

... Taylor expansion ...

The effect of the Taylor expansion step is by no means small, offering an absolute increase of 3.2 points in precision and 4.1 in recall. The fact that this increase is overshadowed by the relative gains offered by structured attention and the structural variable representation serves only to empirically verify the necessity of the architectural desiderata we argued for in Section 4.2, and to show that their effect is significantly stronger than that of SOTA architectural adjustments and micro-optimizations.

... Missing venue ...

Good catch, thank you! We had to cook up the citation manually, seeing as Graph2Tac was released concurrently with our submission.

We hope this clarifies things a bit. Once again, thank you very much for your time and effort.