Numerical Tours
← All Python tours
Download notebook Open in Colab

Learning Sparse Dictionaries

Learn an overcomplete collection of image atoms by alternating sparse coding and dictionary updates. Inspect the learned atoms and the objective history to connect the optimization problem with recurring local image structures.

Run this tour

Run the cells in order with a Python 3 kernel. The first cell locates the companion data and toolbox and installs missing dependencies when needed. All worked examples include their implementation directly in this notebook. Random seeds make comparisons reproducible; you can change them to explore other samples.

# Locate the companion toolbox locally, or fetch it for a standalone/Colab copy.
from pathlib import Path
import importlib.util
import os
import subprocess
import sys

working = Path.cwd()
candidates = [working, working / "python", working.parent / "python"]
python_dir = next((p for p in candidates if (p / "nt_toolbox").is_dir()), None)
if python_dir is None:
    checkout = working / "numerical-tours-support"
    if not checkout.exists():
        subprocess.run(
            [
                "git",
                "clone",
                "--depth",
                "1",
                "--branch",
                "master",
                "https://github.com/gpeyre/numerical-tours.git",
                str(checkout),
            ],
            check=True,
        )
    python_dir = checkout / "python"
os.chdir(python_dir)
if str(python_dir) not in sys.path:
    sys.path.insert(0, str(python_dir))
requirements = python_dir / "requirements.txt"
if any(
    importlib.util.find_spec(name) is None
    for name in [
        "numpy",
        "scipy",
        "matplotlib",
        "skimage",
        "sklearn",
        "pywt",
        "ipywidgets",
        "cvxpy",
        "skfmm",
        "autograd",
        "progressbar",
        "celer",
    ]
):
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-r", str(requirements)], check=True
    )

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
plt.rcParams.update(
    {
        "figure.figsize": (8, 4),
        "figure.dpi": 100,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "font.size": 11,
        "image.cmap": "gray",
    }
)
%matplotlib inline

