Compact Memory for Continual Logistic Regression

AI summary

Overview

Research area: Continual learning (lifelong learning), with a focus on compact memory construction for logistic regression and shallow models.

Technical level: Advanced. The paper assumes familiarity with Hessian-based analysis, singular value decomposition, probabilistic PCA, and the Knowledge-adaptation prior (K-prior) framework.

Scope: The paper reformulates the search for an optimal memory set for continual logistic regression as Hessian matching, estimates that memory with probabilistic PCA, and evaluates it against Experience Replay on Split-ImageNet.

What This Paper Is About

Continual learning systems forget old tasks when new data arrives, and existing memory mechanisms — weight regularization and replay of raw old examples — still fall short of batch training, which requires keeping all data. The paper builds on the result of Khan and Swaroop (2021) that an optimal compact memory exists for logistic regression (its size depends on the rank of the feature matrix, not the number of data points), but whose construction requires access to the entire past dataset, which is impossible in a continual setting.

The goal is therefore to find a practical, incremental procedure that estimates such a compact memory from only the current task's data plus the previously stored memory.

Key Contributions

1. Hessian-matching reformulation. The authors reframe the search for the optimal memory as matching the Hessian of the K-prior to the Hessian of the accumulated loss, rather than matching gradients (which are too small a quantity to identify a larger memory matrix).

2. A PPCA/EM estimation algorithm for linear regression. They show that the Hessian-matching condition has the same form as maximum-likelihood estimation in probabilistic PCA (Tipping and Bishop, 1999b), yielding a one-line EM update (Algorithm 1) that produces the memory vectors and their weights.

3. An extension to binary logistic regression (Algorithm 2), where exact recursion is infeasible, so the requirement is relaxed to Hessian matching at the current parameter value. The differences from the linear case appear only in four lines and consist of terms involving the derivative of the prediction function, y'(f) = sigma(f)[1 - sigma(f)]. A brief extension to multi-class is discussed.

4. Empirical demonstration that compact memory beats replay by a wide margin on Split-ImageNet, with public code at https://github.com/team-approx-bayes/compact_memory_code.

Main Findings

• The optimal memory is theoretically exact but practically inaccessible. Khan and Swaroop (2021, App. A) show that setting the memory vectors to the left singular vectors of the feature matrix Phi_{1:t} lets the K-prior perfectly recover the full-batch gradient, with memory size equal to the rank K*_t. However, the optimal weights w*_{k|t} require a vector d_x of length equal to the entire data size |D_{1:t}|, so they cannot be computed when old features are unavailable.

• The linear-regression case is exactly solvable. Theorem 1 states that with theta_t <- theta*_t, U_t <- U*_t, and W_t <- S*_t, the K-prior is exactly equal to the accumulated loss and takes the quadratic form ½ (theta - theta_t)^T H(theta_t)(theta - theta_t), with Hessian H(theta_t) = Phi_{1:t} Phi_{1:t}^T + delta_t I.

• Hessian matching is equivalent to PPCA maximum likelihood. The condition C = T T^T / N, where C = U_{t+1} W_{t+1} U_{t+1}^T + epsilon_{t+1} I and T concatenates the new features with the rescaled old memory, is the same as the Hessian-matching equation; maximum-likelihood PPCA therefore performs Hessian matching and admits an EM implementation.

• For logistic regression the recursion is relaxed. Because a recursion like the linear-regression one (Eq. 14) is infeasible, the method instead matches Hessians at theta_{t+1} (Eq. 20), giving an equation of the same form as the linear case (Eq. 21) with different diagonal matrices W~_{t+1}, W~_t, and B_{t+1}.

• Large accuracy gains over Experience Replay on Split-ImageNet. With a memory size equivalent to 0.3% of the data size, the method reaches 60% accuracy versus 30% for replay.

• The batch gap nearly closes at a still-small memory. Increasing memory to 2% of the data size raises accuracy to 74%, compared to batch accuracy of 77.6% on the same task.

• Implementation details. The noise parameter epsilon_{t+1} is fixed to a constant rather than estimated for simplicity; Algorithm 1 assumes K_{t+1} = K_t, and growing the memory requires one extra initialization step, in practice using a subset of the new features.

• Reporting caveat. In the paper content available for this summary, Split-ImageNet is the only benchmark with reported accuracy numbers; other datasets, baselines, and model architectures used in the experiments are not reported in that text.

Methodology in Plain English

