Algorithmic Language Models with Neurally Compiled Libraries

We appreciate your review, but your summary suggests that you may have misread our paper. To clarify, we are using differentiable programs as a basis for training on answers only, not training on intermediate parsed values. In this regime, parsing may be learned implicitly, which is much more difficult than training it directly. The purpose of doing this is to test the limits of differentiability alone. To the best of our knowledge, this has not been done in other LLM Augmentation research and represents a major qualitative difference. Given this difference, we would request you to reconsider your evaluation. We have made modifications to the writing to hopefully make the main points of our paper clearer, and hope that they help you.

In the results we present, all parameters of LLaMA are trained, not just the last layer. In general we have found this to be far more effective.

We agree about the strengths of discrete program search methods, however the point of this paper is not to use search for explicit programs but rather to learn how to compose existing programs for natural language tasks, something which is uniquely possible with a differentiable program library that augments an LLM. The purpose of experiments which select only a single program is to establish if the gradient signal from the differentiable computer is sufficient to do this in practice. However, experiments which compose multiple library functions are more compelling.

We have removed speculative statements or supported them with experiments or citations.

For calculating text accuracy, we use a binary correctness: A list is either sorted correctly or isn’t, and the accuracy is just the percentage of correctly sorted problems.

Clarify the starting library for each problem:

We have updated the text to make the starting libraries clearer, but in summary we first experiment with a calculator-only augmentation, and again with an augmentation that uses a small program library. In the case of the calculator, the model classifies which operation to perform. In the case of the library, the model classifies which functions to call.

Why differentiate sorting using differentiable register machines instead of many other works that potentially provide better gradient estimations?

We agree that discussing gradient estimation techniques is valuable, so we have added an appropriate section to our related work.

Again, we thank you for reading and reviewing our paper, and will be happy to help in clarifying it further if needed.

I have read the authors' responses. To clarify, I did not misunderstand the paper as the authors claimed. The current formulation of the method is still mapping input texts to programs whose execution results are mapped to texts as outputs, as I summarized. Program parsers and neural networks are two separate parts (regardless of how they are trained) and we can easily replace the parsers with either pretrained LLMs + prompts or many other traditional parsers (discrete or not) to solve the problems in this paper. That was my point.

I would request the authors to check my comments again. Some quick experiments to run to make this paper more convincing, at least from my perspective:

• Comparison with LLM Tool argument baselines, e.g., Chain-of-Code.
• You can come up with some foreign language so that the pre-trained models won't work and joint learning would be valuable.
• Comparison with simple parsers: Given the small library and program space, I would guess you can get pseudo-labels to train parsers by enumerating programs.
• Evaluate the methods on more algorithmic tasks.

For the future, I would suggest the authors think about how to integrate programs and neural networks more seamlessly, e.g., call programs more than once, programs calling neural networks, and non-obvious mapping from texts to programs.

The authors' proposal is definitely original, the paper outlines it in a mostly clear manner, and there is reason to believe that such augmentations, once refined and properly scaled, could indeed prove to be invaluable in making foundation models capable of algorithmic reasoning.

We greatly appreciate the reviewer’s assessment, and while we acknowledge the preliminary nature of the work, we also believe that, even in the current state, there is significant value to the research community in the sense that other research work may be able to build on what we have done so far.

We'd like to point out that the main point of our paper is in identifying precisely how such techniques could be refined and scaled -- in particular, we establish/confirm that gradient descent based training is highly reliant on shallow and highly parallel gradient paths that are less noisy, which suggests that our approach would work far better if done with a highly parallel differentiable computer rather than the sequential one we have opted for. Accordingly, we believe the paper has value to the ICLR community in the current form, in that it gives clear next steps for this sub-field.

In this regard, I cannot provide any more suggestions for improvement than the authors already do in section 6

Indeed, we are well-aware of the potential ways this paper could be improved, and hope to make significant progress in doing so. At the same time, some limitations are fundamental (e.g. sequential computation <> differentiability), while other limitations are out of scope of a single paper (e.g. numeric representations, tokenization). With this in mind, we find it more appropriate to publish our intermediate work in the interest of exploring other aspects of this problem, namely highly parallel differentiable computing.

