Learning Structure-Aware Representations of Dependent Types

Authors: konstantinos  kogkalidis
orestis  melkonian
jean-philippe  bernardy
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 18:47:17

Abstract: Agda is a dependently-typed programming language and a proof assistant, pivotal in proof formalization and programming language theory.
This paper extends the Agda ecosystem into machine learning territory, and, vice versa, makes Agda-related resources available to machine learning practitioners.
We introduce and release a novel dataset of Agda program-proofs that is elaborate and extensive enough to support various machine learning applications -- the first of its kind.
Leveraging the dataset's ultra-high resolution, which details proof states at the sub-type level, we propose a novel neural architecture targeted at faithfully representing dependently-typed programs on the basis of structural rather than nominal principles.
We instantiate and evaluate our architecture in a premise selection setup, where it achieves promising initial results, surpassing strong baselines.

Main paper: 
 introduction: 
 Automation on theorem proving over arbitrary domains has been a staple task for computer science and artificial intelligence, and a concept that in fact predates the inception of both fields. Over the last decade, the success of gradient-based optimization and large-scale neural models has brought this once elusive goal within reach. Capitalizing on that fact, community efforts have led to a wider availability of language-specific theorem-proving datasets. This, in turn, has led to a constant stream of milestones and achievements. Rather than threatening theorists worldwide with premature retirement, this development empowers mathematicians and computer scientists on the individual scale, and potentiates an unprecedented acceleration for the formal sciences on the collective level.Theorem-proving tasks can take on many forms, one of which is premise selection. In a premise selection setup, the model is presented with a goal statement and a collection of theorems or lemmas, some of which are potentially relevant for the proof process, whereas others are not. The model is then tasked with selecting which of those lemmas are relevant. A successful model can relieve burden from the human operator by dramatically reducing or even trivializing the proof search space. From an epistemic perspective, premise selection epitomizes the good aspects of neurosymbolic AI. Human and machine get to interact and collaborate towards a common goal. Meanwhile, the model's predictions are made safe by virtue of being confined within (and verified by) a trustworthy white box: the proof assistant. Depending on the logic employed by the prover, the line between proof search and program synthesis may become blurry. Concretely, the Curry-Howard correspondence asserts an equivalence between intuitionistic logics and constructive type theories; proofs in such logics are indistinguishable from programs in some corresponding functional programming language, and vice versa. This marvel has been the foundational core behind the development of type-theoretic proof assistants, such as Coq, Lean and Agda, inter alia. On the one hand, these are programming languages, except strict enough to find use in the formalization of science or the certification of engineering. On the other hand, these are also theorem provers, except with a less obscure syntax, that offer additional niceties such as compilation-level reporting, termination checking, and a natural ability to extract executable programs from the proofs produced.In the machine learning world, domain-, language-and task-specific representation learning for theorem proving has gained significant community attention in recent years. Modeling strategies can be broadly grouped in two general categories. The dominant approach relies on sequence-based encoder and/or decoder architectures, operating on the user-facing representation of formal objects (theorems, formulas, lemmas, etc.) ; this approach has become popular due to its direct compatibility with generic pretrained large language models. The alternative approach instead seeks to explicate the relations within and between these formal objects, e.g., with tree or graph architectures.In this paper, we identify and address two weak points in the existing literature.Machine Learning for Agda First, we note a disparity in the availability of machine learning resources and results across proof assistants. A plurality of datasets, models and tools are available for a handful of languages, the majority of them for Coq, HOL and Lean, whereas none exist for Agda. We believe that this state of affairs is an effect of historical momentum and methodological precedence, rather than inherent merit. Agda spearheads developments in constructive type theory with an array of innovative features, gaining adoption and a rapidly maturing ecosystem. Aiming to aid Agda with meeting and expanding her potential, we implement an Agda program that extracts intermediate compilation steps, and produces them in the form of human-readable partial proof states. Applying the extraction algorithm, AGDA2TRAIN, on established high-quality libraries, we obtain and release an elaborate and extensive dataset of Agda program-proofs, at a scale and form that can support various machine learning applications -the first of its kind.Modeling Type Structure Second, we note a significant gap between the rigor of the structures modeled and the structural lenience of the architectures employed. Seeking to bridge this gap, we design and implement a general-purpose scheme for the faithful representation of expressions involving dependent types. We apply our methodology on the extracted dataset to produce QUILL: a novel neural guidance tool for Agda. Beyond its current instantiation in Agda, the system is universal, being equally applicable to any language that uses type theory as its foundational core. 
