Pathways: Asynchronous Distributed Dataflow for ML

🎯 Pitch

A single-controller system can match the performance of multi-controller SPMD frameworks like JAX on thousands of TPUs—a feat previously thought impossible—by using a novel asynchronous dispatch mechanism that hides host-side scheduling latency. Pathways achieves this while natively supporting heterogeneous, pipelined, and sharded computations that resist expression in the ubiquitous SPMD model, enabling 97% throughput retention when splitting a 136B-parameter Transformer across two data center network-connected TPU islands.

1. Executive Summary

This paper introduces PATHWAYS, a new large-scale orchestration layer for accelerators designed to enable exploration of novel ML research ideas while matching state-of-the-art performance on current models. Evaluated on TPU clusters scaling to 2048 cores with both JAX and TensorFlow models—including Transformer architectures up to 136B parameters—PATHWAYS employs a sharded dataflow graph of asynchronous operators (a single-controller coordination model where each compiled function becomes a computation node in a dataflow DAG, using futures for cross-host communication) combined with parallel asynchronous dispatch (gang-scheduling heterogeneous SPMD sub-computations such that host-side work for multiple nodes runs concurrently rather than serially, exploiting statically known resource requirements of compiled functions). The system achieves performance parity with multi-controller JAX on SPMD computations (reaching identical token throughput—e.g., 618k tokens/s for T5-Base, 84.8k tokens/s for T5-11B—and matching throughput for computations as small as 2.3 ms on 128 TPUs and 35 ms on 2048 TPUs) while delivering comparable throughput for models pipelined across 16 stages or sharded across two islands of accelerators connected over a data center network (131.4k tokens/s for a 3B Transformer, and ~97% of single-island throughput for a 136B model split across two islands). The work establishes that a single-controller architecture can match multi-controller SPMD performance only when realistic computation sizes mask dispatch overhead and parallel asynchronous dispatch amortizes scheduling latency across pipeline stages.

2. Context and Motivation

The Core Problem: ML Systems Are Over-Specialized to SPMD

The fundamental problem PATHWAYS addresses is that the dominant programming model for large-scale ML—SPMD (Single Program Multiple Data)—is both too restrictive for emerging model architectures and poorly aligned with how hardware resources should be managed in a shared cluster.

In the SPMD model, inspired by MPI (Clarke et al., 1994), every accelerator runs identical code in lockstep, and communication between accelerators happens exclusively through collective operations like AllReduce. The paper observes that this approach has been enormously successful: most state-of-the-art ML workloads today use SPMD, and systems like JAX (Bradbury et al., 2018) and PyTorch (Paszke et al., 2019) have demonstrated excellent performance within this paradigm. However, the authors identify a deeper structural problem that becomes visible when we look at where ML workloads are heading rather than where they currently are.

The tension is between two forces. On one side, ML models are becoming increasingly heterogeneous and sparse, moving away from the uniform, dense computation patterns that SPMD handles naturally. On the other side, ML hardware is becoming increasingly heterogeneous and fragmented, with clusters composed of multiple smaller islands of accelerators rather than one giant homogeneous pool. The paper argues that SPMD—despite its proven performance—creates a straitjacket for both trends.

Three Specific Gaps the Paper Identifies

Gap 1: Emerging model architectures don't fit the SPMD mold. The paper points to several concrete developments that strain SPMD:

Pipelining across model layers: Very large language models have been scaled using pipeline parallelism (Narayanan et al., 2019; Rasley et al., 2020; Narayanan et al., 2021), where different stages of the model run on different accelerators. This is fundamentally an MPMD (Multiple Program Multiple Data) pattern—each stage executes a different computation—that researchers have had to awkwardly shoehorn into SPMD frameworks using "ingenious techniques."
Mixture of Experts (MoE) (Shazeer et al., 2017): These models route different inputs to different subsets of model weights based on learned routing functions. The result is computational sparsity—not all accelerators are computing on all data—and heterogeneous computation—different accelerators run different expert sub-networks. As the paper notes, this is "most naturally expressed using fine-grain control flow and heterogeneous computation across accelerators."
Data-dependent control flow at the sub-example level: The paper's authors mention that their ML research colleagues "would like to use sparsity more effectively when training ever larger models, with ever more tasks, but that current frameworks limit their ability to experiment with novel model architectures" (Section 6.3). The vision is models where "different model weights can be updated per example, or even per sub-example (patch of an image, or word of a sentence)"—a level of granularity that SPMD's uniform computation model fundamentally cannot express cleanly.

Gap 2: Hardware heterogeneity makes exclusive island ownership wasteful. Each new generation of accelerators introduces variety—different memory capacities, different interconnect topologies, different computational characteristics. Providing exclusive access to large islands of homogeneous accelerators connected over high-bandwidth interconnects is expensive and, the paper argues, "often wasteful as a single user program must try to keep all of the accelerators continuously busy" (Section 1). When a user can only get access to smaller, more readily available islands of accelerators, the natural solution is to map sub-parts of the overall computation to different islands—again, an MPMD pattern that SPMD systems struggle to support.

Gap 3: Multi-tenancy and resource sharing are second-class concerns in SPMD systems. The multi-controller architecture where each host runs its own copy of the user program and takes exclusive ownership of hardware resources "shifts the responsibility of ensuring high utilization of the expensive accelerators on to the user" (Section 2). It also "complicates the design of features like resource virtualization and multiplexing that are needed to build efficient cluster-wide ML infrastructure." When one user's program leaves accelerators idle (due to, say, host-side data preprocessing or I/O), those cycles are stranded—no other user can access them. The paper envisions a world where foundation models (Bommasani and et. al., 2021) are shared across many downstream tasks, with different users concurrently fine-tuning different heads while sharing the same frozen base layers on the same accelerators. This requires centralized resource management that can time-multiplex accelerators at fine granularity, something the multi-controller model fundamentally resists.

Collectively, these gaps define what Dean et al. anticipated as the "next-generation AI architecture" (Dean, 2021): a system that supports sparsity, heterogeneity, resource sharing, and computation patterns beyond SPMD, while retaining the performance of today's best SPMD systems on today's workloads.

Why This Problem Matters: The Co-Evolution Trap

The problem is not merely academic—it reflects what the paper calls the "co-evolution" of ML models, accelerator hardware, and software systems. The danger is that systems become over-specialized to current workloads and fail to anticipate future needs (Section 1). This is a classic systems trap: a design optimized for yesterday's dominant use case creates path dependency that inhibits the exploration of tomorrow's ideas.

The practical stakes are high:

Research velocity: If experimenting with sparsely-activated models or novel routing architectures requires wrestling with SPMD abstractions, fewer researchers will attempt it. The paper explicitly frames PATHWAYS as a platform to enable research, not just deploy it—the system is "explicitly designed to enable exploration of new systems and ML research ideas" (Abstract).
Resource efficiency at scale: Foundation model training and fine-tuning at Google scale involves clusters of thousands of accelerators. Stranded cycles from poor bin-packing of SPMD jobs onto available islands, or from exclusive resource ownership models that prevent sharing, represent enormous financial and carbon costs. Fine-grained multi-tenancy that overlaps computations from different users on the same accelerators could dramatically improve utilization.
The infrastructure-model gap: Hardware continues to diversify (different TPU generations, different interconnects, different memory tiers), while models continue to diversify (sparse, pipelined, multi-modal). A system that forces a uniform computation model onto this diversity will increasingly be the bottleneck.

The Two-Category Landscape of Prior Approaches

The paper situates existing distributed ML systems into two broad architectural categories, and argues that each has a critical weakness.

Category 1: Multi-Controller SPMD Systems (JAX, PyTorch, MPI-style)

These systems run an identical copy of the user's executable directly on every host, with each copy taking exclusive ownership of that host's accelerators. Communication between hosts happens exclusively through collective operations (AllReduce, etc.) that use dedicated high-speed interconnects like NVLink or ICI without going through host memory.

Where they excel: This architecture achieves very low dispatch latency because each host's copy of the program enqueues computations directly over fast PCIe to its local accelerators—no network round-trips are needed. The paper's Figure 1a illustrates this: a sequence of lockstep steps across hosts involves only local enqueue operations and collective communication over the accelerator interconnect. This is why multi-controller systems dominate state-of-the-art performance on SPMD workloads (Mattson et al., 2020).

Where they fall short: The paper identifies several structural limitations:

No primitives for non-SPMD communication: "Any communication beyond standard collectives in multi-controller systems requires users to implement their own coordination primitives" (Section 2). If you want to dynamically route data between accelerators based on computed values (as in MoE), or pipeline different computations across different sets of accelerators, you're on your own.
Exclusive resource ownership is baked in: The assumption that a user program owns the entire island for its duration makes it "practically infeasible" to build centralized resource management, multi-tenancy, or elasticity—features that "are needed to build efficient cluster-wide ML infrastructure" (Section 2).
The model is fundamentally one-size-fits-all: Every accelerator runs the same program. Expressing MPMD requires contorting the SPMD abstraction, typically through conditionals that selectively execute different code paths on different accelerators—functional but unnatural and limiting for truly heterogeneous computations.

Category 2: Single-Controller Dataflow Systems (TensorFlow v1)

TensorFlow v1 (Abadi et al., 2016) took the opposite approach: a single Python client builds a computation graph, a coordinator runtime partitions it into subgraphs for each worker, and workers exchange data and control messages over the datacenter network (DCN). This offers a "very general distributed dataflow model, including optimized in-graph control flow" (Section 2).

Where it should excel in theory: The single-controller model naturally supports non-SPMD patterns—heterogeneous subgraphs on different workers, data-dependent routing between subgraphs, and centralized resource management—because the coordinator has a global view of the entire computation. This is exactly the flexibility that multi-controller systems lack.

Where it falls short in practice: The paper is blunt about TF v1's failure to deliver on this promise at scale, identifying three implementation-level problems that collectively killed the approach:

Dispatch latency kills performance: In multi-controller systems, dispatching a computation to an accelerator involves only PCIe communication (Figure 1a). In TF v1's single-controller model, the client is "farther away"—dispatch requires messages over the DCN, "typically an order of magnitude slower than PCIe" (Figure 1b). This means that for any computation that runs in less time than a DCN round-trip, the accelerator sits idle waiting for dispatch messages. The paper shows this concretely: in TF v1, a sequence of step k computations across hosts requires a DCN message to trigger each host's work, whereas in JAX the same sequence is triggered by local PCIe enqueues with no DCN messages at all.
Cross-host coordination accumulates in chains: When programs involve many cross-host transfers—for example, pipelined models with many stages—TF v1's architecture serializes the dispatch. Figure 1c shows that "host side work at the destination like dispatching the accelerator computation is triggered only after the transfer is completed." In a pipeline with $N$ stages, this adds $O(N)$ sequential dispatch delays, "leading to inefficient accelerator utilization" (Section 2).
Naive sharding blows up the graph: TF v1's dataflow representation explicitly materializes every shard as a separate node and every cross-shard communication as a separate edge. "An edge between an M-way sharded computation and an N-way sharded computation would require $M + N$ nodes and $M \times N$ edges, rapidly becoming unwieldy" (Section 2). At the scale of thousands of shards, the graph contains millions of edges, creating "substantial overhead in both graph serialization and execution."

The paper's diagnosis is that TF v1 was "over-specialized to assume a single, smallish, exclusively-owned island of accelerators" (Section 2). It provided the right programming model (single-controller, flexible dataflow) but the wrong implementation approach for scale. The dispatch overhead, sequential coordination, and graph explosion meant that TF v1 delivered the flexibility of single-controller but at an unacceptable performance cost—making it unusable for the large SPMD models that represent current best practice, even though those are the exact workloads any new system must support to gain adoption.

Where This Paper Positions Itself: The "Best of Both Worlds" Thesis

PATHWAYS positions itself as a single-controller system that matches multi-controller SPMD performance while enabling MPMD patterns. The paper's thesis is that the failures of TF v1 were implementation failures, not fundamental limitations of the single-controller approach. If you can solve the three problems that killed TF v1—dispatch latency, sequential coordination, and sharded representation blowup—then a single-controller design can deliver both today's performance and tomorrow's flexibility.

The paper explicitly frames this as a reconciliation:

"PATHWAYS combines the flexibility of single-controller frameworks with the performance of multi-controllers. We adopt a single-controller model since we believe it offers much better opportunities than multi-controller for novel and efficient ML computation, both by exploiting computational sparsity and heterogeneity, and by enabling cluster management systems that promote sharing and virtualizing resources." (Section 2)

The design response to each TF v1 failure mode is:

Dispatch latency → Parallel asynchronous dispatch (§4.5): Rather than serializing the dispatch of each computation after its predecessor completes (as in Figure 1b), PATHWAYS runs most host-side work for a chain of computations in parallel, exploiting the fact that compiled functions have statically known resource requirements. This allows the dispatch pipeline to stay full even when individual computations are small, amortizing DCN latency across multiple concurrent dispatches.
Sequential coordination → Gang-scheduled dynamic dispatch with a centralized scheduler per island (§4.4): By having one scheduler that consistently orders all computations across all programs sharing an island, PATHWAYS avoids the need for per-step global barriers or sequential coordination messages. The scheduler can interleave computations from different programs (multi-tenancy) while still ensuring consistent ordering for collectives.
Sharded graph explosion → A sharded dataflow representation (§4.3): PATHWAYS uses PLAQUE, a "sharded dataflow system where each node generates output data tuples tagged with a destination shard" so that a chain of two N-way sharded computations requires only 4 nodes in the IR (Arg → Compute(A) → Compute(B) → Result) regardless of N, rather than $O(N^2)$ edges.

The paper also positions itself relative to an emerging third category: systems like Ray (Moritz et al., 2018) that offer general-purpose distributed computing with actor-based parallelism. The paper compares PATHWAYS to Ray in its microbenchmarks (§5.1) and acknowledges that "it would be feasible to re-implement the full PATHWAYS design using other distributed frameworks such as Ray rather than PLAQUE" but notes that Ray lacks an HBM object store and primitives for efficient remote object transfer over GPU interconnects—capabilities that are critical for performance but could potentially be added.

The Deeper Motivation: Enabling the Research Agenda

A thread running throughout the paper is that PATHWAYS is not just for deploying models—it's for inventing them. Section 6.3 describes a concrete research vision that current systems cannot support:

"We want to enable research that uses fine-grain control flow so that different model weights can be updated per example, or even per sub-example (patch of an image, or word of a sentence)."

This kind of sparsity—where the computation graph itself is data-dependent and changes during training—requires a system that can: (a) express non-uniform per-accelerator computations naturally; (b) dynamically route data between accelerators based on computed values; (c) handle the resulting heterogeneous memory and compute loads; and (d) do all of this at scale without sacrificing the throughput that makes large-scale training feasible. No existing system at the time of PATHWAYS' publication provided all four. The paper's ambition is to build the platform on which this research becomes possible, while ensuring that the platform doesn't penalize researchers who are still working on today's SPMD models.

3. Technical Approach

3.1 Reader Orientation

