Induced Model Matching: Restricted Models Help Train Full-Featured Models

Authors: usama  muneeb
mesrob i ohannessian
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 19:07:24

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.

Main paper: 
 introduction: 
 In many applications, it is both statistically and computationally easier to construct (often small) feature-restricted models. In this paper, we 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 7. Say we have three features (x 1 , x 2 , x 3 ) ∈ R 3 and a binary label y ∈ {0, 1}. Let features be generated from a distribution π and let labels be conditionally generated according to a logistic (full-featured) true model P (y|x 1 , x 2 , x 3 ). Assume we have ample feature-restricted data with only one feature (x 1 , y), which, by marginalization, are samples from the (feature-restricted) true induced model P (y|x 1 ) = x2,x3 π(x 2 , x 3 |x 1 )P (y|x 1 , x 2 , x 3 ). By virtue of this data, we get a very good approximation P of P , which we can use as a target induced model. How do we use P , along with full-featured data (x 1 , x 2 , x 3 , y) to obtain a (full-featured) learned model Q(y|x 1 , x 2 , x 3 )? A 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 a secondary loss that keeps Q close to P . This is reasonable -in Section 2 we relate this to reverse knowledge distillation and in Section restricted set of features. Instead, what this paper proposes is to equalize the field when comparing to P , by inducing a restricted model Q(y|x 1 ) 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 1 we show how this speeds up learning and reduces predictor variance. 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 we offer here is that such noising is best understood as an attempt to incorporate a feature-restricted model's knowledge into the full-featured model being trained. This interpretation reveals the general fundamental question: If we are armed with a very accurate feature-restricted model, what is the right way to incorporate it into the training of a model with a larger feature set?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 a regularizer that matches the predictions of the induced learned model with that of the desired target restricted model. This is the language model instance of the induced model matching (IMM) methodology. 3. In Section 5, we 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 to the noising approach of Xie et al. (2017) and to reverse knowledge distillation. We share two key findings. First, these alternatives may be thought of as approximations to IMM. Second, they have a major caveat. They may not be consistent even in the ideal infinite-data regime, in contrast to IMM, which is consistent. 5. In Section 7, we experimentally demonstrate the effectiveness of 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 8 we discuss limitations and applications beyond those illustrated in the paper. 
related works: 
 We review two lines of work that are closely related to IMM, noising and knowledge distillation, which are respectively implicit and explicit versions of the same idea. We also review key literature showing the continued merit of N -grams as restricted models.Noising and Data Augmentation in Language Modeling The noising methodology in 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, in 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 et 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. (2017) were able to show a clear advantage, beyond the use of other regularization techniques such as dropout. Section 6 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 -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 a special data structure to precompute 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 -grams also exist, such as that of Li et al. (2022) , who let LMs learn what the N -gram could not, i.e., the residual.Knowledge Distillation Knowledge distillation (KD) is a paradigm first proposed to let a powerful teacher model help better train a 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,EQUATIONwhere Cross-Entropy uses hard data, x are contexts, π n is the empirical distribution of contexts, Q is the learned/student prediction model, and P teacher is the teacher model, which can be thought of as providing soft predictions. τ , 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 a powerful student. This "reverse-KD" 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 a 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 et al. (2017) transition from uniform noising to Kneser-Ney noising. Later, reverse-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 a powerful student's learning. A quick comparison to Eq. ( 14) when P teacher is the restricted model P reveals that reverse-KD has more in 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-KD ignores this fact by comparing the student model Q directly to the teacher model in the KL term of Eq. ( 1). 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 in language, unlike in vision, can be harmful with very weak teachers, and thus it is 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. 
problem description: 
 Problem Setting We consider the problem of training a full model Q(y t |x t ) that is able to predict a label of data point t based on its full context. For example, in forward predictive models, x t is the sequence of past tokens and y t is 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 (x t , y t ) of size n, possibly tokenized from a 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) = - t log Q(y t |x t ) ≡ x π n (x) y P n (y|x) log 1 Q(y|x),where π n is the empirical distribution of the context and P n be the empirical distribution of the prediction given the context. Note that (x t , y t ) refer to tokens while (x, y) refer to types. Minimizing this empirical risk can be seen as Q striving to approach the "true" model P (y t |x t ) generating the data, in average context-conditional KL divergence. The idealized risk is thus:EQUATIONwhere now π is also the true (not empirical) context distribution.Induced Models While Q strives to 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 x t can be decomposed into a short context x t and an extended context x t , and that one has access to the "true" model that predicts y t based solely on the short context x t , i.e. P (y|x). 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:EQUATIONWe call P the true induced model. It depends both on the context distribution π and the true model P .Since we do not have the latter two, we cannot explicitly compute P . What motivates us, however, is the possibility to learn it 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 a very good approximation thereof, how can this information be incorporated to improve the learned model Q?4 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 to match P with Q, but rather with Q, the learned induced model that, just like the move from P to P in Eq. ( 3), specializes Q to performing predictions with only the short context:EQUATIONLet'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-conditional KL divergence with the restricted context:EQUATIONSince P and π are not available in practice, the idealized IMM risk cannot be computed. However, as the core motivation of using P 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, knowledge of π in Eq. ( 4) can be intuitively understood as a 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:EQUATIONThe proportionality constant is unimportant as, thanks to the logarithm, it contributes only a constant to the overall risk. By combining these empirical proxies, we obtain the empirical IMM risk:EQUATIONThis mirrors Eq. ( 5), up to P inside the logarithm, which only contributes an additive entropy term that does not depend on Q and is thus irrelevant to optimization.IMM as a regularizer Given a reliable P , we propose to incorporate this knowledge into the full model by using the IMM risk as a regularizer. Using a hyper-parameter λ that can be tuned, our induced model matching methodology consists of training the model Q by minimizingEQUATIONSeparating IMM into components To weave IMM into existing ML optimization pipelines, it is useful to treat it as a separable risk. For this, we can rewrite Eq. ( 7) by expanding the empirical distribution π n . Then, the components of this separable risk are given by the conditional crossentropies between P and Q, because we can write:EQUATIONThe 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 be used to 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 asP (y|x) = 1 n t (1 -ν(x))1{y t = y} + ν(x)b(y),where ν(x) is the missing mass for previous token x and b is the back-off distribution. For logistic regression in Section 7 the context space is continuous, and these components come from a separately trained one-feature predictor. For our RL example, they come from the optimal actions of a POMDP. 
