Austin Shen
B.S. Computer Science, Mathematics & Data Science • University of Michigan • GPA 4.00 / 4.00
I work on neural theorem proving in Lean 4 and on small transformers that learn to compute exactly.
Formal reasoning meets Machine Learning
I'm Austin Shen, a Computer Science, Mathematics, and Data Science triple major at the University of Michigan (GPA 4.00 / 4.00, expected May 2027). I am a James B. Angell Scholar and hold University Honors.
I have worked with Prof. Yunong Shi at Michigan since July 2025, first on LeanQC (autoformalization into Lean 4) and, since March 2026, on a proof-state snapshot engine for the Lean 4 language server and a pre-registered audit of whether execution throughput limits neural theorem provers. Since June 2026 I have worked with Tianyu Hua and Prof. Nick Haber at Stanford on tiny transformers that compute exactly, verified bit-for-bit against oracles.
Before that I interned at BotSmart, where I built an evaluation platform for LLM prompt attacks, and at Innovation AI, where I built the backend for a pre-launch AI writing product.
Tools & Technologies
Languages
Formal Methods
ML / Systems
Web / Cloud
Research Projects
Exact Algorithmic Generalization in Tiny Transformers ("The Neural Engine")
I study when small transformers learn to compute exactly. I designed and built PaperDiT, a 0.2 to 0.6M-parameter cells-as-tokens 2-D transformer (no positional encodings, fixed-point halting) that computes bit-for-bit against executable oracles: 1,000-digit addition from 10-digit-max training, a universal cellular-automaton interpreter running unseen rules at 50× tape width, and damage-robust interactive physics.
- Self-play: bootstrapped arithmetic from zero data with a propose, solve, oracle-verify, retrain loop. 100% on all lengths 1 to 22 (eval cap) from 1-to-10-digit training, 3/3 seeds. The oracle supplies the verification signal.
- Mechanistic interpretability: linear probes show the carry bit is computed, not memorized. Causal activation patching localizes the carry circuit (about 100% output flip at layers 0 to 1, 0% past layer 2).
- Transformer plus exact external memory: a 10-digit-trained GPT with a pointer/register harness reaches 100% exact addition on 483/483 problems through 3,000 digits and at 5,000 and 7,000 digits (700× training length), 3/3 seeds.
- Four live in-browser demos on a hand-written, dependency-free JavaScript transformer engine with verified fp16 export.
Execution Audit & Snapshot Engine for Neural Theorem Proving (Lean 4)
I co-authored a 24-page pre-registered measurement study, the first controlled test of whether execution throughput limits neural theorem provers. Every claimed proof was independently re-verified by whole-file Lean compilation (633/636).
- Patched the Lean 4 language server with custom JSON-RPC snapshot/branch methods for millisecond in-memory proof-state forking: 5.6 to 16.1× (mean 9.1×) end-to-end speedup, and distributed proof search across 4 cluster nodes (about 15.7×).
- Ran 2-policy × 2-backend × 244-problem benchmark grids with 7B provers via vLLM on SLURM A40 GPUs. Proof sets were statistically identical across backends.
- Found a premise-selection inversion: a reranker that improves Recall@32 by 8.3 points loses 23.7 pp of end-to-end proof rate (p=0.022).
LeanQC: Autoformalization & Formal Theorem Proving
Developing a neuro-symbolic framework to formalize informal mathematical statements into Lean 4 code. The DSP-Plus (Draft, Sketch, Prove) architecture uses LeanTree with structured generation and an iterative compiler-feedback pipeline. Ablation studies conducted on MiniF2F and ProofNet (Pass@k metric).
- DSP-Plus architecture: Draft → Sketch → Prove with LeanTree
- Structured generation + iterative compiler feedback
- Benchmarked on MiniF2F and ProofNet (Pass@k)
Preprints & Papers
Keep the Proof State Live: Snapshotting for Efficient Tactic Search in Lean 4
Austin Shen, Yunong Shi
Parallel tactic search in Lean 4 rebuilds each proof state by re-running elaboration, which accounts for over 99% of per-branch wall time. The paper introduces proof-state snapshotting: capture the elaborated proof state once and reuse it across search branches through a small extension to the Lean 4 language server. Across 48 miniF2F-v2 problems the paper reports a 5.6 to 50× wall-time speedup over the standard fallback (average 14×, median 9.7×).
Industry Experience
Software Engineer Intern
BotSmart
Built LLMAE, an automated evaluation platform for LLM prompt attacks (modular Python, Django, Docker). A 5-level configuration space with at least 5 options per level yields 3,000+ attack variants, tested across content-safety categories. The platform was adopted for continued development after the internship.
Full Stack Engineer Intern
Innovation AI
Built the complete backend for a pre-launch AI writing assistant: content storage and retrieval, LLM text refinement, Firestore cross-platform sync, and a Stripe subscription pipeline (webhook verification, customer metadata mapping) tested end-to-end in Stripe's test environment.
Have a project in mind?
I'm open to research collaborations, internship opportunities, and interesting engineering problems.