PaperlayerView exact revision on arXiv

Source-linked AI summary

Marformer: A Transformer for Predicting Missing Data Distributions

Prabhav Singh, Xiheng Tom Wang, Haojun Shi, Jason Eisner

arXiv:2610.12379v1cs.LGstat.ME

TL;DR

Decisions with incomplete data require predictive distributions over missing values, but classical methods fit a joint model and then infer the needed marginals. The Marformer predicts conditional marginals directly in one forward pass, and is competitive with classical methods across synthetic and real-data evaluations.

  • Problem

    Incomplete-data decisions need predictive distributions over missing values, while classical methods obtain conditional marginals through a fitted joint model and subsequent inference.

    Source: §1 · pp. 1–2

  • Method

    The Marformer directly predicts conditional marginals from observed values using a Transformer, with all missing-value predictions produced in one forward pass.

    Source: §1 · p. 2

  • Results

    Across three synthetic domains, the Marformer can match or outperform methods with the correct generative family; on real annotation data it outperforms evaluated baselines at the largest training size and runs at a fraction of Stan’s cost.

    Source: §9 · p. 33

  • Takeaways & Limitations

    The results show that direct prediction from incomplete training data can be competitive with methods given the correct generative family, while avoiding joint-model inference at prediction time.

    Source: §1 · p. 2, §9 · p. 33

    Masked-autoencoding is justified under MCAR training instances; with informative MAR missingness, masked examples may not represent naturally missing values and predictions may be systematically biased.

    Source: §9 · p. 33

Abstract

from arXiv · show

Real decisions are made under incomplete information. If we observe only some of the random variables we need, we can predict the others. The \textbf{conditional marginals} over the missing variables are the key ingredient for computing Bayes risk and Value of Information (VOI), the expected gain from acquiring one more observation before deciding. We present the Marformer, a Transformer trained to directly predict conditional marginals given any set of observed values. Like BERT, which is trained to predict missing words from context, the Marformer constructs a hidden-vector representation for each distribution $p(X_i)$ and iteratively refines it through attention to other distributions $p(X_j)$. Unlike generative approaches, the Marformer does not model the full joint distribution, requires no domain knowledge of the data-generating process, and makes all predictions in a single forward pass. We evaluate across three synthetic domains with missing data---Bayesian networks, discretized multivariate Gaussians, and structured annotation data. The Marformer can match or outperform classical missing-data methods, even when those methods are given the true model family and prior that generated the synthetic data. We also evaluate on a real annotation dataset, where the Marformer outperforms the evaluated baselines at the largest training size. In both cases, the Marformer is substantially faster than the evaluated generative baselines.

1 Introduction

Incomplete-data decisions require predictive distributions over missing values, not just point estimates. The paper introduces the Marformer to predict conditional marginals directly, reporting competitive results against classical methods even when given the true model and prior.

1 Introduction

  • Predictive distributions over missing values support residual-uncertainty estimates, action-risk prediction, and deciding whether more data is worth its cost.Examples include adaptive medical testing, recommender systems, and evaluating AI systems.Source: §1 · pp. 1–2
  • The Marformer fits conditional marginals directly rather than first fitting a joint generative model, requiring no prior over model parameters and enabling fast amortized inference.Classical EM or MAP-EM first fits a joint model and then performs inference to obtain conditional marginals.Source: §1 · p. 2
  • Conditional log-likelihood trains predictions to assign high probability to true values, while a single forward pass predicts all missing marginals and shares computation across them.The inference pass is differentiable, so training directly optimizes its accuracy.Source: §1 · p. 3
  • A relational-attention extension supports domains with one partially observed sample and many interrelated variables, where shared parameters capture recurring relationships.This addresses settings in which variables may be observed only once or never observed.Source: §1 · p. 3
  • Training masks observed values to predict them from the remainder; this recipe is justified for originally MCAR data with additional MAR masking, but not automatically for MAR evaluations.The paper's main synthetic evaluations use MCAR data; its real annotation data has an unverified MCAR assumption.Source: §1 · p. 4
  • The Marformer is competitive with classical approaches even when they know the true model and prior, and generalizes from masked-value training to predicting actually missing values.It can learn from modestly sized datasets.Source: §1 · p. 4

