ARMT-Gemma-3-1B-IT-Kestrel

An Associative Recurrent Memory Transformer (ARMT) built on Gemma-3-1B-IT and fine-tuned for multi-hop fact lookup over long fact lists. It is used in a course lab on the limits of retrieval-augmented generation (the "Kestrel handbook" lab): plain RAG cannot answer A→D questions whose intermediate entities the question never names, while ARMT reads the whole handbook segment by segment and carries the intermediate entities in its memory.

Task format

Question: In which town is the building where the leader of Project X works?
Facts:
<facts, one per line>
Question: In which town is the building where the leader of Project X works?
Answer: <leader>; <building>; <town>

The question comes first and is repeated at the end. For 2-hop and 3-hop questions the model writes the whole chain (leader; building or leader; building; town); the final element is the answer. Single-hop questions (In which year did Project X start?, Where does <person> work?, In which town is <building> located?, Who leads Project X?) are answered directly.

Training

• Base: irodkin/ARMT-Gemma-3-1B-IT-8x1024-portable (segment size 1024, 32 memory tokens per segment).
• LoRA (rank 128, all attention and MLP projections, merged into this checkpoint) plus full fine-tuning of the associative-memory parameters (W_mq, W_mk, W_mv, W_mb, memory tokens).
• Data: freshly generated random worlds with the same name pools as the Kestrel handbook but different links (the Kestrel handbook itself, seed 7, was never seen in training). Facts in random order, with the question's chain facts in hop order.
• Curriculum from 2 to 8 segments of 1024 tokens (up to ~615 facts), batch 32, bf16 autocast; the final stage used 8 segments and mostly 3-hop questions.

Results (exact match, greedy)

Context	1-hop	2-hop	3-hop
8 segments (600 random-order facts, random worlds)	1.00	0.96	0.84
Whole Kestrel handbook, 4,500 facts in sections (~54 segments)	0.92	0.88	0.68

50 questions per cell. The whole-handbook setting is about 7× longer than anything seen in training.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "irodkin/ARMT-Gemma-3-1B-IT-Kestrel"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, dtype=torch.float32).cuda().eval()

prompt = f"Question: {question}\nFacts:\n{facts}\nQuestion: {question}\nAnswer:"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
use_bf16 = torch.cuda.is_bf16_supported(including_emulation=False)
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_bf16):
    output = model.generate(input_ids, attention_mask=torch.ones_like(input_ids), max_new_tokens=32,
                            eos_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(output[0], skip_special_tokens=True))  # generate() returns only the new tokens

Notes:

• Precision: use bf16 autocast or float32. fp16 overflows (Gemma's hidden states reach ~4e4) and produces NaN once the memory is active.
• Long inputs: generate() processes the completed segments one at a time and decodes inside the last one, so it handles inputs far longer than the training length. For a plain forward pass over very long inputs set model.vertical_mode = True (segment-by-segment processing); vertical mode does not clear the memory at the end of forward, so call model.zero_mem() before the next call.
• Memory write strengths: every layer's W_mb module computes the write strength beta = sigmoid(W_mb(memory_token)) of each memory token after each segment; forward hooks on these modules show where in the input the model stores information.