computation: 
 A direct implementation of IMM is prohibitive, since evaluation of the risk, Eq. ( 8), requires making a secondary pass over the entire dataset for each t, 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 the advantage of low gradient variance at 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), as a proof of concept for IMM's scalability.Sampled IMM To alleviate the exact computation of IMM, we could maintain a dictionary/multiset of extended contexts for each short context1 , and sweeping only across those. This amounts to rewriting Eq. ( 6) as follows:Q(y|x) ∝ t ′ : x t ′ =xt Q(y|x t , x t ′ ) = x∈extend(xt) Q(y|x t , x), where extend(x) = t : xt=x { {x t } } (10)is 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. ( 10) remains expensive. In particular, gradients need to 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(x t ), followed by approximating this expectation with k samples. As a result, the IMM component at t becomes:EQUATION≈ -EQUATIONwhere x i ∼ extend(x t ) are samples from the multiset of extended contexts. Implementation details are deferred to Appendix C, including Algorithms 1 and 2 that describe a typical training loop when sampled IMM or serialized IMM are incorporated into SGD respectively.At first glance, it appears that sampled IMM requires maintaining k copies of the model. However, this can be sequentialized, as explained Appendix C.2. This means that the space overhead during training is a factor of 2 compared to the baseline of no-IMM (i.e., a second set 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 to restart the LSTM, which then incurs an O(kL)-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 it only every Ω(1/L) iterations, the overhead factor remains O(k).Serialized IMM Sampled IMM delivers the benefits of IMM, but at a computational cost. To make IMM truly scalable, it is imperative to bring it to 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:- x π(x) y P (y|x) log Q(y|x), where the induced model is Q(y|x) = x π(x|x)Q(y|x, x).The gradient of this IMM risk then becomes:- x π(x) y P (y|x) x π(x|x)∇Q(y|x, x) Q(y|x) = - x x π(x)π(x|x) y P (y|x) ∇Q(y|x, x) Q(y|x)The empirical version of this gradient is:-1 n t y P (y|x t ) ∇Q(y|x t , x t ) Q(y|x) = 1 n t ∇ IMMt - y P (y|x t ) Q(y|x t , x t ) Q(y|x t ) correction log Q(y|x t , x t ) (13)The last expression in Eq. ( 13) is in the form of a "corrected" cross-entropy between P and Q, which we denote by IMM t . The correction factor is not differentiated and is the ratio or Radon-Nykodim 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 is 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 Q † (that eventually becomes stale) is used to calculate Q † , use the ratio Q † Q † for correction factor. Using the current Q tends to cause instability, likely due to the correction no longer obeying expected constraints, e.g., mean 1 for every y. By choosing the update period inversely proportionally to the cost of updating Q † and by keeping Q † in memory, we incur an O(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, a connection that we elaborate on in Section 6. 
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 consistency only, even though IMM also confers a finite-sample advantage (understanding this advantage analytically is worth studying in the future.) We show (1) that IMM is consistent, (2) that noising and reverse-KD are akin to single-sample IMM, (3) that this shows that they minimize an upper 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.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 we recover the true model in the infinite-data regime, i.e., that IMM is consistent for all λ. 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.From single-sample IMM to noising and reverse-KD We first explain how IMM and noising are related. Experimentally, using a 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 t 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(x t ). That context belongs to a different data point t ′ . 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 x t ′ has no role in calculating the loss. Instead, the prediction is randomized through the sum over all possible y 0 in Eq. ( 9). Though not identical, this is a very close map to the noising-based data augmentation proposed by Xie et al. (2017) , namely prediction noising (see Appendix B.1 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 to the reverse-KD objective (Yuan et al., 2020; Qin et al., 2021) (see also Section 2) with P as the teacher:Cross-Entropy(Q) + λ x π(x) y P (y|x) log 1 Q(y|x) key difference (14)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 ( 12), which by Jensen's inequality can be upper bounded as follows:EQUATIONA single-sample approximation of the expectation inside the log is in fact a biased estimator of the left-hand side. However, it is an unbiased estimator of the right-hand side, with the expectation outside of the log. Thus, noising and reverse-KD 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. ( 5)), to show that this suboptimality can indeed occur. Even in the realizable case and with infinite data, IMM is always consistent but there exists a counterexample where noising and reverse-KD fail to consistently recover the true model. The proof is in Appendix B.2.Proposition 6.1. Assume that we are optimizing the idealized noising objective of Eq. ( 14) -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 -i.e., we are in the realizable case. Then, there exists a choice of π and P such that Q ⋆ ̸ = P .7 Experimental Results 
starting toy example: logistic regression: 
 Consider the logistic regression example from the introduction. The main results are given in Figure 1 and the full experimental details can be found in Appendix D.1. 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 = x 1 and x = (x 2 , x 3 ).Setting We sample features uniformly over a cube and assume we have ample data points of the form (x 1 , y). This allows us to build an excellent restricted model to predict the label based on just x 1 , call it P (y|x 1 ), nearly close to the (restricted) Bayes predictor or true conditional probability. Just like in language modeling, to induce a model we need to draw x's from its conditional distribution given x. Q(y|x 1 ), the induced model of Q(y|x 1 , x 2 , x 3 ), can then be interpreted as the average of Q's predictions, when x t = (x 2 , x 3 ) is drawn from its conditional distribution given x t = x 1 . 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 x for a given x. In logistic regression, this is not viable since x 1 is continuous and does not repeat. We rely instead on density estimation. We use a soft nearest-neighbor density estimate f (x 2 , x 3 |x 1 ) ∝ n t=1 δ x2,t,x3,t (x 2 , x 3 )e -α|x1,t-x1| , where 1/α is the bandwidth of the Laplace kernel. (With cross-validation, we determine α = 1 to be a good choice.) If we let w t (x 1 ) = e -α|x1,t-x1| , the resulting induced model by marginalization is:Q(y|x) = f (x|x)Q(y|x, x) ≈ n t=1 w t (x) n t=1 w t (x) Q(y|x, x t )These equations are respectively equivalent to Eqs. ( 4) and ( 6). The IMM risk and the corresponding overall objective remain the same:IMM(Q) = n t=1 y=0,1 P (y|x 1,t ) log 1 Q(y|x 1,t ) , Cross-Entropy(Q) + λ IMM(Q).In Figure 1 , we compare the performance of IMM-trained Q (green) to that without IMM (maroon). We sweep a range of n from 2 to 50, and use a cross-validation optimized λ for each (details in the Appendix D.1.1). 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 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. BERT BASE + MLM +IMM CoLA 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 ± 2.0 91.0 ± 2.0 93.5 ± 1.0 RTE 66 ± 3.0 68 ± 2.0 71 ± 1.0In these language modeling experiments, the restricted model we use is the modified Kneser-Ney bigram (Kneser and Ney, 1995; Chen and Goodman, 1999) 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 et al., 2014; Luong et al., 2015) , better neural models now dominate language modeling (Vaswani et al., 2017; Devlin et al., 2018; Turc et al., 2019) . We investigate how IMM could potentially improve even these more modern models, by using BERT's performance on some GLUE benchmarks as a proof of concept. (Appendix D.2.1 gives evidence that IMM indeed improves the full model's performance on the restricted task.)LSTM RNN Experiments We build on top of the code provided by Xie et al. (2017) 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 is a document dataset and the LLM is solving a next word prediction task. The average cross-entropy of multiple unroll positions of the RNN, or exponentiated as perplexity, is used as the measure of performance. For training and measuring validation perplexity, L = 35 unroll positions are used. During testing, only 1 unroll position is used. The IMM component is always calculated by running the LSTM in 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 Table 1 , 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 original BERT code, MLM was present during pre-training but removed from fine-tuning, possibly because of minimal gain. For us, however, MLM is ideally suited to be augmented with the IMM objective because it is based on cross-entropy and we can similarly generate an induced bigram 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 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 in an academic setting, we use a subset of the GLUE tasks (Warstadt et al., 2018; Dolan and Brockett, 2005; Rajpurkar et al., 2016; Dagan et al., 2005) to demonstrate the gains using IMM.For diversity, our selected GLUE tasks are a mixture of single-sentence and double-sentence tasks. Further experimental details are provided in Appendix D.2.IMM is particularly appealing in situations where incomplete-feature data may be much more available, due to reasons like law, privacy, or limited sensing. If an autonomous driving system is developed with few sensors, and later a new system taking advantage of more sensors is to be designed, the older system may act as a 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 a platform'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 reinforcement-learning (RL) solution.If the personalized engine has full-featured data, it 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 policy-gradient methods such as REINFORCE (Williams, 1992) . In Figure 2 , we illustrate the reward achieved with and without the use of IMM, for learning policies for an agent on a toroidal 11 × 11 grid, with reward peaking at its center. The POMDP only observes one coordinate, whereas the MDP observes both.  
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-based data augmentation techniques in natural language noising, as an attempt to harness the knowledge of the restricted model used for noising. We showed that naïve noising is not the best way to incorporate this knowledge, as it may fail to consistently recover the true model. IMM, on the other hand, directly aligns the learned model with the target model and is 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 is slightly less accurate but has a potential to be as efficient as the no-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-restricted 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 LLMs using current shorter-context LLMs. The principle behind IMM is applicable very generally and we hope this work gives impetus to such explorations. Table 3 compiles our main notation. We summarize the main formulation here and explain how it connects to its instances in the experiments. 
 
 Appendix: • Q is always the full model (LSTM or Transformer). Q(y t |x t ) is thus a learned contextconditional model: takes the full context (x t ) and outputs prediction probabilities (over y t ). In Eq. ( 4),log(Q) is the log-loss, whose average is the cross-entropy.-LSTM/PTB: context = previous words, prediction = next word, and cross-entropy = main (next-word prediction) loss.-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 interchangeable, as their difference doesn't depend on Q.• True context distribution π and true context-conditional model P describe the unknown probability space. We assume tokenized data (x t , y t ) generated according to π(x t ) and P (y t |x t ).-The goal of learning Q is to approximate P .• Empirical context distribution π n and empirical context-conditional model P n 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), e.g., Eq. ( 5) → Eq. ( 9). • Induced model: specializes full context-conditional model (y t |x t ) to short context only (y t |x t ), under a specific context distribution.-If the full context-conditional is P and the context distribution is π, we get the true induced model P . This is the ideal, or Bayes' optimal small model (which is not available in practice). In practice, we approximate it by a target induced model, P (e.g., Kneser-Ney bigram).-If the full context-conditional is Q, and the context distribution is... * ... π, we get the learned induced model Q. This is the best way Q can predict based only on the short context. However, we cannot evaluate it without π. Instead, we use... * ... π n , and get the empirical learned induced model Q. This, we can evaluate and use in our training.In summary, the experiments use Q given by Eq. ( 6), efficiently approximated via sampling in Eq. ( 18). This is then plugged into the empirical IMM objective in Eq. ( 9), which pits Q against P .In Xie et al. (2017) , noising was proposed as an approach to perform data augmentation in language and sequence-to-sequence 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, we consider only a bigram model, with counts c(x, y), where x is the context and y is the prediction. The general noising scheme in Xie et al. (2017) 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 to show that, as a result, the expected counts presented to the learner are no longer c(x, y), but rather change to c(x, y) as follows.(a) If only contexts are noised, then c(x, y) =[1γ(x)]c(x, y) + q(x)x ′ γ(x ′ )c(x ′ , y) (b) If only the predictions are noised, then c(x, y) = [1 -γ(x)]c(x, y) + q(y)γ(x)c(x)(c) Lastly if both predictions and targets are noised, independently, then c(x, y) =[1γ(x)]c(x, y) + q(y)q(x)x ′ ,y ′ γ(x ′ )c(x ′ , y ′ )These noising schemes are primarily motivated (see Xie et al. (2017) , Section 3.3) through the idea that noising leads to classical forms of smoothing, and which in turn may be a desirable property to inject into training more sophisticated language models. This is indeed true in case (a) when γ(x ′ ) = λ is constant and q(•) = c(•)/n is the unigram distribution, leading to simple interpolative smoothing. It immediately fails to be true when q is any other distribution even if γ(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 γ(x ′ ) and q(•) are determined via the missing mass and Kneser-Ney backoff, failing to recreate any likeness of the Kneser-Ney smoothed bigram. The situation is at least as bad in case (c).The only instance that succeeds to 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 γ(x ′ ) 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 γ(x ′ ) = γ 0 N 1+ (x -1 , •)/c(x) with γ 0 = 0.2 or 0.6 (see Xie et al. (2017) , Table 1 and Figures 1 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 ν(x) is the true missing mass in the Kneser-Ney model, then we can understand this choice as γ(x) = λν(x). As a result, we have:EQUATIONUpon close examination, we identify this as an interpolation between the data on the left and the Kneser-Ney bigram on 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.1Assume that we are optimizing the idealized noising objective of Eq. ( 14) -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 -i.e., we are in the realizable case. Then, there exists a choice of π and P such that Q ⋆ ̸ = P .Proof. Let us first rewrite ( 14) below, with the addition of constants (an entropy to the first term and P (y|x) in the logarithm of te second term) that do not change the minimizer of this objective:D(P ∥Q) f (Q) +λ x π(x) y P (y|x) log P (y|x) Q(y|x) g(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 f and g are convex in Q and f ≥ 0. Since f (Q) is minimized at P , its gradient should be 0 atP . g(Q) is 0 at P . If λg(Q † ) is -ϵ < 0 at some Q † , then along the line connecting P to Q † , λg(Q) is below the line 0 to -ϵ.To compensate, f would need to be above the line 0 to +ϵ, which would violate the fact that the gradient of f is 0 at P .We now numerically show that g(Q) can indeed be negative at some Q † , for a given construction with specific choices of π and P .Consider a trigram scenario, with x = (x -1 , x -2 ) and where we understand the short context as x = x -1 and the extended context as x = x -2 . Let us rewrite the regularization term g(Q) explicitly splitting the short and long contexts:g(Q) = x-1,x-2 π(x -1 , x -2 ) y P (y|x -1 ) log P (y|x -1 , x -2 ) Q(y|x -1 , x -2 ) (17)We search for distributions over the prediction y, short context x, and extended context x. Consider the following resulting tensor:P (y, x -1 , x -2 ) =     0.396 0.003 0.1 0.05 0.004 0.297 0.1 0.05    Based on this, we get the context distribution π(x, x), the conditional P (y | x -1 , x -2 ), as well as the restricted model P π(x -1 , x -2 ) = 0.4 0.3 0.2 0.1P (y | x -1 , x -2 ) =     0.99 0.01 0.5 0.5 0.01 0.99 0.5 0.5    These give us the following induced model:P (y | x -1 ) = 0.57 0.5 0.43 0.5Now consider the following choice of Q † :Q † (y | x -1 , x -2 ) =     0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5    Plugging the above choice of π, P , P (y | x -1 , x -2 ) and Q into Eq. ( 17) gives us a negative value of g(Q † ) = -1.1, thus completing the construction of the counterexample and the proof. Experimental Evidence of Inconsistency One may dismiss Proposition 6.1 as being too specific of a counterexample. After all, in Figure 1 , 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.1: 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 critical in the reverse knowledge-distillation literature (see for example Sec. 3 of Qin 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 6 in Appendix D.1.1).To experimentally verify this, we ran the logistic regression example with fixed λ = 1.5 (optimal at data size 5). The results are below. IMM maintains performance comparable to Figure 1 , whereas noising experiences a widening gap, and soon underperforms even the baseline.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 not directly apply here, however, a common insight may be that the process of inducing a model can also be thought of as ensembling, but over contexts rather than over independent resampling. Otherwise, considering that the implicit noising objective that we identify in Eq. ( 17) is equivalent to the explicit objective of reverse knowledge distillation, the arguments in this section equally apply as caveats for reverse-KD when the weak teacher is a restricted-context model.Algorithm 1 Sampled IMM with SGD for a Model Q with parameters WInput: Tokenized data (x t , y t ) for t = 1, 2, ..., n, k = sampling rate Output: IMM-trained model Q 1: repeat 2: Q(y t |x t ) ← FEEDFORWARD(Q, x t , x t ) 3: ∇ W Cross-Entropy(Q) ← BACKPROPAGATE(Q, Cross-Entropy(Q)) 4: Q(y t |x t ) ← 0 5:for all x ′ ∈ Sample(extend(x t ), k) do 6:Q(y t |x t ) ← Q(y t |x t ) + FEEDFORWARD(Q, x t , x ′ ) 7:end for 8:IMM t (Q) = -y P (y|x t ) log Q(y|x t ) 9: ∇ W IMM t (Q) ← BACKPROPAGATE(Q, IMM t (Q)) 10: APPLYGRADIENTS(∇ W Cross-Entropy(Q) + λ∇ W IMM t (Q)) 11: until convergence C.1 k-approximated IMMFor the approximation, a given short context x t , we take only a fixed number k of samples X 1 , • • • , X k uniformly from extend(x t ). The approximated Q (Eq. ( 11)) can then be written asEQUATIONIndeed this is what is happening on lines 5-8 of Algorithm 1. Sampling k 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. ( 8), because it requires differentiating Eq. ( 7), where a naïve implementation would need all k instances of the LLM used to estimate Q(y|x) in the memory simultaneously. Fortunately, we provide a solution to this, explain in detail in Section C.2 below.Eq. ( 7) is a cross-entropy and the Q(y|x) 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 to approximate Q(y|x). In Algorithm 2 Serialized IMM with SGD for a Model Q with parameters W Input: Tokenized data (x t , y t ) for t = 1, 2, ..., n, r = refresh frequency Output:IMM-trained model Q 1: repeat 2: Q(y t |x t ) ← FEEDFORWARD(Q, x t , x t ) 3: ∇ W Cross-Entropy(Q) ← BACKPROPAGATE(Q, Cross-Entropy(Q)) 4:if Repeats % r = 0 then 5:Q † ← Q 6: Q † ← 0 7:for all t = 1, 2, ..., n do 8:for all y do 9:Q † (y|x t ) ← Q † (y|x t ) + FEEDFORWARD(Q † , x t ){Or use other induction mechanism, e.g., density estimator for logistic regression.} 10:end for 11:end for12: NOGRADIENT(Q † , Q † ) 13:end if 14:IMM t (Q) := -y Q † (y|xt) Q † (y|xt) P (y|x t ) log Q(y|x t ){This is not the IMM loss, but its gradient averages to the correct gradient, see Section 5.} 15:∇ W IMM t (Q) ← BACKPROPAGATE(Q, IMM t (Q)) 16: APPLYGRADIENTS(∇ W Cross-Entropy(Q) + λ∇ W IMM t (Q))17: until convergence our experiments, it 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 k the computation of the derivative of a log where the argument of the log is an average (or a sum) of k entities. This is expected, because log(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 a 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(f + g) can be represented as a weighted average of the derivatives of log(f ) and log(g), where the "crosstalk" weights do not require differentiation:∇ log(f + g) = ∇f + ∇g f + g = f f +g ∇ log f + g f +g ∇ log g.In summary, this allows us to compute the derivatives separately for each random sample and accumulate them appropriately.Recall that we can write an IMM component at data point t (Eq. ( 12)) with k randomly sampled extended contexts as:EQUATIONwhere a comes from Eq. ( 18).The above entity cannot be directly decomposed into k 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 k terms.Denoting the set of parameters using W , we can then write ∇ W IMM(Q) as below.EQUATIONWe notationally move the derivative outside, but it is crucial that C t,i (y) = Q(y|xt,xi) k i=1 Q(y|xt,xi) be treated as a constant with respect to the parameters of the LLM. We do this to show that we can effectively compute individual cross-entropies, with the C t,i (y) terms acting like a "crosstalk" between the k random samples. This exactly parallels the decomposition that we highlighted above for the case of ∇ log(f + g). Note that this decomposition is not possible directly on the cross-entropy (that involves the log of a sum), but possible when we consider its gradient.We cannot compute the above total gradient, i.e. ∇ W IMM(Q) unless for every t and y we have C t,i (y) for i ∈ [1, k]. To this end, we implement this by partially running the feedforward computation graphs of all k random samples only until the point which gives us Q(y 0 |x t , x i ) for i ∈ [1, k]. Once we have all the Q(y|x t , x i ), we can normalize all of them by their element-wise sum to get the weights C t,i (y). These weights are then fed back into the graph, and the forward pass is completed to get a "tampered" cross-entropy for i ∈ [1, k]. 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 ∇ W IMM(Q).We now provide some additional detail of both the logistic regression and language modeling experiments. One minor difference between the two is that in the logistic regression case we interpolate the main loss (1λ) and the IMM risk (λ), to efficiently cross-validate for the choice of λ as we explain in Section D.1.1. In the language modeling experiments, we simply add the regularizer with the factor λ, while maintaining the main loss as is (factor 1).We provide the complete details of the logistic regression toy example2 . Consider a dataset with feature space x = (x 1 , x 2 , x 3 ) ∈ R 3 containing two linearly separable classes, defined by a linear discriminant g(x) = ax 1 + bx 2 + cx 3 + d. We will think of features as decomposable into x, the short context of restricted features, and x, the extended context of remaining features. For this example, let x = x 1 and x = (x 2 , x 3 ). The choice of word "context" stems from features in language modeling.Target restricted model Say features are sampled uniformly over a cube and that we have ample observation of the labels y = 1{g(x) > 0} along only the short context, i.e., ample data points of the form (x 1 , y). This allows us to build an excellent restricted model to predict the label based on just x 1 , call it P (y|x 1 ) or P (y|x). For the sake of this example, let us assume that this is the (restricted) Bayes predictor or true conditional probability.We then observe a few samples of all features along with labels, i.e., data points of the form (x 1 , x 2 , x 3 , y). How can we use the restricted model P (y|x 1 ) as part of the training of the full model Q(y|x 1 , x 2 , x 3 )?Induced model Say we have n data points. Let us index them by t = 1, • • • , n, written either as (x t , y t ) or (x 1,t , x 2,t , x 3,t , y t ). The key concept of the method is that it is not Q's behavior that should target that of P (y|x 1 ), 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 x t is drawn from its conditional distribution given x t . 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 x for a given x. In logistic regression, this is not viable since x 1 does not repeat. We instead use a soft nearest-neighbor density estimate f (x 2 , x 3 |x 1 ) ∝ n t=1 δ x2,t,x3,t (x 2 , x 3 )e -α|x1,t-x1| , where 1/α is the bandwidth of the kernel. (With cross-validation, we determine α = 1 to be a good choice in this example.) If we let w t = e -α|x1,t-x1| , the resulting induced model by marginalization is:EQUATIONEQUATIONJust like in language modeling, to induce a model we need to to draw x's from its conditional distribution given x. Unlike in the discrete case where we could directly approximate the conditional distribution using the empirical counts, we need to rely here on an estimated density.Matching the target The main objective of logistic regression is to minimize n t=1 log Q(y t |x 1,t , x 2,t , x 3,t ), which is equivalent to Cross-Entropy(Q) versus the empirical 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 a KL-divergence, and introduced as a regularizer. The result is a secondary objective, which we call the IMM risk, expressed as:EQUATIONThe overall objective then becomesEQUATION)where λ is the regularization trade-off and can be determined via cross-validation (details in the Appendix D.1.1).IMM improves restricted-context prediction In Figure 4 , for n = 10 and 30 Monte-Carlo runs, we show how the induced model of Q (i.e. Q) fairs in its prediction, compared to the target model P . Note that without IMM (i.e. with λ = 0), the secondary objective is large and thus Q's performance is worse. With increasing λ, IMM improves this performance both in average and in variance. We deduce that there is information in the accurate restricted model P (y|x) that is not known to Q, unless we explicitly incorporate it. This shows that Q gets better at predicting from a restricted context, naturally bringing up the question: does it also get better at predicting from the full context, i.e., the main objective? 1+λ ratio (refer to Appendix D.1.1). We stop at λ 1+λ = 0.9 because λ 1+λ = 1 would zero out the contribution of the main objective (and replacing the main objective completely is never the intention).In Figure 1 , we compare the performance of IMM-trained Q (green) to that without IMM (maroon). We sweep a range of n from 2 to 50, and use a crossvalidation optimized λ for each (details in the Appendix D.1.1). 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 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.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 x 1 -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 ± agrees with these weights at each x 1 . Intuitively, this biases the separating plane to have the right inclination/alignment with the x 1 -axis, which subsequently speeds up the learning process.As mentioned in Section 7.1, IMM always improves baseline performance. The main parameters to choose are α and λ. For α, we performed a few experiments on held-out data and determined the value of 1 to be generally adequate. For λ, we saw that adapting to the dataset size made a more significant difference. This makes sense, as for smaller datasets, we want more IMM contribution to compensate for the lack of data (i.e., we want a larger λ).To determine a good schedule for λ as a function of dataset size n, we sweep the ratio λ 1+λ (which represents the IMM coefficient as a 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 a minimum of 2 to a 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 our plots in Figure 6 , we select λ/(1 + λ) of 0.8, 0.7, 0.6 for dataset sizes, n of 2, 10 and 20 respectively. These suggest a linearly decaying optimal λ/(1 + λ), in this range of dataset sizes. For this example, doing an interpolation fit, we create our automatic λ schedule rule to be λ 1+λ = -0.0111n + 0.818.Code is available at https://github.com/uicdice/imm-language-modeling 2017), 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 in the training set (full word history) that share that short context (same previous word). In the language of the paper, this is exactly the induced bigram of the learned LSTM. We report the 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.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.We consider an 11 × 11 toroidal grid in the (x, y) plane with a predefined reward landscape that has a single peak at the center of the grid, see the heatmap in Figure 7 . 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. We model this as both a 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 y dimension. We first solve the POMDP using the POMDPs.jl (Egorov et al., 2017) package package and the Fast Informed Bound (FIB) Solver (Kochenderfer et al., 2022) to obtain a very accurate policy based on partial information. (These POMDP policies are represented in the form of alpha vectors.)MDP Training without IMM Our MDP training algorithm is REINFORCE (Williams, 1992) . REINFORCE, at every epoch, samples a set of observations, which can also be called a batch. In REINFORCE, at every epoch we 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 x (the known dimension of the POMDP agent). Unlike language modeling experiments (where we randomly sample for the extended context), we use the POMDP belief. Here, we thus use a 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 a 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-augmented REINFORCE3 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 at 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 are reported in Figure 2 .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 have a 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 is 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 8 . For RL, this is illustrated in Figure 9 . 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.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.to noising. However, P addresses a markedly different task than Q, i.e., that of predicting with a 38th Conference on Neural Information Processing Systems (NeurIPS 2024).This requires space O(n log n) only, because each key in the dictionary is a short context, and each value is 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 us n and each pointer requires O(log n) bits.Code is available at https://github.com/uicdice/imm-logistic-regressionCode is available at https://github.com/uicdice/imm-reinforce 


Decision: 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.

Reviews: review: paper_summary: paper_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...)$. 
\
Instead 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
\
$$\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)}$$
$$\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 } \\ 
\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)
$$
- 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. 
- 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. 
- Empirically, IMM always does better than with no-IMM, but the performance difference is especially large in data-limited regimes.
strength_weakness: strength_weakness: - 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. 
- It would also be nice to include some settings where the feature-restricted model is not close to optimal as an ablation study. 
- 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?
questions: 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: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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.
* It would be interesting to include further analysis regarding how the restricted model’s performance affects this method.
questions: 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$? 
* 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. 
* Just for my understanding, are you assuming a fixed-context length model?
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

