IMI AI/ML Fellowships, 2026: HIGH RISK RESEARCH - Advancing the Frontier.
2026 Cohort

Overview

This document describes the research topics we are opening for fellowships. Each topic suggests several projects indicative of the scope and complexity we expect you to undertake.

10
Places Open
6 or 12
Month Duration
Global
Fully Remote
Apr 1 & Sep 1
Deadlines (2026)

"High risk research" means ambitious projects that may not produce results, but will have lasting impact if they work.

We believe that publication in AI/ML has become overly biased towards short-term and iterative results, despite major gaps in our understanding of how to build optimal models.

Our focus is on frontier research: order-of-magnitude improvements, rather than disposable papers barely better than baseline.

The intent of these fellowships is to give awardees uninterrupted time to focus on harder open problems, with adequate compute, talented peers, and weekly 1:1 mentorship from senior researchers but without chasing specific metrics.

You will be expected to spend about 80% of your time on your own research, and up to 20% of your time either assisting other fellows or participating in wider research programs at IMI. We focus largely on applied research and its applications to online security problems, but often publish and support frontier research aligned with our broader interests. Serving hundreds of millions of people gives us a unique perspective as to what works at scale.

Eligibility

Previous fellows and research staff have come from disparate backgrounds, including early career researchers previously at MSR, FAIR, Mila, MPI, etc. and self taught senior engineers transitioning into research.

We do not discriminate on the basis of pedigree or age. If you have done interesting work, that is enough. Conversely, if you don't already have publications that show you can do good self-directed research then this is not the right opportunity.

Location

You may reside anywhere in the world, excluding sanctioned jurisdictions. This will be a remote fellowship.

Deliverables

We do not have hard targets, and focus on quality. Fellows generally get 1-2 papers with code done in a year, targeting NeurIPS, ICML, ICLR, etc. At the end of your fellowship, if the threshold for refereed publication is unmet we will expect a final report, which may be published as a blog post.

Deadline & Review

Admitting fellows in two cohorts. Deadlines for consideration: April 1 and September 1, 2026. 3 week decision period. Rolling thereafter.

Compensation

Competitive location-adjusted stipend, conference and travel support for conferences with accepted papers.

After the Fellowship

Past fellows have stayed on in full-time roles at IMI or co-founded AI unicorns. This is a selective process, so we claim little credit for their success. However, we will support your next steps.

Application Process

Applying & Selection

Ready to advance the frontier?

Send a brief bio / CV link via our official application portal.

Submit Application

What to Include in Your Submission:

  • 1 Topic Interest: The topic you are interested in working on.
  • 2 Prior Work: A few lines on any relevant prior work you've done.
  • 3 Links: Your GitHub / Google Scholar / X links.
  • 4 Availability: Desired start date, duration (6 or 12 months), and other obligations (if any) during that period.
  • 5 Project Analysis: A brief analysis of one of the projects outlined below.

Note on Project Analysis: Each project intentionally includes some gaps or glosses. List the ones you see, and how you'd solve them. Alternatively, if you dislike the projects outlined under a particular topic, write up your own idea and why it is more promising, along with your estimate of time and compute required.

Selection Criteria

Novelty and importance, clarity of approach, feasibility given time/compute, alignment with topics. Panel review and two interviews.

Selection is based on merit. We promote equality of opportunity, and welcome applications from anyone with talent, skills and potential. IMI is an equal opportunity employer, and does not discriminate on the basis of age, disability, sex, orientation, race, religion or belief.

Research Program

Research Topics and Suggested Projects

Note that these suggested projects are intentionally over-specified to give you a chance to show research taste. Make sure to mention in your analysis which parts you think are less likely to work, or should be pruned entirely.

Topic 01 Compute & Architecture

Faster learners

Thesis: Training power usage should be many orders of magnitude lower. No single change will get us there, but a plausible research program is to:

  1. Make most tokens unnecessary,
  2. Quickly get to a good solution analytically or with amortized learners,
  3. Learn to use external structure rather than training every fact into the weights, and
  4. Combine second order methods with low precision training to converge in fewer, larger steps.
Project 1.1

Faster training via approximate analytic solutions

Motivation

Training big networks via SGD is hideously inefficient. Can approximate analytic solutions provide a 100x reduction in training time?

Idea

Use curvature information from a tiny fraction of the corpus, solve the resulting quadratic problem once per block, then de-linearize by injecting the update through small gates so that the network stays in its local linear regime.

Suggested Approach

Estimate a good solution by solving a sketched natural-gradient step in closed form. Stream the corpus once to estimate a block-diagonal or Kronecker-factored Fisher (or Gauss-Newton) in each layer using random projections, compute the ridge solution for linearized parameters, de-linearize by composing layers near identity.

Why It Works

Layerwise K-FAC/Shampoo already works, randomized second-order solvers are mature, neural nets are close to linear early in training. A high-quality one-shot quasi-Newton init could remove 90-99% of steps.

Key Risks

Linearization error, imperfect Fisher, stability in deep stacks. (+ still needs a brief fine-tune at the end). LLMs also benefit from feature learning outside the lazy/NTK regime, where analytic linearized steps help least.

