Abstraction Augmented Continual Learning

What This Skill Does

Swap standard cross-entropy loss in continual learning with a dual-objective loss that simultaneously optimizes concrete instances and abstract versions (entity-masked or key-component-removed versions), enabling models to learn generalizable structure without maintaining large replay buffers.

The Component Swap

Old component: Standard supervised fine-tuning with only concrete instance labels, which causes entity-specific gradients to dominate and overwrites prior knowledge during streaming updates.

# Traditional continual learning: optimize only concrete instances
class StandardContinualLearning(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def training_step(self, batch):
        inputs, labels = batch  # Concrete instances only
        logits = self.model(inputs)
        loss = F.cross_entropy(logits, labels)
        return loss

New component: Dual-objective loss combining concrete and abstract objectives, with balanced weighting to prevent catastrophic forgetting.

class AbstractionAugmentedTraining(nn.Module):
    def __init__(self, model, abstraction_fn, alpha=0.5, use_replay=True):
        super().__init__()
        self.model = model
        self.abstraction_fn = abstraction_fn  # Entity masking, key removal, etc.
        self.alpha = alpha  # Weight for abstract vs. concrete
        self.use_replay = use_replay  # Optional: small replay buffer for stability

    def create_abstract_version(self, inputs):
        """
        Generate abstract version by removing entity-specific details.

        Examples:
        - Entity-masked: replace specific names with [ENTITY] tokens
        - Key-removal: remove non-structural information (dates, numbers)
        - Relational abstraction: keep only relationships, abstract entities
        """
        return self.abstraction_fn(inputs)

    def training_step(self, concrete_batch, abstract_batch=None, replay_batch=None):
        """
        Total loss combines:
        1. Concrete instance objective: optimize on actual data
        2. Abstract objective: generalize structural patterns
        3. Replay objective (optional): prevent catastrophic forgetting
        """
        inputs_concrete, labels_concrete = concrete_batch

        # Concrete instance objective
        logits_concrete = self.model(inputs_concrete)
        loss_concrete = F.cross_entropy(logits_concrete, labels_concrete)

        # Abstract objective: learn invariances
        if abstract_batch is None:
            # Auto-generate abstract version from concrete
            inputs_abstract = self.create_abstract_version(inputs_concrete)
            labels_abstract = labels_concrete  # Same semantic target
        else:
            inputs_abstract, labels_abstract = abstract_batch

        logits_abstract = self.model(inputs_abstract)
        loss_abstract = F.cross_entropy(logits_abstract, labels_abstract)

        # Combined loss: balance concrete and abstract
        loss_balanced = self.alpha * loss_concrete + (1 - self.alpha) * loss_abstract

        # Optional: small replay buffer prevents catastrophic forgetting
        if self.use_replay and replay_batch is not None:
            inputs_replay, labels_replay = replay_batch
            logits_replay = self.model(inputs_replay)
            loss_replay = F.cross_entropy(logits_replay, labels_replay)
            loss_balanced = loss_balanced + 0.1 * loss_replay

        return loss_balanced

    def on_new_stream_element(self, new_inputs, new_labels):
        """
        Called when a new data point arrives in the stream.
        Abstraction-Augmented Training pipeline:

        1. Create abstract version of new input
        2. Train on both concrete and abstract
        3. Store only a small fraction for later replay (optional)
        """
        # Create abstract version
        abstract_inputs = self.create_abstract_version(new_inputs)

        # Joint loss
        loss = self.training_step(
            concrete_batch=(new_inputs, new_labels),
            abstract_batch=(abstract_inputs, new_labels)
        )

        return loss

Key architectural difference: Instead of optimizing purely on concrete instances (which leads to overfitting to specific entities and catastrophic forgetting), the new loss jointly optimizes the model to predict correctly on both concrete examples AND their abstracted versions. This forces the model to learn the underlying relational structure, not just entity-specific patterns.

Performance Impact

Relational Cycle Benchmark (1.5B-7B parameter models):

On Qwen model:

• Cumulative accuracy improvement: +2.05 percentage points
• Unknown edge accuracy: +6.3 percentage points (major for generalization)
• Forgetting on unknown edges: 6.3% lower than best experience replay baseline
• Loss landscape variance: 17.0% reduction (more stable learning)

On SmolLM:

• Cumulative accuracy gain: +5.76 percentage points
• Unknown edge forgetting: 80% improvement (relative) over standard continual learning

Key insight: Improvements are larger on tasks with strong latent structure (relational data) and grow with model scale.

When to Use

• Continual learning on streaming data with relational structure
• Tasks where entities have meaningful relationships (knowledge graphs, narrative understanding, proverb interpretation)
• When catastrophic forgetting is problematic and replay buffers are memory-constrained
• Small-to-medium language models (1.5B-7B parameters; scalability to larger models TBD)
• Scenarios where abstract patterns (semantic, logical, narrative) are as important as concrete facts
• When you can systematically derive abstractions (entity masking, key removal, semantic categories)

When NOT to Use

• Tasks without latent relational structure (e.g., pixel classification)
• Large models (>10B parameters): scalability not demonstrated
• When concrete instance accuracy is more important than structural generalization
• If abstraction function cannot be reliably defined for your domain
• Real-time systems where the cost of generating abstractions is prohibitive
• Tasks where entity identity is critical (e.g., facial recognition, biometric verification)

Implementation Checklist

1. Define abstraction function:

# Example: entity masking for knowledge graph / dialogue
def entity_mask_abstraction(text):
    """Replace entity names with [ENTITY] tokens."""
    entities = extract_entities(text)  # Your entity extractor
    abstract_text = text
    for entity in entities:
        abstract_text = abstract_text.replace(entity, "[ENTITY]")
    return abstract_text

# Example: key removal for numerical reasoning
def numeric_abstraction(text):
    """Remove specific numbers, keep structure."""
    import re
    return re.sub(r'\d+', '[NUM]', text)

# Plug into AAT:
aat = AbstractionAugmentedTraining(
    model=your_model,
    abstraction_fn=entity_mask_abstraction,  # Or numeric_abstraction
    alpha=0.5
)

2. Integrate into training loop:

# Old training loop:
# for inputs, labels in stream:
#     logits = model(inputs)
#     loss = F.cross_entropy(logits, labels)
#     loss.backward()

# New training loop with AAT:
for inputs, labels in stream:
    loss = aat.training_step(
        concrete_batch=(inputs, labels)
    )
    loss.backward()
    optimizer.step()

3. Optional: lightweight replay buffer:

# Prevent catastrophic forgetting with minimal memory
replay_buffer = deque(maxlen=1000)  # Store ~1K examples, not 100K

for inputs, labels in stream:
    # Train on current + replay
    replay_batch = random.sample(list(replay_buffer), k=min(32, len(replay_buffer)))
    loss = aat.training_step(
        concrete_batch=(inputs, labels),
        replay_batch=replay_batch
    )
    loss.backward()

    # Add to replay
    replay_buffer.append((inputs, labels))

4. Verification:

• Measure cumulative accuracy on sequential evaluation
• Compare: baseline (no AAT) vs. AAT model
• Measure forgetting: accuracy on task N after learning tasks N+1, N+2, ...
• Ablate: remove abstract objective to confirm +2-5% delta

5. Hyperparameter tuning if needed:

• Alpha (concrete vs. abstract weight): start at 0.5; increase concrete weight if abstract obj hurts accuracy
• Replay buffer size: 10-50% of single task size (if using)
• Abstraction complexity: simple abstractions (entity mask) > complex ones (semantic decomposition)
• Learning rate: slightly higher than baseline (more gradients to optimize)

6. Known issues and workarounds:

• Abstraction quality matters: poorly defined abstractions hurt more than they help
  • Workaround: validate abstraction function on clean data first
• Model size: demonstrated on 1.5B-7B; larger models may not scale
  • Workaround: start small, verify on your model before scaling
• Streaming synchronization: if abstract version can't be generated in real-time
  • Workaround: pre-compute abstractions for common patterns
• Domain specificity: abstraction function is task-dependent
  • Workaround: develop abstraction function during task analysis phase
• Balance sensitivity: alpha=0.5 works for balanced tasks; adjust for imbalanced data

Related Work

Builds on experience replay (buffer-based continual learning) and auxiliary task learning (MTL). The dual-objective loss relates to invariance learning (data augmentation) and structural regularization (contrastive learning). Inspired by cognitive science: humans learn both concrete facts and abstract patterns; generalizable knowledge requires both.