| # Reason-First Program |
|
|
| **Concept-Guided Program Space Exploration** |
|
|
| When an LLM fills a stub function, there exists a *possible program space* of valid implementations. Any implementation that satisfies the spec, the formal constraints, and the observable effects is "aligned" — its exact shape shouldn't matter. What *does* matter is the set of **concepts** that characterize distinct regions of that space. |
|
|
| This framework discovers those concepts, projects programs into a concept-guided embedding space, and provides a query language for steering generation. |
|
|
| ## Architecture |
|
|
| ``` |
| ┌─────────────────────────────────────────────────────────────────┐ |
| │ @reason_first(spec="...") │ |
| │ def process(items): ... │ |
| │ ┌──────────┐ │ |
| │ │ STUB │ │ |
| │ └────┬─────┘ │ |
| │ │ │ |
| │ ┌───────────────┼───────────────┐ │ |
| │ ▼ ▼ ▼ │ |
| │ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ |
| │ │ Model A │ │ Model B │ │ Model C │ Stage 1 │ |
| │ │ T=0.2..1.2 │ │ T=0.2..1.2 │ │ T=0.2..1.2 │ Sampling │ |
| │ └──────┬───────┘ └──────┬───────┘ └──────┬───────┘ │ |
| │ └────────────────┼────────────────┘ │ |
| │ ▼ │ |
| │ ┌─────────────────────────┐ │ |
| │ │ PROGRAM SPACE │ │ |
| │ │ (valid implementations)│ │ |
| │ │ DA@K diversity metrics │ │ |
| │ └────────────┬────────────┘ │ |
| │ │ │ |
| │ ┌───────────────┼───────────────┐ │ |
| │ ▼ ▼ ▼ │ |
| │ ┌─────────────┐ ┌──────────────┐ ┌─────────────┐ │ |
| │ │ Behavioral │ │ SAE-Based │ │ Abstraction │ Stage 2 │ |
| │ │ (execution) │ │(hidden state)│ │ (AST frag.) │ Discovery │ |
| │ └──────┬──────┘ └──────┬───────┘ └──────┬──────┘ │ |
| │ └───────────────┼────────────────┘ │ |
| │ ▼ │ |
| │ ┌──────────────────────────┐ │ |
| │ │ CONCEPT SET │ │ |
| │ │ {recursive, mutation, │ │ |
| │ │ hash_based, sorting, │ │ |
| │ │ fast_execution, ...} │ │ |
| │ └────────────┬─────────────┘ │ |
| │ │ │ |
| │ ┌────────────┼────────────┐ │ |
| │ ▼ ▼ ▼ │ |
| │ ┌──────────────┐ ┌─────────┐ ┌──────────┐ │ |
| │ │ CB-AE │ │ GCAV │ │ MSRS │ Stage 3 │ |
| │ │ Bottleneck │ │ Vectors │ │ Orthog. │ Embedding │ |
| │ └──────┬───────┘ └────┬────┘ └────┬─────┘ │ |
| │ └──────────────┼───────────┘ │ |
| │ ▼ │ |
| │ ┌──────────────────────────┐ │ |
| │ │ CONCEPT EMBEDDING SPACE │ │ |
| │ │ (interpretable axes, │ │ |
| │ │ verifiable alignment) │ │ |
| │ └────────────┬─────────────┘ │ |
| │ │ │ |
| │ ▼ │ |
| │ ┌──────────────────────────┐ │ |
| │ │ QUERY LANGUAGE │ Stage 4 │ |
| │ │ recursive=0.8, │ Steering │ |
| │ │ mutation=-0.5, │ │ |
| │ │ fast_execution=0.7 │ │ |
| │ └──────────────────────────┘ │ |
| └─────────────────────────────────────────────────────────────────┘ |
| ``` |
|
|
| ## Quick Start |
|
|
| ```python |
| from reason_first_program import ( |
| ProgramSpace, Program, UnifiedConceptDiscovery, SteeringEngine |
| ) |
| from reason_first_program.stub import Stub, StubConstraints |
| from reason_first_program.program_space import execute_program |
| |
| # 1. Define a stub |
| stub = Stub( |
| name="two_sum", |
| source='def two_sum(nums: list[int], target: int) -> list[int]:\n ...', |
| signature="(nums: list[int], target: int) -> list[int]", |
| constraints=StubConstraints(decorator_spec="Find two indices that sum to target"), |
| test_inputs=[{"nums": [2, 7, 11, 15], "target": 9}], |
| ) |
| |
| # 2. Build program space (add implementations) |
| space = ProgramSpace(stub) |
| for src in implementations: |
| p = Program(source=src, full_source=src, stub_id=stub.stub_id) |
| p = execute_program(p, stub, stub.test_inputs) |
| space.add(p) |
| |
| # 3. Discover concepts |
| discovery = UnifiedConceptDiscovery() |
| concepts = discovery.discover(space) |
| |
| # 4. Query & steer |
| engine = SteeringEngine(concepts) |
| results = engine.select(space, "uses_dict=1.0, uses_mutation=-1.0", top_k=3) |
| ``` |
|
|
| ## Using the @reason_first Decorator |
| |
| ```python |
| from reason_first_program import reason_first |
|
|
| @reason_first( |
| spec="Sort items by priority, breaking ties by recency", |
| postconditions=["output is sorted", "all input items present in output"], |
| ) |
| def process_queue(items: list) -> list: |
| #> stable sort; O(n log n); must preserve Item identity |
| ... |
| ``` |
| |
| ## Concept Discovery Methods |
|
|
| ### 1. Behavioral Concepts (AST + Execution) |
| Discovers concepts from code structure and execution traces: |
| - **Algorithmic patterns**: `uses_recursion`, `uses_iteration`, `uses_list_comprehension` |
| - **Data structure usage**: `uses_dict`, `uses_set`, `uses_heap` |
| - **Control flow**: `uses_early_return`, `single_return`, `uses_generator` |
| - **Mutation pattern**: `uses_mutation` (in-place) vs immutable |
| - **Performance**: `fast_execution` (below-median execution time) |
|
|
| ### 2. SAE-Based Concepts (Hidden States) |
| Trains a Sparse Autoencoder on code LLM hidden states (requires GPU): |
| - Based on [CB-SAE](https://arxiv.org/abs/2512.10805) and [DN-CBM](https://arxiv.org/abs/2407.14499) |
| - Auto-names neurons against a code concept vocabulary |
| - Discovers implicit representational concepts the LLM uses internally |
|
|
| ### 3. Structural Abstraction Concepts (AST Fragments) |
| Finds common program fragments across implementations: |
| - Based on [LILO](https://arxiv.org/abs/2310.19791) / [ReGAL](https://arxiv.org/abs/2401.16467) |
| - Patterns like `for_range`, `for_enumerate`, `while_loop`, `nested_function` |
|
|
| ## Diversity Metrics |
|
|
| ```python |
| report = space.diversity_report() |
| # { |
| # "total_programs": 10, |
| # "valid_programs": 10, |
| # "functionally_unique_clusters": 3, |
| # "da_at_5": 2.8, # Expected distinct algorithms in 5 samples |
| # "entropy_diversity": 2.1, # Entropy-based diversity index |
| # } |
| ``` |
|
|
| From [AlgoDiv (2503.00691)](https://arxiv.org/abs/2503.00691): `DA@K = Σ_m (1 - C(N-s_m, K) / C(N, K))` |
|
|
| ## Embedding & Steering |
|
|
| ### GCAV Concept Vectors |
| Based on [GCAV (2501.05764)](https://arxiv.org/abs/2501.05764): `e' = e + ε · v_concept` |
|
|
| ```python |
| from reason_first_program.embeddings import GCAVEmbedding |
| |
| gcav = GCAVEmbedding() |
| gcav.train_all(concepts, features, programs) |
| steered = gcav.multi_steer(features, { |
| "uses_recursion": 0.8, "fast_execution": 0.6, "uses_mutation": -0.3 |
| }) |
| ``` |
|
|
| ### MSRS Orthogonal Steering |
| Based on [MSRS (2508.10599)](https://arxiv.org/abs/2508.10599): orthogonal subspaces prevent concept interference. |
|
|
| ## Query Language |
|
|
| ```python |
| engine = SteeringEngine(concepts) |
| |
| # Weight-based queries |
| results = engine.select(space, "uses_recursion=0.8, uses_mutation=-0.5") |
| |
| # Constraint queries |
| results = engine.select(space, "uses_dict > 0.5 AND fast_execution > 0.7") |
| |
| # Concept negation |
| results = engine.select(space, "NOT uses_mutation") |
| |
| # Programmatic |
| query = engine.query_language.build(uses_recursion=0.8, fast_execution=0.6) |
| results = engine.select(space, query) |
| ``` |
|
|
| ## Concept Lattice (FCA) |
|
|
| Builds a [Formal Concept Analysis](https://en.wikipedia.org/wiki/Formal_concept_analysis) lattice: |
| - **Extent**: set of programs sharing certain properties |
| - **Intent**: set of concepts shared by those programs |
|
|
| ```python |
| lattice = concepts.concept_lattice() |
| for extent, intent in lattice[:5]: |
| print(f"Concepts {sorted(intent)} → {len(extent)} programs") |
| ``` |
|
|
| ## Theoretical Foundations |
|
|
| | Method | Paper | Key Contribution | |
| |--------|-------|-----------------| |
| | Program space diversity | [AlgoDiv](https://arxiv.org/abs/2503.00691) | DA@K metric, AlgoSim clustering | |
| | SFS scattering | [SFS](https://arxiv.org/abs/2411.05010) | Diverse algorithmic direction discovery | |
| | SAE concept discovery | [CB-SAE](https://arxiv.org/abs/2512.10805) | Concept Bottleneck Sparse Autoencoders | |
| | Concept naming | [DN-CBM](https://arxiv.org/abs/2407.14499) | Discover-then-Name concept pipeline | |
| | Concept vectors | [GCAV](https://arxiv.org/abs/2501.05764) | Concept Activation Vectors for steering | |
| | Multi-concept steering | [MSRS](https://arxiv.org/abs/2508.10599) | Orthogonal subspace steering | |
| | Execution semantics | [TRACED](https://arxiv.org/abs/2306.07487) | Execution-trace-aware representations | |
| | Library learning | [DreamCoder](https://arxiv.org/abs/2006.08381) / [LILO](https://arxiv.org/abs/2310.19791) | Abstraction discovery from program corpora | |
| | Latent program space | [LPN](https://arxiv.org/abs/2411.08706) | Continuous latent program representations | |
| | Typed holes | [Hazel](https://arxiv.org/abs/2409.00921) | Spec-constrained completion spaces | |
|
|
| ## Running the Experiment |
|
|
| ```bash |
| pip install -e . |
| python -m reason_first_program.experiment |
| ``` |
|
|
| Demonstrates the full pipeline on hand-crafted `two_sum` and `flatten` implementations: |
| 1. Program space construction & diversity metrics |
| 2. Behavioral + structural concept discovery (31 concepts for two_sum, 30 for flatten) |
| 3. Concept-guided embedding & alignment verification |
| 4. Query language selection & concept boundary exploration |
| 5. Formal concept lattice construction |
| |
| ## Installation |
| |
| ```bash |
| # Core (concept discovery + steering, no GPU needed) |
| pip install -e . |
| |
| # With LLM sampling support |
| pip install -e ".[llm]" |
| |
| # With SAE concept discovery (requires GPU) |
| pip install -e ".[sae]" |
| |
| # Everything |
| pip install -e ".[all]" |
| ``` |
| |
| ## License |
| |
| Apache 2.0 |
| |
