Julia の SparseArrays で協調フィルタリング

Julia の標準ライブラリには疎行列を扱う SparseArrays が含まれています。
これを使ってメモリ効率の良い協調フィルタリングを実装してみます。

アルゴリズム

今回実装するのはコサイン類似度を用いたアイテムベース協調フィルタリングです。

ユーザ $u$ に対するアイテム $i$ の評価値を $r_{u,i}$ 、その推定値を $\hat{r}_{u,i}$ と表します。

$\hat{r}_{u,i}$ は $u$ が評価済みのアイテム集合 $I_u$ 、およびアイテム $x, y$ の類似度 $\operatorname{sim}(x,y)$ から求めます。

\hat{r}_{u,i} = \bar{r_u} + \frac{1}{|I_u|} \sum\limits_{i' \in I_u} \operatorname{sim}(i, i') (r_{u, i'} - \bar{r_u})

$\operatorname{sim}(x,y)$ はアイテム $x, y$ に対する中心化された評価ベクトルのコサイン類似度として求めます。

\operatorname{sim}(x,y) = \frac{
    \sum\limits_{u \in U_x \cap U_y} (r_{u,x} - \bar{r_x})(r_{u,y} - \bar{r_y})
  }{
    \sqrt{
      \sum\limits_{u \in U_x} (r_{u,x} - \bar{r_x})^2
    }
    \sqrt{
      \sum\limits_{u \in U_y} (r_{u,y} - \bar{r_y})^2
    }
}

ここでアイテム $x$, $y$ を評価しているユーザの集合をそれぞれ $U_x$, $U_y$ としています。

疎行列のメリット

全ユーザの全アイテムに対する推定評価値を一気に計算すると、ユーザ数 × アイテム数分のメモリ空間を消費してしまいます。
このため実際の推薦システムでは、評価行列と類似行列を先に計算しておき、必要になったタイミングで推定値を計算することが多いです。

評価行列は (ユーザ, アイテム) の評価値が存在する場合のみ値をもち、類似行列も (アイテムA, アイテムB) をともに評価しているユーザが存在する場合のみ値をもちます。
これらを疎行列で保持することで、メモリ消費を抑えることができます。

実装

評価行列

ユーザ u のアイテム i に対する評価値 r を、次のように評価行列 R で表します。

# R :: SparseMatrixCSC{Float64,Int64}
R[u, i] = r

類似行列

計算を簡単にするため、R の値をアイテム毎の平均評価値により中心化して D としておきます。

# bi :: SparseVector{Float64,Int64}
bi = sum(R, dims = 1)[:] ./ [nnz(R[:, i]) for i in 1:R.n]

# D :: SparseMatrixCSC{Float64,Int64}
us, is, vs = findnz(R)
D = sparse(us, is, vs - bi[is])

さらに、D の各列ベクトル (= アイテムの評価ベクトル) のノルムを計算し、D を正規化して wD としておきます。

# w :: SparseVector{Float64,Int64}
w = 1. ./ sqrt.(sum(D .^ 2, dims = 1)[:])
w[isinf.(w)] .= 0.
w = sparsevec(w)

# wD :: SparseMatrixCSC{Float64,Int64}
wD = w' .* D

これらの準備によりアイテムの類似行列 S は行列積として一気に計算できます。

# S :: SparseMatrixCSC{Float64,Int64}
S = wD' * wD

評価値の推定

ユーザ u に対する各アイテム推定評価値も、同じ要領で一気に計算できます。

まず u の各アイテムに対する評価値を中心化しておきます。

# ru :: SparseVector{Float64,Int64}
ru = R[u, :]

# b :: Float64
b = sum(ru) / nnz(ru)

# du :: SparseVector{Float64,Int64}
is, vs = findnz(ru)
du = sparsevec(is, vs .- bu)

以上で求めた D, bu , du によりユーザ u に対する各アイテムの評価値の予測値を計算できます。

# pu :: SparseVector{Float64,Int64}
n = 1 ./ nnz(ru)
pu = b .+ n .* S * du

これをソートし、上位のアイテムを求めれば推薦結果を出力できます。

is, vs = findnz(pu)
perm = partialsortperm(vs, 1:min(k, length(is)), rev = true)
recommendations, scores = is[perm], vs[perm]

実験

MovieLens の small dataset を使って簡単に実験してみました。
https://grouplens.org/datasets/movielens/

sim(R) の部分で 20万 × 20万 程度の疎行列同士の積を計算していますが、累積のメモリ割当は500MB程度なので、意図通り動作していそうです。

コード

using CSV
using DataFrames
using SparseArrays
using Statistics

function sim(R::SparseMatrixCSC{Float64,Int})
    bi = sum(R, dims = 1)[:] ./ [nnz(R[:, i]) for i in 1:R.n]

    us, is, vs = findnz(R)
    D = sparse(us, is, vs - bi[is])

    w = 1. ./ sqrt.(sum(D .^ 2, dims = 1)[:])
    w[isinf.(dnorms)] .= 0.
    w = sparsevec(w)

    wD = D .* w

    return wD' * wD
end

function recommend(R::SparseMatrixCSC{Float64,Int}, S::SparseMatrixCSC{Float64,Int}, u::Int64, k::Int64)
    ru = R[u, :]
    b = sum(ru) / nnz(ru)

    is, vs = findnz(ru)
    du = sparsevec(is, vs .- b)

    n = 1 ./ nnz(ru)
    pu = b .+ n .* (S * du)

    is, vs = findnz(pu)
    perm = partialsortperm(vs, 1:min(k, length(is)), rev = true)
    recommendations, scores = is[perm], vs[perm]
end

function main()
    println("Loading...")
    ratings = CSV.File("ratings.csv", types = [Int64, Int64, Float64, Int64]) |> DataFrame
    movies = CSV.File("movies.csv", types = [Int64, String, String]) |> DataFrame
    R = sparse(ratings[:userId], ratings[:movieId], ratings[:rating],
               maximum(ratings[:userId]), maximum(ratings[:movieId]), max)

    println("Calcurating...")
    S = sim(R)

    while true
        print("Input userId: ")
        user = parse(Int, readline())

        rated = ratings[ratings[:userId] .== user, :]
        len = size(rated, 1)

        println("Movies rated high:")
        rated_high = sort(rated, [:rating], rev = true)
        rated_high = join(rated_high[1:min(10, len), :], movies, on = :movieId)
        println(rated_high)

        println("Movies rated low:")
        rated_low = sort(rated, [:rating])
        rated_low = join(rated_low[1:min(10, len), :], movies, on = :movieId)
        println(rated_low)

        println("Movies recommended:")
        items, scores = recommend(R, S, user, 10)
        recommended = DataFrame(movieId = items, score = scores)
        recommended = join(recommended, movies, on = :movieId)
        println(recommended)
    end
end

main()

出力結果

$ julia main.jl
Loading...
(610, 193609)
Calcurating...
  5.256181 seconds (817.86 k allocations: 519.393 MiB, 3.65% gc time)

Input userId: 1

Movies rated high:
10×6 DataFrame
│ Row │ userId │ movieId │ rating  │ timestamp │ title                                     │ genres                         │
│     │ Int64  │ Int64   │ Float64 │ Int64     │ String                                    │ String                         │
├─────┼────────┼─────────┼─────────┼───────────┼───────────────────────────────────────────┼────────────────────────────────┤
│ 1   │ 1      │ 47      │ 5.0     │ 964983815 │ Seven (a.k.a. Se7en) (1995)               │ Mystery|Thriller               │
│ 2   │ 1      │ 50      │ 5.0     │ 964982931 │ Usual Suspects, The (1995)                │ Crime|Mystery|Thriller         │
│ 3   │ 1      │ 101     │ 5.0     │ 964980868 │ Bottle Rocket (1996)                      │ Adventure|Comedy|Crime|Romance │
│ 4   │ 1      │ 151     │ 5.0     │ 964984041 │ Rob Roy (1995)                            │ Action|Drama|Romance|War       │
│ 5   │ 1      │ 157     │ 5.0     │ 964984100 │ Canadian Bacon (1995)                     │ Comedy|War                     │
│ 6   │ 1      │ 163     │ 5.0     │ 964983650 │ Desperado (1995)                          │ Action|Romance|Western         │
│ 7   │ 1      │ 216     │ 5.0     │ 964981208 │ Billy Madison (1995)                      │ Comedy                         │
│ 8   │ 1      │ 231     │ 5.0     │ 964981179 │ Dumb & Dumber (Dumb and Dumber) (1994)    │ Adventure|Comedy               │
│ 9   │ 1      │ 260     │ 5.0     │ 964981680 │ Star Wars: Episode IV - A New Hope (1977) │ Action|Adventure|Sci-Fi        │
│ 10  │ 1      │ 333     │ 5.0     │ 964981179 │ Tommy Boy (1995)                          │ Comedy                         │

Movies rated low:
10×6 DataFrame
│ Row │ userId │ movieId │ rating  │ timestamp │ title                                        │ genres                                          │
│     │ Int64  │ Int64   │ Float64 │ Int64     │ String                                       │ String                                          │
├─────┼────────┼─────────┼─────────┼───────────┼──────────────────────────────────────────────┼─────────────────────────────────────────────────┤
│ 1   │ 1      │ 3176    │ 1.0     │ 964983504 │ Talented Mr. Ripley, The (1999)              │ Drama|Mystery|Thriller                          │
│ 2   │ 1      │ 1219    │ 2.0     │ 964983393 │ Psycho (1960)                                │ Crime|Horror                                    │
│ 3   │ 1      │ 2253    │ 2.0     │ 964981775 │ Toys (1992)                                  │ Comedy|Fantasy                                  │
│ 4   │ 1      │ 2338    │ 2.0     │ 964983546 │ I Still Know What You Did Last Summer (1998) │ Horror|Mystery|Thriller                         │
│ 5   │ 1      │ 2389    │ 2.0     │ 964983094 │ Psycho (1998)                                │ Crime|Horror|Thriller                           │
│ 6   │ 1      │ 2617    │ 2.0     │ 964982588 │ Mummy, The (1999)                            │ Action|Adventure|Comedy|Fantasy|Horror|Thriller │
│ 7   │ 1      │ 70      │ 3.0     │ 964982400 │ From Dusk Till Dawn (1996)                   │ Action|Comedy|Horror|Thriller                   │
│ 8   │ 1      │ 223     │ 3.0     │ 964980985 │ Clerks (1994)                                │ Comedy                                          │
│ 9   │ 1      │ 296     │ 3.0     │ 964982967 │ Pulp Fiction (1994)                          │ Comedy|Crime|Drama|Thriller                     │
│ 10  │ 1      │ 316     │ 3.0     │ 964982310 │ Stargate (1994)                              │ Action|Adventure|Sci-Fi                         │

Movies recommended:
  0.112969 seconds (188.18 k allocations: 13.848 MiB)
10×4 DataFrame
│ Row │ movieId │ score   │ title                             │ genres                                   │
│     │ Int64   │ Float64 │ String                            │ String                                   │
├─────┼─────────┼─────────┼───────────────────────────────────┼──────────────────────────────────────────┤
│ 1   │ 2899    │ 4.42608 │ Gulliver's Travels (1939)         │ Adventure|Animation|Children             │
│ 2   │ 2048    │ 4.42509 │ Great Mouse Detective, The (1986) │ Action|Animation|Children|Crime          │
│ 3   │ 2033    │ 4.42276 │ Black Cauldron, The (1985)        │ Adventure|Animation|Children|Fantasy     │
│ 4   │ 1804    │ 4.42236 │ Newton Boys, The (1998)           │ Crime|Drama                              │
│ 5   │ 157     │ 4.42173 │ Canadian Bacon (1995)             │ Comedy|War                               │
│ 6   │ 1226    │ 4.42106 │ Quiet Man, The (1952)             │ Drama|Romance                            │
│ 7   │ 2654    │ 4.41799 │ Wolf Man, The (1941)              │ Drama|Fantasy|Horror                     │
│ 8   │ 1024    │ 4.41686 │ Three Caballeros, The (1945)      │ Animation|Children|Musical               │
│ 9   │ 2090    │ 4.41226 │ Rescuers, The (1977)              │ Adventure|Animation|Children|Crime|Drama │
│ 10  │ 2993    │ 4.41144 │ Thunderball (1965)                │ Action|Adventure|Thriller                │