cs.LGOct 21, 2026

Optimal Multitask Representations Under a Rate Constraint: An Information Bottleneck Analysis

Authors: Nicholas Stiles, Gregory W. Wornell

Organizations: Massachusetts Institute of Technology

Abstract

Which pretraining objective yields the best representation for downstream tasks under an information budget? We formalize this problem as a multitask information bottleneck, using a mutual information constraint to model incomplete information retention. For Gaussian inputs and tasks, we characterize optimal representations and show that a weakly task-aligned specialist can outperform a reconstruction generalist at low rates but lose at high rates. We also construct task families for which every fixed Gaussian prediction objective incurs arbitrarily large relative regret, driven by reversals in optimal allocation priorities. Conversely, the average task information matrix defines an objective with a uniform relative-regret bound controlled by maximum task SNR. CelebA and dSprites experiments illustrate that compression can change which pretraining objective performs best, and even task-aware training can reverse which downstream task is better preserved as rate increases.

Figures & tables

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Sep 29, 2026

Mutual Information Constrained Chernoff Bottleneck

The classical information bottleneck (IB) measures the relevance of a representation UU of XX to a target YY by I(U;Y)I(U;Y), which does not directly characterize the error of downstream decisions. For a binary hypothesis YY inferred from many separately encoded observations, the optimal error exponent is the Chernoff information between the two conditional distributions of UU given YY. We study the mutual information constrained Chernoff bottleneck, which seeks an encoder that maximizes this Chernoff information subject to a rate constraint I(U;X)≤RI(U;X) \leq R. We show that its optimal value C(R)C(R) increases strictly up to R=H(V)R = H(V), where VV merges the symbols of XX with equal likelihood ratio, remains at the uncompressed exponent beyond, and, unlike the IB curve, need not be concave. We further show that k+1k+1 outputs suffice to attain C(R)C(R), where kk is the cardinality of VV. We propose an alternating algorithm that updates the encoder via a generalized Blahut--Arimoto algorithm and the Chernoff parameter ss via a nonlinear equation, and prove that its iterates remain feasible, with nondecreasing and convergent Chernoff information. Comparing with the classical IB for nn samples, we show that, whenever the two objectives disagree, the IB encoder has the smaller equivocation H(Y∣Un)H(Y|U^{n}) at n=1n = 1, whereas the Chernoff bottleneck has the smaller error probability and equivocation for all large nn. Numerical experiments confirm the theory, and on real topic-detection data from the 20 Newsgroups corpus, compressing each word to only 17%17\% of its entropy retains 90%90\% of the error exponent and nearly the accuracy of the uncompressed classifier.
Apr 8, 2026

Learning is Forgetting: LLM Training As Lossy Compression

Despite the increasing prevalence of large language models (LLMs), we still have a limited understanding of how their representational spaces are structured. This limits our ability to interpret how and what they learn or relate them to learning in humans. We argue LLMs are best seen as an instance of lossy compression, where over training they learn by retaining only information in their training data relevant to their objective(s). We show pre-training results in models that are optimally compressed for next-sequence prediction, approaching the Information Bottleneck bound on compression. Across an array of open weights models, each compresses differently, likely due to differences in the data and training recipes used. However even across different families of LLMs the optimality of a model's compression, and the information present in it, can predict downstream performance on across a wide array of benchmarks, letting us directly link representational structure to actionable insights about model performance. In the general case the work presented here offers a unified Information-Theoretic framing for how these models learn that is deployable at scale.
Jun 23, 2026

Semantic Allocation in Ordered Bottlenecks: Predictive Residual Inference for Visual Representation Learning

Ordered bottlenecks aim to provide utility at flexible budgets by assigning coarse information to early tokens and task-relevant detail to later ones. Prior work, including tail dropping (TD), typically enforces ordering by means of a masking-based ordering pressure (MBOP): Late tokens are masked more frequently than early tokens and are therefore encouraged to store less essential fine details. We introduce predictive residual inference for ordered representations (PRIOR), a framework designed to address inherent weaknesses of MBOP. MBOP is prone to weak late-token utility because it lacks an explicit refinement objective and uses gradient exposure as a proxy for importance. Furthermore, representations may become particularly brittle in optimization-sensitive settings, such as when using discrete or quantized token representations. PRIOR replaces activation-rate control with log2-scaled levels and level-wise predictors. These predictors separate already explained from unexplained information, focusing each level on residual error. We compare PRIOR against MBOP-TD and independent tail-biased dropout (MBOP-ITD) in contrastive learning and image reconstruction tasks. Unlike the baselines, PRIOR learns well-ordered representations across experiments: low budgets provide coarse descriptors, while high budgets add refinements. Simultaneously, full-budget performance with PRIOR is higher in all but one experimental setting, where performance remains comparable. MBOP baselines are severely limited in discrete and quantized settings, while PRIOR approaches the performance of continuous counterparts. Taken together, these findings establish PRIOR as an effective framework for ordered representation learning.