The starting point is the K-prior, a regularizer that keeps a new model close to the old one in two ways: directly in parameter space (a quadratic penalty pulling theta toward theta_t) and in prediction space (comparing the old and new model's predictions on a small set of stored "memory vectors" U_t, each with a weight w_t).

Ideal memory vectors and weights exist that would let this regularizer reproduce the full-batch gradient exactly, but computing the ideal weights needs statistics from every past example. The authors instead ask a different question: what memory would make the curvature (Hessian) of the regularizer match the curvature of the accumulated loss? For linear regression this matching condition turns out to be exactly the condition solved by probabilistic PCA. So they build a matrix T by stacking the new task's features alongside the old memory vectors scaled by the square root of their weights, and run the standard PPCA EM update on it. After convergence, the columns of the updated matrix are normalized to give new unit-norm memory vectors and their weights.

For logistic regression the same recipe applies with a correction: the prediction function's derivative y'(f) is folded into the relevant diagonal matrices, since curvature for logistic regression is weighted by sigma(f)(1 - sigma(f)) rather than being constant. The full procedure alternates between training the model on the new task plus the current K-prior, and updating the memory via this PPCA step.

Why This Matters

Impact on research. The paper supplies the first practical construction, to the authors' knowledge, of the optimal memory whose existence was proved by Khan and Swaroop (2021) but left without an estimation method. It reframes memory construction as a curvature-matching problem, connects it to classical PPCA and online SVD, and shows that a memory far smaller than the data can retain most of the benefit of batch training. This places a concrete, analyzable result in a setting that has resisted exact solutions even for one- or two-layer models.

Real-world applications (potential, as implied by the paper's framing of reducing batch-training cost):

• On-device or edge model updates where storing user data is infeasible and only a tiny memory budget is available.
• Streaming or sensor pipelines where tasks arrive continuously and old raw data cannot be retained or revisited.
• Privacy- or regulation-constrained deployments where replaying raw historical examples is undesirable and a compact derived memory is preferable.
• Distributed or federated settings where communication of raw data is expensive and a small memory summary is more practical.

Industry relevance. The authors motivate the work partly by the cost and environmental impact of batch training, which requires access to all data at all times. A demonstration that 2% of the data volume can recover most of batch accuracy points at substantial savings in storage, retraining cost, and data retention obligations, if the result transfers beyond logistic regression.

Future Directions

• Extension to deep learning. The authors explicitly frame the work as opening a direction for building compact memory that could be useful for continual deep learning in the future; whether the Hessian-matching and PPCA machinery scales to deep networks is untested in the available content.
• Extending beyond binary logistic regression. The paper states that the derivation extends straightforwardly to other generalized linear models and briefly discusses multi-class, but does not present full derivations or results for those settings in the available text.
• Closing the remaining accuracy gap. Accuracy reaches 74% with 2% memory against 77.6% for batch training, leaving a gap whose source and possible reduction are open questions.
• Refining the approximations. The method fixes the noise parameter epsilon to a constant even though PPCA can estimate it, and assumes a constant memory size across tasks; both choices are stated simplifications that could be revisited.

Target Audience

Researchers working on continual learning, especially those interested in theory-driven memory design, online second-order methods, and connections between probabilistic models (PPCA/SVD) and regularization-based approaches. It is also relevant to engineers implementing memory-constrained continual learning pipelines, and to readers already familiar with the K-prior framework who want a practical estimation algorithm for it. Readers without background in Hessian analysis or probabilistic PCA will find the derivations demanding despite the accessible high-level motivation.

Authors’ abstract

Despite recent progress, continual learning still does not match the performance of batch training. To avoid catastrophic forgetting, we need to build compact memory of essential past knowledge, but no clear solution has yet emerged, even for shallow neural networks with just one or two layers. In this paper, we present a new method to build compact memory for logistic regression. Our method is based on a result by Khan and Swaroop [2021] who show the existence of optimal memory for such models. We formulate the search for the optimal memory as Hessian-matching and propose a probabilistic PCA method to estimate them. Our approach can drastically improve accuracy compared to Experience Replay. For instance, on Split-ImageNet, we get 60% accuracy compared to 30% obtained by replay with memory-size equivalent to 0.3% of the data size. Increasing the memory size to 2% further boosts the accuracy to 74%, closing the gap to the batch accuracy of 77.6% on this task. Our work opens a new direction for building compact memory that can also be useful in the future for continual deep learning.