Induced Model Matching: Restricted Models Help Train Full-Featured Models
Authors: Usama Muneeb
Mesrob I Ohannessian
Venue: NeurIPS 2024 spotlight (2024)
Keywords: restricted
model
learning
Primary area: other

Abstract: We consider scenarios where a very accurate (often small) predictive model using restricted features is available when training a full-featured (often larger) model. This restricted model may be thought of as ``side-information'', and can come either from an auxiliary dataset or from the same dataset by forcing the restriction. How can the restricted model be useful to the full model? To answer this, we introduce a methodology called Induced Model Matching (IMM). IMM aligns the context-restricted, or induced, version of the large model with the restricted model. We relate IMM to approaches such as noising, which is implicit in addressing the problem, and reverse knowledge distillation from weak teachers, which is explicit but does not exploit restriction being the nature of the weakness. We show that these prior methods can be thought of as approximations to IMM and can be problematic in terms of consistency. Experimentally, we first motivate IMM using logistic regression as a toy example. We then explore it in language modeling, the application that initially inspired it, and demonstrate it on both LSTM and transformer full models, using bigrams as restricted models. We lastly give a simple RL example, which shows that POMDP policies can help learn better MDP policies. The IMM principle is thus generally applicable in common scenarios where restricted data is cheaper to collect or restricted models are easier to learn.

