Ex 2.2.19, follow-up: where the calibration KL sits¶

A well-calibrated model is as sure of its answer as the examples warrant: sure when they point to one op, spread out when they fit several. The calibration KL of ex-2.2.19 is how far the seven-op model is from that, over all held-out contexts, and this note asks where that distance comes from. The contexts whose examples fit several ops are a quarter of the held-out set and hold 42% of the KL, and there the model is too sure of itself: it puts 0.54 on its favorite answer where the examples warrant 0.45. Where the examples point to one op, which is where the op confusion matrices look, it is about as sure as it should be, and those contexts hold 9% of the KL. Training for 400 epochs instead of 200 leaves the total where it was but moves it: at every seed the longer run is a little better calibrated where the op is clear and worse where it is not. By op, the KL is spread over all seven, highest on mix.

Scope¶

Ex-2.2.19 measured two things about how a model spreads its answers. The calibration KL¹ is taken over every held-out context. The op confusion matrix (E3) shows which op answers the model leans toward, but only on confident contexts, where the Bayes posterior on the true op is above 0.99. We ask whether a matrix like the confusion one can be made over the whole held-out set, showing how much each op contributes to the KL (backlog item).

This note looks at the KL two ways, by how clearly the examples pick out the op and by op, on the 200- and 400-epoch runs of the seven-op set at four seeds each (4 runs per length). It is post hoc, with no hypotheses or gates, and trains nothing. Each length has 14,000 held-out contexts, the same contexts in every run.

By confidence¶

The first way groups contexts by the Bayes posterior on the true op: the probability a perfect reader of the examples would give the op that generated them. Where it is near one the examples fit one op alone; where it is below a half they fit several about as well. The four groups run from a posterior below 0.5 to above 0.99, and the last is the confident set of the confusion matrices. The first figure shows how the contexts and the KL spread over them.

The largest share of the KL is in the least confident contexts. The group below 0.5 is 26% of the contexts and holds 42% of the KL at 200 epochs and 45% at 400. The confident set is 23% of the contexts and holds 9% and 7%. In the figure, the line sits above the bar on the left and below it on the right.

Per context, the KL falls steadily as the posterior sharpens (next figure), from about 0.53 nats below 0.5 to 0.13 above 0.99. The two lengths cross: from 200 to 400 epochs, the KL per context below 0.5 rises by 0.039 to 0.072 over the four seeds, and in the confident set it falls by 0.007 to 0.029. Every seed moves the same way at both ends. The total barely changes (0.323 against 0.335), which is what ex-2.2.19 reported as calibration not separating the two lengths.

The last figure shows how sure the model is against how sure it should be: the mass each puts on its single most likely answer.

Below a posterior of 0.5, Bayes puts 0.45 on its top answer and the model puts 0.54 at 200 epochs and 0.56 at 400, more at every seed. Its top answer is also the Bayes top answer in only 55% of those contexts, so it leans the wrong way about as often as the right one. In the other three groups the model is within 0.03 of Bayes. So on confident contexts the longer run gains elsewhere than the top answer, in the small amounts of mass on the other colors: at 400 epochs, less of the KL of the confident set goes to colors that no op gives (18%, against 21% at 200).

So the model does know when the examples are clear. Its confidence tracks the posterior from one end of the axis to the other, and where the examples fit one op it is as sure as it should be, and the little KL there is small amounts of mass spread over wrong colors. Where the examples fit several ops it does spread its answer, but it leans toward one op more than the examples justify, and often toward the wrong one. Longer training makes that lean stronger while tidying the clear contexts, so the total stands still. The largest part of the KL that remains is this one habit.

By op¶

The second way makes the matrix. Rows are the true op, and a row adds up to the part of the KL on the contexts of that op. Columns say where the model has mass in excess of the Bayes answer distribution, and each context divides its KL among the columns in proportion to that excess. Excess on a color that several ops of the set give is shared among them by the Bayes posterior over ops, since the posterior the model itself holds over ops is not something we can measure: the arrays hold its answer distribution and nothing of how it got there.

