Group Variance Regularization

This guide documents the group variance regularization feature in NLSQ’s hybrid streaming optimizer. This feature prevents per-group parameters (such as per-angle contrast and offset in XPCS laminar flow fitting) from absorbing physical signals that should be captured by shared parameters.

Motivation

In multi-angle or multi-channel fitting problems, a common model structure is:

\[y_\phi(q, t) = C_\phi \cdot g_2(q, t; \Gamma) + O_\phi\]

where:

  • \(C_\phi\) is the per-angle contrast

  • \(O_\phi\) is the per-angle offset

  • \(\Gamma\) are the shared physical parameters (e.g., relaxation rates)

Without regularization, the per-angle parameters \(C_\phi\) and \(O_\phi\) can absorb angle-dependent physical signals that should be captured by the shared parameters \(\Gamma\). This leads to:

  • Artificially uniform \(\Gamma\) values across angles

  • Loss of physical information in the fitted parameters

  • Systematic bias in extracted quantities

Group variance regularization addresses this by penalizing variance within parameter groups, encouraging \(C_\phi\) values to remain similar to each other (and likewise for \(O_\phi\)), while allowing \(\Gamma\) to capture the true angle-dependent physics.

Mathematical Formulation

The regularized loss function becomes:

\[\mathcal{L} = \text{MSE} + \lambda \sum_g \text{Var}(\theta_g)\]

where:

  • \(\text{MSE}\) is the mean squared error of the fit

  • \(\lambda\) is the regularization strength (group_variance_lambda)

  • \(\theta_g\) is the vector of parameters in group \(g\)

  • \(\text{Var}(\theta_g) = \frac{1}{n_g}\sum_{i \in g}(\theta_i - \bar{\theta}_g)^2\)

The variance penalty is added in two phases:

Phase 1 (L-BFGS Warmup)

During L-BFGS warmup, the loss function directly includes the variance penalty:

def loss_fn(params, x_batch, y_batch):
    predictions = model(x_batch, *params)
    mse = jnp.mean((y_batch - predictions) ** 2)

    variance_penalty = 0.0
    for start, end in group_slices:
        group_params = params[start:end]
        variance_penalty += jnp.var(group_params)

    return mse + lambda_ * variance_penalty

Phase 2 (Gauss-Newton)

During Gauss-Newton optimization, the regularization modifies the normal equations by adding gradient and Hessian contributions:

Gradient contribution (added to \(J^T r\)):

\[\nabla_{\theta_g} \text{Var}(\theta_g) = \frac{2}{n_g}(\theta_i - \bar{\theta}_g)\]

Hessian contribution (added to \(J^T J\)):

\[H_{\text{var}} = \frac{2}{n_g}\left(I - \frac{1}{n_g}\mathbf{1}\mathbf{1}^T\right)\]

where \(I\) is the identity matrix and \(\mathbf{1}\) is a vector of ones. This Hessian has the form of a centering matrix that penalizes deviations from the group mean.

Configuration

Group variance regularization is configured through HybridStreamingConfig:

from nlsq import HybridStreamingConfig

config = HybridStreamingConfig(
    # Enable the feature
    enable_group_variance_regularization=True,
    # Regularization strength
    group_variance_lambda=0.1,
    # Define parameter groups as (start, end) slices
    # Example: 23 contrast params [0:23], 23 offset params [23:46]
    group_variance_indices=[(0, 23), (23, 46)],
)

Configuration Parameters

enable_group_variance_regularizationbool, default=False

Enable the variance penalty. When False, standard MSE optimization is used without regularization.

group_variance_lambdafloat, default=0.01

Regularization strength. Larger values more strongly penalize variance within parameter groups.

  • 0.001-0.01: Light regularization, allows moderate group variation

  • 0.1-1.0: Moderate regularization, constrains groups to be similar

  • 10-1000: Strong regularization, forces groups to be nearly uniform

A practical formula for XPCS applications:

\[\lambda \approx 0.1 \times \frac{n_{\text{data}}}{n_\phi \times \sigma_{\text{exp}}^2}\]

where \(\sigma_{\text{exp}}\) is the expected experimental variation (e.g., 0.05 for 5% variation).

group_variance_indiceslist of tuple, default=None

List of (start, end) tuples defining parameter groups. Each tuple specifies a slice [start:end] of the parameter vector.

If None when regularization is enabled, no groups are regularized (effectively disabling the feature).

Example: XPCS Laminar Flow Fitting

For XPCS laminar flow analysis with 23 angular positions:

import jax.numpy as jnp
from nlsq import curve_fit, HybridStreamingConfig

# Model: g2(q, t) for 23 angles
# Parameters layout:
#   [0:23]   - contrast C_phi for each angle
#   [23:46]  - offset O_phi for each angle
#   [46:]    - shared physical parameters (Gamma, etc.)


def laminar_flow_model(x, *params):
    n_phi = 23
    contrast = jnp.array(params[:n_phi])
    offset = jnp.array(params[n_phi : 2 * n_phi])
    gamma = params[2 * n_phi :]  # Shared physics parameters

    # Compute g2 correlation function
    # ... physics implementation ...

    return contrast * g2_theory + offset