PATHWAYS is a distributed runtime system—an orchestration layer sitting between ML user code (Python) and accelerator hardware (TPUs)—that takes a computation graph expressed as a collection of compiled functions and executes it across thousands of accelerators using a centralized controller, asynchronous dataflow coordination, and gang-scheduling of SPMD sub-computations. The core problem it solves is the single-controller performance gap: prior single-controller systems like TensorFlow v1 offered the right programming flexibility for non-SPMD workloads but suffered from dispatch latency, sequential coordination stalls, and graph representation blowup that made them uncompetitive with multi-controller SPMD systems on standard large-scale training. PATHWAYS' solution has three mutually-reinforcing design choices—parallel asynchronous dispatch (running host-side work for multiple computation nodes concurrently by exploiting statically known resource requirements), a centralized per-island scheduler (ordering all computations across all client programs to enable gang-scheduled collectives without global barriers), and a sharded dataflow representation (using Plaque, where a single edge in the IR represents communication between sharded computations regardless of shard count)—that together close the performance gap while retaining the programmability benefits of a single controller.

3.2 Big-Picture Architecture (Diagram in Words)

The PATHWAYS system has five major components, arranged in a client-server topology with centralized resource management:

Client library — receives user Python code (JAX or TensorFlow), traces it into a device-location-agnostic intermediate representation (IR) expressed as a custom MLIR dialect, progressively lowers this IR through compiler passes to produce a low-level program with physical device locations and explicit data-transfer operations, and submits the resulting program to the coordinator for execution. The client uses a sharded buffer abstraction to represent logical buffers distributed across multiple devices as single entities, amortizing reference counting and bookkeeping at the granularity of logical buffers rather than individual shards.
Resource Manager (global) — a centralized service that tracks all available accelerators across all islands, accepts requests for "virtual slices" of devices with specified 2D or 3D mesh shapes, and assigns physical devices to satisfy those requests. It maintains a one-to-one mapping between virtual and physical devices in the current implementation, and enables dynamic addition/removal of backend compute resources.
Scheduler (per island) — a centralized per-island component that consistently orders all computations enqueued by all client programs sharing that island's accelerators. It receives single messages describing entire subgraphs of statically-schedulable computations and sequences their execution, potentially interleaving computations from different concurrently-executing programs. The scheduler operates at millisecond timescales and currently uses FIFO ordering.
Executor (per device/shard) — the per-accelerator runtime responsible for dispatching compiled XLA computations to TPU cores, managing input and output buffer futures, and coordinating data transfers over the accelerator interconnect (ICI) and datacenter network (DCN).
Plaque (sharded dataflow coordination substrate) — an existing production sharded dataflow system that provides the low-level cross-host coordination over DCN. It converts the low-level PATHWAYS IR into a Plaque dataflow program, handles sparse data exchanges between dynamically chosen subsets of shards, enqueues local compiled functions on each accelerator, enqueues network sends for buffer futures to remote accelerators, and communicates with the per-island scheduler to determine consistent ordering.

Information flow for a typical execution: A user writes Python code decorating compiled functions with device placements → the client traces this into a PATHWAYS IR (device-location-agnostic, with one node per sharded computation regardless of shard count) → compiler passes lower the IR to include physical device locations, transfer subgraphs (scatter/gather operations when data exchange between different shardings is needed), and collective operations → the low-level program is converted to a Plaque dataflow program → the Plaque program enqueues local compiled functions at each accelerator, enqueues network sends for output buffer futures to remote accelerators, and communicates with the per-island scheduler to determine consistent ordering → compiled functions execute on TPU cores, with outputs flowing as futures through the dataflow graph → results are returned to the client as logical (sharded) buffers.

3.3 Roadmap for the Deep Dive

The explanation proceeds bottom-up through PATHWAYS' architecture, building from the lowest-level coordination primitives to the highest-level system properties:

First, the Plaque coordination substrate (§4.3), because it is the foundation on which all cross-host communication is built—understanding its sharded dataflow model, sparse communication support, and progress tracking is prerequisite to understanding how PATHWAYS avoids the graph explosion and sequential coordination problems that killed TF v1.
Second, gang-scheduled dynamic dispatch (§4.4), which explains how PATHWAYS enables SPMD collectives to execute correctly on shared (multi-tenant) accelerators—the centralized per-island scheduler, the consistent ordering requirement for TPU non-preemptible kernels, and why gang-scheduling matters even for GPUs.
Third, parallel asynchronous dispatch (§4.5), which is PATHWAYS' key performance innovation—how it exploits the statically known resource requirements of compiled functions to run host-side dispatch work for multiple nodes concurrently, amortizing DCN latency and preventing the serial dispatch bottleneck that limited TF v1.
Fourth, the Resource Manager (§4.1) and Client (§4.2), which together provide the centralized resource allocation and the programming interface—virtual devices, sharded buffer abstraction, and the progressive IR lowering pipeline that connects user code to the runtime.
Fifth, data management (§4.6), covering the sharded object store, HBM buffer tracking, opaque handles for migration, and garbage collection—the mechanisms that make intermediate program values available across host and accelerator memory boundaries.

This ordering follows the dependency chain: Plaque provides the messaging fabric → the scheduler uses it to order computations → parallel dispatch uses it to pipeline host-side work → the resource manager and client build on top to provide the user-facing abstractions → data management ties everything together by handling the lifecycle of tensors in flight.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems design paper whose core idea is that a single-controller architecture can match multi-controller SPMD performance if three specific implementation challenges are solved: dispatch latency, sequential coordination, and sharded representation blowup. The paper validates this claim through microbenchmarks (§5.1–5.2) and end-to-end model training (§5.3) on up to 2048 TPUs, demonstrating that PATHWAYS achieves ~100% accelerator utilization parity with JAX on SPMD computations while enabling pipelining, multi-island training, and multi-tenant resource sharing.

Plaque: The Sharded Dataflow Coordination Substrate

Plaque is an existing production sharded dataflow system at Google, used for customer-facing services where high-fanout or high-fanin communication is necessary, and both scalability and latency are important. PATHWAYS converts its low-level IR directly to a Plaque program, and all cross-host coordination that uses the datacenter network (DCN) flows through Plaque's dataflow graph. Understanding Plaque's properties is essential because they dictate what PATHWAYS can and cannot express efficiently.

Sharded representation without $O(N^2)$ edges. The most critical property Plaque provides is that each node in the dataflow IR represents a sharded computation—a single logical operation that may be instantiated across many physical accelerators—rather than representing each physical shard as a separate node. The paper gives a concrete example: "a chained execution of 2 computations A and B with N computation shards each should have 4 nodes in the dataflow representation: Arg → Compute(A) → Compute(B) → Result, regardless of the choice of N" (Section 4.3). This is in direct contrast to TensorFlow v1, where the same computation would require $2N$ nodes and $N^2$ edges (every shard of A would have an edge to every shard of B if there is a data exchange), causing the graph to "rapidly become unwieldy" at thousands of shards.

How it works at runtime. In Plaque's runtime, each IR node generates output data tuples tagged with a destination shard identifier. When performing data-parallel execution, $N$ data tuples flow between each adjacent pair of IR nodes—one per shard—but these $N$ tuples are managed within the runtime's internal data structures, not represented as $N$ separate edges in the IR graph. The coordination runtime tracks progress per shard and detects completion when all expected messages for a shard have been received.

This mechanism is exactly what enables PATHWAYS to represent computations spanning thousands of shards compactly. The IR stays small (proportional to the logical computation graph, not the physical shard count), serialization and deserialization overhead is constant with respect to shard count, and the runtime handles the mapping from logical to physical data movement internally.

Sparse communication support. Plaque supports "sparse data exchanges along sharded edges, in which messages can be sent between a dynamically chosen subset of shards" (Section 4.3). This uses standard progress tracking mechanisms from stream processing systems (Akidau et al., 2013; Murray et al., 2013) to detect when all messages for a shard have been received. The paper explicitly states that "efficient sparse communication is a requirement to avoid the DCN becoming a bottleneck for data-dependent control flow on accelerators, which is one of the key capabilities that we want PATHWAYS to enable" (Section 4.3).

This is the mechanism that will eventually support Mixture-of-Experts style routing (Section 6.3), where different examples are routed to different subsets of accelerators based on data-dependent routing decisions. Without sparse communication support, the system would need to either (a) send data to all possible destinations and discard at the receiver (wasting DCN bandwidth proportional to the full fanout), or (b) implement dynamic routing in user code with explicit send/recv pairs (returning to the TF v1 bottleneck where host-side dispatch triggers only after transfer completion).

Latency-throughput dual mode. The coordination substrate must handle two regimes: low-latency for critical-path scheduling messages and data handles, and high-throughput for bulk data movement. Plaque supports both by sending "critical messages with low latency" while batching "messages destined for the same host when high throughput is required" (Section 4.3). This is important because PATHWAYS separates control messages (which determine what executes when) from data messages (which move tensor data). Control messages—scheduling requests, future handoffs, completion signals—must be low-latency to minimize idle time between computations. Data messages—gradient transfers between islands, activation forwarding in pipelines—need high bandwidth.

General-purpose extensibility. Because Plaque is a general-purpose dataflow engine, PATHWAYS can also use it for "background housekeeping tasks such as distributing configuration information, monitoring programs, cleaning them up, delivering errors on failures, and so on" (Section 4.3). This means PATHWAYS doesn't need a separate control channel or coordination service—all communication, both performance-critical and operational, flows through the same dataflow substrate.

Design choice: why Plaque rather than building from scratch. The paper implicitly argues that a sharded dataflow system with the properties described—compact sharded IR representation, sparse communication with progress tracking, dual latency/throughput mode, and general-purpose extensibility—is a non-trivial piece of infrastructure to build. By adopting Plaque as an existing production system, PATHWAYS inherits all of these properties and focuses its novel engineering effort on the ML-specific concerns: compiler integration, gang scheduling, parallel dispatch, and accelerator memory management.

Design choice: why not Ray. The paper explicitly considers whether Ray (Moritz et al., 2018) could substitute for Plaque: "we believe that it would be feasible to re-implement the full PATHWAYS design using other distributed frameworks such as Ray rather than PLAQUE to realize the low-level coordination framework" (Section 4.3). In such an implementation, "PATHWAYS executors and schedulers would be replaced by long-running Ray actors that would implement PATHWAYS scheduling on top of the underlying Ray cluster scheduling, and executors could use PyTorch for GPU computation and collectives." However, the paper notes that Ray lacks certain capabilities that PATHWAYS requires for comparable performance: an HBM object store (Ray stores objects in host DRAM, not accelerator memory), and primitives to efficiently transfer remote objects over the GPU interconnect (GPUDirect or equivalent). These are not fundamental limitations—they could be added to Ray—but they represent engineering work that PATHWAYS avoids by using Plaque, which already provides an object store that tracks HBM buffers per shard (Section 4.6).

Gang-Scheduled Dynamic Dispatch

The requirement for gang-scheduling arises from a fundamental constraint of TPU hardware: TPUs are single-threaded and only run non-preemptible kernels, meaning that if two communicating computations are enqueued in different orders on different devices, the system will deadlock—device A waits for a message from device B on collective operation $X$ , but device B is executing a different collective operation $Y$ because it received a different enqueue order, and neither can proceed. The paper states this precisely: "the system will deadlock if communicating computations are not enqueued in a consistent order" (Section 2).

Why gang-scheduling matters even for GPUs. Although GPUs can execute concurrent computations (unlike TPUs, which are restricted to a single program at a time), the paper argues that "gang scheduling allows more efficient execution of collectives" (Section 2), citing Feitelson and Rudolph (1992). The reasoning: when all participants in a collective operation start at approximately the same time (gang-scheduled), the collective can proceed without waiting for stragglers. If computations are independently scheduled, a collective might have some participants arrive early and idle while waiting for late arrivals, reducing utilization. So while gang-scheduling is mandatory for TPU correctness, it is also beneficial for GPU efficiency.

The per-island centralized scheduler. PATHWAYS includes "a centralized scheduler per island that consistently orders all of the computations in the island" (Section 4.4). This is a single component per island (not per host or per device) that receives scheduling requests from all Plaque dataflow programs executing on behalf of all clients sharing that island's accelerators. The scheduler's job is to sequence computations such that when a collective operation is enqueued on multiple devices, all devices receive the enqueue instruction in the same relative order with respect to other computations on those devices.

How the scheduler integrates with Plaque. The Plaque dataflow program for a PATHWAYS computation is responsible for three actions per enqueued computation: (i) enqueueing the execution of local compiled functions at each accelerator, with buffer futures as inputs; (ii) enqueueing network sends to remote accelerators for the buffer futures output by function executions; and (iii) communicating with the scheduler to determine a consistent order of function executions across all programs running on the island (Section 4.4). The scheduler receives these communication requests and returns ordering decisions that the Plaque program respects when actually dispatching the accelerator computations.

Single-message subgraph scheduling. When a subgraph of a computation can be scheduled statically (i.e., all resource requirements are known before any predecessor executes), "the program sends a single message (describing the entire subgraph) to the scheduler, which is able to sequence the execution of all the active shards in the subgraph back to back" (Section 4.5). This single-message design minimizes network traffic between the Plaque program and the scheduler—rather than sending one message per shard per computation (which would be $O(N \times \text{steps})$ ), the system sends one message per subgraph. Importantly, "this does not require the scheduler to actually enqueue all the subgraph's shards as a batch: computations may still be interleaved with those submitted by other concurrently executing programs" (Section 4.5). The scheduler maintains the freedom to interleave at the granularity of individual computations, not entire subgraphs, which is what enables fine-grained multi-tenancy.

Current scheduling policy and future extensibility. The paper states that the "current implementation simply enqueues work in FIFO order, but more sophisticated schedulers might for example reorder computations based on estimated execution times" (Section 4.4). This is conservative: FIFO is simple, predictable, and avoids starvation. More sophisticated policies (shortest-job-first to reduce head-of-line blocking, priority-based for latency-sensitive inference vs. throughput-oriented training, proportional share for fairness) are left as future work. The key architectural point is that a centralized scheduler has the information needed to implement these policies, whereas a multi-controller system (where each host independently decides what to run) fundamentally cannot make globally optimal scheduling decisions.

Design choice: why centralized per-island rather than distributed. A distributed consensus protocol for ordering could theoretically work—every host could participate in a Paxos or Raft group to agree on computation ordering. The paper argues implicitly against this by noting that scheduling must happen "at a time-scale of milliseconds" (Section 4.4). Distributed consensus at millisecond granularity across hundreds of hosts would be challenging (and would add its own DCN latency to every scheduling decision). A centralized scheduler for a single island—where all devices in the island are connected via low-latency ICI—can make decisions in microseconds and communicate them to all hosts within the island quickly.

