ryanyen22/reason-first-program

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

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

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 and DN-CBM
• 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 / ReGAL
• Patterns like for_range, for_enumerate, while_loop, nested_function

Diversity Metrics

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): DA@K = Σ_m (1 - C(N-s_m, K) / C(N, K))

Embedding & Steering

GCAV Concept Vectors

Based on GCAV (2501.05764): e' = e + ε · v_concept

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): orthogonal subspaces prevent concept interference.

Query Language

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 lattice:

• Extent: set of programs sharing certain properties
• Intent: set of concepts shared by those programs

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	DA@K metric, AlgoSim clustering
SFS scattering	SFS	Diverse algorithmic direction discovery
SAE concept discovery	CB-SAE	Concept Bottleneck Sparse Autoencoders
Concept naming	DN-CBM	Discover-then-Name concept pipeline
Concept vectors	GCAV	Concept Activation Vectors for steering
Multi-concept steering	MSRS	Orthogonal subspace steering
Execution semantics	TRACED	Execution-trace-aware representations
Library learning	DreamCoder / LILO	Abstraction discovery from program corpora
Latent program space	LPN	Continuous latent program representations
Typed holes	Hazel	Spec-constrained completion spaces

Running the Experiment

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

# 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