Length generalization

We agree with the reviewers assessment and highlight length generalization is an essential reason why neurally compiled libraries would be desirable. We may update the paper with results in length-generalization

Am I to understand that none of the differentiable intepreter's parameters are trained? If not, which ones are?

Yes, it is accurate to say that the differentiable interpreter is not trained, as in general it is not parameterized. Technically the program library itself and aspects such as lookup tables could be parameterized, but we opt for them to remain static since this maintains guarantees of their functionality. Accordingly, the only components that are trained are the transformer itself and specialized classification layers that convert the final embedding into classifications of inputs and selected programs.

Related to the above, it is still not clear to me which algorithms constitute the compiled library for the experiments in section 4 and, most importantly, how are these algorithms compiled into the model before the training runs.

We have updated the manuscript to clarify what the initial library looks like in each experiment.

In section 4.1, the authors assert that a differentiable circuit ALU generalises "beyond lookup tables". How can this be deduced by figure 3, where text accuracy for differentiable circuits is actually lower than for a lookup tables?

The text accuracy is a side-effect of trained parsing ability, not the augmentation. When we say that a circuit generalizes further, we are actually talking about length generalization! A table generalizes only to a fixed length (it only supports modular arithmetic), but a circuit can generalize to arbitrarily long numbers without modification (if/when the inputs are parsed correctly and the operator is chosen correctly!).

How did the authors pick the training hyperparameters used in section 4.2?

The hyperparameters used are based on Meta’s recommendations for fine-tuning, as well as manual experimentation that confirmed these as effective choices.

In section 4.3, to the authors consider generalising to sorting longer lists at test time (e.g., training on list of size 8 and testing on strings of size 12)

This is a valuable suggestion.

Again in section 4.3, the authors attribute the weak performance of their augmented model to the difficulty of parsing the natural language description of the problem. Can they provide any evidence (e.g. in the appendix) for this statement via some ablation?

We are currently running experiments which do sorting on prompts that are less sensitive to external factors, e.g. tokenization.

We thank you again for your diligent review and questions.

Thank you for recognizing our direction as important and interesting. We understand why you may conclude our work is preliminary, but hope to convince you it has value to the ICLR community in the current form.

Reasoning, arithmetic and algorithmic abilities are still weak spot of LLM . 'Algorithmic Language Models with Neurally Compiled Libraries' suggest interesting and promising approach to improve capabilities of LLM.

This is precisely the point of the paper -- reasoning abilities in LLMs are currently lacking, and will rely on more fundamental changes, namely algorithmic ability, to improve. LLM reasoning is one of the most important research topics currently, and we hope our paper progresses this line of research in a unique way.

Neurally Compiled Library is primary experimental work. Therefore I believe work would greatly benefit from extending it's evaluation on more, preferable public datasets.

We agree about evaluating on public datasets. In particular, we would like to extend our evaluation to the CLRS benchmark [1]. In the current form, we use synthetic datasets which are essentially a subset of CLRS.

We plan to add baselines before the end of the review period.

Model background section would benefit from either added citations or clear indication what was proposed in current work. (I.e Differentiable Registers, Probabilistic execution, Differentiable Interpreter sections)

We have further modified the methods section of the paper to make it clearer what is being introduced in this work. Also while it is fair to say that the paper is heavily reliant on methods introduced in Bunel 2016, we’d like to point out the critical difference that Bunel 2016 focused on learning contextual programs (e.g. that perform better on biased data), while we focus on using unaltered programs as a method for augmenting LLMs. So while the math is very similar, there is a significant qualitative difference in the purpose it is used for (copied from response to reviewer EFCQ).

On page 4 'instruction counter' and 'program counter' are mentioned. Could you please clarify what is the difference?

We have used these terms interchangably, but have updated the paper to prefer 'instruction counter', which is more accurate.

Discussion of sorting performance

Please see our overall response, however we believe the problem of learning to parse/select the differentiable insertion sorting algorithm from answer supervision alone is quite difficult. Regardless, we plan to provide baselines for comparison and ideally an ablation study testing if our intuitions about why performance is bad (e.g. tokenization) are correct.