Interaction with multi-tenancy (Section 5.2). The scheduler's ability to interleave computations from different client programs is demonstrated in Figure 8 and Figure 9. When multiple clients submit programs concurrently, the scheduler time-multiplexes the accelerators among them, interleaving gang-scheduled computations at sub-millisecond granularity. Figure 9 shows traces where 4 clients share accelerators with proportional-share ratios of 1:1:1:1 and 1:2:4:8—the scheduler enforces these ratios by controlling how many computations from each client are enqueued in each scheduling epoch. This is a concrete demonstration that gang-scheduling and multi-tenancy are compatible: the scheduler maintains consistent ordering for each program's collectives while also ensuring fairness across programs.

Parallel Asynchronous Dispatch

This is PATHWAYS' most novel performance mechanism and the key to matching multi-controller dispatch latency with a single-controller architecture. To understand it, we first need to understand the baseline (sequential) dispatch model and why it fails.

The sequential dispatch problem (Figure 4a). Consider a three-node dataflow graph where node A runs on host A, node B on host B, and node C on host C. In the standard asynchronous dispatch model (used by TF v1 and the fallback mode of PATHWAYS), the sequence of events is:

Host A enqueues node A on its local accelerator, receives a future for A's outputs.
Host A transmits the future to host B over DCN.
Host B receives the future, allocates B's input buffers, transmits the buffer addresses back to host A, and performs preparatory work to launch node B.
When node A completes, its outputs are transferred via the accelerator interconnect (ICI) directly into B's input buffers (zero-copy within the island).
Host B starts node B.
Host B transmits B's output futures to host C over DCN.
Host C receives, allocates, transmits buffer addresses, prepares to launch node C.
When B completes, its outputs transfer to C's input buffers via ICI.
Host C starts node C.

This works well when each computation's execution time exceeds the coordination latency. The key insight from Figure 4a is that host-side work for a node (allocation, buffer address exchange, launch preparation) happens only after the predecessor has been enqueued and its output futures have been transmitted. If computation times are very short relative to coordination latency—"which is the case shown in the figure"—the pipeline stalls: "the asynchronous pipeline stalls and the host-side work becomes the critical bottleneck for executing the overall sequence of computations" (Section 4.5).

For SPMD computations with small per-step work (e.g., micro-benchmarks with scalar AllReduce, or models with very thin layers), this sequential dispatch would leave accelerators idle while waiting for DCN messages to propagate host-side work through the pipeline. Multi-controller systems avoid this entirely because each host's copy of the program can independently enqueue the next computation over PCIe without any DCN coordination (Figure 1a).

The key insight: compiled functions have statically known resource requirements. PATHWAYS observes that compiled functions—the XLA computations that dominate ML workloads—have a crucial property: their input and output types and shapes, loop bounds, and memory requirements are known before the computation executes (Appendix B). This means that a successor node's input shapes "can in practice be computed before the predecessor computation was even enqueued" (Section 4.5). If we know what resources node B will need before node A finishes (or even before node A starts), we can do the host-side work for node B in parallel with node A's execution.

The parallel dispatch mechanism (Figure 4b). In parallel asynchronous dispatch, PATHWAYS exploits this property to "run most of the host-side work for a computation's nodes in parallel, rather than serializing the work for a node to happen after its predecessors have been enqueued" (Section 4.5). Concretely, for a chain of nodes A → B → C:

The client (or Plaque program) sends a single message describing the entire subgraph A → B → C to the scheduler.
Host-side work for all three nodes—allocation, buffer address exchange, launch preparation—proceeds in parallel on their respective hosts.
The scheduler sequences the enqueue order consistently across all devices.
When A completes, B can start immediately because all preparatory work is already done; similarly for B → C.

The effect is that coordination latency is amortized across the entire subgraph rather than being paid per edge. Instead of latency $t_{\text{coord}}$ between every pair of nodes, the system pays latency once per subgraph. For a subgraph with $k$ nodes, this reduces total coordination latency from $k \times t_{\text{coord}}$ to approximately $t_{\text{coord}}$ (plus the remaining unavoidable data transfer time between device memories).

The fallback to sequential dispatch. Parallel dispatch is only safe when resource requirements are statically known. If a computation has data-dependent control flow—where the shapes, loop bounds, or memory requirements depend on values computed by a predecessor—PATHWAYS "falls back to the traditional model when a node's resource requirements are not known until a predecessor computation has completed" (Section 4.5). This is a graceful degradation: the system uses the fast path when it can prove safety and the slow path when it cannot. Since "almost all of today's high performance ML computations are expressed as long stretches of compiled functions and only occasionally (if ever) branch based on data that is computed by a compiled function" (Appendix B), the fast path covers the vast majority of execution time.

Quantitative evidence from microbenchmarks (Figure 7). The paper evaluates parallel dispatch using a pipelined benchmark where simple computations are chained across different sets of 4 TPU cores, each on a different host, with data transferred via ICI between stages. Figure 7 shows that with sequential dispatch, throughput drops as more pipeline stages are added (because each stage adds coordination latency to the critical path). With parallel dispatch, throughput initially increases (amortizing fixed client overhead), then decreases as transfer costs grow with stage count, then increases again (amortizing scheduling overhead), eventually expected to decrease as transfer costs dominate again. The key message is that parallel dispatch prevents coordination latency from being the bottleneck for deep pipelines.

Why this matters for the single-controller thesis. The dispatch latency problem was the primary reason single-controller systems were considered inherently slower than multi-controller systems. Parallel asynchronous dispatch challenges this assumption by showing that the latency gap is not fundamental—it arises from when host-side work is scheduled relative to computation, not from the number of controllers. By moving host-side work earlier (before predecessor completion, exploiting static knowledge of resource requirements), PATHWAYS achieves the same effect as multi-controller systems (host-side work overlapping with computation) without requiring a copy of the user program on every host.

Resource Manager and Virtual Devices

The Resource Manager is the centralized component that "is responsible for the centralized management of devices across all of the islands" (Section 4.1). It serves as the indirection layer between client requests for compute resources and the physical allocation of those resources—a role that is impossible in multi-controller systems where each host independently owns its accelerators.

Virtual slices and virtual devices. A client requests compute resources by asking for "virtual slices" of an island with specific properties: desired 2D or 3D mesh shapes that suit the client's communication pattern, optional constraints on device types, locations, or interconnect topology. Each virtual slice contains "virtual devices" that the client uses to express how computations are laid out on the mesh. For example, a client might request a 2D mesh of 8 × 8 virtual TPUs to run a model with data parallelism along one dimension and model parallelism along the other.

The Resource Manager "dynamically assigns physical devices for virtual devices satisfying the desired interconnect topology, memory capacity, etc." (Section 4.1). This layer of indirection between virtual and physical devices is what enables the system to transparently support features like suspend/resume and migration in the future—the client's program references virtual devices, and the system can remap those to different physical devices without the client's cooperation.

Current allocation heuristic. The initial Resource Manager implementation "uses a simple heuristic that attempts to statically balance load by spreading computations across all available devices, and keeps a one to one mapping between virtual and physical devices" (Section 4.1). The one-to-one mapping means that each virtual device maps to exactly one physical device—there is no over-subscription or sharing of a physical device by multiple virtual devices from the same or different clients. This is conservative but correct: it avoids the complexity of memory management and performance isolation when multiple virtual devices share a physical accelerator.

The paper acknowledges that future implementations could adopt "a more sophisticated allocation algorithm, for example taking into account the resource requirements of all client computations and the current state of the system to approximate an optimal allocation of physical devices to computations" (Section 4.1).

Dynamic resource changes. The Resource Manager allows "backend compute resources to be added and removed dynamically, with the resource manager tracking available devices" (Section 4.1). This supports cluster elasticity—adding new TPU pods to the system, or taking pods offline for maintenance, without requiring PATHWAYS to restart. The indirection between virtual and physical devices means that when physical resources change, the system can re-lower programs to map virtual devices to the new physical configuration.

Design choice: why centralized resource management. Multi-controller systems leave resource allocation to the cluster scheduler (e.g., Kubernetes, Slurm) at job granularity—the scheduler assigns a set of physical devices to a job, and the job runs on those devices until completion. This prevents fine-grained sharing (a job that underutilizes its allocated devices cannot donate cycles to another job) and complicates features like elasticity (changing a running job's allocation requires coordination with the job's own resource management logic). PATHWAYS' centralized Resource Manager, by contrast, can make allocation decisions at finer granularity and adjust them over time, because all client programs interact with the same resource management authority.

Client and Intermediate Representation (IR) Pipeline

The client is the programming interface through which user code becomes executing PATHWAYS programs. It performs three major functions: compilation and registration, IR construction and lowering, and sharded buffer management.

Compilation and registration. When a user runs a traced program, "the PATHWAYS client library first assigns virtual devices to any computations that have not been run before, and registers the computations with the resource manager, triggering the servers to compile the computations in the background" (Section 4.2). This means that XLA compilation—which can be expensive for large models—happens asynchronously on the server side, not on the critical path of program execution. When the same compiled function is used in subsequent programs, the compiled version is already available.

IR construction and progressive lowering. The client "constructs a device location-agnostic PATHWAYS intermediate representation (IR) for the program, expressed as a custom MLIR dialect" (Section 4.2). MLIR (Lattner et al., 2021) is a compiler infrastructure that supports defining custom intermediate representations and progressively lowering them through a series of transformation passes. PATHWAYS uses this to separate concerns:

Device-location-agnostic IR: The initial IR represents the computation graph without knowledge of which physical devices will execute each computation. It contains nodes for each compiled function and edges representing data dependencies between them. This IR does not include transfer, scatter, or gather operations—it represents what the user intends to compute, not how it is physically realized.
Lowering passes: The IR is "progressively 'lowered' via a series of standard compiler passes, which eventually output a low-level representation that includes the physical device locations" (Section 4.2). These passes consult the Resource Manager's mapping from virtual to physical devices and insert the necessary communication operations:
- Transfer subgraphs for moving output data from the physical location of a source computation shard to the physical location of a destination computation shard.
- Scatter operations when a data exchange requires distributing data from fewer to more shards.
- Gather operations when a data exchange requires collecting data from more to fewer shards.
- Collective operations (AllReduce, etc.) over the ICI interconnect for SPMD sub-computations.
The low-level program "takes into account the network connectivity between physical devices" (Section 4.2)—it knows which devices are connected via ICI (low-latency, high-bandwidth, within an island) and which are connected only via DCN (higher-latency, lower-bandwidth, across islands), and chooses the appropriate communication mechanism for each edge.

Re-lowering on remapping. The low-level program is efficient to repeatedly run "in the common case that the virtual device locations do not change." If the Resource Manager changes the mapping between virtual and physical devices (due to elasticity, migration, or load rebalancing), "the program can be re-lowered" (Section 4.2). This is the mechanism that enables transparent migration: the client's high-level IR is independent of physical placement, so changing physical placement only requires re-running the lowering passes.

Sharded buffer abstraction. This is a critical scalability mechanism: "the PATHWAYS client uses a sharded buffer abstraction to represent a logical buffer that may be distributed over multiple devices" (Section 4.2). Rather than tracking each physical shard of a tensor as a separate object with its own reference count and lifecycle, the client tracks a single logical buffer. This "helps the client scale by amortizing the cost of bookkeeping tasks (including reference counting) at the granularity of logical buffers instead of individual shards" (Section 4.2).

Consider a data-parallel computation across 2048 TPU cores. With per-shard tracking, every all-reduced gradient would require 2048 reference count updates, allocation records, and lifecycle transitions. With the sharded buffer abstraction, the entire gradient is one logical buffer—one reference count, one lifecycle event. The runtime handles the per-shard details internally. This is what makes the single-controller model feasible at scale: the client's control plane complexity is proportional to the logical computation graph size (number of operations), not the physical shard count (number of devices).

Program tracing vs. individual function calls. The paper describes two modes of client interaction:

Default (no tracing): Each compiled function is converted into a standalone PATHWAYS program containing just one sharded computation. "A separate Python call and RPC from client to coordinator is required for each function" (Section 3). This is the simplest mode—each call to a jax.pmap-decorated function becomes its own PATHWAYS program—but incurs the RPC overhead for every function invocation.
Program tracer (Figure 2): A user can wrap a block of Python code that calls many compiled functions inside a @pw.program decorator. "The tracer generates a single PATHWAYS program where each compiled function is represented by a computation node in a dataflow graph" (Section 3). A single RPC submits the entire traced program, avoiding per-function RPC overhead and enabling the compiler to optimize across function boundaries within the program.

The program tracer is the mechanism that enables PATHWAYS to "execute back-to-back accelerator computations directly from C++ while JAX OpByOp transitions to Python for every computation" (Section 5.1), which is why PATHWAYS Chained outperforms JAX OpByOp at moderate scales in Figure 5.

Integration with JAX and TensorFlow. PATHWAYS supports both JAX and TensorFlow as frontends. For JAX, "PATHWAYS can be used as a plug-in replacement for the JAX backend, allowing JAX code to run unmodified" except that SPMD computations now have access to as many cores as are provisioned, and "can communicate over both ICI and DCN," allowing JAX programs to scale "for the first time to multiple TPU pods, containing many thousands of TPU cores" (Section 3). For TensorFlow, PATHWAYS can execute TF graphs as computation nodes, as demonstrated in the 3B Transformer pipeline experiment (Table 2).

Data Management: Sharded Object Store and Garbage Collection

PATHWAYS maintains a distributed object store that spans both host DRAM and accelerator HBM, enabling intermediate program values to persist across computation boundaries without unnecessary copying.

Sharded object store with HBM integration. "Each host manages a sharded object store that is similar to Ray's object stores (Moritz et al., 2018), but extended to also track buffers held in accelerator HBM at each shard" (Section 4.6). This is the critical difference from Ray: Ray's object store resides in host DRAM, so transferring a tensor from one GPU to another requires the tensor to be moved from GPU HBM to host DRAM (to create the Ray object), then from host DRAM to the other GPU's HBM—two PCIe transfers. PATHWAYS' object store tracks buffers directly in HBM, enabling zero-copy transfer between accelerators within the same island over ICI, and between islands over DCN with only one DCN transfer (no intermediate host DRAM staging).

Opaque handles for remote references. "Client programs can hold references to objects in remote host or accelerator memory, and the client and servers refer to them using opaque handles that allow the system to migrate them if needed" (Section 4.6). These opaque handles are what flow through the Plaque dataflow graph as futures (Section 4.4). When a Plaque program enqueues a network send for a buffer future, it is transmitting an opaque handle, not the buffer data itself. The receiving host can then use the handle to reference the data, potentially triggering a transfer if the data resides on a different physical device.

Intermediate value retention. "Intermediate program values are also kept in the object stores, for example while the system is waiting to transfer them between accelerators, or pass them to a subsequent computation" (Section 4.6). This means that if computation A produces a tensor that computation B will consume, the tensor stays in the object store (potentially in HBM if B is on the same island) rather than being copied back to the client or to host DRAM. This is what enables the efficient chaining of back-to-back computations observed in Figure 5.

Ownership and garbage collection. Objects are "tagged with ownership labels so that they can be garbage collected if a program or client fails" (Section 4.6). The ownership label identifies which program (and which client) created the object. If a program completes successfully, its intermediate objects are released. If a client disconnects or crashes, all objects owned by that client's programs are reclaimed. This is essential for multi-tenancy: without garbage collection, a misbehaving or crashed client could leak HBM memory, eventually starving other clients.

