Title: Induced Model Matching: Restricted Models Help Train Full-Featured Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: restricted

model

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

Induced Model Matching:

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.

1 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 ) = ∑ x 2 ,x 3 π ( 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 6 to noising. However, ˆ P addresses a markedly different task than Q , i.e., that of predicting with a

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.

0

10

20

30

40

50

50

60

70

80

90

100

IMM with Target Model

ˆ

P

ˆ

P

Interpolation with Target Model

ˆ

P

No IMM, Noising or Interpolation

Dataset Size

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:

In Section 3, we rigorously frame this question through the notion of an induced model .

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.

In Section 5, we propose computational approaches to make IMM practical and pave the way for further scaling IMM.

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.

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.

2 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,

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

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

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:

where 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:

We 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:

Let's first idealize and assume availability of P and the context distribution π , required to compute Q . Equipped with Q , we can introduce the idealized induced model matching (IMM) risk, a secondary risk that is the average context-conditional KL divergence with the restricted context:

Since 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:

The 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 :

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

IMMas 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 minimizing

Separating 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:

The simplest version of ˆ P could be based on empirical counts, which may be valid if the context space is discrete and small. In language modeling, smoothing methods can 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 as

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.

5 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 context 1 , and sweeping only across those. This amounts to rewriting Eq. (6) as follows:

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:

where 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 in 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:

The gradient of this IMM risk then becomes:

The empirical version of this gradient is:

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-Nikodym

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.

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 y t ′ has no role in calculating the loss. Instead, the prediction is randomized through the sum over all possible y 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:

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:

A 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

7.1 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 δ x 2 ,t ,x 3 ,t ( x 2 , x 3 )e -α | x 1 ,t -x 1 | , 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 -α | x 1 ,t -x 1 | , the resulting induced model by marginalization is:

These equations are respectively equivalent to Eqs. (4) and (6). The IMM risk and the corresponding overall objective remain the same:

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

̸

7.2 Language Modeling Experiments

In these language modeling experiments, the restricted model we use is the modified Kneser-Ney bigram (Kneser and Ney, 1995; Chen and Goodman, 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.

7.3 Reinforcement Learning: POMDPs helping MDPs

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.

0

100

200

300

400

500

Number of epochs

0

5

10

15

20

Average Reward

IMM using maximal utility action

No IMM

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

Acknowledgments and Disclosure of Funding

This paper is based upon work supported in part by the National Science Foundation, through the NSF CAREER Program under Award No. CCF-2146334 (From Rare Events to Competitive Learning Algorithms), the NSF HDR TRIPODS Phase II Program under Award No. ECCS-2217023 (IDEAL Institute), and the NSF TRIPODS Phase I Program under Award No. CCF-1934915 (UIC Foundations of Data Science Institute). Computational infrastructure was supported in part by the NSF MRI Program Award No. CNS-1828265 (COMPaaS DLV).

References

Bahdanau, D., Cho, K., and Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473 .

Chelba, C., Norouzi, M., and Bengio, S. (2017). N-gram language modeling using recurrent neural network estimation. arXiv preprint arXiv:1703.10724 .

Chen, S. F. and Goodman, J. (1999). An empirical study of smoothing techniques for language modeling. Computer Speech & Language , 13(4):359-394.

Dagan, I., Glickman, O., and Magnini, B. (2005). The pascal recognising textual entailment challenge. In Machine Learning Challenges Workshop , pages 177-190. Springer.

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 .

Dolan, W. B. and Brockett, C. (2005). Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005) .

Egorov, M., Sunberg, Z. N., Balaban, E., Wheeler, T. A., Gupta, J. K., and Kochenderfer, M. J. (2017). POMDPs.jl: A framework for sequential decision making under uncertainty. Journal of Machine Learning Research , 18(26):1-5.

Gao, F., Zhu, J., Wu, L., Xia, Y ., Qin, T., Cheng, X., Zhou, W., and Liu, T.-Y . (2019). Soft contextual data augmentation for neural machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 5539-5544.

Hinton, G., Vinyals, O., Dean, J., et al. (2015). Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2(7).

Kaplun, G., Malach, E., Nakkiran, P., and Shalev-Shwartz, S. (2022). Knowledge distillation: Bad models can be good role models. Advances in Neural Information Processing Systems , 35:28683-28694.

Kneser, R. and Ney, H. (1995). Improved backing-off for m-gram language modeling. In 1995 International Conference on Acoustics, Speech, and Signal Processing , volume 1, pages 181-184. IEEE.

Kochenderfer, M. J., Wheeler, T. A., and Wray, K. H. (2022). Algorithms for decision making . MIT press.

Lee, H., Hou, R., Kim, J., Liang, D., Hwang, S. J., and Min, A. (2023). A study on knowledge distillation from weak teacher for scaling up pre-trained language models. arXiv preprint arXiv:2305.18239 .

Li, H., Cai, D., Xu, J., and Watanabe, T. (2022). Residual learning of neural text generation with n -gram language model. arXiv preprint arXiv:2210.14431 .

Liu, J., Min, S., Zettlemoyer, L., Choi, Y ., and Hajishirzi, H. (2024). Infini-gram: Scaling unbounded n-gram language models to a trillion tokens. arXiv preprint arXiv:2401.17377 .

Luong, M.-T., Pham, H., and Manning, C. D. (2015). Effective approaches to attention-based neural machine translation. arXiv preprint arXiv:1508.04025 .

Ng, N., Cho, K., and Ghassemi, M. (2020). SSMBA: Self-supervised manifold based data augmentation for improving out-of-domain robustness. arXiv preprint arXiv:2009.10195 .

Pascanu, R., Mikolov, T., and Bengio, Y. (2013). On the difficulty of training recurrent neural networks. In International conference on machine learning , pages 1310-1318. PMLR.

Qin, Y., Lin, Y., Yi, J., Zhang, J., Han, X., Zhang, Z., Su, Y ., Liu, Z., Li, P., Sun, M., et al. (2021). Knowledge inheritance for pre-trained language models. arXiv preprint arXiv:2105.13880 .

Rajpurkar, P., Zhang, J., Lopyrev, K., and Liang, P. (2016). Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250 .

Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. (2014). Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research , 15(1):1929-1958.

Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction . MIT press.

Turc, I., Chang, M.-W., Lee, K., and Toutanova, K. (2019). Well-read students learn better: The impact of student initialization on knowledge distillation. arXiv preprint arXiv:1908.08962 , 13.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems , 30.

Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. (2018). Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 .

Wang, S. I., Wang, M., Wagner, S., Liang, P., and Manning, C. D. (2013). Feature noising for log-linear structured prediction. In Proceedings of the 2013 conference on empirical methods in natural language processing , pages 1170-1179.

Warstadt, A., Singh, A., and Bowman, S. R. (2018). Neural network acceptability judgments. arXiv preprint arXiv:1805.12471 .

Wei, J. and Zou, K. (2019). Eda: Easy data augmentation techniques for boosting performance on text classification tasks. arXiv preprint arXiv:1901.11196 .

Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8:229-256.

Xie, Z., Wang, S. I., Li, J., L´ evy, D., Nie, A., Jurafsky, D., and Ng, A. Y. (2017). Data noising as smoothing in neural network language models. arXiv preprint arXiv:1703.02573 .

Yuan, L., Tay, F. E., Li, G., Wang, T., and Feng, J. (2020). Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 3903-3911.

A Explanation of main notations

Figure 3 gives a schematic overview of the process of IMM.

Induced Model Matching

Cross-Entropy

P

(

y

|

x

)

Q

(

y

|

x

)

P

(

y

|

x

)

ˆ

P

(

y

|

x

)

Marginalization

Marginalization

ˆ

Q

(

y

|

x

)

𝜋

!

empirical context

distribution

P

n

(

y

|

x

)

data

π

true context

distribution

true feature-restricted

induced model

full-featured true

predictive model

full-featured learned

predictive model

learned feature-restricted

induced model

data

proxy

Table 3 compiles our main notation. We summarize the main formulation here and explain how it connects to its instances in the experiments.

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

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

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

B Caveats of prior methods: noising and reverse-KD

B.1 A quick review and analysis of noising-based data augmentation

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 ) =