Possible Procedure

  1. 1
    Initialize: pre-LN or RMSNorm, gated residuals (alpha_l ~ 0.05-0.1), DeepNet/muParam scaling, light spectral constraints.
  2. 2
    Calibrate: stream 1-5% of tokens (more if necessary) to estimate K-FAC/Shampoo factors per block, collect residual RMS and logit scale stats. Dropout off.
  3. 3
    Solve: compute blockwise Gauss-Newton/natural-gradient delta W_l with damping and a small proximal term. Allocate per-block KL budgets, scale each delta W_l to fit.
  4. 4
    Compose delta theta: enforce a global KL cap via delta theta^T F delta theta <= epsilon, then apply alpha_l gates (start 0.05-0.2, larger for top blocks).
  5. 5
    Guard: run linearization-error checks per block, downscale any block that fails. Run spectral checks for MHA/MLP.
  6. 6
    Line search: Armijo on held-out loss plus KL cap, adjust alpha_l (not just a single alpha). Prefer to increase alpha_l for blocks passing both KL and linearization checks.
  7. 7
    Refresh: optionally relinearize top K blocks, refresh their factors, and take a second (smaller) trust-region step.
  8. 8
    Fine-tune: second-order optimizer (K-FAC/Shampoo grafted onto AdamW) with big batches for a few thousand steps. Re-enable dropout. Unfreeze MoE router here if we decide to tackle that. Maintain per-step KL regularization to keep the student initially near the trust region.
Project 1.2

Learn-to-database (explicit hierarchical memory)

Motivation

Storing facts in weights is wasteful and gets stale. Can we externalize knowledge as an explicit, updatable hierarchy the model learns to write to and query?

Idea

Make the model 1) induce its own ontology of concept nodes, typed edges, and table schemas, and 2) plan short programs over it. The LM becomes a planner/fuser, facts live in a compact, interpretable store.

Suggested Approach

Build a small, learned hierarchy: nodes with hyperbolic/tree codes and prototypes, sparse typed edges with evidence, leaves as text passages and table rows. Train write (create/link) and plan (DSL) heads, retrieve tiny typed subgraphs/rows, fuse with gated cross-attention, penalize param-only answers.

Why It Works

Transformers already form latent hierarchies. Making them explicit improves sample efficiency, updatability, and interpretability. RAG/FiD, Poincare embeddings, schema induction, and programmatic retrieval provide working parts.

Key Risks

Ontology drift/fragmentation, noisy links, planner brittleness, latency/complexity. Model may bypass the store. Some fast/robust knowledge must remain parametric.

Possible Procedure

  1. 1
    Initialize: pre-LN + gated residuals, add planner head, write head, and light fusion blocks. Storage: concept nodes with hyperbolic codes, relation vocab R, induced table schemas, ANN for nodes/leaves.
  2. 2
    Calibrate: cluster 1-5% of corpus to seed nodes/relations, label nodes by summarizing prototypes, estimate top-k recall/coverage, set per-source budgets and alpha_ext gates.
  3. 3
    Write: train create/link policy (straight-through Gumbel). Attach support spans and timestamps, regularize node/edge growth and degree, merge near-duplicates.
  4. 4
    Plan: teach a tiny DSL (select, hop, filter, join, aggregate) with weak traces (BM25/DPR/SQL/paths). Router picks graph/text/table. Penalize long programs.
  5. 5
    Retrieve/Fuse: execute plans, return compact subgraphs/rows/snippets. Fuse with cross-attn, gating evidence vs param prediction, require citations of nodes/edges used.
  6. 6
    Externalize: anti-memorization drops and KL penalties so factual answers rely on the store. Increase alpha_ext and budgets where evidence consistently helps.
  7. 7
    Compose: alternate phases: (A) train writer/planner/retriever with LM frozen, (B) train LM-on-evidence with store frozen, (C) brief joint pass.
  8. 8
    Guard: structural checks (type purity, degree/entropy caps, cycle detection), drift alarms, and timeouts with fallback to text-only retrieval.
  9. 9
    Line search: tune program depth/width and alpha_ext on held-out loss + latency + citation quality, promote sources that pass support/consistency checks.
  10. 10
    Refresh & Fine-tune: incremental writes and nightly re-embed/re-index, merge/retire nodes, maintain hot shards in memory, version edges for freshness. Preference training that rewards supported, cited answers, keep base LM small, adapt planner, writer, fusion. Maintain offloading penalties early.
Sub-project 1.2.1

Write-head sparsity curriculum (explicit, updatable hierarchies)

Motivation

Unconstrained writes bloat the store and kill interpretability. Can we grow a compact, legible hierarchy by starting ultra-sparse and relaxing only when utility is proven?

Idea

Treat create node/edge/schema as gated actions with an explicit cost. Start with near-zero write capacity, force reuse/links, relax L0/MDL penalties as evidence accumulates that new structure reduces loss.

Suggested Approach