Back-pressure for memory management. PATHWAYS implements "simple back-pressure to stall a computation if it cannot allocate memory because other computations' buffers are temporarily occupying HBM" (Section 4.6). If the HBM on a particular device is full, the scheduler or executor can delay launching new computations on that device until sufficient memory is freed. This is a basic form of flow control that prevents out-of-memory errors from crashing programs. The paper does not describe a sophisticated memory scheduling policy (e.g., eviction, defragmentation, swapping to host DRAM), but the back-pressure mechanism provides a foundation that such policies could build on.

Comparison to Ray and TensorFlow. The paper explicitly compares the data management architectures in the microbenchmarks (Section 5.1):

Ray: "must transfer the result of a computation from GPU to DRAM before returning the object handle to the client." This is because Ray's object store is in DRAM—there is no concept of an object that resides purely in GPU memory with a remote reference.
TensorFlow v1: "transfers the data back to the client." This is even worse: the client (which may be on a different machine) receives the actual tensor data, not just a handle. For large tensors, this saturates the DCN and creates a serialization bottleneck.
PATHWAYS: keeps data in the distributed object store (potentially in HBM) and passes opaque handles. The client never needs to materialize intermediate tensors. This is why PATHWAYS' OpByOp throughput degrades less severely than TF or Ray at larger scales in Figure 5—the data stays near the accelerators rather than being funneled through a central point.

This design is what makes the single-controller model scalable: the controller (client + coordinator) deals only with metadata (handles, scheduling decisions, buffer lifetimes), while the actual data movement happens directly between accelerators over high-bandwidth interconnects, orchestrated by the Plaque dataflow runtime.

Summary of Key Design Choices and Their Justifications

The paper makes specific architectural choices, each with a clear motivation grounded in the limitations of prior systems:

Single-controller over multi-controller: enables centralized resource management, multi-tenancy, and natural expression of non-SPMD computations (heterogeneous sub-computations, sparse routing). The risk—dispatch latency from remote clients—is mitigated by parallel asynchronous dispatch.
Plaque (sharded dataflow) over bespoke coordination: leverages an existing production system with proven sharded representation (avoiding $O(N^2)$ edge blowup), sparse communication (for future data-dependent routing), and dual latency/throughput messaging. The alternative of building custom coordination would duplicate substantial infrastructure.
Per-island centralized scheduler over distributed consensus: enables consistent ordering for gang-scheduled collectives with microsecond-level decision latency, mandatory for TPU deadlock avoidance and beneficial for GPU collective efficiency. Distributed consensus at millisecond timescales would add latency and complexity.
Parallel asynchronous dispatch over purely sequential dispatch: exploits the static resource requirements of compiled functions (known shapes, loop bounds, memory needs) to run host-side work for multiple nodes concurrently, amortizing DCN coordination latency across pipeline stages. This is what closes the gap with multi-controller systems on small-to-medium computations.
Sharded buffer abstraction over per-shard tracking: amortizes client-side bookkeeping (reference counting, lifecycle management) at the granularity of logical buffers rather than physical shards, preventing the client from becoming a bottleneck at thousands of accelerators.
HBM-integrated object store with opaque handles over host-DRAM-only storage: enables zero-copy data transfer between accelerators within an island (via ICI) and avoids the serialization bottleneck of returning intermediate results to the client.
MLIR-based progressive lowering over fixed compilation: separates concerns between device-location-agnostic user programs and physical-device-aware execution, enabling transparent re-lowering when the virtual-to-physical device mapping changes (for elasticity, migration).
FIFO scheduling as initial policy (not a design limitation): provides simplicity and correctness as a baseline, with the architecture designed to support more sophisticated policies (shortest-job-first, priority-based, proportional share) in the future, since the centralized scheduler has the global information needed to make such decisions.

4. Key Insights and Innovations

Innovation 1: The Single-Controller Performance Ceiling Is an Implementation Artifact, Not a Fundamental Constraint

The dominant architectural assumption in distributed ML systems prior to PATHWAYS was that single-controller and multi-controller models represent an inherent tradeoff: you could have programming flexibility (heterogeneous computations, centralized resource management, non-SPMD patterns) or you could have performance at scale (low dispatch latency, efficient collectives), but not both. TensorFlow v1's failure to compete with multi-controller JAX and PyTorch on large-scale SPMD training had cemented this assumption as community wisdom—the single-controller model was considered too slow for production-scale workloads, and the field broadly accepted that the future of distributed ML infrastructure lay in refining multi-controller systems despite their programmability limitations.

PATHWAYS' most fundamental intellectual contribution is reframing this as an engineering problem rather than an architectural one. The paper identifies three specific, solvable mechanisms that caused TF v1's performance gap—dispatch latency from serial coordination, sequential host-side work blocking the pipeline, and graph representation blowup from naive sharding—and demonstrates that each admits a targeted solution that does not require abandoning the single-controller model. The critical conceptual move is recognizing that the latency advantage of multi-controller systems comes from when host-side work happens relative to computation, not from where the controller resides. Multi-controller systems run host-side work locally over PCIe, overlapping it with accelerator execution; single-controller systems can achieve the same overlap through parallel asynchronous dispatch that exploits the statically known resource requirements of compiled functions.

This is a fundamental insight with implications beyond PATHWAYS itself. It means that future systems need not choose between flexibility and performance—they can pursue both simultaneously by investing in the right coordination mechanisms. The paper's evidence for this reframing is the performance parity demonstrated in Section 5.3 (Table 1: identical throughput across T5-Base through T5-11B between PATHWAYS and JAX) and the convergence analysis in Figure 6 (a computation of only 2.3 ms masks all single-controller overhead at 128 TPU scale). These are not marginal improvements—they represent complete elimination of the performance gap that had previously disqualified single-controller architectures from serious consideration for large-scale training.

The negative implication is equally important: TF v1's failure was not evidence that single-controller designs cannot work at scale, but rather evidence that specific implementation choices in TF v1 (serialized dispatch, per-shard graph materialization, no gang-scheduling) were incompatible with scale. This is a diagnostic contribution that separates the architectural concept (single controller) from the implementation (how coordination is implemented), preventing the field from erroneously discarding a promising design space based on one flawed instantiation.

Innovation 2: Gang-Scheduling as a Centralized Mechanism That Enables Multi-Tenancy, Not Just Correctness

The conventional understanding of gang-scheduling in distributed ML is that it is a necessary evil for correctness on TPUs (to prevent deadlock from inconsistent collective enqueue ordering) and a performance optimization on GPUs (to reduce straggler effects in collectives). PATHWAYS repurposes gang-scheduling as something much more architecturally significant: the mechanism that makes fine-grained multi-tenancy compatible with SPMD collectives.

This is a conceptual shift. Prior work on accelerator sharing either accepted coarse granularity (dedicating whole accelerators to single jobs for seconds or more, as in conventional cluster schedulers like Gandiva (Xiao et al., 2018) or Themis (Mahajan et al., 2020)), or pursued fine-grained sharing by breaking the SPMD model entirely (running multiple independent non-communicating computations on the same accelerator). PATHWAYS shows that with a centralized per-island scheduler that consistently orders computations across all client programs, you can interleave gang-scheduled SPMD sub-computations from different programs on the same set of accelerators at sub-millisecond granularity, with proportional-share fairness policies. Figure 9 demonstrates this concretely: traces of cores show interleaving of computations from 4 independent clients with configurable share ratios (1:1:1:1 and 1:2:4:8), while maintaining the consistent ordering that gang-scheduling requires.

The architectural significance is that gang-scheduling transforms from a constraint into a capability. In multi-controller systems, gang-scheduling is implicit (all participants run the same program and enter collectives at the same time by construction) but this locks out multi-tenancy—you cannot interleave someone else's computation because you cannot control the ordering. In PATHWAYS, gang-scheduling is explicit (the centralized scheduler orders everything) and this explicitness is precisely what enables controlled interleaving: because the scheduler sees all programs, it can order computations from different programs while ensuring each program's collectives remain consistently ordered.

This insight generalizes: any mechanism that provides a global ordering point can simultaneously serve as a policy enforcement point. The scheduler achieves both correctness (ordering for deadlock avoidance) and resource management (fairness, utilization, priority). Figure 8 provides the quantitative validation: PATHWAYS achieves at least the same aggregate throughput as JAX when multiple clients submit programs concurrently, with no overhead to context switch—meaning the centralized scheduling does not introduce a performance tax relative to exclusive ownership.

Innovation 3: The Compiled Function Abstraction as the Boundary Between Static and Dynamic Scheduling

PATHWAYS' architecture introduces a clean phase separation between what can be scheduled statically (before execution) and what must be scheduled dynamically (during execution), using the compiled function as the boundary. This is a systems design insight that resolves a tension present in prior dataflow systems: static scheduling enables aggressive optimization (parallel dispatch, single-message subgraph scheduling, compile-time memory allocation) but cannot handle data-dependent computation shapes; dynamic scheduling handles arbitrary control flow but incurs runtime coordination overhead.

The innovation is not splitting static and dynamic scheduling (which is an old idea) but identifying that the compiled function—a concept already present in ML frameworks for JIT compilation—is the natural abstraction at which to draw this boundary, and building the entire system architecture around respecting this boundary with a graceful fallback. When a subgraph consists entirely of compiled functions with known resource requirements, PATHWAYS uses the fast path: single-message subgraph scheduling (Section 4.4), parallel asynchronous dispatch (Section 4.5), and compile-time buffer allocation (Appendix B). When a computation has data-dependent requirements—shapes, loop bounds, or memory needs that depend on predecessor outputs—PATHWAYS "falls back to the traditional model" of sequential dispatch (Section 4.5).

This matters because it means the system's performance properties are predictable and compositional. Users (and automated tools converting model code to PATHWAYS programs) can look at a computation graph and know which parts will use the fast path and which will use the fallback. Since "almost all of today's high performance ML computations are expressed as long stretches of compiled functions" (Appendix B), the fast path covers the overwhelming majority of execution. The fallback exists for the rare data-dependent cases but doesn't penalize the common case.

This is conceptually similar to how database systems separate query planning (static) from query execution (dynamic), or how compilers separate ahead-of-time optimization from just-in-time compilation. PATHWAYS applies this separation to distributed ML scheduling, and the concrete result is that the single-controller dispatch overhead (Figure 5: OpByOp throughput degradation at scale) can be avoided for all realistic workloads without sacrificing support for dynamic control flow.

Innovation 4: The Training-Inference Resource Boundary as a Unified Execution Model

A less explicit but architecturally significant innovation in PATHWAYS is the unification of what were previously separate system designs for different ML lifecycle phases under a single execution model. By providing a centralized resource manager, gang-scheduled dispatch, and multi-tenancy, PATHWAYS creates an architecture that can simultaneously handle training workloads (high throughput, large SPMD computations, long-running jobs) and inference-serving patterns (multiple concurrent clients, fine-grained request interleaving, shared model layers) within the same framework.

Section 6.3 sketches the research vision this enables: foundation models where "several researchers might concurrently fine-tune a foundation model for different tasks, using the same accelerators to hold the fixed foundation model layers" and "training or inference over shared sub-models can benefit from techniques that allow examples from different tasks to be combined in a single vectorized batch to get better accelerator utilization." This is a pattern that current systems handle with separate infrastructure: training uses one system (multi-controller, exclusive ownership, batch-optimized) while serving uses another (model servers like Clipper (Crankshaw et al., 2017), request-level scheduling, latency-optimized). PATHWAYS' architecture dissolves this boundary: the same centralized scheduler that gang-schedules SPMD training steps can interleave inference requests from multiple clients, and the same resource manager that allocates virtual slices for training can allocate fractions of accelerators for shared model layers.

The significance is that this eliminates a form of system-induced path dependency: if training and inference use different systems, models must be ported between them, and optimizations developed in one context do not transfer to the other. With a unified execution model, the same compiled functions run in both contexts, and the system can adapt resource allocation based on workload mix—shifting accelerators from training to inference as demand changes, or colocating fine-tuning and inference on shared base layers. This is not demonstrated experimentally in the paper (the evaluation focuses on training throughput), but it is a direct architectural consequence of the single-controller, centralized-resource-management design, and it represents a conceptual advance over the bifurcated training/serving infrastructure that dominates current ML deployment.

5. Experimental Analysis

Evaluation Methodology

Dataset. All experiments use the MATH benchmark (Hendrycks et al., 2021), consisting of high-school competition-level math problems. The authors use the specific split from Lightman et al. (2022): 12,000 training questions and 500 test questions. The paper characterizes the benchmark as requiring "multi-step logical deduction rather than novel factual recall," making it well-suited to studying test-time compute strategies that amplify existing reasoning capability.
Base model(s). All experiments use PaLM 2-S* (Codey) (Anil et al., 2023). The authors select this model because it is "representative of the capabilities of many contemporary LLMs" and sits in a regime of non-trivial but far-from-saturated MATH performance (roughly 10–19% pass@1 depending on prompt and sampling configuration), leaving substantial room for test-time compute to improve results. For the FLOPs-matched comparison (Section 7), a second model with approximately 14× more parameters is used as the pretraining-scaled baseline, evaluated with greedy decoding only (no extra test-time compute).
Metrics. The primary metric is MATH test accuracy (%), defined as the fraction of the 500 test questions for which the selected final answer matches the ground truth. Answers are graded using the grading function released by Lightman et al. (2022) (Appendix G). When analyzing difficulty-dependent behavior, the paper reports accuracy within each of the five difficulty quintiles separately. For the FLOPs-matched comparison, results are also reported as relative improvement/disadvantage compared to the 14× larger model with greedy decoding.
Baselines. The paper compares against several baselines across its experiments:
- Majority voting: select the most common final answer among N sampled solutions without any learned verifier.
- ORM best-of-N weighted: score N solutions with an outcome reward model (trained on complete-solution correctness) and apply best-of-N weighted selection (aggregating scores for the same final answer and selecting the answer with the highest total).
- PRM best-of-N weighted: score N solutions with the process reward model and apply best-of-N weighted selection.
- Parallel sampling (for revisions): generate N independent solutions from the revision model and select the best via verifier or majority voting—this serves as the baseline against which sequential and hybrid sequential/parallel strategies are compared.
- For the FLOPs-matched comparison, the baseline is the ~14× larger model with greedy decoding (no test-time augmentation of its own).
Generation budget / compute accounting. The universal unit of test-time compute is one "generation," defined as one complete sampled solution from the base LLM. For best-of-N and standard beam search, the budget equals the number of samples or beams N. For lookahead search with k lookahead steps, the cost is N × (k+1) to account for the additional rollout computation at each beam expansion (Section 5.3). Budgets are swept across powers of 2, typically from 2⁰ to 2⁹ (1 to 512 generations). For the revision experiments, the budget is the total number of solution chains × the length of each chain—a strategy producing √N parallel chains each of length √N consumes N total generations. The difficulty estimation cost (2048 samples per question) is explicitly not included in any budget calculations—a significant caveat the authors acknowledge (Section 3.2).
Cross-validation / statistical protocol. To avoid contaminating strategy selection with test-set performance when evaluating compute-optimal policies, the authors use two-fold cross-validation within each difficulty bin on the 500-question test set. The best-performing strategy is selected on one fold and evaluated on the other, with results averaged (Section 3.2). Difficulty bins are constructed using model-specific pass@1 rates (oracle) or PRM final-answer score distributions (predicted), dividing the test set into five quintiles of approximately 100 questions each. With cross-validation, strategy selection is based on approximately 50 questions per fold per bin.

