VKВлад КуликовОткрытые лекции
Теория информации для машинного обучения9 / 16

Модуль 9. Скорость–искажение и Information Bottleneck: как управлять потерей информации

Сначала попробуйте решить упражнение. Ответы, подробные решения и проектные ориентиры раскрываются отдельно и сохраняют своё состояние при переключении RU/EN.

Скачать файлы курса (.zip) 33,9 КБ

Архив содержит девять Python-файлов решений, общих для RU/EN, с сохранением структуры каталогов.

Упражнение 1. Скорость–искажение для источника Бернулли

Пусть X∼Bernoulli⁡(p)X\sim\operatorname{Bernoulli}(p), p≤1/2p\le1/2, и

d(x,x^)=1{x≠x^}.d(x,\widehat x)=\mathbf1\{x\ne\widehat x\}.
  1. Докажите нижнюю границу

    I(X;X^)≥h2(p)−h2(D)I(X;\widehat X)\ge h_2(p)-h_2(D)

    для 0≤D≤p0\le D\le p.

  2. Постройте тестовый канал, достигающий равенства.

  3. Объясните, почему R(D)=0R(D)=0 при D≥pD\ge p.

Показать ответ и решение

Пусть

Pe=P(X≠X^)≤D.P_e=P(X\ne\widehat X)\le D.

Поскольку XX бинарна, знание X^\widehat X и индикатора ошибки E=1{X≠X^}E=\mathbf1\{X\ne\widehat X\} однозначно определяет XX. Поэтому

H(X∣X^)≤H(E∣X^)≤H(E)=h2(Pe).H(X\mid\widehat X) \le H(E\mid\widehat X) \le H(E) =h_2(P_e).

При 0≤D≤p≤1/20\le D\le p\le1/2 функция h2h_2 возрастает, поэтому

H(X∣X^)≤h2(Pe)≤h2(D).H(X\mid\widehat X) \le h_2(P_e) \le h_2(D).

Следовательно,

I(X;X^)=H(X)−H(X∣X^)≥h2(p)−h2(D).\boxed{ I(X;\widehat X) =H(X)-H(X\mid\widehat X) \ge h_2(p)-h_2(D) }.

Чтобы достичь равенства, удобнее задать обратный тестовый канал. Пусть

X^∼Bernoulli⁡(q),Z∼Bernoulli⁡(D),\widehat X\sim\operatorname{Bernoulli}(q), \qquad Z\sim\operatorname{Bernoulli}(D),

независимы, и

X=X^⊕Z.X=\widehat X\oplus Z.

Тогда P(X≠X^)=DP(X\ne\widehat X)=D, а

P(X=1)=D+q(1−2D).P(X=1)=D+q(1-2D).

Выбираем

q=p−D1−2D,q=\frac{p-D}{1-2D},

что допустимо при 0≤D≤p≤1/20\le D\le p\le1/2 (граничный случай p=D=1/2p=D=1/2 рассматривается отдельно). Тогда маргинальное распределение XX равно Bernoulli⁡(p)\operatorname{Bernoulli}(p), а

H(X∣X^)=H(Z)=h2(D).H(X\mid\widehat X)=H(Z)=h_2(D).

Следовательно,

I(X;X^)=h2(p)−h2(D).I(X;\widehat X)=h_2(p)-h_2(D).

Прямой канал P(x^∣x)P(\widehat x\mid x) получается из построенного совместного распределения по формуле Байеса.

Если D≥pD\ge p, декодер может всегда выдавать X^=0\widehat X=0. Тогда скорость равна нулю, а среднее искажение равно P(X=1)=p≤DP(X=1)=p\le D. Поэтому

R(D)=0,D≥p.R(D)=0, \qquad D\ge p.

Упражнение 2. Предиктивное искажение IB

Для марковской цепи Y→X→TY\to X\to T докажите

I(X;Y)−I(T;Y)=EX,TDKL(PY∣X∥PY∣T).I(X;Y)-I(T;Y) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X}\|P_{Y\mid T}\bigr).

Покажите, что минимизация I(X;T)−βI(T;Y)I(X;T)-\beta I(T;Y) эквивалентна лагранжиану скорость–искажение с этим искажением.

Показать ответ и решение

Из цепного правила:

I(X,T;Y)=I(T;Y)+I(X;Y∣T).I(X,T;Y) =I(T;Y)+I(X;Y\mid T).

С другой стороны,

I(X,T;Y)=I(X;Y)+I(T;Y∣X).I(X,T;Y) =I(X;Y)+I(T;Y\mid X).

Для марковской цепи Y→X→TY\to X\to T выполняется

I(T;Y∣X)=0.I(T;Y\mid X)=0.

Поэтому

I(X;Y)−I(T;Y)=I(X;Y∣T).I(X;Y)-I(T;Y)=I(X;Y\mid T).

Условную взаимную информацию можно записать как среднюю KL-дивергенцию:

I(X;Y∣T)=EX,TDKL(PY∣X,T∥PY∣T).I(X;Y\mid T) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X,T}\|P_{Y\mid T}\bigr).

Марковость даёт PY∣X,T=PY∣XP_{Y\mid X,T}=P_{Y\mid X}, следовательно,

I(X;Y)−I(T;Y)=EX,TDKL(PY∣X∥PY∣T).\boxed{ I(X;Y)-I(T;Y) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X}\|P_{Y\mid T}\bigr) }.

Обозначив правую KL-дивергенцию через dIB(X,T)d_{\mathrm{IB}}(X,T), получаем

I(X;T)−βI(T;Y)=I(X;T)+βE[dIB(X,T)]−βI(X;Y).\begin{aligned} I(X;T)-\beta I(T;Y) &=I(X;T) +\beta\mathbb E[d_{\mathrm{IB}}(X,T)] -\beta I(X;Y). \end{aligned}

Последний член не зависит от PT∣XP_{T\mid X}. Поэтому минимизаторы совпадают с минимизаторами лагранжиана скорость–искажение

I(X;T)+βE[dIB(X,T)].I(X;T)+\beta\mathbb E[d_{\mathrm{IB}}(X,T)].

Упражнение 3. Вариационные границы VIB

  1. Докажите

    I(T;Y)≥H(Y)+E[log⁡rψ(Y∣T)].I(T;Y) \ge H(Y)+\mathbb E[\log r_\psi(Y\mid T)].
  2. Докажите

    EXDKL(q(T∣X)∥r(T))=I(X;T)+DKL(q(T)∥r(T)).\mathbb E_XD_{\mathrm{KL}}(q(T\mid X)\|r(T)) = I(X;T)+D_{\mathrm{KL}}(q(T)\|r(T)).
  3. Объясните, когда обе границы становятся точными одновременно.

Показать ответ и решение

1. Нижняя граница на релевантность

Для любого декодера rψ(y∣t)r_\psi(y\mid t):

E[−log⁡rψ(Y∣T)]=H(Y∣T)+ETDKL(PY∣T∥rψ(⋅∣T))≥H(Y∣T).\begin{aligned} \mathbb E[-\log r_\psi(Y\mid T)] &= H(Y\mid T) + \mathbb E_T D_{\mathrm{KL}} \bigl(P_{Y\mid T}\|r_\psi(\cdot\mid T)\bigr)\\ &\ge H(Y\mid T). \end{aligned}

Следовательно,

I(T;Y)=H(Y)−H(Y∣T)≥H(Y)+E[log⁡rψ(Y∣T)].\begin{aligned} I(T;Y) &=H(Y)-H(Y\mid T)\\ &\ge H(Y)+\mathbb E[\log r_\psi(Y\mid T)]. \end{aligned}

Граница становится точной тогда и только тогда, когда

rψ(y∣t)=P(y∣t)r_\psi(y\mid t)=P(y\mid t)

почти всюду по PTP_T.

2. Верхняя граница на скорость

Пусть

q(t)=∫q(t∣x)p(x) dx.q(t)=\int q(t\mid x)p(x)\,dx.

Тогда

EXDKL(q(T∣X)∥r(T))=EX,Tlog⁡q(t∣x)r(t)=EX,Tlog⁡q(t∣x)q(t)+ETlog⁡q(t)r(t)=Iq(X;T)+DKL(q(T)∥r(T)).\begin{aligned} &\mathbb E_X D_{\mathrm{KL}}(q(T\mid X)\|r(T))\\ &=\mathbb E_{X,T} \log\frac{q(t\mid x)}{r(t)}\\ &=\mathbb E_{X,T} \log\frac{q(t\mid x)}{q(t)} + \mathbb E_T\log\frac{q(t)}{r(t)}\\ &=I_q(X;T)+D_{\mathrm{KL}}(q(T)\|r(T)). \end{aligned}