2 Problem Setting

This section defines Marformer’s task of estimating conditional marginals for unobserved attributes from observed subsets, including when those conditionals are valid under missingness. It describes how incomplete observations yield training examples, how masked supervision is handled and predictions evaluated, and how training and test missingness differ across transductive and non-transductive setups.

2.2 Conditional Marginals and Their Uses

  • Conditional marginals for any unobserved attribute given any observed subset support point prediction, joint prediction by chaining, Bayes-risk action evaluation, and value-of-information acquisition.The paper focuses on making these estimates straightforward, fast, and accurate with a generic model rather than a domain-specific one.Source: §2.2 · pp. 4–6

2.3 Predicting From Conditional Marginals on Incomplete Data

  • For incomplete test data, conditioning on the observed pattern can change missing-value predictions; under IMAR, the missingness likelihood is constant and can be dropped, but general MNAR prediction requires it.That likelihood cannot be learned from observed instances alone.Source: §2.3 · pp. 6–7
  • MAR makes joint prediction of all missing variables ignorable, and datasets produced by active feature acquisition are always MAR.The acquisition policy cannot use currently missing attributes when choosing what to observe or when to stop.Source: §2.3 · p. 7

2.4 Training Conditional Marginals on Incomplete Data

  • CMLE trains on examples pairing an observed subset with a held-out attribute and value, weighting the fit toward input conditions that occur more often in training.A single incomplete training instance can yield multiple masked examples, whose terms can be computed together in one forward pass.Source: §2.4 · pp. 7–8
  • Masking observed values can bias conditional targets even under MAR or MCAR; with MCAR training instances, IMAR masking yields unbiased supervision, while MAR data require imputation before IMAR masking.The paper simplifies its training setup by requiring MCAR training instances, avoiding the need for weighting or imputation.Source: §2.4 · pp. 8–9
  • Training and test are assumed to share the complete-data distribution, but their missingness mechanisms may differ: training missingness informs masking, while test missingness informs prediction.The experiments leave deriving training examples for a query-distribution mismatch to future work, noting that the training and target query distributions are already reasonably close.Source: §2.4 · p. 9

2.5 Experimental Setup

  • At training time, observed attributes are masked to create conditional prediction examples; the experiments independently mask each observed attribute with 15% probability.The masking rate aims to keep masked training inputs similar to unmasked test inputs while producing several examples per sampled instance.Source: §2.5 · pp. 9–10
  • Test predictions target actually missing attributes and are evaluated against complete instances using log-loss, a proper scoring rule for conditional distributions.Each missing attribute yields an individually evaluated example, with test examples weighted equally.Source: §2.5 · pp. 9–10
  • Classical generative baselines use the same training and test sets and evaluation metric, but fit a joint model rather than converting observations into conditional training examples.The experiments’ MCAR training data are also MAR, so joint-model fitting uses an ignorable likelihood contribution.Source: §2.5 · p. 10
  • Transductive training adds incomplete test inputs as training inputs, while withholding their missing outputs; non-transductive Marformer testing is faster and has comparable accuracy at large training sizes.The transductive setup enables a fair comparison with Bayesian MCMC, which jointly samples over parameters and missing values using training and test data.Source: §2.5 · pp. 10–11

3 A Neural Architecture for Conditional Marginal Prediction

The Marformer encodes observed and missing variables as per-variable representations, then refines them through attention to predict all conditional marginals in parallel. Relational attention supports prediction for rarely or never observed attributes, while training regularizers encourage copying observations and using relational evidence.

