Recurrent Routing & Self-Rethinking

• The paper introduces recurrent routing where iterative, controller-driven selection of computational units refines intermediate representations.
• It employs LSTM and GRU-based controllers to dynamically route and combine expert units in both ReSet and GRAPHMOE architectures.
• Empirical results reveal improved accuracy, robustness, and efficiency, including up to 40% CU skipping and significant gains on commonsense benchmarks.

Recurrent routing and self-rethinking are architectural strategies for neural networks aimed at enabling dynamic, input-adaptive computation—moving beyond fixed, feedforward transformation. These techniques allow a model to iteratively revisit, refine, and recombine intermediate representations by leveraging controllers that repeatedly route information through pools of computational units or experts. The result is a form of iterative, example-specific reasoning within both vision and LLMs, with demonstrated gains in robustness, expressivity, and computational efficiency across multiple benchmarks (Kemaev et al., 2018, Tang et al., 14 Jan 2025).

1. Architectural Foundations of Recurrent Routing

In both vision (ReSet) and language (GRAPHMOE) domains, the central innovation is replacing fixed sequential block structures with modules that contain a pool of expert or computational units (CUs) and a learnable routing controller. The controller selects, at each iteration, a subset or mixture of these units for processing the current state, and the process repeats over several rounds.

In ReSet (Kemaev et al., 2018), ResNet’s standard residual block sequence is segmented into “ReSet blocks,” each orchestrating $k$ recurrent routing steps through a small pool of $n$ learnable CUs—each a residual block with independent parameters. The routing for each input is governed by a lightweight controller (CNN- or RNN-based), producing a distribution over the pool at each iteration. The selected units’ outputs are combined and added residually to the current feature state, allowing for state refinement across iterations.

GRAPHMOE (Tang et al., 14 Jan 2025) extends this paradigm to LLMs, embedding multi-expert MoE transformations in Transformer decoder layers. Each MoE layer is augmented with a “virtual node” implemented via a low-rank GRU. Routing is performed recurrently: at each round, the router uses the refined current state—incorporating prior expert mixtures—to select a top-$k$ subset of experts. The outputs from selected experts are aggregated and fed through the GRU, producing a residual message that updates the layer’s hidden state for the next routing round.

2. Mathematical Formalism and Routing Controllers

The routing controller, central in both models, computes a distribution over available computational units conditional on the current feature and, in recurrent variants, the evolving hidden/controller state.

• For ReSet, the RNN-based controller at iteration $t$ produces $[\pi_t, h_{t+1}] = C(x_t, h_t)$, where $\pi_t \in \mathbb{R}^n$ are expert-selection logits. After softmax normalization,

$$y_{t,i} = \frac{\exp(\pi_{t,i})}{\sum_{j=1}^n \exp(\pi_{t,j})}$$

The next state is composed as $x_{t+1} = x_t + \sum_{i=1}^n y_{t,i} F_i(x_t)$, allowing both soft mixtures and cycles.

• In GRAPHMOE, the router $R^\ell$ projects the feature to $n$ dimensions, and after applying softmax and top-$n$0 sparsification, yields mixture weights $n$1:

$n$2

The expert outputs are $n$3, where each $n$4 is either a LoRA-adapted FFN (language) or a learnable CU (vision).

A recurrent controller (LSTM in ReSet, GRU in GRAPHMOE) provides the capacity for multi-step, state-aware routing, enabling the model to refine decisions with each pass.

3. The Self-Rethinking Mechanism

Self-rethinking refers to the capacity of the model to iteratively “revisit” its previous routing and computation, refining its intermediate representation before handing it off to the next layer or stage.

In ReSet (Kemaev et al., 2018), repeated routing enables cycles—the network may reuse or skip CUs multiple times within the block. The LSTM controller’s hidden state $n$5 encodes information about past decisions, supporting dynamic adjustment and sharpening. Empirical analysis shows that early iterations in ReSet blocks are similar across inputs, while later iterations diverge semantically.