Hard-concrete gates per write op (create/link/split), group-lasso over relation types, and an MDL-style budget. A staged schedule: 1) link-only, 2) controlled splitting, 3) schema induction, 4) relaxed growth. Each write gets a justification (support spans, proto summary), enabling audit and rollback.

Why It Works

Early noise is what makes ontologies sprawl. A sparsity curriculum preserves tree-like structure, encourages reuse, and keeps concepts interpretable while still allowing growth when it pays off.

Key Risks

Over-sparsity (missed concepts), late discovery tax (hard to recover if you never split), planner gaming the penalties, merge instability.

Possible Procedure

  1. 1
    Initialize: seed a tiny tree (hyperbolic codes) and relation vocab. Enable retrieval/fusion, disable creates (writes off). Set very high L0/MDL penalties and tiny per-batch budgets (0-1 creates, 1-5 links).
  2. 2
    Calibrate: measure coverage gaps, residual/error hotspots, node heterogeneity (prototype entropy), and average degree. Set a target growth curve (sublinear nodes vs tokens).
  3. 3
    Link-only: allow edges to existing nodes with top-k sparsity, dedup/merge aggressively. New edges must carry citations and type, enforce degree/branching caps.
  4. 4
    Controlled splitting: permit K new nodes per N tokens when triggers fire (high residual on a node, multi-modal embeddings, planner dead-ends). Child nodes inherit parent, get proto summaries, and must improve held-out loss to persist.
  5. 5
    Schema induction: when repeated attribute patterns appear, propose columns/typed relations with group-lasso over relation types. Writes to schema have higher cost, require multi-example support.
  6. 6
    Relaxation: slowly decay L0 and MDL penalties, increase budgets in domains where new structure consistently helps. Keep periodic prune sweeps (drop low-utility nodes/edges, auto-merge siblings).
  7. 7
    Guard: pre-commit checks (type purity, cycle limits, citation density), post-commit A/B (with/without the write) to confirm utility, rollback on regressions. Maintain bounded hyperbolic radius to keep tree-like geometry.
  8. 8
    Line search (budget control): adjust penalties to hit the target growth slope while minimizing held-out loss + write cost + latency. Prefer granting budget to sources/relations with highest utility per write.
  9. 9
    Refresh: regular scheduled merge/consolidate, re-embed nodes, retire stale edges, version writes with timestamps.
  10. 10
    Fine-tune: alternate freeze phases: 1) train planner/retriever under strict budgets, 2) briefly unfreeze write-heads to apply a small number of high-confidence writes, 3) joint stabilization pass. Keep anti-memorization so factual wins come from the store.

Signals it's working: Sublinear graph growth, rising citation rate, stable small branching factors, and large accuracy drops when the store is ablated in the domains it covers.

Project 1.3

Train on the residual (semantic LZ + template VM)

Motivation

Pretraining rereads the same patterns and reasoning chains. If models learn reusable templates, we should only pay for residual bits, not repetitions. Maybe direct + hierarchical compressibility can cut token exposure.

Idea

Turn pretraining into compression. Learn a hierarchical dictionary of semantic spans and reasoning templates. Encode the corpus into 1) template calls with slot fills and 2) residual tokens the dictionary can't predict. Train primarily on the residual, keep the dictionary executable so a few examples teach many instances.

Suggested Approach

1) A neural compressor that induces macros (semantic LZ) and templates (reasoning VM) with MDL pressure, and 2) a residual trainer that samples only novel bits with importance weighting. Refresh the dictionary online so the residual steadily shrinks.

Why It Works

Language is highly compressible, gradient contributions are heavy-tailed. Grammar induction, dataset distillation, and macro tokenization already show large sample cuts. If chain-of-thought compresses into a small set of schemas, we need far fewer demonstrations.

Key Risks

Over-compression (miss rarities), drift as the model/dictionary co-evolve, template brittleness, compressor overhead. Requires tight guards to avoid bias and loss spikes.

Possible Procedure

  1. 1
    Initialize: pre-LN/RMSNorm Transformer with gated residuals, add two heads: a utility u(x,t) predicting per-token residual bits, and a template planner that emits short programs over a tiny DSL (steps, slots, control).
  2. 2
    Calibrate: stream 1-2% of tokens. Fit u using loss/entropy/grad-sketch proxies. Seed a macro dictionary with frequent spans and a small set of reasoning schemas from mined CoT traces. No dropout.
  3. 3
    Solve (encode): train a compression model with an MDL objective:
    • Direct compressibility (semantic LZ): segment text, learn macros that maximize code-length reduction across paraphrases and formats.
    • Hierarchical compressibility (Template VM): induce typed templates (e.g. 1) compare 2) sort 3) argmax, 1) retrieve 2) filter 3) aggregate) with slots, compile to executable programs.
    • Choose between macro, template, or raw tokens per segment via straight-through gating, penalize long codes and deep programs.
  4. 4
    Compose (dataset): re-encode the corpus as codes:
    • Keep only (a) template calls + slot fills and (b) residual tokens above a u-threshold, drop predictable tokens.
    • Importance-weight residuals to match the original distribution, reserve a small shadow random batch each step to measure gradient mismatch.
  5. 5
    Guard: track gradient-match vs full sampling on shadow batches, bound NLL drift per domain, auto-bail to denser sampling if mismatch or rare-phenomena error rises. Enforce diversity caps so the dictionary doesn't collapse.
  6. 6
    Line search (budget): on held-out loss + code length, adapt thresholds: raise macro/template use until mismatch grows, lower when residual error spikes. Learn per-domain quotas and max program depth.
  7. 7
    Refresh: every K steps, re-minimize MDL on a fresh stream: merge/split macros, specialize/generalize templates, retire low-utility entries. Maintain a novelty buffer that bypasses compression for surprising spans until absorbed.
  8. 8
    Fine-tune: short full corpus sweep at low rate to debias, then large-batch training dominated by residuals. Keep off-policy checks with the shadow batch. At inference-time, optional on-the-fly template execution for chain-of-thought.