(b) If only the predictions are noised, then ˜ c ( x, y ) =

(c) Lastly if both predictions and targets are noised, independently, then ˜ 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:

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

̸

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:

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 at P . 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:

We search for distributions over the prediction y , short context x , and extended context x . Consider the following resulting tensor:

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

These give us the following induced model:

Now consider the following choice of Q † :

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.

B.3 Knowledge Distillation Theory

We would like mention that recent theory elucidates how weak teachers can be useful in reverse KD (Kaplun et al., 2022). That work is specific to the classification setting and assumes that the teachers are good samplers, i.e., noisy versions of the Bayes decision, and shows that reverse-KD can effectively ensemble multiple teachers and remove the noise. This setting and assumptions do 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.

C Implementation Details

Algorithm 1 Sampled IMM with SGD for a Model Q with parameters W . To connect with the notation of the paper, this version most closely resembles the language modeling version (with the details of the sequentialization of log-sum gradients skipped for simplicity). The logistic regression equivalent uses a soft nearest neighbor density estimate (ref D.1) instead of creating a multiset and sampling from it.

Input:

Data ( x t , y t ) for t = 1 , 2 , ..., n , k = sampling rate

Output:

IMM-trained model Q

1: repeat

2: Q ( y | x t ) ← FEEDFORWARD ( Q,x t , x t )

