🏆 KDD 2026 Think-like-LSTM (FraLLM)

This repository provides the official implementation of the KDD 2026 paper Think-like-LSTM: Memory-Augmented Large Language Models via Dynamic Fine-Tuning for Financial Risk Assessment.

[Paper (PDF)] [ACM Digital Library]

FraLLM is a memory-augmented large language model framework for learning from long-range temporal behavior. It follows a Think-like-LSTM design: sequence knowledge is learned dynamically across time steps, while a compact memory token recurrently carries historical information forward. This public release uses privacy-preserving, domain-neutral event representations.

The framework addresses two central challenges in long-range sequence modeling:

• Temporal dynamics: a Dynamic Fine-tuning Paradigm represents sanitized event sequences as timestamped text prototypes and trains the LLM with next-text prediction across consecutive time steps.
• Long-term dependencies: a Memory Token Mechanism recurrently aggregates historical prototypes into a continuously updated token, allowing the LLM to retain long-horizon behavior without repeatedly processing the entire history.
• Controlled public release: task-specific prompts, business schemas, and private preprocessing logic are replaced by generic interfaces without changing the memory-learning code path.

🧠 Method Overview

The public implementation processes each entity's sanitized event history as an ordered sequence of temporal chunks:

1. Privacy-reviewed preprocessing converts each event chunk into a generic textual prototype.
2. A predicting LLM learns to forecast the prototype of the next time step, providing linguistic supervision over the evolution of transaction patterns.
3. A gated Memory Token Adapter combines the previous memory, the latest response representation, and the elapsed-time encoding.
4. The updated memory token replaces a designated token embedding in the next LLM input, carrying long-term state through the sequence.
5. The latest memory token can be consumed by an authorized downstream task adapter, which is outside this public snapshot.

The released implementation contains the dynamic fine-tuning and memory-token components. In particular:

• models/MemoryLLM.py wraps a local causal LLM with LoRA adapters and injects the aligned memory token into its input embeddings.
• models/modules.py implements the gated memory updater, cosine time encoder, and memory/time merge layer.
• utils/DataLoader.py constructs chronological text inputs and token-level response masks for next-text prediction.
• dynamic_fine_tune.py trains the model chunk by chunk and saves memory-augmented checkpoints and loss traces.

🔒 Public Release and Privacy Notice

This repository demonstrates the model architecture and training flow without reproducing protected production details. In particular:

• The prompt in utils/DataLoader.py is a generic public template, not the prompt used in the private experimental or production environment.
• Raw records, original feature names, business rules, private preprocessing code, and serving configuration are intentionally omitted.
• Data adapters accept only sanitized or synthetic identifiers and text. Replace the public artifacts only with data approved for the target environment.
• All documented paths are relative to this repository root; no parent-directory layout is required or described.

The public loader uses Python pickle artifacts. Only load files from trusted sources, because unpickling untrusted content can execute arbitrary code.

🛠️ Installation

1. Create an environment

Python 3.10 is recommended. A CUDA-capable GPU is strongly recommended for the 7B model.

conda create -n frallm python=3.10 -y
conda activate frallm

Install a PyTorch build compatible with your CUDA environment, then install the remaining dependencies:

pip install torch
pip install transformers accelerate peft numpy pandas tqdm

For reproducible GPU installation, follow the platform-specific command provided by the PyTorch installation guide.

2. Prepare the backbone LLM

The code loads model weights with local_files_only=True. Place one of the supported Qwen2.5 checkpoints under LLMs/:

LLMs/
├── Qwen2.5-0.5B-Instruct/
└── Qwen2.5-7B-Instruct/

For example, the 7B checkpoint can be downloaded from Hugging Face:

hf download Qwen/Qwen2.5-7B-Instruct \
  --local-dir LLMs/Qwen2.5-7B-Instruct

Use the 0.5B model for a lower-resource smoke test and the 7B model for the released experiment configuration.

📦 Data Preparation

The private datasets used in the paper are not distributed in this repository. The public code expects synthetic or privacy-reviewed artifacts only.

The training entry point expects preprocessed artifacts in the following layout when commands are run from this directory:

.
├── FT_data/
│   ├── LLM_temporal_chain_data.pkl
│   ├── temporal_chain_text.pkl
│   ├── entity_text.pkl
│   └── event_text.pkl
├── dynamic_fine_tune.py
├── LLMs/
├── models/
└── utils/