Induced Model Matching: Restricted Models Help Train Full-Featured Models
Usama Muneeb
Mesrob I. Ohannessian Electrical and Computer Engineering University of Illinois Chicago
Electrical and Computer Engineering University of Illinois Chicago
umunee2Quic.ede
mesrobQuic.edd
Abstract
We consider scenarios where very accurate (often small) predictive model using restricted features available when training full-featured (often larger) model. This restricted model may be thought of as "side-information" and can come either from an auxiliary dataset or from the same dataset by forcing the restriction. How the restricted model be useful to the full model? To answer this, we introduce methodology called Induced Model Matching (IMM). IMM aligns the context-restrieted or induced, version of the large model with the restricted model. We relate IMM to approaches such as noising, which is implicit addressing the problem, and reverse knowledge distillation from weak teachers, which is explicit but does not exploit restriction being the nature the weakness. We show that these prior methods can be thought of as approximations problematic in terms of consistency. Experimentally. we first motivate IMM using logistic regression
IMM and can be
toy example. We then explore it in language modeling, the application that initially inspired it, and demonstrate it on both LSTM and transformer full models, using bigrams as restricted models. We lastly give a simple RL example, which shows that POMDP policies can help learn better MDP policies. The IMM principle thus generally applicable in common scenarios where restricted data is cheaper to collect or restricted models are easier to learn.
Introduction
In many applications, it is both statistically and computationally easier to construct (often small) feature-restricted models. In this paper, address the question of whether this could be beneficial in the construction of an (often larger) full-featured model.
To motivate, consider the following toy logistic regression problem, with further details in Section 22, X3) E R and a binary label y E {0, Let features be
7 Say we have three features ( generated from distribution and let labels be conditionally generated according to a logistic (full-featured) true model P(ylæ1 22, 22, 23). Assume we have ample feature-restricted data with only one feature (x1: y). which, by marginalization, are samples from the (feature-restricted) true induced model P(y](11) Emggrs T 22(zz1) P #2, By virtue of this data, we get very good approximation P of P, which we can use as target induced model. How do we use P, along with full-featured data (I1, 22, va, y) to obtain (ull-featured) learned model Qy)æ11 22, za)? few possible approaches are as follows. We could ignore P. and learn Q simply by minimizing cross-entropy on the data. This is wasteful, because P contains valuable information. Alternatively, we could learn Q by minimizing cross-entropy in addition to secondary loss that keeps Q close to
P. This is reasonable in Section we relate this to reverse knowledge distillation and in Section 6 noising. However, P addresses markedly different task than Q, i.e., that of predicting with
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

restricted set of features, Instead, what this paper proposes is to equalize the field when comparing to P. by inducing a restricted model ((yla)) from Q during training, and using a secondary loss to match P to Q, rather than to Q. We call this induced model matching (IMM). In Figure we show how this speeds up learning and reduces predictor variance.
100
68
Inl wit Trge NcMM
Not Pooess Targe Mode
Inrppaeee Tager Uati
Cotaset SEzn
Figure 1: Comparing test accuracy of logistic model, trained using interpolation, noising, and IMM (with Bayes Optimal P). Bars are 10th 90th percentiles of 300 runs.
This example gives an overview of the entire process behind IMM. Most of the paper, however, is dedicated to language modeling, where very natural restricted models exist, namely N-grams. The inspiration of this research is in fact rooted in certain language model data augmentation techniques that noise the data using N-grams Xie et al. (2017). The fresh perspective that offer here that such noising is best understood as an attempt to incorporate a feature- into the fall-featured model being trained. This interpretation reveals the general fundamental question: If we are armed with very accurate feature-restricted model, what is the right way to incorporate it into the training of a model with larger feature set?
model's knowledge
Our contributions and organization are as follows:
1. In Section 3 we rigorously frame this question through the notion of an induced model.
2. In Section 4 we use this framework to propose a strategy to incorporate the knowledge of the restricted model during learning. This consists of using regularizer that matches the predictions of the induced learned model with that of the desired target restricted model. This the language model instance of the induced model matching (IMM) methodology.
3. In Section 5 propose computational approaches to make IMM practical and pave the way for further scaling IMM.
4. In Section 6 we loop back to relate IMM the noising approach of al. (2017) and to reverse knowledge distillation. We share two key findings. First, these alternatives may be thought of approximations IMM. Second, they have a major caveat. They may not be consistent even in the ideal infinite-date regime, in contrast to IMM, which is consistent.
5. In Section 7 we experimentally demonstrate the effectiveness IMM through: (1) details of the logistic regression example, (2) experiments on an LSTM RNN for language modeling and on BERT for classification, showing improvements on multiple tasks, and (3) a simple reinforcement learning example that illustrates the further potential of IMM.
In Section we discuss limitations and applications beyond those illustrated in the paper.
Related Works
We review two lines of work that closely related to IMM, noising and knowledge distillation, which are respectively implicit and explicit versions the same idea. We also review key literature showing the continued merit of N N-grams restricted models.
Noising and Data Augmentation Language Modeling The noising methodology language modeling was proposed initially in order to provide a Natural Language Processing (NLP) parallel to the many data augmentation techniques that existed in other machine learning applications (e.g.. sampling translated and rotated variants, image processing and computer vision).

---

Most data augmentation techniques in NLP have used noising or smoothing in one form or another. Feature noising by Wang al. (2013) was one of the earliest attempts at noising when it came to structured prediction tasks. Later, Xie et al. (2017) successfully used restricted language models, namely unigrams and bigrams, to perform data augmentation through noising. By making a rough connection between noising and smoothing, Xie et al. (017)) were able to show a clear advantage, beyond the use of other regularization techniques such as dropout. Section more closely examines the claims of that paper and connects it to the framework of the present one.
The current approach is not related to all data augmentation techniques, rather specifically to those that use noising with a restricted model. Indeed, it can complement other data augmentation techniques. For example. in our BERT experiments, this technique complements the MLM objective, where the masking itself can be thought of as data augmentation. Similarly, it is capable of complementing alternative masking and corruption techniques, such as EDA (Wei and Zou 2019) and SSMBA (Ng et al., 2020), which are known to outperform BERT"s default masking for the MLM objective. Similarly, IMM's gains are in addition to those from other regularization techniques, e.g., weight decay and dropout (Srivastava et al., 2014), as these can be simultaneously utilized.
N N-grams and their Merits N-grams are based on co-occurrence counts, which makes them easy to learn but limits their usefulness for long histories. However, common techniques such as smoothing and backoff make it possible to build excellent short-context models such as bigrams and trigrams, which can rival or exceed neural models that use the same context size (Chelba et al. 2017). This hints at there being value in using these N-gram models to improve long-context models. Some of the earliest attempts at this interpolate the output of the smaller model with modern models. The continued relevance of this approach is evidenced in a very recent paper (Liu et al.) 2024) that proposes special data structure to precompute N N-grams to arbitrary lengths and uses interpolation to improve larger language models (LMs). A key motivator of the current paper is the approach of Xie et al. (2017), which instead noises the training data of LMs using N-grams, with improved outcomes over interpolation. Note that other approaches that take advantage of N N-grams also exist, such as that of Li al. 2022), who let LMs learn what the N N-gram could not, i.e., the residual.
Knowledge Distillation Knowledge distillation (KD) is a paradigm first proposed to let powerful teacher model help better train weaker student model, by complementing the hard labels of existing data with soft labels (Hinton et al. 2015). In the present notation, the resulting average loss takes the following general form.
Cross-Entropy (Q) + + (=)())DK (p (teaachrr
(I)
where Cross-
Cross-Entropy Entropy uses hard data, x are contexts, Tn is the empirical distribution of contexts, Q is the learned/student prediction model, and pleacher is the teacher model, which can be thought of as providing soft predictions. T, the softmax temperature, is used for smoothing the outputs.
Most relevant to IMM is the recent discovery that, paradoxically, KD with a weak teacher can be helpful to powerful student. This "reverse-KID or "distillation from weak teachers" phenomenon was first demonstrated in vision by Yuan et al. (2020). who interpret and ascribe its performance to smoothing targets in context-dependent way using the teacher, in contrast to typical label smoothing that can be interpreted as using a uniform distribution. Though this comes years after the noising papers, it parallels closely how Xie al. (2017) transition from uniform noising to Kneser-Ney noising. Later, reverse-KD KD was also shown to be effective in language models (Qin et al. 2021).
On the surface, IMM is similar to reverse-KD as it uses a weak target teacher to regularize powerful student's 's learning. A quick comparison to Eq. (14) when pleacher is the restricted model P reveals that reverse-KD KD has more common with noising, which we show can be sub-optimal. Indeed, the major difference between reverse- KD and IMM is the fact that the weakness of the teacher in this case is of a very particular nature - it stems from the reliance on a restricted context. Reverse-KI KD ignores this fact by comparing the student model Q directly to the teacher model in the KL term of Eq. In contrast, IMM harnesses this fact, by comparing the student indirectly to the teacher (target model) at its own level, i.e., using the induced model Q. Very recent work by Lee et al. (2023) shows that reverse-KD KD in language, unlike in vision, can be harmful with very weak teachers, and thus it not advisable to use reverse-KD with the simple restricted models, e.g., bigrams, that we consider here. Our work shows that they can be effectively incorporated, if we use IMM instead.
3

---

3 Problem Description
Problem Setting We consider the problem of training full model Q(y((re) that is able predict a label of data point based on its full context. For example, in forward predictive models, 2; is the sequence of past tokens and y/ the next token. Throughout the paper, Q always denotes the full model and is deemed to belong to a reasonably large functional class. Given context-prediction data (2,, yt) of size n. possibly tokenized from single text, let the primary loss used for training this model be the log-loss. Thus training minimizes the cross-entropy risk:
Cross-Entropy(Q) = --= E log Q(V(lz))= Σ (2(x) Σ Pn log Qylæ))
where Tn is the empirical distribution of the context and Pn be the empirical distribution of the prediction given the context. Note that (It, Mr) refer to tokens while (z, y) refer to types. Minimizing this empirical risk can be seen as Q striving to approach the "true" model P(y (z)) generating the data, in average context-conditional KL divergence. The idealized risk is thus:
Plylz)
Σ =())
Σ
P(ylz) log
Ex(z) DKL (P(+)))Q((12)
(2)
Q
where now is also the true (not empirical) context distribution.
Induced Models While Q strives approximate all of P, we may have additional knowledge about P that captures some of its properties. We consider in particular knowledge of the following form. Assume that the full context H, can decomposed into a short context Ty and an extended context Σ., and that one has access to the "true" model that predicts y/ based solely on the short context Σz, i.e. (T)(T). To make the notation easier to follow, we invite the reader to consult the glossary of notations in Appendix A] How is this restricted model related to P and . By a simple marginalization argument, we can see that this model is:
P(y P(y, Σ ( P()
(3)
We call P the true induced model. depends both on the context distribution and the true model P. Since we do not have the latter two, cannot explicitly compute P. What motivates us, however, is the possibility to learn more accurately than P. either by virtue of its smaller parametrization or thanks to cheaply procured auxiliary context-restricted data.
Problem Statement Given knowledge of the true induced model P. or very good approximation thereof, how can this information be incorporated to improve the learned model Q?
Induced Model Matching (IMM)
Construction of the IMM risk To address the problem, we introduce a secondary loss that matches the learned model's prediction with that of the induced model, whence the name Induced Model Matching or IMM. The key insight here is not match P with Q, but rather with Q, the learned induced model that, just like the move from P to P in Eq. B) specializes Q performing predictions with only the short context:
Q(yIT)= = Σ (
(4)
Let's first idealize and assume availability of P and the context distribution æ, required to compute Q. Equipped with Q, we can introduce the idealized induced model matching (IMM) risk, a secondary risk that is the average context-conditiona KL divergence with the restricted context:
P
Σ T())
P(y log
(5)
Q T
DEL (1

---

Since P and æ are not available in practice, the idealized IMM risk cannot be computed. However, as the core motivation of using is the potential ability to learn it very accurately from data, we assume instead that we have access to a target induced model P as a proxy to P. As for calculating Q mechanism for filling-in for the extended context, based on the short context. As such, we have the following natural empirical version which can be thought of as averaging Q over all extended contexts in the dataset, while keeping the short context fixed:
knowledge of in Eq. 4 can be intuitively understood
Qu1 & L
(6)
The proportionality constant is unimportant as, thanks to the logarithm, it contributes only constant the overall risk. By combining these empirical proxies, we obtain the empirical IMM risk:
IMM(Q)= Σ nn (x)
Σ y
P(y log
()y
This mirrors Eq. 5) up to P insidee the logarithm, which only contributes an additive entropy term that does not depend Q and is thus irrelevant to optimization.
IMM as a regularizer Given reliable P. we propose to incorporate this knowledge into the full model by using the IMM risk as a regularizer. Using hyper-parameetr À that can be tuned, our induced model matching methodology consists of training the model Q by minimizing
Cross-Entropy(Q) À IMM(Q).
8
Separating IMM into components To weave IMM into existing ML optimization pipelines, is useful to treat it as a separable risk. For this, we can rewrite Eq. (7) by expanding the empirical distribution T.. Then, the components of this separable risk are given by the conditional cross- entropies between P and Q, because we can write:
IMM(Q) = = - Σ Σ Ply
IMM+(Q)
The simplest version of P could be based on empirical counts, which may be valid if the context space is discrete and small. In language modeling, smoothing methods can used generate better versions of P, such as the modified Kneser-Ney smoothed version of the bigram counts Kneser and Ney 1995 Chen and Goodman 1999) that can be expressed as
= E v) =z)b(())
where v(z) is the missing mass for previous token I' and is the back-off distribution. For logistic regression in Section the context space is continuous, and these components come from separately trained one-feature predictor. For our RL example, they come from the optimal actions of a POMDP.
Computation
direct implementation of IMM is prohibitive, since evaluation of the risk, Eq. requires making secondary pass over the entire dataset for each f, when inducing the model in Eq. (6). We address this with two approaches. The first approximates the objective by replacing this secondary pass with sampling whereas the second incorporates this secondary pass into the primary pass by serializing the gradient calculation. Sampled IMM has advantage of low gradient variance the expense of added computation. Serialized IMM has higher gradient variance but has only constant-factor computational overhead. All of our experiments use sampled IMM, with the exception of a demonstration of serialized IMM for logistic regression (in Appendix D.5 proof of concept for IMM's scalability.
5

---

Sampled IMM
To alleviate the exact computation of IMM, we could maintain a dictionary/mulliset of extended contexts for each short contex and sweeping only across those. This amounts to rewriting Eq. (6) as follows:
E Qu -
E Q(y(Te,)), where extend(x) (=extten
+
(10)
the multiset of extended contexts of the short context of x.
Since a given short context may appear in a large number of extended contexts, the dictionary/multiset approach of Eq. (1O) remains expensive. In particular, gradients need be computed with each of these contexts. Instead, we propose an approximation based on writing the (normalized) sum as an expectation over samples from extend (T₂). followed by upproximating this expectation with samples. As a result, the IMM component at becomes:
(QM
Σ
P log
E [x-extend( [Q(y
(11)
k Σ Qui
Σ
P(y|T) log
(12)
where ~ (xx, are samples from the multiset of extended contexts. Implementation details including Algorithms/ [ [and[] that describe typical training loop when sampled IMM or serialized IMM are incorporated into SGD respectively.
deferred to Appendix C
At first glance, it appears that sampled IMM requires maintaining k copies of the model. However, this can be sequentialized, as explained in Appendix C22 This means that the space overhead during training is factor of 2 compared to the baseline of no-IMM (i.e., a second of gradients). However, the time overhead is generally a k-fold increase over the baseline, but can be worse for recurrent models such as LSTMs. This is because we typically unroll the model over L tokens and can perform forward/backward passes over this unroll, in L steps. If we apply IMM at every unroll position, we need restart the LSTM. which then incurs an (--ol-fold increase. We partially mitigate this in our current implementation by applying IMM periodically (not at every iteration), see Appendix D.2.2 For example, if we apply only every () iterations, the overhead factor remains (k)).
Serialized IMM
Sampled IMM delivers the benefits of IMM, but at computational cost. To make IMM truly scalable, it imperative bring near-constant factor overhead to the baseline of no-IMM. The main bottleneck is the calculation of the learned induced model Q. The following alternative approach bears a resemblance to the sequentialization aspect covered in Appendix C.2 Eq. (20.. Up to a constant, we can write the idealized IMM risk of Eq. (5) as an averaging only over short contexts:
-[ P(y(T) log ()) where the induced model is Q11FT)
Σ
The gradient of this IMM risk then becomes:
ΣT(T) Σ P Σ. VQVWL2))
(Q(U1z,
(<))
Σ P(y
Q(y(T)
Q(y(TT
The empirical version of this gradient is:
IMM
(Q(ulz,T Q
Q Ie Fe) Q
P(y
Ply
log Quyl,, F.)) (13)
n
n
9
9
correction
The last expression in Eq. 13 is in the form of a "corrected" cross-entropy between 1 and Q which we denote by IMM, The correction factor is not differentiated and the ratio or Nikodyy
'This requires space On log n) only, because each key in the dictionary is short context, and each value a set of pointers to the data, with set size equal to the number of times that short context appears. Adding up the number of occurrences of all short histories gives and each pointer requires Oloo n)) bits.

---

derivative of Q relative to Q. What this accomplishes is to delegate the averaging of the gradients to the primary pass over the dataset. If the correction term is given, computing this gradient costs the same as computing the baseline of no-IMM gradients. There are two caveats: first, because the averaging is now happening in the primary pass, the variance of this gradient higher than sampled IMM and, second, the correction factor still requires knowing the learned induced model Q. To address the higher variance of the gradients, techniques such as momentum approaches along with learning rate schedulers can be used. To address knowing Q, we suggest the heuristic of updating Q only periodically. Then, if model Qi (that eventually becomes stale) is used to calculate Q), use the ratio Q / for correction factor. Using the current Q tends to cause instability, likely due to the correction no longer obeying expected constraints, c.g., mean 1 for every y. By choosing the update period inversely proportionally to the cost of updating Q and by keeping Q1 in memory, we incur an (1) factor increase in time and space compared to no-IMM, which makes this variant of IMM highly scalable. Note that without the correction factor, IMM turns into reverse- KD and noising, connection that we elaborate on in Section
Analysis
In what follows, we assume to be in the realizable case, i.e., that the model space of Q is expressive enough to contain the true model P. We also take an infinite-data perspective and focus on consis- tency only, even though IMM also confers a finite-sample advantage (understanding this advantage analytically worth studying in the future.) We show (1) that IMM is consistent, (2) that noising and eeverse-KD are akin to single-samp bound on the IMM risk, and finally (4) that this introduces sub-optimality, which could even threaten consistency. This gives basic analytical backing to why IMM is preferable.
IMM, (3) that this shows that they minimize an upper
Consistency of IMM Observe that if, in the main objective, Eq. (8) cross-entropy and IMM were replaced with their true counterparts, Eqs. (2) and (5) respectively, then Q = P remains the minimizer of the objective. This observation shows that regime, i.e., that IMM is consistent for all A. We next aim to show that the key caveat of noising and reverse-KD is that they may be inconsistent, unless À is made to vanish.
recover the true model in the infinite-das -ata
From single-sample IMM to noising and reverse-KD We first explain how IMM and noising are related. Experimentally, using single sample (k- = 1) in the IMM approximation of Eq. (12 produces perplexities that are near-identical to noising. We explain this phenomenon as follows. Say data point & is considered during optimization, in an SGD mini-batch. When a single random extended context is sampled, it is equivalent to swapping that data point's context with another based on the multiset extend (Zx.. That context belongs to a different data point L'. Since sampling is done uniformly from the multiset, this simply presents data to the SGD algorithm in a different random order, possibly with repetition. More pertinently to noising, the actual prediction yr has role in calculating the loss. Instead, the prediction is randomized through the sum over all possible y in Eq. ( Though not identical, this is very close map the noising- based data augmentation proposed by Xie et al. (2017), namely prediction noising (see Appendix BB for a review of this framework and why prediction noising is its main aspect.) It is in fact even closer to an alternative proposed later by
Gao et al. (2019), which uses soft prediction vectors just like P in this single-sample approximation of IMM. As a result of this argument, we can think of noising as minimizing the following idealized
objective, which coincidentally is equivalent the reverse-Ki objective Yuan et al., 2020 Qin et al.
2021 (see also Section 2) with P as the teacher:
Cross-Entropy(() + A E (x) Σ
P(y
log (yy
(14)
key difference
Inconsistency of noising and reverse-KD How is this single-sample IMM objective related to performing IMM? Recall that we can write a single component of the empirical IMM risk with sampling as in 122 which by Jensen's inequality can be upper bounded as follows:
P(yl=) log (Ex [Q(v) X))) < - Σ E( [log (Q(v),,, X))]
(15)
7

---

A single-sample approximation of the expectation inside the log is in fact biased estimator of the left-hand side. However, an unbiased estimator of the right-hand side, with the expectation outside of the log. Thus, noising and reverse- KDD upper bound the IMM risk. Minimizing an upper bound instead of the desired risk could introduce suboptimality. The following proposition uses the difference between these methods, which pit the target model against the full learned model Q and not the induced learned model Q like IMM (contrast Eq. (14) and Eq. ()), to show that this suboptimality can indeed occur. Even in the realizable case and with infinite data, IMM always consistent but there exists a counterexample where noising and reverse-Ki KD fail to consistently recover the true model. The proof in Appendix B.2
Proposition 6.1. Assume that we are optimizing the idealized noising objective of Eq. (4)) i.e., we are operating in the infinite-data regime- and let Q" be its global minimizer. Assume further that the model class for Q contains the true model P - iee, we are in the realizable case. Then, there exists a choice of π and P such that Q* P.
7
Experimental Results
7.1 Starting Toy Example: Logistic Regression
Consider the logistic regression example from the introduction. The main results are given in Figure and the full experimental details can be found in Appendix [D.]] Here we highlight how the problem fits the IMM framework and where it deviates from language modeling. First, note that the context decomposition is x I1 and g (I2, T33.
Setting We sample features uniformly over a cube and assume we have ample data points of the form (x1.y). This allows us to build an excellent restricted model to predict the label based on just II, call P(y(æ1), ). nearly close to the (restricted) Bayes predictor or true conditional probability. Just like in language modeling, to induce model we need to draw z's from its conditional distribution given F. Q(y(x1)) the induced model of Q(y z2, 23), can then be interpreted as the average of Q's predictions, when 2, (I2, .z)) is drawn from its conditional distribution given V, 21 Since we typically don't have access to this distribution, we approximate it empirically. In language modeling, we could just sample from the empirical distribution of for given T. In logistic regression, this not viable since Ij is continuous and does not repeat. We rely instead on density estimation. We use (al.e/))aaaaa. -a/_,,-- where soft nearest-neighbor density estimate f(z 23 (211) x Ei=1 Sua,..zs. (D2, 1 1/a is the bandwidth of the Laplace kernel. (With cross-validation, we determine a 1 to be good choice.) If we let wr(xx) c-a -211, the resulting induced model by marginalization is:
Q1)l))= / f(
Q(y|z,,) x2)
We (I)
These equations are respectively equivalent to Eqs. 44) and 66. The IMM risk and the corresponding overall objective remain the same:
QM
p(y(m,,) log Q((()))))
Cross-Entropy(Q) IMM(Q)-
Results In Figure D we compare the performance of IMM-trained Q (green) to that without IMM (maroon). We sweep a range of from to 50, and use cros-vvlidation optimized X for each (details in the Appendix I.I.I). A.I.). The key observations are that: (1) IMM always improves on the baseline performance, (2) the variance of the outcomes also typically diminished, (3) the improvement greater with less data, but the gain across data sizes is equivalent to access to an average of 30% extra data. This and similar experiments suggest that gains are highest when the dataset size is comparable to the number of parameters. This simple scenario demonstrates how IMM effectively harnesses the benefit of accurate feature-restricted models when training full-featured models. For reference, we also include the results of noising and interpolation. IMM is always better than noising, but interestingly interpolation is better with very few samples, though much worse with more. We attribute this to the fact that with less data it is harder to obtain an accurate induced model.
8

---

Table 1: Perplexity for LSTM Language Model using the Penn TreeBank dataset. The numbers on None and KN Naising are from Xie et al. (017)] and can be replicated using their original code (we use the model with latent dimension 1500). Like the baseline, for each row, we report the best value across as many restarts.
Improvement
Validation
Test
None (only regular dropout)
81.6
77.5
KN Noising (reproducible)
76.7
73.9
IMM with KN Bigram
76.0
73.3
Table 2: Results on the BERT RASE Language Model. The baseline numbers can be replicated using the original BERT code by Google, as well our provided repository. Matthew's 's Correlation Coefficient used for CoLA, F1 score for MRPC and Accuracy for QNLI and RTE. Like the baseline, reported numbers are averages across multiple restarts.
BERT
BASE
MLLM
+IMM
CoLAA
52.1 4.0
55.0 3.0
60.0 + 1.0
MRPC
88.9 2.0
89 1.0
90 = 1.0
QNLI
90.5
91.0
9 1.0
RTE
66 3.0
2.0
71 = 1.0
7.2 Language Modeling Experiments
In these language modeling experiments, the restricted model we use is the modified Kneser-Ney bigram (Kneser and Ney 1995 Chen and Goodman [999) of the dataset in question. To see why this is a good choice, we refer the reader to benchmarking done by Chelba et al. (2017); neural models (single layer and 2-layer LSTM RNNs) could not improve upon the perplexity of an interpolated Kneser-Ney N-gram of a similar order. After the introduction of the attention mechanism (Bahdanau
al. 2014 Luong et al. 2015). better neural models now dominate language modeling (Vaswani et al. 2017 Devlin et al. 2018 Ture et al. (019). We investigate how IMM could potentially improve even these more modern models, by using BERT's performance on some GLUE benchmarks as proof of concept. (Appendix D gives evidence that IMM indeed improves the full model's s performance on restricted task.)
LSTM RNN Experiments We build on top of the code provided by Xie al. (017 using the same topology for the LSTM RNN. The chief purpose of these experiments is to contrast directly with noising introduced in that paper. The PTB dataset document dataset and the LLM is solving next word prediction task. The average cross-entropy of multiple unroll positions of the RNN, or exponentiated as perplexity, is used as the measure performance. For training and measuring validation perplexity, L = 35 unroll positions are used. During testing, only unroll position is used. The IMM component always calculated by running the LSTM an evaluation mode (i.e., without any dropout). In addition, while regular evaluation updates the state in stateful models like the LSTM, the IMM branch never changes the state and uses the state set by the primary branch when traversing the dataset. In [ we report perplexity values using k-sampled IMM with k = 10. The table also includes the best noising results that we could reproduce based on the code of Xie et al. (2017) after communication with the authors (these are 0.5 more than the paper's numbers).
BERT Experiments We introduce the IMM objective in BERT's fine-tuning phase by reintroducing the Masked Language Model (MLM) objective that is originally only present during pre-training. In Google's 's original BERT code, MIM was present during pre-training but removed from fine-tuning, possibly because of minimal gain. For us, however, MLM is idcally suited to be augmented with the similarly generate induced bigram
IMM objective because it is based on cross-entropy and for predicting masked words in a sequence of tokens. We report numbers on these datasets in Table 2 The second column shows the numbers after adding back the MLM objective, which doesn't 'ttt produce much gain on its own. The third column adds IMM within MLM, significantly boosting the gains.
Since some GLUE (Wang et al. 2018) datasets (used in the original BERT paper) are too large to be trained an academic setting, we use a subset of the GLUE tasks (Warstadt et al., 2018 Dolan and Brocketi 2005 Rajpurkar et al. 2016 Dagan et al., 2005) to demonstrate the gains using IMM. For diversity, our selected GLUE tasks are mixture of single-sentence and double-sentence tasks.
Further experimental details provided in Appendix D.2

---

7.3 Reinforcement Learning: POMDPs helping MDPs
IMM is particularly appealing in situations where incomplete- data may be much more available, due to reasons like law, privacy, or limited sensing. If an autonomous driving system developed with few sensors, and later a new system taking advantage of more sensors is to be designed, the older system may act as restricted model helping the new design. If a diagnostic system is built based on a limited set of genetic markers, and later more markers are determined relevant, then the legacy system can be used without referring to the legacy data, which may need to be kept private. If platform's 's recommendation and ad engine is trained from ample general-public data, and later a personalized engine is to be developed, then the older engine can inform the personalized one through IMM.
In stateful environments, problems like the latter often require a reinforeement-leariine (RL) solution. the personalized engine has fell-featured data, can use the data to train an MDP (Markov Decision Process) policy to optimize expected reward (Sutton and Barto 2018). In contrast, the general-public data, despite being abundant, may lack in features and may only allow for solving a POMDP (Partially Observable Markov Decision Process). We can show that IMM can allow a good POMDP solution to significantly improve MDP training, by modifying we illustrate the reward achieved with and without the use of IMM, for learning policies for an agent on a toroidal 11 x 11 grid, with reward peaking at its center. The POMDP only observes one coordinate, whereas the MDP observes both.
methods such as REINFORCE
(Williams) 1992). In Figure 2
20
10-
BMM maximal lity action
100
200
300
400
500
Number epachs
Figure 2: Average reward of MDP trained without and with IMM incorporating POMDP. Details in Appendix
Conclusion
In this paper, we addressed the question of how to incorporate accurate restricted models in the training of full-featured models using the principle of induced model matching (IMM). This was inspired by interpreting some noising-basee data augmentation techniques in natural language noising, as an attempt to harness the knowledge of the restricted model used for noising. We showed that naive noising is not the best way to incorporate this knowledge, it may fail to consistently recover the true model. IMM, on the other hand, directly aligns the learned model with the target model and consistent. The results shows that IMM always outperforms noising, and improvements even decay gracefully with lower restricted model quality (see Appendix D.4) One limitation of our approach is that computing the induced model exactly is not always viable. To remedy this, we proposed sampled IMM, which yields accurate but somewhat computationally demanding learning, and serialized IMM, which slightly less accurate but has potential to be as efficient as the IM-IMM baseline. We then experimentally demonstrated the gains that IMM can offer, in a logistic regression toy example, when training LSTM language models and fine-tuning pretrained transformers, and in a simple reinforcement learning scenario. We believe that scaling from feature-restrictte to full-featured models is an important yet under-studied sub-problem of knowledge transfer. In addition to the proof-of-concept experiments in this paper, many others of particular contemporary relevance may be devised. For example, lengthening the context of large language models remains an open problem; we believe that IMM can be part of betters solutions, by correctly informing new longer-context I.LMs using current shorter-context L.Mss. The principle behind IMM is applicable very generally and we hope this work gives impetus to such explorations.
10

---

Acknowledgments and Disclosure of Funding
This paper is based upon work supported in part by the National Science Foundation, through the NSF CAREER Program under Award No. CCF-2146333 (From Rare Events to Competitive Learning Algorithms), the NSF HDR TRIPODS Phase II Program under Award No. ECCS- SSCS-2217023 (IDEAL Institute), and the NSF TRIPODS Phase I Program under Award No. CCF-1934915 (UIC Foundations of Data Science Institute). Computational infrastructure was supported part by the NSF MRI Program Award No. CNS-1828265 (COMPaaS DLV).
References
Bahdanau, D., Cho, K., and Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate. arXiv preprint 1409.044
Chelba, C., Norouzi, M., and Bengio, S. (2017). N-gram language modeling using recurrent neural network estimation. arXiv preprint arXiv:1703.10724.
Chen, S. F. and Goodman, J. (1999). An empirical study of smoothing techniques for language modeling. Computer Speech & Language, 13(4):359-394
Dagan, I., Glickman, O., and Magnini, B. (2005). The pascal recognising textual entailment challenge. In Machine Learning Challenges Workshop. pages 177-190. Springer.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 10X.vv11100048000
Dolan, W. and Brockeu, C. (2005). Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (aaessinnng((/PP/ (/WP2005).
Egorov, M., Sunberg, Z. N., Balaban, E., Wheeler, T. A., Gupta, J. K., and Kochenderfer, M. J. (2017). POMDPs.jl: framework for sequential decision making under uncertainty. Journal of Machine Learning Research, 18(26):1-5.
Gao, F., Zhu, J., Wu, L., Xia, Y., Qin, T., Cheng, X., Zhou, W., and Liu, T.-Y. (2019). Soft contextual data augmentation for neural machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5539-5544
Hinton, G., Vinyals, O., Dean, J., et al. (2015). Distilling the knowledge in neural network. arXiv preprint I0I.::150000231... 2(7).
Kaplun, G., Malach, F., Nakkiran, P., and Shalev-Shwartz, S. (2022). Knowledge distillation: Bad models can be good role models. Advances Neural Information Processing Systems, 35:28683-2866
Kneser, R. and Ney, H. (1995). Improved backing-off for m-gram language modeling. In 1995 International Conference on Acoustics, Speech, and Signal Processing, volume 1, pages 181-184. IEEE.
Kochenderfer, M. J., Wheeler, T. A., and Wray, K. H. (2022). Algorithms for decision making. MIT press.
Lee, H., Hou, R., Kim, J., Liang, D., Hwang, S. J., and Min, A. (2023). A study on knowl- edge distillation from weak teacher for scaling up pre-trained language models. arXiv preprint arXiv:2305.188 18239.
Li, H., Cai, D., Xu, J., and Watanabe, T. (2022). Residual learning of neural text generation with n-gram language model. arXiv preprint arXiv:2210.1443
Liu, J., Min, S., Zettlemoyer, L., Choi, Y., and Hajishirzi, H. (2024). Infini-gram: Scaling unbounded n-gram language models trillion tokens, arXiv preprint 27X31:240..1377.
Luong, M.-T., Pham, H., and Manning, C. (2015). Effective approaches to attention based neural machine translation. arXiv preprint 150.:0400000.

---

Ng, N., Cho, K., and Ghassemi, M. (2020). SSMBA: Self Self-suue manifold based data augmenta- tion for improving out-of-domain robustness. arXiv preprint arXiv:2009.10195.
Pascanu, R., Mikolov, T., and Bengio, (2013). On the difficulty training recurrent neural networks. In International conference on machine learning, pages 1310-1318. PMLR.
Qin, Y., Lin, Y., Yi, J., Zhang, J., Han, X., Zhang, Z.., Su, Y., Liu, Z., Li, P., Sun, M., al. (2021). Knowledge inheritance for pre-trained language models. arXiv preprint arXiv:2 138.80
Rajpurkar, P., Zhang, J., Lopyrev, K., and Liang. P. (2016). Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: 05250
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. (2014). Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1955)
Sutton, R. and Barto, A. G. (2018). Reinforcement learning: An introduction. MIT press.
Turc, I., Chang, M.-W., Lee, K., and Toutanova, K. (2019). Well-read students learn better: The impact of student initialization knowledge distillation. arXiv preprint arXiv: 08962, 13.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.., Gomez, A. N., Kaiser, t.., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
Wang. A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. (2018). Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: vrXiv:1804.000 07461.
Wang, S. I., Wang, M., Wagner, S., Liang, P., and Manning, C. D. (2013). Feature noising for log- log-linear structured prediction. In Proceedings of the 2013 conference empirical methods in natural language processing, pages 1170-1179.
Warstadt, A., Singh, A., and Bowman, S. R. (2018). Neural network acceptability judgments. arXiv preprint arXiv: 1805. 12471.
Wei, J. and Zou, K. (2019). Eda: Easy data augmentation techniques for boosting performance on text classification tasks. arXiv preprint arXiv:
Williams, R. J. (1992). Simple statistical followng-ooloowng algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-25 256.
Xie, Z., Wang. S. I., Li, J., Lévy, D., Nie, A.. Jurafsky, D., and Ng. A. Y. (2017). Data noising as smoothing in neural network language models. arXiv preprint IrXiv:1703.0257:
Yuan, L., Tay, F. E., Li, G., Wang, T,, and Feng, (2020). Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEF/CVF conference on computer vision and paltern recognition, pages 3903-3911.
12

---

A
Explanation of main notations
Figure gives a schematic overview the process of IMM.
full-featured learned predictive model
full-featured true
predictive mode!
P(y
Pa (y)
Cross-Entropy
Pizatinaaiization
Th empirical context distribution
distribution
Pyy i true feature- restricted induced model
P(y T
Induced Model Matching
(3)
learned restrree-estrictted Induced model
Figure 3: Schematic overview of IMM.
Table 3: Glossary
Prediction Variables
y (generic) y (at datapoint & of
Context Variables
(generic) (at datapoint & of m))
denoles short context and
g denotes extended context
Thus, æ (I)
extend (z) multiset of co-occurring with F
Prediction based on Full Context
((ylzr (using true model) Quy[=) (using learned model)
In the case of BERT experiments, Q(v(=) denotes prediction vector for certain masked position.
Dataset Size
is the dataset size The subscript may be used denote empirical distributions
Context Distribution
(truc context distribution)
R, (empirical context distribution)
Learned Model
alone, without parenthesis denotes the learned context-conditional model (which includes its parameters), c.g., LSTM Transformer
Prediction based on Restricted Context
P() using true induced model P(y[T), using target induced model, a proxy to P. This is
the Kneser-Ney bigram our experiments.
Q(y[T), using learned induced model, i.e., based on true Q(y[=), using empirical learned induced model, iee. based empirical T.. proxy to Q
Table 3 compiles our main notation. We summarize the main formulation here and explain how connects instances the experiments.
Q always the full model (LSTM or Transformer). Q (xy) is thus a learned context- conditional model: takes the full context (Ix) and outputs prediction probabilities (over y.). In Eq. (4), log(Q) the log-loss, whose average the cross-entropy.
LSSMMPPBB: context = previous words, prediction next word, and cross-entropy main (nrxt-word prediction) loss.
13

---

BERT/GLUE: context = unmasked words, prediction masked word, and cross- entropy -MLM (masked word prediction) portion of the fine-tuning loss.
Note: KL divergence and cross-entropy are interchangeabl,, as their difference doesn't t depend on Q.
True context distribution я and true context-conditional model P describe the unknown probability space. We assume tokenized data (#t, Mt) generated according to (xx)) and (()
- -The goal of learning Q is to approximate P.
Empirical context distribution nn and empirical context-conditional model Pn are histograms of the context and context-conditional prediction, per training counts.
-Using these instead of the true is equivalent to replacing expectations (over the data distribution) with sums (over the training set), c.g., Eq. (5) + Eq. (
Induced model: specializes full context-conditional model (ye(24)) short context only (T)), under a specific context distribution.
If the full context-conditional is and the context distribution л, we get the true induced model P. This is the ideal, or Bayes' optimal small model (which is available practice). In practice, we approximate by a target induced model. P (e.g., Kneser-Ney bigram).
If the full context-conditional is Q, and the context distribution is...
get the learned induced model Q. This is the best way can prediet based only on the short context. However, we cannot evaluate it without и. Instead, we use...
...
... Tn. and get the empirical learned induced model Q. This, evaluate and use in our training.
In summary, the experiments use Q given by Eq. ( efficiently approximated via sampling in Eq. 18) This is then plugged into the empirical IMM objective in Eq. ( which pits Q against P.
B
Caveats of prior methods: noising and reverse-KD
B.1 A quick review and analysis of noising-based data augmentation
In Xie al. (2017), noising was proposed an approach to perform data augmentation in language and sequence-to-sequeee
models. Two types of noising were suggested: context noising and target noising. We first show that context noising is not justifiable except in select cases. To simplify, consider only a bigram model, with counts c(x, y). where & is the context and y is the prediction. The general noising scheme in Xie et al. (211)) would, with some probability ((x) that may depend on the context, substitute either the context, the prediction, or both with a sample from a distribution q()). It is straightforward show that, result, the expected counts presented the learner longer c(æ, y), but rather change to ö(x, 3) as follows.
(a) If only contexts are noised, then ö(I, 3)
[1 (=))(((( + v(x:) En(r))((rs)))
(b) If only the predictions are noised, then ह(τ; y)
- (=>(2)((22,))))) v) (vv)(2)c(zz)
(c) Lastly if both predictions and targets are noised, independently, then a(x, y)
(=)]c((..)) (())((1) arlg
Σ (xx) (a)))((a'
These noising schemes are primarily motivated (see Xie et al. (017). Section 3.3) through the idea that noising leads to classical forms of smoothing, and which in turn may be a desirable property
14

---

to inject into training more sophisticated language models. This is indeed true in case (a) when ((22) = À constant and g()) = c(-//n is the unigram distribution, leading simple interpolative smoothing. It immediately fails to be true when g is any other distribution even if 7(x') X, as one needs to normalize with c(x) to get a conditional distribution. The failure of this interpretation is even more apparent in case (a), when (xx') and g(-) are determined via the missing mass and Kneser-Ney backoff, failing to recreate any likeness of the Kneser-Ney smoothed bigram. The situation at least as bad in case (c).
The only instance that succeeds uphold "noising as smoothing" is case (b), which recreates both interpolative smoothing in the unigram case, and gives an output related to the Kneser-Ney bigram with the choices of yya') and q(-) from the paper. This last case is of particular interest to us, namely because even though it may appear that we recreate the Kneser-Ney bigram itself, the choice of (xx') = No -))(1)))((() with To = 0.2 or 0.6 (see Xie et al. (2017). Table and Figures and 2) makes it evident that this is under-weighed to represent the missing mass, which typically corresponds to larger discounts (0.75 or higher), due to the heavy Zipfian nature of language. What can we deduce from this? If v(x) is the true missing mass in the Kneser-Ney model, then we can understand this choice as 7(x) = (r). As a result, we have:
((x.y) (1 - 1))(),(,))) v) + +
v()))c(E,)) =(x)g())((=))
(16)
Upon close examination, we identify this as an interpolation between the data on the left and the Kneser-Ney bigram the right, which suggests a form that is similar to IMM, in that the typical calculation of the log-loss is additionally joined by the log-loss pitted against a target model, the Kneser-Ney bigram.
B.2 Proof of Proposition 6.1
Assume that we are optimizing the idealized noising objective of Eq. 14 i.e., we are operating in the infinite-dai regime- and let Q* be its global minimizer. Assume further that the model class for Q contains the true model P p i.e., we are in the realizable case. Then, there exists choice of and P such that Q' +
Proof. Let us first rewrite 14 below, with the addition of constants (an entropy to the first term and P(y2z) in the logarithm second term) that do not change the minimizer of this objective:
P
D(PIQ)+ +A (x) Σ P(ulz) log S(Q)
Q
s(Q)
To simplify the notation, let f refer to main objective term and g refer to the noising regularization, which we have equivalently identified with the single-sample approximation of IMM.
We note that both ff and are convex in Q and f 2 0. Since f(Q) is minimized at P. its gradient should be 0 at P. g(Q) is 0 at P.
If Ag((t) is -c < 0 at some Q", then along the line connecting P to Qt, Ag(Q) is below the line 0 -e. To compensate, would need to be above the line 0 +e, which would violate the fact that the gradient of f is at P.
We now numerically show that g(Q) can indeed be negative at some Q1, for given construction with specific choices of π and P.
Consider a trigram scenario, with r = (x-1, 2-) and where we understand the short context as == __ and the extended context as I = x-2. Let us rewrite the regularization term g(Q) explicitly splitting the short and long contexts:
ply
g(Q)
Σ 1-1---
T(
Σ P(y x1 log
(17)
Q(y 11, 11, 2-
15

---

We search for distributions over the prediction y, short context ž, and extended context z. Consider the following resulting tensor:
0.396 0.1
0.003 0.05
Ply,I-1,--1
0.004 0.1
0.297
0.05
Based on this, we get the context distribution ((,, 2) the conditional P(y I-11 -2), as well as the restricted model P
0.4
0.3
T T(--1,--2)
0.2
0.1
0..99
0.01
0.5
0.5
Ply x-1,x-2) =
T0.01
0.99
0.5
0.5
These give the following induced model:
0.57
0.5
Ply x-1)
0.43
0.5
Now consider the following choice of Q
0.5
0.5
0.5
0.5
Q'(y x=1,
0.55
0.5
0.5
0.5
Plugging the above choice of n, P, P(y x-1, and into Eq. (77) gives negative value of g(QT)= 1.1, thus completing the construction of the counterexample and the proof.
Table 4: Deterioration of noising vs. consistency of IMM when À 1.5 is fixed in the logistic regression example.
This is concrete manifestation of Propositon 6.1
Dataset Size
Baseline
Noising
IMM
IMM-Noising Gap
N=5
13.14416.
76.85 +18.58/11113
19.99+14.99/112.33.
3.11
N=10
84.17-13.511-1..66
84.37 -7.4077773
88.65 49.65/77788
4.28
N=15
89,86 68.86/6.80
86.17 16.17/55.53
92.52 +6.19/-4.81
6.35
N=20
92.35 47.35/55.32
86,99 +5.69/-5.37
94.30 +4.63/-3.70
7.30
N=30
94.94
88.68 +4.35/-4.32
95.69 43.39/-2.64
7.01
N=40
96.33 43.00/2..33
89.70 43.73/-3.97
96.59 42.59/-2.08
6.89
N=50
22.47-22002
90.78 +3.78/-3.55
97.34 42.34/-1.99
6.56
Experimental Evidence of Inconsistency One may dismiss Proposition 6.1 as being too specific of counterexample. After all, in Figure the gap between all methods does vanish. It is important therefore to emphasize that this only happens because we are being very favorable to noising. Specifically, we are decaying (the amount of noising) optimally with increasing data. This is necessary for noising, for the precise reason of Proposition 6. even if the target model is perfect, because noising incorrectly tracks the target model, without decaying its influence it will not only not narrow the gap, but would in fact derail the learned model. Decaying À is also acknowledged as
16

---

critical in the reverse knowledge-distillation distillation literature (see for example Sec. of of et al. (2021)) However, tuning À is optional for IMM, thanks to Q (with more extended context samples) accurately tracking the target (see the flat curves on the right column of Figure Appendix (DII.]
To experimentally verify this, we ran the logistic regression example with fixed À = 1.5 (optimal at data size 5). The results below. IMM maintains performance comparable to Figure T whereas noising experiences widening gap, and soon underperforms even the baseline.
B.3 Knowledge Distillation Theory
We would like mention that recent theory elucidates how weak teachers can be useful in reverse KD (Kaplun et al. 2022). That work is specific to the classification setting and assumes that the teachers are good samplers, i.e., noisy versions of the Bayes decision, and shows that reverse-KD can effectively ensemble multiple teachers and remove the noise. This setting and assumptions do directly apply here, however, a common insight may be that the process of inducing a model can also be thought of ensembling. but over contexts rather than considering that the implicit noising objective that we identify in Eq. (17) equivalent to the explicit objective of reverse knowledge distillation, the arguments in this section equally apply caveats for reverse-KD when the weak teacher is a restricted-contex model.
independent resampling. Otherwise,
C Implementation Details
Algorithm Sampled IMM with SGD for a Model Q with parameters W. To connect with the notation of the paper, this version most closely resembles the language modeling version (with the details of the sequentialization of log-sum gradients skipped for simplicity). The logistic regression equivalent uses soft nearest neighbor density estimate (ref instead of creating a multiset and sampling from it.
Input: Data (ze; ye) for / = 1,2, h: sampling rate
Output: IMM-trained model Q
1: repeat
2: Q(ylx;)) < FREDFORWARD Iz)
3: Vw
< BACKPROPAGATE(C. Cross- (ooos-EEntrpyyQ..00
4:
5:
for all i smple Enumeratee/Smpleeeettnddf33 do
6:
(w((T.) Q(y(T.)) FEEDFFORAARRD(.
7:
end for
8:
QU)(T,) QuylT.)(k
9:
IMM,(Q) = -- P(yla,) log Q(vlz:)
10:
W (wIMM(() + BACKPROPAGATE(Q, (Q))
11:
APPLYGRADIER
Cross Entropyy)))))((()) + AV (MMM(
12:
until convergence
C.1
k-approximated IMM
For the approximation, given short context take only fixed number & of samples Xi, Xx. uniformly from extend(T,)) The approximated Q (Eq. (II can then be written as
Q Qu sample extend(z,))
(18)
Indeed this what happening on lines 5-8 of Algorithm Sampling l extended contexts is only part of the solution. Another algorithmic innovation that we need is to address the task of computing the derivative of our main objective, Eq. ( because requires differentiating Eq. (7). where
naïve implementation would need all & instances of the L.M used to estimate Q(y(a)) in the memory simultaneously. Fortunately, provide solution this, explain in detail in Section C below.
17

---

Algorithm 2 Serialized IMM with SGD for Model Q with parameters W. As in Algorithm | this is notationally closest the language modeling version. For implementation details for the logistic regression example, see D.5 as well as D.I (for the soft nearest- density estimate formulation).
Input:
Data (Iee: yt for = 1, ... n, refresh frequency IMM-trained model Q repeat
Output:
2:
Q(ylx;) + FEEDFORWARD(
3: Vw-Cross- (oooss-Entooy(()) f BACKPROPAGATE((, Cross- Ennooss-ntrooy((1.1 yr))
4: if Repeeats % r 0 then
5:
Q Q Q
6:
Q 0
7:
for all & = 1.2,.....
8:
for all do
9:
(+(x)T)) - (x(T))) + FEEDFORWARD(Q).22)
(Or use other induction mechanism, e.g., density estimator for logistic regression.) 10: end for
end for
12:
NOGRADIENT(QR, Q+)
13:
end if
(MM((Q) - Ey (+(va=e) P(y)T,) log Q (r)) (u(T)))
14:
(This is nor the IMM loss, but its gradient averages to the correct gradient, see Section 5
(I) + BACKPROPAGATE((. (Q))
16:
CrOpEEENNTTTTCCOSO-EtSOOtte until convergence
APPLYGRADIENTS( V Cross- (ntropy))))))) + AV (Q).
C.2 IMM Gradient Computation
Eq. is a cross-entropy and the Q(ylz) term that we approximated in Eq. (18) occurs inside the log term of the cross-entropy. Naïvely backpropagating through this term makes the memory complexity of backpropagation scale with the number of random samples k used approximate Q(ylz). In our experiments, was problematic for the GPU (Nvidia V100 with 32 GB memory) to perform backpropagation for k > 6 for the LSTM RNN.
Contemporary deep learning frameworks do not have the ability to sequentialize along & the compu- tation of the derivative of log where the argument of the log is average (or sum) of k entities. This is expected, because log(A) A + B) cannot be decomposed, and therefore, sequentializing this by splitting the cost function (this cross-entropy term) is not possible.
Despite this, we propose solution that may be of interest generally when such differentiation needs to be performed. This stems from the simple observation that the derivative of log(ʃ + g) can be represented as a weighted average the derivatives of log(S) and log(g), where the "crosstalk" weights do not require differentiation:
V + Vg
V log(f +
[ V log + 7tg
v loggg.
f
Jtg
In summary, this allows us to compute the derivatives separately for each random sample and accumulate them appropriately.
Sequentializing IMM gradient computations using crosstalk
Recall that we can write an IMM component at data point (Eq. extended contexts as:
12
with & randomly sampled
18

---

IMM,(Q) Cross- Cnoss P(yl=))
Σ P(y|a) log
Ex (V Te, X) Fi]]
k
23
Σ P(yl=)ln log
Quylt,, ~ extend(T+))
(19)
h
1=1
Σ P(y|z) log
Σ k
where * comes from Eq. 18
The above entity cannot be directly decomposed into terms because it involves the log of a sum. Since it's the IMM component of the loss, during backpropagation, we will need its derivative with respect to the parameters of the neural model (LSTM RNN or BERT). In the next section, we see that the derivative can be decomposed into terms.
Denoting the set of parameters using W, we can then write Vw (M/ as below.
VwIMM.((Q) = Σ
P(y
Vw log
Q [.,
K
V
F Σ i=1 Q..
P(y|T
1o -i=1 Q.y)
- Vw V (3)
Σ P(y
# Σ Q(u F2, I.)
Vw Qw I.)
Σ P(y|
D F,,
Σ P(ylz)
V Vi ((v Q.I
1
(20)
QuI
VwQ(ulz,,
Σ P(y
Q(y(Te,I.))
Qu
P(y|
Vw log Q(y/ E.)
E Q(y(T
Q(vlT..z.)
P(U
Vw log Q(y(=,,
Q(1
Vw
P(y())) Crd(y) log (,/ R.)
crosstalk
Q(y We notationally move the derivative outside, but it is crucial that C..(v) Σt Q(y =., be treated
as constant with respect to the parameters of the LLM. We do this to show that we can effectively compute individual cross-entropies, with the Cr,(y) terms acting like a "crosstalk" between the random samples. This exactly parallels the decomposition that we highlighted above for the case of log(ʃ + g). Note that this decomposition not possible directly on the cross-entropy (that involves the log of sum), but possible when we consider its gradient.
19

---

Overcoming limitations of contemporary frameworks to compute the crosstalk vector
We cannot compute the above total gradient, i.e. Vw IMM(Q) unless for every and y we have Ct. (v) for E [1, kt.. To this end, we implement this by partially running the feedforward computation graphs of all random samples only until the point which gives us Q(y [x,)) x;)) for [1, k). Once we have all the 2;)) we can normalize all of them by their element-wise sum to get the weights Cy) These weights are then fed back into the graph, and the forward pass completed to get a "tampered" cross-entropy for E [1, M.. The k backward passes using these tampered cross-entropies
give us the terms inside the sum in the last step of (20). which after summation produce V. IMM)
D Additional Experimental Details
We now provide some additional detail of both the logistic regression and language modeling experiments. One minor difference between the two that in the logistic regression case we interpolate the main loss (1 - A) and the IMM risk (A), to efficiently cross-validate for the choice of A as we explain in Section D.I.I In the language modeling experiments, we simply add the regularizer with the factor X, while maintaining the main loss as is (factor 1).
D.1 Logistic Regression Toy Example
We provide the complete details of the logistic regression toy example Consider dataset with feature space z = (21, 22, x3) E I3 containing two linearly separable classes, defined by a linear discriminant g(z) = az + bxz + cr3 + We will think of features decomposable into I, the short context of restricted features, and g. the extended context of remaining features. For this example, let 2z). The choice of word "context" stems from features in language modeling.
===a 21 and æ
Target restricted model Say features are sampled uniformly over a cube and that we have ample observation of the labels y = [9(x)) > 0} along only the short context, i.e., ample data points of the form (x1. This allows us to build an excellent restricted model to predict the label based on just
T1, call it Plyla (x)) or P(y(z). For the sake of this example, let us assume that this is the (restricted) Bayes predictor or true conditional probability.
Main question We then observe few samples of all features along with labels, i.e., data points of the form (I1, I2, 23, How can we use the restricted model P(ylar) as part of the training of the full model Q(y[z1 22.
Induced model Say have data points. Let us index them by t 1, .n, written either as (It, 31) or (21,,, T3,t; Mt). The key concept of the method is that it is not Q's behavior that should target that of P(ylz1); but rather the behavior of Q if it were itself restricted. We call this restricted version Q the induced model of Q and, by marginalization, we interpret it as the average of Q's predictions, when 2 drawn from its conditional distribution given Tt. Since we typically don't have access to this distribution, we approximate it empirically. In language modeling, we could just sample from the empirical distribution of g for a given E. In logistic regression, this is not viable since Ij does not repeat. We instead use a soft nearest-neighbor density estimate f(12, 23 (1) kernel. (With cross cross-validatie validation, we determine a = 1 to be a good choice in this example.) If we let w. = -alzie mx,,, the resulting induced model by marginalization is:
Σm= 822, 12....25,4 (22, ( -o|_..._-- where 1/a is the bandwidth of the
/ ((1)))((()
Qlylz)
(21)
7
wx(x) 1= we (x) (
(22)
Just like in language modeling, to induce model we need to draw z's from its conditional distribution given F. Unlike in the discrete case where we could directly approximate the conditional distribution using the empirical counts, need rely here on an estimated density.
"Code is available altng:///ghiabbboooomiideeemmmme
20

---

Matching the target The main objective of logistic regression is to minimize 22.1: Di=1 log 2 which is equivalent to Cross-Entropy(QQ) versus the em- pirical distribution. We now additionally want the induced model Q to match the target restricted model P. This gives this method the name Induced Model Matching (IMM). This requirement can be captured through KL-divergence, and introduced as regularizer. The result is secondary objective, which we call the IMM risk, expressed as:
IMM(Q)
Σ Σ 1= y=0,1
P(ylæ1,)) log Q1ya111)
(23)
The overall objective then becomes
Pross-Entropy(Q) + IMM(Q).
(24)
where À is the regularization trade-off and can be determined via cross-validation (details in the Appendix D.I.I).
IMM improves restricted-conter prediction In Figure 4 for 10 and 30 Monte-Carlo runs, we show how the induced model of Q (i.e. Q) fairs its prediction, compared the target model P. Note that without IMM (i.e. with À 0), the secondary objective is large and thus Q's performance worse. With increasing X, IMM improves this performance both in average and in variance. We deduce that there information in the accurate restricted model P(y(T) that not known Q, unless explicitly incorporate it. This shows that Q gets better at predicting from restricted context, naturally bringing up the question: does also get better at predicting from the full context, i.e., the main objective?
IMM (Q)
0.4
11.6
0.8
Figure 4: Performance restricted task. i.e. IMM(Q) measured models Q trained using Eq. 24 the X objective with varying 14x ratio (refer to Appendix D.1.1). We stop at I+X 0.9 because 14x would zero the contribution the main objective (and replacing the main objective completely is never the intention).
IMM improves full-context prediction In Figure I we compare the performance of IMM-trained Q (green) that without IMM (maroon). We sweep range of n from 2 50, and use cross- validation optimized X for each (details in the Appendix D.I.T). The key observations are that: (1) IMM always improves on the baseline performance, (2) the variance of the outcomes is also typically diminished, (3) the improvement is greater with less data, but the gain across data sizes is equivalent access to an average of 30% extra data. This and similar experiments suggest that gains are highest when the dataset size is comparable to the number of parameters. This simple scenario demonstrates how IMM effectively harnesses the benefit of accurate feature-restrieted models when training full-featured models.
21

---

Visualizing the effect of IMM In Figure 5 we illustrate the 3-dimensional logistic regression problem. The features are samples uniformly in this box. The Bayes-optimal restricted model only uses the -ooordndne,,,, and so assigns probabilities proportionally to the blue/red areas in the illustrated slice. IMM then encourages the full logistic model be consistent with these weights, i.e., making sure the proportion of points labeled + agrees with these weights at each T1. Intuitively, this biases the separating plane to have the right inclination/aligmment with the 2 -axxx,, which subsequently speeds up the learning process.
X3 0
2
1X2
1
X1
Figure 5: A visualization of the inductive bias brought upon by IMM in the logistic regression example.
D.1.1 Tuning Lambda
As mentioned Section 7.1 IMM always improves baseline performance. The main parameters to performed few experiments held-out data and determined the
choose and X. For value of to generally adequate. For X, we saw that adapting to the dataset size made a more significant difference. This makes sense, for smaller datasets, we want more IMM contribution to compensate for the lack of data (i.e., we want larger A).
To determine a good schedule for À function of dataset size n, we sweep the ratio I+X (which represents the IMM coefficient
fraction of the combined IMM and regular loss coefficients) in the range (0, 1). We repeat this for learning rates 0.1 as well as 1.0, for a range of dataset sizes, from minimum of 2 to maximum of 50.
We provide our tuning plots in Figure 6 Since a learning rate of 1.0 was found to give the best results, we choose it for all the reported experiments. Looking at plots in Figure [ we select X/(1+ X) of 0.8, 0.7, 0.6 for dataset sizes, n of 2, 10 and respectively. These suggest a linearly decaying optimal A/(1 + X), in this range of dataset sizes. For this example, doing interpolation fit, create our automatic À schedule rule to be I
00.0111n 0.818.
22

---

Dataset Size: 2, Learning Rate: 0.1
Dataset Size: 2, Learning Rate: 1.0
75-
I
I
70
70-
Accuracy
65-
Accuracy
63
60-
53-
53-
50
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
Dataset Size: 10, Learning Rate: 0.1
Dataset Size: 10, Learning Rate: 1.0
95
95
I
I
90
90
55
T
0.2
0.6
0.8
0.2
0.4
0.8
1+A
1+1
Dataset Size: Learning Rate: 0.1
Dataset Size: Learning Rate: 1.0
08
96
95
94-
94-
90
92:
90-
96-
94-
82
0.2
0.8
0.8
14A
142
Figure 6: Effect of varying (1 + X) different dataset sizes and learning of 0.1 and 1.0. Every point an average of 30 runs and the error bars are computed using 10 and 90 percentiles.
23

---

Table 5: Restricted vs. Full Model the Bigram Task (i.e.. predicting the word given only the previous one.) The small model is the Kneser-Ney bigram on PTB training set. The full model is the 1,500-dimensional 2-layer LSTM from Xie et al. (207)) trained without noising on the PTB training set. To make the LSTM predict the next word based only on the previous one, we complete the history by averaging predictions over all extended contexts the training set (full word history) that share that short context (same previous word). In the language of the paper, this exactly the induced bigram of the learned LSTM. report the cross cross-entropy and perplexity of each model on the PTB test set. Note that the small model clearly outperforms the LSTM on this restricted task.
Dataset
LSTM
LSTM
Kneser-Ney
w/o IMM
w/ IMM
Train
260.16
237.55
92. 19
Validation
339.24
- 302.30
199.28
Test
309.56
+ 278.48
185.71
D.2
Language Modeling Experiments
Code is available at
D.2.1 Evidence that IMM acts through improving restricted task
We give an illustrative example that shows the advantage that a restricted model may have, and which we would like to incorporate into the full-featured model. Clearly, our approach is only valuable this the case. The simple experiment that propose is to compare how well full-featured model (LSTM) performs a restricted task, namely bigram prediction, vs. a restricted model built for that task (Kneser-Ney Bigram).
How can an LSTM perform bigram prediction when given only a single word? It can do so by averaging its output over all possible long histories consistent with that word, which is exactly the notion of induced model of Eq. 3 In Section 4 we show how to do this empirically.
When we do this comparison for the non-noised LSTM in (Xie et al., (017). we obtain the results tabulated in Table 5 The Kneser-Ney bigram (restricted model) outperforms the LSTM (full model) on the restricted task. This may appear surprising, however, had this not been the case, the improvements that Xie et al.] (2017) obtained through noising would have been hard to explain! More importantly, using IMM improves the performance of the LSTM on this task, mirroring the logistic regression case - cf. Figure 4
D.2.2 Parameter details
Choice of A For Language Modeling experiments, due to computational cost, we did not perform an exhaustive search for the best X. We tried a few values of À and tested on held-out data, and determined the value of 0.2 to work well in most cases, and adhered to it in all of the experiments.
Restarts We would like to clarify that the baseline for LSTM RNN (Xie et al., 2017 reports the
best perplexity across multiple restarts while the baseline for BERT (Devlin et al. average across multiple restarts. We do likewise our experiments. We have additionally included error bars for BERT experiments as that's where most of the gain is seen for BERT.
2018 reports the
Clipping On the technical front, we note that in either LSTM or BERT, gradient clipping (Pascanu et fal.2011) is used, then gradients should be clipped separately for the primary objective and the IMM objective (i.e. they should be clipped before they are added together). The reason that they are obtained from two loss functions that have very different scales.
Choice of To reduce overhead for LSTM, we perform IMM only every j'th batch using s-samples. Otherwise we use 1-sample IMM. For BERT, we perform k-sample IMM every minibatch. For the selection of the parameter (number of long histories used to approximate the induced model), we experimented with & E 51,5, 10, 20}. For the LSTM experiments, there was not much gain going from 10 to 20, so we settled for k = 10 (and j = 5) as a good compromise between the accuracy and time complexity of the induced model computation. For the BERT experiments, we used k 5.
24

---

IMM through reintroduced MLM loss in BERT fine-tuning BERT on GLUE differs from LSTM on PTB in two ways: (1) it is not a next-word predictive model, and (2) its main loss during fine-tuning is the task (e.g., CLS) loss. We do not touch the latter. Instead, we reintroduce the ML. loss, to have a fine-tuning predictive task close to our formulation. Since MLM is based on cross-entropy, IMM can be directly added to it as a secondary loss. Additionally, only for the computation of the IMM component, at each location a fine-tuning training sentence, we mask that word and all future words during the forward passes being done to compute the IMM component. This forces BERT to act like a causal model, i.c., it predicts the next word based only on the past, as in our main formulation. This is done ensure that the induced model of BERT is the right induced model to match against an accurate causal model. Incidentally, for this accurate target model we used the Kneser-Ney bigram constructed from the Wikipedia data, on which BERT is initially trained.
In summary, in the BERT experiments we fine-tune with the sum of 3 losses: the task loss, the MLM cross-entropy loss, and the IMM risk for the MLM task.
D.3
RL Experiment
We consider an 11 x 11 toroidal grid in the (z, v) plane with predefined reward landscape that has a single peak at the center of the grid, see the heatmap Figure The toroidal configuration allows the agent to wrap around the grid in case it exceeds its confines, which allows us to rely on uniformity in the action space, in any state.
-14
0.8
0.6
0.4
0.2
0.0
10
Figure 7: Heat map representing the reward function, which depends on state only, on the 11 toroidal grid of the RL experiment.
We model this as both Markov Decision Process (MDP) as well as a Partially Observable Markov Decision Process (POMDP). The MDP agent knows both the x and the y coordinates of position, while the POMDP agent knows only x but maintains a uniform belief over the dimension. We first solve the POMDP using the POMDPs. (Egorov al. (017) package package and the Fast Informed Bound (FIB) Solver (Kochenderfer et al. (2022) to obtain very accurate policy based on partial information. (These POMDP policies represented in the form of alpha vectors.)
MDP Training without IMM Our MDP training algorithm is REINFORCE (Williams 1992, REINFORCE, at every epoch, samples set of observations, which can also be called a batch. In REINFORCE, at every epoch sample fresh data from the policy. We elect to limit the observation length at every epoch, to emulate limited exploration, and hence the number of epochs can be seen as the effective dataset size.
Adding IMM to REINFORCE For every observation, we consider its restricted context to be (the known dimension of the POMDP agent). Unlike language modeling experiments (where randomly sample for the extended context), the POMDP belief. Here, we thus use
25

---

uniform distribution over the unknown dimension y (the analog of the extended context). Using this distribution, we compute our policy's induced model. We then match this induced model against the POMDP agent's action, which is an action based solely on the knowledge of x: (and only belief over y). We manually set À to 0.25.
We set the per epoch observation length to 50. Since then the number of epochs represents our dataset size control, we vary it and compare the performance of REINFORCE and IMM-an -aggmeeeetedd REINFORCH by evaluating their average reward over a given rollout horizon. Similarly to the logistic regression experiments, we see that IMM is uniformly beneficial, with gains their best when the dataset size is roughly in the regime of the number of parameters. We perform 30 Monte Carlo runs at each number of epochs, and report 10th and 90th percentiles. Here also, IMM reduces the variance of the reward. Results reported in Figure 2
Experimental Results
D.4 Effect of Restricted Model Quality
While the paper is presented under the assumption that we have access to a very good target model P. we also ran experiments to show the effect of model quality on IMM. Of course, for the language model, we don't 't have gold reference and the fact that IMM achieves gains and improves on noising is evidence that even non-perfect models help. However, for the logistic regression and RL experiments we do have gold references, which we can weaken. We weaken the restricted Bayes-optimal predictor in the regression example by adding Bernoulli noise to the label (interpolate with a coin flip), 20% for the medium quality model and 50% for the low quality model. For the RL experiment, we replace the max on the POMDP utility with a softmax, and adjust the quality via the temperature of the softmax (small is higher quality). Even with reduced model quality, we find out that IMM always improves learning. However, for this to happen, we need to tune more carefully than when the model perfect, leading to less reliance on the IMM risk when the dataset size increases (otherwise the noise in the lower quality models would offset the convergence). For logistic regression, this is illustrated in Figure g For RL, this is illustrated in Figure 9
80
10
BVM
1754 model
Figure 8: Test accuracy of logistic regression model trained without IMM and with IMM using restricted model of varying quality levels, with X determined through dataset size (refer to Appendix D.L.Iy At every dataset size, we perform 300 runs and plot the 10th and 90th percentiles for error bars.
All experiments in this section use the Sampled IMM variant. In Appendix D.5 show that Serialized IMM achieves competitive statistical performance with
D.5 Serialized IMM Performance
To demonstrate the validity of the Serialized IMM variant, which comes with computational advan- tages that allow IMM scale, implemented it for the logistic regression example. However, we first elaborate on the details of the Sampled IMM, to show the dramatic computational advantage. By combining the density estimation approach for obtaining the induced model, outlined in Section 7.I with the sequentialization technique explained in Appendix C.2
we obtain the following gradient
Code is available htttts///gthbbboooo//ccccee//eee
26

---

15-
Avaroge Htmoodd
10
masimal Llt wction paicy using policy 10m
No MM
egachs
Figure 9: Average reward of MDP trained without and with IMM incorporating POMDP solutions of various "qualities" " Details in Appendix D
calculation for the IMM risk (compare to Eq. (20)):
(M) Σ P(y log ((yl=+)
V W IMM, = - Σ Ply V log (2)
n
Σ
P( log
Wh.
Quy Jee Ier)
6' Wi.e
22
Quyl Ley
Ply
p==1
Wayt
Weye!
EP=1
De We,e
D W4 Q Te, 241] 11
Σ Ply
='=1 1 ( Ir.
V | Q Tt, Ze))
P 1T
x=
tt Ep)
[wr.e Q T Zt Zr)
P [T)
21== W (1) Fr,
Wi Q T4, I.r)
[Q(y(Tn.z))
Ply TT
Dr= W7 Q Ere Ze)
Q(y]Tr.
Wi. Q1 (i, Imr)
P IT
V log Q 2T,
Ep=1 Wa Q TT, 2p)
F=
erestalk
IMM(M)
Σ P()
We.e
V log x[))
EV=1 We. (1 Ee)
double summation
crssstlk
Here Wt,t' uses the index of 21, rather than its value, as argument. By using a density estimator that looks at the entire dataset, we have sampling scheme that effectively has set k costing an n-fold computational time increase. While we can address this through shorteuts (e.g., sparsifying w). this extreme version of computing the induced model is revealing in terms of where the bottleneck is, i.c., computing Q
27

---

For the Serialized IMM variant, by comparing the above to Eq. (3 we see the large gain thanks to the double summation collapsing into single summation. We still need to compute Q for the correction factor, but we can do this every n iterations. By incurring the (n) slowdown only of the time, go back to only constant factor overhead. The additional noise in the gradient and the fact that Q is stale between updates mean that Serialized IMM cannot be expected to perform the same level as Sampled IMM. Furthermore, the Serialized IMM does not use the smooth nearest- nearest-neighbor averaging, except for the calculation of Q. Despite this, we see in Figure 10 Serialized IMM remains better than noising and, for larger dataset size, starts to approach Sampled IMM. Note that for ease of reproducibility, a fixed 0.7 was used for these results and no particular optimization tweaks were made to compensate for gradient noise, suggesting that performance could be improved by making such changes.
50
70
Sumple B4M
50
NV
Figure 10: Serialized IMM vs. Sampled IMM Noising for logistic model. Serialized IMM provides ((k)- times computational speedup over Sampled IMM (where l is the number of samples and, in this case, k = dataset size). The serialized version lacks the smoothing of the sampled version yet upproaches it for larger dataset sizes and always improvex noising.
E Potential Negative Societal Impact
Improving language models can have negative societal impact, if these language models are fine-tuned for tasks that are not aligned with positive human values. We hope that progress achieved by new methodologies such as IMM will not be put into misuse of this kind. On the contrary, we expect that any improved statistical efficiency achieved by methods like IMM to bring value when data is not as abundant, thus helping make machine learning and language modeling more impactful in underserved domains, where data is more scarce.
28

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
IMM as generalization of noising is analytically shown in the paper. along with noising's s caveats that are solved by IMM. It's application is experimentally demonstrated the paper, in Language Modeling and Reinforcement Learning experiments, as well proof of concept, using Logistic Regression.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
IMM, in some applications like noising, requires random sampling, because exact evaluation is not feasible. This affects the accuracy of the induced model, but regardless, IMM is able to outperform the baseline (i.e. noising).
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
The only proof is the one that discusses the caveat of noising through a counter example. All ussumptions and complete proof is present.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Code for all experiments has been anonymously released with clear instructions for re- producibility. A summary of relevant experimental details is also contained within the paper.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Code for experiments has been anonymously released with clear instructions for re- producibility. A summary of relevant experimental details is also contained within the paper.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
A summary of relevant experimental details is also contained within the paper.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Error bars have been provided along with the number of Monte Carlo runs performed to obtain each bar.
29

---

8. Experiments Compute Resourees
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
We mention the GPU used (Nvidia V100 with 32 GB memory)
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics atttps://neurr htt.ghtttt.s///a/r.......p/ppl/p
Answer: [Yes]
Improving language models have negative societal impact, if these language models fine-tuned for tasks that are not aligned with positive human values. We hope that progress achieved by new methodologies such as IMM will not be put into misuse of this kind.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Appendix
B addresses this.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
None of our datasets are scrapped. Pre-trained checkpoints are from reliable sources (Google trained checkpoint).
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
All our code (both original and derivative) Apache 2.0 licensed and a copy of this license is included in all code repositories. Datasets are public domain.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The only new asset is code. Code for all experiments has been anonymously released with clear instructions for reproducibility. summary of relevant experimental details is also contained within the paper.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
No human subjects were involved.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
30

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
No human subjects were involved.
31


Reviews:
[{"review_id": "ocY5d3dHTs", "reviewer": "Reviewer_vmaB", "summary": "- Algorithm: This paper proposes a framework for how a good but restricted feature model, e.g. $\\bar{P}(y \\mid x_1)$, can be used as guidance when training a full-feature model, e.g. $Q(y \\mid x_1, x_2, x_3...)$. \n\\\nInstead of the knowledge distillation objective from weak teachers, which directly adds a regularizer comparing the probabilities of the restricted feature model \\bar{P} and full-feature Q with each other directly, IMM proposes comparing P with Q marginalized over the other features\n\\\n$$\\mathsf{Reverse-KL}: \\lambda \\sum_{y-} \\pi(x_1, x_2, x_3) \\sum \\bar{P}(y \\mid x_1) \\log \\frac{1}{Q(y \\mid x_1, x_2, x_3)}$$\n$$\\mathsf{IMM}: \\lambda \\sum_{y-} \\pi(x_1, x_2, x_3) \\sum \\bar{P}(y \\mid x_1) \\log \\frac{1}{\\hat{Q}(y \\mid x_1)} \\text{ where } \\\\ \n\\hat{Q}(y \\mid x_1) = \\sum_{x_2, x_3} \\pi(x_2, x_3 \\mid x_1) Q(y \\mid x_1, x_2, x_3)\n$$\n- The marginalization can be quite expensive in practice, but they demonstrate (in small RL and Language tasks) that they are able to approximate the marginalization efficiently by sampling. The single-sample IMM objective is connected to noising/reverse-KD. \n- They mathematically show that in the infinite data regime with the perfect true distribution $\\bar{P}(y | x_1)$ and the true model P(y \\mid x_1, x_2, x_3) is in the hypothesis class of Q, IMM objective always recovers P. On the other hand, reverse KL and noising cannot. \n- Empirically, IMM always does better than with no-IMM, but the performance difference is especially large in data-limited regimes.", "questions": "-  Is there any practical importance to the infinite-data regime analysis? Are there circumstances where training with the IMM objective harms performance compared to just cross-entropy with increasing data?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "While the objective may still be too computationally expensive for very large-scale tasks (i.e., sufficiently sampling a suffix tree, high inference cost), the paper is thought provoking, and well-written/motivated. I would recommend it for acceptance.", "weaknesses": "- One weakness of the paper is obviously the small scale of most of the experiments (learning policies over 11x11 grid, subset of GLUE, toy linear model). But the experiments sufficiently demonstrate the proof of concept. \n- It would also be nice to include some settings where the feature-restricted model is not close to optimal as an ablation study. \n- In Figure 1, why is noising able to achieve almost the same performance as IMM starting from around 5% dataset size, if noising is similar to  single-sample IMM? Generally, what is the importance of using more samples when marginalizing Q for improved performance? Does performance strictly increase with increasing samples?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for deeply understanding and appreciating our paper. We address some of the points raised:\n\n+ **Scale of the experiments** — The computational overhead of the model and the fact that we didn't have the infrastructure for large models, meant that we dedicated ourselves to simple reproducible benchmarks to demonstrate the merit of IMM. (Many of the experiments stem from our original goal of understanding the underpinnings of the practice of noising.) In the general rebuttal, we propose approaches to more radically overcome computational hurdles in the future.\n+ **Suboptimal feature-restricted models** — We address this! We do it experimentally in Appendix D4 (Figures 6 and 7), where we artificially weaken the models in the logistic regression and RL experiments. (Additionally, the Kneser-Ney bigram in language modeling is very good, but it certainly is not the ideal bigram of English.) The takeaway is that, by tuning $\\lambda$, IMM always helps and never hurts, even with good but suboptimal restricted models, and its gains  are commensurate with the quality of the model.\n+ **Noising vs IMM at small dataset sizes** in Figure 1 — High level explanation: In this regime, IMM needs to marginalize/induce the model accurately and it can't do that at _extremely_ small data sizes. Lower-level explanation: although noising is a biased version of IMM, it also has less variance, because it doesn't need the induced model. With very little data,, that's useful. This explanation also suggests a possible best-of-both-worlds solution, by performing a bias-variance tradeoff. This is something we've thought about but not yet investigated since the issue only surfaces at such extremes. The performance of IMM increases with more samples for marginalization, because the variance is reduced. However, the returns diminish, and we find that k=5 to 10 samples are sufficient to get good performance.\n+ **Infinite data analysis** — This analysis is just an analytical tool to contrast IMM and noising, and to reveal the latter's shortcomings.  If one truly had infinite data, neither of those techniques would be necessary. The only circumstance where IMM will hurt is if the target model is very wrong and the $\\lambda$ is not tuned to account for it. (If the target model is correct, $\\lambda$ can be arbitrary.) In contrast, the only circumstance where noising/reverse-KD will _not_ hurt is if the target model is correct and the $\\lambda$ shrinks with more data. This is what we do in Figure 1 (i.e., for noising, we use the high quality feature-restricted model and we find the best $\\lambda$ at each data size). Despite this, the benefits of noising disappear quickly with increasing data size.\n\nWe hope these address your points. If there is anything else we could elucidate, please don't hesitate to ask us during the discussion period. We know you've already been very generous with your score, but if you could consider a higher score, AC's decision would be easier, and we would be extremely grateful! Thank you again for your time and effort.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XlQNbzPhHC", "reviewer": "Reviewer_2S2s", "summary": "The authors introduce the problem of “Induced Model Matching” where there exists a small and restricted model that only takes into account some of the features and is able to predict relatively well the label given these features. The key question of this paper is how one can leverage such a small model when training a full-feature, larger model. The authors answer this question by providing an algorithm which aims to match the restricted features version of the large model to the restricted features model. The authors suggest a regularization term to the loss, called IMM, which accomplishes this. The authors present a toy experiment of regression as well as a larger-scale experiment with language, showing that their method is able to perform better when using the IMM regularization.", "questions": "* It might be realistic to assume that $\\hat{P}$ has some loss $\\varepsilon$. Do you have any insights  as to how one might expect IMM to degrade as a function of $\\varepsilon$? \n* Did you consider the setting in which there are many restricted feature models? How do you think your setting can be generalized even further? Do you think there are any settings in which using IMM regularization might hurt the model? For example, if the task is recalling something from a long time ago in the context, it might seem reasonable to assume this method would not work. Generally speaking, what kinds of problems do you think your method is best suited for? Information theoretically, there might be examples in which having some base solver might make the overall problem easier, but as mentioned, there might also be scenarios in which this method actually harms the training of the overall model. \n* Just for my understanding, are you assuming a fixed-context length model?", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This is a very interesting topic, and seems like it could be useful in real-world scenarios. The idea appears novel and also is an interesting angle when compared to knowledge distillation. \n* It is very well-analyzed, and thoroughly explains how the method compares to existing methods. \n* The authors provide a clear presentation of their ideas as well as clear definitions. It is also nice to also have included a glossary in the appendix. Overall the paper is very clear to follow and precise about definitions.\n* The experiments show that, indeed, the proposed regularization term benefits the large model.", "weaknesses": "* As the authors mention, it is difficult to compute the regularization term. The computational cost is a drawback, although having increased performance is still a nice result. Given the computational cost, this brings into question when this method becomes useful in practice. I will leave my thoughts on this in the questions section.\n* It would be interesting to include further analysis regarding how the restricted model’s performance affects this method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for deeply understanding and appreciating our paper. We address some of the points raised:\n\n+ **Computational cost** — Our general attitude is: given a fixed amount of data and a feature-restricted model, how can we do most with it? We are thus mostly concerned with statistical performance, however we acknowledge that computational performance is critical for applicability. In the general rebuttal, we propose approaches to more radically overcome computational hurdles in the future.\n+ **When $\\hat P$ is of lower quality / how this quality affects things** —  We address this experimentally. Due to lack of space, we moved this to Appendix D4 (Figures 6 and 7). There, we artificially weaken the models in the logistic regression and RL experiments. (Additionally, the Kneser-Ney bigram in language modeling is very good, but it certainly is not the ideal bigram of English. We could also do an ablation by worsening it.) The takeaway is that, by tuning $\\lambda$, IMM always helps and never hurts, even with good but suboptimal restricted models, and its gains  are commensurate with the quality of the model.\n> We do not yet have a theoretical analysis of restricted-model quality vs. benefit. However, the insight of why this works is that the dual interpretation of $\\lambda$ is as constraining the learned model to have its induced version $\\overline{Q}$ in a $\\delta$-ball around $\\hat P$, with larger $\\lambda$ meaning smaller $\\delta$. If $\\hat P$ is $\\epsilon$-away from the true $\\overline{P}$, we need to make $\\delta\\geq \\epsilon$, to make sure that we're not harming the learned model by keeping it artificially away from the truth. When the model is high quality ($\\epsilon$ small), we can thus make $\\delta$ small, and thus $\\lambda$ can be large. However, when the quality is low ($\\epsilon$ large) then we need to make $\\delta$ large too, which means that $\\lambda$ has to be small, and therefore the relative benefit from IMM diminishes.\n+ **Did you consider many restricted-feature models** — We did think about this, but we haven't worked on it yet. Similarly to multi-task learning, we imagine that adding multiple IMM losses, each for a different restriction, could be able to handle this.\n+ **Other generalizations** — One setting we contemplated to try IMM in is in building physics-informed models. A closed-form physics model (e.g., climate prediction) can be thought of as a restricted model, because it often relies on very specific features. A full-featured model could use many more features than the physics model. IMM can be used as a method to incorporate the physics model, by making sure that the full-featured model's induced version matches the physics model. Another generalization that we're studying is to take IMM beyond feature-restriction, to also cover task-restriction: incorporating models that can perform a sub-task very well.\n+ **Settings where IMM would not help** — As we have noted in Appendix D4, if the restricted model is bad and $\\lambda$ is not tuned properly, IMM could hurt. However, let us assume that the restricted model is good, perhaps even the ideal restricted model. In this case, we speculate that IMM can only help. It may be counterintuitive, but even if the restricted features are uninformative for the task (as in your example), IMM can help. Indeed, if the task is impossible to perform with the restricted features, then the best restricted model will be no better than chance. IMM will try to make the learned induced model mimic this, and perform no better than chance too. We believe this will have the effect of informing the full model's training that the restricted features are pointless, which is useful information! Without this guidance, the full model could waste resources/data to discover this fact. In contrast, with this guidance, it can search the subset of hypothesis space that does not use the restricted features. As this is a smaller space, IMM can be thought of as regularizing the learning even in this case. Therefore, we believe our method is suited for any problem in which we can have a reasonably  accurate approximation of the restricted model.\n+ **Are we assuming a fixed context length model** — The full context can have variable length, however, we are assuming that the short context has a fixed length. That said, this is mostly about how we're formalizing this in the paper. The methodology itself can work with either, as long as there is a consistent way to split each full context into a short and extended context.\n\nWe hope these address your points. If there is anything else we could elucidate, please don't hesitate to ask us during the discussion period. We know you've already been very generous with your score, but if you could consider a higher score, AC's decision would be easier, and we would be extremely grateful! Thank you again for your time and effort.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yPoBvezsbI", "reviewer": "Reviewer_ADbo", "summary": "The paper considers the learning problem when in addition to the training set an additional _restricted_ model is available. The restricted model is trained on a different dataset, potentially containing a only subset of features. It is proposed to augment the training loss with the special induced model matching (IMM) loss that encourages similarity between full and restricted predictive distributions where the extra features are marginalized on the empirical data distribution. IMM is compared to similar techniques such as noising and weak-teacher knowledge distillations where the benefits of IMM are demonstrated theoretically. Experiments are conducted on a toy logistic regression problem, language modelling tasks and reinforcement learning in a toy environment.", "questions": "1. I struggle to find a clear intuitive explanation for why induced model matching helps. Can authors think of a simple 2-dimensional logistic regression problem (for example) and include a graphical illustration of full and restricted models (and their induced variants)?\n\n2. Do I understand it right that the LM experiments the restricted n-gram model has been \"trained\" exactly on the same dataset on which the full model is being trained?\n2.1 If so, then, again, the only reason why IMM helps I can think of is some kind of \"smoothing\" or regularization of the full model, which probably learned an overly \"kinky\" distribution from limited data. In the eyes of the authors is that the right intuition?\n\n3. Is there an experiment where the restricted model indeed comes from a different and much larger dataset (which authors mention as indeed an interesting scenario for applying IMM)? It would be interesting to look at the \"trade value\" of a\"full\" data point compared to a\"restricted\" data point.\n\n4. Could authors hint at the scenarios in which IMM is not useful or even harmful? I would be especially interested if that could be the property of the model classes or data distributions from the full and the restricted model come.", "limitations": "Adequately discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The proposed idea is interesting and it seems to be an improvement over the existing techniques such as noising and reverse-KD. I see value in being able to systematically incorporate restricted knowledge from a different distribution. Empirically IMM seems to be helping on a number of LM tasks.", "weaknesses": "The paper adopts rather exotic notation that doesn't help reading the paper. I had to constantly look into the glossary in the appendix to go through equations. Descriptions of each experiment are also not self-contained, even when reading the appendix. The general idea of IMM is not the most intuitive to me and I think the paper would benefit from a simple graphical illustration of the principle. In my opinion, insufficient clarity in presenting the method and describing experiments is the main weakness of the paper. \n\nIf I understand the setup of LM experiments correctly (full and restricted models are trained on the same dataset) then the idea of the restricted model capturing a richer distribution is under-explored and the only way I can see IMM improving performance is as a regularizer (see my questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very constructive review! We address some of the points raised:\n\n+ **Notation** — Thank you for your suggestions on improving the notation. We have been working on a few potential alternatives to further clarify our presentation. Here's what we suggest to make things more readable.\n  - Instead of using $y_{-}$ or $y_{-t}$ for context, we'll use $x$ nd $x_t$, and reserve $y$ and $y_t$ only for predictions.\n  - In order to denote the short/extended components of a context, instead of using $\\textsf{sh}(x)$ and $\\textsf{ex}(x)$, we'll use $\\overline{x}$ and $\\underline{x}$ respectively. This will parallel using $\\overline{P}$ when denoting the induced model, which only depends on the short history.\n  - While we tried to be very explicit in our notations, we believe these changes will declutter equations and make everything more legible. We hope that you approve. We are open to other suggestions.\n+ **Graphical Representation** — Thank you for the suggestion. We propose to offer a high-level visual representation of the approach. [We have drafted the following figure](https://i.imgur.com/NwQgOPV.png), which also incorporates the notational changes suggested above. We hope to improve this and include it in the paper.\n+  **Experiment Descriptions** — In the main body of the problem, it is difficult to include all the details of the experiments. Appendix D has all the details. However, we will revisit how we have divided this information across the main body and the appendix. We will make every effort to make each experimental setup self-contained in the main body, with precise references to the appendix.\n\n\nAlso your questions:\n\n1. **Visualizing restricted model in logistic regression** — [The following figure](https://i.imgur.com/G4lLm2u.png) illustrates the 3-dimensional logistic regression problem in the paper. The features are samples uniformly in this box. The Bayes-optimal restricted model only uses the x-coordinate, and so assigns probabilities proportionally to the blue/red areas in the illustrated slice. IMM then encourages the full logistic model to be consistent with these weights, i.e., making sure the proportion of points labeled $\\pm$ agrees with these weights at each x. Intuitively, this biases the separating plane to have the right inclination/alignment with the x-axis, which subsequently speeds up the learning process.\n2. **When the data provides the restricted model** — Your intuition is correct for our own language modeling experiments. We do use the data itself to create our bigram for PTB, and for BERT we use Wikipedia, which it is pretrained on. The extensive literature that tries to incorporate N-grams into more sophisticated language models argues indeed that N-grams tend to capture structural detail that may be missed by the larger model, albeit in their restricted setting. However, some newer papers train N-grams on much larger datasets (since it's faster to do so), which adds more benefit (see Liu et al. 2024).\n3. **When the restricted model comes from elsewhere** — In our other experiments, logistics regression and RL, the restricted model is more powerful/rich. As mentioned above, for logistics regression it is the Bayes-optimal restricted model. For RL, it is the exact solution for the POMDP, which can be thought of as full exploration of the reward landscape, but with only one-coordinate observation. We do not have a data-point by data-point value comparison, however, we do have in Appendix D4 (Figures 6 and 7) the effect of artificially weakening those perfect models. The bottom line is that as long as we have that extra data to build a decent approximation for the true restricted model, then IMM improves performance, in a way commensurate with the quality of the model (thus the amount of restricted data.) We have not done a theoretical analysis of this tradeoff, but it's an excellent suggestion for future investigation.\n4. **When is IMM harmful** — In our experience, we can identify two scenarios in which IMM can be harmful:\n  + If, simultaneously, (a) the restricted model is of bad quality, and (b) the value of $\\lambda$ is not properly tuned. This is hinted at in the importance of tuning lambda properly in the low-quality model experiments of Appendix D4. Therefore, you are right that if the model class of the restricted model is not powerful enough, it may not capture the true restricted model properly, and could lead to harm. (However, this goes against the main premise of the paper, i.e., the availability of good target models)\n  +  Also, when the amount of data is so limited or the distribution of the extended context given the short context is hard to learn (e.g, lack of structure such as smoothness or latent low-dimensionality), then the learned induced model $\\hat Q$ will not be accurate, and the performance can suffer. We see a hint of this in Figure 1, where with very few data points IMM is bested by noising (though it is not worse than no-IMM).\n\nWe did our best to address your concerns. If there is anything else we could elucidate, please don't hesitate to ask us during the discussion period. We appreciate your leaning toward acceptance. We hope that you will engage us further, and that based on this we will earn a higher score from you! Thank you so much for your time and effort.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ai5o9NOwB0", "reviewer": "Reviewer_vLY4", "summary": "This paper proposes a method to train models which have access to all of the features by making their marginal distribution match a known weaker model which only uses a subset of those features to predict. The authors relate this to knowledge distillation and noising, and come up with an approximate objective which computes this marginalization over the bigger model. Even though it might be intractable to do analytically, they provide a sampling-based method for the computation. They evaluate the model on logistic regression and language modelling", "questions": "My questions all relate to my understanding of the limitations of the IMM approach.  I would be happy to reconsider my assessment if the above questions are addressed satisfactorily.\n\n1. Could the authors clarify the approximate computational and memory complexity of this method in terms of the number of sequences n, sequence length L, and any other relevant parameters? It would be helpful to see the analysis for both RNNs and Transformers.\n\n2. Could the authors comment on the asymptotic space complexity of the auxiliary data structure? For large corpora, it seems like this could require a huge amount of disk space and be challenging to query efficiently. Is there a data structure that could mitigate these concerns?\n\n3. Could the authors provide a FLOPs-corrected comparison between the IMM method and the vanilla training approach? Given the higher computational cost per forward and backward pass for IMM, it would be informative to see how the two methods compare when given a similar computational budget.\n\n4. Can the authors comment on practical cases where the IMM approach would be preferred over spending the additional $k$ FLOPS on e.g. processing $k$ more examples?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The relation of the proposed method to the previous work in noising is insightful and not immediately obvious, giving an interesting perspective.\n- The connection to knowledge distillation is valuable to highlight.\n- The code for the method is available.\n- The text is generally very well-written and explains the proposed method very well.", "weaknesses": "- The empirical results are somewhat dated. The models considered are quite old (RNNs, BERT) and it would be nice to see a more modern set of models such as decoder-only transformers or state-space models. However, using the old models does allow comparison to the noising approach. While the direct comparison to noising is appreciated, in practice most people would likely be comparing to the vanilla training approach as a baseline.\n\n- The main weakness of the paper is the computational overhead incurred by this method, which is not adequately addressed in the paper. With the introduction of language modeling, the question of how this method would scale to larger datasets and models is important, and the scaling properties do not appear favorable.\n- For the RNN experiment, each forward and backward pass seems to require k additional passes due to the gradient accumulation for the samples drawn for the model matching objective. This incurs a significant overhead factor of k.\n- The overhead factor may actually be worse when considering the need to replace the context for each of the k summands in the IMM objective. This requires swapping out the current context and computing a new set of hidden vectors, potentially up to the length of the input. A back-of-the-envelope calculation suggests the FLOPs could increase from O(L) to O(L + kL^2). However, there may be some subtlety I'm missing here that reduces the overhead.\n- For a large dataset, the lookup of the context could be quite difficult to implement. In the worst case, for n sequences of length L, storing the auxiliary data structure could require space on the order of nL^2, which is a factor of L larger than the original dataset. Supporting random access to this data structure may not be feasible for large corpora like Common Crawl without significant infrastructure overhead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your deep reading of our paper and for being open-minded about revisiting your assessment! Regarding the key points you raise:\n\n+ **Empirical results dated** — You are correct that the motivation for comparison with LSTMs was because of  the noising baselines existing there. We used the encoder/decoder BERT to illustrate validity for Transformers, however we didn't have the infrastructure for large-scale models like LLMS. We hope that the body of all the experiments we've included, non-language models too, show the wide merit of the approach.\n\n+ **Computational overhead** — IMM, as implemented in the paper, does have a computational cost relative to the baseline of training only with cross-entropy. We detail this overhead in answering your questions below. You are right that this is a limitation, however our focus was on the statistical aspect. Given limited data (so additional samples are not an option) along with a good restricted model (potentially derived from the same data, or from a corpus of restricted data), IMM achieves improved accuracy/perplexity/reward. That said, in the general rebuttal we outline a potential solution to tackle the computational overhead.\n\n\n\nRegarding your questions:\n\nFirst, to clarify, your use of $L$ should be replaced with the unroll length of the LSTM, which is the maximum depth through which gradients are backpropagated (set to 35). Also, to be consistent with the paper, we use $n$ for the length of the data (in tokens/words).\n\n1. **Time complexity** — The best way to quantify the computational overhead is relative to the baseline of using traditional cross-entropy.\n  +  Let $L$ again represent the unroll length of the LSTM, then since histories have to be swapped at every unroll location and the LSTM re-evaluated, the overhead factor is $\\mathcal{O}(kL)$. Therefore your expression is partly correct, apart from $L$ representing unroll length and $n$ data size, i.e. we go from $\\mathcal{O}(n)$ to $\\mathcal{O}(nkL)$ per epoch. We partially mitigate this in our current implementation by applying IMM periodically (not at every iteration), see Appendix D2.2. For example, if we do it only every $\\Omega(1/L)$ iterations, the overhead factor becomes $\\mathcal{O}(k)$ and we go from $\\mathcal{O}(n)$ to $\\mathcal{O}(nk)$ .\n  + For BERT/Transformers (and in fact in our non-language modeling experiments), the baseline does not have a recurrence that require new passes, and therefore the overhead is $\\mathcal{O}(k)$, so we go from $\\mathcal{O}(n)$ to $\\mathcal{O}(nk)$.\n  + A $k$-fold increase is not ideal, but it is acceptable as our choices of $k$ range from $5$ to $10$.\n\n2. **Space complexity** — There are two overheads:\n  + *Model memory*: As explained in Appendix C.2, we sequentialize the gradient computation across the $k$ sample. This means that the space overhead during training is a factor of $2$ compared to the baselines (i.e., a second set of gradients), which are themselves of the order of the number of parameters of each model.\n  +  *Lookup overhead*: You are correct that a naive implementation of the data structure could take space $\\mathcal{O}(nL)$. However, the better implementation is with a dictionary/hash table of lists containing indices/pointers to a reverse linked list representing the data set. By referencing a position, we can recover a length-$L$ history at any position, at runtime. This requires space $\\mathcal{O}(n)$ only,  because each key in the dictionary is a short history, and the number of long histories is the same as the number of times that short history appears. Adding up the number of occurrences of all short histories, gives us $n$.\n\n3. **FLOPS-Correction**  and 4. **Comparison** — For our reply, we would like to address an important misunderstanding here.\n  + IMM is not expected to be used in situations where we have access to an interminable data stream, where we can simply continue to train the cross-entropy baseline until the clock runs out. Even in the best of cases (see the general rebuttal) IMM could take twice as much time as the baseline, and it's unclear when the performance gain could exceed having twice as much data (often the gains seem to be equivalent to 20-30% more data.)\n  + Rather, IMM is expected to be used when the data is what it is, and one is required to additionally incorporate the side-information presented in the form of the feature-restricted model. The goal is to obtain a statistical advantage, which IMM does, at the expense of computational overhead. In the premise of IMM, it simply is not an option to collect that extra full-featured data. Therefore FLOPS-corrections and comparisons are not germane to the context, and if done they will unfairly show that IMM is not competitive. We hope that you do accept this explanation as to why this request cannot be addressed equitably.\n\nWe believe that we have addressed all your major points. If there is anything else we could elucidate, please don't hesitate to ask us during the discussion period. We appreciate your deep insight about the computational aspect of the problem. We hope that we have quantified our current overhead and given insight about how this overhead can be reduced in the future. However, considering that our focus was primarily statistical, we hope that you will judge the merits of the paper on that basis. We believe we have something very valuable for the community. We would immensely appreciate it if you could recommend acceptance of the paper! Thank you once again for all your time and effort.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank you all for your insightful and positive reviews. We are encouraged by your appreciation of our work and for your constructive criticism. We are lucky to have received such high quality feedback.

We have individually addressed all the points that you've raised. There is, however, one common theme that arose across the reviews: **whether IMM can be made more computationally efficient.**

+ While the paper itself is focused mostly on the statistical aspect of the problem, we recognize that scalability of IMM is critical for adoption. As such, we have been working on this issue, and believe that there is an elegant solution that could potentially achieve this scalability. Unfortunately, this work is not complete yet and we don't have experiments paralleling those in the paper. However, we are happy to share with you our insights.

+ The main bottleneck with IMM is the need to calculate the learned induced model. In the paper, we are doing this by sampling $k$ long histories and averaging the outputs. This has certain advantages, such as giving us a low-variance estimate of the gradient. However, it can be computationally expensive, especially in recurrent models which require new passes over the substituted histories.

+ The following alternative idea is similar to the sequentialization aspect covered in Appendix C.2, Equation (19). Using $\overline{x}$ to refer to short history and $\underline{x}$ to refer to long history, we can write the idealized IMM loss as ([see this figure for reference of this new streamlined notation](https://i.imgur.com/NwQgOPV.png)):

$ - \sum_\overline{x} \pi(\overline{x}) \sum_y \overline{P}(y|\overline{x}) \log \overline{Q}(y|\overline{x})$

where the learned induced model is:

$ \overline{Q}(y|\overline{x}) = \sum_\underline{x} \pi(\underline{x}|\overline{x}) Q(y|\underline{x},\overline{x})$.

The gradient of this IMM loss then becomes:

$ - \sum_\overline{x} \pi(\overline{x}) \sum_y \overline{P}(y|\overline{x}) \frac{\sum_\underline{x} \pi(\underline{x}|\overline{x}) \nabla Q(y|\underline{x},\overline{x})}{\overline{Q}(y|\overline{x})}$

$ = - \sum_\overline{x} \sum_\underline{x} \pi(\overline{x}) \pi(\underline{x}|\overline{x}) \sum_y \overline{P}(y|\overline{x}) \frac{ \nabla Q(y|\underline{x},\overline{x})}{\overline{Q}(y|\overline{x})}$

The empirical version of this gradient is:

$ - \frac{1}{n} \sum_t \sum_y \overline{P}(y|\overline{x}_t) \frac{ \nabla Q(y|\underline{x}_t,\overline{x}_t)}{\overline{Q}(y|\overline{x}_t)}$

As you can see, there is no sampling of any histories in this expression! It can even be rewritten in the form of cross-entropy, making the connection with noising and reverse-KD even more apparent, as it introduces a correction factor:

$ - \frac{1}{n} \sum_t \sum_y \overline{P}(y|\overline{x_t})  \underbrace{\frac{Q(y|\underline{x_t},\overline{x_t})}{\overline{Q}(y|\overline{x_t})}}_{\textsf{correction}} \nabla \log Q(y|\underline{x_t},\overline{x_t})$

The fascinating thing about this approach is that it only has a constant factor overhead on the baseline of cross-entropy training. There are two caveats, however:
+ This has a higher variance than the sampling approach of the paper. Reducing the variance is challenging, and we are investigating it using momentum-based approaches.
+ It requires the maintenance of an induced-model estimate. We can solve this relatively easily for the bigram case, by accumulating prediction vectors into a matrix.

We are sharing this insight with you in the hope of convincing you that there are indeed avenues to making IMM computationally effective. We don't know whether we'd be able to share with you parallel experiments by the end of the discussion period, but we hope that you will take this into consideration and become even more confident in the merits of this work. Thank you very much for all your time and effort!

*Note: The attached PDF and the links in the rebuttals are all anonymized and contain the same two figures.*

Meta-review: This paper explores an interesting setting in which a (very accurate) model relying on a small subset of features is available as side information to a modeler who seeks to now train a larger fuller-featured model on a full dataset. The paper proposes a new approach called IMM, showing results on simple models like logistic regression and extending the evaluation to include results with deep nets and language models. Reviewers all vote for acceptance, finding the paper to be well-written, appreciating the connections to the knowledge distillation literature, appreciating the strength of the analysis, and finding the idea to be practical.
Decision: Accept (spotlight)