Expected outcome: Train on 3-10% of original tokens at steady state for compressible domains. Large wall-clock/power savings in early training, better factual and reasoning generalization per token via explicit reuse of templates, fast adaptation by updating the dictionary.

Topic 02 Reasoning & RL

Learning from failure

Thesis: People learn from our mistakes. Models don't, but pass@512 works better than @1, indicating the answer is often reachable already. Let's make max(pass@N) == pass@1:

  1. Turn search into supervision,
  2. Convert in-context fixes into weights,
  3. Compile failures into repairs, and
  4. Keep gains without regressions/forgetting.

Rough ideas: 1) bound updates with KL/Fisher caps, run canaries, use selective counterfactual replay, periodically distill validated patches into main. 2) align fail vs win traces to learn detectors and conditional rewrites, store them as routed patches and planner rules. 3) trigger on surprise from tool returns, crystallize traces into reusable templates and low-rank adapters with trust-region commits. 4) harvest pass@N trees with verifiers and partial checks, learn prefix value functions and distill winner prefixes into the policy.

Project 2.1

Prefix advantage distillation

Motivation

Many problems succeed with pass@64-512 but miss on pass@1 because early plan choices diverge. Can we shift policy to prefer winners without leaking finals or eval labels?

Idea

Distill the advantage of early prefixes from winning traces over near-misses. Learn to emit better plan tokens, decompositions, and invariants in the first few steps, under a KL trust region and without training on final answers.

Suggested Approach

Mine pass@N logs with a verifier. Extract prefixes up to the first tool call or K tokens. Run preference learning on winner vs loser prefixes plus step-wise advantage-weighted updates from partial checks. Update the planner head and small early-layer adapters, freeze late layers.

Why It Works

Early decisions carry most of the causal weight on success, pass@N pools expose reliable winner patterns, DPO/IPO and advantage-weighted regression work with preferences and partial rewards, avoiding ground-truth leakage.

Key Risks

Spurious correlations in prefixes, verbosity inflation, domain shift from synthetic near-misses, regressions if KL is loose.

Possible Procedure

  1. 1
    Initialize: pre-LN transformer with a planner head for plan tokens/templates, per-block low-rank adapters in early layers, verifier API and prefix extractor, global KL and latency budgets.
  2. 2
    Calibrate: collect pass@N logs on a training pool, label winners with the verifier, record prefixes, partial checks, and problem type features. Train a light prefix-success predictor to guide sampling and hard negative mining.
  3. 3
    Solve: build a preference set of (winner_prefix, loser_prefix, context). Train with DPO or IPO on prefixes. Add step-wise advantage-weighted regression using partial checks as rewards with a control variate. Apply per-step KL penalties and EWC on core capabilities, constrain updates to planner head and early adapters.
  4. 4
    Compose: at inference, lightly bias decoding toward learned plan tokens and templates (logit offsets, small temperature on planner head). Keep standard decoding otherwise. Cache priors per problem type.
  5. 5
    Guard (no-cheat constraints):
    • Train only on training pool logs, never eval logs.
    • Use verifier and partial checks that exist at inference (unit tests, proofs, execution), not hidden labels.
    • Restrict supervision to prefixes and plan tokens, not full final strings.
    • Enforce per-batch and cumulative KL caps, run regression canaries and length monitors.
  6. 6
    Line search: tune prefix length K, preference temperature, advantage weights, and KL caps on held-out pass@1 uplift vs pass@N, verbosity, and latency. Prefer shorter prefixes with large uplift.
  7. 7
    Refresh: periodically mine fresh pass@N logs, add hard negatives near the decision boundary, merge redundant templates via MDL, retire low-utility plan tokens.
  8. 8
    Fine-tune: brief consolidation pass to stabilize planner and adapters under tight KL. Verify gains persist with normal decoding and no reliance on best-of sampling.

Expected outcome: Pass@1 gains from the same pass@N budget, localized updates that do not overfit finals, minimal compute overhead beyond log mining and short preference training.

Project 2.2

Verifier-guided tree policy iteration

Motivation

Best-of-N search induces an implicit decision tree where early nodes determine success. Can we turn those trees into a value signal over plan tokens and perform policy iteration to improve first-shot behavior?