If 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).
questions: 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)?

2. 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?
2.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?

3. 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.

4. 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: limitations: Adequately discussed.
paper_summary: paper_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
strength_weakness: strength_weakness: - 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.

- 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.
- 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.
- 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.
- 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.
questions: 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.

1. 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.

2. 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?

3. 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.

4. 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: limitations: Yes

Comments: title: Did we address your points?
Thank you!
Thank you so much! + Update on a computationally efficient version.
Thank you very much!
An important update on this computationally efficient version.
Thanks!
Response
Answer (2/2)
Answer (1/2)
Thanks, one more question
comment: Dear Reviewer,

Thank you again for taking the time to review our paper. We understand that life is busy and that reviewing can be hectic. We believe that we addressed all the points that you raised.

Additionally, we would like to bring your attention to the fact that (in the general rebuttal) we proposed a formulation for making IMM computationally efficient, and that we just finished implementing it successfully for the logistic regression experiment. We hope that you will also take that into consideration when making your final recommendation.

If we did address all your points, we very much hope that you will raise your score to reflect it. It would helps us get closer to sharing this work with the community and we would appreciate that greatly!

Thank you again.

Sincerely,
Authors
Dear Reviewer,

Thank you very much for being very vocal about the qualities of our paper. We remain ever appreciative! We're sorry we didn't reply sooner. We were working on a computationally efficient version of IMM. We finally managed to make the formulation that we proposed in the general rebuttal work for the logistic regression example. We report on this in a comment to the general rebuttal. We hope that you will read it.

