docs: bring algorithm pages, README and CONTRIBUTING in line with the code

CONTRIBUTING.md

@@ -41,7 +41,7 @@ source .venv/bin/activate
 Install the project together with development dependencies:
 
 ```bash
-uv sync --all-groups
+uv sync --all-extras
 ```
 
 ---
@@ -63,8 +63,8 @@ We use **Ruff** for linting and formatting.
 Before committing, run:
 
 ```bash
-ruff check .
-ruff format .
+ruff check bensemble/ tests/ benchmarks/
+ruff format bensemble/ tests/ benchmarks/
 ```
 
 ---

README.md

@@ -128,14 +128,14 @@ print(f"Prediction: {mean[0].item():.2f} ± {std[0].item():.2f}")
 
 ## Algorithms & Demos
 
-We implement a wide range of Bayesian and Ensembling approaches. Check out the interactive demos in the `notebooks/` directory:
+We implement a wide range of Bayesian and Ensembling approaches. Check out the interactive demos in the `examples/` directory:
 
 | Method | Description |
 | :--- | :--- |
 | **Deep Ensembles** | Naive yet powerful ensembling of independent networks with explicit uncertainty decomposition. |
 | **Monte Carlo Dropout** | Implicit ensembling by keeping dropout active at test time. |
 | **Neural Ensemble Search (NES)** | Automatically searches for diverse architectures (NES-RS/NES-RE). |
-| **NES via Bayesian Sampling** | Extracts diverse subnetworks from a Supernet using Stein Variational Gradient Descent (SVGD). |
+| **NES via Bayesian Sampling** | Selects a diverse ensemble from a pool of trained candidates using a validation-loss posterior and SVGD-inspired repulsion. |
 | **Variational Inference** | Approximates posterior using Gaussian distributions with the *Local Reparameterization Trick*. |
 | **Variational Rényi** | Generalization of VI minimizing $\alpha$-divergence (VR-VI) for better robustness. |
 | **Laplace Approximation** | Fits a Gaussian around the MAP estimate using Kronecker-Factored Curvature (K-FAC). |

blog/bensemble-blogpost.md

@@ -265,7 +265,7 @@ This is approximated by Monte Carlo using samples $\boldsymbol{\theta}^{(k)}$ dr
 
 Qualitatively, this gives a knob that controls how aggressive or conservative the variational approximation is. Once trained, sampling networks is as simple as drawing from the Gaussian $q(\boldsymbol{\theta})$ and plugging the sampled weights into the base model, just as in PVI.
 