3: ∇ W Cross-Entropy ( Q ) ← BACKPROPAGATE ( Q, Cross-Entropy ( Q,y t ))

4: ˆ Q ( y | x t ) ← 0

5: for all i, x sample ∈ Enumerate ( Sample ( extend ( x t ) , k )) do

6: ˆ Q ( y | x t ) ← ˆ Q ( y | x t ) + FEEDFORWARD ( Q,x t , x sample )

7: end for

8: ˆ Q ( y | x t ) ← ˆ Q ( y | x t ) /k

9: IMM t ( Q ) = - ∑ y ˆ P ( y | x t ) log ˆ Q ( y | x t )

10: ∇ W IMM t ( Q ) ← BACKPROPAGATE ( Q, IMM t ( Q ))

11: APPLYGRADIENTS ( ∇ W Cross-Entropy ( Q ) + λ ∇ W IMM t ( Q ))

12: until convergence

C.1 k-approximated IMM

For 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 as

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

Algorithm 2 Serialized IMM with SGD for a Model Q with parameters W . As in Algorithm 1, this is notationally closest to the language modeling version. For implementation details for the logistic regression example, see D.5, as well as D.1 (for the soft nearest-neighbor density estimate formulation).

C.2 IMMGradient Computation

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

In summary, this allows us to compute the derivatives separately for each random sample and accumulate them appropriately.

Sequentializing IMM gradient computations using crosstalk

Recall that we can write an IMM component at data point t (Eq. (12)) with k randomly sampled extended contexts as:

where ⋆ 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.

We notationally move the derivative outside, but it is crucial that C t,i ( y ) = Q ( y | x t ,x i ) ∑ k i =1 Q ( y | x t ,x i ) 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.

Overcoming limitations of contemporary frameworks to compute the crosstalk vector

Wecannot 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 | 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 ) .

D Additional Experimental Details

We now provide some additional detail of both the logistic regression and language modeling experiments. One minor difference between the two 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 ).

D.1 Logistic Regression Toy Example

We provide the complete details of the logistic regression toy example 2 . 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.

Main question 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 δ x 2 ,t ,x 3 ,t ( x 2 , x 3 )e -α | x 1 ,t -x 1 | , 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 -α | x 1 ,t -x 1 | , the resulting induced model by marginalization is:

Just 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:

The overall objective then becomes

where λ is the regularization trade-off and can be determined via cross-validation (details in the Appendix D.1.1).

IMMimproves 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?

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

1

2

3

4

5

IMM(Q)

IMM improves full-context prediction 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.

x1

x2

0

0

0

1

1

1

-1

2

2

D.1.1 Tuning Lambda

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 . 0111 n +0 . 818 .

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

50

55

60

65

70

75

80

0.1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

45

50

55

60

65

70

75

80

1.0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

75

80

85

90

95

0.1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

75

80

85

90

95

1.0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

82

84

86

88

90

92

94

96

98

0.1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ

1+

λ

86

88

90

92

94

96

98

1.0

D.2 Language Modeling Experiments

Code is available at https://github.com/uicdice/imm-language-modeling

D.2.1 Evidence that IMM acts through improving restricted task

