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

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.

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

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.