Learning Structure-Aware Representations of Dependent Types

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.

* 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.

**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.

> `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.

**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.

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).

Hi again. We return with a **RoPE-based baseline**, as requested. | metric | stdlib[id] | stdlib[ood] | unimath | type-topo | | -------- | ----------- | -------------- | ---------- | ------------- | | AP | 22.56 | 12.83 | 18.00 |9.13 | | R@P | 11.74 | 8.95 | 8.33 |3.82 | | R@1 | 11.78 | 16.16 | 9.75 |3.22 | 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).

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. ---

Thank you for the nice questions :)