Basic linear algebra for complex numbers

[boeddeker]

🚀 Feature

Support basic linear algebra for complex numbers.

Motivation

I talked with @sw005320 about https://github.com/nttcslab-sp/dnn_wpe and it turns out, that the matrix inversion implemented with real numbers is unstable. In a beamforming example @Emrys365 observed a performance difference of 5 dB in a signal to distortion ratio (SDR) where he replaced the inversion with numpy code (torch: 5dB, numpy 10dB).

I tried torch.inverse and torch.solve and interestingly they are working in 1.6.0.dev20200623+cpu (Not mentioned in pytorch/pytorch#33152).
Is it possible, to support torch.matmul and some other linear algebra functions?

I also tried to use backward after torch.solve and the code fails with the exception msg, that matmul is not implemented.
Does someone know, how the gradient is defined in torch for complex numbers?
Is it grad_real + j grad_imag or grad_real - j grad_imag?
And how can I add/fix the gradient, when I find a broken implementation?

Pitch

• Using native functions for linear algebra of complex numbers instead of https://github.com/kamo-naoyuki/pytorch_complex
  • torch.solve and torch.inverse miss a backward function
  • torch.matmul does not work

Alternatives

• https://github.com/kamo-naoyuki/pytorch_complex has precision problems with the matrix inverse.

Additional context

Currently, I am considering to jump between pytorch_complex and torch.autograd.Function:

def hermite(a):
    return a.transpose(-2, -1).conj()

def matmul(t1, t2):
    real1, imag1 = t1.real, t1.imag
    real2, imag2 = t2.real, t2.imag
    o_real = torch.matmul(real1, real2) - torch.matmul(imag1, imag2)
    o_imag = torch.matmul(real1, imag2) + torch.matmul(imag1, real2)
    return o_real + 1j * o_imag

class Solve(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, b):
        x, _ = torch.solve(b, A)
        ctx.save_for_backward(A, x)
        return x

    @staticmethod
    def backward(ctx, grad_output):
        A, x = ctx.saved_tensors        
        gb, _ = torch.solve(grad_output, hermite(A))
        gA = - matmul(gb, hermite(x))
        return gA, gb

[sw005320]

@vincentqb and @mthrok,
This is one of the most serious issues when we try to use pytorch to some speech and audio processing during the JSALT workshop activities, as I mentioned before. Now, with the investigation from @boeddeker and @Emrys365, this problem becomes more serious, and also what we need is more clear. Of course, it's better to be supported by PyTorch for these basic operations in the long term future, but it would be great if torchaudio can handle this issue in the short term.

[vincentqb]

pytorch/pytorch#33152
cc @anjali411

[anjali411]

Hi @boeddeker thanks for creating this issue! torch.matmul is already supported for complex on cuda on master (#40431). I will add torch.matmul for complex numbers on CPU once the the TH to ATen port is complete (#40927).

For the expected autograd behavior for complex, checkout this doc note (to be merged) here: pytorch/pytorch#41012

[boeddeker]

Hi @anjali411, thanks for the feedback. Good to hear, that matmul is supported on the GPU.
Is torch.solve also planned to be implemented on the GPU?

How do I figure out, if a gradient is already implemented?
In pytorch/pytorch#33152 are no gradients mentioned (except gradient check).
Since the gradient torch.solve currently fails, I expect, that it uses the real gradient definition.

I am not sure, what the documentation in pytorch/pytorch#41012 tries to tell me.
I only knew two ways to define the complex derivative.
Assuming L is a loss value (real valued) and z = x + j y is a complex number, then the gradient can be:

• \partial L / \partial x + j \partial L / \partial y

or

• \partial L / \partial x - j \partial L / \partial y

You may add a factor in front of the derivative, but that would not change the chain rule.
Depending on which definition is used, some real valued gradients may be correct or most are incorrect.

In Table 4.1, 4.2 and 4.3 of [1] are the equations that can be used to obtain the gradients for both definitions and Equation 3.2 and 3.3 show the notation that is used in that work.
The difference between both definitions is only a conjugate transpose. So maybe a conj is sometimes missing.

[1] https://arxiv.org/abs/1701.00392

[anjali411]

Hi @anjali411, thanks for the feedback. Good to hear, that matmul is supported on the GPU.
Is torch.solve also planned to be implemented on the GPU?

Hi @boeddeker Yeah it's a part of the plan, but we can certainly make it high priority if it's commonly used. can you create an issue for that on pytorch repo?

In general, some of the ops are not fully tested for complex yet, but we will be working on making the support for complex more extensive in 1.7. are there any other functions that are commonly used by Audio Community and are not supported for complex numbers in PyTorch yet? We can certainly prioritize adding such functions.

How do I figure out, if a gradient is already implemented?

Here's a list of the functions for which backward is supported and tested: https://github.com/pytorch/pytorch/blob/dac63a13cbe105c81db4539147687d593704fe83/test/test_autograd.py#L4336

In pytorch/pytorch#33152 are no gradients mentioned (except gradient check).
Since the gradient torch.solve currently fails, I expect, that it uses the real gradient definition.

I am not sure, what the documentation in pytorch/pytorch#41012 tries to tell me.
I only knew two ways to define the complex derivative.
Assuming L is a loss value (real valued) and z = x + j y is a complex number, then the gradient can be:

• \partial L / \partial x + j \partial L / \partial y

or

• \partial L / \partial x - j \partial L / \partial y

The autograd support for complex numbers is still in the prototype stage but we follow the JAX's convention for autograd for complex numbers, wherever supported.

[boeddeker]

Hi @anjali411, I forgot to say, thank you for pushing complex support in pytorch.

Hi @boeddeker Yeah it's a part of the plan, but we can certainly make it high priority if it's commonly used. can you create an issue for that on pytorch repo?

Sure, I opened an issue in pytorch.

In general, some of the ops are not fully tested for complex yet, but we will be working on making the support for complex more extensive in 1.7. are there any other functions that are commonly used by Audio Community and are not supported for complex numbers in PyTorch yet? We can certainly prioritize adding such functions.

Currently, the exception RuntimeError: _th_addmm_out not supported on CPUType for ComplexDouble blocks me to do further test (On my notebook I have no GPU to do the tests on the GPU).

In the long term it would be nice to have the eigenvalue decomposition implemented with gradient support, but the implementation and the tests are quite complicated. So other operations are currently more interesting (e.g. fft with complex numbers (not tested if already implemented)).

The autograd support for complex numbers is still in the prototype stage but we follow the JAX's convention for autograd for complex numbers, wherever supported.

Thanks for the link. The information that I am interested in, is not directly on that page, but I found how pytorch defined the gradient. In Pytorch the optimizers have to add a conj when they do an optimizer step. I personally prefer the other definition, but both work.

[anjali411]

@sw005320 @boeddeker can you links to some functions/code blocks from the ESPnet repository that use complex numbers for a reference point? That would help in avoiding any last minute surprises for some functionalities that are not yet supported for complex numbers (for eg. torch.clamp etc) but we can prioritize them if we know about them ahead of time.

Also, I will we working on writing a tutorial for complex numbers and was curious to see if you have any ideas on application (s) I could cover in the tutorial that would demonstrate the importance of complex numbers as well as interesting as a unit application.

[boeddeker]

@anjali411 sure, I can give you two pointers.

• WIP: support torch fgnt/nara_wpe#46 is an example for a dereverberation algorithm (WPE) implemented with torch.
  The algorithm works in the STFT domain. I tested the code with real numbers, so I expect it should also work with complex numbers, when the support for the operations is available.
  • Application: Speedup the computation by using the GPU
• In https://github.com/espnet/espnet/tree/master/espnet/nets/pytorch_backend/frontends are some components that are used for an End-2-End training of speech enhancement/source separation and speech recognizer. The complex operations are in the beamformer and dereverberation algorithm (WPE, https://github.com/nttcslab-sp/dnn_wpe). To compensate the complex support they use https://github.com/kamo-naoyuki/pytorch_complex (Interpret complex operations as real operations) at the moment.
  • Application: Support End-2-End training of frontend and backend [1] in pytorch.
  • I haven't carefully analyzed, but using pytorch_complex increases the memory consumption (e.g. 1 matrix multiplication is splitted in multiple operations with the real and imag component)
  • Maybe @sw005320 can point you to an example (I got only some files to train an NN that are not under version control)

Maybe one motivation for complex support is the feature extraction (Beside to the E2E training).
In the audio domain (or at least in our subfield) most of the time the feature extraction involves complex numbers.
So currently the code is a mix of numpy and torch. It would be nice when everything can be done in torch.
e.g. When the feature extraction gets expensive and you do it on demand, the GPU can speedup the input pipeline.

[1] Beamnet: End-to-end training of a beamformer-supported multi-channel ASR system

[anjali411]

@anjali411 sure, I can give you two pointers.

• fgnt/nara_wpe#46 is an example for a dereverberation algorithm (WPE) implemented with torch.
  The algorithm works in the STFT domain. I tested the code with real numbers, so I expect it should also work with complex numbers, when the support for the operations is available.

which functions in this PR need to work with complex numbers? do all the functions also need autograd support?

• Application: Speedup the computation by using the GPU
• In https://github.com/espnet/espnet/tree/master/espnet/nets/pytorch_backend/frontends are some components that are used for an End-2-End training of speech enhancement/source separation and speech recognizer. The complex operations are in the beamformer and dereverberation algorithm (WPE, https://github.com/nttcslab-sp/dnn_wpe). To compensate the complex support they use https://github.com/kamo-naoyuki/pytorch_complex (Interpret complex operations as real operations) at the moment.
  This is actually helpful for referencing which functions are required for complex tensors. We already support many of these functions. I would encourage you to check out the complex support for these functions and give feedback!

For example, I would also like to draw your attention to the view functions we added very recently torch.view_as_real, torch.view_as_complex that are equivalent to the functions in this file: https://github.com/kamo-naoyuki/pytorch_complex/blob/master/torch_complex/utils.py.

I haven't carefully analyzed, but using pytorch_complex increases the memory consumption (e.g. 1 matrix multiplication is splitted in multiple operations with the real and imag component)

what do you mean by pytorch_complex? do you mean floating point tensors of shape (..., 2) to represent complex values?

Also I think it would be useful to setup a meeting and discuss more on how we can integrate torch complex tensors in ESPnet.

[boeddeker]

which functions in this PR need to work with complex numbers? do all the functions also need autograd support?

The first argument of the function wpe_v6in nara_wpe is complex valued. So most functions have to work with complex numbers.
Currently, the code fails with missing support for matmul.

When combining WPE with a NN the inverse_power will be estimated from an NN. So gradient support for mul, matmul, conj, transpose/permute, getitem, sub and solve would be necessary (I hope I didn't forget a function).

For example, I would also like to draw your attention to the view functions we added very recently torch.view_as_real, torch.view_as_complex that are equivalent to the functions in this file: https://github.com/kamo-naoyuki/pytorch_complex/blob/master/torch_complex/utils.py.

Thanks for this hint. I think they are slightly different. In torch_complex they have two tensors, one for real and one for imag and they are converted to a Tensor "(..., 2)", while torch.view_as_real expects a complex tensor (e.g. dtype=complex) as input. In torch_complex is the dtype never complex (so they support also old pytorch versions).

what do you mean by pytorch_complex? do you mean floating point tensors of shape (..., 2) to represent complex values?

Ah, sorry, that was a type, I meant torch_complex from kamo-naoyuki.
He splits the complex numbers in real and imag and does all operations just with real operations.
Example for multiplication: z = x + jy, w = u + jv, z * w = (xu - yv) + j(x * v + (y * u))
I haven't carefully thought about the intermediate tensors that are necessary to save for backward, but it may be higher than doing the complex multiplication with a single function (vs 4 multiplications and 2 additions).

Also I think it would be useful to setup a meeting and discuss more on how we can integrate torch complex tensors in ESPnet.

Sure, we can setup a meeting. My timezone is UTC+2 (Germany).
Do you have a suggestion when and where we meet?
@sw005320 Who should/wants to participate from ESPnet?

[mthrok]

@anjali411 @boeddeker I will coordinate the meeting then.

[mpariente]

I should have seen this earlier. Did the meeting happen? What were the conclusions?

I currently use this method to emulate complex tensors in pytorch and I see now that this is what is done in torch_complex as well.

I think svd and pinverse would be very good to have as well as it's more robust than inverse.

[mthrok]

Hi @mpariente

Thanks for your comment.

The meeting has happened and we followed the discussion above and narrowed down the requirements from ESPNet to the followings

• Gradient support for
  • mul
  • matmul
  • conj
  • transpose / permute
  • getitem
  • sub
  • solve
  • symeig / eigh
  • exp

We are interested in hearing complex number use cases and we are also interested in your work of asteroid. So how about we set up a meeting?

[mpariente]

Hi, thanks for your answer.

This would already be great to have !
I guess z.abs() would be z * z.conj() so there is no need to explicitly require the gradients?
How about log?

Yes, we could set up a meeting. I'm tagging @popcornell and @faroit here as they could also be interested.