Searching Latent Program Spaces

We thank the reviewer for their detailed review and appreciate their recognition of the creativity and efficiency of LPN relative to TTT.

Table 3 reports ARC-AGI performance as the percentage of tasks solved. We will update the table caption to clarify this.

Below, we respond to the other questions posed.

Q1. “... I'm skeptical that it could compositionally generalize, because fundamentally the "program" is being "executed" by a neural network, and has to be represented by a fixed length vector. Did you ever see evidence of compositional generalization?”

We agree that since the latent program is ultimately executed by a neural network, there are inherent limitations to the compositionality achievable in this representation. LPN’s ability to perform compositionally generalization depends on the nature of the considered generalization. We investigated a variety of types of compositional generalization and found that the search component of LPN enables generalization of type Compose-Different-Concepts and Switch-Concept-Order outlined in [1].

We ran experiments combining programs from the training set in novel ways not seen during training. For example, we tested a task that required both filling black cells within objects with a specific color and cropping a grid based on a rectangle, which are two operations that were only ever seen separately during training. While the encoder alone failed to generalize (only executing the cropping task), LPN with gradient-based search was able to successfully compose the two after 500 gradient steps. This demonstrates that the smooth, learned latent space supports compositional generalization beyond what was explicitly trained. We agree with the reviewer that illustrating the model's generalization capabilities adds clarity, and we will include this analysis in the appendix.

However, we find no evidence in our experiments to suggest that LPN is capable of generalising for length-based recurrence or applying concepts in structurally different ways (e.g., chaining programs longer than those seen during training). We will also include such negative examples in the appendix. As noted on line 394, future work could explore discrete latent representations, potentially enabling more robust compositionality. However, learning such representations is non-trivial and must be weighed against increased training complexity.

[1] Shi, Kensen, et al. "Compositional generalization and decomposition in neural program synthesis." arXiv preprint arXiv:2204.03758 (2022).

Q2. Did you ever try also also optimizing the parameters of the encoder/decoder during test-time optimization?

We did consider the possibility of hybridizing TTT and LPN by combining LPN’s latent space search with decoder parameter updates at test time. This could help when the decoder cannot fully execute a discovered latent program. While this may improve performance for complex programs, it would also introduce parameter updates, diverging from LPN's core advantages of efficiency and reduced overfitting. In both the pattern and systematic sequence tasks, parameter fine-tuning degraded performance. Only on the ARC-AGI benchmark have we observed increased performance, likely due to decoder capacity limitations. We chose not to include this hybrid in our initial work to maintain focus on LPN's unique strength: efficiently adapting without weight updates. We hypothesize that as the diversity of programs during training increases, the need for parameter updates will diminish, reinforcing the value of search-based adaptation.

Q3. Did you ever try having the encoder take multiple examples as input?

We did explore feeding multiple input-output pairs to the encoder in early versions of LPN. However, this approach had key downsides: It led the encoder to overfit to the number of examples seen during training, reducing its ability to generalize across varying specification sizes (as seen in Figure 3). It broke the permutation invariance of the I/O pairs within the specification, which is a desirable property. Encoding pairs independently allows for parallel computation and faster training. That said, the reviewer raises a valid point: ideally, representations of individual pairs could inform one another. A promising direction for future work would be to, e.g., encode each pair independently, then apply attention-based aggregation to allow cross-pair interaction and refinement.

Additional Ablation

We also draw the reviewers’ attention to an additional ablation conducted for reviewer qA5h. This evaluates different LPN test-time inference strategies on ARC-AGI, providing further insight into the roles of the program prior and gradient-based search in driving performance.

We thank the reviewer for their detailed feedback and address each comment below.

We acknowledge that most evaluations focus on grid-based reasoning tasks. However, we direct the reviewer’s attention to the sequence-based program tasks discussed in Section B.7, which do not rely on a 2D grid structure. These tasks further illustrate the tendency of TTT to overfit and demonstrate the consistent improvement of LPN with additional gradient steps.

