Minimal multi-headed self-attention

multihead.jl

using Flux
using Fluxperimental: @compact

nf = 10
nb = 32
nt = 100
d_attn = 64
d_value = 128
d_head = 16
d_out = 256

X = randn(Float32, nf, nt, nb)
# X: (feature, time, batch)

# Minimal:
attn(i,o,a,dv,H)=@compact(out=Dense(H*dv=>o),heads=[@compact(D=l->Dense(i=>l),K=D(a),V=D(dv),Q=D(a))do x;k,v,q=K(x),V(x),Q(x);x=sum(k.*q,dims=1)./√a;softmax(x,dims=2).*v;end;for _∈1:H])do x;vcat([h(x) for h∈heads]...)|>out;end

# Expanded:
function multihead_self_attention(num_features, num_out, d_attn, d_value, num_heads)
    @compact(
        out = Dense(num_heads * d_value => num_out),
        heads = [
            @compact(
                K = Dense(num_features => d_attn),
                V = Dense(num_features => d_value),
                Q = Dense(num_features => d_attn)
            ) do x
                k, v, q = K(x), V(x), Q(x)
                x = sum(k .* q; dims=1) ./ sqrt(d_attn)
                softmax(x; dims=2) .* v
            end for _ in 1:num_heads
        ]
    ) do x
        out(vcat([h(x) for h in heads]...))
    end
end

attn(nf, d_out, d_attn, d_value, d_head)(X)
multihead_self_attention(nf, d_out, d_attn, d_value, d_head)(X)