When the Next Step Is Not One Step: Distribution-Aware Execution Modeling for Concurrent Go Programs

This paper introduces a distribution-aware execution modeling approach for concurrent Go programs that leverages scheduler nondeterminism to fine-tune a 7B model on empirical event distributions, achieving state-of-the-art accuracy and improved calibration on real-world bug predictions while also providing formal guarantees for detecting specific goroutine leaks.

The Gist

Imagine you are trying to teach a robot to predict the next move in a game of chess. If the game is standard chess, the rules are fixed: if you move a pawn here, the opponent must respond in a specific way. The robot just needs to memorize the pattern.

But now, imagine a game of chess played in a chaotic, noisy room where three different people are all trying to move the pieces at the same time, and a random wind blows the board around. Sometimes, if you move a pawn, the wind might knock it over. Sometimes, a player might grab a piece before you can. Sometimes, the opponent might decide to move a different piece entirely.

This is the problem with computer programs that run multiple tasks at once (concurrent programs).

The paper you provided tackles this exact chaos. Here is the breakdown in simple terms:

The Problem: "One Answer" vs. "Many Possible Answers"

In traditional computer science, when a program runs, we usually assume it follows a straight line. If you give it the same input, it gives the same output.

• The Old Way: Researchers trained AI models to predict the one next step a program would take. They treated the program like a straight line.
• The Reality: In concurrent programs (like those written in the Go language), the "scheduler" (the part of the computer that decides which task runs when) is like a chaotic referee. If you run the same program twice, it might do A then B, or B then A. Both are correct. Both are valid.

If you train an AI to guess just one answer for a situation where there are actually three valid answers, the AI gets confused. It's like asking a weather forecaster to predict "It will rain" when the reality is "It might rain, it might snow, or it might be sunny," and the AI just picks one and hopes for the best.

The Solution: Predicting the "Weather Forecast"

The authors realized they shouldn't treat the chaos as a mistake. Instead, they treated the chaos as data.

1. Run it many times: They took a program and ran it hundreds of times.
2. Count the outcomes: They noticed that while the order changed, some patterns emerged. For example, "Event A" happened 60% of the time, "Event B" happened 30%, and "Event C" happened 10%.
3. Teach the AI the distribution: Instead of teaching the AI to guess "Event A," they taught it to guess the whole forecast: "There is a 60% chance of A, 30% of B, and 10% of C."

They used a special math trick (called a "KL objective") to train a 7-billion-parameter AI model to match these real-world percentages rather than just guessing a single winner.

The Results: Did it Work?

They tested this on real-world, messy code from famous systems like Kubernetes and Google's gRPC.

• The AI vs. The Experts: The fine-tuned AI (trained on less than 1,000 examples) got the next step right 36.2% of the time.
• The Competition: This beat a very powerful, pre-trained AI (Gemini 3.5 Flash) that hadn't been trained on this specific type of problem at all (which only got 34.8% right).
• The "Calibration" Win: Even more importantly, the new AI was better at knowing when it was unsure. If the situation was chaotic, the AI said, "I'm not sure, it could be anything." If the situation was predictable, it said, "I'm pretty sure." The old way of training made the AI confidently wrong more often.

The Limits: Where the Ceiling Is

The paper is very honest about what the AI can't do yet:

• The Accuracy Ceiling: The AI tops out around 35–36% accuracy. It can't get much higher because some events are so rare (like a specific type of glitch) that the AI never sees them enough to learn them.
• The "One Step" Problem: The AI is great at predicting the very next move. But if you ask it to predict the next 10 moves in a row, it falls apart after about one step. It's like a person who can tell you what happens in the next second of a movie, but if you ask them to predict the whole plot, they start making things up.

The "Leak" Discovery

The authors also found a specific "signature" for a type of computer bug called a "goroutine leak" (where a task gets stuck and never finishes).