proof assistants: 
 The leading proof assistants (Coq, Lean, Agda) all follow the same core principles with respect to the kinds of objects they define and manipulate. Modulo theoretical flavorings and syntactic sweeteners, these are always the terms of a typed λ-calculus. Leveraging the propositions-as-types interpretation [Sørensen and Urzyczyn, 2006, Wadler, 2015] , the proof assistant's duty is chiefly to check that a program adheres to a specific type, or, equivalently, that a proof attests to a particular proposition. Typically, proofs are constructed in an incremental fashion: the user can defer some sub-proof to the future by instead supplying a hole. We provide an illustrative example in Appendix A, Figure 4 . A proof is only complete if it is without holes. The proof assistant may assist the user by providing a set of hypotheses (or premises), which can be used to fill in a given hole. The problem of singling out relevant premises is completely specified by (i) the type of the hole itself (henceforth called the goal), and (ii) the types of premises available. Except for trivial problems, the set of premises that can be used in any given hole is enormous, and the number of ways in which they can be combined is even larger. Being able to filter and rank premises thus goes a long way towards efficient proof search. To aid in the process of proof construction, some assistants additionally implement tactics. Tactics can range from simple (introducing a single λ-abstraction) to complex (type-guided combinatorial search, domain specific solvers, etc.). Tactics also get to heavily benefit from a premise selection pass, given their combinatorial nature; however, they do not (yet) play a major role in the Agda world, and are not of primary interest to us here.Navigating the search space of possible proofs by purely symbolic means is generally infeasible. As an alternative, several lines of work have explored the application of statistical learning methods as a way to either truncate the search space, or to guide the theorem prover.An early step of theorem proving into machine learning territory is HolStep [Kaliszyk et al., 2016] .HolStep provides string representations of intermediate proof states extracted from 11,400 narrowdomain but non-trivial HOL proofs. The dataset is used to benchmark character-and token-level encoders in the task of telling whether a statement is relevant in the proof of a conjecture or not. Another contribution along the same lines is HOList [Bansal et al., 2019] . The dataset encompasses 30,000 HOL proofs, and the task is framed as selecting a tactic and its arguments, given a goal. The modeling approach rests on DeepHOL: a reinforcement learning loop incorporating convolutional encoders applied on string-formatted formulas.On the type-theoretic front, two related and concurrent contributions are GamePad [Huang et al., 2018] and CoqGym [Yang and Deng, 2019] . Both implement a minimal interactive interface with Coq, the latter also exposing a dataset of 71,000 human-written proofs. The main task is once more framed as tactic selection. GamePad uses a sequential recurrence to encode formula terms, emulating variable binding using a random embedding-association table. CoqGym, on the other hand, employs a tree-recursive encoder for representing goals and premises, and a semantically constrained tree-recursive decoder for decoding into the tactic language.Beginning with the early works of Urban and Jakubův [2020] on Mizar and Polu and Sutskever [2020] on Metamath, large language models (LLMs) have by now permeated the theorem proving literature. Interfaces between LLMs include LeanStep [Han et al., 2021] , lean-gym [Polu et al., 2023] and LeanDojo [Yang et al., 2023] for Lean, as well as Lisa [Jiang et al., 2021] and Magnushammer [Mikuła et al., 2023] for Isabelle. LeanStep and lean-gym export and utilize compiler-level type information from the intermediate proof steps of about 128,000 proofs; yet these are still formatted as strings and processed as such by the corresponding LLM. LeanDojo exports similarly-structured data from 98,734 theorems, additionally providing interactive capabilities, while Lisa exports shallow states from 183,000 theorems, represented as sequential MDPs. Magnushammer, finally, provides access to a shallow premise dataset from 570,000 mixed human and synthetic proofs states. The models underlying the above works all build on the same key idea: an LLM is exposed to a textual representation of an incomplete proof, and is then tasked with autoregressively continuing.Evidently, the bulk of the modeling work has so far been outsourced to sequential architectures. Despite their merits, we argue in Section 4.2 that such models are ill-equipped to capture the structures being modeled, except superficially and on a data-driven basis. The structural effects of (co-)reference, variable binding and abstraction, equivalence under substitution etc. are all dismissed in favor of a uniform and simplified representation format. Other than aforementioned exceptions, works that stand out in this respect are those of Wang et al. [2017] and Paliwal et al. [2020] , who opt for a messagepassing approach on formula trees (using graph-like edges to handle co-referencing and variable binding). Along the same lines, Li et al. [2020] use a 2-level hierarchical encoder to distinguish between intra-and inter-formula attention. Contemporary work [Blaauwbroek et al., 2024] further explores this perspective, fixing a symbol embedding table for a collection of symbols and employing message-passing networks over formula graphs to encode definitions in Coq. Representations are built in topological order, such that referring objects can utilize the previously computed representations of their referents. These representations finally find use in allowing the dynamic prediction of tactic arguments, given a chosen tactic.Our contribution follows along and expands upon the structurally-disciplined line of approaches. Our data-generation process explicates the structure of Agda files at the finest possible resolution: that of the internal type structure. Like Blaauwbroek et al. [2024] , our modeling approach iteratively builds dynamic object representations that are derived on the basis of structural rather than nominal principles. Unlike Blaauwbroek et al. [2024] , our representations are also: (1) complete, in the sense of allowing the representation of any item in the object language, and (2) respectful of α-equivalence, in the sense of being invariant to variable renaming and substitution. 
data: 
 Data is extracted from Agda files in a type-mindful manner. We first allow Agda to type-check program-proofs as usual, and then consult her for their internal representations, which we store mostly open import Relation.Binary.PropositionalEquality using ( ≡ ; refl; cong; trans) data N : Set where zero : N suc :N → N + : N → N → N zero + n = n suc m + n = suc (m + n) +-comm : (m n : N) → m + n ≡ n + m +-comm zero zero = refl +-comm zero (suc n) = cong suc (+-comm zero n) +-comm (suc m) zero = cong suc (+-comm m zero) +-comm (suc m) (suc n) = cong suc (trans (+-suc m n) (+-comm (suc m) n))where +-suc : unchanged. These representations are formally safe and unambiguous, but also more honest than the surface pretty prints we get to witness as users, as syntactic sugar has been translated away.∀ m n → m + suc n ≡ suc (m + n) +-suc zero n = refl +-suc (suc m) n = cong suc (+-suc m n)Figure 1 presents our running example: a mechanised proof that addition on natural numbers is commutative. The program first imports the standard library's propositional equality (≡) and some of its properties (refl, cong, trans) . It then goes on to define Peano natural numbers as an inductive datatype (N) and arithmetic addition as a recursive function (+) via pattern matching: a natural is either a 0 (zero) for which 0 + n = n, or the successor (suc) of a number m, for which (m + 1) + n = (m + n) + 1. With the above in hand, we can inductively prove that addition is commutative (+-comm), i.e., ∀ m n. m + n = n + m, following a case-by-case analysis:1. If m and n are both 0, we can prove our goal using reflexivity: both sides of the equation are syntactically equal.2. If m is 0 and n is the successor of some other n ′ , then we can invoke the proof of the commutativity of m + n ′ .3. Similar to (2): if n is 0 and m is the successor of some other m ′ , then we can invoke the proof of the commutativity of m ′ + n.4. If both m and n are successors of m ′ and n ′ , respectively, then we appeal to a helper lemma about how the successor function distributes across addition (+-suc), before finally invoking the proof of the commutativity of m + n ′ .Note how each case above corresponds to a pattern-matching clause in the function definition and inductive hypotheses manifest as recursive function calls; that is the beauty of dependently-typed functional programming through the lens of the Curry-Howard correspondence.We implement a generic dataset extractor for arbitrary Agda programs as a compilation backend, AGDA2TRAIN, that now "compiles" a source Agda file to a corresponding JSON structure1 . A collection of such structures together make up our extracted dataset.Each extracted JSON file mirrors Agda's internal compilation state for a given file. This state is however neither static nor singular; it changes depending on what the current focal point of the compilation is. From the user's perspective, each state corresponds to a different moment in the coding process, where a subset of the file has been written, another exists only as a hole, and another is yet to be written altogether. We navigate through all such possible states. Concretely, we iterate through each of the definitions in a file, pausing not at the possible prefixes of a proof (seen as a textual object), but rather at all its possible sub-terms (seen as a syntactic tree). This allows us to optimally maximize the utilization of each file as a data resource, akin to a preemptive data augmentation routine (except type-safe, courtesy of Agda). At each step, we record the missing term that fills the hole, as well as the full typing context available at the time.The typing context of a hole consists of its type (the goal), the collection of antecedent definitions in scope, plus a context of locally bound variables. For each definition and variable, we record its name, its type as well as its proof term.Term Structure Terms, whether type-or proof-level, are structurally explicated to the fullest by the JSON structure. Since Agda is a dependently typed language, its terms go far beyond typical λ-expressions that canonically consist of just abstractions, applications and variables. Terms capture Agda's internal grammar in its entirety, allowing the specification of many more constructs, such as datatype declarations, records, function clauses defined via pattern matching, patterns etc.. For the sake of readability (and to appease practitioners more interested in the sequential format), the dataset also includes pretty-printed displays of each AST.For additional details on the extraction, including an example JSON matching the program of Figure 1 , refer to Appendix A. For the reference JSON grammar, refer to Appendix C.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: We do not make any theoretical contributions.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
