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.

  • $\begingroup$ In the optimization effort, what parameter are they trying to maximize or minimize by varying the window width? $\endgroup$ Commented Feb 19, 2022 at 21:19
  • $\begingroup$ And with the Gaussian window, there are two width parameters to think about. 1. the $\sigma$ parameter of the bell curve. 2. the actually cutoff point where the window goes to zero. $\endgroup$ Commented Feb 19, 2022 at 21:21
  • $\begingroup$ Yes, the σ parameter of the bell curve. $\endgroup$
    – JXuan
    Commented Feb 21, 2022 at 6:55
  • $\begingroup$ Well, if you have an optimum $\sigma$ for the optimum Gaussian, then you can fit a Hann or Hamming to it so that the second-derivative of the main lobe in the middle is the same for both the Hann and the optimal Gaussian. $\endgroup$ Commented Feb 21, 2022 at 8:00
  • $\begingroup$ I am not planning to find the optimum σ for Gaussian (this is what was done in the paper) but would like to optimize Hann. Are you suggesting to use the optimum σ of Gaussian to optimize Hann indirectly? $\endgroup$
    – JXuan
    Commented Feb 22, 2022 at 9:27

2 Answers 2


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:

enter image description here

Code at Github.

  • $\begingroup$ Great!! I did not know it's already implemented in pytorch. Many many thanks! $\endgroup$
    – JXuan
    Commented Feb 21, 2022 at 6:56
  • $\begingroup$ @JXuan Glad it helped, also added $\sigma$ case. If the problem's solved, consider voting & accepting. $\endgroup$ Commented Feb 21, 2022 at 20:11
  • $\begingroup$ Thanks for adding the σ case. For STFT, the window size is supposed to be the power of 2, instead of any real numbers, isn't it? If so, one could not just autodiff the win of Hann/Hamming window directly (because it is not differentiable)? $\endgroup$
    – JXuan
    Commented Feb 22, 2022 at 9:24
  • $\begingroup$ I got a reply from the contributors who implemented the torch stft and window functions in the Pytorch forum. They said 'The window size is not immediately trainable, although you could implement your own method of updating it.' $\endgroup$
    – JXuan
    Commented Feb 22, 2022 at 10:49
  • $\begingroup$ @JXuan Window can be of any length, power of 2 is typically for computational reasons. Yes, tuning something like width of Hanning will be harder, but is certainly doable. ssqueezepy has a torch implementation, but a key step isn't differentiable (buffer via CuPy) so it must be implemented via PyTorch - alternatively, code here can be ported easily to torch (though needs optimizing) $\endgroup$ Commented Feb 23, 2022 at 1:56

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.

  • $\begingroup$ Thanks for contributing the idea! What kind of quality loss do you mean specifically? $\endgroup$
    – JXuan
    Commented Feb 25, 2022 at 8:36
  • $\begingroup$ By "quality" I mean that "distinguishing attribute" with the DFT. $\endgroup$ Commented Feb 25, 2022 at 15:56

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.