From ba696ea3dfec4cbe693bf06a84c75dc196077f5b Mon Sep 17 00:00:00 2001
From: Chin-Yun Yu <chin-yun.yu@qmul.ac.uk>
Date: Sat, 14 Sep 2024 17:02:47 +0800
Subject: [PATCH] feat: reduce computations in backprop of `lfilter` (#3831)

* feat: efficient backprop for a_coeffs

* update lfilter docstring

* doc: double quotation

* use dafx paper for reference as it includes more details

* revise docstring
---
 docs/source/refs.bib                   |  8 +++++++
 src/libtorchaudio/lfilter.cpp          | 31 ++++++++++++--------------
 src/torchaudio/functional/filtering.py |  3 ++-
 3 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/docs/source/refs.bib b/docs/source/refs.bib
index 0d853de3cb..9fc6c108c2 100644
--- a/docs/source/refs.bib
+++ b/docs/source/refs.bib
@@ -633,3 +633,11 @@ @article{forgione2021dynonet
   year={2021},
   publisher={Wiley Online Library}
 }
+
+@inproceedings{ycy2024diffapf,
+    title={Differentiable All-pole Filters for Time-varying Audio Systems},
+    author={Chin-Yun Yu and Christopher Mitcheltree and Alistair Carson and Stefan Bilbao and Joshua D. Reiss and György Fazekas},
+    booktitle={International Conference on Digital Audio Effects (DAFx)},
+    year={2024},
+    pages={345--352},
+}
diff --git a/src/libtorchaudio/lfilter.cpp b/src/libtorchaudio/lfilter.cpp
index e9b8d4eea2..454b2cbcda 100644
--- a/src/libtorchaudio/lfilter.cpp
+++ b/src/libtorchaudio/lfilter.cpp
@@ -151,7 +151,6 @@ class DifferentiableIIR : public torch::autograd::Function<DifferentiableIIR> {
     auto a_coeffs_normalized = saved[1];
     auto y = saved[2];
 
-    int64_t n_batch = x.size(0);
     int64_t n_channel = x.size(1);
     int64_t n_order = a_coeffs_normalized.size(1);
 
@@ -161,26 +160,24 @@ class DifferentiableIIR : public torch::autograd::Function<DifferentiableIIR> {
 
     namespace F = torch::nn::functional;
 
-    if (a_coeffs_normalized.requires_grad()) {
-      auto dyda = F::pad(
-          DifferentiableIIR::apply(-y, a_coeffs_normalized),
-          F::PadFuncOptions({n_order - 1, 0}));
-
-      da = F::conv1d(
-               dyda.view({1, n_batch * n_channel, -1}),
-               dy.reshape({n_batch * n_channel, 1, -1}),
-               F::Conv1dFuncOptions().groups(n_batch * n_channel))
-               .view({n_batch, n_channel, -1})
-               .sum(0)
-               .flip(1);
-    }
+    auto tmp =
+        DifferentiableIIR::apply(dy.flip(2).contiguous(), a_coeffs_normalized)
+            .flip(2);
 
     if (x.requires_grad()) {
-      dx =
-          DifferentiableIIR::apply(dy.flip(2).contiguous(), a_coeffs_normalized)
-              .flip(2);
+      dx = tmp;
     }
 
+    if (a_coeffs_normalized.requires_grad()) {
+      da = -torch::matmul(
+                tmp.transpose(0, 1).reshape({n_channel, 1, -1}),
+                F::pad(y, F::PadFuncOptions({n_order - 1, 0}))
+                    .unfold(2, n_order, 1)
+                    .transpose(0, 1)
+                    .reshape({n_channel, -1, n_order}))
+                .squeeze(1)
+                .flip(1);
+    }
     return {dx, da};
   }
 };
diff --git a/src/torchaudio/functional/filtering.py b/src/torchaudio/functional/filtering.py
index b4140447ef..97ff6040dc 100644
--- a/src/torchaudio/functional/filtering.py
+++ b/src/torchaudio/functional/filtering.py
@@ -999,7 +999,8 @@ def _lfilter_core(
 
 def lfilter(waveform: Tensor, a_coeffs: Tensor, b_coeffs: Tensor, clamp: bool = True, batching: bool = True) -> Tensor:
     r"""Perform an IIR filter by evaluating difference equation, using differentiable implementation
-    developed independently by *Yu et al.* :cite:`ismir_YuF23` and *Forgione et al.* :cite:`forgione2021dynonet`.
+    developed separately by *Yu et al.* :cite:`ismir_YuF23` and *Forgione et al.* :cite:`forgione2021dynonet`.
+    The gradients of ``a_coeffs`` are computed based on a faster algorithm from :cite:`ycy2024diffapf`.
 
     .. devices:: CPU CUDA