Main Quantitative Results

Search Against PRM Verifiers (Section 5)

Aggregate comparison across search algorithms. Figure 3 (left) compares best-of-N weighted, beam search (M=4 and M=√N), lookahead search (k=1 and k=3), and majority voting across generation budgets from 1 to 256 (maximum budget of 256 generations for search experiments, per Section 5.3):

At low budgets (2–8 generations), beam search (M=4) significantly outperforms best-of-N weighted. At 4 generations, beam search achieves approximately 27% accuracy versus roughly 16% for best-of-N weighted—an 11 percentage point gap.
At high budgets (64–256), the advantage diminishes or reverses. Best-of-N weighted reaches approximately 38% at 512 generations; beam search (M=4) plateaus around 34%. Beam search (M=√N) performs similarly to M=4.
Lookahead search generally underperforms all other methods at the same generation budget due to its higher per-step cost. The 3-step lookahead variants (both M=4 and M=√N) converge to similar performance as other methods at very high budgets but never surpass them. Lookahead with k=1 on M=√N also trails the non-lookahead beam search variants.
Majority voting substantially trails all verifier-based methods, reaching only about 29% at 512 generations.

The key aggregate finding is non-monotonic: beam search provides the strongest optimization but its advantage diminishes at high budgets, while the simpler best-of-N weighted continues to improve and eventually catches up or surpasses it.

Difficulty-bin analysis for search. Figure 3 (right) breaks down beam search (M=4) vs. best-of-N weighted by difficulty quintile at four budget levels (4, 16, 64, 256 generations), revealing the core difficulty-dependent patterns:

Bin 1 (easiest, highest base-model pass@1): Beam search accuracy actually decreases from roughly 78% at 4 generations to roughly 77% at 256 generations, while best-of-N weighted increases from 68% to 88%. This is the clearest evidence of PRM over-optimization—beam search finds solutions that score well under the verifier but are incorrect, and the effect worsens with more search budget.
Bin 2: Beam search improves modestly (roughly 14% → 32%) but best-of-N weighted improves faster (roughly 14% → 60%), maintaining a clear advantage at all budgets.
Bin 3 (medium difficulty): Beam search consistently outperforms best-of-N weighted across all budgets, reaching roughly 34% vs. 23% at 256 generations. This is where PRM-guided search provides genuine benefit—the verifier helps navigate toward correct solutions the model wouldn't find by random sampling.
Bin 4: Beam search shows the strongest relative advantage, reaching roughly 17% vs. 10% for best-of-N at 256 generations.
Bin 5 (hardest): Both methods hover near 1–3% regardless of budget. No method makes meaningful progress on problems fundamentally outside the base model's capability reach.

Compute-optimal search results. Figure 4 shows the performance of selecting the best search strategy per difficulty bin at each budget level, using both oracle and predicted difficulty bins:

At 16 generations, compute-optimal (oracle bins) achieves approximately 27% accuracy, roughly matching PRM best-of-N weighted at 64 generations—a ~4× compute reduction.
At 256 generations, compute-optimal oracle reaches approximately 39.5%, surpassing PRM best-of-N weighted at the same budget (roughly 37%).
Compute-optimal with predicted (non-oracle) difficulty bins tracks the oracle version closely, particularly at lower budgets. The two curves "largely overlap" per the authors, with the predicted version reaching approximately 37% at 256 generations—a ~2 percentage point gap from oracle.
Both compute-optimal variants consistently outperform ORM best-of-N weighted (which peaks around 34% at 512 generations) and majority voting (around 29% at 512 generations).

PRM vs. ORM comparison. Figure 14 (Appendix F) shows that PRM best-of-N weighted scales better than ORM best-of-N weighted. At 2048 samples, PRM achieves approximately 40% accuracy versus roughly 35% for ORM and roughly 30% for majority voting. The gap between PRM and ORM widens with the number of samples, confirming that step-level training provides beneficial representation learning even when final-step aggregation reduces the PRM to ORM-like behavior at selection time.

Revision Model Results (Section 6)

Revision model pass@1 trajectory. Figure 6 (left) tracks the revision model's per-step pass@1 across a chain of up to 64 revisions. Starting from approximately 18.2% pass@1 at step 1 (the initial answer), accuracy improves to roughly 24–25% by steps 15–20 and remains in the 23–25% range out to 64 steps. The model generalizes beyond its 4-step training horizon—the paper notes this as "evidence that the model has learned a generalizable revision skill" (Section 6.1).

Sequential vs. parallel comparison. Figure 6 (right) compares fully sequential revisions against fully parallel sampling at 64 generations, under both verifier-based (best-of-N weighted) and majority-based answer selection:

Sequential + best-of-N weighted: approximately 41.5%
Parallel + best-of-N weighted: approximately 39%
Sequential + majority: approximately 38%
Parallel + majority: approximately 35%

Sequential outperforms parallel under both selection mechanisms. The verifier-based gap (roughly 2.5 percentage points) is narrower than the majority-based gap (roughly 3 percentage points), suggesting the verifier partially compensates for the lack of sequential refinement. However, the fact that sequential revisions benefit majority voting as well indicates the improvement is not purely an artifact of the verifier seeing revision history—the revision model genuinely produces better answers over successive steps.

Sequential-to-parallel ratio sweep. Figure 7 (left) sweeps the sequential-to-parallel ratio from fully parallel (leftmost) to fully sequential (rightmost) for a fixed total generation budget at multiple budget levels (8, 16, 32, 64, 128, 256):

At 256 generations, the optimal ratio is around 2¹ to 2³ (2:1 to 8:1 sequential-to-parallel), achieving approximately 43–44% accuracy.
Fully parallel (leftmost) yields approximately 40%.
Fully sequential (rightmost) yields approximately 42%.
At lower budgets (8–32 generations), the curves are monotonically increasing with the sequential-to-parallel ratio—fully sequential is optimal, meaning that when the total compute is limited, the model benefits most from iterative refinement of a single chain rather than exploring multiple chains.

Difficulty-dependent optimal ratio. Figure 7 (right) breaks down the sequential-to-parallel sweep at a fixed budget of 128 generations by difficulty bin:

Bin 1: Performance is essentially flat across all ratios, around 90–92%. Easy questions are insensitive to allocation strategy—the model already solves most of them correctly.
Bin 2: Slight but clear advantage for higher sequential ratios, approximately 63% at fully sequential vs. 58% at fully parallel.
Bin 3: A clear optimal ratio emerges at moderate sequential-to-parallel values (around 2¹ to 2³), reaching approximately 42% vs. 35% at both extremes. This is where the complementary benefits of exploration (parallel sampling) and exploitation (sequential refinement) are most apparent.
Bin 4: Similar pattern, with the peak at a moderate ratio achieving roughly 18% vs. 14% at fully parallel.
Bin 5: All ratios produce roughly 2–3% accuracy. No allocation strategy helps on problems the base model fundamentally cannot solve.

Compute-optimal revisions. Figure 8 shows performance when the optimal sequential-to-parallel ratio is selected per difficulty bin (both oracle and predicted bins) compared against parallel best-of-N weighted and parallel-only baselines:

At 64 generations, compute-optimal oracle achieves approximately 40%, matching parallel best-of-N weighted at 256 generations—a ~4× compute reduction.
At 256 generations, compute-optimal oracle reaches approximately 44%, compared to roughly 41% for parallel best-of-N weighted and 37% for parallel-only.
Compute-optimal with predicted bins performs slightly below oracle at high budgets (approximately 41% at 256 generations vs. 44% for oracle) but still substantially outperforms the parallel baseline (37%).
A notable trend: the parallel baseline appears to plateau around 36–37% at high budgets, while compute-optimal scaling continues to improve, suggesting the gains from adaptive allocation compound at higher budgets rather than saturating.

Verifier transfer across distributions. Appendix J (Figure 15a) shows that the base-model PRM underperforms a revision-specific ORM when scoring revision model outputs. At 64 generations, sequential revisions with the base-model PRM achieve roughly 40%, while sequential revisions with the revision-specific ORM achieve roughly 42%. This confirms distribution shift as a practical concern: verifiers trained on one model's outputs do not transfer optimally to another model's outputs.

Revision history in verifier context. Appendix J (Figure 15b) ablates whether the revision verifier benefits from seeing previous revisions in its context. Including revision history provides a small improvement over no-history (approximately 1–2 percentage points at 64 generations), but both variants substantially outperform the parallel baseline. This demonstrates that the sequential revision benefit is not solely an artifact of the verifier having access to more context—the revision model genuinely produces progressively better answers.

FLOPs-Matched Comparison: Test-Time vs. Pretraining Compute (Section 7)

The FLOPs-matched comparison asks: given a fixed total FLOPs budget (pretraining + inference), is it better to train a larger model or to keep the smaller model and spend the extra FLOPs on inference-time computation using compute-optimal strategies? The comparison uses three values of the ratio R = D_inference / D_pretrain: 0.16 (R ≪ 1), 0.79 (R ≈ 1), and 22 (R ≫ 1).

Revisions vs. pretraining (Figure 9, left; Figure 1, top-right bar chart). Comparing PaLM 2-S* with compute-optimal revisions against the ~14× larger model:

Easy questions (bin 1 in Figure 9, "Easy" bar in Figure 1): Test-time compute outperforms pretraining at all R values, with the margin narrowing as R increases: +11.8% at R=0.16, +3.5% at R=0.79, and −11.9% at R=22 (the negative value at R≫1 indicating the larger model is preferred).
Medium questions (bins 2–3 in Figure 9, "Medium" bar in Figure 1): Strong advantages at low R (+27.8% at R=0.16), moderate at R=0.79 (+16.7%), and still positive at R=22 (+5.4%).
Hard questions (bins 4–5 in Figure 9, "Hard" bar in Figure 1): Mixed: advantage at R≪1 (+21.6%), but implied negative at R≈1, and a large disadvantage at R≫1 (−37.2%).

Across all difficulty levels, revisions provide test-time compute advantages at R≪1, and the advantage persists for easy and medium questions across all R values tested. Only on hard questions at high R does pretraining clearly dominate.

PRM search vs. pretraining (Figure 9, right; Figure 1, bottom-right bar chart). The results are starkly different and less favorable to test-time compute:

Easy questions: +19.1% at R=0.16, +2.2% at R=0.79, +2.0% at R=22. Test-time compute with PRM search maintains an advantage on easy questions across all R values, though the margin narrows substantially from low to high R.
Medium questions: 0.0% (parity) at R=0.16, −35.3% at R=0.79, −30.8% at R=22. Test-time compute with PRM search is at best equal to the larger model and at worst substantially worse on medium-difficulty questions.
Hard questions: −3.6% at R=0.16, −35.3% at R=0.79, −52.9% at R=22. The larger model dominates on hard questions across all R values, with the disadvantage growing substantially at higher R.