• They proved mathematically that if a task gets stuck in a specific type of waiting loop, the chance of it ever "waking up" is zero.
• This isn't something the AI learned by guessing; it's a rule of the universe (the Go programming language's rules). The AI correctly learned that "Wake Up" is impossible in this specific scenario, which is a good sign that it's understanding the logic, not just memorizing numbers.

Summary

The paper says: "Stop trying to force a chaotic, multi-path system into a single straight line. Instead, show the AI the whole map of possibilities. It won't be perfect, and it can't predict long chains of events yet, but it becomes much better at understanding the nature of the chaos and knowing when it's guessing."

They released their code, data, and tools so others can try to build on this "chaos-aware" approach.

Technical Summary

Technical Summary: Distribution-Aware Execution Modeling for Concurrent Go Programs

Problem Statement
Training language models to predict the next step in a concurrent program is fundamentally different from sequential code due to scheduler nondeterminism. In sequential execution, a specific state and statement yield a single deterministic next state. In concurrent Go programs, however, the same execution prefix can legitimately lead to multiple different next events (e.g., a block, a start, or an unblock) depending on how the scheduler interleaves goroutines.

Standard Code World Models (CWMs) trained with cross-entropy loss assume a single correct label per input. When applied to concurrent traces, this forces the model to memorize one arbitrary outcome of a random process, failing to capture the underlying dynamics. This limitation is critical because concurrency bugs (deadlocks, data races, goroutine leaks) are inherently schedule-dependent and lack sequential analogues. Existing benchmarks and models often fail to address this nondeterminism, treating it as noise rather than a structural property of the execution.

Methodology
The authors propose reframing next-event prediction as distribution estimation rather than point prediction. Instead of training a model to guess a single next event, the approach aggregates observed next events from multiple runs of the same program to form an empirical distribution.

• Data Collection: The study utilizes 130 concurrent Go programs, including hand-crafted patterns, synthesized code, and 66 real-world concurrency bug kernels from production systems (CockroachDB, Kubernetes, gRPC, etcd). Each program is executed five times under a runtime tracer to capture scheduler events.
• Target Definition: For a given program prefix, the target is not a single label but an empirical distribution $\hat{p}_g$ over six event types: GoBlock, GoCreate, GoEnd, GoSched, GoStart, and GoUnblock. This distribution is derived from the frequency of observed next events across repeated runs.
• Model Training: The authors fine-tune a 7B parameter model (Qwen2.5-Coder) using two objectives:
  1. Cross-Entropy (CE): A baseline where examples are duplicated proportionally to event frequency to approximate the distribution.
  2. KL Divergence (KL): A distribution-aware objective that minimizes the Kullback-Leibler divergence between the model's predicted distribution and the empirical target. The KL term is applied specifically at the token position discriminating event types.
• Evaluation: The model is evaluated on a held-out set of 66 real-world production bugs (GoKer) it never saw during training. Metrics include top-1 accuracy, Expected Calibration Error (ECE), and multi-step coherence (simulating execution by feeding predictions back as input).

Key Contributions

1. Distribution-Aware Training: The paper demonstrates that treating scheduler nondeterminism as a signal allows a model to learn the "valid futures" permitted by the Go scheduler. A 7B model fine-tuned on fewer than 1,000 traces achieves 36.2% accuracy on held-out production bugs, outperforming the same model zero-shot (28.6%) and Gemini 3.5 Flash zero-shot (34.8%).
2. Calibration Improvement: While distribution training (KL loss) achieves accuracy comparable to cross-entropy (35.8% vs. 36.2%), it significantly improves model calibration. The Expected Calibration Error drops from 0.205 to 0.169, and model entropy correlates with program nondeterminism, allowing the model to express uncertainty where appropriate.
3. Formal Leak Signature: The authors derive a formal signature for a specific class of goroutine leaks: select-blocked goroutines. They show that for these cases, the probability of a GoUnblock event is exactly zero ($P(\text{GoUnblock}) = 0$) at all trace depths. This is a consequence of Go scheduler semantics (if no reachable goroutine can satisfy the select condition, unblocking is impossible) rather than a learned statistical pattern.
4. Identification of Limitations: The study explicitly maps the boundaries of the approach. Accuracy plateaus near 35–36% regardless of model size or objective. The model fails to predict rare event types (GoEnd, GoSched) due to class imbalance and distribution shift. Furthermore, multi-step predictions lose scheduler coherence after approximately one step, as the model was trained only on single-step transitions.

Results and Analysis

• Accuracy Ceiling: The performance ceiling (~36%) is attributed to two factors: class imbalance (rare events like GoEnd and GoSched are never predicted correctly) and distribution shift (real-world bugs exercise event types like GoCreate and GoSched more frequently than the hand-crafted training set).
• Event-Specific Performance: The model learns common lifecycle events well (GoStart at 47%, GoCreate at 44%) but fails on rare events. Notably, GoCreate achieves 44% accuracy despite being only 1.5% of the training data, suggesting the model reasons from program structure rather than just frequency.
• Multi-Step Coherence: When used autoregressively, the model produces valid transitions for roughly one step before violating Go invariants. This confirms that single-step training does not transfer to multi-step simulation without trajectory training.
• Reasoning: Enabling "thinking" modes in the baseline model (Gemini) did not improve accuracy, indicating the task relies on structural signals rather than multi-step deductive reasoning.

Significance and Claims
The authors explicitly state they are not claiming a deployable bug detector or a production-ready execution simulator. The significance of the work lies in:

• Formulation: Providing a principled framework for modeling concurrent execution as distribution estimation.
• Dataset and Baselines: Releasing a dataset of concurrent traces, trained adapters, and tooling to establish baselines for this specific problem.
• Clarification of Limits: Demonstrating that while nondeterminism-aware training works and improves calibration, significant hurdles remain (rare events, distribution shift, multi-step coherence) that must be addressed for the approach to be more effective.

The paper concludes that the current approach proves the viability of distribution-aware training for concurrent code but highlights that future work must focus on trajectory training, richer state representations (including channel buffers and mutex ownership), and rebalancing training data to match real-world event distributions.