Neural network optimisation for sparse deconvolution

Hi, long-time listener—first-time poster.

I have been wasting time waiting for a trans-continental video call by thinking about a problem a friend worked on during his PhD. He was recovering spike timings from extracellular neural recordings using sparse deconvolution and matching pursuit. The problem sounds simple—recover a sparse set of spikes from a noisy recording. One common formulation is:

min_x 1/2(|y-Ax|_2)^2 + λ|x|_1

y is the measured signal, A contains the waveform, and x is the event locations.

Matching pursuit is a greedy method of solving this sparse reconstruction problem. Find the dictionary element that best matches the current residual, estimate its amplitude, subtract its contribution and then just like washing your hair repeat k times.

It turns out this problem has made frequent guest appearances throughout my career. From extracellular neural recordings, full-waveform LiDAR returns to ultrasonic echoes for non-destructive testing.

Another guest character that has started making frequent appearances in my life is PyTorch. This led me to the obvious episode idea—crossover guest appearances.

After a suggestion from Paul Bellette, I started using PyTorch to build a small differentiable matching pursuit experiment. Training—after all—is just solving an optimisation problem. With all the libraries and hardware around, society has built a very good hammer, so I am going to turn this problem into a nail.

I set up a small PyTorch model as an explicitly unrolled pursuit:

residual - r₀ = input y ├correlate(r₀, waveform) ├soft spike selection ├ estimate amplitude ├ add spike to reconstruction └ subtract reconstructed spike ▼ r₁ -> correlate -> select spike -> estimate amplitude -> add spike -> subtract ▼ ... ▼ rₖ - reconstruction error

The interesting bit is making this differentiable. A conventional matching-pursuit would use a hard operation such as argmax to choose the next spike location. Instead, each stage uses a softmax.

The problem is this forces there to be exactly k spikes. I added one extra location at the end of the softmax representing “no spike”. The unrolled model is solving: find up to K spikes

Because this now remains differentiable, the optimiser can simultaneously learn the waveform dictionary and find the spike times directly from the reconstruction error. No spike-timing labels. No pre-specified waveform template. Just a reconstruction loss and some weak penalties:

Loss = reconstruction error + λ₁(spike activity) + λ₂(refractory penalty) + λ₃(waveform smoothness)

That is probably the part I find most interesting. PyTorch does not particularly care whether what you built is a neural network. You build a computational graph, add a loss and some parameters, then let it do it's thing.“Training” is just one application of that thing.