Idea

Build a prefix value model that predicts expected verifier success given a partial plan. Estimate Q over plan decisions from search trees and partial checks, then perform KL-regularized policy improvement on the planner head.

Suggested Approach

Construct shallow trees from pass@N rollouts, annotate nodes with verifier outcomes and partial rewards, compute soft Q via backward value propagation. Train a small critic over prefixes, improve the planner policy with advantage-weighted cross-entropy while freezing late layers.

Why It Works

Tree policy iteration is a principled way to learn from search, verifier and partial checks provide dense rewards, KL-regularized policy updates are stable and sample-efficient.

Key Risks

Value leakage if trees are overfit, instability from off-policy bias, verbosity creep, and interference with non-searched domains.

Possible Procedure

  1. 1
    Initialize: base model with planner head and early-layer adapters, add a prefix critic (small transformer or MLP on hidden prefix states), verifier API and partial-check probes, set KL and latency budgets.
  2. 2
    Calibrate: collect pass@N trees by saving top-k beams and sampled branches per problem. Annotate each node with partial checks passed and leaf success. Estimate off-policy correction weights for branches.
  3. 3
    Solve: compute soft Q at each node via backward propagation of verifier success with entropy regularization. Train the prefix critic to predict Q from hidden states and plan tokens. Improve the planner with advantage-weighted cross-entropy or PPO-style updates under KL to the current policy, apply EWC to protect core skills.
  4. 4
    Compose: at inference, combine planner logits with critic advantages for the first K steps (small additive bias). Keep decoding otherwise unchanged. Cache advantages by problem type for fast reuse.
  5. 5
    Guard (no-cheat constraints):
    • Construct trees only on training pools, never use eval trees.
    • Rely on verifiers and partial checks available at inference.
    • Bias only early plan decisions, do not regress on final answers.
    • Cap per-update and cumulative KL, enforce length and latency bounds, run regression canaries.
  6. 6
    Line search: tune K, critic regularization, advantage temperature, and KL caps on held-out pass@1 vs pass@N gap, stability, and verbosity. Prune nodes with low contribution to avoid overfitting.
  7. 7
    Refresh: rebuild trees on new snapshots, update the critic with fresh logs, remove stale branches, merge redundant plan tokens, maintain a buffer of hard problems for continual improvement.
  8. 8
    Fine-tune: periodic short run to co-train planner and critic with small KL, freeze late layers, verify gains on unseen tasks and that pass@N remains stable.

Expected outcome: First-shot accuracy approaches best-of-N without extra sampling at inference, stable, interpretable improvements via early plan value shaping, modest compute overhead with strong sample efficiency.

Topic 03 Neuro-Symbolic & Compilers

Making learned programs explicit

Thesis: Neural nets can implicitly learn small, reusable programs (sin/cos, sorting, date arithmetic, unit conversion) encoded as superpositions in their weights. We should extract these programs into explicit, verifiable modules and make models call them, reducing parameter bloat, improving generalization, and enabling updates without retraining. If this works well we'll next tackle composability, i.e. a 'neural linker' step.

Project 3.1

Neural-to-Program compilation (N2P)

Motivation

Storing algorithms in weights is opaque, hard to update, and redundancy-prone. If a net "knows" trig, a calendar, or a regex engine, we should externalize them as code and prune the corresponding circuits.

Idea

Build a compiler that maps network behaviors on targeted subspaces into a small typed intermediate representation of numerical and symbolic modules, with verification tests and equivalence checks. Replace the discovered circuit with a call to the extracted module.

Approach

  • Scope: Identify candidate latent programs by probing for low-entropy, low-rank, highly repeatable behaviors at specific layers/heads like 'angle to sin(angle)', 'string to match(pattern)'.
  • Specification mining: Generate test benches using counterfactual inputs, invariance checks (e.g. sin(x+2pi)=sin(x)), and smoothness/periodicity detectors, then fit candidates from a library (trig, polynomials, finite-state transducers, arithmetic, set ops).
  • Synthesis: Use hybrid methods—sparse identification of nonlinear dynamics (SINDy-style), symbolic regression, and enumerative search over a compact IR (typed SSA with vector ops, conditionals, and bounded loops).
  • Verification: Run equivalence tests against the original subnetwork over adversarial and randomized inputs, then certify error bounds and input domain. If verified, freeze/retire the circuit and insert a differentiable "call-module" stub with gradients routed to arguments, not the replaced weights.
  • Maintenance: Version modules, track provenance and input domains, and maintain a regression suite. Allow hot-swaps (e.g. better trig approximations) without touching base weights.

Neural circuits often realize simple, reusable functions. Explicit modules are easier to verify, cache, optimize, and upgrade. Synthesis with invariants curbs overfitting, then calls avoid recomputation and shrink parameter/activation footprints.

Target Goals
  • 30-60% of occurrences of common algorithmic skills offloaded to modules with certified error <= 1e-6 on their domains.
  • 10-30% parameter and 10-25% activation reduction at equal or better quality on tasks invoking those skills.
  • Measurable gains in out-of-distribution robustness for offloaded skills (e.g. long-range dates, big numbers).
