ib-components

Modular building blocks for Information Bottleneck research in PyTorch

---

Overview

ib-components is a modular library of Information Bottleneck building blocks for PyTorch. Unlike existing implementations (VIB-pytorch, information-bottleneck) that bundle end-to-end training scripts tied to specific experiments, ib-components provides the components themselves as a library — compose them into your own architectures without forking.

The Information Bottleneck Principle

The IB principle (Tishby et al., 1999) formalizes the trade-off between compression and predictive power:

The Variational Information Bottleneck (Alemi et al., 2017) makes this tractable for deep networks:

---

Installation

# Core library
pip install ib-components

# With Neural ODE support (AttractorDynamics)
pip install "ib-components[ode]"

# With HuggingFace Transformers integration
pip install "ib-components[transformers]"

# All optional dependencies
pip install "ib-components[all]"

# Development install
git clone https://github.com/Fengrru/ib-components.git
cd ib-components
pip install -e ".[dev]"

Requirements: Python ≥ 3.9, PyTorch ≥ 2.0

---

Quick Start

VIB Compression (4 lines)

import torch
from ib_components.compressor import LayerwiseVIBCompressor

compressor = LayerwiseVIBCompressor(hidden_size=768, latent_size=192, beta=0.001)

x = torch.randn(32, 128, 768)           # (batch, seq_len, hidden)
x_recon, info = compressor(x)            # reconstruct + latent info
losses = compressor.compute_ib_loss(info)  # IB loss

print(f"Total: {losses['loss']:.4f} | Recon: {losses['recon_loss']:.4f} | KL: {losses['kl_loss']:.4f}")

Hierarchical Multi-Layer Compression

from ib_components.compressor import HierarchicalVIBCompressor

compressor = HierarchicalVIBCompressor(num_layers=12, hidden_size=768, latent_ratio=0.25)

hidden_states = [torch.randn(32, 128, 768) for _ in range(12)]
reconstructed, infos = compressor(hidden_states)
total_loss = compressor.compute_total_loss(infos)

Riemannian Optimization on Sphere Manifold

from ib_components.compressor import project_to_tangent, sphere_exponential_map

point = torch.randn(1, 64)
point = point / torch.norm(point)  # normalize to unit sphere
grad = torch.randn(1, 64)

# Riemannian step — stays exactly on the sphere
riem_grad = project_to_tangent(grad, point)
riem_next = sphere_exponential_map(point, -0.1 * riem_grad)

# Euclidean step — drifts off the manifold
euc_next = point - 0.1 * grad

print(f"Riemannian norm: {torch.norm(riem_next):.6f}")  # 1.000000
print(f"Euclidean norm:  {torch.norm(euc_next):.6f}")    # ~1.07 (drifted)

Adaptive Beta Scheduling

from ib_components.improvements import AdaptiveBetaScheduler

scheduler = AdaptiveBetaScheduler(beta_min=0.0001, beta_max=0.01, schedule="sigmoid")

# High surprise → low beta (preserve information)
# Low surprise → high beta (aggressive compression)
for surprise in [0.0, 0.25, 0.5, 0.75, 1.0]:
    print(f"Surprise={surprise:.2f} → β={scheduler(surprise):.6f}")

---

Features

Components

Module	Component	Description
compressor	LayerwiseVIBCompressor	Single-layer VAE-style compressor with reparameterization
compressor	HierarchicalVIBCompressor	Multi-layer cascaded compressors with per-layer beta
compressor	RiemannianOptimizer	Sphere manifold optimization (tangent + exp map)
compressor	SphereManifold	Geodesic distance, parallel transport, normalization
improvements	AdaptiveBetaScheduler	Surprise-driven β adjustment (sigmoid/linear/exponential)
improvements	KLAnnealingScheduler	Linear/cosine/exponential KL warmup
improvements	InformationGainDrive	MINE-based mutual information estimation
improvements	CrossLayerConsistency	Frobenius norm alignment between adjacent layers
improvements	AttractorDynamics	Neural ODE with learnable attractor centers

