🤖 machine learning

One Model, Two Roles: Emergent Specialization in a Shared Recurrent Transformer

This paper demonstrates that a shared-weight recurrent Transformer can spontaneously develop distinct functional roles—specifically a committed proposal state and an uncertain intermediate state—when provided with a clear signal to differentiate between update types, such as asymmetric input injection or a separate level token.

Original authors: Jucheng Shen, Barbara Su, Anastasios Kyrillidis

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jucheng Shen, Barbara Su, Anastasios Kyrillidis

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a single, very smart assistant (the AI model) who is trying to solve a complex puzzle, like a Sudoku or a maze. Usually, to solve these, you might think you need two different people: one to brainstorm wild ideas and another to double-check the facts.

This paper asks a fascinating question: Can one single person do both jobs effectively if they just know which job they are doing at that moment?

The researchers built a system called AIR (Asymmetric Input Recurrence) to test this. Here is how it works, using simple analogies:

The Setup: One Brain, Two Hats

The AI has two "mental states" it cycles through repeatedly:

The "Scratchpad" (State $z_L$ ): This is for messy, local thinking. It's where the AI tries out small changes, gets confused, and figures out what might work.
The "Proposal" (State $z_H$ ): This is for making a firm decision. It's where the AI commits to a full solution, even if it's not perfect yet.

In most AI models, these two roles are handled by two completely different sets of brain cells (parameters). But in this experiment, the researchers forced the exact same brain cells to do both jobs.

The Secret Sauce: The "Signal"

How does the brain know which hat to wear? The researchers gave it a tiny signal:

When the brain is doing the Scratchpad job, it gets a "hint" (the puzzle input) injected into its mind.
When it's doing the Proposal job, it gets no hint. It has to rely on its own memory of the previous steps.

It's like a chef who gets a fresh list of ingredients when they are planning a recipe, but has to cook the final dish without looking at the list again. The presence or absence of that list tells the chef which mode to be in.

What Happened? (The Results)

The researchers found that yes, the single brain successfully split into two distinct roles, even though it was the same hardware.

The "Commitment" Split:
- The Proposal state ( $z_H$ ) became very confident. It looked at the whole puzzle and said, "Here is the answer!" even if the answer was wrong. It was fully "committed."
- The Scratchpad state ( $z_L$ ) remained unsure. It left some cells blank (like a "to-be-determined" note) and kept shifting its focus around. It was "locally uncertain."
The "Attention" Split:
- When the brain was in Scratchpad mode, it looked very closely at its immediate neighbors (like checking the 3x3 box in Sudoku). It was local.
- When the brain was in Proposal mode, it looked at the whole board at once. It was global.

The "Freeze" Test

To prove these two states were actually talking to each other and relying on each other, the researchers did a "freeze" experiment. They literally stopped one part of the brain from moving while the other kept working.

In Sudoku: If they froze the "Proposal" brain, the "Scratchpad" brain stopped changing its mind (it became calm). But if they froze the "Scratchpad," the "Proposal" brain went crazy and changed its mind constantly.
In Mazes: The relationship was different; freezing one made the other work harder to compensate.

This proved that the two roles weren't just random; they were a coordinated team.

The Big Lesson

The paper concludes that you don't necessarily need two different AI models to do complex reasoning. You just need one model and a clear signal that tells it, "Right now, you are the dreamer," or "Right now, you are the decider."

If you give the model a way to distinguish between the two tasks (like the "hint" or a special token), it naturally evolves to specialize, creating a "division of labor" inside a single shared brain.

In short: One brain can wear two hats, as long as it knows which hat it's wearing at any given second.

Technical Summary: One Model, Two Roles: Emergent Specialization in a Shared Recurrent Transformer

Problem Statement
The paper investigates whether a single, shared-weight recurrent Transformer can develop distinct internal computational roles for two latent states without being partitioned into separate parameter modules. While prior work like HRM (Hierarchical Reasoning Model) uses separate networks for different latent states and TRM (Transformer Reasoning Model) uses a single shared network, it remains unclear if a shared model can spontaneously specialize into distinct functional roles (e.g., a "proposal" state vs. a "refinement" state) when provided with a stable signal distinguishing the update types. The central question is whether specialization requires structural modularity (separate parameters) or if it can emerge from functional differentiation driven by a clear state-identity signal within a monolithic shared architecture.