We give an illustrative example that shows the advantage that a restricted model may have, and which we would like to incorporate into the full-featured model. Clearly, our approach is only valuable if this the case. The simple experiment that we propose is to compare how well a full-featured model (LSTM) performs a restricted task, namely bigram prediction, vs. a restricted model built for that task (Kneser-Ney Bigram).

How can an LSTM perform bigram prediction when given only a single word? It can do so by averaging its output over all possible long histories consistent with that word, which is exactly the notion of induced model of Eq. (3). In Section 4, we show how to do this empirically.

When we do this comparison for the non-noised LSTM in (Xie et al., 2017), we obtain the results tabulated in Table 5. The Kneser-Ney bigram (restricted model) outperforms the LSTM (full model) on the restricted task. This may appear surprising, however, had this not been the case, the improvements that Xie et al. (2017) obtained through noising would have been hard to explain! More importantly, using IMM improves the performance of the LSTM on this task, mirroring the logistic regression case - cf. Figure 4.

D.2.2 Parameter details

Choice of λ For Language Modeling experiments, due to computational cost, we did not perform an exhaustive search for the best λ . We tried a few values of λ and tested on held-out data, and determined the value of 0 . 2 to work well in most cases, and adhered to it in all of the experiments.

Restarts We would like to clarify that the baseline for LSTM RNN (Xie et al., 2017) reports the best perplexity across multiple restarts while the baseline for BERT (Devlin et al., 2018) reports the average across multiple restarts. We do likewise in our experiments. We have additionally included error bars for BERT experiments as that's where most of the gain is seen for BERT.

Clipping On the technical front, we note that in either LSTM or BERT, if gradient clipping (Pascanu et al., 2013) is used, then gradients should be clipped separately for the primary objective and the IMM objective (i.e. they should be clipped before they are added together). The reason is that they are obtained from two loss functions that have very different scales.

Choice of k To reduce overhead for LSTM, we perform IMM only every j 'th batch using k -samples. Otherwise we use 1 -sample IMM. For BERT, we perform k -sample IMM every minibatch. For the selection of the parameter k (number of long histories used to approximate the induced model), we experimented with k ∈ { 1 , 5 , 10 , 20 } . For the LSTM experiments, there was not much gain going from 10 to 20, so we settled for k = 10 (and j = 5 ) as a good compromise between the accuracy and time complexity of the induced model computation. For the BERT experiments, we used k = 5 .

→

IMM through reintroduced MLM loss in BERT fine-tuning BERT on GLUE differs from LSTM on PTB in two ways: (1) it is not a next-word predictive model, and (2) its main loss during fine-tuning is the task (e.g., CLS) loss. We do not touch the latter. Instead, we reintroduce the MLM loss, to have a fine-tuning predictive task close to our formulation. Since MLM is based on cross-entropy, IMM can be directly added to it as a secondary loss. Additionally, only for the computation of the IMM component , at each location of a fine-tuning training sentence, we mask that word and all future words during the forward passes being done to compute the IMM component. This forces BERT to act like a causal model, i.e., it predicts the next word based only on the past, as in our main formulation. This is done to ensure that the induced model of BERT is the right induced model to match against an accurate causal model. Incidentally, for this accurate target model we used the Kneser-Ney bigram constructed from the Wikipedia data, on which BERT is initially trained.

In summary, in the BERT experiments we fine-tune with the sum of 3 losses: the task loss, the MLM cross-entropy loss, and the IMM risk for the MLM task.

D.3 RL Experiment

We consider an 11 × 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.

0

2

4

6

8

10

0

2

4

6

8

10

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

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.

Experimental Results 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 REINFORCE 3 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.

D.4 Effect of Restricted Model Quality

While the paper is presented under the assumption that we have access to a very good target model ˆ P , we also ran experiments to show the effect of model quality on IMM. Of course, for the language model, we don't 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.

0

10

20

30

40

50

50

60

70

80

90

100

IMM with high quality model

IMM with medium quality model

IMM with low quality model

No IMM

Dataset Size

All experiments in this section use the Sampled IMM variant. In Appendix D.5, we show that Serialized IMM achieves competitive statistical performance with

D.5 Serialized IMM Performance