LLM_temporal_chain_data.pkl contains sanitized temporal chains, while temporal_chain_text.pkl contains their target text prototypes. entity_text.pkl and event_text.pkl are generic ID-to-text mappings. The exact private preprocessing pipeline is intentionally not included; prepare compatible synthetic artifacts or implement an authorized adapter locally.

🚀 Dynamic Fine-Tuning

Run all commands from the repository root containing dynamic_fine_tune.py.

python dynamic_fine_tune.py \
  --dataset_name FT_data \
  --data_path ./FT_data \
  --code_path . \
  --LLM_model_name Qwen2.5-7B-Instruct

The script uses Hugging Face Accelerate for device preparation. For a customized distributed setup, configure Accelerate before launching:

accelerate config
accelerate launch dynamic_fine_tune.py \
  --dataset_name FT_data \
  --data_path ./FT_data \
  --code_path . \
  --LLM_model_name Qwen2.5-7B-Instruct

💾 Outputs

A training run writes the following artifacts:

logs/<model>/<dataset>/<timestamp>.log
saved_models/<model>/<dataset>/memoryLLM_<epoch>.pt
saved_loss/<model>/<dataset>/<epoch>.json

• Checkpoints contain the state dictionary of the memory-augmented LLM, including the LoRA and memory modules.
• Loss JSON files contain the average epoch loss and per-batch losses.
• Training logs record the resolved configuration, data size, parameter counts, progress, and runtime.

Note: dynamic_fine_tune.py clears saved_models/<model>/<dataset>/ at the beginning of each run. Move checkpoints that must be retained before launching another run with the same model and dataset names.

📊 Evaluation and Released Logs

The paper evaluates financial risk assessment with AUC, Kolmogorov-Smirnov (KS), and Recall@10. Please refer to the publication for the complete baseline, ablation, efficiency, and deployment results; this README intentionally avoids reproducing private experimental details.

Reference downstream FRA logs for Qwen2.5-7B-Instruct are included here:

Dataset	Log
FinTrans-S	logs/Qwen2.5-7B-instruct/FRA/FinTrans-S.ndjson
FinTrans-M	logs/Qwen2.5-7B-instruct/FRA/FinTrans-M.ndjson
FinTrans-L	logs/Qwen2.5-7B-instruct/FRA/FinTrans-L.ndjson

The current public snapshot includes the dynamic fine-tuning stage and reference FRA logs. It does not include private datasets, original prompt templates, private preprocessing, or a standalone downstream FRA training and evaluation entry point.

📁 Repository Structure

.
├── dynamic_fine_tune.py       # Dynamic, chunk-wise fine-tuning entry point
├── FT_data/                   # User-provided sanitized artifacts (not included)
├── models/
│   ├── MemoryLLM.py           # Memory-augmented causal LLM with LoRA
│   └── modules.py             # Memory updater, time encoder, and merge layer
├── utils/
│   ├── DataLoader.py          # Temporal text and token preparation
│   ├── load_configs.py        # Command-line configuration
│   └── utils.py               # Reproducibility and parameter utilities
├── imgs/
│   ├── comparison.png         # Think-like-LSTM motivation
│   └── method.png             # FraLLM framework
└── logs/
    └── Qwen2.5-7B-instruct/
        └── FRA/               # Released FinTrans-S/M/L experiment logs

📝 Citation

If you find this work useful, please cite:

@inproceedings{zhang2026thinklikelstm,
  title     = {Think-like-LSTM: Memory-Augmented Large Language Models via Dynamic Fine-Tuning for Financial Risk Assessment},
  author    = {Zhang, Siwei and Xiong, Yun and Chen, Xi and Tang, Yateng and Jia, Zi'an and Zheng, Xuehao and Xu, Jiarong},
  booktitle = {Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2},
  pages     = {8510--8521},
  year      = {2026},
  month     = aug,
  publisher = {Association for Computing Machinery},
  doi       = {10.1145/3770855.3818491},
  url       = {https://doi.org/10.1145/3770855.3818491}
}

🙏 Acknowledgements

This implementation is built with PyTorch, Hugging Face Transformers, PEFT, and Accelerate, using the Qwen2.5 family as the LLM backbone. We thank the maintainers and contributors of these projects.