pavi-rajes/dr-gru-dynamic-routing

Dynamic Routing GRU

A recurrent neural network trained to perform context-dependent sensory routing — the same computational problem solved by prefrontal and parietal circuits in the primate brain during flexible attention tasks.

The model learns to alternate between visual and auditory stimulus discrimination across block switches, using only reward feedback to infer the current context. No explicit context signal is ever provided as input.

What the model does

Each trial the network receives 7 input channels:

Channel	Description
vis1, vis2	Two visual stimuli
aud1, aud2	Two auditory stimuli
reward	Reward feedback
own_lick	Efference copy of own lick
trial_phase	Current phase (quiescent / stimulus / response)

The session alternates between visual blocks (lick to vis1, withhold to vis2/aud1/aud2) and auditory blocks (lick to aud1, withhold to others), with ~5 instruction trials at each block start. The network must detect the block switch from reward statistics and route accordingly.

Evaluation (5 seeds × 20 sessions, noiseless)

Metric	Mean	Std
d'_intra — within-block stimulus discrimination	4.634	0.021
d'_inter — across-block context switching	1.864	0.128
Hit rate	0.998	0.003
False alarm rate (intra-modal)	0.000	0.000
Catch lick rate	0.001	0.001

d' is signal detection theory discriminability (d' > 1.5 = reliable discrimination, d' ≈ 4.65 = ceiling).

Architecture

• GRU, 200 hidden units, 7 inputs, 1 output (lick probability)
• Recurrent noise σ_rec = 0.05 (training only; noiseless at eval)
• Weight constraints: W_in (candidate gate) non-negative; W_ir / W_iz non-negative for sensory channels 0–3
• Update gate bias initialized to +1.0

Training

• Curriculum: Stage 0 (200 steps, target-only warm-up) → Stage 3 (block-switching, up to 15 000 steps)
• Loss: BCE on response window + λ_rate=1e-3 activity regularization + λ_Δw=1e-4 weight-change regularization
• FA penalty: Option B — negative RPE signal (reward channel dips on incorrect lick), biologically matching dopaminergic prediction-error encoding
• Optimizer: Adam, lr=1e-3, gradient clip max_norm=1.0, batch_size=16
• All 5 seeds reached context-learning criterion in ~700 training steps

Quickstart

import torch
from huggingface_hub import hf_hub_download
from dr_gru.models.gru import DynamicRoutingGRU

# Download checkpoint
ckpt_path = hf_hub_download(
    repo_id="pavi-rajes/dr-gru-dynamic-routing",
    filename="checkpoints/seed0_context_learning.pt",
)

ckpt  = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model = DynamicRoutingGRU(N=200, n_in=7, n_out=1, sigma_rec=0.05)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

print(f"d'_intra = {ckpt['metrics']['d_prime_intra']:.3f}")
print(f"d'_inter = {ckpt['metrics']['d_prime_inter']:.3f}")

Clone the full repo for rollout and PCA visualization:

git clone https://huggingface.co/pavi-rajes/dr-gru-dynamic-routing
cd dr-gru-dynamic-routing
python demo_load_model.py --ckpt checkpoints/seed0_context_learning.pt

Available checkpoints

All checkpoints are snapshotted at the moment each seed first achieved d'_intra > 1.5 and d'_inter > 1.5 for two consecutive evaluations.

File	d'_intra	d'_inter	Train step
checkpoints/seed0_context_learning.pt	4.642	2.048	700
checkpoints/seed1_context_learning.pt	4.653	1.933	700
checkpoints/seed3_context_learning.pt	4.598	1.888	775
checkpoints/seed4_context_learning.pt	4.653	1.681	700
checkpoints/seed5_context_learning.pt	4.626	1.769	700

Citation

@misc{dr_gru_2026,
  author  = {Ravi, Pavithra},
  title   = {Dynamic Routing GRU: Context-Dependent Sensory Routing via Reward Feedback},
  year    = {2026},
  url     = {https://huggingface.co/pavi-rajes/dr-gru-dynamic-routing},
}