From 02fb92af2d6a4709dde46309e32b43625985d94d Mon Sep 17 00:00:00 2001 From: Charlie Fay Date: Sat, 29 Aug 2026 22:37:48 +0200 Subject: [PATCH] Keep gradients finite when electrodes are driven to zero amplitude The default sqrt size equation (Tehovnik 2007) computes sqrt(x / current_spread) whose derivative 1/(2*sqrt(x)) is infinite at x = 0. An electrode driven to exactly zero amplitude therefore emits a NaN gradient that propagates into every parameter of the model, while the forward pass continues to look correct. Nothing signals the corruption until the whole network is NaN, which presents as a diverged model rather than a bug. This is reachable in ordinary use rather than a corner case: sparse stimulation is the goal for a charge-limited device, sigmoid encoder outputs underflow to zero, and any charge penalty drives amplitudes down. With one electrode in four at zero, 25% of the amplitude gradients are NaN. Clamping the argument keeps the derivative finite. The bound is far below any physical stimulation current, so the forward output is unchanged: the worst absolute difference over 30 random varied-amplitude batches is exactly zero. The sigmoid size equation is smooth at zero and is unaffected. Adds a regression test covering 100%, 50% and 25% of electrodes off, plus a check that the forward percept is untouched. The gradient tests fail on the previous behaviour and the forward test passes either way. Co-Authored-By: Claude Opus 5 --- dynaphos/simulator.py | 15 +++++++++- test/test_gradients.py | 62 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 test/test_gradients.py diff --git a/dynaphos/simulator.py b/dynaphos/simulator.py index 30b388b..4758e6b 100644 --- a/dynaphos/simulator.py +++ b/dynaphos/simulator.py @@ -12,6 +12,12 @@ get_deg2pix_coeff, set_deterministic, print_stats, sigmoid, to_numpy, Map) +# Lower bound on the argument of the sqrt size equation. Far below any +# physical stimulation current, so the forward pass is unaffected; its only +# purpose is to keep the gradient finite at zero amplitude. +MIN_SIZE_ARG = 1e-30 + + class State: def __init__(self, params: dict, shape: Tuple[int, ...], verbose: Optional[bool] = False): @@ -145,7 +151,14 @@ def __init__(self, params: dict, shape: Tuple[int, ...], p = self.params['size'] if p['size_equation'] == 'sqrt': # Tehovnik 2007 def f(x): - return torch.sqrt(torch.div(x, p['current_spread'])) + # The derivative of sqrt is infinite at zero, so an electrode + # driven to exactly 0 A would emit a NaN gradient that + # propagates silently through the whole model while the + # forward pass still looks correct. Clamping the argument + # keeps the derivative finite and leaves the forward output + # unchanged for any physically meaningful current. + return torch.sqrt( + torch.div(x, p['current_spread']).clamp_min(MIN_SIZE_ARG)) elif p['size_equation'] == 'sigmoid': # Bosking et al., 2017 def f(x): return 0.5 * p['MD'] * sigmoid(p['slope_size'] * diff --git a/test/test_gradients.py b/test/test_gradients.py new file mode 100644 index 0000000..c942a54 --- /dev/null +++ b/test/test_gradients.py @@ -0,0 +1,62 @@ +"""Gradients must stay finite when electrodes are switched off. + +The simulator is used for end-to-end optimisation, where an encoder routinely +drives some electrodes to zero: sparse stimulation is the goal for a +charge-limited device, sigmoid outputs underflow, and any charge penalty pushes +amplitudes down. The sqrt size equation has an infinite derivative at zero, so +without a lower bound those electrodes emit NaN gradients that propagate into +every parameter while the forward pass still looks correct. +""" + +import os + +import numpy as np +import pytest +import torch + +from dynaphos.cortex_models import \ + get_visual_field_coordinates_probabilistically +from dynaphos.simulator import GaussianSimulator +from dynaphos.utils import load_params + +CONFIG = os.path.join(os.path.dirname(__file__), '..', 'config', 'params.yaml') + + +def build(batch_size=2, num_phosphenes=64, resolution=32, seed=0): + params = load_params(CONFIG) + params['run']['resolution'] = [resolution, resolution] + params['run']['batch_size'] = batch_size + params['run']['seed'] = seed + coordinates = get_visual_field_coordinates_probabilistically( + params, num_phosphenes, np.random.default_rng(seed)) + return GaussianSimulator(params, coordinates, rng=np.random.default_rng(seed)) + + +@pytest.mark.parametrize('fraction_off', [1.0, 0.5, 0.25]) +def test_gradient_finite_with_electrodes_off(fraction_off): + simulator = build() + amplitude = torch.full((2, simulator.num_phosphenes), 80e-6) + n_off = int(simulator.num_phosphenes * fraction_off) + amplitude[:, :n_off] = 0.0 + amplitude.requires_grad_(True) + + simulator.reset() + percept = simulator(amplitude) + percept.sum().backward() + + assert torch.isfinite(percept).all(), 'forward pass produced non-finite values' + assert torch.isfinite(amplitude.grad).all(), ( + f'{torch.isnan(amplitude.grad).float().mean():.0%} of amplitude ' + f'gradients are NaN with {fraction_off:.0%} of electrodes off') + + +def test_forward_unchanged_by_the_clamp(): + """The bound sits far below any physical current, so percepts must not move.""" + simulator = build() + generator = torch.Generator().manual_seed(0) + amplitude = torch.rand( + (2, simulator.num_phosphenes), generator=generator) * 100e-6 + simulator.reset() + percept = simulator(amplitude) + assert torch.isfinite(percept).all() + assert percept.max() > 0, 'expected a visible percept at up to 100 uA'