3.1 Marformer Architecture

  • Each input variable is encoded as a vector, and Transformer layers update per-variable residual streams through attention before the final layer supplies conditional marginals.Observed attributes should retain their values, while predictions for missing attributes can sharpen from a uniform input distribution.Source: §3.1 · pp. 11–12
  • Relational attention lets attributes consult related attributes, enabling prediction by analogy even for attributes never predicted during training.The model can learn to consult ratings of the same document, for example, and apply that pattern to a new attribute.Source: §3.1 · pp. 12–13
  • Relational attention can build entity representations from connected attributes, so rarely seen entities can be represented through their relations.An entity’s residual stream can aggregate its annotations, which other attributes can then consult when making predictions.Source: §3.1 · p. 13

3.2 Marformer Training

  • The copy regularizer trains each attribute’s residual stream to provide a conditional distribution whether its value was observed or inferred.Deep supervision is discussed as a possible addition, but whether it improves training remains open.Source: §3.2 · pp. 14–15
  • Entity dropout omits attribute-specific embeddings during training, encouraging predictions to rely on type embeddings and relational evidence.The annotation experiments use a fixed dropout rate of 0.7 on the generalized axis and 0 on the other axis.Source: §3.2 · p. 15

4 Domain 1: Graphical Models

Domain 1 tests whether the Marformer can predict missing values in randomly generated Bayesian networks without knowing their conditional-independence structure. It is compared with joint modeling that knows the true family and prior; at larger training sizes, larger Marformers achieve lower KL divergence and are faster to train.

4 Domain 1: Graphical Models

  • The experiment evaluates missing-variable prediction in Bayesian networks without giving the Marformer their conditional-independence structure, unlike the classical methods.The synthetic domains use MCAR missingness and known generating distributions, allowing evaluation against true conditional marginals.Source: §4 · p. 16

4.1 True Generative Distribution

  • The true distribution is a 10-variable binary Bayesian network with logistic conditional probabilities, and each variable is observed independently with probability one-half.The MCAR mechanism is ignorable for all methods.Source: §4.1 · p. 16

4.2 Datasets

  • Each run uses 250 test instances and training sets from 10 to 2,000 incomplete instances, with ten generated distributions evaluated.Training datasets are nested across sizes, and each model is trained independently.Source: §4.2 · pp. 16–17

4.3 Marformer Configuration

  • The Marformer experiments compare Tiny, Small, and Large configurations, mask each observed training value independently with probability 0.15, and use no dropout.The variants range from 2 layers and 32 dimensions to 6 layers and 128 dimensions.Source: §4.3 · p. 17

4.4 Joint Model Estimation

  • MAP-EM estimates the joint model’s parameters even though it is given the true DAG, logistic CPT family, and prior; ten random restarts address sensitivity to initialization.Its predictions are then computed exactly and compared with the true marginals.Source: §4.4 · p. 18

4.5 Experimental Results

  • Larger Marformers reliably achieve lower KL divergence as training data grows, while joint modeling performs worse than every variant except Tiny despite knowing the true family and prior.With 1,010 training instances, Large is more accurate at most predictions; at 2,000, Small is as well.Source: Figure 4 · p. 17, §4.5 · pp. 18–19

4.6 Runtime

  • MAP-EM is slower to train than the Marformer even on these small networks, and its inference cost grows exponentially with network size.The classical experiments were feasible only because the networks were small; the Marformer’s runtime grows quadratically with the graph.Source: §4.6 · p. 19, Figure 6 · p. 19

5 Domain 2: Discretized Multivariate Gaussians

This domain tests prediction of missing ordinal ratings produced by discretizing correlated latent Gaussian variables. With enough training data, Marformer variants outperform the MAP joint-model baseline, while MAP performs better with less data.

5 Domain 2: Discretized Multivariate Gaussians

  • The domain is an ordered-probit setting: correlated latent continuous variables are discretized into ordinal ratings, and experiments use 10 variables rather than the 2 shown.The discretization uses fixed, irregular bins, whose widths also affect the observable probabilities.Source: §5 · p. 20, Figure 7 · p. 20