Regarding your suggestion about timely larger-scale applications, we understand that in this day and age, it feels like any ML research that isn't directly tested in the latest language, vision, or robotics applications is not as applicable. However, for small academic research groups without an existing pipeline into these disciplines, especially PhD students working just with their advisors, it is often not feasible to implement such large-scale tests. Instead, we rely on smaller-scale problems that still have the essence of those larger problems, where success strongly suggests potential success if scaled. An analogy here would be a mechanical engineering researcher proposing a design change that could make combustion engines more efficient. We would likely all agree that this researcher should not be expected to build a whole engine and a car in order to demonstrate that their design idea is valid.

Perhaps, however, you are not suggesting that we actually build a large language model with IMM, or anything of that scale, and that you rather only want us to evoke these potential applications. If that is what you mean, we have in fact several ideas where IMM could be very impactful. Here are two: 

1. **Infusing Models with Physics**: We already suggested this to a fellow reviewer.  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.
2. **Augmenting Sensing**: This scenario was partly behind the RL example in the paper. Imagine an autonomous driving setting where a model is built using a certain set of sensors (e.g., depth camera and ultrasound sensors). The policy may have been trained using a lot of expensively collected imitation learning data, and it performs driving very well with just these sensors. What if we decide to add additional sensors (e.g., Lidar)? Do we need to re-collect all that driving data now with the new sensors on? Or can we use the previous model/policy as we train a new one that _also_ uses the Lidar? We hope that it is evident that this is a problem of a feature-restricted model informing a fully-featured one, and that thus IMM can play a big role in enabling it and saving the cost of collecting brand new data.