Contextualizing ARC-AGI Results

Regarding comparisons with other methods on ARC-AGI, we emphasize the importance of controlling variables carefully when evaluating selected test-time adaptation methods (i.e., In-context, LPN, TTT). High-performing methods such as the o3-preview-low model or ARC-AGI 2024 winners differ significantly in training, test-time data, and computational assumptions. Our selection of the ARC-AGI benchmark for evaluation was specifically aimed at isolating and understanding the dynamics of test-time adaptation strategies, rather than maximizing performance through specific methodological biases or increased compute.

However, we agree that providing context about the diverse performance levels and computational assumptions across different methods is essential. To enhance clarity, we will include a summary table outlining various methodologies, their model sizes, training strategies, test-time adaptation approaches, and respective performances.

(ARC-AGI Components): whether the methodology directly targets the ARC-AGI benchmark

[1] ARC-specific data augmentations for training and test time, ARC-specific tokenization scheme

[2] ARC-specific DSL

Answers to Questions

Q1. The core experiments use small models which raise question about how LPN scales to the size of state-of-the-art LLMs or more complex program spaces. Demonstrating consistent LPN benefits with larger models would validate scalability and robustness claims.

We first note that, although the LPN models trained for ARC-AGI might be small compared to large language models, they are still relatively large, with approximately 178M parameters. These models require pixel-perfect execution on 900-pixel grids across hundreds of programs, demanding significant capability from the neural decoder. Regarding scaling LPN to state-of-the-art LLMs, we agree that this is an exciting direction for future research. However, state-of-the-art LLMs typically have around 200 billion parameters, and training such models exceeds the scope of this initial study, for engineering and financial reasons. Nevertheless, our experiments demonstrate LPN's consistent advantages across different scales, from 1M parameters (pattern tasks) to 178M parameters (ARC-AGI).

Q2. How LPN methods would compare with the results presented in [1]. [1] "Combining Induction and Transduction for Abstract Reasoning" Li et al. 2024

We appreciate the reviewer’s request for comparison with Li et al. (2024), which also evaluates methods on the ARC-AGI benchmark. We have included their models in the summary table provided above, alongside different methods and assumptions. It's important to note that their work uses significantly different compute and data assumptions than ours. Specifically, their experiments rely on a pre-trained language model (LLaMA 3.1-8B-Instruct) for both induction and transduction strategies. Methodologically, their transduction approach aligns with our transduction baseline in that both use similar inference strategies: they do not explicitly represent programs but instead predict outputs directly from input-output examples. However, our work differs in its assumptions regarding model size and pre-training data. These distinctions are critical to contextualizing performance and understanding the broader implications of test-time adaptation methods like LPN.

We thank the reviewer for their thorough evaluation, highlighting the method's novelty, clear writing, and detailed ablations. Below, we address each question and comment provided.

We note the reviewer’s suggestion to improve the explanation of baselines in the main text. We acknowledge this and will revise the main text to include sufficient details about the baselines.

Contextualizing ARC-AGI Results

The reviewer accurately notes that certain approaches achieve higher scores on the ARC-AGI benchmark but typically involve substantial, often ARC-specific additions (e.g., pre-trained and fine-tuned language models, synthetic data generation at test time). Our primary goal was to fairly compare test-time adaptation methods (In-context, TTT, and LPN) while holding other factors constant. Hence, direct comparisons with methods employing ARC-specific enhancements or additional pre-trained models would not be equitable.

However, we agree that contextualizing our method's performance relative to various approaches is valuable, regardless of differing computational requirements or assumptions. We have thus compiled a comprehensive table summarizing a range of methods, pre-training and post-training strategies, and ARC-AGI scores. This provides a clearer context regarding LPN's performance relative to existing methods.

[See Review H6z6 for table]

Answers to questions

Q1. “any theories why more grad steps doesn't work?”