5.1 True Generative Distribution

  • The true distribution has 10 attributes with five categories, generated by thresholding a multivariate normal latent vector into roughly equal-probability bins.The latent Gaussian covariance is positive definite, and the category thresholds are derived from the marginal latent distributions.Source: §5.1 · pp. 20–21
  • The joint-model baseline knows the true covariance sparsity pattern, while Marformer is not told the pattern, the covariance structure, or even category order.The covariance precision matrix uses a path or cycle structure.Source: §5 · p. 20, §5.1 · p. 21

5.2 Datasets

  • The study generates 100 test instances and evaluates training sets of 100 to 2,400 incomplete instances, repeating experiments across 10 sampled distributions.The training-size values are 100, 500, 800, 1,200, 1,600, 2,000, and 2,400.Source: §5.2 · p. 21

5.3 Marformer Configuration

  • Marformer uses the same training procedures and metrics as the prior domain, encoding each five-category variable with a shared type and a five-logit representation.Observed categories map to one-hot vectors, while missing values map to the zero vector.Source: §5.3 · p. 21

5.4 Joint Model Estimation

  • The MAP baseline optimizes the observed-data likelihood plus priors over covariance and bin probabilities, estimating rectangular-region probabilities with Genz’s Monte Carlo integration.It uses stochastic gradient ascent with AdamW, allows 10 random restarts, and predicts test values with Gibbs sampling.Source: §5.4 · pp. 21–22

5.5 Experimental Results

  • With enough training data, all Marformer variants achieve significantly lower KL divergence than MAP; MAP performs better with less data and improves when initialized at ground-truth parameters.Per-prediction cross-entropy is also lower for Marformer on most marginal predictions with enough data, especially lower-entropy marginals.Source: §5.5 · p. 22
  • With enough training data, Marformer—particularly larger variants—achieves lower cross-entropy than the joint model for most marginal predictions, especially lower-entropy marginals.Those marginals have a clearer right answer.Source: §5.5 · p. 22

6 Annotation as a Domain

This section studies conditional-marginal prediction for sparse annotation, where calibrated predictions over missing ratings can support decisions about which annotations to acquire. It describes the synthetic annotation setup and evaluation, then reports Marformer’s performance, judge generalization, calibration, and runtime results.

6.1 True Generative Distribution

  • The CPM generates ordinal ratings from item embeddings and compositional judge preferences, assuming semantic item structure, criterion-specific evaluation directions, and partially shared judge preferences.The setting also requires generalization to new items or judges that may enter mid-study.Source: §6.1 · p. 23

6.2 Dataset Construction and Missingness

  • The synthetic dataset uses MCAR missingness and evaluates 316 fixed hidden targets in a 100-item, 10-judge test corner as training data grows.Transductive models can use observed ratings from test entities during training; Marformer NT excludes those entities until prediction.Source: §6.2 · pp. 24–25

6.3 Marformer Configuration for the Annotation Domain

  • The annotation Marformer links rating variables to item, judge, and criterion entities, and trains with artificial masks that emphasize masked observed ratings.Each mask hides 15% of observed ratings; masked and unmasked observed targets receive loss weights 15 and 1, respectively.Source: §6.3 · p. 25

6.4 Joint-Model Estimation with Stan

  • Stan Oracle uses the true CPM family and generating prior families but infers latent parameters from ratings; three restricted CPM variants test misspecification.Stan averages posterior predictive distributions over HMC samples, integrating uncertainty in latent parameters.Source: §6.4 · p. 26

6.5 Other Baselines

  • The empirical baselines are pooled conditional frequencies, Naive Bayes with criterion and cross-judge evidence, and a jointly trained log-linear model without conditional-independence assumptions.All three use observed test-corner ratings during training, whereas a separately evaluated Marformer NT lacks that transductive advantage.Source: §6.5 · pp. 26–27