preliminaries: 
 We begin by recalling that a snapshot of an interactive Agda session can be thought of as an incomplete but type-checked file. The file consists of a collection of definitions (henceforth lemmas), i.e., all proofs or programs within scope (either imported or locally defined). Of these, zero or more might contain holes, i.e., remain unfinished. The distinction between lemmas and holes essentially boils down to the presence or absence of a proof-level term. For the time being, we ignore those, and focus solely on type declarations. Lemmas and holes are then syntactically indistinguishable, since their types adhere to the same grammar.2 Capitalizing on this symmetry, we represent both kinds of objects using a single parameter-sharing encoder. Given a collection of ASTs as input (some being lemmas, others holes), the encoder will be tasked with returning a vectorial representation for each.Conceptual Our primary aim is generalizability, which entails first completeness (i.e., ability to model any data point in the domain, regardless of distributional shifts), and then appropriate inductive biases (i.e., ability to effectively extrapolate from the patterns in the training data). The commonly employed sequential approach attains completeness by tokenizing expressions according to the co-occurrence frequencies of their constituent substrings, either at the "word" or at the "subword" level. In doing so, however, it provides the model with the wrong inductive biases, since the expressions under scrutiny are not formed by just any string grammar, but by precisely constrained inductions. At first glance, one might think to resolve the problem by simply applying a tree-or graph-based architecture on Agda ASTs. The challenge lies in the fact that Agda types, being dependent, are considerably expressive. While simple types are composed of a fixed set of operators and propositional constants (base types), dependent types may contain quantifications (introducing variable bindings), and refer both to other lemmas in scope, and to any variables previously introduced.Case in point, all AST-encoding approaches to date assume a fixed symbol vocabulary. By treating (even just) the primitive nodes denoting constants and variables as strings, they implicitly concede that all occurrences of the same string carry the same denotational meaning. Inadvertently, this also implies that the substitution of such strings with others alters the final representations of the objects these strings occur within, even if such substitutions are semantically void. Worse yet, it means that completeness is sacrificed. Shifts in naming conventions, use of libraries that are stylistically distant or targeted at novel domains, etc., might necessitate the evocation of symbols not present in the trained embedding tables, collapsing representations into meaningless generics (e.g., multiple different entries might be made the same due to an excess of [UNK] tokens). A non-nominal resolution of these complex referencing patterns necessitates an extension of the traditional AST-encoding approaches. Crucially, we need our encoder to generically refer to both locally defined variables (intra-AST) as well as other scope entries in their entirety (inter-AST); see Figure 2 for a visual example.[SOS] Set N [SOS] → → (N) (N) (N) + [SOS] Π Π $ $ (m) $ (n) (+) $ $ (n) $ (m) (+) [≡] (n : N) (m : N) +-commPractical Beyond theoretical concerns, tree-and graph-based approaches are also lacking in several practical aspects. Tree-like approaches involve a bottom-up reduction of the AST into a single representation, often times achieved by translating leaves to vectors and operators to neural functions. Despite offering suitable biases, they struggle to capture long-distance relations, exhibit a temporal complexity that scales linearly with tree depth, and are challenging to parallelize across and within ASTs. In comparison, graph-based models are easier to parallelize, but also offer an appealingly native solution for non-tree edges. Moreover, by fixing the number of message-passing rounds, graph architectures become invariant to AST depth, thresholding temporal complexity. This, however, comes at a cost, as the model's perceptive field becomes significantly limited. Updated node features are obtained from a fixed-size local neighborhood, and their aggregation has no means to account for fine structure at larger scales, leading to a coarse and "lossy" (or oversmoothed) representation of the AST as a whole. As an alternative, Transformer-based encoders implementing full self-attention allow concurrent pairwise interactions between all nodes, regardless of their relative placements. This is particularly advantageous, as it equips the model with the means of capturing patterns and dependencies at arbitrary scales. But while this is true for sequences, a vanilla Transformer lacks the representational biases to perceive the complex inductive structure of an Agda file. Furthermore, its memory footprint scales linearly with the batch size, and quadratically with the longest sample's length. In a premise selection setup, each premise constitutes an "independent" sample, with the total of ASTs being processed making up a (subset of a) batch. This mandates operating on a file-by-file basis, which translates into a requirement for variable-sized batches, during both training and inference. Worse yet, AST counts and sizes tend have a high variance, and can often be overlong, e.g., when dealing with complex files or types. In brief, a vanilla Transformer is practically unusable; we will need to do better.Having identified a fully self-attentive encoder as an attractive modeling option, we set to amend the shortcomings of the Transformer with a modified architecture that is better attuned to the task. In the following paragraphs, we isolate and explicate its building blocks.Efficient Attention To first address the scaling issue, we adopt the linearized attention mechanism of Katharopoulos et al. [2020] . By foregoing the softmax normalization around the query-key dotproduct matrix, this variant allows a sub-quadratic formulation of the attention function. Concretely, given an AST with n input nodes, transformed into queries q : R n×d , keys k : R n×d and values v : R n×e , we obtain an attention output a : R n×e given by:a i = ϕ(R i q i ) • n j=1 ϕ(k j R j ) ⊗ v j ϕ(R i q i ) • n j=1 ϕ(k j R j ) (1)In the above equation, • and ⊗ denote the inner and outer product, applied here over dimensions d and (d, e) respectively. Following the empirical insights of Arora et al. [2023] , we define the feature map ϕ : R d → R 1+d+d 2 as:EQUATIONwhere ; denotes vector concatenation and vec flattens the (d × d)-sized outer product into a d 2 -sized vector. This formulation emulates a softmax-like exponentiation of the vector dot-product via a second degree Taylor polynomial. Finally, R is a tensor packaging n matrices from the orthogonal group O(d) that serve to introduce a positional specification-we spell out their functionality in the next paragraph. In practice, the above operations are broadcasted over multiple ASTs t and attention heads h, yielding an asymptotic complexity of O tnehd 2 . By setting d ≪ N , where N is the maximum AST size encountered, complexity remains tractable throughout training and inference.Structured Attention To enhance the model with the tree-inductive biases necessary to process ASTs, we employ the binary-tree positional encoding scheme of Kogkalidis et al. [2024] . The scheme alters the attention coefficients by rotating and reflecting query and key vectors with positiondependent orthogonal matrices. One such matrix is assembled for every unique AST position using a chain of matrix products made of two primitives. Each primitive, in turn, corresponds to choice of a tree branch option (i.e., left or right). By applying the appropriate orthogonal matrix on each query q i and key k j prior to Taylor expansion, we practically modulate the dot-product between by the norm-preserving bilinear scalar function R ⊤ i R j that represents the relative path between them upon the tree structure. This setup inductively parameterizes the construction of relative paths, allowing the representation of ASTs of arbitrary sizes and shapes using just two primitives.Static Embeddings Initially, we embed AST nodes depending on their content. Primitive symbols with a static, uniform meaning are handled by a simple embedding table. The table contains a distinct entry for each of four type operators (dependent function Π, simple function →, λ-abstraction and function application), three terminal nodes, the content of which we designate as "don't-care" (sorts, levels and literals), and two meta-symbols (the start-of-sequence token [SOS] , and an out-of-scope symbol [OOS], used to denote masked or otherwise unavailable references).Variable Representations Dependent function types and λ-abstractions introduce and bind variables to arbitrary names, which can then be referred to by nodes further down within the same AST. A canonical way to make the representations name-agnostic is de Bruijn indexing [de Bruijn, 1972] , where names are replaced with an integer indicating the relative binder-distance to the variable's introduction. For representation learning purposes, substituting variable names with de Bruijn indices is not much good in itself. Despite the normalized vocabulary, there are still enumerably infinite indices to learn embeddings for. Furthermore, a single index carries no "meaning" of its own, other than as an address-referencing instruction. Even so, its meaning is highly contextual and indirect: the translation from binder-distance to a concrete "address" within the AST depends on the number and placement of variable binders within the referring node's depth-first ancestry. In order to facilitate learning, we take the extra step of actually resolving the indexing instruction, substituting the de Bruijn index with a pointer to the node introducing the referent. To represent the pointer, we then use the positional encoding scheme described earlier. We procure and compose the two orthogonal matrices representing the variable's and the binder's absolute positions to represent their relation, and apply the resulting matrix against a trainable embedding. The final vector enacts a representation of the variable's offset relative to its binder. This constitutes a name-free indexing scheme with a clean and transparent interpretation, abolishing the need for an ad-hoc nominal embedding strategy.Reference Representations Topologically sorted, the lemmas within a file (and, by proxy, their ASTs) form a poset. For lemmas a and b, a ≤ b denotes that the transitive closure of lemmas referenced by b contains a. Based on this ordering, the collection of ASTs can be partitioned into a strictly ordered collection of mutually exclusive subsets, each subset containing the ASTs of lemmas of a specific dependency level. We follow along this sequential order, using the encoder to iteratively populate a dynamic embedding table of AST representations, using each AST's [SOS] token as an aggregate summary. Any time we encounter a node referencing another lemma, we can query the embedding table for the corresponding AST representation. Parallelism is still fully maintained within each partition, where entries do not depend on one another. From a high-level perspective, this makes for an adaptive-depth encoding strategy, implemented by a fixed-depth network. Putting Things Together Following standard practices, we compose an attention block with a SwiGLU [Dauphin et al., 2017 , Shazeer, 2020] , inserting residual connections and prenormalizing with RMSNorm [Zhang and Sennrich, 2019] to obtain a single layer. We compose multiple such layers together, sharing and reusing positional encodings across layers and attention heads.3 Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We provide a detailed account of our experimental setup in the Appendix, and make our code available for review.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiments & results: 
 To quantitatively assess the model's capacity for representation learning, we instantiate it composed with a linear layer on top, and train it in a premise selection setting derived by minimally processing the dataset of Section 3. We call the composite model QUILL, reflecting on it being a featherweight tool (n.b., Agda is a hen) intended to assist in the proof writing process4 . QUILL takes a collection of ASTs as input, but produces now a scalar value for each hole-lemma pair, indicating the latter's relevance to the former.Setup We extensively detail our exact experimental setup in Appendix B. Concisely, we train using InfoNCE on a size-constrained subset of the Agda standard library. Following some necessary processing and filtering steps, this amounts to 481 training and 92 evaluation files, counting 15,244 and 4,120 holes respectively. Out of the 92 evaluation files, 5 files are marked outliers with respect to the distribution of file sizes, containing 25% of the total holes and significantly exceeding the size of largest training file-we take them apart and treat them as a separate OOD evaluation set. To obtain more substantiated insights on OOD performance, we further consider two additional evaluation sets, extracted from portions of the Unimath [Rijke et al.] and TypeTopology [Escardó and contributors] Agda libraries. These are radically distant to the standard library in terms of content, abstraction level, proof style, library structure and file size. Both libraries, and especially so TypeTopology, are infamous for their austere and minimal (affectionately so-called spartan) style, making them ideal test sets at the furthest ends of the modeling domain. The portions we consider are extracts of 137 and 28 files with 5,247 and 1,983 holes for Unimath and TypeTopology respectively.To evaluate performance, we gather the pairwise matching scores for all hole-lemma pairs of the stdlib:ID evalution set, group them by ground truth label and visualize them in Figure 3 . On the standardized axis, we see negative pairs centered around the sample mean at 0. Even though there is overlap between the scores of positive and negative lemmas, 80% of the positive pairs are ranked significantly higher than 80% of negative pairs, with a wide separation between the lowest positive and highest negative quartile. What this means in practice is that the trained model enacts a relatively reliable discriminator, capable of discarding "strongly negative" lemmas (or, vice versa, filtering for "potentially positive" ones).For a more quantitative evaluation, we compare the ranking of lemmas produced by QUILL with the set of relevant items defined by the gold standard labels. We borrow metrics from recommender system literature: average precision (AVEP) and R-precision (R-PREC). Intuitively, AVEP provides an estimate for precision as one iterates through suggestions until all relevant lemmas have been retrieved. R-PREC, on the other hand, assesses the model's effectiveness in placing the set of relevant lemmas (of any size R) at the top R positions among the ranked candidates. We report results in Table 1 (means and 95% confidence intervals of 4 training repetitions).Within the in-distribution evaluation set, we see relevant lemmas interspersed among as many irrelevant ones at the top of the ranked suggestions. This is particularly compelling when compared to the average ratio of relevant lemmas to total candidates, which lies below 1.5%. Unsurprisingly, performance degrades the further away one drifts from the original training distribution. The stdlib:OOD set contains ASTs and scopes that are approximately two orders of magnitude larger than the training mean. On top of ASTs being harder to informatively represent, the average ratio of relevant lemmas to available options is now about 2‰. The other two libraries essentially constitute zero-shot transfer learning setups, intended to stress test the model's representations.Ablations To experimentally validate our design decisions, we conduct an ablation study where we substitute network components for less exotic alternatives. We consider three ablations, first in isolation (i.e., keeping the rest of the model intact), and then in combination. First, we completely skip the Taylor expansion step, opting for an offset-by-one ELU feature map instead. We note a modest but tangible performance hit, but drastically improved memory complexity. Second, we remove the tree-structure bias in the attention function, substituting the multiplicative tree-positional encoding with the additive sequential sinusoids of the vanilla Transformer. The performance effect is notable, as the model is burdened with the added cognitive load of internally "parsing" the sequentalized type expressions. Finally, we remove the structural de Bruijn resolution, treating all varables as syntactically equivalent and mapping them to the same meta-embedding. This is a radical and destructive simplification, collapsing multiple types into a single representation -nonetheless, the model is still able to recover a non-trivial percentage of lemmas, retaining about 75% of the original AVEP. Applied together, the ablations above reduce the model to a standard (linear) Transformer. The baseline Transformer collapses, utterly and completely failing at capturing the task.Put together, our experiments assert the robustness and generality of our approach and affirm the positive effect of each of our design decisions. These findings suggest that our model's success stems from a combination of structure-faithful representations, and targeted, deliberate architectural adjustments that honor them. 