Поэтому KL к выбранному априорному распределению является верхней границей на I(X;T)I(X;T).

3. Когда обе границы точны

Одновременно требуется:

rψ(y∣t)=P(y∣t)r_\psi(y\mid t)=P(y\mid t)

и

r(t)=q(t).r(t)=q(t).

Первое означает идеальный вариационный декодер, второе — совпадение априорного распределения с агрегированным апостериорным распределением. В реальной параметрической модели оба условия могут быть недостижимы.


Упражнение 4. VAE: правдоподобие задаёт искажение

Рассмотрите декодер

pθ(x∣z)=N(μθ(z),σ2I).p_\theta(x\mid z) = \mathcal N(\mu_\theta(z),\sigma^2I).
  1. Выведите отрицательный логарифм правдоподобия и покажите его связь с MSE.

  2. Объясните, почему изменение σ2\sigma^2 меняет относительный вес искажения и скорости даже при β=1\beta=1.

  3. Выведите разложение

    R=I(X;Z)+DKL(q(Z)∥p(Z)).\mathcal R = I(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z)).
Показать ответ и решение

Для dd-мерного наблюдения

pθ(x∣z)=(2πσ2)−d/2exp⁡[−∥x−μθ(z)∥22σ2].p_\theta(x\mid z) =(2\pi\sigma^2)^{-d/2} \exp\left[-\frac{\|x-\mu_\theta(z)\|^2}{2\sigma^2}\right].

Поэтому

−log⁡pθ(x∣z)=d2log⁡(2πσ2)+∥x−μθ(z)∥22σ2.\boxed{ -\log p_\theta(x\mid z) = \frac d2\log(2\pi\sigma^2) + \frac{\|x-\mu_\theta(z)\|^2}{2\sigma^2} }.

При фиксированной σ2\sigma^2 минимизация отрицательного логарифма правдоподобия эквивалентна минимизации MSE. Но коэффициент перед MSE равен 1/(2σ2)1/(2\sigma^2). Уменьшая σ2\sigma^2, мы усиливаем относительный вес члена реконструкции по сравнению с KL-скорость даже при формальном β=1\beta=1.

Для разложения скорости добавим и вычтем log⁡q(z)\log q(z):

R=EX,Zlog⁡q(z∣x)p(z)=EX,Zlog⁡q(z∣x)q(z)+EZlog⁡q(z)p(z)=Iq(X;Z)+DKL(q(Z)∥p(Z)).\begin{aligned} \mathcal R &=\mathbb E_{X,Z} \log\frac{q(z\mid x)}{p(z)}\\ &=\mathbb E_{X,Z} \log\frac{q(z\mid x)}{q(z)} + \mathbb E_Z\log\frac{q(z)}{p(z)}\\ &=I_q(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z)). \end{aligned}

То есть KL-скорость одновременно платит за информацию о конкретном объекте и за несовпадение агрегированного апостериорного распределения с априорным распределением генеративной модели.


Упражнение 5. Квантизация: какая задача скорость–искажение задана?

Для квантизации только весов перечислите:

  1. что является алфавитом источника;

  2. что входит в фактическую скорость;

  3. три возможные меры искажения;

  4. почему гауссовскую границу по MSE нельзя напрямую превратить в границу на точность в последующей задаче.

Показать ответ и решение

Возможный аккуратный ответ.

Алфавит источника

Источник нужно определить явно. Это могут быть:

  • отдельные веса;
  • блоки весов;
  • строки или столбцы матрицы;
  • группы каналов вместе с их статистиками активаций.

В классической постановке также требуется вероятностный ансамбль таких объектов, а не одна фиксированная матрица без модели источника.

Фактическая скорость

В него могут входить:

  • квантованные значения или индексы;
  • масштабы и нулевые точки;
  • кодовые книги;
  • границы групп;
  • служебные данные разреженности;
  • заполнение и выравнивание;
  • заголовки и формат контейнера.