(In a previous revision of this paper, we had these examples, but had to remove them due to space. In our conclusion, we also mention the potential usefulness of IMM as a means to extend the context window of LLMs.)

We hope that these problems are thought-provoking and enough to convince you that the usefulness of IMM is not just in a few lines of math and basic experiments, and that there is merit to those ideas beyond the pages of the paper, especially to those who can more readily perform wider experiments and implementations. However, unless the paper is shared with the community, those ideas will remain not as publicized. We hope that you will help us share it.

Thank you again for everything.

Sincerely,
Authors
Dear Reviewer,

We are humbled by and very appreciative of your decision to move our work firmly into the accept zone. Thank you so much! Computation is not an afterthought for us. We will itemize both time and memory complexity in detail in the paper, just as we did in our rebuttal.

More importantly, to show our dedication to making IMM computationally efficient, we were working very hard on implementing the computationally efficient version that we proposed in the general rebuttal. We had to surmount several technical hurdles, but we were finally successful in applying it to the logistic regression example. In a comment to the general rebuttal, we report on these results. The new method is only 50-70% slower than the baseline, but it continues to have an edge on noising and tracks the version of the paper very closely, especially as the dataset size increases. We also highlight how the same steps could be taken in the language modeling examples.

We don't claim to have fully solved the problem, but we now have both a formulation for computationally efficient IMM and a proof of concept. We plan to describe these in the current paper and to add it to our published code. Thank you very much for pushing us to do this. The case for IMM is so much stronger as a result.

