Optimize window length (STFT) via gradient descent (in neural networks)

Question

The authors from this paper optimized a Gaussian window size via gradient descent (the σ parameter of the bell curve) together with the other parameters of neural networks.

I don't use Gaussian window but use Hann window instead. I would like to know how to optimize stft window size with Hann/Hamming window via gradient descent?

The problem is that unlike the Gaussian window, Hann window does not have continuous parameter σ as a proxy for gradient descent. Is there a way to rewrite Hann window or are there parameters that one could use to control the window size and that are differentiable? Currently the $n$ is positive integer and is not differentiable.

torch.hann_window uses:

$$w[n] = \begin{cases} \tfrac{1}{2} \left( 1 − \cos \left(\tfrac{2\pi}{N-1}n\right) \right) \qquad & 0 \le n \le N-1 \\ 0 & \text{otherwise} \\ \end{cases}$$

I scratched my head for quite some time but could not figure out how to differentiate it.

Any hints from you are highly appreciated.

Answer 1

The easiest way is to take STFT using operators with autodiff support, e.g. via PyTorch. Then simply set the window as an updatable parameter, initializing as Hanning etc:

import torch.nn as nn
from scipy.signal import windows

win = torch.from_numpy(windows.hann(128))
win_t = nn.Parameter(win, requires_grad=True)

Then we'd do e.g.

def forward(self, x):
    x = STFT(x, win_t)
    x = self.conv(x)
    ...

It might help to apply a symmetry constraint, by e.g. defining weights of the window as only half of the initial window, but operating with the full window that uses the updated weights.

See this post for example with fully differentiable CWT applied on inversion. That network could be embedded in a conv net and its filter weights updated if we wrap them with nn.Parameter.

Optimizing for width

If the goal's to optimize a parameter of a predefined window function, then generate the window through the differentiable parameter:

def gauss(N, sigma):
    t = torch.linspace(-.5, .5, N)
    return torch.exp(-(t / sigma)**2)

N = x.shape[-1]
sigma = nn.Parameter(torch.tensor(0.1))
win_t = gauss(N, sigma)

Optimizing for length

To optimize for the number of samples of the window, we must ensure the sampling is differentiable.

• Hann does arange(N) / N, but N must be continuous, which requires rounding. torch.round isn't differentiable, but torch.clamp is.
• To ensure the process works, we define an optimization objective via normalized cross-correlation (NCC): the optimal window length will be that which matches a reference. We implement NCC via conv1d.
• Optimize via plain gradient descent, ensure to zero gradients after each step to avoid gradding the grads

Suppose we start with N=129 and optimum is N_ref=160. Results:

Code at Github.

Answer 2

Consider $N$ as a positive real number and use your formula for the Hann window, $w[n]$. Then the partial derivative of $w[n]$ with respect to $N$ will be:

$$\frac{\partial\,w[n]}{\partial N} = \begin{cases} - \dfrac{\pi\,n\,\sin\left(\frac{2\pi}{N - 1}n\right)}{{\left(N - 1\right)}^2} \qquad & 0 \le n \le N-1 \\ 0 & \text{otherwise.} \\ \end{cases}$$

At integer $N$, $w[n]$ has a very simple, sparse length-$N$ discrete Fourier transform (DFT). With non-integer $N$ that quality is lost, but also I don't know that you would need that quality in our application.