conclusion: 
 In the Future Pending community feedback, there is a number of directions we would be eager to pursue going forward. On the dataset front, we would like to optimize the extraction's performance and coverage beyond the standard library. We further acknowledge that while our structured format is highly elaborate, different levels of abstraction and detail come with their own pros and cons, and are inherently biased towards certain tasks and architectures. That said, the optimal setup would likely be to provide multiple representations and allow the end-user the choice of which one to use; we plan to extend the array of options provided by the extraction. On the integration front, the pipeline implemented so far is unidirectional, in the sense of not providing any hooks for feedback or interaction. Given our strong experimental results, we are currently looking for the best integration angle that can allow at least one full back-and-forth between the language and the machine learning backend. In combination's with Agda's auto, this should enable an end-to-end evaluation in terms of proof completion, but also permit the system's use in real-life scenarios, outside of asynchronous offline settings. Finally, a plethora of improvements and optimizations are possible on the modeling front. Among those, we are primarily interested in seeing our approach applied and evaluated on other languages and especially unifying meta-frameworks such as Dedukti [Assaf et al., 2016] to maximize interoperability.In Sum We have presented a method to expose structured representations of program-proofs extracted from arbitrary Agda libraries and files. Applying the extraction algorithm on the expansive Agda standard library, we have generated and released a high-quality static resource, intended to enable the offline (pre-)training of ML-based proof guidance tools. To our knowledge, this resource is the first of its kind for Agda. Across proof assistants, the resource is also one of the very few that explicates proof structure at the sub-type level rather than just pretty-printed surface representations.Noting the particularities of dependently typed programs, and pinpointing weak points of prior related endeavors, we argued for the difficulty but also the need for structurally faithful models and representations. Utilizing the ultra-high resolution of our extracted resource, alongside recent advances in efficient and structurally biased self-attention, we set out to design and implement such a model. Trained for premise selection, the model achieves remarkably high precision while relying just on type structure alone, experimentally confirming both the utility of the resource and the soundness of our approach. Contrary to recent trends in ML in general and ML-guided theorem proving in particular, our model is tiny, requires no pretraining or server hosting, and makes no use of autoregressive language modeling. We open-source our code and make the extracted resources and trained model weights accessible to the community. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We repeat experiments multiple times, varying the seeds, and report means as well as 68% confidence intervals.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: We do report hardware infrastructure, but not memory consumption or clock times. This is an omission our part, since we did not track these statistics during training.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 • Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite all software libraries and datasets we use, and comply with their licenses.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [No] Justification: We do provide assets, but not as part of the submission at the current stage. Documentation is still pending.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: No human subjects were involved in this study.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: See above.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This paper presents work aiming to advance the fields of Representation Learning and Interactive Theorem Proving. We do not see any potential risks for misuse or negative social impact.right tails: a handful of files contain more than 1,000 definitions in total. Holes are more evenly distributed, except for a high peak at 0 due to files without any holes.Figure 5b similarly presents the empirical distribution of the AST lengths of the types of lemmas and (context-merged) holes. Both distributions are bell-shaped, with holes being in general larger. As before, a handful of ASTs are overlong, counting thousands of tokens.Figure 5c , finally, presents the empirical distribution of lemma occurrences. We obtain this by considering the set of fully qualified lemma names as they occur in the extracted files, and associating each name with two occurrence counts: the number of times the name is encountered across files as an import (i.e., excluding the single time it is defined), and the number of times this lemma is actually selected as a relevant premise by a hole. We only count a single occurrence per hole (i.e., we increment the counter by just 1, even if that lemma is used multiple times). As expected, the plot suggests that most lemmas are only ever used locally, with linearly less lemmas being imported exponentially more times. Analogously, most lemmas rarely ever find their way to a hole, with linearly less lemmas being useful in the derivation of exponentially more holes. This goes to show that premise selection is hard: a few lemmas might be relevant "universally", but the relevance of most lemmas is sparse and contextually decided.The dataset is by no means set in stone and is amenable to various generalizations. Since Agda is a dependently-typed language, terms appear in all kinds of places, e.g., as type-level indices of an inductive datatype. Currently, only terms appearing in functions (i.e., proofs) are considered, but there is nothing preventing us in principle from utilizing terms that appear elsewhere.While the extractor's performance is sufficient for our studies so far, the current implementation is in no way the most efficient possible. First, there is a lot of context sharing between adjacent holes that could be leveraged to drastically reduce disk usage. Second, CPU time and RAM consumption quickly exhibit exponential behavior due to multiple sub-invocations of Agda's type checker, making it impossible to extract data from a handful of complex Agda modules. Finally, the dataset does not currently contain structured decompositions for literals, sorts (i.e., kinds of types), and universe levels. We have conciously ignored them for the time being, seeing as they are largely boilerplate for proof search purposes, but acknowledge that they still might prove useful when reasoning about type hierarchies.We read and tokenize files offline to facilitate training. Figure 2 depicts the structure of the JSON file of Figure 4 (in turn the extract of the Agda code of Figure 1 ) as it is perceived post-tokenization by the neural model. Definitions (proof-level terms) are ignored, and type ASTs are binarized. References (presented in gray font) are resolved, and treated as pointers to either other ASTs or nodes within the same AST, depending on whether the referenced object is another lemma or a variable. Type operators and meta-symbols (presented in normal font) are treated as static elements of a fixed vocabulary. The presentation is the product of some artistic license: in reality, sets are quantified by universe levels, and the equivalence operator is a reference to the actual definition of equivalence.We fix an 85%-15% training-evaluation split over the 818 files extracted from the Agda standard library. After tokenization, we reject 18 files (2.2%) containing mutual inductions, i.e., cyclic reference structures which our model does not yet support. Another 196 files (23.97%) contain no holes, and are thus of no use in this present setup. Finally, we identify 36 files (4.4%) representing size anomalies, each accumulating more than 2 14 tokens. According to our original split, 31 of those were marked for training and 5 for evaluation. We ignore the former, but use the latter as a standalone left-out test, intended to stress the model's generalization. Finally, we inspect the fully qualified names of scope entries and find that the training sets contains 10,320 unique definitions. The evaluation set contains 3,495, with 732 of them being "novel" in the sense of never occurring in the training set (the rest being imported from the training set, or, conversely, having been imported at least once by the training set). All holes in the evaluation set are by construction novel, seeing as they stem from locally defined scope entries. Table 2 : Hyperparameter setupWe generate training batches by randomly selecting a fixed number of files. From each file, we extract the entire scope along with a predetermined number of holes. We use the encoder to obtain vector representations for the types of both the holes and the lemmas within scope. Next, we consider all pairs in the Cartesian product of lemmas and holes within each file but discard causally illegal pairs-those with a hole referencing a lemma that has not yet been defined. For each pair, the gold data provide us with a label indicating whether the indexed lemma actually occurs when filling in the indexed hole. We use this label as an imperfect proxy for the lemma's relevance, i.e., whether it should be used to fill in the indexed hole. Given hole, we want QUILL to be able to identify all relevant lemmas; that is, to successfully discriminate between positive and negative pairs. Translating this requirement to an objective function, we make a minor modification to infoNCE [Oord et al., 2018 ] so as to capture the existence of possibly multiple positive matches. Concretely, denoting with P + h and P - h the sets of positive and negative pairs anchored to a common hole h, we compute a per-hole loss term L h as:EQUATIONwhere f is simply a weighted dot-product reducing each vector pair to a scalar value. We derive a per-batch loss by aggregating the loss terms of all the holes considered.We conduct 5 training iterations using the same model and optimization hyperparameters, reported in Table 2 . We conduct no hyperparameter sweep of any kind, instead opting for efficiency with a sensibly sized model and otherwise default optimization values that maximize GPU utilization. We train on a single A100 for 100 epochs, monitoring training and evaluation performance. Keeping athome deployment in mind, we opt for a modestly sized model. QUILL enumerates fewer than 6 million parameters in total-more than 3 orders of magnitude smaller than the LLM-based architectures currently in use. Once trained and binarized, it takes up a mere 26MB of disk space. Training completes in about 16 hours, but converges halfway through by epoch 55±7, overfitting slowly thereafter.The dataset has a non-trivial JSON structure, so we deem it worthy to document it more precisely here: Fig. 6 presents the JSON schema in a more readable BNF-grammar form. We slightly simplify things by embedding definitions, types, patterns, and term expressions, all in a single syntactic category of "terms" (TERM). While this is not completely accurate (e.g., it does not make sense to have a datatype declaration appear as a function's clause pattern), it strikes a good balance between accuracy and accessibility; the JSON validation that accompanies the JSON schema will make sure that such conditions are always met.A JSON file (FILE) contains scope entries and holes for each sub-term appearing in the source program, where the scope is split into three categories: the global scope of imported definitions; the local scope of definitions declared earlier in the same file; and the private scope of helper definitions that is only visible to the current definition (e.g., arising from a where clause). Each scope entry (SC ITEM) has its type and the actual definition that inhabits this type. Holes (HOLE) are only included for the local or private scope entries, since the modules imported in the global scope already have a corresponding JSON dataset with their corresponding holes. For each hole, we record the context of (typed) bound variables, its type, and the term that successfully fills the hole (from which we also compute the set of used lemmas, specifically for the task of premise selection).Definitions can either declare an algebraic datatype (ADT) consisting of various types of constructors; a record (REC) with fields of certain type; or a function (FUN) that can be defined via pattern matching across multiple clauses (CLAUSE).The rest of the cases correspond to the usual presentation of terms in λ-calculus: variable references (REF DB/REF SC); abstractions (PI/LAM); applications (APP); and primitives for literals/sorts/levels (LIT/SORT/LEVEL). We do not retain the structure of literals, sorts, and levels; there is no real use for it on premise selection, but it can easily be added in the future if the need arises.Last, notice how top-level terms are printed to human-readable strings (P TERM), or even captured at different points of their evaluation (TERM + ). Each of these reduction levels are options (marked with ? ), since evaluating an arbitrary term might prove too costly; only those computed within a reasonable time frame are populated.judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.38th Conference on Neural Information Processing Systems (NeurIPS 2024).Capitalizing on Agda's extensible backend infrastructure, we develop AGDA2TRAIN as a standalone Haskell package, independent of the core Agda compiler. The package is publicly accessible at https://github.com/ omelkonian/agda2train.To be precise, a hole is natively specified as a goal type and a context, i.e., a "telescope" structure of local variable bindings. We use the proof-theoretic equivalence between hypothetical reasoning and dependent functions to merge context and goal into a single type by folding the assumptions provided in the context to a nesting of dependent functions. Put simply, if the context is x : A and the goal is B, the folded type of the hole is x:A B. This effectively homogenizes the representation format, but at the cost of elongating goal types.The orthogonal matrices are applied per-head, and are thus lower-dimensional than the encoder as a whole. This makes them superficially incompatible with the input embeddings. We fix the issue by lowering the embedding dimensionality, and expanding it later via a trainable linear map.Source code, configuration files and scripts supporting experiment execution are available at https: //github.com/konstantinosKokos/quill.https://agda.readthedocs.io/en/v2.6.3/https://github.com/agda/agda-stdlib/releases/tag/v1.7.2 


Decision: There has recently been a great deal of interest in using neural networks (and LLMs in particular) for theorem proving.  There are several formal languages in common use, including Coq, Agda, and Lean, all of which are based on the calculus of constructions, which is a lambda calculus with dependent types.  

Current SOTA when applying language models to code of any kind is to model only the surface syntax of terms.  Despite the fact that the underlying structure of source code is an abstract syntax tree, LMs typically ignore tree structure.  Moreover, LMs pay undue attention to syntactic features like the spelling of variable names, even though formal theories always treat terms as equivalent under alpha-renaming -- which means that it is the structure of the term, not the spelling of variable names, that matters.  As a result, there is a severe mismatch between the way in which LMs are usually applied, and the underlying formal theory used for theorem proving.  

This paper takes some initial steps in correcting this mismatch.  The authors extract the complete internal representation of Agda terms into a JSON format.  This includes not only the tree-structure of terms, but also typing information, goals, lexical scope, etc.  The authors release the processed JSON as a new dataset.  The authors then design a custom transformer which operates on this structure, using tree positional encodings.  They show that the structured representation improves over a vanilla transformer on the task of premise selection.

The reviewers universally recommend "accept" for this paper, and I concur.   The new dataset alone will be useful to other practitioners in this area.  In my opinion, however, the most important contribution of this paper is not necessarily the dataset, or even the details of the architecture and experiments.  The main contribution of this paper is to explain and discuss why the structure of terms is so important in formal languages.  This is an area which up until now has been almost completely ignored, and which will guide future research in LM-based theorem proving. 

The main weakness of the paper, as with many papers that use new datasets and architectures, is the size of the model that the authors are able to train.

Reviews: review: paper_summary: paper_summary: This paper introduces a learning algorithm for Agda, a functional language used  in proof assistance that is known for its dependent types.
The paper includes the design of the algorithm, and an experimental evaluation,
strength_weakness: strength_weakness: The major weakness of the paper is that it assumes familiarity with Agda, The authors drop Figure 1 in the paper but make no effort to explain it. They assume the reader to know what are dependent types. One or two paragraphs would make a difference here.

This makes it hard to understand the results (especially for non-experts like me). You define a positive to be an hole that is filled by a lemma in the proof? Does this mean that a hole/lemma pair could have been a positive, it just wasn´t considered? Also, is this for any proof? If this is a proof assistant, shouldn´t it be about a specific hole in a specific sequence of steps? , In B.1 you do mention scope and selecting a fixed number of holes. can you be more specific? Also how many positive and negative examples do you actually have

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

Finally in "Practical' you seem to consider either Tree Based Models or Transformers. Any specific reason to ignore GNNs?
questions: questions: I mentioned above my main suggestions. Other details:

Fig 1, please explain

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

Page 6: it seems that the Taylor expansion is how you achieve global features, which seems to be goal here. Yet, later on you drop Taylot expansion and still obtain good results.

J. Rute, M. Olsˇak, L. Blaauwbroek, F. I. S. Massolo, J. Piepenbrock, and V. Pestun. Graph2tac: ´
442 Learning hierarchical representations of math concepts in theorem proving, 2024.
443 N -> Missing venue
limitations: limitations: The authors do not see potential negative societal impact, and I would agree.
paper_summary: paper_summary: The paper reported the creation of a new dataset to predict a term for filling the hole of a proof in Agda. The paper also proposed an attention-based neural architecture, trained it on the made dataset, and shows it outperforms the ordinary Transformer.
strength_weakness: strength_weakness: - The importance of creating a new dataset could be discussed more appropriately and carefully. Specifically, I'd like to see why the dataset is necessary or valuable although there exist the datasets for automated theorem proving already (like the ones mentioned in the related work section and MiniF2F).  I don't think it is sufficient just to say "there doesn't exist a dataset for Agda" because then one just could adopt the non-Agda datasets.
- The experimental setting is unclear for me. Specifically, how is Transformer set up? Is it fine-tuned on the Agda dataset? If so, it is trained against the data in JSON format or text format? If the JSON format is adopted, I suspect it might be unfavorable for Transformer because it needs to parse JSON data by itself.
- I'm unsure Transformer is appropriate as a baseline because, even if the models are restricted to be Transformer-based, there exist many previous works that claim their models outperform the vanilla Transformer as, e.g., the following [1]

[1] Stanislas Polu, Ilya Sutskever: Generative Language Modeling for Automated Theorem Proving. CoRR abs/2009.03393 (2020).

- The presentation could be improved. The following are unclear for me.
  - What is the outputs of the proposed model? An AST of a proof term? Or, just a text?
  - What is the design principles of the proposed model? Why did the authors think efficient and structured attention work well on the task of interest?
  - In Figure 3, what the x- and y-axis represent, respectively? 
  - What is "the added cognitive load" mentioned in Section 5.2 (page 9)?
questions: questions: * Why is creating a dataset on Agda important?
* Why did the paper choose the vanilla Transformer as a baseline?
* Was the Transformer fine-tuned? If so, how?
limitations: limitations: Certain limitations are described in the paper.
paper_summary: paper_summary: This paper is the first to extract data from Agda, a dependently typed programming language. The extracted data can be utilized by machine learning practitioners. Additionally, the paper proposes a novel neural structure designed to faithfully represent dependently typed programs. The experiments demonstrate the effectiveness of the proposed neural architecture.
strength_weakness: strength_weakness: I don’t see any major weaknesses in this paper, except that it is somewhat difficult to follow. It would definitely benefit from more illustrative figures to explain various concepts. Concepts in Agda, such as dynamic compilation, “holes,” typing contexts, and term structures, are hard to understand. Including figures with concrete examples of these components, in addition to textual descriptions, would be very helpful.
questions: questions: - Typos in line 162 `that that`.
limitations: limitations: The authors have adequately addressed the limitations and posed no negative societal impact.
paper_summary: paper_summary: The article contributes to the field by detailing a method to capture and utilize the internal compilation states of Agda, a dependently typed programming language, through JSON files. These JSON files represent different compilation states, reflecting various stages of the coding process. Key contributions of the article include:

- It explains how each JSON file mirrors Agda's internal state during compilation, which is dynamic and changes based on the current focus of the compilation process. This approach captures the evolution of the file content at different coding stages.

- The method involves iterating through each definition within a file, focusing on all possible sub-terms rather than just potential prefixes of a proof. This detailed navigation allows for a more comprehensive utilization of the file as a data resource, enhancing the understanding of Agda’s compilation process.

- By treating each file as a data resource and navigating its states preemptively, the approach acts like a type-safe data augmentation routine. This maximizes the resource utilization of the file content.

- The article also highlights how each state's typing context is recorded, including the type, relevant definitions, and local variables. Additionally, it thoroughly explicates the structural aspects of terms in Agda, going beyond simple λ-expressions to include a variety of constructs like datatype declarations and pattern matching, which are all captured within the JSON structure.

- For better readability and to aid practitioners interested in a more linear format, the dataset includes pretty-printed displays of each abstract syntax tree (AST).

The article provides foundational insights into leveraging Agda’s internal mechanisms for better data handling and representation in JSON format, which could be useful for both researchers and developers working with dependently typed languages.
strength_weakness: strength_weakness: - Although it's within the setting of premise selection, I still hope to provide a usable baseline for search-based ATP.  
- Can you provide some visual results on representation?
questions: questions: - Can you provide some visual results on representation?  
- Are you asking if it's possible to commit to open-sourcing the code in the future?
limitations: limitations: - Premise selection is an interesting area, but it cannot truly generate theorems. It could be considered for expansion into a more general setting.

Comments: comment: 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).
Thank you for your response. I think these models[1][2] can serve as baselines. Could their experiments be supplemented?