# Configure with group variance regularization
config = HybridStreamingConfig(
    enable_group_variance_regularization=True,
    group_variance_lambda=0.1,
    group_variance_indices=[
        (0, 23),  # Regularize contrast group
        (23, 46),  # Regularize offset group
    ],
    # Other settings
    precision="float64",
    gauss_newton_tol=1e-10,
)

# Initial parameters
p0 = (
    [0.3] * 23  # Initial contrast (same for all angles)
    + [0.0] * 23  # Initial offset (zero for all angles)
    + [1.0, 0.1]  # Initial shared parameters
)

# Fit with regularization
popt, pcov = curve_fit(
    laminar_flow_model,
    x_data,
    y_data,
    p0=p0,
    method="hybrid_streaming",
    hybrid_config=config,
)

# Extract results
fitted_contrast = popt[:23]
fitted_offset = popt[23:46]
fitted_physics = popt[46:]

# Check group variance
print(f"Contrast std: {jnp.std(fitted_contrast):.4f}")
print(f"Offset std: {jnp.std(fitted_offset):.4f}")

Choosing Lambda

The regularization strength \(\lambda\) controls the trade-off between:

  • Fit quality: Lower \(\lambda\) allows per-group parameters to fit the data more closely

  • Physical constraint: Higher \(\lambda\) forces per-group parameters to be more uniform

L-Curve Method

Perform fits with varying \(\lambda\) and plot the L-curve:

import matplotlib.pyplot as plt

lambdas = [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]
mse_values = []
variance_values = []

for lam in lambdas:
    config = HybridStreamingConfig(
        enable_group_variance_regularization=True,
        group_variance_lambda=lam,
        group_variance_indices=[(0, 23), (23, 46)],
    )

    popt, _ = curve_fit(
        model, x, y, p0=p0, method="hybrid_streaming", hybrid_config=config
    )

    # Compute unregularized MSE
    residuals = y - model(x, *popt)
    mse = float(jnp.mean(residuals**2))
    mse_values.append(mse)

    # Compute total group variance
    var_total = jnp.var(popt[:23]) + jnp.var(popt[23:46])
    variance_values.append(float(var_total))

# Plot L-curve
plt.figure()
plt.loglog(mse_values, variance_values, "o-")
for i, lam in enumerate(lambdas):
    plt.annotate(f"{lam}", (mse_values[i], variance_values[i]))
plt.xlabel("MSE (data fidelity)")
plt.ylabel("Group variance (regularization)")
plt.title("L-curve for lambda selection")
plt.show()

Choose \(\lambda\) at the “corner” of the L-curve where both MSE and variance are reasonably low.

Physical Constraints

If you have prior knowledge of expected parameter variation:

# Expected 5% variation in contrast across angles
expected_sigma = 0.05 * mean_contrast

# Set lambda to penalize deviations beyond expected variation
lambda_ = 0.1 * n_data / (n_angles * expected_sigma**2)

Validation

After fitting, validate the regularization effect:

# Check that group variance is reduced
contrast_std = jnp.std(popt[:23])
offset_std = jnp.std(popt[23:46])

print(f"Contrast coefficient of variation: {contrast_std/jnp.mean(popt[:23]):.1%}")
print(f"Offset standard deviation: {offset_std:.4f}")

# Compare with unregularized fit
config_unreg = HybridStreamingConfig(
    enable_group_variance_regularization=False,
)
popt_unreg, _ = curve_fit(
    model, x, y, p0=p0, method="hybrid_streaming", hybrid_config=config_unreg
)

contrast_std_unreg = jnp.std(popt_unreg[:23])
print(f"Contrast std (regularized): {contrast_std:.4f}")
print(f"Contrast std (unregularized): {contrast_std_unreg:.4f}")

Implementation Details

Source Files

The implementation spans two files:

nlsq/hybrid_streaming_config.py

  • Configuration dataclass with three new fields

  • Validation for group_variance_lambda > 0

  • Validation for valid (start, end) tuples

nlsq/adaptive_hybrid_streaming.py

  • _create_warmup_loss_fn() (line ~817): Adds variance penalty to loss

  • _gauss_newton_iteration() (line ~1690): Adds gradient and Hessian terms

Numerical Stability

The Hessian contribution \(H_{\text{var}} = \frac{2}{n}(I - \frac{1}{n}\mathbf{1}\mathbf{1}^T)\) is a rank-deficient matrix (rank = n-1). This is intentional: the null space corresponds to uniform shifts of all parameters in the group, which do not change the variance. The existing regularization factor in the Gauss-Newton solver handles this gracefully.

Computational Cost

The additional cost per iteration is:

  • Gradient: O(n_params) for each group

  • Hessian: O(n_group^2) for each group

For typical XPCS applications with 23 angles and 2 groups, this adds negligible overhead compared to the Jacobian computation.

See Also