To demonstrate the validity of the Serialized IMM variant, which comes with computational advantages that can allow IMM to scale, we implemented it for the logistic regression example. However, we first elaborate on the details of the Sampled IMM, to show the dramatic computational advantage. By combining the density estimation approach for obtaining the induced model, as outlined in Section 7.1 with the sequentialization technique explained in Appendix C.2, we obtain the following gradient

0

100

200

300

400

500

Number of epochs

0

5

10

15

20

Average Reward

IMM using maximal utility action

IMM using softmaxed POMDP policy (temperature=0.5)

IMM using softmaxed POMDP policy (temperature=1.0)

No IMM

calculation for the IMM risk (compare to Eq. (20)):

Here w t,t ′ uses the index of x 1 , rather than its value, as argument. By using a density estimator that looks at the entire dataset, we have a sampling scheme that effectively has set k = n , costing an n -fold computational time increase. While we can address this through shortcuts (e.g., sparsifying w ), this extreme version of computing the induced model is revealing in terms of where the bottleneck is, i.e., computing ˆ Q .

For the Serialized IMM variant, by comparing the above to Eq. (13), we see the large gain thanks to the double summation collapsing into a single summation. We still need to compute ˆ Q for the correction factor, but we can do this every n iterations. By incurring the O ( n ) slowdown only 1 n of the time, we go back to only a constant factor overhead. The additional noise in the gradient and the fact that ˆ Q is stale between updates mean that Serialized IMM cannot be expected to perform at the same level as Sampled IMM. Furthermore, the Serialized IMM does not use the smooth nearest-neighbor averaging, except for the calculation of Q . Despite this, we see in Figure 10 that Serialized IMM remains better than noising and, for larger dataset size, starts to approach Sampled IMM. Note that for ease of reproducibility, a fixed λ = 0 . 7 was used for these results and no particular optimization tweaks were made to compensate for gradient noise, suggesting that performance could be improved by making such changes.

0

10

20

30

40

50

50

60

70

80

90

100

Sampled IMM

Serialized IMM

Noising/Reverse-KD

No IMM or Noising

Dataset Size

E Potential Negative Societal Impact

Improving language models can have negative societal impact, if these language models are fine-tuned for tasks that are not aligned with positive human values. We hope that progress achieved by new methodologies such as IMM will not be put into misuse of this kind. On the contrary, we expect that any improved statistical efficiency achieved by methods like IMM to bring value when data is not as abundant, thus helping make machine learning and language modeling more impactful in underserved domains, where data is more scarce.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

IMM as a generalization of noising is analytically shown in the paper, along with noising's caveats that are solved by IMM. It's application is experimentally demonstrated in the paper, in Language Modeling and Reinforcement Learning experiments, as well as a proof of concept, using Logistic Regression.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

IMM, in some applications like noising, requires random sampling, because exact evaluation is not feasible. This affects the accuracy of the induced model, but regardless, IMM is able to outperform the baseline (i.e. noising).

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

The only proof is the one that discusses the caveat of noising through a counter example. All assumptions and a complete proof is present.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Code for all experiments has been anonymously released with clear instructions for reproducibility. A summary of relevant experimental details is also contained within the paper.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Code for all experiments has been anonymously released with clear instructions for reproducibility. A summary of relevant experimental details is also contained within the paper.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

A summary of relevant experimental details is also contained within the paper.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Error bars have been provided along with the number of Monte Carlo runs performed to obtain each bar.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

We mention the GPU used (Nvidia V100 with 32 GB memory)

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

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.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Appendix E addresses this.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

None of our datasets are scrapped. Pre-trained checkpoints are from reliable sources (Google trained checkpoint).

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

All our code (both original and derivative) is Apache 2.0 licensed and a copy of this license is included in all code repositories. Datasets are public domain.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

The only new asset is code. Code for all experiments has been anonymously released with clear instructions for reproducibility. A summary of relevant experimental details is also contained within the paper.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

