File size: 12,929 Bytes

c311f46

# 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