The asymmetry between revisions and PRM search is a critical finding: revisions provide stronger FLOPs-matched benefits than search, particularly on medium-difficulty questions. The paper does not fully explain this asymmetry, but it likely reflects the difference between improving the proposal distribution (revisions generate better candidates, helping across difficulty levels) versus selecting among existing candidates (PRM search is limited by the base model's capacity to produce any correct solution, which is near zero for the hardest problems).

Figure 9 detail. The line plots show accuracy per difficulty bin as test-time compute scales from 1 to 256 generations for revisions (left) and 1 to 512 for search (right). The ~14× larger model's greedy performance is shown as horizontal dashed lines (stars) placed at three x-axis positions corresponding to how much test-time compute the smaller model is entitled to at each R value. Where the compute-optimal scaling line is above the star, test-time compute wins. For revisions on bin 1 (purple, topmost line), the scaling line is above all three stars, confirming test-time compute advantage across all R. For revisions on bin 5 (blue, bottommost line), the scaling line is below all three stars and essentially flat near 0–5%, confirming that no amount of test-time compute helps on the hardest problems regardless of R.

Large-scale model results (Section 5.3). Training throughput for 64B and 136B Decoder-only Transformer models over two islands of accelerators (1024 cores for 136B, 512 cores for 64B) shows PATHWAYS achieves ~97% of the throughput compared to a single island with twice as many devices. Cross-island DCN transfers for global gradient reduction (1030 GB for 136B, 457 GB for 64B) incur minimal overhead. The execution trace (Figure 12, Appendix D) shows DCN transfers are effectively overlapped with computation—the paper states "DCN transfers incur minimal overhead" (Section 5.3).

Ablation Studies and Robustness Checks

PRM aggregation strategy (Appendix E, Figure 13): Comparing "min," "prod," and "last" step-wise aggregation for the PRM reveals that "last" performs best: approximately 37% at 256 samples vs. 35% for "min" and 27% for "prod." ORM achieves approximately 34%. This is a non-obvious finding because prior work (Lightman et al., 2023; Wang et al., 2023) found "min" to be superior. The authors hypothesize that their use of soft Monte Carlo labels rather than binary correctness labels changes the distribution of per-step scores in ways that make the final-step prediction most reliable. An interesting consequence: using "last" aggregation effectively reduces the PRM to ORM-like behavior at selection time, yet the PRM still outperforms a separately trained ORM (Figure 14), suggesting that the step-level training acts as beneficial representation learning.

PRM vs. ORM scaling (Appendix F, Figure 14): The PRM consistently outperforms the ORM, with the performance gap widening at higher sample counts: at 2048 samples, PRM best-of-N weighted reaches approximately 40% vs. ORM's 35% and majority voting's 30%. This validates the design choice to use process-level supervision even though final-step aggregation is used at selection time.

Parallel vs. sequential asynchronous dispatch (Figure 7): Forcing sequential asynchronous dispatch in a pipelined benchmark (where computations are chained across hosts and must wait for predecessor enqueue before preparing the successor) substantially reduces throughput compared to parallel dispatch. The benefit is most pronounced as the number of pipeline stages increases—parallel dispatch amortizes the fixed scheduling overhead that would otherwise accumulate per stage. This ablation validates that parallel async dispatch, rather than any other architectural feature, is the mechanism responsible for closing the coordination latency gap with multi-controller systems.

Oracle vs. predicted difficulty bins (Figures 4, 8, and Appendix C, Figures 11–12): Both oracle and predicted bins yield qualitatively similar trends across difficulty levels in both search and revision settings. For search (Figure 4), the two curves "largely overlap," with predicted bins closely tracking oracle bins across all budgets. For revisions (Figure 8), predicted bins show slightly lower performance at high budgets (approximately 41% vs. 44% at 256 generations) but maintain the same qualitative advantage over the parallel baseline. This is the critical robustness check for deployability: the compute-optimal framework works without ground-truth labels, using the PRM's own score distribution as a difficulty proxy.

Revision model verifier choice (Appendix J, Figure 15a): The base-model PRM underperforms the revision-specific ORM when scoring revision model outputs by approximately 2 percentage points at 64 generations, confirming that verifier distribution shift is a measurable practical concern—not just a theoretical one.

Revision history in verifier context (Appendix J, Figure 15b): Including previous revisions in the ORM's context provides a small benefit (approximately 1–2 percentage points at 64 generations) over the no-history ablation. Critically, both variants substantially outperform the parallel baseline, demonstrating that the sequential revision benefit is not merely an artifact of the verifier having access to more context for answer selection.

Majority voting for revisions (Appendix B, Figure 10): The sequential-to-parallel ratio trends observed with verifier-based selection are replicated when using majority voting: easy questions are insensitive to ratio, hard questions show an optimal intermediate ratio, and fully sequential marginally outperforms fully parallel in aggregate. This replicates the core finding without reliance on a learned verifier, strengthening the claim that sequential revisions genuinely improve answer quality.

ReSTᵉᵐ revision model (Appendix K, Figure 16): An attempt to further optimize the revision model using ReSTᵉᵐ (Singh et al., 2024) produces a substantial negative result: additional sequential revisions with the ReSTᵉᵐ-trained model severely degrade performance. At 256 generations, fully sequential performance drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio—the model performs worse with more revisions, the opposite of the intended effect. The authors hypothesize that on-policy data collection in ReSTᵉᵐ exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly. This is an important negative result that highlights the sensitivity of revision training to the data generation procedure and serves as a caution that not all self-improvement approaches are beneficial.

Multi-tenancy throughput (Figure 8): PATHWAYS achieves at least the same aggregate throughput as JAX when multiple clients concurrently submit different programs, across a range of per-program compute times (0.04 ms to 2.4 ms). This demonstrates zero overhead to context-switch between programs from different clients when resources concurrently fit in HBM. Notably, for very small computations (0.04 ms), PATHWAYS' maximum throughput exceeds JAX's, achieving higher TPU utilization—the paper attributes this to PATHWAYS workers being able to accept more computations from remote clients than JAX can dispatch using Python locally.

Proportional-share fairness (Figure 9): The per-island scheduler can enforce proportional-share ratios between concurrent clients (demonstrated with 1:1:1:1 and 1:2:4:8 ratios among 4 clients). Traces show interleaving of gang-scheduled computations at sub-millisecond granularity while maintaining the specified share ratios. This validates that centralized gang-scheduling is compatible with fine-grained multi-tenancy and fairness policies.

Pipeline scaling efficiency (Table 2): Increasing the number of pipeline stages from 4 to 16 for a 3B Transformer model reduces throughput only marginally, from 133.7k to 131.4k tokens per second (~1.7% reduction). This validates that PATHWAYS' parallel asynchronous dispatch effectively amortizes coordination overhead across pipeline stages—the reduction would be substantially larger if coordination latency accumulated per stage. The pipelined model is competitive with SPMD (125.7k tokens/s), and in this instance the pipeline outperforms SPMD because "collective communication within the SPMD computation incurs higher overhead than pipeline bubble overhead" (Section 5.3).

Multi-island training overhead (Section 5.3): PATHWAYS achieves "the same throughput (131.4k tokens/sec) using a single island of 128 cores on configuration (B), or 4 islands of 32 cores each on configuration (C)." The execution trace (Figure 10) shows DCN transfers occurring between pipeline stages assigned to different islands, with communication time "effectively overlapped with computation."

Critical Assessment

This section evaluates whether the reported experiments genuinely support the paper's central claims, identifying where the evidence is strong, where it is qualified or conditional, and where it falls short of what would be needed to fully substantiate the claims.

Claim: "Compute-optimal scaling improves efficiency by more than 4× over best-of-N." The evidence for this claim is Figures 4 and 8, which show that compute-optimal strategies achieve comparable accuracy with ~4× fewer generations than best-of-N weighted (e.g., 16 generations matching 64 in Figure 4; 64 generations matching 256 in Figure 8). The claim is supported with the important qualification that the difficulty estimation cost—2048 samples per question to determine the difficulty quintile—is not included in the generation budget. The authors acknowledge this explicitly (Section 3.2: "our experiments do not account for this cost largely for simplicity"). In a realistic deployment where difficulty estimation is amortized over one-time batch evaluation (e.g., processing thousands of problems where the 2048-sample estimate is computed once per problem), the amortized cost may be negligible. In a single-query setting, however, the cost would dominate and the 4× figure would not hold. The paper's "predicted difficulty" variant (using PRM scores rather than ground-truth labels) brings the approach closer to deployability but does not reduce the sample cost. The reported 4× gain is therefore best understood as an upper bound on achievable efficiency contingent on amortized or cheaper difficulty estimation. A critical missing experiment is evaluating how accuracy degrades when using a drastically smaller number of samples (e.g., 4, 8, or 16) for difficulty estimation—this would directly quantify the exploration-exploitation tradeoff the paper identifies as future work.

Claim: "Test-time compute with a smaller model can outperform a ~14× larger model." The evidence in Figures 9 and the Figure 1 bar charts supports this claim with sharply qualified conditions. The claim holds convincingly for easy-to-medium difficulty questions at low inference-to-pretraining ratios (R ≪ 1), where revisions show +27.8% relative improvement on medium questions and search shows +19.1% on easy questions. However, the claim weakens or reverses as difficulty increases and as R grows. On hard questions (bins 4–5) at R ≫ 1, PRM search underperforms the larger model by −52.9%. On medium questions at R ≫ 1, PRM search underperforms by −30.8%. Even revisions, the stronger of the two methods, show −37.2% on hard questions at R ≫ 1.

Several experimental design choices make the comparison favorable to test-time compute in ways that are not fully acknowledged: (1) The ~14× larger model uses greedy decoding with no test-time augmentation of its own—a stronger baseline would give the larger model a modest test-time budget (e.g., best-of-8 majority voting). (2) The larger model scales parameters only (not data), following LLaMA-style training rather than Chinchilla-optimal training (Hoffmann et al., 2022) where both parameters and data are scaled. A compute-optimally trained larger model would be a stronger baseline. (3) The comparison is made on a single model family (PaLM 2 variants), and the scaling relationship between model size and MATH performance may not generalize to other model families. The claim is best restated as: on problems within the base model's capability reach, and when inference volume is low relative to pretraining, test-time compute can substitute for model scale, but this substitution breaks down on hard problems and at high inference-to-pretraining ratios.

Claim: "Efficacy depends critically on prompt difficulty." This is the most robustly supported claim in the paper, replicated across search methods (Figure 3, right), revision strategies (Figure 7, right), selection mechanisms (majority voting in Appendix B, verifier-based in Figure 7), and FLOPs-matched comparisons (Figure 9). The difficulty-dependent behavior is qualitatively different—and sometimes opposite—for the same strategy at different difficulty levels: beam search hurts easy-problem performance at high budgets (Figure 3, right, bin 1) while helping on medium problems (bin 3); fully sequential revisions are optimal on easy problems while balanced sequential-parallel ratios are optimal on hard problems (Figure 7, right). This non-monotonic relationship is strong evidence that difficulty is a genuine moderating variable, not a confound.

The primary limitation is that all experiments use a single benchmark (MATH) and a single model family (PaLM 2-S*). The paper's claim that PaLM 2-S* is "representative of the capabilities of many contemporary LLMs" (Section 4) is asserted but not tested. Different model families may exhibit different difficulty-dependent scaling curves—a model with better calibration might show less over-optimization, while a model with different error patterns might benefit differently from beam search vs. revisions. The binary classification of MATH problems as mathematically solvable via string matching also limits generality: extending the difficulty-dependent scaling framework to open-ended generation, dialogue, or creative tasks—where correctness is ambiguous—would require fundamentally different verifier training and evaluation protocols that the paper does not address.

Claim: "Beam search over-optimizes the PRM, causing performance degradation at high budgets on easy problems." The evidence in Figure 3 (right, bin 1) shows beam search accuracy decreasing slightly from 78% at 4 generations to 77% at 256 generations while best-of-N weighted increases from 68% to 88%. This is clear qualitative evidence of over-optimization, and qualitative examples in Appendix M (Figures 29, etc.) show degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM. However, the quantitative effect is relatively small—a ~1 percentage point degradation—on a subsample of approximately 100 questions in bin 1. The paper does not report confidence intervals, and with such a small sample the degradation may not be statistically significant. The more compelling evidence for over-optimization is the relative performance: beam search flatlines while best-of-N continues to improve, meaning the optimization pressure from beam search prevents the system from capturing gains that are available with simpler methods. A missing experiment is testing whether ensembling multiple independently trained PRMs reduces over-optimization, which would directly test the paper's claim that "improving verifier robustness is the key bottleneck."

Missing experiments that would strengthen the paper. Several experiments would substantially increase confidence in the paper's conclusions: (1) Replication on a non-MATH benchmark (e.g., GSM8K, MBPP, or a code generation task) to test whether difficulty-dependent scaling patterns generalize beyond competition math. (2) Testing on a model from a different family (e.g., a LLaMA variant) to assess whether the specific difficulty thresholds and optimal strategies transfer. (3) A combined PRM-search + revisions experiment (the paper acknowledges this was not done in Section 8) to assess whether the complementary strengths of the two approaches yield additive gains. (4) A full cost-benefit analysis that amortizes difficulty estimation over varying numbers of queries per problem. (5) Comparison against a compute-optimally trained larger model (scaling both parameters and data) rather than a parameter-only-scaled model, to strengthen the FLOPs-matched claims. (6) Latency-aware evaluation that accounts for the wall-clock implications of sequential revisions (which are inherently serial) vs. parallel best-of-N (which can be batched). (7) Ablation on the number of difficulty bins (e.g., 3, 5, 7, 10) to assess whether the coarse binning loses information and whether finer-grained adaptation would improve results further.

6. Limitations and Trade-offs

Difficulty Estimation Cost Is Excluded from Headline Efficiency Gains

The assumption or constraint. The compute-optimal scaling framework relies on estimating each prompt's difficulty before allocating the inference budget. The paper's method for doing so—generating 2048 samples per question and computing either ground-truth pass@1 (oracle) or the PRM's average final-answer score (predicted), then binning into quintiles—consumes more compute than the largest test-time budgets studied (256–512 generations). The authors acknowledge this explicitly:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity" (Section 3.2)

The consequence. The headline claim of a "more than 4× better efficiency over a standard best-of-N baseline" is computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment where difficulty is estimated per problem, the total cost would be (2048 samples for estimation) + (strategy execution budget), which could exceed the cost of simply running best-of-N with a larger budget—undermining the claimed efficiency gain. The 4× figure is an upper bound contingent on difficulty being known nearly for free.

For a single-query setting (one problem, one answer needed), the estimation cost completely dominates: you would spend 2048 generations just to determine the difficulty quintile, then run perhaps 16–64 generations of the selected strategy, for a total cost of 2080–2112 generations—far worse than running best-of-256 for 256 generations. Only in a batch setting where difficulty is estimated once and amortized over many repeated queries on the same problem (e.g., evaluating a known test set, or generating synthetic training data for the same fixed questions) can the estimation cost become negligible.

What evidence exists in the paper. The gap between oracle and predicted difficulty curves in Figures 4 and 8 demonstrates that the approach works without ground-truth labels (using the PRM's score distribution as a proxy), but neither variant addresses the sample cost of estimation. The paper does not report any experiment that varies the number of difficulty estimation samples or measures how the compute-optimal strategy degrades when drastically fewer samples are used (e.g., 4, 8, or 16 instead of 2048). This is the critical missing experiment: the exploration-exploitation tradeoff between estimation and solving is identified as future work (Section 3.2) but never quantified.

Mitigation status. The paper explicitly identifies this as a limitation and suggests "future work on training models to predict difficulty directly from the question text" (Section 8), or adaptive schemes that amortize estimation into the solution process. No such mechanism is developed or evaluated in the current work. Without this mitigation, the compute-optimal framework as presented is not directly deployable in its claimed form—the efficiency gains exist in the analysis but cannot be realized without solving the estimation-cost problem.

The Method Fails on the Hardest Problems (Difficulty Bin 5)

The assumption or constraint. The paper's approach amplifies the base model's existing capabilities through better search and revision strategies, but does not create capability where none exists. If the base model's pass@1 on a problem class is near zero—meaning it essentially never produces a correct solution in its output distribution—then no amount of test-time search or revision can find or refine a correct answer. The paper is transparent about this:

"On the hardest questions (bin 5), no method makes meaningful progress—the base model simply lacks the capability to produce correct solutions regardless of how the budget is allocated." (Section 5.2)

The consequence. The approach offers no path forward for genuinely novel or out-of-distribution reasoning that exceeds the base model's training distribution. Across all methods—search, revisions, and their compute-optimal combinations—bin 5 accuracy hovers at 1–3% regardless of budget (Figure 3, right; Figure 7, right). In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for both revisions and search, meaning that test-time compute provides zero marginal return on these problems. Scaling pretraining (the ~14× larger model) is the only effective intervention—the larger model achieves non-negligible accuracy on hard problems where test-time compute adds nothing.

This creates a fundamental capability ceiling: test-time compute can amplify capability but cannot create it. For problems the base model cannot begin to solve, no allocation strategy, verifier quality improvement, or revision depth will help. This boundary condition sharply limits the scope of the approach: it applies to problems within the model's "rough capability range" (problems where pass@1 is non-trivially above zero) and offers no benefit for problems outside that range.

What evidence exists in the paper. The evidence is consistent and stark across every experiment that disaggregates by difficulty:

Figure 3 (right): Bin 5 shows 1–3% accuracy for all search methods at all budgets (4–256 generations).
Figure 7 (right): Bin 5 shows ~2–3% accuracy for all sequential-to-parallel ratios at 128 generations.
Figure 9: Bin 5 compute-optimal scaling lines are flat near 0–5% for both revisions and search. The ~14× larger model's greedy performance on bin 5 exceeds what any amount of test-time compute achieves with the smaller model.
The FLOPs-matched bar chart (Figure 1, bottom-right): Hard questions show −52.9% relative disadvantage for PRM search at R ≫ 1.

Mitigation status. The paper candidly acknowledges this limitation but offers no mitigation within the current framework. The authors state in Section 7 that "test-time compute amplifies existing capability but does not create it from nothing." Future work on combined pretraining + inference optimization might push the capability boundary, but the fundamental limitation remains: for problems where the base model's pass@1 is near zero, pretraining is the only viable path. This is not a bug to be fixed but rather a structural property of any test-time compute approach—optimization can only select among or refine what the base model can produce.

The Method Combines Test-Time Strategies Only with a Weak Pretraining Baseline

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time strategies against a model with approximately 14× more parameters, but this larger model is (a) trained by scaling parameters only (not data), departing from compute-optimal pretraining (Hoffmann et al., 2022), and (b) evaluated with greedy decoding only—no test-time augmentation of its own. The paper acknowledges the first point:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)

The consequence. The comparison is favorable to test-time compute in ways that are not fully accounted for. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data according to the optimal ratio) would likely outperform a parameter-only-scaled model, making the pretraining baseline stronger. More importantly, giving the larger model even a modest test-time compute budget—best-of-8 majority voting, or a simple best-of-N with an ORM—would substantially improve its performance and narrow or reverse the reported advantages of the smaller model with compute-optimal test-time strategies.

