iteryne 0.1.0

Model-Agnostic Meta-Learning (MAML) for any PyTorch nn.Module.

Iteryne

Model-Agnostic Meta-Learning (MAML) for any PyTorch nn.Module.

iteryne is a small, well-tested implementation of MAML (Finn, Abbeel & Levine, 2017) built on PyTorch's native torch.func. It meta-trains a model's initial parameters so that a few gradient steps on a new task's small support set generalize to that task's query set.

Features

• Model-agnostic. Works on any nn.Module (MLPs, CNNs, BatchNorm, ...) with no rewriting, via torch.func.functional_call.
• First- and second-order. Full MAML (gradient-through-gradient) and FOMAML (the first-order approximation) share one code path; flip first_order=True.
• Variants included. Meta-SGD (learnable per-parameter inner learning rates) and ANIL (adapt only the head).
• Two levels of API. A high-level MAML wrapper plus MetaTrainer, and an exposed functional core (adapt, inner_step, functional_forward).
• Typed and tested. py.typed, mypy-strict, pytest suite including a gradcheck of the second-order meta-gradient.

Install

pip install iteryne
# or, from a clone:
pip install -e ".[dev]"

Requires Python >= 3.10 and PyTorch >= 2.1.

Quickstart

import torch
from torch import nn
from iteryne import MAML, MetaTrainer, SinusoidTaskSampler

model = nn.Sequential(nn.Linear(1, 40), nn.ReLU(), nn.Linear(40, 40), nn.ReLU(), nn.Linear(40, 1))
maml = MAML(model, inner_lr=0.01, inner_steps=1, first_order=False)
meta_opt = torch.optim.Adam(maml.parameters(), lr=1e-3)

trainer = MetaTrainer(maml, meta_opt, nn.MSELoss(), SinusoidTaskSampler(seed=0))
trainer.fit(num_iterations=2000, meta_batch_size=25)

# Fast-adapt to a new task with a few gradient steps:
task = SinusoidTaskSampler(seed=123).sample(1)[0]
learner = maml.clone()
learner.adapt_on(nn.MSELoss(), task.support_x, task.support_y)
preds = learner(task.query_x)

Manual training loop

MetaTrainer is optional. The core pattern:

maml = MAML(model, inner_lr=0.01, inner_steps=5, first_order=False)
meta_opt = torch.optim.Adam(maml.parameters(), lr=1e-3)

meta_opt.zero_grad()
for task in task_batch:
    learner = maml.clone()
    for _ in range(maml.inner_steps):
        learner.adapt(loss_fn(learner(task.support_x), task.support_y))
    loss_fn(learner(task.query_x), task.query_y).backward()  # accumulates meta-grad
meta_opt.step()

The single backward() call computes the full second-order meta-gradient when first_order=False, and the first-order (FOMAML) meta-gradient when True.

Variants

from iteryne import MetaSGD, ANIL

meta_sgd = MetaSGD(model, inner_lr=0.01)            # learns per-parameter inner LRs
anil = ANIL(model, head=model[-1], inner_lr=0.01)   # adapts only the head

How it works

Parameters are carried as a dict[str, Tensor] and applied with functional_call, so the module itself never has to be modified. The inner step

p' = p - alpha * grad(support_loss, p)

is computed with torch.autograd.grad(..., create_graph=not first_order). With create_graph=True the adapted parameters stay connected to the originals, so a later backward() differentiates through the inner step (full MAML). With create_graph=False the gradient term is detached, leaving the identity edge p' = p - alpha * g so the same backward() yields the first-order meta-gradient (FOMAML). One flag, one code path.

BatchNorm and other buffers are carried through the functional forward but are not adapted in the inner loop; for few-shot work consider track_running_stats=False.

Documentation

Build the docs site locally:

pip install -e ".[docs]"
mkdocs serve

Citation

@inproceedings{finn2017maml,
  title     = {Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks},
  author    = {Finn, Chelsea and Abbeel, Pieter and Levine, Sergey},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2017}
}