-Variational Rényi inference is implemented in the `VariationalRenyi` class in Bensemble. Visit our [variational Rényi demo](https://github.com/intsystems/bensemble/blob/master/notebooks/variatinal_renyi_demo.ipynb) for an example on how to use it.
+Variational Rényi inference is available in Bensemble by passing `alpha` to `VariationalLoss` on top of the same Bayesian layers. See the [Variational Rényi page](https://intsystems.github.io/bensemble/algorithms/variational-renyi/) for details.
 ### Laplace approximation
 
 [Laplace approximation (LA)](https://openreview.net/pdf?id=Skdvd2xAZ) starts from a different point. Instead of designing a Bayesian method from scratch, you begin with a network that has already been trained in the usual deterministic way, with weight decay capturing the prior. Let
@@ -329,7 +329,7 @@ $$
 
 The end result is a factorized Gaussian over weights plus Gamma distributions over hyperparameters. From that, sampling full networks is straightforward: draw weights from the Gaussians, plug them into a standard multilayer perceptron, and you have a concrete ensemble member.
 
-PBP is implemented in the `ProbabilisticBackpropagation` class in Bensemble. For an example on how to use it, check out our [probabilistic backpropagation demo](https://github.com/intsystems/bensemble/blob/master/notebooks/pbp_probabilistic_backpropagation_test.ipynb).
+PBP is implemented in the `PBPEngine` class in Bensemble. For an example on how to use it, see the [Probabilistic Backpropagation page](https://intsystems.github.io/bensemble/algorithms/pbp/).
 
 ### Neural Ensemble Search
 

docs/docs/algorithms/laplace.md

@@ -4,13 +4,13 @@ A scalable Laplace approximation using Kronecker-Factored Approximate Curvature.
 
 Since this is a post-hoc method, a standard deterministic network is first trained to find the MAP estimate $W^{\text{MAP}}_l$. We then capture the covariance of activations ($A_l$) and pre-activation gradients ($G_l$).
 
-The posterior for layer $l$ is then approximated as a matrix normal distribution $\mathcal{MN}(W^{\text{MAP}}_l, A_l, G_l)$. Samples are generated efficiently using the Cholesky decomposition of Kronecker factors:
+The posterior for layer $l$ is then approximated as a matrix normal distribution $\mathcal{MN}(W^{\text{MAP}}_l, A_l, G_l)$. Samples are generated from the two Kronecker factors independently:
 
 $$
 W_{\text{sample}} = W_{\text{MAP}} + L_V Z L_U^T
 $$
 
-where $L_V, L_U$ are Cholesky factors of the inverse regularized covariances and $Z$ is sampled from the standard matrix normal distribution.
+where $L_V, L_U$ are symmetric square roots of the inverse regularized covariances, computed by eigendecomposition with the eigenvalues clamped from below for stability, and $Z$ is sampled from the standard matrix normal distribution.
 
 ---
 

docs/docs/algorithms/nesbs.md

@@ -2,18 +2,24 @@
 
 To reduce the prohibitive computational cost of standard NES, one can use Neural Ensemble Search via Bayesian Sampling.
 
-It utilizes training a Supernet with uniform path sampling to share weights across different model architectures. A variational posterior over architectures $p_\alpha(\mathcal{A}) \approx p(\mathcal{A}|\mathcal{D})$ is learned via ELBO minimization.
+The original method trains a Supernet with weight sharing and learns a variational posterior over architectures. `bensemble` implements a discrete, pool-based version of that idea instead: `NESBayesianSampler` draws `pool_size` architectures from the `SearchSpace`, trains each one independently with the user's `train_fn`, and scores it on a validation set. The scores define a posterior over the pool,
 
-Ensemble member architectures can then be sampled from the variational posterior using two methods:
+$$
+p(\mathcal{A}_i \mid \mathcal{D}) \propto \exp\!\left(-\frac{s_i - \min_j s_j}{T}\right),
+$$
+
+where $s_i$ is the validation loss of candidate $i$ and $T$ is the `temperature`.
+
+Ensemble members are then selected from the pool in one of two ways:
 
-- **Monte-Carlo Sampling**: Simple random sampling from the posterior.
-- **SVGD-RD**: Stein Variational Gradient Descent with Regularized Diversity. This uses controlled optimization of the set of architectures with the following objective:
+- **Monte-Carlo Sampling** (`sample_mc`): draw `ensemble_size` candidates from the posterior.
+- **SVGD-inspired sampling** (`sample_svgd`): a greedy, particle-style selection over the pool. Each candidate's posterior probability is traded off against a repulsion term measuring how similar its validation predictions are to those of the members already chosen, so the selected set is pushed towards architectures that disagree with each other.
 
 $$
 q^* = \arg\min_{q\in\mathcal{Q}} \text{KL}(q\|p) + n\delta\mathbb{E}_{x, x' \sim q}[k(x, x')]
 $$
 
-This repulsive force mathematically ensures that the sampled architectures are highly diverse.
+The objective above is the one the original paper optimizes with Stein Variational Gradient Descent; here it motivates the repulsion heuristic rather than being solved exactly.
 
 ---
 

docs/docs/algorithms/variational-renyi.md

@@ -2,7 +2,7 @@
 
 This method generalizes the standard ELBO using $\alpha$-Rényi divergence.
 
-Unlike VI implementation with LRT, here explicit weights $w \sim \mathcal{N}(\mu, \text{softplus}(\rho))$ are sampled using weight perturbation during the forward pass. The objective is defined as:
+It uses the same Bayesian layers as [Variational Inference](variational-inference.md), so weights are still sampled with the Local Reparameterization Trick; only the objective changes. Pass `alpha` to `VariationalLoss` and feed it $K$ stochastic forward passes stacked along the first dimension. The objective is defined as:
 
 $$
 \mathcal{L}_{\text{VR}}(\theta, \alpha) = -\frac{1}{1-\alpha} \log \frac{1}{K} \sum_{k=1}^K \left( \frac{p(\mathcal{D}, w_k)}{q_\theta(w_k)} \right)^{1-\alpha}

docs/docs/index.md

@@ -51,6 +51,6 @@ hide:
 
     ---
 
-    Algorithms to automatically search for diverse architectures using NNI and Stein Variational Gradient Descent.
+    Algorithms to automatically search for diverse architectures, including random search, regularized evolution and SVGD-inspired Bayesian sampling.
 
 </div>

docs/docs/user-guide/basic-concepts.md

@@ -9,4 +9,4 @@ Everything revolves around the `Ensemble` class. It acts as a manager for multip
 - **Explicit Ensembles**: A collection of different models (e.g., from NAS or Deep Ensembles).
 - **Implicit Ensembles**: A single model that behaves like an ensemble (e.g., MC Dropout or Bayesian layers).
 
-Regardless of the source, an `Ensemble` always returns a tensor of shape `[M, Batch, Output]`, where `M` is the number of ensemble members.
+Regardless of the source, `ensemble.predict_members(x)` returns a tensor of shape `[M, Batch, Output]`, where `M` is the number of ensemble members, and calling `ensemble(x)` returns the combined prediction of shape `[Batch, Output]` (the mean by default).

docs/docs/user-guide/calibration-and-metrics.md

@@ -27,7 +27,7 @@ scaler = VectorScaling(num_classes=3).fit(val_logits, val_labels)
 probs = torch.softmax(scaler(test_logits), dim=-1)
 ```
 
-For an ensemble, fit the scaler on the logits you actually evaluate, whether that is the per-member output of `predict_members` or its mean.
+Both scalers expect a 2-D `[N, num_classes]` tensor. For an ensemble, fit the scaler on the logits you actually evaluate — typically the member average, `ensemble.predict_members(x).mean(0)` or simply `ensemble(x)` — rather than on the stacked `[M, N, num_classes]` member outputs.
 
 ## Scoring rules