No human subjects were involved.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\n- The relation of the proposed method to the previous work in noising is insightful and not immediately obvious, giving an interesting perspective.\n- The connection to knowledge distillation is valuable to highlight.\n- The code for the method is available.\n- The text is generally very well-written and explains the proposed method very well.\n\nWEAKNESSES:\n- The empirical results are somewhat dated. The models considered are quite old (RNNs, BERT) and it would be nice to see a more modern set of models such as decoder-only transformers or state-space models. However, using the old models does allow comparison to the noising approach. While the direct comparison to noising is appreciated, in practice most people would likely be comparing to the vanilla training approach as a baseline.\n\n- The main weakness of the paper is the computational overhead incurred by this method, which is not adequately addressed in the paper. With the introduction of language modeling, the question of how this method would scale to larger datasets and models is important, and the scaling properties do not appear favorable.\n- For the RNN experiment, each forward and backward pass seems to require k additional passes due to the gradient accumulation for the samples drawn for the model matching objective. This incurs a significant overhead factor of k.\n- The overhead factor may actually be worse when considering the need to replace the context for each of the k summands in the IMM objective. This requires swapping out the current context and computing a new set of hidden vectors, potentially up to the length of the input. A back-of-the-envelope calculation suggests the FLOPs could increase from O(L) to O(L + kL^2). However, there may be some subtlety I'm missing here that reduces the overhead.\n- For a large dataset, the lookup of the context could be quite difficult to implement. In the worst case, for n sequences of length L, storing the auxiliary data structure could require space on the order of nL^2, which is a factor of L larger than the original dataset. Supporting random access to this data structure may not be feasible for large corpora like Common Crawl without significant infrastructure overhead.\n\nQUESTIONS:\nMy questions all relate to my understanding of the limitations of the IMM approach.  I would be happy to reconsider my assessment if the above questions are addressed satisfactorily.\n\n1. Could the authors clarify the approximate computational and memory complexity of this method in terms of the number of sequences n, sequence length L, and any other relevant parameters? It would be helpful to see the analysis for both RNNs and Transformers.\n\n2. Could the authors comment on the asymptotic space complexity of the auxiliary data structure? For large corpora, it seems like this could require a huge amount of disk space and be challenging to query efficiently. Is there a data structure that could mitigate these concerns?\n\n3. Could the authors provide a FLOPs-corrected comparison between the IMM method and the vanilla training approach? Given the higher computational cost per forward and backward pass for IMM, it would be informative to see how the two methods compare when given a similar computational budget.\n\n4. Can the authors comment on practical cases where the IMM approach would be preferred over spending the additional $k$ FLOPS on e.g. processing $k$ more examples?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe proposed idea is interesting and it seems to be an improvement over the existing techniques such as noising and reverse-KD. I see value in being able to systematically incorporate restricted knowledge from a different distribution. Empirically IMM seems to be helping on a number of LM tasks.\n\nWEAKNESSES:\nThe paper adopts rather exotic notation that doesn't help reading the paper. I had to constantly look into the glossary in the appendix to go through equations. Descriptions of each experiment are also not self-contained, even when reading the appendix. The general idea of IMM is not the most intuitive to me and I think the paper would benefit from a simple graphical illustration of the principle. In my opinion, insufficient clarity in presenting the method and describing experiments is the main weakness of the paper. \n\nIf I understand the setup of LM experiments correctly (full and restricted models are trained on the same dataset) then the idea of the restricted model capturing a richer distribution is under-explored and the only way I can see IMM improving performance is as a regularizer (see my questions).\n\nQUESTIONS:\n1. I struggle to find a clear intuitive explanation for why induced model matching helps. Can authors think of a simple 2-dimensional logistic regression problem (for example) and include a graphical illustration of full and restricted models (and their induced variants)?\n\n2. Do I understand it right that the LM experiments the restricted n-gram model has been \"trained\" exactly on the same dataset on which the full model is being trained?\n2.1 If so, then, again, the only reason why IMM helps I can think of is some kind of \"smoothing\" or regularization of the full model, which probably learned an overly \"kinky\" distribution from limited data. In the eyes of the authors is that the right intuition?\n\n3. Is there an experiment where the restricted model indeed comes from a different and much larger dataset (which authors mention as indeed an interesting scenario for applying IMM)? It would be interesting to look at the \"trade value\" of a\"full\" data point compared to a\"restricted\" data point.\n\n4. Could authors hint at the scenarios in which IMM is not useful or even harmful? I would be especially interested if that could be the property of the model classes or data distributions from the full and the restricted model come.\n\nLIMITATIONS:\nAdequately discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* This is a very interesting topic, and seems like it could be useful in real-world scenarios. The idea appears novel and also is an interesting angle when compared to knowledge distillation. \n* It is very well-analyzed, and thoroughly explains how the method compares to existing methods. \n* The authors provide a clear presentation of their ideas as well as clear definitions. It is also nice to also have included a glossary in the appendix. Overall the paper is very clear to follow and precise about definitions.\n* The experiments show that, indeed, the proposed regularization term benefits the large model.\n\nWEAKNESSES:\n* As the authors mention, it is difficult to compute the regularization term. The computational cost is a drawback, although having increased performance is still a nice result. Given the computational cost, this brings into question when this method becomes useful in practice. I will leave my thoughts on this in the questions section.\n* It would be interesting to include further analysis regarding how the restricted model’s performance affects this method.\n\nQUESTIONS:\n* It might be realistic to assume that $\\hat{P}$ has some loss $\\varepsilon$. Do you have any insights  as to how one might expect IMM to degrade as a function of $\\varepsilon$? \n* Did you consider the setting in which there are many restricted feature models? How do you think your setting can be generalized even further? Do you think there are any settings in which using IMM regularization might hurt the model? For example, if the task is recalling something from a long time ago in the context, it might seem reasonable to assume this method would not work. Generally speaking, what kinds of problems do you think your method is best suited for? Information theoretically, there might be examples in which having some base solver might make the overall problem easier, but as mentioned, there might also be scenarios in which this method actually harms the training of the overall model. \n* Just for my understanding, are you assuming a fixed-context length model?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- Algorithm: This paper proposes a framework for how a good but restricted feature model, e.g. $\\bar{P}(y \\mid x_1)$, can be used as guidance when training a full-feature model, e.g. $Q(y \\mid x_1, x_2, x_3...)$. \n\\\nInstead of the knowledge distillation objective from weak teachers, which directly adds a regularizer comparing the probabilities of the restricted feature model \\bar{P} and full-feature Q with each other directly, IMM proposes comparing P with Q marginalized over the other features\n\\\n$$\\mathsf{Reverse-KL}: \\lambda \\sum_{y-} \\pi(x_1, x_2, x_3) \\sum \\bar{P}(y \\mid x_1) \\log \\frac{1}{Q(y \\mid x_1, x_2, x_3)}$$\n$$\\mathsf{IMM}: \\lambda \\sum_{y-} \\pi(x_1, x_2, x_3) \\sum \\bar{P}(y \\mid x_1) \\log \\frac{1}{\\hat{Q}(y \\mid x_1)} \\text{ where } \\\\ \n\\hat{Q}(y \\mid x_1) = \\sum_{x_2, x_3} \\pi(x_2, x_3 \\mid x_1) Q(y \\mid x_1, x_2, x_3)\n$$\n- The marginalization can be quite expensive in practice, but they demonstrate (in small RL and Language tasks) that they are able to approximate the marginalization efficiently by sampling. The single-sample IMM objective is connected to noising/reverse-KD. \n- They mathematically show that in the infinite data regime with the perfect true distribution $\\bar{P}(y | x_1)$ and the true model P(y \\mid x_1, x_2, x_3) is in the hypothesis class of Q, IMM objective always recovers P. On the other hand, reverse KL and noising cannot. \n- Empirically, IMM always does better than with no-IMM, but the performance difference is especially large in data-limited regimes.\n\nSTRENGTHS:\nWhile the objective may still be too computationally expensive for very large-scale tasks (i.e., sufficiently sampling a suffix tree, high inference cost), the paper is thought provoking, and well-written/motivated. I would recommend it for acceptance.\n\nWEAKNESSES:\n- One weakness of the paper is obviously the small scale of most of the experiments (learning policies over 11x11 grid, subset of GLUE, toy linear model). But the experiments sufficiently demonstrate the proof of concept. \n- It would also be nice to include some settings where the feature-restricted model is not close to optimal as an ablation study. \n- In Figure 1, why is noising able to achieve almost the same performance as IMM starting from around 5% dataset size, if noising is similar to  single-sample IMM? Generally, what is the importance of using more samples when marginalizing Q for improved performance? Does performance strictly increase with increasing samples?\n\nQUESTIONS:\n-  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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-16"}