6.6 Experiments and Results

  • Evaluation covers item and judge generalization on the fixed test corner, measuring log-loss, MBR-L2 for conditional-mean predictions, and smooth expected calibration error.The experiments sweep training data along the item or judge axis.Source: §6.6 · p. 27
  • At Itrain = 4000, log-loss is 0.683 for the Marformer versus 0.600 for Stan; MBR-L2 shows the same qualitative gap, while both improve in judge generalization.Stan uses the correctly specified generative model; the Marformer learns conditional structure from observed ratings. Small differences between individual item-training sizes do not establish a systematic advantage for either Marformer variant.Source: Figure 8 · p. 28, §6.6 · p. 28
  • Conditional predictions also enable chained sampling and VOI-based acquisition.The standard transductive Marformer is evaluated alongside Marformer NT, which excludes whole test entities during training.Source: §6.6 · pp. 27–28
  • At Itrain = 4000, smECE is 0.020 for the Marformer, 0.022 for Marformer NT, and 0.025 for Stan CPM, indicating comparable calibration at that size.At Itrain = 200–500, the Marformer has higher smECE than Stan Oracle. The conditional log-loss objective rewards matching distributions on artificially masked observations, and learned readouts recover near-final predictions by layer 3, while raw parameter coordinates continue improving through the final layers.Source: §6.6 · p. 28, Figure 9 · p. 29
  • At Itrain = 4000, estimated Marformer runtime is 60.7 minutes versus about 757 minutes for Stan Oracle, a roughly 12× difference.This is a hypothetical estimate assuming approximately constant per-epoch cost and stopping 50 epochs after the selected epoch; it does not establish the cost of recovering the reported checkpoint. At Itrain = 500, totals are 2.5 versus about 52 minutes; on the judge axis, Marformer takes 1.6 minutes at Jtrain = 40 and 2.9 at 50, while Stan takes about 55 minutes at 40.Source: §6.6 · p. 29

7 Applying Marformer to Real Annotation Data

On real LLM-Rubric annotation data, Marformer predicts missing human ratings using observed ratings that include dense LLM judgments. At the largest training size it leads the evaluated baselines on log-loss and MBR-L2, while calibration and the study’s missingness assumptions qualify the findings.

7 Applying Marformer to Real Annotation Data

  • The evaluation assumes retained human ratings are MCAR, an assumption the authors do not verify, and has no known ground-truth distribution or oracle model.The study evaluates transductively and focuses on generalization across items with the full judge set fixed.Source: §7 · p. 29

7.1 LLM-Rubric

  • The dataset has 225 conversations, nine rubric dimensions, and ratings from 24 human judges plus one LLM judge; human ratings are sparse while the LLM slice is fully observed.The 25 test conversations’ 666 collected human ratings are scored, with their LLM ratings visible as context; uncollected human ratings are not evaluation targets.Source: §7.1 · p. 30
  • At Itrain = 175, log-loss is 1.007 for Marformer versus 1.093 for Stan CPM and 1.112 for Pooled Unigram; MBR-L2 is 0.59 versus 0.78 for Stan CPM.The Stan CPM baseline has no privileged access to the correct generative family in this real-data evaluation.Source: Figure 10 · p. 31, §7.1 · p. 31
  • At Itrain = 175, Marformer is well-calibrated, but at Itrain = 10 its calibration is worse than Stan CPM’s.At the smaller training size, calibration degrades for all methods; MIWAE and ReMasker are overconfident at the scale extremes.Source: Figure 11 · p. 31, §7.1 · p. 32
  • At Itrain = 175, recorded training takes approximately 22 minutes for Marformer and 162 minutes for Stan CPM.Marformer’s time includes its 300-epoch run and best-validation checkpoint selection, without refitting.Source: §7.1 · p. 32

8 Other Related Work

Related work spans classical and generative imputation, Transformer-based tabular methods, masked autoencoding, and active feature acquisition. Marformer is distinguished by directly predicting conditional marginals with shared parameters and relational structure, rather than providing an acquisition policy.