One theory is that using more gradient steps during training reduces the pressure on the encoder to produce a strong initialization for the decoder since extra gradient steps can compensate for a weaker starting point. Thus, beyond a certain threshold, the encoder may receive a less effective training signal, resulting in poorer initialization given the same computational budget. Our experiment (Figure 2) underscores the encoder's importance, indicating that weaker encoder performance might negatively impact test-time search. Table 1 supports this observation: training with 5 gradient steps (Grad 5) results in 0.0% encoder-only performance, whereas training with 0 steps (Grad 0) achieves 3.2%. However, with at least 5 test-time gradient steps, the Grad 5 model performs better. Grad 1 may strike the optimal balance, training a strong encoder for good initialization while still allowing effective latent space refinement through gradient steps at test time. To mitigate this issue, we applied progressive training, from 0 steps gradually increasing to N steps (see line 366), which ensures the encoder initially receives strong training signals. Gradient steps are only increased once the model converges, preventing encoder underfitting. Additionally, we hypothesize that doing more gradient steps in the latent space at test-time might eventually overfit the latent to the specification (train input-output pairs) and fail to generalize to the test input. For instance, in Table 3, there is some evidence that suggests minor overfitting occurs on the in-distribution dataset when going from 2e13 to 2e15 flops, with additional search marginally lowering performance. In the future, it would be valuable to study the extent to which specification overfitting occurs in LPN.

Q2. on page 4, you say "we diverge from such transduction-based methods, as they cannot inherently adapt at test-time or ensure specification consistency" — isn't this exactly what TTT does, which transduction-based methods are compatible with.

The reviewer makes a valid point. Our original wording intended to emphasize that transduction-based methods (e.g., TTT) typically lack an inherent mechanism for adapting to new data without modifying the model's parameters directly (unlike in-context learning and LPN). We acknowledge this was unclear and will clarify this distinction in the final paper, explicitly highlighting the benefits of methods capable of adapting to new data and maintaining consistency without relying on parameter-based fine-tuning.

Q3. for table 3 is "in distribution" the re-arc dataset? is it the training set, or a validation set from RE-ARC? and then ood is ARC-AGI validation? maybe I missed this in the main text, but if it's not there it should be described there (and even better in the table too).

The in-distribution dataset refers to the ARC-AGI training set, and the out-of-distribution (OOD) dataset is the ARC-AGI evaluation set. Neither dataset is used directly for training; instead, we use data generated by RE-ARC during training. However, RE-ARC tasks are built upon primitives from the ARC training set, making them effectively in-distribution relative to the ARC training set. While we attempt to explain this distinction in the description of Table 3, we acknowledge that it can be clearer and will provide additional clarification in the final version of the manuscript.

Additional LPN experiments on ARC-AGI

The reviewer notes the lack of analysis regarding the specific components contributing to LPN's performance on ARC-AGI. To address this, we conducted additional ablations exploring different test-time inference strategies to better understand the impacts of the program prior and gradient-based search. We evaluated performance on a single seed on ARC-AGI eval dataset (out of distribution relative to re-arc), keeping constant the number of sampled latents during inference, varying budgets in [10,50,100,200,400].

Specifically, we compared the following inference strategies.

Encoder Sampling: Samples multiple latents from the encoder, evaluates the likelihood of the specification given each latent, and selects the best. This tests the value of structured, gradient-based search versus brute-force sampling near the encoder output.

Gaussian Init + Gradient Search: Uses the same gradient-based search as core LPN but initializes from a sample drawn from the prior (Gaussian). This isolates the encoder’s role in both suggesting good initial points and improving overall performance.

Encoder Init + Local Search: Starts from the encoder output but replaces gradient updates with a mutation operator (continuous noise perturbation). This examines the benefit of gradient-based updates versus unstructured local search, while still allowing movement beyond the encoder distribution.

Gaussian Init + Local Search: Combines prior-based initialization with local search. This baseline helps evaluate the importance of both encoder initialization and gradient guidance.