Methodology: Asymmetric Input Recurrence (AIR)
To address this, the authors introduce Asymmetric Input Recurrence (AIR), a minimal testbed derived from HRM and TRM.

Architecture: AIR maintains two latent states, $z_H$ (high-level) and $z_L$ (low-level), updated by a single shared Transformer model $f(\cdot; \theta)$ .
Update Mechanism: The recurrence follows a pattern of interleaved sub-steps (e.g., $L, L, H, L, L, H$ $L, L, H, L, L, H$ ). The critical asymmetry lies in the input injection:
- L-updates: The encoded input $\tilde{x}$ is injected into the update rule ( $z_L \leftarrow f(z_L + z_H + \alpha_L \tilde{x})$ ).
- H-updates: The encoded input is not injected ( $z_H \leftarrow f(z_H + z_L + \alpha_H \tilde{x})$ with $\alpha_H=0$ ).
Training: The model is trained using one-step truncated gradients (backpropagating only through the final L and H updates of a cycle) on two synthetic reasoning tasks: Sudoku-Extreme (9x9 puzzles) and Maze (30x30 grids).
Analysis: The authors employ several techniques to probe the model:
- Decoded Rollouts: Visualizing intermediate states to observe commitment levels.
- Freeze Experiments: Freezing one state to measure content changes in the other.
- Ablations: Testing symmetric input injection (removing the asymmetry) and alternative state-identity signals (level tokens).
- Mechanistic Analysis: Analyzing attention patterns (locality vs. globality) and violation sensitivity.

Key Contributions and Results

Existence of Functional Specialization:
Decoded rollouts reveal a stable, emergent split in the shared model's behavior:
- $z_H$ (High-level): Behaves as a "fully committed proposal state." It decodes to a complete candidate solution at every sub-step, even if incorrect.
- $z_L$ (Low-level): Behaves as a "shifting-commitment scratchpad." It retains local uncertainty, with undecided positions (BLANK in Sudoku, PAD in Maze) that shift across sub-steps.
- Control: In symmetric variants (where both states receive identical input injection), this distinct role split nearly vanishes, confirming that the specialization is not inherent to having two states but relies on the ability to distinguish update types.
State-Identity Signal Induces Specialization:
The paper demonstrates that the shared model requires a reliable signal to distinguish update types to induce specialization.
- Input Asymmetry: Asymmetric input injection ( $\Delta > 0$ ) significantly improves accuracy (e.g., ~59% on Sudoku vs. ~51% for symmetric variants).
- Level Tokens: The authors tested alternative signals using learned level tokens. Only a "prepend-and-strip" strategy (where a token is added, attended to, and removed) successfully recovered the performance gap. Simple element-wise addition or persistent tokens failed, suggesting the signal must be structurally separable and not mixed into content tokens.
- Single-State Ablation: A single-state model with periodic input injection failed to match the performance of the two-state asymmetric model, indicating that distinct memory buffers are necessary for stable specialization.
Interdependent State Dynamics:
Freeze experiments reveal a specific relationship between the two states:
- Sudoku: Freezing $z_H$ suppresses $z_L$ 's activity (reducing content changes), while freezing $z_L$ increases $z_H$ 's activity.
- Maze: Freezing either state increases content changes in the other.
- Symmetric Control: Symmetric models lack this organized, role-consistent interaction pattern, showing perturbation-sensitive but unstructured responses.
Mechanistic Evidence (Attention Analysis):
Attention analysis confirms the functional split:
- Locality: $L$ -updates consistently attend more locally (to immediate neighborhoods) than $H$ -updates, which exhibit more global attention.
- Violation Sensitivity: On Sudoku, deeper layers of $L$ -updates show specific attention routing toward violated constraints, a signal not observed in $H$ -updates or in the Maze task.

Significance and Claims
The paper claims that explicit parameter modularity (separate networks) is not strictly necessary for functional specialization in recurrent reasoning. Instead, a clear state-identity signal—supplied either through asymmetric input injection or a structurally separable level token—is sufficient to induce stable, related functional roles within a shared-parameter recurrent Transformer.

The authors position this work as foundational research into the inductive biases of small recurrent reasoning systems. They emphasize that their conclusions are drawn from controlled synthetic settings (Sudoku and Maze) and do not claim universality across natural language or partial observability. The primary contribution is demonstrating that "one model, two roles" is a viable paradigm for efficient reasoning architectures, provided the model can distinguish the computational context of its updates.