Key Risks
  • Misspecification: Wrong library or IR misses real behavior, producing brittle extractions.
  • Distributed representations: Useful programs spread across layers/heads, hard to isolate.
  • Verification gaps: Passing tests but failing on rare regimes.
  • Integration tax: Latency/ABI overhead and gradient mismatch at the call boundary.
Project 3.2

Differentiable ABI and self-offloading training

Motivation

Even with a library, models won't use it unless calling is easy and rewarded. We need a clean ABI (types, shapes, side effects) and a training regime that prefers external calls over re-learning in weights.

Idea

Introduce a differentiable application binary interface (dABI) and "call" tokens that route subproblems to external modules. Train with an MDL-style objective: prefer short "call+args" programs over long parametric computation. Use write-through learning so successful in-context calls are consolidated into persistent call policies.

Approach

  • ABI design: Typed arguments (scalars, vectors, strings, sets), effect annotations (pure/impure), shape contracts, and gradient rules (exact, straight-through, or stop-grad).
  • Router: A lightweight planner head predicts when to call, which module, and with what arguments. Provide partial credit via differentiable surrogates (e.g. relaxed arg parsing, soft alignment of spans to args).
  • Costs and rewards: Penalize param-only solutions when a verified module exists, then reward correct calls with small KL bonuses and latency/energy credits, then enforce per-batch "offload budgets" to shift usage gradually.
  • Pruning: After stable adoption, gradually L0-prune circuits shadowed by calls, then keep a safety adapter to catch drift and trigger re-training.
  • Continual learning: When no module fits, log traces and auto-propose new candidates for Project 1 to compile. Close the loop: discovered modules immediately become callable via the dABI.

Models already plan tool usage, then making library calls first-class and cheap creates selection pressure to externalize. A typed ABI contains complexity, reduces integration bugs, and allows acceleration (e.g. vectorized trig kernels).

Target Goals
  • >= 80% correct module call rate on benchmarks containing extractable skills.
  • Net wall-clock speedups of 1.3-2.0x on workloads heavy in offloaded operations.
  • Stable or improved task accuracy with >= 70% reduction in gradients flowing through replaced circuits.
Key Risks
  • Over-calling: Router overuses modules where param paths suffice.
  • Cold-start: Early errors discourage calls, then needs careful curriculum and safety nets.
  • Gradient pathologies: Approximate gradients through discrete calls can bias training.
  • Library sprawl: Too many niche modules increase complexity and latency.
Topic 04 Interpretability & Alignment

Self-documenting weights / live stack traces for reasoning

Thesis: Today's 'explanations' are mostly rationalizations written after the fact. Let's bake interpretability into the forward pass: layers and heads get compact, typed docstrings and citations that (a) summarize what each component is doing on this input, (b) constrain what information is allowed to flow next ('explain to execute'), and (c) compose into a real-time execution path/stack trace. Target two domains initially (multi-hop QA over provided context, grade-school math word problems) to demonstrate faithful, causal traces that show the facts and logic used at each step.

Project 4.1

HeadDocs (static capabilities + dynamic call-site docstrings and citations)

Motivation

Attention heads and MLPs often specialize, but their behavior is hidden. We want two levels of documentation: (1) a static capability docstring per head/block, and (2) a dynamic "call-site docstring" on each token that declares what the unit is doing now and which facts it uses.

Idea

Add a small documentation head and a citation head to each attention head and MLP block. The documentation head emits a short program sketch from a tiny vocabulary (the codebook), and the citation head names the spans/tokens used as evidence. Train them to be (i) predictive of the unit's effect, (ii) minimal via MDL pressure, and (iii) causally faithful by enforcing that downstream computation depends on the cited evidence and declared operation.

Suggested Approach

  • Representations:
    • Static capability vector z_h per head/block with a canonical docstring (learned once, versioned).
    • Dynamic call-site docstring d_h,t per token/time: a short sequence in a constrained DSL (e.g. COPY_FROM(span), MATCH(pattern), COREF_ANTE(span), AGG(NUM, window), DATE_ADD, ARGMAX(key), ROUTE_TO(node), ASSERT(invariant)).
    • Citations c_h,t: a sparse set of input spans (start, end, source_id) with confidence.
  • Losses and Constraints:
    • Reconstruction: small decoders reconstruct the head's output from (d_h,t, c_h,t, selected source embeddings), the docstring must be sufficient to predict the unit's effect.
    • Causal faithfulness: ablate or perturb cited spans, require the predicted effect to change accordingly (contrastive consistency). Enforce sparsity on citations and docstring length (MDL).
    • Alignment: cluster behaviors offline and map clusters to codebook tokens, with human-curated seeds for a few canonical operations (coref, copy, local-n-gram, delimiter detection, number aggregation).
  • Runtime: A trace aggregator collects (d_h,t, c_h,t, z_h) across layers and compiles a readable stack trace with timestamps and evidence links.