Encoder Init + Gradient Search (LPN): The standard LPN inference method.

Below is the table comparing these different inference strategies with varying search budgets.

Summary of Results

Repeated sampling from the encoder offers slight improvements over a single sample (7.75%), but the gains are minimal. Comparing Gaussian Init + Gradient Search to Encoder Init + Gradient Search highlights the encoder's importance: it not only boosts performance with few gradient steps but also significantly improves results at high budgets (e.g., 400 steps). Encoder Init + Local Search performs similarly to gradient-based search at low budgets (10 steps), but diverges sharply as the budget increases, demonstrating the effectiveness of gradient guidance. Gaussian Init + Local Search performs poorly, with results near 0%, showing the combined importance of good initialization (via the encoder) and gradient-based search. Overall, Encoder Init + Gradient Search (LPN) consistently outperforms all other methods, confirming that LPN’s core components, i.e., encoder-informed initialization and gradient-based optimization, are essential for strong performance.

We thank the reviewer for their detailed review, for highlighting that our code is fully open-source for reproducibility, and for noting the strength of the out-of-distribution (OOD) performance of LPN.

One highlighted strength was the recognition of LPN as a type of variational autoencoder (VAE), a point with which we fully agree. We clarify this further in Appendix D, explicitly framing LPN as a VAE and describing the latent optimization step as semi-amortised variational inference (SVI). Specifically, at test time, SVI refines the aggregate posterior for individual data points to close the amortization gap remaining after training. While we reference this discussion in line 182, we agree that explicitly stating in the main text that LPN performs variational inference would enhance clarity.

To address the reviewer's concern about our experiments focusing narrowly on visual grid-like tasks, we highlight an additional experiment in Appendix B.7. Here, we explicitly compare LPN and test-time-training (TTT) on sequence-based programs without grid-based features, confirming that our method is not limited to grids. This experiment further demonstrates TTT's tendency to overfit and supports our claim that LPN effectively improves performance through latent-space gradient updates at test time.

Regarding the comment that our definition of "program" implicitly focuses on transformations of 2D grids and thus may not fully support a general claim, we agree that our experiments do not show LPN learning a fully general space of all programs. Indeed, LPN intentionally leverages biases from its training distribution to efficiently represent a specific class of programs tied to a particular input-output distribution. We believe LPN is general in the sense that it makes minimal assumptions beyond having data in the form of N input-output pairs. Our experiments demonstrate that test-time adaptation improves performance consistently across all three tested environments and provides robust generalization to out-of-distribution tasks. However, we acknowledge the reviewer's concern. We will update the manuscript, in both the introduction and discussion, to clarify that our method targets program synthesis specifically within programming-by-example tasks, exploring grid-based and integer-sequence domains. Training LPN on a truly general program space would require extensive pre-training on diverse datasets and remains an exciting direction for future research, given sufficient compute resources.

Regarding the reviewer’s concern about the claim that LPN represents a "program space," we agree that there are important distinctions between LPN's neural programs (latent vector plus neural decoder) and traditional code-based programs. Although both represent valid programs that map inputs to outputs, they differ notably in terms of compositionality, differentiability, and expressivity. Specifically, the space learned by LPN covers only a small subset of the broader space expressible through traditional code. While LPN sacrifices compositionality, it gains differentiability, enabling efficient test-time search, which is challenging for methods that generate code. We will explicitly clarify this trade-off in the final paper, emphasizing that LPN learns a narrow, implicit program space optimized for efficient search and strong performance within the training distribution.

Additional Ablation

We also draw the reviewers’ attention to an additional ablation conducted for reviewer qA5h. This evaluates different LPN test-time inference strategies on ARC-AGI, providing further insight into the roles of the program prior and gradient-based search in driving performance.

Dear reviewer,

Please read the rebuttal and discuss with the authors if you have any questions. Thanks!

AC