You've been very generous with your revised opinion and we hope that you will continue to support the paper. (We have revised this paper a few times, and it would be great if our final revision is for the camera-ready version at this year's NeurIPS.)

Thank you again for everything.

Sincerely,
Authors
Dear Reviewer,

We thank you very much for approving of the new notation. We will definitely use it in the paper, since it's so much more streamlined. We also have a few ideas on how to improve the visualization of the toy example. In particular, we can display the Bayes restricted model and a model-in-training along with its induced model, and show how IMM encourages one to tend toward the other (the alignment that we mentioned).

We appreciate a lot the additional point to your score! We're sorry for not replying sooner. We were spending a lot of time researching a computational efficient variant of IMM. We were finally able to implement it for the toy example, and our results are summarized in a comment to the general rebuttal. Please feel free to consult it as you make your final recommendations.

Sincerely,
Authors
We have been hard at work at researching the computationally efficient implementation of the approach suggested above. Let's call it *no-sampling IMM*, to contrast it with the paper's *sampled IMM*. We were finally able to implement this approach in the logistic regression example in a satisfactory way, by addressing the two caveats mentioned above:

+ **High-variance**: The gradient of *no-sampling IMM* does not have the $k$-sample averaging and has thus higher variance. To compensate for this, we experimented with SGD with momentum. This helped most of the time, but unfortunately also frequently caused instability. We traced this instability to two sources: staleness in the correction factor, which we address in the second point, and exploding gradients. We controlled the latter using gradient-clipping. When using momentum, since the learning rate effectively increases, the optimization also sometimes got stuck in plateaus. We solved this by using learning-rate decay schedules. 

