Full 2D Attention for ARC-AGI

A vision transformer with native 2D spatial awareness for the Abstraction and Reasoning Corpus.

Quick Start (Step by Step)

Step 1: Install Python dependencies

pip install torch numpy tqdm matplotlib

Step 2: Run Checkpoint 1 (verify code works)

cd arc2d
python run_tests.py

What to expect: All 8 modules should show ✓ PASS. Time: Under 2 minutes, no GPU needed.

Step 3: Run Checkpoint 2 (verify model can learn)

python train_single_task.py

What to expect: Loss decreases from ~2.4 to near 0. Accuracy > 90%. Time: 1-5 minutes on CPU.

Step 4: Run Checkpoint 3 (verify bias tables learn)

PYTHONPATH=. python -m visualization.bias_tables

What to expect: Heatmap images saved in bias_table_plots/ showing spatial patterns. Time: Under 1 minute.

Step 5: Download ARC data

git clone https://github.com/fchollet/ARC-AGI.git

Step 6: Pretrain on ARC tasks

PYTHONPATH=. python training/pretrain.py --data_dir ARC-AGI/data/training --size small

Or use synthetic tasks for a quick test without ARC data:

PYTHONPATH=. python training/pretrain.py --synthetic --epochs 20

Time: 2-4 hours on GPU (Small config).

Step 7: Evaluate on ARC tasks

PYTHONPATH=. python evaluation/evaluate.py \
    --model_path checkpoints/best_model_small.pt \
    --eval_dir ARC-AGI/data/evaluation \
    --ttt_steps 100 \
    --n_views 10

Time: ~30 seconds per task.

Project Structure

arc2d/
├── config.py                    # Model configurations (Small/Medium/Large)
├── run_tests.py                 # Checkpoint 1: run all self-tests
├── train_single_task.py         # Checkpoint 2: single task overfitting
├── requirements.txt             # Python dependencies
├── README.md                    # This file
│
├── model/                       # The neural network
│   ├── attention.py             # ★ Full 2D Attention (core innovation)
│   ├── ffn.py                   # ConvGLU feed-forward network
│   ├── block.py                 # Transformer block (attention + FFN)
│   └── model.py                 # Complete model
│
├── data/                        # Data loading and processing
│   ├── canvas.py                # Canvas placement utilities
│   ├── augmentations.py         # Geometric + color augmentations
│   └── dataset.py               # ARC JSON file loader
│
├── training/                    # Training scripts
│   └── pretrain.py              # Stage 1: offline pretraining
│
├── evaluation/                  # Evaluation scripts
│   └── evaluate.py              # TTT + multi-view inference
│
└── visualization/               # Visualization tools
    └── bias_tables.py           # Checkpoint 3: visualize bias tables

Architecture Summary

Input: ARC grid (integers 0-9) placed on a fixed-size canvas

Three components:

1. 2D Attention — per-head learnable bias tables indexed by (Δrow, Δcol)
2. ConvGLU — depth-wise 3×3 convolution on the gate path
3. Output head — per-pixel classification into 11 color classes

No positional embeddings. Position enters exclusively through the 2D bias tables.

Training: Pretrain on 400 ARC tasks + augmentation, then test-time training per task.