NeuroGenesis
Atlas AI Labs
A self-optimizing compiler that encodes IR with a Transformer and learns optimization rules with RL.
- Rewrite accuracy
- 99.37%
- IR before → after, Sprint 2
- Edit-distance proxy
- 0.0063
- Sprint 2
- Research sprints
- 6
01
Problem
- Who
- Compiler engineers and ML-for-systems researchers.
- Why it mattered
- Optimisation passes are hand-written by experts; NeuroGenesis asks whether a model can discover, formalise and choose rewrite rules itself from program IR.
- Before
- Fixed, hand-tuned pass pipelines.
02
Constraints
Technical
Prototype on synthetic SSA-like IR (64-token programs) with a proxy cost model before moving to real IR.
Infrastructure
Every sprint runs end to end on CPU in a notebook.
03
Architecture
- Component
- Key decision (click)
Represent
- Synthetic program IRADD, MUL, LOAD, STORE, LOOP, IF…
- Structural featuresOp frequency, dependency-graph degree, control flow
Learn
- Gated attentionLearned residual mixing
Discover
- Pattern miningk-gram sequences + subgraph edges
- Rule formalisationPattern → rule, confidence, support, condition
Optimise
- Self-rewriting optimiserEvolving vs static pipeline (Sprint 6)
Programs are encoded as IR tokens plus structural features; a small transformer learns the before → after rewrite; mining and formalisation turn what it learns into explicit rules with confidence and support; an RL policy then picks which rule to apply from the execution-cost reward.
04
Key decisions
D1Inject compiler locality into attention
- Alternatives considered
- Plain content-only attention
- Why
- Peephole optimisations depend on neighbouring instructions, so a near-diagonal bias matches how rewrites work.
D2Reward = relative drop in execution cost
- Alternatives considered
- Supervised labels only
- Why
- Lets the system choose rules by their effect, not by imitation.
- Trade-off
- On random programs the reward is sparse (most programs give 0).
05
Build highlights
- 01
Sprint 0–1: baseline and features
A 2-layer transformer learns the synthetic rewrite (loss 2.23 → 0.043); fusing op-frequency, dependency-graph and control-flow features takes validation loss to 0.012.
- 02
Sprint 2: optimisation learning engine
With locality-biased, gated attention the model reproduces the optimised IR at 99.37% token accuracy (edit-distance proxy 0.0063).
- 03
Sprint 3–4: from patterns to rules
Sequence and subgraph mining over 200 programs, then rules formalised as pattern → rewrite with confidence, support and an applicability condition.
- 04
Sprint 5: execution feedback (RL)
A REINFORCE policy chooses between candidate rules (MUL·CONST → SHIFT, ADD·CONST → NOP) using a cost model where MUL = 2 and DIV = 3.
06
Results
99.37%
Token accuracy
Sprint 2, IR before → after
0.012
Validation loss
Sprint 1, token + feature fusion
0.0069
Average RL reward
Sprint 5, 100 random programs
How this was measured: Notebook runs on synthetic 64-token IR programs on CPU, with a proxy execution-cost model; not yet measured on real compiler IR.
07
Lessons
What worked
- A locality bias in attention made the before → after rewrite easy to learn.
- Fusing structural graph features with tokens cut validation loss sharply.
What didn't
- Frequency-based mining surfaced common unchanged patterns ahead of the hidden rewrite rules.
- On random programs the RL reward was sparse, so the policy learned slowly.
What I'd do differently
- Train and evaluate on real IR (e.g. LLVM) with measured runtime, and weight mined rules by whether they change the program.
08