The paper's key quantitative claim—that "a smaller model augmented with compute-optimal test-time strategies can outperform a ~14× larger pretrained model on easy-to-medium difficulty problems"—is contingent on the specific weakness of the larger model baseline. Against a strong baseline (compute-optimally trained + modest test-time augmentation), the advantage would likely shrink, and the crossover point where pretraining becomes preferable (currently at R ≈ 1 for PRM search on medium questions) would shift leftward.

What evidence exists in the paper. The paper reports the 14× comparison across three R values and five difficulty bins (Figures 1 and 9), showing advantages at low R and on easy problems. However, the paper does NOT include:

A comparison against a compute-optimally trained larger model (scaling data with parameters)
A comparison where the larger model receives any test-time compute budget
Training curves or validation loss for the larger model that would allow readers to assess whether it is undertrained relative to compute-optimal scaling
Any experiment that varies the pretraining-data-to-parameters ratio for the larger baseline to test sensitivity

Mitigation status. The paper acknowledges the parameter-only scaling choice (Section 7) and frames it as "representative of a canonical approach" (LLaMA-style training, Touvron et al., 2023), leaving the compute-optimal pretraining comparison to future work. This is a reasonable scoping decision—comparing against the most widely-used pretraining paradigm—but it means the reported test-time-compute advantages are not evidence that test-time compute dominates compute-optimal pretraining; they are evidence that it dominates a specific (suboptimal) pretraining recipe. For practitioners deciding where to invest compute, this distinction is material.

The Method Has a 38% Correct-to-Incorrect Reversion Rate That Is Never Solved

The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target answer. At inference time, the model may encounter correct answers in its own revision chain (produced during earlier revision steps) and must decide what to do with them. Because the training data provides no signal for what to do when the current answer is already correct, the model often "revises" correct answers into incorrect ones. The paper quantifies this failure mode:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach" (Section 6.1)

The consequence. The revision chain is inherently unstable: as the model iterates, it may produce a correct answer at step 3, "revise" it into an incorrect answer at step 4, produce another correct answer at step 7, and so on. This means that simply taking the last revision (the natural endpoint of a sequential chain) is unreliable—the answer quality oscillates rather than monotonically improving. The paper mitigates this with post-hoc selection mechanisms (majority voting or verifier-based selection across the entire chain of revisions, picking the best answer from any point), but these are band-aid solutions that do not address the root cause: the model does not know when to stop revising.

This reversion phenomenon creates several downstream problems:

Wasted compute: Late-stage revisions that corrupt earlier correct answers consume generation budget without improving final accuracy. The compute spent on revisions after the first correct answer appears is partially wasted.
Unreliable chain length: There is no principled way to decide when to stop revising. The paper sweeps chain lengths empirically, but the optimal length depends on when (and whether) the chain hits a correct answer—which is unknown at runtime without ground truth.
Fragility to training data: The 38% reversion rate is a direct consequence of the training data construction (only incorrect-to-correct trajectories). A training procedure that included correct-to-correct trajectories (teaching the model to recognize and preserve correct answers) would likely reduce the reversion rate, but the paper does not explore this.

What evidence exists in the paper. The 38% figure is reported in Section 6.1 based on analysis of revision chains. The revision model's per-step pass@1 trajectory (Figure 6, left) shows accuracy improving from ~18% to ~24–25% by steps 15–20, then fluctuating in the 23–25% range out to 64 steps—consistent with oscillation rather than monotonic improvement. The comparison between sequential and parallel sampling (Figure 6, right) shows sequential outperforming parallel, but the gap (roughly 2.5 percentage points at 64 generations with verifier-based selection) is modest given that sequential revisions cost the same budget but produce correlated (non-independent) samples. The within-chain selection mechanism (picking the best answer from any step) is what prevents the reversion from making sequential revisions worse than parallel sampling, but the paper does not ablate what performance would be without this mitigation (i.e., always taking the last revision).

Mitigation status. The paper uses majority voting or verifier-based selection across the chain to mitigate the reversion problem, but does not attempt to solve it at the model or training level. A "more principled solution—such as training the model to recognize when no revision is needed—is not explored" (the paper does not discuss this). The ReSTᵉᵐ experiment (Appendix K, Figure 16) shows that an alternative training approach makes the problem worse: the ReSTᵉᵐ-trained model shows severe degradation with sequential revisions (fully sequential drops to ~33.5% at 256 generations vs. ~38.5% at the optimal ratio), likely because on-policy data collection amplifies the reversion pathology. This negative result suggests that the revision approach is sensitive to training methodology in ways that are not fully understood, and that naive attempts to improve the revision model can backfire.

All Results Are on a Single Benchmark with a Single Model Family

The assumption or constraint. Every experiment in the paper uses the MATH benchmark (500 test questions, competition-level math) with variants of PaLM 2-S* as the base model. The authors assert—but do not test—that this model is "representative of the capabilities of many contemporary LLMs" (Section 4). The paper does not evaluate on any other reasoning benchmark (e.g., GSM8K, MBPP, HumanEval, ARC, FOLIO), any non-math domain (code, logic, scientific QA), or any model from a different family (e.g., GPT, LLaMA, and certainly not any open-weight model that would have been available at the time).

The consequence. The paper's core findings—the difficulty-dependent optimal strategies, the 4× efficiency gain from compute-optimal allocation, the over-optimization behavior of beam search, the complementary strengths of revisions and search—may be specific to the mathematical reasoning domain, the PaLM 2 architecture, or their interaction. Several aspects of the results could be model-specific or domain-specific:

PRM quality and over-optimization behavior: The PRM is trained on PaLM 2-S*'s output distribution. A model with different calibration properties, different error patterns (e.g., syntax errors vs. logical errors), or different sensitivity to prompt formatting might exhibit different over-optimization thresholds. The beam-search degradation on easy problems (Figure 3, right, bin 1) might be larger or smaller depending on the base model's typical failure modes.
Revision model effectiveness: The revision model's ability to learn from in-context incorrect examples depends on the base model's in-context learning capabilities and its error patterns. Different model families have different in-context learning behaviors and might benefit more or less from edit-distance-based training data construction.
Difficulty distribution: The MATH benchmark has a specific difficulty distribution (competition problems ranging from easy AMC to hard AIME). The five difficulty quintiles and their associated optimal strategies are calibrated to this distribution. A different benchmark with a different difficulty spectrum might shift the thresholds at which beam search becomes preferable to best-of-N, or at which sequential revisions become preferable to parallel sampling.
Verifiability of correctness: MATH problems have ground-truth answers that can be checked with exact string matching (via the grading function from Lightman et al., 2022). This enables clean PRM training (via Monte Carlo rollout correctness) and difficulty estimation (via pass@1). Many important real-world tasks—code generation where correctness is functional but not exact-match, open-ended generation, multi-step planning with ambiguous success criteria—lack such clean signals. Extending the compute-optimal framework to these domains would require different verifier training and evaluation protocols.

What evidence exists in the paper. None that would support generalization. The paper presents no ablation across benchmarks, no evaluation on non-math tasks, and no test on a non-PaLM model. The claim of representativeness is an assertion, not a finding. The paper does not discuss how domain-specific properties of MATH (symbolic reasoning, exact answer matching, multi-step logical deduction rather than factual recall) might interact with the test-time strategies.

Mitigation status. The paper acknowledges the single-benchmark limitation only implicitly—there is no explicit caveat about generalizability in the limitations or conclusions sections. The authors do note that their choice of MATH is deliberate because "test-time compute is expected to help most when the model already possesses the necessary knowledge and the challenge is drawing complex inferences" (Section 4)—suggesting they view their results as applying to reasoning tasks generally, not just math. However, whether domains like code generation (where execution feedback provides a different kind of verifier signal) or logical reasoning (where the reasoning steps have different structure) exhibit similar difficulty-dependent scaling patterns is an open question. The paper leaves this entirely to future work.

The Method Never Combines PRM Search with Revisions

The assumption or constraint. The paper studies two complementary mechanisms for improving test-time performance—PRM-guided search (modifying the verifier to better select among candidates) and iterative revisions (modifying the proposal distribution to generate better candidates)—but studies them entirely independently. The compute-optimal policy selects between search strategies for one set of experiments (Section 5) and between sequential-to-parallel ratios for another (Section 6), but never combines PRM tree-search with the revision model as the proposal distribution. The paper explicitly acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions" (Section 8)

The consequence. The paper's results represent a lower bound on what a fully integrated system could achieve. The two mechanisms have complementary strengths that the difficulty-bin analysis reveals: PRM search helps most on medium-difficulty problems where guided exploration is valuable (Figure 3, right, bins 3–4), while revisions help most on easy-to-medium problems where initial answers are roughly correct and benefit from refinement (Figure 7, right, bins 1–3). Combining them—for example, using the revision model as the proposal distribution within beam search, or using the PRM to guide which revision branches to pursue—could yield gains beyond either method alone, particularly on medium-difficulty problems where both mechanisms show individual benefits.

Several natural combinations are not explored:

Beam search with revision-model proposals: At each step of the beam search tree, the revision model conditions on previous (potentially rejected) branches as context, potentially producing higher-quality candidate steps than the base model would generate independently.
PRM-guided revision decisions: Rather than blindly generating a long revision chain and selecting the best answer post-hoc, use the PRM's per-step scores to decide in real-time whether a revision is on track (continue the chain) or off track (restart from scratch). This could reduce the 38% correct-to-incorrect reversion rate by stopping chains before they corrupt good answers.
Hierarchical allocation: Use a few parallel beams (exploration) to generate diverse initial solution approaches, then apply sequential revisions (exploitation) within each beam to refine the best candidates. This directly combines the difficulty-dependent strengths: exploration for hard problems, exploitation for easy refinements.

What evidence exists in the paper. The paper provides no experiments, ablations, or even design sketches for combined approaches. The independent results are suggestive: revisions outperform search in the FLOPs-matched comparison (Figure 9, left vs. right), and search outperforms revisions on medium-difficulty problems (Figure 3, right, bin 3 vs. Figure 7, right, bin 3), so a combination might capture the best of both. But this is speculation without data. The paper does not report whether the revision model's outputs are scoreable by the PRM (Appendix J, Figure 15a, shows distribution shift between base-model PRM and revision-model ORM, suggesting potential incompatibility), or whether the PRM's step-level guidance would be useful when the proposal distribution changes.

Mitigation status. The paper explicitly identifies the combination as future work (Section 8), framing the current study as a systematic analysis of the two axes independently before integrating them. This is a reasonable research strategy—decompose the problem, understand each mechanism in isolation, then combine—but it means the paper's results cannot be taken as the best achievable performance. A practitioner deploying this system would want to combine search and revisions, and the paper provides no guidance on how to do so or what gains to expect. The absence of this natural combination also limits the strength of the FLOPs-matched comparison: if combined search + revisions narrow the gap with the larger model further than either mechanism alone, the paper's already-qualified claims about test-time compute substituting for pretraining would need to be revised upward for the relevant difficulty tiers.

7. Implications and Future Directions

How This Work Changes the Landscape

PATHWAYS reframes a fundamental architectural debate in distributed ML systems—single-controller versus multi-controller—from an unresolvable tradeoff into an engineering problem with specific, solvable bottlenecks. Prior to this work, the community had broadly accepted that the programming flexibility of single-controller architectures (heterogeneous computation, centralized resource management, natural expression of non-SPMD patterns) came at the cost of unacceptable performance degradation at scale, as demonstrated by TensorFlow v1's inability to compete with multi-controller JAX and PyTorch on large-scale SPMD training. PATHWAYS demonstrates that this performance gap is not inherent to the single-controller model but rather an artifact of three specific implementation deficiencies: serialized dispatch latency from remote clients, sequential host-side coordination that accumulates in pipeline chains, and naive sharded graph representations that explode to millions of edges at scale.

This is a diagnostic contribution of substantial practical consequence. By identifying the precise mechanisms that caused TF v1's failure—and showing that each admits a targeted solution (parallel asynchronous dispatch, centralized per-island gang-scheduling, and sharded dataflow via Plaque)—PATHWAYS prevents the field from erroneously discarding single-controller designs based on one flawed instantiation. The evidence is quantitative and decisive: Table 1 shows identical training throughput between PATHWAYS and multi-controller JAX across model sizes from 270M to 11B parameters (618k tokens/s for T5-Base, 84.8k tokens/s for T5-11B on both systems), and Figure 6 demonstrates that a computation of only 2.3 ms on 128 TPUs—or 35 ms on 2048 TPUs—fully masks all single-controller overhead. These are not marginal improvements; they represent complete elimination of the performance gap that had previously disqualified single-controller systems from serious consideration for large-scale training.

The conceptual reframing is equally significant: the latency advantage of multi-controller systems comes from when host-side work happens relative to computation, not from where the controller resides. Multi-controller systems run host-side work locally over PCIe, overlapping it with accelerator execution. PATHWAYS achieves the same overlap through parallel asynchronous dispatch that exploits the statically known resource requirements of compiled functions—moving host-side work for successor nodes to before predecessor completion rather than after. This insight generalizes beyond PATHWAYS: any single-controller system with access to static resource requirements can amortize coordination latency across pipeline stages, closing the gap with multi-controller dispatch without requiring a copy of the user program on every host.

The work also reconciles a latent tension in the ML infrastructure landscape: the bifurcation between training systems (multi-controller, exclusive ownership, throughput-optimized) and serving systems (single-controller or request-router, shared resources, latency-optimized). PATHWAYS demonstrates that a single architecture can simultaneously support both patterns—gang-scheduled SPMD training steps interleaved with fine-grained inference requests from multiple clients, all managed by the same centralized scheduler and resource manager. Figure 8 and Figure 9 provide the concrete evidence: PATHWAYS time-multiplexes accelerators between concurrent client programs at sub-millisecond granularity with zero context-switch overhead, and the scheduler can enforce proportional-share fairness policies (1:2:4:8 ratios demonstrated among 4 clients) while maintaining the consistent ordering required for gang-scheduled collectives. This unification dissolves a form of system-induced path dependency—models no longer need to be ported between separate training and serving infrastructure—and opens the door to workload-adaptive resource allocation where accelerators shift between training and inference based on demand.

