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:

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)

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.

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},
}