+ **Maintaining an induced model**: In order not to compute the full induced model at every iteration (which would bring back the cost that we are trying to eliminate), we experimented with updating the model every $n$ iterations. By doing so, even if we recompute a full induced model using $\mathcal{O}(n^2)$ computation (which is the extreme of $k=n$, using _all_ samples in the density estimator, see Sec. 6.1), it will only cost $\mathcal{O}(n)$ additional computation, which is the complexity of the baseline. However, this causes the estimate of $\overline{Q}$ to become stale (since $Q$ gets updated between recomputations of $\overline{Q}$). Say model $Q_\dagger$ was used to calculate $\overline{Q_\dagger}$. We found that if we use the updated/fresh $Q$ for the correction factor, as in $\frac{Q}{\overline{Q_\dagger}}$, then it causes instability in the gradient, likely due to these correction factors not obeying expected constraints, e.g., they should have mean $1$ for every $y$, but they won't if one is fresh the other stale. This reasoning gave us the following solution: we found that if we use the stale model instead, i.e., if we use $\frac{Q_\dagger}{\overline{Q_\dagger}}$ for the correction factor, then the instability disappears.

Here is a comparison of the average performance (over 20 runs) of *no-sampling IMM* and *sampled IMM*. The latter is as in the paper. The former uses the following hyperparameters: $\lambda$ is tuned the same way as for the sampled version, $\overline{Q}$ is updated every $n$ samples, momentum is set at 0.85, clipping is done with the inf-norm at threshold 2, a plateau-based learning rate scheduler attenuates learning by 0.9 at each plateau, the max number of epochs is set to 20,000, with early stopping based on convergence.

This computationally efficient version continues to have the edge on noising. The match between the two versions improves with more data. We hypothesize that this is due to the fact that *no-sampling IMM* approximates the ideal IMM loss directly empirically, whereas *sampled IMM* does this through estimating $\overline Q$. In terms of clock time, *no-Sampling IMM* is considerably faster than *sampled IMM* and comparable to the baseline (we clocked it to take 50-70% more time at most than the baseline. (e.g., on a multicore node with a V100 GPU, a typical run for the *no-sampling IMM* vs. baseline is 4.5 vs. 2.8 seconds at size 10, and 20 vs. 12 seconds at size 50.)

| Dataset Size |   Baseline           | *sampled IMM*            | *no-sampling IMM*                 |
|--------------|----------------------|---------------------|---------------------|
| 5          | 73.14  | 79.96 | 77.78  |
| 10         | 84.17  | 88.65 | 86.42  |
| 15         | 89.86  | 92.52 | 91.02  |
| 20         | 92.35  | 94.30 | 95.27  |
| 30         | 94.94  | 95.69 | 95.65  |
| 40         | 96.33  | 96.59 | 96.52  |

Finally, although we did not implement this in the language model case yet, we are confident that the observations we made will carry over (momentum, gradient clipping, learning rate schedules, intermittently updating $\overline{Q}$, using the stale model). Notably, to induce a bigram, one only needs to do a single pass through the data. This would give $\overline{Q_\dagger}$ . As for $Q_\dagger$, the most space-efficient way to access it is to keep a copy of that model in memory (the logits over the entire dataset would require more space to store.)

---

Our paper was concerned primarily with the statistical aspects of the problem. We are very thankful to all of you for pushing us to think harder about making IMM also computationally attractive. We don't claim to have completely addressed it, however we **(1)** gave a high-level  formulation to ease the computation, and **(2)** demonstrated its validity in an example from the paper (albeit a toy one). We hope that you will take this earnest effort into consideration for your final recommendations. Thank you.
Thank you for the thorough clarification! My questions have been answered. I leave my final score as a 7, after consideration. The paper is well written and proposes interesting ideas and algorithms. In the future, if authors can demonstrate the usefulness of this method on more timely larger-scale applications, I think that could make the paper even stronger.
Thanks for the detailed response. 
After reading the rebuttal and other rebuttals and reviews, I am still a bit skeptical of the computational feasibility of the approach. Slowing down training by a factor of 5-10 is quite a sizeable decrease, and this should be addressed in the main text and not glossed over. That said, I appreciate the authors' willingness to engage with the topic in the rebuttal and especially to identify alternative, equivalent, objectives that are more tractable (such as in the main response). 

As the authors point out, I understand the focus is on the relatively small-data regime, where IMM does seem to show some promise. I will raise my recommendation, although I urge the reviewers to try their best to make the IMM approach more computationally tractable for their next revision if they want it to be adopted more widely.
I thank authors for their clarifications and I'm raising my score for 1 point. I hope that the improved notation and the toy task figure (which I would still improve) will make it into the future version of the paper.
> [...] but my understanding is that IMM is much more computationally expensive.

We address the computational aspect of IMM in the general rebuttal. In essence, the alternative that we propose has the potential to make IMM computationally equivalent to noising.


> I wonder whether there are circumstances where IMM has any gains over cheaper methods trained on a little bit more data [...]

Our general premise is that we have what data we have, and no more. Noising and IMM both have this data and the same feature-restricted) target model. (The sampling to estimate $\hat{Q}$ is done with the same data set.) With these exact same information resources, IMM always statistically outperforms noising. If the goal is to make the most out of what information we have, or when acquiring that little bit more data is expensive, IMM is the way to go.