Perhaps most profoundly, PATHWAYS repositions gang-scheduling from a necessary evil to an enabling mechanism. The conventional understanding is that gang-scheduling is required for correctness on TPUs (to prevent deadlock from inconsistent collective enqueue ordering) and beneficial for GPU efficiency (to reduce straggler effects in collectives). PATHWAYS shows that because gang-scheduling provides a global ordering point through the centralized per-island scheduler, it simultaneously enables fine-grained multi-tenancy, fair sharing, and policy-controlled resource allocation—capabilities that multi-controller systems, where each host independently decides what to run, fundamentally cannot provide. The scheduler becomes the policy enforcement point for the entire island, making resource management decisions at millisecond granularity with global information that no distributed protocol could replicate at comparable latency.

The research directions this work makes more or less attractive shift accordingly. More attractive are: (1) investigating centralized scheduling policies (beyond the current FIFO baseline) that optimize for objectives like shortest-remaining-processing-time, priority-aware preemption, or utilization-aware placement; (2) building ML-native resource managers that make allocation decisions based on model-specific characteristics (layer-wise memory requirements, communication patterns, sparsity structure) rather than generic cluster metrics; (3) extending the compiled-function abstraction to support data-dependent control flow while preserving the static analyzability that enables parallel dispatch (e.g., through shape polymorphism or bounded dynamic shapes); and (4) exploring the training-inference convergence that PATHWAYS enables—shared foundation model layers serving multiple downstream fine-tuning tasks on the same physical accelerators.

Less attractive are: (1) efforts to bolt multi-tenancy or resource sharing onto multi-controller SPMD systems through external cluster schedulers—the fundamental limitation is the lack of a global ordering point, not scheduler sophistication; (2) distributed consensus protocols for accelerator scheduling—the millisecond-timescale latency requirements and the centrality of gang-scheduling make centralized per-island schedulers the more natural design point; and (3) approaches that treat training and inference as requiring fundamentally different system architectures, since PATHWAYS demonstrates that a unified model is viable without performance compromise.

Follow-Up Research This Work Enables

Scheduling policies beyond FIFO for ML-optimized resource allocation. The current PATHWAYS scheduler uses simple FIFO ordering (Section 4.4), but the architecture provides all the necessary information for more sophisticated policies: estimated execution times (from compiled function resource requirements), inter-computation dependencies (from the dataflow graph), and multi-client priority or fairness constraints (from the resource manager). A natural follow-up would implement and evaluate shortest-job-first scheduling to reduce head-of-line blocking when small inference requests are queued behind large training steps, or deadline-aware scheduling that prioritizes computations on the critical path of latency-sensitive serving workloads. The specific experiment: compare FIFO against SJF and earliest-deadline-first on a mixed workload of T5-Base training steps (90.4k tokens/s, per Table 1) and variable-rate inference requests, measuring both tail latency (p99 inference response time) and training throughput. The hypothesis—motivated by the traces in Figure 9 showing sub-millisecond interleaving—is that SJF can reduce inference tail latency by 2–5× without meaningful training throughput loss, because the scheduler can insert short inference computations into gaps between training microbatches that would otherwise be idle pipeline bubbles.

Compiled functions with bounded dynamic shapes to extend parallel dispatch coverage. Parallel asynchronous dispatch currently requires statically known resource requirements, falling back to sequential dispatch when shapes, loop bounds, or memory needs depend on predecessor outputs (Section 4.5). This fallback covers the rare data-dependent cases in today's models but will become the common case for the sparse, routed architectures the paper envisions (Section 6.3)—Mixture of Experts, capsule networks, data-dependent control flow at the sub-example level. An important extension would be a compiled function variant that declares upper bounds on dynamic resource requirements: "this function will use at most 2 GB of HBM, produce output tensors of at most dimension [batch_size, 4096], and run for at most 100,000 iterations." With bounded dynamic shapes, the system could still use parallel dispatch (allocating the worst-case resources before predecessor completion) while the function internally adapts to actual data-dependent requirements. The specific experiment: implement a bounded-dynamic compiled function in XLA, measure the resource waste from worst-case allocation (HBM overhead, idle compute from over-provisioning) against the latency benefit of retaining parallel dispatch, and identify the Pareto frontier where bounded-dynamic dispatch is preferable to sequential fallback. This is directly motivated by the paper's observation that "efficient sparse communication is a requirement to avoid the DCN becoming a bottleneck for data-dependent control flow" (Section 4.3).

Difficulty estimation models trained on PRM score distributions to eliminate the 2048-sample overhead. The compute-optimal framework's most immediate deployment bottleneck is the cost of estimating prompt difficulty (Section 3.2): generating 2048 samples per question is more expensive than the largest test-time budgets studied. A natural follow-up would train a lightweight classifier—potentially a small transformer or even a linear probe on the base model's embeddings—to predict the difficulty quintile directly from the question text, using the PRM's average score distribution across 2048 samples as training targets. The specific experiment: train a 100M-parameter difficulty predictor on the MATH training set (12,000 questions with PRM-score-distribution labels), evaluate on the 500-question test set, measure how closely the predictor's difficulty bin assignments match the oracle bins (from Figure 4), and quantify the end-to-end accuracy of compute-optimal scaling when using predictor bins versus PRM-sample bins. The paper's Figure 4 shows that predicted (PRM-based) bins closely track oracle bins, suggesting that the difficulty signal is present in the PRM's scoring behavior; the question is whether that signal can be extracted from a single forward pass rather than 2048 sample-and-score operations. If the predictor achieves >90% bin agreement with the PRM-based method, the compute-optimal framework becomes immediately deployable without the estimation cost that currently prevents real-world use.

Combining PRM tree-search with the revision model as the proposal distribution. The paper studies search and revisions independently but explicitly acknowledges they were never combined (Section 8: "we did not experiment with PRM tree-search techniques in combination with revisions"). The complementary difficulty-dependent strengths—beam search helps most on medium problems (Figure 3, right, bins 3–4) where guided exploration is valuable, revisions help most on easy-to-medium problems (Figure 7, right, bins 1–3) where local refinement improves roughly-correct answers—suggest that a combined system could outperform either alone, particularly on bin 3 problems where both mechanisms show individual benefits. The specific experiment: use the revision model as the proposal distribution within beam search, where at each step of the search tree the model conditions on previous (potentially rejected) branches as revision context, producing higher-quality candidate steps than independent generation. Compare combined search+revision against compute-optimal search-only and compute-optimal revision-only from Figures 4 and 8 at matched generation budgets (16, 64, 256 generations), disaggregated by difficulty bin. The key measurement is whether the combined approach achieves accuracy exceeding the maximum of either method alone on bins 3–4—specifically, whether it can push bin 3 accuracy beyond the ~34% that beam search achieves alone (Figure 3, right) or the ~42% that revisions achieve at optimal ratio (Figure 7, right). If combined accuracy is additive (rather than subadditive due to overlapping benefits), this would substantially strengthen the FLOPs-matched case for test-time compute over pretraining on medium-difficulty problems.

Replication on code generation benchmarks to test difficulty-dependent scaling generalizability. Every experiment in the paper uses the MATH benchmark with PaLM 2-S*; the difficulty-dependent optimal strategies—beam search on medium problems, sequential revisions on easy problems, verifier over-optimization at high budgets—may be specific to mathematical reasoning or to the PaLM 2 model family. Code generation is the most natural generalization target because it shares MATH's key properties (multi-step logical reasoning, verifiable correctness via unit tests, clean difficulty signals through pass@1 on test cases) while testing a different reasoning modality. The specific experiment: replicate the Figure 3 (search) and Figure 7 (revisions) analyses on HumanEval or MBPP using a comparable-sized code model, training a PRM with Monte Carlo rollouts where correctness is determined by unit test execution rather than string matching, binning problems into difficulty quintiles by base-model pass@1, and measuring whether the same difficulty-dependent patterns emerge—beam search over-optimization on easy problems, beam search advantage on medium problems, optimal intermediate sequential-to-parallel ratios on hard problems. A negative result (different patterns on code vs. math) would establish important boundary conditions on the generality of the compute-optimal framework; a positive result (replicated patterns) would substantially strengthen the paper's claims and motivate investment in code-specific PRMs and revision models.

Dynamic difficulty assessment and mid-course strategy adjustment. The paper's difficulty bins are static—estimated once before strategy execution—and the strategy is fixed for the entire compute budget. A more sophisticated approach would dynamically assess difficulty during execution: start with a small number of parallel samples (say, 4), use the PRM's score distribution on those initial samples as a quick difficulty signal, and then allocate the remaining budget accordingly—switching to beam search if the problem appears medium-difficult, continuing parallel sampling if it appears easy, or terminating early with a low-confidence answer if it appears too hard. This amortizes difficulty estimation into the solution process itself and addresses the exploration-exploitation tradeoff the paper identifies as future work (Section 3.2). The specific experiment: implement an adaptive policy that starts with 4 parallel samples, estimates difficulty from the PRM scores on those samples, selects the compute-optimal strategy for the estimated difficulty at the remaining budget, measures end-to-end accuracy accounting for the estimation samples, and compares against the static compute-optimal policy from Figure 4/Figure 8. The hypothesis is that dynamic assessment can recover most of the compute-optimal gains with substantially lower total cost (since the estimation samples contribute to the final answer), making the framework deployable without the separate 2048-sample estimation phase.

Practical Applications and Downstream Use Cases

Unified training and serving infrastructure for foundation model fine-tuning. The vision sketched in Section 6.3—multiple researchers concurrently fine-tuning a shared foundation model for different downstream tasks, with the frozen base layers held once in accelerator HBM and shared across all fine-tuning jobs—becomes architecturally feasible with PATHWAYS' multi-tenancy and centralized resource management. In a deployment with a 136B-parameter foundation model (at the scale PATHWAYS validates in Section 5.3, achieving ~97% single-island throughput over two islands), the base transformer layers could be loaded once onto a set of accelerators, and multiple clients could submit fine-tuning programs that share those layers as read-only context while training separate task-specific heads on the same or different accelerators. The centralized scheduler interleaves fine-tuning steps from different clients at sub-millisecond granularity (as demonstrated in Figure 9), and the sharded object store's HBM integration (Section 4.6) keeps the shared base layers resident without repeated data loading. The concrete benefit: a 4–8× reduction in total accelerator-hours for multi-task fine-tuning compared to dedicating separate accelerators to each task, since the memory and compute cost of the base layers is amortized across clients. This is a direct realization of the "foundation models" vision (Bommasani and et. al., 2021) that the paper explicitly cites as motivation.

Elastic training that adapts to available cluster resources without checkpoint-restart. PATHWAYS' dynamic resource management (Section 4.1)—the ability to add and remove backend compute resources dynamically, combined with the virtual-to-physical device indirection and re-lowering of programs when the mapping changes—enables training jobs that expand or contract their accelerator allocation mid-training without the traditional stop-checkpoint-reconfigure-restart cycle. In a shared cluster where accelerator availability fluctuates (due to preemption by higher-priority jobs, maintenance, or variable user demand), a PATHWAYS training program could release a subset of its virtual devices when preempted and continue training on the remaining devices at a proportionally reduced batch size, then re-expand when resources become available again. The re-lowering mechanism (Section 4.2) handles the remapping from virtual to physical devices transparently. The concrete benefit: training jobs that would traditionally be killed and restarted on preemption can instead degrade gracefully, reducing wasted compute from discarded progress. For a 3B-parameter model training at 131.4k tokens/s on 128 TPUs (Table 2), a 50% resource reduction would yield approximately 65–70k tokens/s rather than zero (if the job were killed), recovering roughly half the throughput rather than losing all progress. This elasticity is enabled by the centralized resource manager—it is fundamentally incompatible with multi-controller systems where exclusive resource ownership is baked into the execution model.

Cross-pod model parallelism for models exceeding single-island memory capacity. The paper demonstrates that PATHWAYS can train models split across multiple TPU pods connected over DCN, achieving the same throughput (131.4k tokens/s for a 3B model) whether using a single island of 128 cores or 4 islands of 32 cores each (Section 5.3), and achieving ~97% throughput for a 136B model split across two islands compared to a single island with twice the devices. This directly enables training runs that exceed the HBM capacity, ICI bandwidth, or core count of a single TPU pod. In a deployment scenario where the largest available homogeneous island has 512 TPUs but the model requires 1024 TPUs' worth of HBM to fit parameters, optimizer state, and activations at the minimum viable batch size, PATHWAYS' cross-pod pipelining (Figure 10) and data-parallel gradient reduction over DCN (Figure 12) make the training feasible without waiting for a larger pod to become available. The concrete benefit: model scale is decoupled from maximum island size. Researchers can train models requiring aggregate HBM or compute exceeding any single pod, using PATHWAYS' sharded dataflow to coordinate communication across DCN-connected islands, with the DCN transfer overhead partially hidden by overlapping with computation (as the trace in Figure 10 shows for the 3B pipelined model, and Figure 12 shows for the 64B data-parallel model).

When to Prefer This Method

The paper positions PATHWAYS explicitly against two architectural alternatives: multi-controller SPMD systems (JAX, PyTorch) and older single-controller systems (TensorFlow v1). The decision rule is:

Prefer PATHWAYS over multi-controller SPMD systems when: the computational workload includes non-SPMD patterns that multi-controller systems handle poorly—pipeline parallelism with many stages (Table 2 shows efficient scaling to 16 stages), Mixture-of-Experts-style computational sparsity requiring data-dependent routing, or multi-task workloads where shared model layers are concurrently accessed by multiple clients (Section 6.3). The condition is that the non-SPMD flexibility justifies any residual coordination overhead, which the paper shows is zero for computations exceeding 2.3 ms on 128 TPUs and 35 ms on 2048 TPUs (Figure 6). Also prefer PATHWAYS when centralized resource management is required—multi-tenancy with gang-scheduled collectives (Figure 8), proportional-share fairness (Figure 9), elastic resource allocation, or transparent migration—since multi-controller systems lack the global ordering point that enables these features. The paper's key evidence that the performance cost of switching is negligible is Table 1: identical throughput across all tested model sizes between PATHWAYS and JAX on standard SPMD training.
Prefer multi-controller SPMD systems over PATHWAYS when: the workload is pure SPMD with no foreseeable need for non-SPMD patterns, exclusive resource ownership is acceptable (dedicated islands, no sharing), and the deployment does not require centralized resource management features. In this regime, multi-controller systems provide equivalent performance (per Table 1) with lower system complexity—no Plaque dependency, no scheduler to configure, no virtual-to-physical device indirection. The paper's Figure 5 shows that for extremely small computations (scalar AllReduce in OpByOp mode), multi-controller JAX maintains higher throughput than single-controller systems as the number of hosts scales, though this gap vanishes as soon as computations are fused (JAX-F vs. PW-F parity up to 1000 TPU cores).
Prefer TensorFlow v1 or similar older single-controller systems: the paper makes no argument for this case. The entire design is motivated by TF v1's failure modes—dispatch latency, sequential coordination, and sharded graph blowup—and PATHWAYS is positioned as a direct replacement that fixes these problems. There is no workload or deployment scenario described in the paper where TF v1 would be preferable to PATHWAYS, except possibly scenarios involving only a "single, smallish, exclusively-owned island of accelerators" (Section 2) with no performance sensitivity, where TF's mature ecosystem might outweigh PATHWAYS' architectural advantages.