Three columns catch excess on colors outside the set. The seven-op set is the no-four set of ex-2.2.18: the eleven ops of ex-2.2.16 less screen, multiply, exclusion, and hsvmix. Those four dropped ops are left out of the corpus the model trains on, and nothing is edited: the model is the plain trained model. Their answers are still defined on the color table, so a color that only a dropped op gives has the column dropped; a color that no op gives on the query pair has none; and mass on tokens that are not colors has non-color. (difference is in the set, so it has a row like any other.)

The column split is somewhat arbitrary. The KL counts the mass that the model is short of on the Bayes answers, and the excess says where that mass went; giving the KL to the excess in proportion is one way to join them. The rows add up either way.

In both maps a higher number is more KL, so worse. A high row says the model is poorly calibrated on the contexts of that op, with too much or too little mass on any answer; it does not say which op the model answered with. A high column says the model puts too much mass on the answers of that op; a column never counts too little.

By row, the KL is spread over all seven ops. The three HSV-channel ops hold 48% of it at 200 epochs, for three ops of seven. The highest KL per context is on mix (0.42), and the lowest on lighten and darken (0.24 and 0.24). On confident contexts, the confusion matrix of ex-2.2.19 put most of the extra mass of the shorter run in the HSV-channel rows. Over all contexts, the KL does not gather there to the same degree.

By column, the true op takes 27% of the KL, the other six ops of the set 51%, colors no op gives 17%, and the dropped ops 5%. Mass on tokens that are not colors is negligible. Below a posterior of 0.5 the other ops take 60%, consistent with a model that leans further than Bayes toward one op when the examples allow several. No single other op stands out in any row. From 200 to 400 epochs the share of the true op rises to 31% and that of colors no op gives falls to 14%: the longer run moves mass off the colors no op gives and onto the answers of the true op.

What it means for what follows¶

The KL that remains at 200 and 400 epochs is mostly in contexts whose op is unclear, which the confusion matrices leave out. So a lower KL would come mainly from better hedging between ops on those contexts, and less from the HSV-channel ops, whose shortfall shows in exact match on confident contexts.

Is Bayes a fair standard there? The Bayes answer distribution is the distribution of the answer given the examples under the process that made the corpus: an op drawn uniformly, examples with replacement op noise and cube noise at the rates the corpus used, and a clean query. It uses nothing about a context beyond its examples, which the model also sees, and a held-out context is a fresh draw from the same process. A model that has learned the process is calibrated on it, and cross-entropy training pulls toward that distribution, so zero KL is reachable in principle and the standard is no oracle. The geometry of the cube is part of the standard too: the answer table is each op applied to every ordered pair of grid colors, so an op that clips at a face of the cube, or rounds to several neighboring colors, has those answers at those probabilities in the posterior and in the answer distribution alike. What the model has to learn, and has not fully learned at either length, is how much the examples leave open.

The longer run trades calibration on unclear contexts for calibration on clear ones. That fits a model that keeps sharpening its answers with training, and it could be overfitting to the training corpus; the held-out arrays alone cannot tell the two apart, since there is no training-set KL to set them beside. One can be had without retraining: the end checkpoints of every run are stored, so an evaluation pass over a sample of training contexts, a few minutes of L4 time per run, would give it (backlog item). Either way, a level total KL between two lengths does not mean the two are calibrated alike.

Method¶

The runs are the 200- and 400-epoch runs of the seven-op set in ex-2.2.19, at seeds 600, 601, 602, 603; the 400-epoch run at seed 600 is the ex-2.2.18 no-four run. Each run stored, per held-out context, the true op, the posterior over the seven ops, the query pair, and the model answer distribution over the grid (in float16).

The model answer distribution is the softmax over the whole vocabulary at the query =, the position that predicts the answer, taken at the color tokens. Mass on tokens that are not colors is kept as loss rather than renormalized away, and the greedy answer plays no part. The Bayes answer distribution is rebuilt from the posterior and the answer table, and the KL from it agrees with the stored KL to within 0.0013 on average per run, the difference coming from the float16 storage.