ARC-AGI Program Synthesis System

The system uses a hybrid approach combining LLM-driven program generation with a growing library of successful programs to solve ARC-AGI tasks.

Core Workflow:

1. Initialization: Starts with an empty library of Python programs.
2. Prompting: For each task, the LLM is prompted to generate Python programs that solve the training examples. The prompt includes:
   • Input/output grids in multiple formats (dimensions, Base64, ASCII, and Python nested lists).
   • Chain-of-Thought (CoT) reasoning instructions.
   • The current best program(s) from the existing library.
3. Scoring: Every program in the library is evaluated using two metrics:
   • Primary Accuracy Score: Number of correct training examples.
   • Secondary Accuracy Score: Average cell-level accuracy on training examples.
4. Library Expansion: The LLM-generated program with the best primary and secondary scores is added to the library.
5. Iteration: The system loops through tasks, expanding the library, for a user-defined number of rounds.

This approach aims to bridge the gap between the scalability of LLMs (generating diverse programs) and the knowledge transfer of systems like DreamCoder (using a library of programs).

Instead of random search, you can optimize latents directly using gradient ascent. Define a loss function based on lpn._compute_log_probs (the log probability of the ground truth output given the context). Use jax.value_and_grad to compute gradients with respect to the latents and update them using an optimizer like optax.adam.

def log_probs_fn(latents, input_seq, output_seq):
    # Repeat latents for all pairs in the task
    latents = latents[..., None, :].repeat(output_seq.shape[-2], axis=-2)
    row_logits, col_logits, grid_logits = lpn.decoder.apply(
        {"params": loaded_state.params["decoder"]},
        input_seq, output_seq, latents, dropout_eval=True
    )
    return lpn.apply({"params": loaded_state.params}, row_logits, col_logits, grid_logits, output_seq, method=lpn._compute_log_probs)

# Optimization loop
log_probs, grads = jax.value_and_grad(log_probs_fn)(search_latents, input_seq, output_seq)
search_latents = search_latents + lr * grads

To find optimal latent contexts that maximize the log-probability of a specific task, you can perform gradient ascent on the log_prob_other_pairs function. This is useful for exploring the model's learned task manifold.

Workflow:

1. Define a function that computes log_prob for a given context using the decoder.
2. Use jax.value_and_grad to get the gradient of the log-probability with respect to the context.
3. Apply an optimizer (like optax.adam) to update the current_context in the direction of the gradient.
4. Map the resulting contexts back to CDF space (using norm.cdf) for visualization.

# 1. Define log_prob function
@jax.jit
def log_prob_other_pairs(context):
    input_seq, output_seq = lpn._flatten_input_output_for_decoding(pairs[1:], grid_shapes[1:])
    row_logits, col_logits, grid_logits = lpn.decoder.apply(
        {"params": loaded_state.params["decoder"]},
        input_seq=input_seq,
        output_seq=output_seq,
        context=context[..., None, :].repeat(output_seq.shape[-2], axis=-2),
        dropout_eval=True,
    )
    return lpn._compute_log_probs(row_logits, col_logits, grid_logits, output_seq)

# 2. Gradient Ascent Loop
optimizer = optax.adam(learning_rate=0.01)
current_context = start_context
opt_state = optimizer.init(current_context)

for _ in range(26):
    log_prob, grad = jax.value_and_grad(log_prob_other_pairs)(current_context)
    updates, opt_state = optimizer.update(-grad, opt_state, current_context)
    current_context = optax.apply_updates(current_context, updates)

To generate predictions for the ARC-AGI challenge, use the Evaluator class.

1. Initialize Evaluator with your LPN model instance.
2. Configure inference_mode (e.g., "gradient_ascent") and provide inference_mode_kwargs to control the optimization process during inference.
3. Load the challenge JSON file.
4. Call evaluator.json_submission(challenges, loaded_state.params) to produce the generations.
5. Save the resulting dictionary to a JSON file.

from src.evaluator import Evaluator
import json

# Configure Evaluator
evaluator = Evaluator(
    lpn,
    inference_mode="gradient_ascent",
    inference_mode_kwargs={
        "num_steps": 200,
        "lr": 1.0,
        "lr_schedule": True,
        "optimizer": "adam",
        "optimizer_kwargs": {"b2": 0.9},
        "accumulate_gradients_decoder_pairs": True,
    },
)

# Run submission
with open("/path/to/arc-agi_test_challenges.json", "r") as f:
    challenges = json.load(f)

generations = evaluator.json_submission(challenges, loaded_state.params, progress_bar=True)

with open("submission.json", "w") as f:
    json.dump(generations, f, indent=4)

To use a pre-trained Latent Program Network (LPN), you can download artifacts from Weights & Biases (WandB). The process involves:

1. Initializing a WandB run.
2. Downloading the model artifact.
3. Extracting the configuration from the artifact metadata.
4. Instantiating the EncoderTransformer, DecoderTransformer, and the LPN model.
5. Loading the trained parameters from the state.msgpack file into a Flax TrainState.

Note: This requires omegaconf, hydra, jax, flax, and wandb to be installed.

import os
import jax
import jax.numpy as jnp
import wandb
import omegaconf
import hydra
from flax.training.train_state import TrainState
from flax.serialization import from_bytes
from lpn.src.models.lpn import LPN
from lpn.src.models.transformer import EncoderTransformer, DecoderTransformer

# 1. Download artifact
artifact_path = "TheThinker/ARC/quiet-thunder-789--checkpoint:v0"
run = wandb.init()
artifact = run.use_artifact(artifact_path, type='model')
run.finish()
cfg = omegaconf.OmegaConf.create(artifact.metadata)
artifact_dir = artifact.download()

# 2. Instantiate model components
encoder = EncoderTransformer(hydra.utils.instantiate(cfg.encoder_transformer))
decoder = DecoderTransformer(hydra.utils.instantiate(cfg.decoder_transformer))
lpn = LPN(encoder=encoder, decoder=decoder)

# 3. Initialize state for loading (requires dummy inputs matching model config)
key = jax.random.PRNGKey(0)
grids = jax.random.randint(key, (1, 3, decoder.config.max_rows, decoder.config.max_cols, 2), minval=0, maxval=decoder.config.vocab_size)
shapes = jax.random.randint(key, (1, 3, 2, 2), minval=1, maxval=min(decoder.config.max_rows, decoder.config.max_cols) + 1)
variables = lpn.init(key, grids, shapes, dropout_eval=False, prior_kl_coeff=0.0, pairwise_kl_coeff=0.0, mode="mean")

# 4. Setup optimizer and TrainState
optimizer = optax.adamw(0.0)
train_state = TrainState.create(apply_fn=lpn.apply, tx=optimizer, params=variables["params"])

# 5. Load weights
with open(os.path.join(artifact_dir, "state.msgpack"), "rb") as data_file:
    byte_data = data_file.read()
loaded_state = from_bytes(train_state, byte_data)