Why: 'Documentation' used to predict and constrain the forward computation is far harder to fake than post-hoc gloss. Minimal codebook tokens with explicit citations make trace outputs short, comparable, and auditable.

Target Goals
  • Multi-hop QA: Most correct answers accompanied by traces whose citations, when masked, cause the answer to fail (strong causal test).
  • GSM8K Math: Most correct solutions accompanied by stepwise traces where masking cited numbers/ops flips the outcome.
  • Overhead: Low latency on a 7B-class model, static per-head docs are stable across corpora with >= 0.8 Jaccard overlap of codebook tokens.
Key Risks
  • Doc collapses to vague tags, codebook drift.
  • Overhead in both compute and latency if not carefully architected.
  • Faithfulness gaps if constraints are too weak or too soft.

Possible Procedure

  1. 1
    Initialize: Choose base model (e.g. 7B pre-LN Transformer), instrument attention heads/MLPs with documentation head (linear to small LM over fixed codebook of 64-128 tokens), citation head (pointer network over input tokens/spans, k=1-3 pointers), and lightweight reconstructor. Define codebook with 12-20 primitive tokens (COPY, COREF_ANTE, LOCAL_NGRAM, OPEN_QUOTE, CLOSE_QUOTE, MATCH_DIGITS, SUM_NUMS, MAX_BY_KEY, TABLE_HEADER_ALIGN, DATE_PARSE, DATE_ADD, FORMAT). Leave "OTHER_x" slots for emergent clusters.
  2. 2
    Calibrate: Collect behavior sketches: run base model on QA and math corpora, log attention patterns and interventions (activation patching, heads ablation). Cluster head behaviors to propose initial head-to-codebook mappings, hand-label only 10-20 examples to seed.
  3. 3
    Train (Phase A: doc-only): Freeze base LM. Train documentation and citation heads and reconstructors: cross-entropy loss on doc tokens with entropy regularization, length penalty (hard cap 3-5 tokens), sparse pointer loss for citations (top-k), MSE/cosine reconstruction loss, and causal imitation on counterfactual subsets.
  4. 4
    Guard (doc quality): Disallow trivial docs: if doc length > 1 but adds sub-epsilon reconstruction gain vs "OTHER", increase MDL penalty. Merge near-duplicate codebook entries, bound per-head doc entropy over time.
  5. 5
    Train (Phase B: explain-to-execute coupling): Unfreeze small gates around each unit: when a docstring declares COPY_FROM(span), soft-mask attention to non-cited spans; when AGG(NUM), encourage number-feature heads to turn on. Performance loss KL-regularized to base, plus penalty for off-trace token usage.
  6. 6
    Runtime trace: Implement a trace aggregator producing per-token JSON trace (step_id, layer, head, doc_tokens, citations, confidence, before/after norms). UI: render collapsible stack trace with clickable spans showing source context.
  7. 7
    Evaluate: Fidelity via "mask-the-citation" causal tests (drop cited spans, rerun, record delta; report per-domain causal F1). Measure compactness (average doc length, citation count) and stability (doc overlap on new corpora, drift alarms).
  8. 8
    Refresh: Periodically recluster behaviors, reassign "OTHER" tokens to concrete roles, prune unused codebook entries. Small preference-tuning round to favor shorter, more faithful docs without hurting task metrics.

Literature & Weaknesses Addressed

  • Chain-of-thought: Boosts performance, but often unfaithful and verbose. We avoid free-form text and enforce causal use via masking.
  • Attention != explanation: We add reconstruction plus interventional tests to ensure attention/citations matter causally.
  • Self-Explaining Neural Networks / rationale extraction: Many produce proxies not used by the model. We couple docs to execution.
  • Sparse autoencoders on residual streams: Promising for features but not tied to live traces. We integrate with citations and reconstruction.
Project 4.2

Trace-Execute (minimal reasoning DSL that gates computation and yields a live stack trace)

Motivation

Even with HeadDocs, end-to-end reasoning is unclear. We want a compact 'reasoning atoms' DSL that the model emits interleaved with generation. These atoms both (a) constrain what the next step can use and (b) act as the executable plan, producing a faithful, verifiable stack trace.

Idea

Add a parallel trace channel that emits short sequences of typed atoms with arguments (spans, numbers, patterns). Enforce "explain to execute": the atom determines which submodules/heads are allowed to contribute in the next step and which evidence is admissible. Keep the DSL small and domain-scoped to avoid bloat and ensure strong fidelity.

Suggested Approach

  • DSL (first pass, 8-12 atoms):
    • QA: FETCH(span_id), HOP(via_anchor), COMPARE(span_a, span_b, key), AGG(list, op), SELECT(condition), ASSERT(supports(answer)), CITE(span_id).
    • Math: PARSE_NUM(span), APPLY(op, args), KEEP(track_id), CHECK(invariant), FORMAT(result).
  • Execution Coupling:
    • When an atom fires, apply compile-time masks: limit attention to nominated spans, route to dedicated submodules (e.g. arithmetic/date kernels), and activate heads matching the atom (HeadDocstoDSL mapping).
    • At each step, trace channel proposes 0-2 atoms with confidences. If confidence < tau, default to unconstrained decoding but flag "no-trace" for transparency.
  • Supervision: Weak signals: distant supervision from retrieval citations, gold answer spans, program-of-thought datasets (math), and tool returns. No reliance on ground-truth explanations at scale. Consistency: require emitted atoms and HeadDocs to agree.
  • Guardrails: Budget per step: at most one FETCH and one APPLY per token budget. Abstain when inputs are ambiguous, do not fabricate arguments.

