Epoch-Frozen Masked Low-Rank Candidate Encoder
Implementation & benchmark of arXiv:2607.08971 — Stochastic Linear Bandits with Partially Observed Actions
Source paper: Stochastic Linear Bandits with Partially Observed Actions arXiv:2607.08971 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Add an online low-rank reconstruction layer before a recommender or ranking MLP. It estimates a latent action-feature subspace from masked candidate vectors, freezes that subspace for an epoch, reconstructs each partially observed candidate in the latent coordinates, and feeds only those coordinates to the predictor. The method is most promising when the ambient candidate dimension d is large but the effective rank m is small and missingness is not too severe.
Formulas
Mathematical statement
The paper's TOFU-POV guarantee is $\widetilde{O}(\kappa^{2}m/(p^{4}K)+m\sqrt{T}+\kappa m\sqrt{T}/(p^{2}\sqrt{K}))$, where $T$ is the number of rounds, $m$ is the latent action-subspace dimension, $d$ is the ambient dimension, $p$ is the independent coordinate-observation probability, $K$ is the number of offered actions per round, and $\kappa$ is the condition number of the action covariance. We use its constructive mechanism: estimate the latent subspace from masked actions, freeze the representation within each epoch, and impute the current action by projection into that subspace. For masked vector $y=P_{\Omega}x$, with mask coordinates independently distributed as $\mathrm{Bernoulli}(p)$, form an inverse-probability covariance estimate $\widehat{\Sigma}_{ij}=\frac{1}{Np^{2}}\sum_{n=1}^{N}y_{n,i}y_{n,j}$ for $i\neq j$ and $\widehat{\Sigma}_{ii}=\frac{1}{Np}\sum_{n=1}^{N}y_{n,i}^{2}$. Let $U\in\mathbb{R}^{d\times m}$ contain the top-$m$ eigenvectors of $\widehat{\Sigma}$. For a new mask $\Omega$, reconstruct latent coordinates by ridge least squares $z=\arg\min_z\|P_{\Omega}Uz-P_{\Omega}x\|_{2}^{2}+\lambda\|z\|_{2}^{2}$, equivalently $z=(U_{\Omega}^{\top}U_{\Omega}+\lambda I)^{-1}U_{\Omega}^{\top}x_{\Omega}$, and use $\widehat{x}=Uz$ or directly use $z$ as the neural input. The beneficial property is that prediction complexity scales with $m$ rather than $d$, while epoch freezing prevents representation noise from changing every gradient step.
Implementation notes
Integrate this at the candidate-feature input of a retrieval or recommendation model, before the ranking MLP and before any expensive cross-feature interaction. At the start of each epoch, collect the masked candidate vectors from the previous epoch; estimate the inverse-probability covariance using the two cases in the formula above, symmetrize it, and compute its top-m eigenvectors with randomized eigendecomposition. Freeze U for the entire next epoch. For every candidate x with observed coordinates x_O and mask O, extract U_O, solve z=(U_O^T U_O+lambda I)^(-1)U_O^T x_O using a batched Cholesky solve, and pass [z; context_embedding] to the existing ranking MLP. If the model needs reconstructed features, use x_hat=Uz, but initially test the cheaper latent-only path. Train the ranking network with the original click or reward loss; do not backpropagate through U during the frozen epoch. Store only U and the m-dimensional candidate codes for retrieval caches. The paper-derived quantities are the mask correction, PCA subspace, ridge imputation, and epoch freezing. Estimate p from the empirical fraction of observed coordinates, choose lambda=1e-3 times the median diagonal of U_O^T U_O, and select m by validation or an explained-variance threshold. First run a MovieLens or synthetic contextual-recommendation experiment with d=512, K=32, p in {0.25,0.5,0.75}, and true rank 8 or 16. Compare a full-feature MLP, zero-imputation, and un-frozen PCA. Success means equal or better NDCG/click regret with at least 2x lower ranking-layer FLOPs and visibly more stable validation loss under changing masks.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significant win ✓
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 7/10
Built the epoch-frozen masked low-rank encoder with inverse-probability covariance estimation, top-eigenvector subspace extraction, batched ridge reconstruction, and latent-only ranking MLP input. The math check recovered the true rank-4 subspace with principal cosines 0.9994–0.9998, 3.2% relative reconstruction error, and 6.0% relative covariance-estimation error. In the synthetic ranking test, the frozen latent model improved accuracy over zero-imputation at p=0.25 (0.650 vs 0.450) and p=0.5 (0.633 vs 0.575), with a 4x lower ranking input FLOP proxy; it lost at p=0.75, so the effect is promising but conditional rather than universally better.
- Agent confidence
- 7/10
- Baseline
- Zero-imputed ambient MLP accuracy: p=0.25 0.450, p=0.5 0.575, p=0.75 0.633; ranking FLOP proxy 1.0x.
- Idea
- Frozen latent MLP accuracy: p=0.25 0.650, p=0.5 0.633, p=0.75 0.500; ranking FLOP proxy 0.25x. Math: principal cosines 0.9994–0.9998 and relative reconstruction error 0.0319.
Limitations: This is a small synthetic contextual-ranking benchmark with one seed, 48-dimensional ambient features, rank 6, 8 candidates, and 7 training epochs; no MovieLens data, NDCG, wall-clock speed, or large-scale retrieval cache was tested. The unfrozen control refit on the same epoch data, so it does not fully model representation drift from changing masks. Validation-loss stability was measured but was not consistently better for the frozen method.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.08971)
Beats the tuned baseline by 91.1% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Bench_custom_masked_candidate_low_rank
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 0.545
- Idea mean
- 0.0487
- Effect (Δ)
- -0.4963 (−91.1%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±39.6%
- Mechanism
- Confirmed ✓
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The experiment used a custom synthetic retrieval regression track with d=128, latent rank 8, fixed p=0.5, 400 training examples, 400 test examples, and 25 epochs. It did not test MovieLens, NDCG, varying missingness or rank, deployed wall-clock latency, cache memory, or the built-in tracks. Complexity is reported as a parameter/FLOP proxy rather than hardware timing.
How to run: python3 bench_runner.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Worked ✓
Agent confidence: 8/10
Implemented the epoch-frozen masked low-rank candidate encoder and evaluated it end-to-end against zero-imputed ambient features on a structurally matched custom retrieval regression track. Across 8 paired seeds, the idea achieved MSE 0.048654 versus baseline 0.545000, with all 8 paired wins and permutation p=0.0081. The measured first-layer parameter ratio was 0.06977, matching the predicted complexity ratio, so the mechanism signature was confirmed.
- Baseline
- Zero-imputed ambient MLP, best swept lr=0.001: test MSE mean 0.5450004563, std 0.2628610195.
- Idea
- Epoch-frozen inverse-probability covariance PCA with rank 8 and ridge masked encoding, best swept lr=0.006: test MSE mean 0.0486544219, std 0.0216828888; observed first-layer parameter ratio 0.0697674.
Benchmark result
Beats the tuned baseline by 91.1% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Bench_custom_masked_candidate_low_rank
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 0.545
- Idea mean
- 0.0487
- Effect (Δ)
- -0.4963 (−91.1%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±39.6%
- Mechanism
- Confirmed ✓
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The experiment used a custom synthetic retrieval regression track with d=128, latent rank 8, fixed p=0.5, 400 training examples, 400 test examples, and 25 epochs. It did not test MovieLens, NDCG, varying missingness or rank, deployed wall-clock latency, cache memory, or the built-in tracks. Complexity is reported as a parameter/FLOP proxy rather than hardware timing.
How to run: python3 bench_runner.py
Artifacts
- bench_report.json 3.6 KB View Raw
- bench_runner.py 5.0 KB View Raw
- custom_candidate_track.py 1.1 KB View Raw
- experiment.py 4.6 KB View Raw
- report.md 1.7 KB View
- report_bench_2026-09-04T204320.md 5.6 KB View
- results.json 2.4 KB View Raw