**Reference:**  
[1]Wang, Y. S., Lee, H. Y., & Chen, Y. N. (2019). Tree transformer: Integrating tree structures into self-attention. arXiv preprint arXiv:1909.06639.  
[2]Nguyen, X. P., Joty, S., Hoi, S. C., & Socher, R. (2020). Tree-structured attention with hierarchical accumulation. arXiv preprint arXiv:2002.08046.
Thank you for the nice questions :)
Thanks for the nice replies
Thanks once more for your engagement, we appreciate the ongoing discussion.

You could certainly use a boolean mask to block out specific attention coefficients (*e.g.* you could attend to just "descendents" in the AST, or even something more elaborate), but it's not immediately clear whether this could faithfully capture the AST structure. Going for something more radical (*e.g.* attending to just some local neighborhood of nodes, given an appropriate definition of local) would result in the problems we attribute to GNNs (lines 211 - 217). Conceptually, this *should* also prove inadequate for capturing the long distance dependencies commonly found in "real life" dependent types (even though empirically this *could* also turn out not to be the case).

Did you have something specific in mind? Is there some experiment that you'd like us to do along these lines (keeping the time window in mind, of course)
Thank you for your experiments.  I believe there are other ways to inject structural bias into an autoregressive model through RoPE, such as using a Mask matrix to block certain connections between tokens.
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)
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.
We remark once more that our model is not directly compatible with or comparable to text-based LLMs (autoregressive, like GPT, or encoder-only, like BERT). The model does *not* operate on textual token sequences, but on complex structures made of a (really) small number of primitive symbols (namely: `Π`, `->`, `λ`, `@`, `Sort`, `Level`, `Literal`, `deBruijn`), plus references.

