Compute Allocation in Bilinear Contrastive Representation Learning via Scaling Laws

The gist Scaling laws study how model performance varies with model size, data size, learning steps and compute This paper develops a theoretical scaling-law model for contrastive learning in a sketched linear setting

Model Setup and Objective

The paper studies a sketched linear model for contrastive learning under a paired Gaussian latentvariable setup. The learner observes only sketched views of two correlated variables and trains a bilinear contrastive score by full-batch empirical gradient descent. The setup involves generating two matched views from a shared latent variable z, x = z + epsilonx, y = z + epsilony, where the marginal covariance is H:= E[xx⊤] = E[yy⊤] = Λz + Λepsilon and the cross-covariance is C:= E[xy⊤] = Λz. The learner observes only the sketched pair (xe = Sx, ye = Sy), simulating limitations in observation.

Risk Decomposition

The paper decomposes the population risk into irreducible risk, approximation error, GD bias, GD variance, and a cross term. Proposition 1 states that the population risk has decomposition ED [RM(AL)] = R(W⋆) z Irreducible + RM(A⋆) − R(W⋆) z Approx + 1/2 ED h B⋆L 2Σ,Σ i z Bias + 1/2 ED h VL 2Σ,Σ i z Var − ED h Σ,Σ i z Cross. The cross term is bounded by the sum of the bias and variance terms for an upper bound analysis.

Scaling Law Derivation

The main result is Theorem 1, which provides an explicit scaling law for empirical gradient descent in sketched contrastive learning under aligned power-law spectra and a contrastive source condition. The theorem states that the population risk satisfies ED[RM(AL) Ecov(R)] = −Θ(1) z Irreducible + Θ M1−2δ z Approx + Θ R 1−2δ 2b z Bias + Θ D×(R, M) N z Var + O R 1−2δ 4b p D×(R, M)/N z Cross. The product effective dimension dprod(R, M) is defined as X i,j≤M min[1,(Rµiµj) 2] product effective dimension.

Interpretation and Compute Allocation

The scaling law gives a three-way tradeoff: increasing M reduces approximation error; increasing R lets GD learn more signal but activates more finite-sample noise; and increasing N suppresses this noise. Proposition 2 provides a compute allocation heuristic, suggesting rules for balancing optimization time L, sketch/model size M, and sample size N under a compute proxy C ≍ LNM2. The leading allocation rules depend on the regime defined by q:= 2δ − 1.

Key Scaling Components

The approximation term is bounded by Theorem 2 as Approx ≈ M1−2δ, which depends on the sketch dimension M and the power-law exponents a and b. The GD-bias term BiasL scales as Θ R 1−2δ 2b, where R = Leffγ. The GD-variance term VarL scales as Θ D×(R, M) N, where D×(R, M) is a piecewise scale in Theorem 1.

Experimental Validation

Synthetic experiments verify the theoretical predictions. The approximation experiment shows the empirical approximation error follows the predicted sketch-dimension rate of −2.01 compared to −2.00. The GD-bias experiment shows the measured bias decreases monotonically with optimization horizon L, with an empirical slope of approximately −0.86 compared to the theorem-guided slope of −0.91. The GD-variance experiment shows the variance grows as additional direction-pairs are learned from finite-sample fluctuations.

Conclusion

The paper provides a first provable scaling-law analysis for a contrastive learning objective under risk decomposition. The main takeaway is that the model must learn interactions between two views, which changes how optimization and finite-sample noise scale. The explicit scaling laws offer guidance for compute allocation. Several directions remain open, such as extending the model to heterogeneous views or analyzing mini-batch SGD.

How it works

1. Latent-variable contrastive setup involves generating two noisy matched views x and y from a shared latent variable z, x = z + epsilonx, y = z + epsilony.

2. The bilinear score is trained by empirical GD on paired sketched samples to minimize RbM(A).

3. The population risk is decomposed into irreducible risk, approximation error, GD bias, GD variance, and a cross term.

4. Theorem 1 gives the scaling law for empirical GD in sketched contrastive learning under aligned power-law spectra and a contrastive source condition.

5. Compute allocation rules are derived in Proposition 2 by balancing the leading excess-risk proxy, which depends on M, N, and R.