Why: A small, typed set of atoms is enough to make traces legible and enforce faithfulness. Gating makes cheating costly—if the atom says FETCH(span), only that span is accessible, so the model can't secretly use other facts.

Goals (3-6 Months)
  • QA: >= 30% reduction in unsupported answers vs baseline; >= 85% of correct answers have at least one FETCH/CITE whose removal flips the outcome.
  • Math: >= 70% of correct solutions accompanied by APPLY/ASSERT traces passing automated step checks; <= 10% verbosity overhead.
  • Live Trace Viewer: Interactive "stack" for both domains; qualitative user studies indicate > 0.8 usefulness ratings for debugging.
Key Risks
  • Coverage gaps: DSL may miss useful operations.
  • Over-constraint: Masks too tight can hurt accuracy early on.
  • Atom spam: Model may emit atoms unnecessarily; need budget and MDL penalties.

Possible Procedure

  1. 1
    Initialize: Add a parallel trace head (small decoder) conditioned on main model's hidden states, share early layers. Define DSL schema and static mapping from HeadDocs codebook tokens to eligible atoms. Implement attention masks for FETCH/HOP and router for APPLY.
  2. 2
    Calibrate: Build small training pools (Hotpot-like QA with support spans; GSM8K math with synthetic step checks). Measure baseline unsupported answer rate and spontaneous citations.
  3. 3
    Warm-start (supervised): Train trace head to imitate weak labels (support spans, tool steps) where available, apply MDL length penalty and abstain option. Enforce atom-to-mask in a fraction of steps, keep base LM frozen.
  4. 4
    Couple (explain-to-execute): Gradually increase fraction of steps where masks are enforced, using a KL trust region. Penalize off-trace usage when gradients/attention mass flow outside allowed spans.
  5. 5
    Consistency with HeadDocs: Jointly train with Project 1 so active heads' dynamic docs match atom class; add small cross-entropy alignment loss. Break ties by preferring HeadDocs if atom confidence is low.
  6. 6
    Guard: Atom budget per K tokens; high cost for atoms that don't change downstream behavior (measured by ablation). Drift alarms if unsupported answers creep up.
  7. 7
    Evaluate: Fidelity via mask-the-citation tests at step level. Measure fraction of runs where removing fetched spans or preventing APPLY flips final answer. Aim for <= 15% latency overhead.
  8. 8
    Refresh & Distill: Add 2-3 atoms based on observed gaps (e.g. MATCH_DATE, JOIN_TABLE) only if clearing MDL/utility threshold. Occasionally fine-tune to reduce mask reliance while preserving stable traces.
Sub-project 4.2.1 (Optional)

Anti-rationalization curriculum

Motivation

Models can learn to emit plausible but non-causal traces. Force faithfulness by training under randomized or hidden-information regimes.

Idea

Randomized ablation and counterfactuals during training. Replace would-be cited spans with paraphrases or foils and require the trace (and answer) to update. Inject "trace traps" where the only way to succeed is to follow the declared atoms.

Possible Procedure

  1. 1
    Generate foils: Paraphrase or swap entities in candidate spans, mark them.
  2. 2
    Train: In 20-30% of batches, replace a cited span with a foil, require the model to either (a) change its answer accordingly or (b) abstain/emit uncertainty, penalize traces that remain unchanged.
  3. 3
    Evaluate: Measure "trace flip rate" under counterfactuals, aim for >= 0.7 on curated sets.
Signs It's Working
  • High causal flip rates when cited spans are ablated.
  • Short, stable per-head docs that generalize across datasets.
  • Trace channel coverage rising over time without atom spam.
  • Human auditors can follow the stack trace to verify answers quickly.
How This Differs From Prior Work
  • Chain-of-thought & rationales: Can be unfaithful and verbose. Our atoms/docstrings are minimal, typed, and coupled to execution via masks.
  • Attention-as-explanation: Is weak; we combine attention with reconstruction and interventional tests.
  • Sparse autoencoders: Produce features but not live, causal traces; we add citations and enforce use.
  • Program-of-thoughts: Can overfit to template programs; our DSL is tiny, with abstention, and is only partially enforced to avoid brittleness.

Integration Points & Non-Bloat Scoping

  • Scope strictly to two domains (multi-hop QA over provided context, GSM8K-like math).
  • Keep codebook <= 128 tokens, DSL <= 12 atoms, k <= 3 citations per step.
  • Enforce masks in a gradually increasing fraction of steps, avoid full hard enforcement early.
  • Defer ambitious expansions (full formal proofs, rich ontologies) to later phases, ship a working live trace with causal guarantees first.