Configuration

Type-safe dataclass configs with validation and YAML support:

from ib_components.config import VIBConfig, ExperimentConfig

# Direct instantiation
config = VIBConfig(hidden_size=768, latent_size=192, beta=0.001)

# From YAML file
experiment = ExperimentConfig.from_yaml("config.yaml")

# To YAML
experiment.to_yaml("experiment_config.yaml")

---

Benchmarks

Synthetic data, hidden_size=256. Run with python benchmarks/benchmark_compare.py.

Metric	Value	Notes
Compression ratio	4× (256→64)	75% memory saved: 12 MB → 3 MB
Riemannian norm drift	0.000001	Euclidean step drifts to 1.258 (430,000× worse)
VIB KL loss	139.9	Compression vs reconstruction trade-off
Compressor parameters	445K	Lightweight, ~0.13% of a 340M-param model

100 gradient steps on S²⁵⁵. Riemannian stays on the manifold; Euclidean drifts ~25% per step.

4× compression: 256-dim input → 64-dim latent, saving 75% memory.

---

Architecture

ib_components/
├── src/ib_components/
│   ├── compressor/              VIB encoder-decoder + Riemannian manifold
│   │   ├── vib_compressor.py    LayerwiseVIBCompressor, HierarchicalVIBCompressor
│   │   └── riemannian.py        RiemannianOptimizer, SphereManifold, tangent/expmap
│   ├── improvements/            Plug-and-play enhancements
│   │   ├── adaptive_beta.py     AdaptiveBetaScheduler, KLAnnealingScheduler
│   │   ├── info_gain.py         InformationGainDrive (MINE mutual information)
│   │   ├── cross_layer_consistency.py  CrossLayerConsistency (Frobenius alignment)
│   │   └── attractor_dynamics.py      AttractorDynamics (Neural ODE vector field)
│   ├── config.py                Type-safe dataclass configurations
│   ├── exceptions.py            Custom exception hierarchy
│   ├── logging.py               Structured logging (ComponentLogger)
│   └── utils/
│       ├── helpers.py           YAML config, logging, device, JSON, env, seeding
│       └── checkpointing.py     CheckpointManager, IncrementalCheckpointer
├── tests/                       Comprehensive test suite
├── benchmarks/                  Performance comparisons
├── examples/                    Runnable tutorials
└── docs/                        Sphinx documentation

Each component follows the same contract:

1. Subclass torch.nn.Module — drop-in PyTorch integration
2. Accept enabled: bool — set False for identity pass-through
3. Return consistent outputs — dict of loss terms or raw tensors

---

Development

# Setup
git clone https://github.com/Fengrru/ib-components.git
cd ib-components
make dev

# Commands
make test       # Run tests (with coverage, ≥80% threshold)
make lint       # Run linting (black + isort + flake8)
make format     # Auto-format code
make typecheck  # Run mypy type checking
make check      # Run all checks (lint + typecheck + test)
make docs       # Build Sphinx documentation
make bench      # Run performance benchmarks

---

References

Paper	Year	Relevance
The Information Bottleneck Principle	1999	Foundational IB theory
Deep Variational Information Bottleneck	2017	VIB loss formulation
MINE: Mutual Information Neural Estimation	2018	Information gain estimation
Neural Ordinary Differential Equations	2018	Attractor dynamics
Optimization on Matrix Manifolds	2008	Riemannian optimization theory
KL Annealing for VAEs	2016	KL warmup strategy

---

Citation

@software{ib_components2025,
  author       = {IB Components Team},
  title        = {ib\_components: Modular building blocks for Information Bottleneck research},
  year         = {2025},
  url          = {https://github.com/Fengrru/ib-components},
  version      = {0.1.3},
}

---

Contributing

Contributions welcome! See CONTRIBUTING.md for development setup, coding standards, and the PR checklist.

License

MIT — see LICENSE for details.