Поэтому номинальная разрядность не всегда равна реальному числу бит на параметр.

Возможные меры искажения

Например:

  1. ∥W−W^∥22\|W-\widehat W\|_2^2 — MSE по весам;

  2. E∥WX−W^X∥22\mathbb E\|WX-\widehat W X\|_2^2 — ошибка выходов слоя на калибровочных активациях;

  3. рост NLL или KL между предиктивными распределениями исходной и квантованной моделей;

  4. изменение метрики последующей задачи.

Почему гауссовская граница не даёт границу на точность

Формула D≥σ22−2RD\ge\sigma^2 2^{-2R} предполагает i.i.d.-гауссовский источник, квадратичное искажение и оптимальный длинноблочный код. Точность в последующей задаче является нелинейной функцией всей сети, данных и границы решений. Даже одинаковый weight MSE может иметь очень разный функциональный эффект в разных направлениях параметров.


Упражнение 6. Python: дискретный IB и разложение скорости VAE

Запустите M9_ib_rate_distortion_experiments.py. Скрипт:

  1. строит теоретические кривые скорость–искажение;

  2. решает игрушечную задачу Information Bottleneck итерациями самосогласования;

  3. проверяет точное разложение скорости VAE на взаимную информацию и несовпадение с априорным распределением.

Ответьте:

  • почему при малом β\beta представление IB тривиально;
  • почему I(T;Y)I(T;Y) насыщается раньше, чем I(X;T)I(X;T);
  • как плохое априорное распределение увеличивает KL-скорость, не добавляя информации об объекте.
Показать ответ и решение

Скрипт строит три независимых проверки.

Кривые скорость–искажение

Кривая для источника Бернулли убывает от h2(p)h_2(p) при D=0D=0 до нуля при D=pD=p. Гауссовская кривая расходится при D↓0D\downarrow0 и становится нулевой при D≥σ2D\ge\sigma^2.

Игрушечная задача Information Bottleneck

При малом β\beta штраф I(X;T)I(X;T) доминирует. Оптимум почти не использует TT:

I(X;T)≈0,I(T;Y)≈0.I(X;T)\approx0, \qquad I(T;Y)\approx0.

При росте β\beta представление начинает разделять группы xx с различными P(Y∣x)P(Y\mid x). Релевантность I(T;Y)I(T;Y) насыщается около I(X;Y)I(X;Y): больше информации о YY в XX просто нет. Но I(X;T)I(X;T) может продолжать расти, если код начинает сохранять детали XX, не добавляющие предиктивной релевантности.

Разложение скорости VAE

Численная проверка должна дать

EXDKL(q(Z∣X)∥p(Z))=I(X;Z)+DKL(q(Z)∥p(Z))\mathbb E_XD_{\mathrm{KL}}(q(Z\mid X)\|p(Z)) = I(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z))

с точностью до ошибки вычислений с плавающей точкой.

Если заменить хорошо подобранное априорное распределение на плохое, I(X;Z)I(X;Z) не изменится: энкодер и совместное распределение остались теми же. Зато DKL(q(Z)∥p(Z))D_{\mathrm{KL}}(q(Z)\|p(Z)) возрастёт, а вместе с ним вырастет и KL-скорость. Это лишняя цена несовпадения кода, а не дополнительная информация о данных.


Программная часть решения · один исходный файл для RU/EN

PythonM9_ib_rate_distortion_experiments.pyСкачать .py
from __future__ import annotations

"""Reproducible numerical experiments for Module 9.

The script:
1. plots Bernoulli and Gaussian rate-distortion curves;
2. solves a finite Information Bottleneck problem by self-consistent updates;
3. verifies the exact VAE-rate decomposition on a discrete toy encoder.

All logarithms used for information quantities are base 2 unless stated otherwise.
"""

from pathlib import Path
import math

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "generated_experiment"
OUT.mkdir(exist_ok=True)
EPS = 1e-15


def normalize(a: np.ndarray, axis: int = -1) -> np.ndarray:
    total = a.sum(axis=axis, keepdims=True)
    if np.any(total <= 0):
        raise ValueError("Cannot normalize non-positive mass")
    return a / total