Furthermore, we do not really expect the differentiable insertion sort algorithm to succeed, because the gradient path produced by the algorithm is so long and noisy. We use this as evidence for our main insight, namely that a differentiable augmentation of an LLM should focus on highly parallel algorithms.

our results highlight the overall shortcomings of gradient-descent based learning, given that even when the ideal algorithm is already present, there are still scenarios where the model may not learn a perfectly generalizable solution.' This is very interesting finding. I think both paper and community would benefit if more details on training difficulties, what didn't work, etc would be given.

We agree, and will expand on this in our paper. Please see [2] which has a very similar finding.

[1] Veličković, P., Badia, A. P., Budden, D., Pascanu, R., Banino, A., Dashevskiy, M., ... & Blundell, C. (2022, June). The CLRS algorithmic reasoning benchmark. In International Conference on Machine Learning (pp. 22084-22102). PMLR.

[2] Zhang, H., Li, L. H., Meng, T., Chang, K. W., & Broeck, G. V. D. (2022). On the paradox of learning to reason from data. arXiv preprint arXiv:2205.11502.

Thank you for your diligent and thoughtful review of our paper.

Poor introduction: most of the introduction, apart from the very last paragraph, is dedicated to motivating the work. The very last paragraph has 1 sentence which describes the methodology.

We have modified the last paragraph of our introduction to be dedicated to describing the methodology at a high level. At the same time, we believe the motivation sections are very important, as it seems to have sold all four reviewers on the premise of our paper.

Clarify contributions

We have edited several sections of the methodology section to make it clearer what our contributions are. Also, we have significantly expanded section 3.4. While it is fair to say that the paper is heavily reliant on methods introduced in Bunel 2016, we’d like to point out the critical difference that Bunel 2016 focused on learning contextual programs (e.g. that perform better on biased data), while we focus on using unaltered programs as a method for augmenting LLMs. So while the math is very similar, there is a significant qualitative difference in the purpose it is used for.

Clarify experimental setups

In 4.1, 4.2, and 4.3, we have added more detail regarding which libraries were used as input. We believe there is sufficient detail regarding the neural architectures (e.g. "Our minimum viable transformer uses the LLaMA architecture with tiny scale parameters.. a scaled LLaMA with a dimension of 128, 4 transformer heads, 2 key-value heads, and 4 transformer layers. This model is trained from scratch using the AdamW optimizer [50], a batch size of 32 and a learning rate of 1 × 10−2"). However, we will add more detail to the appendix.

baselines and results

We plan to add more baselines before the end of the review period, and have clarified our claims relative to our empirical results. Please also see the overall response.

How do you differentiate between the selection of different programs? Do you run all of them at once? If not, how do you expect the LLM to learn which program to select if it’s trained only on the error of the outputs?

In the cases where there are multiple functions available in the library, the model selects one of these functions as a softmax classification, and effectively all available functions are run at once. In principle this allows the model to potentially learn selection from only answer supervision (which it is able to do for arithmetic for instance, where it learns to classify operators without operator supervision), however it is also feasible that some intermediates, e.g. selected library functions, could be supervised.

What are the novel ideas which one can take away from the methodology of this paper?

Augmenting a LLM with a neurally compiled library has never been tried before, and to our knowledge this is the first attempt at doing so, and compared to similar approaches (e.g. [1]), represents an advance in the state of the art.

Beyond the novelty of the methods themselves, there are several insights to be gained from the work, even in its preliminary form. We discuss these in the final section of our paper, but one that stands out is: differentiable augmentations will be more effective for highly-parallel and shallow programs, e.g. future work should focus on machines/algorithms that exploit this, rather than a sequential RNN-like augmentation that we attempt.

[1] Bounsi, W., Ibarz, B., Dudzik, A., Hamrick, J. B., Markeeva, L., Vitvitskyi, A., ... & Veličković, P. (2024). Transformers meet Neural Algorithmic Reasoners. arXiv preprint arXiv:2406.09308.