$\newcommand{\dotp}[2]{\langle #1, #2 \rangle}$ $\newcommand{\enscond}[2]{\lbrace #1, #2 \rbrace}$ $\newcommand{\pd}[2]{ \frac{ \partial #1}{\partial #2} }$ $\newcommand{\umin}[1]{\underset{#1}{\min}\;}$ $\newcommand{\umax}[1]{\underset{#1}{\max}\;}$ $\newcommand{\uargmin}[1]{\underset{#1}{argmin}\;}$ $\newcommand{\norm}[1]{\|#1\|}$ $\newcommand{\abs}[1]{\left|#1\right|}$ $\newcommand{\choice}[1]{ \left\{ \begin{array}{l} #1 \end{array} \right. }$ $\newcommand{\pa}[1]{\left(#1\right)}$ $\newcommand{\diag}[1]{{diag}\left( #1 \right)}$ $\newcommand{\qandq}{\quad\text{and}\quad}$ $\newcommand{\qwhereq}{\quad\text{where}\quad}$ $\newcommand{\qifq}{ \quad \text{if} \quad }$ $\newcommand{\qarrq}{ \quad \Longrightarrow \quad }$ $\newcommand{\ZZ}{\mathbb{Z}}$ $\newcommand{\CC}{\mathbb{C}}$ $\newcommand{\RR}{\mathbb{R}}$ $\newcommand{\EE}{\mathbb{E}}$ $\newcommand{\Zz}{\mathcal{Z}}$ $\newcommand{\Ww}{\mathcal{W}}$ $\newcommand{\Vv}{\mathcal{V}}$ $\newcommand{\Nn}{\mathcal{N}}$ $\newcommand{\NN}{\mathcal{N}}$ $\newcommand{\Hh}{\mathcal{H}}$ $\newcommand{\Bb}{\mathcal{B}}$ $\newcommand{\Ee}{\mathcal{E}}$ $\newcommand{\Cc}{\mathcal{C}}$ $\newcommand{\Gg}{\mathcal{G}}$ $\newcommand{\Ss}{\mathcal{S}}$ $\newcommand{\Pp}{\mathcal{P}}$ $\newcommand{\Ff}{\mathcal{F}}$ $\newcommand{\Xx}{\mathcal{X}}$ $\newcommand{\Mm}{\mathcal{M}}$ $\newcommand{\Ii}{\mathcal{I}}$ $\newcommand{\Dd}{\mathcal{D}}$ $\newcommand{\Ll}{\mathcal{L}}$ $\newcommand{\Tt}{\mathcal{T}}$ $\newcommand{\si}{\sigma}$ $\newcommand{\al}{\alpha}$ $\newcommand{\la}{\lambda}$ $\newcommand{\ga}{\gamma}$ $\newcommand{\Ga}{\Gamma}$ $\newcommand{\La}{\Lambda}$ $\newcommand{\Si}{\Sigma}$ $\newcommand{\be}{\beta}$ $\newcommand{\de}{\delta}$ $\newcommand{\De}{\Delta}$ $\newcommand{\phi}{\varphi}$ $\newcommand{\th}{\theta}$ $\newcommand{\om}{\omega}$ $\newcommand{\Om}{\Omega}$

Instead of using a fixed data representation such as wavelets or Fourier, one can learn the representation (the dictionary) to optimize the sparsity of the representation for a large class of exemplar.

%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np

from nt_toolbox.signal import load_image, imageplot, plot_dictionary
from nt_toolbox.general import crop

Dictionary Learning as a Non-convex Optimization Problem

Given a set $Y = (y_j)_{j=1}^m \in \RR^{n \times m} $ of $m$ signals $y_j \in \RR^m$, dictionary learning aims at finding the best dictionary $D=(d_i)_{i=1}^p$ of $p$ atoms $d_i \in \RR^n$ to sparse code all the data.

In this numerical tour, we consider an application to image denoising, so that each $y_j \in \RR^n$ is a patch of size $n=w \times w$ extracted from the noisy image.

The idea of learning dictionaries to sparse code image patch was first proposed in:

Olshausen BA, and Field DJ., <http://www.nature.com/nature/journal/v381/n6583/abs/381607a0.html Emergence of Simple-Cell Receptive Field Properties by Learning a Sparse Code for Natural Images.> Nature, 381: 607-609, 1996.

The sparse coding of a single data $y=y_j$ for some $j=1,\ldots,m$ is obtained by minimizing a $\ell^0$ constrained optimization $$ \umin{ \norm{x}_0 \leq k } \frac{1}{2}\norm{y-Dx}^2 . $$ where the $\ell^0$ pseudo-norm of $x \in \RR^p$ is $$ \norm{x}_0 = \abs{\enscond{i}{x(i) \neq 0}}. $$

The parameter $k>0$ controls the amount of sparsity.

Dictionary learning performs an optimization both on the dictionary $D$ and the set of coefficients $ X = (x_j)_{j=1}^m \in \RR^{p \times m} $ where, for $j=1,\ldots,m$, $ x_j $ is the set of coefficients of the data $y_j$. This joint optimization reads $$ \umin{ D \in \Dd, X \in \Xx_k } E(X,D) = \frac{1}{2}\norm{Y-DX}^2 = \frac{1}{2} \sum_{j=1}^m \norm{y_j - D x_j}^2. $$

The constraint set on $D$ reads $$ \Dd = \enscond{D \in \RR^{n \times p} }{ \forall i=1,\ldots,p, \quad \norm{D_{\cdot,i}} \leq 1 }, $$ (the columns of the dictionary are unit normalized). The sparsity constraint set on $X$ reads $$ \Xx_k = \enscond{X \in \RR^{p \times m}}{ \forall j, \: \norm{X_{\cdot,j}}_0 \leq k }. $$

We propose to use a block-coordinate descent method to minimize $E$: $$ X^{(\ell+1)} \in \uargmin{X \in \Xx_k} E(X,D^{(\ell)}), $$ $$ D^{(\ell+1)} \in \uargmin{D \in \Dd} E(X^{(\ell+1)},D). $$

One can show the convergence of this minimization scheme, see for instance

P. Tseng, <http://www.math.washington.edu/~tseng/papers/archive/bcr_jota.pdf Convergence of Block Coordinate Descent Method for Nondifferentiable Minimization>, J. Optim. Theory Appl., 109, 2001, 475-494.

We now define the parameter of the problem.

Width $w$ of the patches.

w = 10

Dimension $n= w \times w$ of the data to be sparse coded.

n = w * w

Number of atoms $p$ in the dictionary.

p = 2 * n

Number $m$ of patches used for the training.

m = 20 * p

Target sparsity $k$.

k = 4

Patch Extraction

Since the learning is computationnaly intensive, one can only apply it to small patches extracted from an image.

f = crop(load_image("nt_toolbox/data/barb.bmp"), 256)
n0 = f.shape[0]

Display the input image.

plt.figure(figsize=(5, 5))
imageplot(f)
No description has been provided for this image

Random patch location.

q = 3 * m

# Random sampling of coordinates of the top left corner or the patches
x = (np.random.random((1, 1, q)) * (n0 - w)).astype(int)
y = (np.random.random((1, 1, q)) * (n0 - w)).astype(int)

Extract lots of patches $y_j \in \RR^n$, and store them in a matrix $Y=(y_j)_{j=1}^m$.

[dY, dX] = np.meshgrid(range(w), range(w))
dX = np.tile(dX, (q, 1, 1)).transpose((1, 2, 0))
dY = np.tile(dY, (q, 1, 1)).transpose((1, 2, 0))
Xp = np.tile(x, (w, w, 1)) + dX
Yp = dY + np.tile(y, (w, w, 1))

# Extract patches
Y = f.flatten()[Yp + Xp * n0]
Y = Y.reshape((w * w, q))

We remove the mean, since we are going to learn a dictionary of zero-mean and unit norm atom.

Y -= Y.mean(axis=0)

Only keep those with largest energy.

energies = np.sum(Y**2, axis=0)
indexes = np.argsort(energies)[::-1]
Y = Y[:, indexes[:m]];

We consider a dictionary $D \in \RR^{n \times p} $ of $p \geq n$ atoms in $\RR^n$. The initial dictionary $D$ is computed by a random selection of patches, and we normalize them to be unit-norm.

def projC(Y):
    """Scale the patches to unit norm"""
    norm = np.tile(np.linalg.norm(Y, axis=0), (Y.shape[0], 1))
    Y = np.divide(Y, norm)
    return Y


# Pick p indexes at random to pick patches from Y
sel = np.random.permutation(range(m))[:p]
D = projC(Y[:, sel])
D0 = D.copy()

Display the initial dictionary.

plt.figure(figsize=(10, 6))
plot_dictionary(D)
No description has been provided for this image

Update of the Coefficients $X$

The optimization on the coefficients $X$ requires, for each $y_j = Y_{\cdot,j}$ to compute $x_j = X_{\cdot,j}$ that solves $$ \umin{ \norm{x_j}_0 \leq k } \frac{1}{2} \norm{y-D x_j}^2. $$

This is a non-smooth and non-convex minimization, that can be shown to be NP-hard. A heuristic to solve this method is to compute a stationary point of the energy using the Foward-Backward iterative scheme (projected gradient descent): $$ x_j \leftarrow \text{Proj}_{\Xx_k}\pa{ x_j - \tau D^* ( D x_j - y ) }$ \qwhereq \tau < \frac{2}{\norm{D D^*}}. $$

Denoting $\abs{\bar x(1)} \leq \ldots \leq \abs{\bar x(n)}$ the ordered magnitudes of a vector $ x \in \RR^n $, the orthogonal projector on $\Xx_k$ reads $z = \text{Proj}_{\Xx_k}(x)$ with $$ \forall i=1,\ldots,n, \quad z(i) = \choice{ x(i) \qifq \abs{x(i)} \geq \abs{\bar x(k)}, \\ z(i) = 0 \quad \text{otherwise}. }$ $$

def projX(X, k):
    """Sparsity projection, keeps the k largest coefficients"""
    X = X * (abs(X) >= np.sort(abs(X), axis=0)[-k, :])
    return X

Worked example 1

Perform the iterative hard thresholding, and display the decay of the energy $J(x_j) = \norm{y_j-D x_j}^2$ for several $j$. Remark: note that the iteration can be performed in parallel on all $x_j$.

"""
Sparse code data Y using dictionary D using a forward backward iterative
scheme (projected block coordinate gradient descent).
"""
n_iter = 100
gamma = 1.6 / np.linalg.norm(D, ord=2) ** 2
E = np.zeros((n_iter, m))
X = np.zeros((p, m))
for i in range(n_iter):
    R = np.dot(D, X) - Y
    E[i, :] = np.sum(R**2, axis=0)
    X = projX(X - gamma * np.dot(D.T, R), k)

sel = range(5)
plt.plot(np.log10(E[: len(E) // 2, sel] - np.min(E[:, sel], axis=0)))
plt.axis("tight")
plt.title("$log_{10}(J(x_j) - J(x_j^*))$")
plt.show()
No description has been provided for this image

Update the Dictionary $D$

Once the sparse coefficients $X$ have been computed, one can udpate the dictionary. This is achieve by performing the minimization $$ \umin{D \in \Dd} \frac{1}{2}\norm{Y-D X}^2. $$

One can perform this minimization with a projected gradient descent $$ D \leftarrow \text{Proj}_{\Cc}\pa{ D - \tau (DX - Y)X^* } $$ where $ \tau < 2/\norm{XX^*}. $

Note that the orthogonal projector $\text{Proj}_{\Cc}$ is implemented in the function |ProjC| already defined.

Worked example 2

Perform this gradient descent, and monitor the decay of the energy.

n_iter = 100
E = np.zeros(n_iter)
tau = 1 / np.linalg.norm(np.dot(X, X.T), ord=2)
for i in range(n_iter):
    R = np.dot(D, X) - Y
    E[i] = np.sum(R**2)
    D = projC(D - tau * np.dot(R, X.T))

plt.plot(np.log10(E[: len(E) // 2] - np.min(E)))
plt.axis("tight")
plt.show()
No description has been provided for this image

Worked example 3

Perform the dictionary learning by iterating between sparse coding and dictionary update.

n_iter_learning = 10
n_iter_dico = 50
n_iter_coef = 100
E0 = np.zeros(2 * n_iter_learning)
X = np.zeros((p, m))
D = D0.copy()
for i in range(n_iter_learning):
    # --- coefficient update ----
    E = np.zeros((n_iter_coef, m))
    gamma = 1.6 / np.linalg.norm(D, ord=2) ** 2
    for j in range(n_iter_coef):
        R = np.dot(D, X) - Y
        E[j, :] = np.sum(R**2, axis=0)
        X = projX(X - gamma * np.dot(D.T, R), k)
    E0[2 * i] = np.linalg.norm(Y - np.dot(D, X)) ** 2
    # --- dictionary update ----
    E = np.zeros(n_iter_dico)
    tau = 1 / np.linalg.norm(np.dot(X, X.T), ord=2)
    for j in range(n_iter_dico):
        R = np.dot(D, X) - Y
        E[i] = np.sum(R**2)
        D = projC(D - tau * np.dot(R, X.T))
    E0[2 * i + 1] = np.linalg.norm(Y - np.dot(D, X)) ** 2

index = list(range(E0.shape[0]))
index_coef = list(range(0, E0.shape[0], 2))
index_dict = list(range(1, E0.shape[0], 2))
plt.plot(np.divide(index, 2), E0, label="$|Y-DX|^2$")
plt.plot(
    np.divide(index_coef, 2),
    E0[index_coef],
    "*",
    markersize=7,
    label="After coefficient update",
)
plt.plot(
    np.divide(index_dict, 2),
    E0[index_dict],
    "o",
    markersize=5,
    label="After dictionary update",
)
plt.legend(numpoints=1)
plt.axis("tight")
plt.show()
No description has been provided for this image

Display the dictionary.

plt.figure(figsize=(10, 6))
plot_dictionary(D)
No description has been provided for this image

References and further reading