def xlog2x(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    return np.where(x > 0, x * np.log2(x), 0.0)


def binary_entropy(p: np.ndarray | float) -> np.ndarray:
    p = np.asarray(p, dtype=float)
    p = np.clip(p, EPS, 1 - EPS)
    return -(p * np.log2(p) + (1 - p) * np.log2(1 - p))


def mutual_information(joint: np.ndarray) -> float:
    joint = np.asarray(joint, dtype=float)
    joint = joint / joint.sum()
    px = joint.sum(axis=1, keepdims=True)
    py = joint.sum(axis=0, keepdims=True)
    ratio = np.divide(joint, px @ py, out=np.ones_like(joint), where=joint > 0)
    return float(np.sum(np.where(joint > 0, joint * np.log2(ratio), 0.0)))


def kl_rows(p: np.ndarray, q: np.ndarray) -> np.ndarray:
    """Return all pairwise KL divergences D(p_i || q_j) in bits."""
    p = np.asarray(p, dtype=float)
    q = np.asarray(q, dtype=float)
    p_safe = np.clip(p, EPS, 1)
    q_safe = np.clip(q, EPS, 1)
    return np.sum(p[:, None, :] * (np.log2(p_safe[:, None, :]) - np.log2(q_safe[None, :, :])), axis=-1)


def plot_rate_distortion() -> None:
    p = 0.25
    d_b = np.linspace(0, p, 250)
    r_b = binary_entropy(p) - binary_entropy(d_b)
    r_b[-1] = 0.0

    sigma2 = 1.0
    d_g = np.geomspace(0.01, sigma2, 250)
    r_g = 0.5 * np.log2(sigma2 / d_g)

    fig, ax = plt.subplots(figsize=(8.6, 5.0))
    ax.plot(d_b, r_b, label=r"Bernoulli($p=0.25$), Hamming")
    ax.plot(d_g, r_g, label=r"Gaussian($\sigma^2=1$), MSE")
    ax.set_xlabel("Allowed distortion D")
    ax.set_ylabel("Rate R(D), bits/source symbol")
    ax.set_title("Two rate-distortion benchmarks")
    ax.legend()
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(OUT / "M9_experiment_rate_distortion.png", dpi=180)
    plt.close(fig)


def ib_metrics(px: np.ndarray, py_given_x: np.ndarray, qt_given_x: np.ndarray) -> tuple[float, float]:
    pxt = px[:, None] * qt_given_x
    ixt = mutual_information(pxt)
    pty = np.einsum("x,xt,xy->ty", px, qt_given_x, py_given_x)
    ity = mutual_information(pty)
    return ixt, ity


def solve_ib(
    px: np.ndarray,
    py_given_x: np.ndarray,
    n_t: int,
    beta: float,
    rng: np.random.Generator,
    restarts: int = 24,
    max_iter: int = 5000,
    tol: float = 1e-12,
) -> tuple[np.ndarray, float, float, float]:
    """Solve a finite IB problem by self-consistent updates.

    Returns q(t|x), I(X;T), I(T;Y), and the IB objective I(X;T)-beta I(T;Y).
    Multiple random restarts are used because the fixed-point equations may have
    several stationary solutions.
    """
    best: tuple[np.ndarray, float, float, float] | None = None
    n_x = px.size

    for restart in range(restarts):
        if restart == 0:
            qtx = np.full((n_x, n_t), 1.0 / n_t)
            qtx += 1e-3 * rng.random((n_x, n_t))
        else:
            qtx = rng.dirichlet(np.ones(n_t), size=n_x)
        qtx = normalize(qtx, axis=1)

        for _ in range(max_iter):
            pt = px @ qtx
            pt = np.clip(pt, EPS, None)
            pyt = np.einsum("x,xt,xy->ty", px, qtx, py_given_x)
            pyt = normalize(pyt, axis=1)
            distortion = kl_rows(py_given_x, pyt)

            # Exponents use natural exp. Convert KL bits to nats by ln(2).
            log_new = np.log(pt)[None, :] - beta * math.log(2.0) * distortion
            log_new -= log_new.max(axis=1, keepdims=True)
            new = normalize(np.exp(log_new), axis=1)
            if np.max(np.abs(new - qtx)) < tol:
                qtx = new
                break
            qtx = new

        ixt, ity = ib_metrics(px, py_given_x, qtx)
        objective = ixt - beta * ity
        candidate = (qtx.copy(), ixt, ity, objective)
        if best is None or objective < best[3] - 1e-11:
            best = candidate

    assert best is not None
    return best


def run_ib_experiment() -> list[tuple[float, float, float]]:
    px = np.full(6, 1 / 6)
    # Three progressively more positive regions. Neighboring states have similar
    # predictive distributions, making the bottleneck geometry visible.
    p1 = np.array([0.03, 0.08, 0.30, 0.70, 0.92, 0.97])
    py_given_x = np.column_stack([1 - p1, p1])
    rng = np.random.default_rng(20260713)
    betas = np.array([0.05, 0.10, 0.20, 0.40, 0.75, 1.25, 2.0, 3.5, 6.0, 10.0, 18.0])
    rows: list[tuple[float, float, float]] = []

    for beta in betas:
        _, ixt, ity, _ = solve_ib(px, py_given_x, n_t=6, beta=float(beta), rng=rng)
        rows.append((float(beta), ixt, ity))

    ixy = mutual_information(px[:, None] * py_given_x)

    fig, ax = plt.subplots(figsize=(7.2, 5.4))
    ixt = np.array([r[1] for r in rows])
    ity = np.array([r[2] for r in rows])
    ax.plot(ixt, ity, marker="o", label="IB fixed points")
    for beta, x, y in rows[::2]:
        ax.annotate(f"β={beta:g}", (x, y), xytext=(5, 5), textcoords="offset points", fontsize=8)
    ax.axhline(ixy, linestyle="--", linewidth=1.2, label=f"I(X;Y)={ixy:.3f} bits")
    ax.set_xlabel("I(X;T), bits")
    ax.set_ylabel("I(T;Y), bits")
    ax.set_title("Toy Information Bottleneck frontier")
    ax.legend()
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(OUT / "M9_experiment_ib_frontier.png", dpi=180)
    plt.close(fig)
    return rows


def expected_conditional_kl(px: np.ndarray, qz_given_x: np.ndarray, prior: np.ndarray) -> float:
    """Average KL in bits for normalized probability arrays, without smoothing."""
    joint = px[:, None] * qz_given_x
    positive = joint > 0
    prior_grid = np.broadcast_to(prior, qz_given_x.shape)
    # A prior must cover every latent with positive joint mass. Rows with
    # px[x] = 0 and terms with q(z|x) = 0 contribute nothing, including 0 log 0.
    if np.any(prior_grid[positive] == 0):
        return math.inf
    return float(np.sum(joint[positive] * (
        np.log2(qz_given_x[positive]) - np.log2(prior_grid[positive])
    )))


def run_rate_decomposition(bad_prior: np.ndarray | None = None) -> dict[str, float]:
    px = np.array([0.40, 0.35, 0.25])
    qz_given_x = np.array([
        [0.80, 0.15, 0.05],
        [0.20, 0.65, 0.15],
        [0.05, 0.20, 0.75],
    ])
    qz = px @ qz_given_x
    good_prior = qz.copy()
    if bad_prior is None:
        bad_prior = np.array([0.78, 0.12, 0.10])

    joint = px[:, None] * qz_given_x
    ixz = mutual_information(joint)
    # Ordinary KL is the same calculation with a single source state.
    mismatch_good = expected_conditional_kl(np.ones(1), qz[None, :], good_prior)
    mismatch_bad = expected_conditional_kl(np.ones(1), qz[None, :], bad_prior)
    rate_good = expected_conditional_kl(px, qz_given_x, good_prior)
    rate_bad = expected_conditional_kl(px, qz_given_x, bad_prior)

    if not np.isclose(rate_good, ixz + mismatch_good, atol=1e-12):
        raise AssertionError("Good-prior decomposition failed")
    if not np.isclose(rate_bad, ixz + mismatch_bad, atol=1e-12):
        raise AssertionError("Bad-prior decomposition failed")

    labels = ["I(X;Z)", "prior mismatch", "total KL-rate"]
    good = [ixz, mismatch_good, rate_good]
    bad = [ixz, mismatch_bad, rate_bad]
    x = np.arange(len(labels))
    width = 0.35
    fig, ax = plt.subplots(figsize=(7.7, 5.0))
    ax.bar(x - width / 2, good, width, label="prior = aggregate posterior")
    finite = np.isfinite(bad)
    ax.bar(x[finite] + width / 2, np.asarray(bad)[finite], width,
           label="mismatched prior")
    for position in x[~finite] + width / 2:
        ax.text(position, 0.78, "∞", transform=ax.get_xaxis_transform(),
                ha="center", color="C1", fontsize=22)
    ax.set_xticks(x, labels)
    ax.set_ylabel("Bits")
    ax.set_title("VAE-rate decomposition: information plus prior mismatch")
    ax.legend()
    ax.grid(axis="y", alpha=0.25)
    if not np.all(finite):
        # Text annotations do not participate in Matplotlib's autoscaling.
        ax.set_xlim(-0.6, len(labels) - 0.4)
        fig.text(0.5, 0.02, "∞: prior assigns zero mass to a possible latent",
                 ha="center", fontsize=10)
    fig.tight_layout(rect=(0, 0.06 if not np.all(finite) else 0, 1, 1))
    fig.savefig(OUT / "M9_experiment_rate_decomposition.png", dpi=180)
    plt.close(fig)

    return {
        "I(X;Z)": ixz,
        "good_mismatch": mismatch_good,
        "good_rate": rate_good,
        "bad_mismatch": mismatch_bad,
        "bad_rate": rate_bad,
    }


def main() -> None:
    plot_rate_distortion()
    rows = run_ib_experiment()
    decomposition = run_rate_decomposition()

    print("Toy IB results (bits):")
    print("beta\tI(X;T)\tI(T;Y)")
    for beta, ixt, ity in rows:
        print(f"{beta:>5g}\t{ixt:.6f}\t{ity:.6f}")

    print("\nVAE-rate decomposition (bits):")
    for key, value in decomposition.items():
        print(f"{key}: {value:.9f}")
    print("\nGenerated figures:")
    for path in sorted(OUT.glob("M9_experiment_*.png")):
        print(path.name)


if __name__ == "__main__":
    main()

Бонус. I–MMSE для гауссиана

Для X∼N(0,1)X\sim\mathcal N(0,1) выведите апостериорное среднее, mmse⁡(γ)\operatorname{mmse}(\gamma) и I(X;Yγ)I(X;Y_\gamma). Проверьте тождество и объясните, почему интеграл до ∞\infty расходится, хотя дифференциальная энтропия конечна.

Показать ответ и решение

Пусть

Yγ=γX+N,X,N∼N(0,1)Y_\gamma=\sqrt\gamma X+N, \qquad X,N\sim\mathcal N(0,1)

независимы. Совместное распределение гауссовское, поэтому апостериорное среднее линейно:

E[X∣Yγ]=γ1+γYγ.\mathbb E[X\mid Y_\gamma] = \frac{\sqrt\gamma}{1+\gamma}Y_\gamma.

Апостериорная дисперсия равна

Var⁡(X∣Yγ)=11+γ.\operatorname{Var}(X\mid Y_\gamma) =\frac1{1+\gamma}.

Следовательно,

mmse⁡(γ)=11+γ.\boxed{ \operatorname{mmse}(\gamma)=\frac1{1+\gamma} }.

Формула пропускной способности для гауссовского входа даёт

I(X;Yγ)=12ln⁡(1+γ).\boxed{ I(X;Y_\gamma)=\frac12\ln(1+\gamma) }.

Проверка:

ddγI(X;Yγ)=12(1+γ)=12mmse⁡(γ).\frac d{d\gamma}I(X;Y_\gamma) = \frac1{2(1+\gamma)} = \frac12\operatorname{mmse}(\gamma).

Но

12∫0∞dγ1+γ=∞.\frac12\int_0^\infty\frac{d\gamma}{1+\gamma} =\infty.

Это не противоречит конечности

h(X)=12ln⁡(2πe).h(X)=\frac12\ln(2\pi e).

I–MMSE интегрирует взаимную информацию между непрерывным XX и наблюдением всё большей точности; она растёт без границы. Чтобы связать интеграл с дифференциальной энтропией, требуется вычесть расходящуюся опорную часть — простое равенство без перенормировки неверно.