Classical and generative methods

  • Generative imputation models distributions over missing values, whereas this work focuses on individual conditional marginals aligned with downstream decision tasks.Classical approaches include MICE, SoftImpute, and MissForest; generative examples include GAIN, MisGAN, HI-VAE, VAEAC, and MIWAE.Source: §8 · p. 32

Transformers for tabular modeling

  • Unlike tabular Transformer imputation methods that use feature-specific embeddings or fixed columns, Marformer uses compositional entity embeddings for structured variables.The related methods include NAIM and UnmaskingTrees.Source: §8 · p. 32

Masked Autoencoding

  • Marformer frames masked imputation as conditional distribution prediction, captures structured variables with compositional entity embeddings, and uses an encoder-only design without positional order.This contrasts with ReMasker’s repeated re-masking and reconstruction of observed entries.Source: §8 · p. 32

Active feature acquisition

  • Active feature-acquisition methods learn a sequential policy, while Marformer supplies conditional marginals as a prerequisite for VOI computation rather than the policy itself.Combining Marformer with a VOI-based acquisition policy is identified as future work.Source: §8 · p. 32

9 Discussion and Limitations

The Marformer directly predicts conditional marginals, matching or outperforming methods with the correct generative family across three synthetic domains and beating evaluated baselines on real data at the largest training size. Its limitations include assumptions about missingness, marginal compatibility, calibration, and generalization to underconstrained inputs.

9 Discussion and Limitations

  • Across three synthetic domains, Marformer matches or outperforms methods with the correct family, though it trails Stan Oracle in annotation; on real data, it beats baselines at largest scale.On real data, it runs at a fraction of Stan’s cost.Source: §9 · p. 33
  • Masked-autoencoding justification assumes MCAR training; under MAR, artificial masking may bias learned conditional distributions relative to naturally missing ones.The passage specifies that this can occur when missingness depends on other observed values.Source: §9 · p. 33
  • Conditional marginals are not constrained to form a coherent joint, and miscalibration can yield overconfident VOI estimates and suboptimal acquisition decisions.Source: §9 · p. 33
  • With limited training data, calibration can deteriorate, and test cases unlike training examples may be poorly generalized, especially when dependence is underconstrained by observed data.In the MCAR bandit example, training never predicts X2 from X1, so whether they depend rests on the learner’s inductive bias.Source: §9 · p. 33
  • An ensemble mixture can retain high predictive entropy when the correlation sign is unknown, rather than committing to one plausible relationship.Source: §9 · p. 33

A Results in MAR Setting

The appendix reruns comparisons under MAR missingness, where the Marformer’s naive masking produces biased training examples. Learning curves remain qualitatively similar to MCAR, although losses rise for all methods in Domain 1.

A Results in MAR Setting

  • Under MAR training, performance is qualitatively similar to MCAR, though loss is higher for all methods in Domain 1.Source: Figure 12 · p. 39
  • The Domain 1 comparisons use MAR train and test data: MAP-EM training remains valid, while naive Marformer masking creates biased examples; Domain 2 is also rerun under MAR with independent 0.15 masking.Figures 12 and 13 show the results, including learning curves and prediction-by-prediction comparisons.Source: A Results in MAR Setting · pp. 39–40, Figure 13 · p. 39
  • In Domain 1, the sequential MAR mechanism uses logistic factors with parent-variable indicators for missingness and observed value 1; about half of variables are missing per instance.The factor parameters are drawn i.i.d. from N(0, 1).Source: A Results in MAR Setting · p. 39

B Domain Specifications

The domain algorithms provide complete specifications for data generation and modeling in the Bayesian-network and discretized multivariate Gaussian domains. The supplied excerpt identifies the BN joint-generation expression and the train/test split.

B Domain Specifications

Stan HMC configuration

The appendix reports HMC configuration, compares full baseline curves, and gives the calculation used to obtain Stan’s categorical posterior predictive marginals. The figures describe model performance and runtime accounting.