> [...] or whether IMM generally has any utility in any settings with reasonable amounts of data.

This question can be asked of any variant of data augmentation, as given enough data there is no need to augment it. The question is: in what regime is any such technique useful? In Section 6.1, we mention that IMM appears most successful when data size is comparable to the number of parameters. This agrees with the general rule of thumb of regularization benefiting high-dimensional regimes. The intuition is that the cross-entropy loss does not sufficiently localize the model, so we get better localization when IMM constrains the induced model, which is a low-dimensional projection. Since most modern models (even LLMs) operate in such regimes, it is not a stretch to expect IMM to be widely beneficial. (Otherwise, it would be also hard to imagine the many techniques suggesting using N-grams to augment LLMs having any utility either.)

---

We are convinced that the new perspective that IMM provides, the deeper understanding of noising and reverse knowledge-distillation that it offers, and the consistent statistical edge that it has with limited data, make it worthwhile to share with the community at NeurIPS. Once the idea is out there, we are certain that our effort and that of others can inevitably make IMM more computationally efficient and scalable.

You have been very generous with your assessment of our work. If you believe 7 is the right score, we humbly thank you and only hope that your position can help convince your colleagues to also support us at the same level.
We misunderstood your question to refer to the only statistical advantage that noising has, which is smaller variance at small data sizes. We sincerely apologize for this. We now address all your related questions, point by point.

>  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?  / [...] my question was why the gap between Noising and IMM closes so quickly. By 5+% dataset size, the performance gap seems to have closed to around 3% [...]

The gap narrows in Figure 1 because we are being very favorable to noising. Specifically, we are decaying $\lambda$ (the amount of noising) optimally with increasing data. This is _necessary_ for noising. The reason for this is the same as what occurs in Proposition 5.1: 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 $\lambda$ is also acknowledged as critical in the reverse knowledge-distillation literature (see for example Sec. 3 of Qin et al. 2021). However, tuning $\lambda$ is _optional_ for IMM, thanks to $\hat Q$ (with more extended history samples) accurately tracking the target (see the flat curves on the right column of Figure 4 in Appendix D1.1). 

To fully convince you, we re-ran the experiment with fixed $\lambda=1.5$ (optimal at data size 5). The results are below. IMM maintains performance comparable to Figure 1, whereas noising experiences a widening gap, and soon underperforms even the baseline.

| Dataset Size |   Baseline           | Noising             | IMM                 | IMM-Noising Gap |
|--------------|----------------------|---------------------|---------------------|-----------------|
| 5          | 73.14 +16.17/-14.53  | 76.85 +18.58/-13.15 | 79.96 +14.99/-12.38 | 3.11            |
| 10         | 84.17 +13.51/-10.86  | 84.37 +7.40/-7.63   | 88.65 +9.65/-7.68   | 4.28            |
| 15         | 89.86 +8.86/-6.80    | 86.17 +6.17/-5.53   | 92.52 +6.19/-4.81   | 6.35            |
| 20         | 92.35 +7.35/-5.32    | 86.99 +5.69/-5.37   | 94.30 +4.63/-3.70   | 7.30            |
| 30         | 94.94 +3.97/-3.43    | 88.68 +4.35/-4.32   | 95.69 +3.39/-2.64   | 7.01            |
| 40         | 96.33 +3.00/-2.33    | 89.70 +3.73/-3.97   | 96.59 +2.59/-2.08   | 6.89            |
| 50         | 97.14 +2.47/-2.20    | 90.78 +3.78/-3.55   | 97.34 +2.34/-1.99   | 6.56            |

> Generally, what is the importance of using more samples when marginalizing Q for improved performance? Does performance strictly increase with increasing samples?

The above answer gives one clear benefit of a more accurate $\hat Q$: less sensitivity to tuning $\lambda$. However, since tuning $\lambda$ is possible, the primary advantage is significant and consistent gains, even against optimally-tuned noising. The gains in Figure 1 may seem small in terms of absolute percentages, but they are significant, as the IMM-noising gap is larger than the noising-baseline gap, most of the time and often considerably so. We can quantify these gains in terms of data size increase: e.g., the accuracy of IMM at 15, would need augmenting the data set size to 18 with tuned noising (a 20% increase) or to 20 without noising (a 33% increase). These gains are also more consistent, as the variance of IMM's accuracy is smaller than the variance of noising's accuracy, most of the time (except, as noted, for small data sizes).

IMM strictly improves when we use more extended history samples, but with diminishing returns, allowing us to cap to $k=10$ samples at most.
Thank you for the clarification! I think there was a misunderstanding regarding my question about Noising versus IMM. 

The authors provide intuition for why Noising does better than IMM at small dataset sizes, but rather my question was why the gap between Noising and IMM closes so quickly. By 5+% dataset size, the performance gap seems to have closed to around 3%, but my understanding is that IMM is much more computationally expensive. I wonder whether there are circumstances where IMM has any gains over cheaper methods trained on a little bit more data or whether IMM generally has any utility in any settings with reasonable amounts of data. I'm also not certain how the author's explanation aligns with Figure 1 observations. Noising seems to strictly be doing worse than IMM for any dataset size smaller than 5% and there seems to be no dataset size where IMM does worse than Noising. 

I will most likely keep my score as 7, if the authors would like to devote more time to answering other reviewers' questions.