LKJ Covariance for Variational Adapter Blocks
Implementation & benchmark of arXiv:2608.06116 — Bartlett Couplings of the Onion and Vine LKJ Samplers
Source paper: Bartlett Couplings of the Onion and Vine LKJ Samplers arXiv:2608.06116 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Use an LKJ correlation factor as the correlation component of a variational posterior over a compact adapter, LoRA factor, or Bayesian neural-network parameter block. The model learns marginal scales separately while the correlation matrix remains automatically positive semidefinite and unit-diagonal, avoiding unconstrained covariance matrices, invalid correlations, and fragile covariance decompositions.
Formulas
Mathematical statement
For d latent adapter coordinates, construct C=B B^T from a lower-triangular Bartlett matrix B whose off-diagonal entries are standard normal and whose diagonal satisfies B_ii^2~chi-square(2 eta+i-2), followed by row normalization B_i=B_i/||B_i||_2. This produces C in the elliptope and an LKJ density p(C) proportional to det(C)^(eta-1). Parameterize a variational weight sample as w=mu+diag(s)Bz, where mu is the posterior mean, s_i>0 are learned marginal standard deviations, and z~N(0,I_d); its covariance is diag(s) C diag(s). The LKJ negative log-density term is -(eta-1)log det(C). The entropy result identifies eta-1 as the natural parameter associated with expected log-determinant: eta=1 is uniform over correlation matrices, eta>1 favors less correlated matrices, and eta<1 favors stronger correlations.
Implementation notes
Integrate this at the variational parameterization of a small LoRA adapter or Bayesian linear layer. Start with adapter dimension d=16 or 32 so the dense factor is inexpensive. Store a lower-triangular unconstrained matrix U and a positive scale vector s. For a learned-correlation version, set A[i,j]=U[i,j] for j<i, set A[i,i]=softplus(U[i,i])+epsilon, and row-normalize A to B. Form C=B@B.T only when the regularizer or diagnostics need it. At each Monte Carlo forward pass, draw z~Normal(0,I_d), compute w=mu+s*(B@z), and insert w into the adapter. Add beta times -(eta-1)*logdet(C) to the variational objective; compute logdet(C) with a jittered Cholesky or stable slogdet. The paper supplies the valid-correlation geometry and LKJ density; U, s, eta, and beta are learned or selected experimentally. First experiment: fine-tune a frozen 125M-parameter language model on a small instruction dataset using rank-16 LoRA, comparing diagonal Gaussian, full unconstrained Cholesky, and LKJ-factor posteriors at equal adapter parameter count. Evaluate held-out loss, calibration under posterior sampling, memory, and numerical failures. Success means lower validation loss or better uncertainty calibration than diagonal variational LoRA without Cholesky failures, ideally with fewer covariance parameters than a full posterior.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.