Stan HMC configuration

  • Stan predictive fits use one HMC chain, 300 warmup iterations, and 500 post-warmup samples, with adapt_delta=0.85 and max_treedepth=12.Initialization is random.Source: Stan HMC configuration · p. 40
  • At larger training sizes in both settings, empirical baselines generally have higher log-loss and MBR-L2 squared error than Marformer and Stan variants.Source: Figure 14 · p. 41, Figure 15 · p. 42

C.1 Stan Posterior Predictive Calculation

  • The categorical predictive distribution is calculated for each posterior sample, then averaged across samples to obtain the reported posterior predictive marginal.Source: §C.1 · pp. 41–42
  • The runtime figure reports synthetic annotation minutes; Marformer costs estimate shortened validation search plus measured refit, while Stan costs cover measured HMC fitting.Source: Figure 16 · p. 42

Stan CPM Posterior Hyperparameters

Stan CPM posterior estimates indicate prototype-level variation is the dominant source of judge-to-judge differences, while measurement noise is small. At Itrain = 175, Marformer training is also substantially faster than Stan CPM.

Stan CPM Posterior Hyperparameters

  • The posterior mean for σv is ≈11.6, versus ≈2.1 for σu and ≈1.2 for σuit, indicating prototype-level variation dominates judge-to-judge differences.The posterior summaries use HMC samples from the full LLMRubric training set with weakly informative priors.Source: Stan CPM Posterior Hyperparameters · p. 42, Table 1 · p. 43
  • The small σmeas posterior mean (≈0.38) suggests ordinal thresholds capture most of the rating variability.These posterior distributions informed prior scale choices for synthetic data generation in §6.1.Source: Stan CPM Posterior Hyperparameters · p. 42, Table 1 · p. 43
  • At Itrain = 175, Marformer training took approximately 22 minutes, compared with 162 minutes for Stan CPM.Marformer used best-validation checkpoint selection without refitting.Source: Figure 17 · p. 43

E Logit Lens Analysis

The layerwise readout analysis examines how predictions develop in an eight-layer Marformer on 316 missing test ratings. Tuned probes recover predictions close to the final model earlier than the raw parameter readout, which continues improving through later layers.

E Logit Lens Analysis

  • The analysis fits raw and tuned readouts at each layer, with tuned probes trained to reproduce the final model’s predicted distributions.The frozen Marformer is evaluated on 316 missing test ratings at Itrain = 500; probes are fitted separately using held-out training-missing ratings.Source: E Logit Lens Analysis · p. 44
  • By layer 3, both tuned lenses have log-loss close to the final model’s 0.827; raw Param Lens log-loss falls from 1.117 at layer 6 to 0.827 at layer 8.MBR-L2 shows the same broad pattern: tuned readouts approach the final model early, while the raw readout continues improving through later layers.Source: Figure 18 · p. 44, E Logit Lens Analysis · p. 44
  • After one attention layer, Tuned 80d-Lens log-loss is 0.946, versus 1.264 for raw Param Lens and 0.970 for Tuned Param Lens.The full-state probe has slightly lower loss than the parameter-only probe at this layer.Source: Figure 18 · p. 44, E Logit Lens Analysis · p. 44
  • The authors suggest a possible two-round message-passing interpretation: early attention gathers direct evidence, while later layers may relay indirect evidence.This is presented as a hypothesis consistent with iterative refinement, not as a demonstrated mechanism.Source: E Logit Lens Analysis · pp. 44–45
  • At layer 4, raw Param Lens smECE is 0.073, compared with 0.026 for Tuned Param Lens and 0.043 for Tuned 80d-Lens.At the output, raw Param Lens smECE is 0.042; calibration changes across depth, though all three readouts remain reasonably well calibrated.Source: Figure 19 · p. 45, E Logit Lens Analysis · p. 45
  • The raw layer-0 Param Lens predicts 1/4 for every category, while its zero all-class smECE accompanies log-loss ln 4.This illustrates that calibration alone does not measure how informative a prediction is.Source: Figure 19 · p. 45, E Logit Lens Analysis · p. 45
Loading 2610.12379v1…