Key Scaling Terms

The irreducible risk is constant with respect to M, N, and R The approximation term is the price of using an Mdimensional full-row-rank sketch that satisfies the contrastive source condition The GD-bias term measures incomplete training, decreasing with more effective optimization time R = Leffγ The GD-variance term measures finite-sample noise, which grows with R and is suppressed by N

Interpretation of Terms

The irreducible risk is the best value achievable by the full, unsketched contrastive loss The approximation term decreases as M grows because a larger admissible sketch can represent more effective contrastive spectral directions The GD-bias term decreases with R because more steps make the learning process more complete The GD-variance term is suppressed by N through the factor 1/N, but optimization time R increases the amount of fitted empirical noise Finally, the cross term is controlled by Cauchy–Schwarz and does not affect upper bounds.

Compute Budget Allocation

The effective optimization horizon R = Leffγ is used to determine allocation rules. The leading allocation rules are derived in Proposition 2 based on the regime defined by q:= 2δ − 1. The allocation intuition suggests that increasing the optimization horizon requires at least a proportional increase in sample size.

Limitations

The setup uses a bilinear score and assumes matched dimensions, abstracting away nonlinear encoders like those in CLIP. It lacks explicit sampled negatives, representing them through the quadratic objective instead of explicit negative pairs. The analysis assumes Gaussian data with full-batch GD under a covariance-concentration regime, which is stronger than typical practical training conditions.

References

Arora, S.; Khandeparkar, H.; Khodak, M.; Plevrakis, O.; and Saunshi, N. 2019. A Theoretical Analysis of Contrastive Unsupervised Representation Learning. In Proceedings of the 36th International Conference on Machine Learning, 5628–5637

Atanasov, A.; Zavatone-Veth, J. A.; and Pehlevan, C. 2026. Scaling and Renormalization in High-Dimensional Regression. Journal of Statistical Mechanics: Theory and Experiment, 2026(4): 043404

Bahri, Y.; Dyer, E.; Kaplan, J.; Lee, J.; and Sharma, U. 2024. Explaining Neural Scaling Laws. Proceedings of the National Academy of Sciences, 121(27): e23118781

Bordelon, B.; Atanasov, A.; and Pehlevan, C. 2024. A Dynamical Model of Neural Scaling Laws. In Proceedings of the 41st International Conference on Machine Learning

Caponnetto, A.; and De Vito, E. 2007. Optimal Rates for the Regularized Least-Squares Algorithm. Foundations of Computational Mathematics, 7(3): 331–368

Chen, T.; Kornblith, S.; Norouzi, M.; and Hinton, G. 2020. A Simple Framework for Contrastive Learning of Visual Representations. In Proceedings of the 37th International Conference on Machine Learning, 1597–1607

Chen, Z.; and Zhou, D. 2026. From One-Pass SGD to Data Reuse: Mini-Batch Scaling Laws in Sketched Linear Regression. arXiv preprint arXiv:2605.24316

Cherti, M.

1. Model Training Compute Allocation: The system can dynamically adjust resource allocation by using Proposition 2 as a heuristic to balance optimization time L, model size M, and sample size N under a fixed compute budget C. This provides explicit rules for balancing these resources based on the regime of q := 2 delta minus one.

2. Approximation Error Reduction: The system can optimize sketch dimension M to minimize the approximation error by targeting the scaling law derived in Theorem 2, which shows that Approx is proportional to M raised to the power of one minus two times delta.

3. Gradient Descent Step Scheduling: The system can determine an optimal learning rate schedule by setting Leff equal to R over gamma star, where R is the target horizon from Proposition 2 and gamma star is a stable stepsize constant, ensuring the GD filters are contractions on the covariance event.

4. Compute Budget Allocation Heuristic: The system can select model size M and sample size N based on which of the three leading allocation rules in Proposition 2 applies (product-unsaturated, covariance-limited, or variance-balanced regime) to achieve the lowest leading excess risk proxy.

5. Variance Control via Product Dimension: To manage GD variance, the system can allocate resources such that the product effective dimension dprod(R, M) is kept small relative to N. This prevents the variance term from dominating by ensuring N scales appropriately with R and M as dictated by Theorem 4.