google
Does JAX need a deterministic reasoning kernel for emerging LLM-based decision systems?
Hi JAX team,
I’d like to raise a conceptual question that comes from a recent systems-behavior experiment, and I believe it may have architectural implications for JAX and the emerging field of LLM-based control systems.
This is not a bug report — it is a design discussion.
1. Background: an unexpected form of deterministic behavior in LLMs
In a recent experiment, I used the ChatGPT web client (no API, no temperature control, no RAG, no prompt engineering, no external tools) as a decision-making engine in a multi-step risk-evaluation task.
To my surprise:
- reasoning paths remained highly stable
- decisions were repeatable across long sessions
- no hallucinations occurred
- the system behaved like a finite-state decision kernel rather than a stochastic generator
- "expression-level constraints" seemed to override sampling randomness
Video demonstration (purely technical, non-financial): https://www.youtube.com/watch?v=zgHvAkcBINs
This raises an interesting architectural question:
Does current ML tooling (including JAX) assume too strongly that LLM inference is fundamentally stochastic, when in fact some forms of implicit determinism can emerge when the right semantic constraints exist?
2. Why this might matter for JAX
Many researchers are now using JAX to build:
- agent systems
- simulation environments
- alignment protocols
- structured reasoning pipelines
- differentiable planning loops
But almost all existing frameworks implicitly rely on:
- stochastic sampling
- external control loops
- retrieval-based stabilization (RAG)
Yet the experiment suggests that an LLM can exhibit emergent deterministic structure without RAG, memory, or temperature tuning — if the “expression structure” of the task is properly constrained.
This leads to a hypothesis:
Perhaps JAX could support a new category of tooling for deterministic reasoning kernels — computation graphs designed specifically for stable, auditable LLM decision processes.
3. A possible research question for the JAX ecosystem
Would it be meaningful to explore:
A Deterministic Reasoning Kernel (DRK)
A lightweight, JAX-native abstraction that:
- enforces structural invariants across reasoning steps
- preserves semantic state
- provides deterministic/reproducible execution logs
- integrates with existing JAX transformations
- does not rely on RAG or external memory systems
- allows LLMs to behave as state-machines rather than stochastic samplers
This might be useful for:
- agent frameworks
- safety-critical reasoning
- multi-step decision workflows
- alignment research
- simulation + control tasks
- verifiable computation
4. Why I believe JAX is the right place to discuss this
JAX already excels at:
- functional purity
- reproducibility
- transformation-based computation
- explicit state threading
- predictable execution
Which makes it a strong candidate for supporting stable reasoning pipelines, not just differentiable computation.
My question is:
❓ Would the JAX team be open to a discussion or RFC on deterministic reasoning kernels for LLM-based decision systems?
This is an early-stage idea, but the experimental signal is strong enough that it may be worth exploring from a systems-design perspective.
Thank you for your time — and I appreciate any thoughts from the team or community.
Hi - thanks for reaching out! The idea you're proposing seems like a bit of a mismatch for JAX's scope: the JAX core package does not include any LLM-specific APIs, nor does it even contain frameworks for defining neural networks at all. What you're describing seems like something at a much higher level, and likely specific to particular LLM network architectures. Given that, it may be more relevant in a project like maxtext.
