TRM on SAT Solving

This branch adapts TinyRecursiveModels (TRM) for Boolean Satisfiability (SAT) problems.

Quick Start

1. Preprocess Dataset

uv run dataset/build_sat_dataset.py \
    --input_dir <path_to_sat_data> \
    --output_dir data/sat-uf50-218 \
    --train_ratio 0.8 \
    --seed 42

2. Run Single Experiment

uv run pretrain.py --config-name cfg_sat_pretrain

3. Run All Experiments (H_cycles × L_cycles sweep)

CHECKPOINT_BASE=/path/to/checkpoints ./experiments/sat_cycles_sweep/run_all.sh

Dry run to preview commands:

CHECKPOINT_BASE=/path/to/checkpoints ./experiments/sat_cycles_sweep/run_all.sh --dry-run

4. Analyze Results

python experiments/sat_cycles_sweep/analyze_results.py \
    --checkpoint-base /path/to/checkpoints \
    --output-dir experiments/sat_cycles_sweep/results

File Changes

New Files

File	Description
config/cfg_sat_pretrain.yaml	Base SAT training config
config/arch/sat_trm.yaml	SAT model architecture
config/sat_h*_l*.yaml	Experiment configs (9 files for H×L sweep)
dataset/build_sat_dataset.py	Preprocesses DIMACS CNF to numpy arrays
evaluators/sat.py	SAT-specific metrics (token accuracy, satisfiability rate)
experiments/sat_cycles_sweep/run_all.sh	Runs all experiments

Modified Files

File	Changes
pretrain.py	Updated imports for SAT dataset
puzzle_dataset.py	Added SAT data loading support

Experiment Matrix

The sweep tests H_cycles = {2, 3, 4} × L_cycles = {4, 6, 8} (9 experiments total).

Baseline: H_cycles=3, L_cycles=6

Dataset

• Source: Uniform Random-3-SAT
• 1000 instances, 50 variables, 218 clauses each (3-SAT)

Metrics

• sat/token_accuracy - % correct variable assignments
• sat/satisfiability_rate - % solutions that satisfy all clauses