That said, we can still attempt to train a YaRN-based sequential version, but from-scratch (rather than pretrained).
> `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.
**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.
You're very welcome!

 The model we revert to after all ablations in Table 1 is no longer customized -- it is basically a run-of-the-mill vanilla Transformer. It is a useful baseline exactly because:
1. it is the most common architecture empowering modern state of the art LLMs that find use in ATP tasks, and
2. it is what our model boils down to when we remove all of the task-adapted components
Thanks for the clarification for the baseline. I'm not very convinced whether using a customized Transformer as a baseline is appropriate, but it would be helpful to describe in the main text the detail of the baseline and why it is used unless I missed an explanation in the main text of the submission.
* 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:** 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.
* No, we pick **all** negatives (in the current file). We experimented briefly with picking only hard/easy/random negatives, but this vanilla setting proved to work quite well.
* No, it's trained using the exact same experimental setup, *i.e.*, premise-selection head and the modified infoNCE loss.
- **negative samples selected:** I see. So, actually, the negative samples are randomly selected excluding the positive samples?
- **Supplement the baseline:** Is the vanilla Transformer in the last row of Table 1 trained with autoregressive loss? I don't see any experimental results comparing the vanilla Transformer with Equation 3 in Table 1.
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)
**Thank you for your response.**  
- I would like to know the details of how you construct multiple positive samples.
- What does "negative pairs anchored to a common hole h" mean? Are the negative samples selected using the in-batch trick?
- I would like to know if, as long as a suitable method for constructing positive and negative samples is used, the architecture becomes less important for the premise selection task, and a standard bidirectional encoder model can also be effective. I believe that a vanilla Transformer combined with a contrastive learning objective is also an important baseline.
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.

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