# Custom rrule for Feedback Alignment

**URL:** <https://discourse.julialang.org/t/custom-rrule-for-feedback-alignment/72159>\
**Category:** Machine Learning\
**Tags:** question, flux, chainrulescore\
**Created:** [November 27, 2021, 11:52am UTC](https://discourse.julialang.org/t/custom-rrule-for-feedback-alignment/72159 "2021-11-27T11:52:35Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Rasmus\_Hoier](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rasmus_hoier/32/24036_2.png) [@Rasmus\_Hoier](https://discourse.julialang.org/u/Rasmus_Hoier)\
**Post date:** [November 27, 2021, 11:52am UTC](https://discourse.julialang.org/t/custom-rrule-for-feedback-alignment/72159/1 "2021-11-27T11:52:35Z")

</div>

Hi,  
I am trying to use a custom rrule to implement Feedback Alignment learning (which is just like BP except that the error signals are propagated backwards using random feedback weights B, which is supposed to be more biologically plausible than using W both in the forwards and backwards pass).

I made a simplified layer DenseFA (without nonlinearity and bias for simplicity), along with a custom rrule. the rrule seems to work when called on its own, but when I use Flux gradient function it is no longer used.

> **Boilerplate: imports and layer definition**
>
> ```julia
> using Flux; using Flux: @functor, glorot_uniform
> using ChainRulesCore; using ChainRulesCore: NoTangent, @thunk
> 
> # DenseFA struct
> struct DenseFA{M1<:Matrix, M2<:Matrix}
> W::M1 # Weights used in the forward pass
> B::M2 # Weights used in the backwards pass
> function DenseFA(W::M1, B::M2) where {M1<:Matrix, M2<:Matrix}
> new{M1,M2}(W, B)
> end
> end
> 
> # Initialize DenseFA struct
> function DenseFA(in::Integer, out::Integer; init = glorot_uniform)
> W = init(out, in)
> B = init(out, in)
> return DenseFA(W, B)
> end
> 
> @functor DenseFA
> 
> (a::DenseFA)(X) = my_matmul(a.W, X, a.B)
> 
> function my_matmul(W, X, B) 
> return W * X
> end
> 
> function Base.show(io::IO, l::DenseFA)
> print(io, "DenseFA(", size(l.W, 2), ", ", size(l.W, 1), ")")
> end
> 
> ```

> **The rrule and the gradient computation**
>
> ```julia
> function rrule(::typeof(my_matmul), W::Matrix, X::Matrix, B::Matrix)
> y = my_matmul(W, X, B)
> 
> println("==========Using rrule with fixed random feedback weights==========")
> 
> function times_pullback(ΔΩ)
> ∂W = @thunk(ΔΩ * X')
> ∂X = @thunk(B' * ΔΩ) # Use random feedback weight matrix B
> return (NoTangent(), ∂W, ∂X, NoTangent())
> end
> return y, times_pullback
> end
> 
> Flux.trainable(a::DenseFA) = (a.W,)
> 
> model = Chain(DenseFA(50, 60), DenseFA(60, 40), DenseFA(40, 2))
> 
> # dummy data (Batchsize 64)
> x = rand(Float32, 50, 64);
> y = rand(Float32, 2, 64)
> 
> # compute gradient
> loss = Flux.Losses.mse
> opt = Descent(0.1)
> ps = Flux.params(model);
> gs = gradient(() -> loss(model(x), y), ps)
> 
> ```

I could use advice on why the rrule is not being used and how I can make Flux use it.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 27, 2021, 1:56pm UTC](https://discourse.julialang.org/t/custom-rrule-for-feedback-alignment/72159/2 "2021-11-27T13:56:59Z")

</div>

What you’re missing is qualifying (or importing) rrule. At present this defines a new function of that name, instead (as you did for `show`) you need to add methods to the existing `ChainRulesCore.rrule`.

```julia
julia> rrule
rrule (generic function with 1 method)

julia> function ChainRulesCore.rrule(::typeof(my_matmul), W::Matrix, X::Matrix, B::Matrix)
          y = my_matmul(W, X, B)
          ...

```

---

<div class="post-metadata">

**Author:** ![Rasmus\_Hoier](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rasmus_hoier/32/24036_2.png) [@Rasmus\_Hoier](https://discourse.julialang.org/u/Rasmus_Hoier)\
**Post date:** [November 27, 2021, 8:10pm UTC](https://discourse.julialang.org/t/custom-rrule-for-feedback-alignment/72159/3 "2021-11-27T20:10:36Z")

</div>

Thanks, I had completely overlooked that!  
Now everything works as expected 😃