Julia 言語 非負テンソルトレイン分解の実装

研究で必要だったので実装しました． 参考にした論文はこちらです．

link.springer.com

このアルゴリズムでは，NMFが必要になるので取り急ぎ用意しておきます．

using LinearAlgebra
using Random

function KL(A, B)
    n, m = size(A)
    kl = 0.0
    for i = 1:n
        for j = 1:m
            kl += A[i,j] * log( A[i,j] / B[i,j] ) - A[i,j] + B[i,j]
        end
    end
    return kl
end

function nmf_euc(X, r ; max_iter=200, tol = 0.0001)
    m, n = size(X)
    W = rand(m, r)
    H = rand(r, n)

    error_at_init = norm(X - W*H)
    previous_error = error_at_init
    for iter in 1:max_iter
        H .=  H .* ( W' * X ) ./ ( W' * W * H  )
        W .=  W .* ( X  * H') ./ ( W  * H * H' )

        if tol > 0 && iter % 10 == 0
            error = norm(X - W*H)
            if (previous_error - error) / error_at_init < tol
                break
            end
            previous_error = error
        end
    end

    return W, H
end

function nmf_kl(X, r ; max_iter=200, tol = 0.0001)
    m, n = size(X)
    W = rand(m, r)
    H = rand(r, n)
    one_mn = ones(m, n)

    error_at_init = KL(X, W*H)
    previous_error = error_at_init
    for iter in 1:max_iter
        println(KL(X, W*H))
        H .=  H .* ( W' * ( X ./ (W*H))) ./ ( W' * one_mn )
        W .=  W .* ( (X ./ (W*H) ) * H') ./ ( one_mn * H' )

        if tol > 0 && iter % 10 == 0
            error = KL(X, W*H)
            if (previous_error - error) / error_at_init < tol
                break
            end
            previous_error = error
        end
    end

    return W, H
end

次に論文中のアルゴリズム１を実装します．

function NNTF(A, r; max_iter=200, tol = 0.0001)
    C = A
    r0 = 1
    d = ndims(A)
    n = size(A)
    G = Vector{Array{Float64}}(undef,d)
    for i = 1:(d-1)
        if i == 1
            C = reshape(C, (r0*n[i],:))
        else
            C = reshape(C, (r[i-1]*n[i],:))
        end
        W, H = nmf_euc(C, r[i], max_iter=max_iter, tol=tol)
        H .= norm(W).* H
        W .= W ./ norm(W)
        if i == 1
            G[i] = reshape(W, (1, n[i], r[i]))
        else
            G[i] = reshape(W, (r[i-1], n[i], r[i]))
        end
        C = H
    end
    G[d] = reshape(C, (r[d-1],n[d],1))
   
    return G
end

これで分解はOKなのですが，コアテンソル(train)からテンソルを再構成するための関数を用意しておきます．

outer(v...) = reshape(kron(reverse(v)...),length.(v))
function reconst_train(G)
    D = length(G)
    sizesG = size.(G)
    R = zeros(Int64,D)
    J = zeros(Int64,D)
    for d = 1:D
        R[d] = sizesG[d][1]
        J[d] = sizesG[d][2]
    end
    rs = ( (1:R[d]) for d in 1:D)
    
    term = zeros(Float64, J...)
    for r in product(rs...)
        v = Vector{Array{Float64}}(undef, D)
        for d = 1:D
            if d == 1
                v[d] = G[1][ 1, :, r[2]]
            elseif d == D
                v[d] = G[D][ r[D], :, 1]
            else
                v[d] = G[d][ r[d], :, r[d+1]]
            end
        end
        term += outer(v...)
    end

    return term
end

上の outer(v...)というのは手っ取り速い実装が分からなかったので，stackoverflowの強い人に教えてもらいました．

stackoverflow.com

実験

とりあえず人工データでランクを大きくしていきます．

A = rand(4,5,6,3)
G = NNTF(A,[1,1,1,1]);
@show norm(A - reconst_train(G)) / norm(A)

G = NNTF(A,[2,2,2,2]);
@show norm(A - reconst_train(G)) / norm(A)

G = NNTF(A,[5,5,5,5]);
@show norm(A - reconst_train(G)) / norm(A)

G = NNTF(A,[10,10,10,10]);
@show norm(A - reconst_train(G)) / norm(A)

G = NNTF(A,[25,25,25,25]);
@show norm(A - reconst_train(G)) / norm(A)

#norm(A - reconst_train(G)) / norm(A) = 0.503555751152826
#norm(A - reconst_train(G)) / norm(A) = 0.46242914539607927
#norm(A - reconst_train(G)) / norm(A) = 0.36833865940558996
#norm(A - reconst_train(G)) / norm(A) = 0.261642993995869
#norm(A - reconst_train(G)) / norm(A) = 0.018216782089987324

ランクを大きくすると徐々に再構成誤差が下がっているので，問題なさそうですね．

次に，実データでも実験してみます．とりあえず SIPIデータセットの Baboon を使います．オリジナルのデータサイズは 3×256×256 ですが，これを reshape して (3,16,16,16,16) にして，ターゲットランクを(R,R,R,R,R)としてトレイン分解します．Rを大きくしていくと元の画像が良く復元できていることが分かります．ちなみにR=25でRMSE=0.20くらいでした．