GRAPHMOE (Tang et al., 14 Jan 2025) formalizes self-rethinking as message passing on a pseudo-graph. Each inference round, the virtual node aggregates the expert mixture via a GRU, with the hidden state update: $n$6

$n$7

$n$8

$n$9

$k$0

After $k$1 iterations, the final refined output $k$2 is passed up the stack. This process yields deeper, input-sensitive reasoning within each MoE layer.

4. Training Objectives, Regularization, and Handling Discrete Routing

Both approaches augment standard per-task losses with regularizers and techniques to stabilize or diversify routing:

• ReSet (Kemaev et al., 2018) applies an entropy-based regularization:
  • $k$3 times the entropy of $k$4 (encouraging global diversity of paths)
  • $k$5 times the expectation over $k$6 of $k$7 (encouraging each sample’s routing to be sharp/deterministic)
  • A skip-reward term biases the controller toward skipping computation when appropriate, introducing a trade-off between accuracy and compute.
  • Methods for differentiating through discrete routing include REINFORCE (unstable), Gumbel-Softmax, and Straight-Through Gumbel-Softmax.
• GRAPHMOE (Tang et al., 14 Jan 2025) minimizes

$k$8

where $k$9 is a load-balancing loss that penalizes unbalanced expert usage:

$t$0

$t$1

Implementation requires backpropagation through time over the recurrent routing steps.

5. Empirical Findings and Model Behavior

Empirical analyses highlight multiple benefits and properties of recurrent routing and self-rethinking:

Model	Domain	Routing Mechanism	Self-Rethinking Implementation	Core Empirical Gains
ReSet	Vision	Recurrent CU routing	LSTM controller, recurrent residual CU selection	Improved accuracy and generalization on CIFAR-10.1; up to 40% CU skipping with small accuracy cost; class/sample adaptive routing (Kemaev et al., 2018)
GRAPHMOE	Language	Recurrent MoE routing	Virtual node (GRU) aggregator	+2.1–2.3 points avg. gain across 8 commonsense benchmarks vs. prior LoRA+MoE; SOTA with MixLoRA/MixDoRA; >20% improved expert utilization uniformity (Tang et al., 14 Jan 2025)

• In ReSet, RNN-based controllers outperform CNN-based: they reliably replicate the original ResNet order, support nontrivial input-dependent routing, and can skip or reuse CUs flexibly.
• In GRAPHMOE, accuracy increases with self-rethinking rounds up to a point (2–3 passes), then declines (“overthinking”).
• Routing diversity: In both models, visualization of controller distributions demonstrates that early iterations remain generic, while later rounds show class- or sample-specific routing.

A plausible implication is that recurrent routing enables a trade-off between depth (more reasoning rounds) and computational cost, adjustable even at inference time.

6. Related Methodologies and Integrations

GRAPHMOE (Tang et al., 14 Jan 2025) explicitly integrates LoRA-based parameter-efficient finetuning, broadening applicability to large frozen pretrained models. The MoE expert weights are composed of frozen base weights plus learned low-rank adapter matrices: $t$2 This enables scalable deployment and extension of the self-rethinking mechanism to LLM settings.

ReSet (Kemaev et al., 2018) demonstrates compatibility with established deep architectures (ResNet), requiring only local modifications to residual stage structure.

Both approaches make the routing controller and expert pool sizes tunable architectural hyperparameters, and allow for additional structure such as a skip/identity unit to control computational budget.

7. Significance and Prospects

Recurrent routing and self-rethinking represent a significant shift toward adaptive, iterative processing in neural architectures. By enabling models to “rethink” and refine intermediate representations via input-dependent multi-step computation within a layer, these approaches confer benefits in robustness, accuracy, and efficiency, as well as improving expert utilization in mixture models.

Open directions include investigation of scalability to deeper, more structured pseudo-graphs, refinement of routing regularization strategies, improved handling of overthinking, and further integration with other architectural innovations. The empirical evidence across both vision and language domains confirms the impact and versatility of recurrent routing and self-rethinking as design principles (Kemaev et al., 2018, Tang et al., 14 Jan 2025).